Initial commit
This commit is contained in:
Zhongwei Li
325
templates/mcp-with-workers-ai.ts
Normal file
325
templates/mcp-with-workers-ai.ts
Normal file
@@ -0,0 +1,325 @@
|
||||
/**
|
||||
* MCP Server with Workers AI Integration
|
||||
*
|
||||
* Demonstrates Workers AI integration for image and text generation.
|
||||
* Shows how to use AI binding in MCP tools.
|
||||
*
|
||||
* Based on: https://github.com/cloudflare/ai/tree/main/demos/remote-mcp-github-oauth
|
||||
*
|
||||
* ═══════════════════════════════════════════════════════════════
|
||||
* 🤖 WORKERS AI INTEGRATION
|
||||
* ═══════════════════════════════════════════════════════════════
|
||||
*
|
||||
* This template shows:
|
||||
* 1. AI binding configuration (wrangler.jsonc)
|
||||
* 2. Image generation with Flux
|
||||
* 3. Text generation with Llama
|
||||
* 4. Error handling for AI requests
|
||||
* 5. Streaming responses
|
||||
*
|
||||
* ═══════════════════════════════════════════════════════════════
|
||||
* 📋 REQUIRED CONFIGURATION
|
||||
* ═══════════════════════════════════════════════════════════════
|
||||
*
|
||||
* wrangler.jsonc:
|
||||
* {
|
||||
* "ai": {
|
||||
* "binding": "AI"
|
||||
* }
|
||||
* }
|
||||
*
|
||||
* No API keys required! Workers AI is built into Cloudflare Workers.
|
||||
*
|
||||
* Pricing: https://developers.cloudflare.com/workers-ai/platform/pricing/
|
||||
* - Free tier: 10,000 Neurons per day
|
||||
* - Image generation: ~500 Neurons per image
|
||||
* - Text generation: Varies by token count
|
||||
*
|
||||
* ═══════════════════════════════════════════════════════════════
|
||||
*/
|
||||
|
||||
import { McpAgent } from "agents/mcp";
|
||||
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
|
||||
import { z } from "zod";
|
||||
|
||||
type Env = {
|
||||
AI: Ai; // Workers AI binding (configured in wrangler.jsonc)
|
||||
};
|
||||
|
||||
/**
|
||||
* MCP Server with Workers AI tools
|
||||
*/
|
||||
export class MyMCP extends McpAgent<Env> {
|
||||
server = new McpServer({
|
||||
name: "Workers AI MCP Server",
|
||||
version: "1.0.0",
|
||||
});
|
||||
|
||||
async init() {
|
||||
// ═══════════════════════════════════════════════════════════════
|
||||
// TOOL 1: Generate Image with Flux
|
||||
// ═══════════════════════════════════════════════════════════════
|
||||
// Model: @cf/black-forest-labs/flux-1-schnell
|
||||
// Fast image generation model (2-4 seconds)
|
||||
// Input: Text prompt → Output: Base64-encoded PNG
|
||||
// ═══════════════════════════════════════════════════════════════
|
||||
this.server.tool(
|
||||
"generate_image",
|
||||
"Generate an image from a text prompt using Flux AI model",
|
||||
{
|
||||
prompt: z
|
||||
.string()
|
||||
.describe("Detailed description of the image to generate"),
|
||||
num_steps: z
|
||||
.number()
|
||||
.min(1)
|
||||
.max(8)
|
||||
.default(4)
|
||||
.optional()
|
||||
.describe("Number of inference steps (1-8, default 4). Higher = better quality but slower"),
|
||||
},
|
||||
async ({ prompt, num_steps = 4 }) => {
|
||||
try {
|
||||
// Call Workers AI
|
||||
const response = await this.env.AI.run(
|
||||
"@cf/black-forest-labs/flux-1-schnell",
|
||||
{
|
||||
prompt,
|
||||
num_steps,
|
||||
}
|
||||
);
|
||||
|
||||
// Response is a base64-encoded PNG image
|
||||
const imageBase64 = (response as { image: string }).image;
|
||||
|
||||
return {
|
||||
content: [
|
||||
{
|
||||
type: "text",
|
||||
text: `Generated image from prompt: "${prompt}"`,
|
||||
},
|
||||
{
|
||||
type: "image",
|
||||
data: imageBase64,
|
||||
mimeType: "image/png",
|
||||
},
|
||||
],
|
||||
};
|
||||
} catch (error) {
|
||||
return {
|
||||
content: [
|
||||
{
|
||||
type: "text",
|
||||
text: `Error generating image: ${error.message}`,
|
||||
},
|
||||
],
|
||||
isError: true,
|
||||
};
|
||||
}
|
||||
}
|
||||
);
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════
|
||||
// TOOL 2: Generate Text with Llama
|
||||
// ═══════════════════════════════════════════════════════════════
|
||||
// Model: @cf/meta/llama-3.1-8b-instruct
|
||||
// Fast text generation model
|
||||
// Input: User message → Output: AI-generated text
|
||||
// ═══════════════════════════════════════════════════════════════
|
||||
this.server.tool(
|
||||
"generate_text",
|
||||
"Generate text using Llama AI model",
|
||||
{
|
||||
prompt: z.string().describe("The prompt or question for the AI"),
|
||||
max_tokens: z
|
||||
.number()
|
||||
.min(1)
|
||||
.max(2048)
|
||||
.default(512)
|
||||
.optional()
|
||||
.describe("Maximum number of tokens to generate (default 512)"),
|
||||
},
|
||||
async ({ prompt, max_tokens = 512 }) => {
|
||||
try {
|
||||
// Call Workers AI
|
||||
const response = await this.env.AI.run(
|
||||
"@cf/meta/llama-3.1-8b-instruct",
|
||||
{
|
||||
messages: [
|
||||
{
|
||||
role: "user",
|
||||
content: prompt,
|
||||
},
|
||||
],
|
||||
max_tokens,
|
||||
}
|
||||
);
|
||||
|
||||
// Extract generated text
|
||||
const text = (response as { response: string }).response;
|
||||
|
||||
return {
|
||||
content: [
|
||||
{
|
||||
type: "text",
|
||||
text,
|
||||
},
|
||||
],
|
||||
};
|
||||
} catch (error) {
|
||||
return {
|
||||
content: [
|
||||
{
|
||||
type: "text",
|
||||
text: `Error generating text: ${error.message}`,
|
||||
},
|
||||
],
|
||||
isError: true,
|
||||
};
|
||||
}
|
||||
}
|
||||
);
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════
|
||||
// TOOL 3: List Available AI Models
|
||||
// ═══════════════════════════════════════════════════════════════
|
||||
// Shows all models available in Workers AI
|
||||
// Useful for discovering what's available
|
||||
// ═══════════════════════════════════════════════════════════════
|
||||
this.server.tool(
|
||||
"list_ai_models",
|
||||
"List all available Workers AI models",
|
||||
{},
|
||||
async () => {
|
||||
// This is a static list, but you could dynamically fetch from CF API
|
||||
const models = [
|
||||
{
|
||||
name: "@cf/black-forest-labs/flux-1-schnell",
|
||||
type: "Image Generation",
|
||||
description: "Fast image generation (2-4s)",
|
||||
},
|
||||
{
|
||||
name: "@cf/meta/llama-3.1-8b-instruct",
|
||||
type: "Text Generation",
|
||||
description: "Fast text generation",
|
||||
},
|
||||
{
|
||||
name: "@cf/meta/llama-3.1-70b-instruct",
|
||||
type: "Text Generation",
|
||||
description: "High-quality text generation (slower)",
|
||||
},
|
||||
{
|
||||
name: "@cf/openai/whisper",
|
||||
type: "Speech Recognition",
|
||||
description: "Audio transcription",
|
||||
},
|
||||
{
|
||||
name: "@cf/baai/bge-base-en-v1.5",
|
||||
type: "Text Embeddings",
|
||||
description: "Generate embeddings for semantic search",
|
||||
},
|
||||
];
|
||||
|
||||
return {
|
||||
content: [
|
||||
{
|
||||
type: "text",
|
||||
text: JSON.stringify(models, null, 2),
|
||||
},
|
||||
],
|
||||
};
|
||||
}
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Worker fetch handler
|
||||
*
|
||||
* ═══════════════════════════════════════════════════════════════
|
||||
* 🔧 CONFIGURATION NOTES
|
||||
* ═══════════════════════════════════════════════════════════════
|
||||
*
|
||||
* 1. AI Binding Setup (wrangler.jsonc):
|
||||
* {
|
||||
* "ai": {
|
||||
* "binding": "AI"
|
||||
* }
|
||||
* }
|
||||
*
|
||||
* 2. Deploy:
|
||||
* npx wrangler deploy
|
||||
*
|
||||
* 3. Test Tools:
|
||||
* - generate_image: Creates PNG images from prompts
|
||||
* - generate_text: Generates text responses
|
||||
* - list_ai_models: Shows available models
|
||||
*
|
||||
* 4. Client URL:
|
||||
* "url": "https://YOUR-WORKER.workers.dev/sse"
|
||||
*
|
||||
* ═══════════════════════════════════════════════════════════════
|
||||
*/
|
||||
export default {
|
||||
async fetch(
|
||||
request: Request,
|
||||
env: Env,
|
||||
ctx: ExecutionContext
|
||||
): Promise<Response> {
|
||||
const { pathname } = new URL(request.url);
|
||||
|
||||
// Handle CORS preflight
|
||||
if (request.method === "OPTIONS") {
|
||||
return new Response(null, {
|
||||
status: 204,
|
||||
headers: {
|
||||
"Access-Control-Allow-Origin": "*",
|
||||
"Access-Control-Allow-Methods": "GET, POST, OPTIONS",
|
||||
"Access-Control-Allow-Headers": "Content-Type, Authorization",
|
||||
"Access-Control-Max-Age": "86400",
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
// SSE transport
|
||||
if (pathname.startsWith("/sse")) {
|
||||
return MyMCP.serveSSE("/sse").fetch(request, env, ctx);
|
||||
}
|
||||
|
||||
// HTTP transport
|
||||
if (pathname.startsWith("/mcp")) {
|
||||
return MyMCP.serve("/mcp").fetch(request, env, ctx);
|
||||
}
|
||||
|
||||
// Health check with AI binding info
|
||||
if (pathname === "/" || pathname === "/health") {
|
||||
return new Response(
|
||||
JSON.stringify({
|
||||
name: "Workers AI MCP Server",
|
||||
version: "1.0.0",
|
||||
transports: {
|
||||
sse: "/sse",
|
||||
http: "/mcp",
|
||||
},
|
||||
features: {
|
||||
ai: !!env.AI,
|
||||
},
|
||||
models: [
|
||||
"flux-1-schnell (image generation)",
|
||||
"llama-3.1-8b-instruct (text generation)",
|
||||
],
|
||||
status: "ok",
|
||||
timestamp: new Date().toISOString(),
|
||||
}),
|
||||
{
|
||||
headers: {
|
||||
"Content-Type": "application/json",
|
||||
"Access-Control-Allow-Origin": "*",
|
||||
},
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
return new Response("Not Found", { status: 404 });
|
||||
},
|
||||
};
|
||||
Reference in New Issue
Block a user