|
| 1 | +import { z } from 'zod' |
| 2 | +import type { CloudflareDocumentationMCP } from '../index' |
| 3 | +import { type EmbeddedResource } from "@modelcontextprotocol/sdk/types.js" |
| 4 | +import mime from "mime" |
| 5 | +/** |
| 6 | + * Registers the docs search tool with the MCP server |
| 7 | + * @param agent The MCP server instance |
| 8 | + */ |
| 9 | +export function registerDocsTools(agent: CloudflareDocumentationMCP) { |
| 10 | + // Register the worker logs analysis tool by worker name |
| 11 | + agent.server.tool( |
| 12 | + 'search_cloudflare_documentation', |
| 13 | + `Search the Cloudflare documentation. |
| 14 | + |
| 15 | + You should use this tool when: |
| 16 | + - A user asks questions about Cloudflare products (Workers, Developer Platform, Zero Trust, CDN, etc) |
| 17 | + - A user requests information about a Cloudflare feature |
| 18 | + - You are unsure of how to use some Cloudflare functionality |
| 19 | + - You are writing Cloudflare Workers code and need to look up Workers-specific documentation |
| 20 | +
|
| 21 | + This tool returns a number of results from a vector database. These are embedded as resources in the response and are plaintext doucments in a variety of formats. |
| 22 | + `, |
| 23 | + { |
| 24 | + // partially pulled from autorag query optimization example |
| 25 | + query: z.string().describe(`Search query. The query should: |
| 26 | +1. Identify the core concepts and intent |
| 27 | +2. Add relevant synonyms and related terms |
| 28 | +3. Remove irrelevant filler words |
| 29 | +4. Structure the query to emphasize key terms |
| 30 | +5. Include technical or domain-specific terminology if applicable`), |
| 31 | + scoreThreshold: z.number().min(0).max(1).optional().describe("A score threshold (0-1) for which matches should be included."), |
| 32 | + maxNumResults: z.number().default(10).optional().describe("The maximum number of results to return.") |
| 33 | + }, |
| 34 | + async (params) => { |
| 35 | + // we don't need "rewrite query" OR aiSearch because an LLM writes the query and formats the output for us. |
| 36 | + const result = await agent.env.AI.autorag(agent.env.AUTORAG_NAME).search({ |
| 37 | + query: params.query, |
| 38 | + ranking_options: params.scoreThreshold ? { |
| 39 | + score_threshold: params.scoreThreshold |
| 40 | + } : undefined, |
| 41 | + max_num_results: params.maxNumResults |
| 42 | + }) |
| 43 | + |
| 44 | + const resources: EmbeddedResource[] = result.data.map((result) => { |
| 45 | + const content = result.content.reduce((acc, contentPart) => { |
| 46 | + return acc + contentPart.text |
| 47 | + }, "") |
| 48 | + return { |
| 49 | + type: "resource", |
| 50 | + resource: { |
| 51 | + uri: `docs://${result.filename}`, |
| 52 | + mimeType: mime.getType(result.filename) ?? "text/plain", |
| 53 | + text: content |
| 54 | + } |
| 55 | + } |
| 56 | + }) |
| 57 | + |
| 58 | + return { |
| 59 | + content: resources |
| 60 | + } |
| 61 | + } |
| 62 | + ) |
| 63 | +} |
0 commit comments