more effect and stemming

Author: mattzcarey

site/agents-mcp/src/index.ts

@@ -4,6 +4,7 @@ import { z } from "zod";
 import { createMcpHandler } from "agents/mcp";
 import { fetchAndBuildIndex, formatResults } from "./utils";
 import { search } from "@orama/orama";
+import { Effect } from "effect";
 
 // TODO: instrument this server for observability
 const mcpServer = new McpServer(
@@ -18,8 +19,12 @@ const mcpServer = new McpServer(
 );
 
 const inputSchema = {
-  query: z.string(),
-  k: z.number().optional().default(10)
+  query: z
+    .string()
+    .describe(
+      "query string to search for eg. 'agent hibernate', 'schedule tasks'"
+    ),
+  k: z.number().optional().default(5).describe("number of results to return")
 };
 
 mcpServer.registerTool(
@@ -28,37 +33,45 @@ mcpServer.registerTool(
     inputSchema
   },
   async ({ query, k }) => {
-    try {
+    const searchEffect = Effect.gen(function* () {
       console.log({ query, k });
-      const docsDb = await fetchAndBuildIndex();
-      const results = await search(docsDb, {
-        term: query,
-        limit: k
-      });
-      return {
-        content: [
-          {
-            type: "text",
-            text: formatResults(results, query, k)
-          }
-        ]
-      };
-    } catch (error) {
-      console.error(error);
+      const term = query.trim();
+
+      const docsDb = yield* fetchAndBuildIndex;
+
+      const result = search(docsDb, { term, limit: k });
+      const searchResult = yield* result instanceof Promise
+        ? Effect.promise(() => result)
+        : Effect.succeed(result);
+
       return {
         content: [
           {
-            type: "text",
-            text: `There was an error with the search tool. Please try again later.`
+            type: "text" as const,
+            text: formatResults(searchResult, term, k)
           }
         ]
       };
-    }
+    }).pipe(
+      Effect.catchAll((error) => {
+        console.error(error);
+        return Effect.succeed({
+          content: [
+            {
+              type: "text" as const,
+              text: `There was an error with the search tool. Please try again later.`
+            }
+          ]
+        });
+      })
+    );
+
+    return await Effect.runPromise(searchEffect);
   }
 );
 
 export default {
   async fetch(request: Request, env: Env, ctx: ExecutionContext) {
-    return createMcpHandler(mcpServer as any)(request, env, ctx);
+    return createMcpHandler(mcpServer)(request, env, ctx);
   }
 };

site/agents-mcp/src/utils.ts

@@ -20,6 +20,13 @@ const KV_KEY = "docs-v0";
 const DOCS_REPO_API =
   "https://api.github.com/repos/cloudflare/agents/git/trees/main?recursive=1";
 const TTL_SECONDS = 24 * 60 * 60; // 1 day
+const CHUNK_SIZE = 2000;
+const MIN_CHARS_PER_CHUNK = 500;
+
+const chunker = await RecursiveChunker.create({
+  chunkSize: CHUNK_SIZE,
+  minCharactersPerChunk: MIN_CHARS_PER_CHUNK
+});
 
 const fetchWithRetry = (url: string, useAuth = true) =>
   Effect.tryPromise({
@@ -57,8 +64,8 @@ const fetchWithRetry = (url: string, useAuth = true) =>
     )
   );
 
-const fetchDocsFromGitHub = async (): Promise<Document[]> => {
-  const treeEffect = fetchWithRetry(DOCS_REPO_API).pipe(
+const fetchDocsFromGitHub = Effect.gen(function* () {
+  const treeData = yield* fetchWithRetry(DOCS_REPO_API).pipe(
     Effect.flatMap((response) =>
       Effect.tryPromise({
         try: async () => {
@@ -73,89 +80,90 @@ const fetchDocsFromGitHub = async (): Promise<Document[]> => {
     )
   );
 
-  const treeData = await Effect.runPromise(treeEffect);
-
   const docFiles = treeData.tree.filter(
     (item: GitHubTreeItem) =>
       item.path.startsWith("docs/") && item.path.endsWith(".md")
   );
 
   const docs: Document[] = [];
-  const chunker = await RecursiveChunker.create({
-    chunkSize: 800
-  });
 
   for (const file of docFiles) {
     const contentUrl = `https://raw.githubusercontent.com/cloudflare/agents/main/${file.path}`;
 
-    const contentEffect = fetchWithRetry(contentUrl).pipe(
+    const contentResult = yield* fetchWithRetry(contentUrl).pipe(
       Effect.flatMap((response) =>
         Effect.tryPromise({
           try: () => response.text(),
           catch: (error) => error as Error
         })
-      )
+      ),
+      Effect.flatMap((content) => Effect.promise(() => chunker.chunk(content))),
+      Effect.catchAll((error) => {
+        console.error(`Failed to fetch/chunk ${file.path}:`, error);
+        return Effect.succeed([]);
+      })
     );
 
-    try {
-      const content = await Effect.runPromise(contentEffect);
-      const chunks = await chunker.chunk(content);
-
-      for (const chunk of chunks) {
-        docs.push({
-          fileName: file.path,
-          content: chunk.text,
-          url: contentUrl
-        });
-      }
-    } catch (error) {
-      console.error(`Failed to fetch/chunk ${file.path}:`, error);
+    for (const chunk of contentResult) {
+      docs.push({
+        fileName: file.path,
+        content: chunk.text,
+        url: contentUrl
+      });
     }
   }
 
   return docs;
-};
-
-const getCachedDocs = async (): Promise<Document[] | null> => {
-  const cached = await env.DOCS_KV.get(KV_KEY, "json");
-
-  if (!cached) {
-    return null;
-  }
+});
 
-  return cached as Document[];
-};
+const getCachedDocs = Effect.tryPromise({
+  try: () => env.DOCS_KV.get(KV_KEY, "json") as Promise<Document[] | null>,
+  catch: (error) => error as Error
+});
 
-const cacheDocs = async (docs: Document[]): Promise<void> => {
-  await env.DOCS_KV.put(KV_KEY, JSON.stringify(docs), {
-    expirationTtl: TTL_SECONDS
+const cacheDocs = (docs: Document[]) =>
+  Effect.tryPromise({
+    try: () =>
+      env.DOCS_KV.put(KV_KEY, JSON.stringify(docs), {
+        expirationTtl: TTL_SECONDS
+      }),
+    catch: (error) => error as Error
   });
-};
 
-export const fetchAndBuildIndex = async () => {
-  let docs = await getCachedDocs();
+export const fetchAndBuildIndex = Effect.gen(function* () {
+  const cached = yield* getCachedDocs;
+
+  let docs: Document[];
 
-  if (!docs) {
-    // If not cached, fetch from GitHub, chunk, and cache to KV
-    docs = await fetchDocsFromGitHub();
-    await cacheDocs(docs);
+  if (!cached) {
+    docs = yield* fetchDocsFromGitHub;
+    yield* cacheDocs(docs);
+  } else {
+    docs = cached;
   }
 
-  // Build the search index from docs
-  const docsDb = create({
-    schema: {
-      fileName: "string",
-      content: "string",
-      url: "string"
-    } as const
-  });
+  const docsDb = yield* Effect.sync(() =>
+    create({
+      schema: {
+        fileName: "string",
+        content: "string",
+        url: "string"
+      } as const,
+      components: {
+        tokenizer: {
+          stemming: true,
+          language: "english"
+        }
+      }
+    })
+  );
 
   for (const doc of docs) {
-    await insert(docsDb, doc);
+    yield* Effect.sync(() => insert(docsDb, doc));
   }
 
   return docsDb;
-};
+});
 
 export const formatResults = (
   results: Awaited<ReturnType<typeof search>>,
@@ -167,7 +175,13 @@ export const formatResults = (
 
   let output = `**Search Results**\n\n`;
   output += `Found ${hitCount} result${hitCount !== 1 ? "s" : ""} for "${query}" (${elapsed})\n\n`;
-  output += `Showing top ${k} result${k !== 1 ? "s" : ""}:\n\n`;
+
+  if (hitCount === 0) {
+    output += `No results found. Try using different keywords or modify the spelling.`;
+    return output;
+  }
+
+  output += `Showing top ${Math.min(k, hitCount)} result${Math.min(k, hitCount) !== 1 ? "s" : ""}:\n\n`;
   output += `---\n\n`;
 
   for (const hit of results.hits) {