@mammothb/pi-websearch 0.1.0 → 0.2.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@mammothb/pi-websearch",
-  "version": "0.1.0",
+  "version": "0.2.0",
   "description": "A pi extension that adds a websearch tool for searching the web",
   "keywords": [
     "pi-package"

package/src/lib/mcp-websearch.ts

@@ -0,0 +1,64 @@
+import { parseResponse } from "./parsers";
+import type { SearchArgs } from "./types";
+
+const URL = "https://mcp.exa.ai/mcp";
+const TOOL = "web_search_exa";
+
+function buildMcpRequest(toolName: string, args: SearchArgs) {
+  const value = Object.fromEntries(
+    Object.entries(args).filter(([_, v]) => v !== undefined),
+  );
+  return {
+    jsonrpc: "2.0" as const,
+    id: 1,
+    method: "tools/call" as const,
+    params: { name: toolName, arguments: value },
+  };
+}
+
+export async function call(
+  args: SearchArgs,
+  signal: AbortSignal | undefined,
+  timeoutMs: number,
+): Promise<string | undefined> {
+  const controller = new AbortController();
+  const timeoutId = setTimeout(() => controller.abort(), timeoutMs);
+  // Forward external signal
+  const onAbort = () => controller.abort();
+  if (signal) {
+    if (signal.aborted) {
+      throw new Error("Request aborted");
+    }
+    signal.addEventListener("abort", onAbort, { once: true });
+  }
+
+  try {
+    const response = await fetch(URL, {
+      method: "POST",
+      headers: {
+        "Content-Type": "application/json",
+        Accept: "application/json, text/event-stream",
+      },
+      body: JSON.stringify(buildMcpRequest(TOOL, args)),
+      signal: controller.signal,
+    });
+
+    if (!response.ok) {
+      throw new Error(
+        `Exa MCP returned HTTP ${response.status}: ${response.statusText}`,
+      );
+    }
+
+    return parseResponse(await response.text());
+  } catch (error) {
+    if (controller.signal.aborted && !signal?.aborted) {
+      throw new Error("Request timed out");
+    }
+    throw error;
+  } finally {
+    clearTimeout(timeoutId);
+    if (signal) {
+      signal.removeEventListener("abort", onAbort);
+    }
+  }
+}

package/src/lib/parsers.ts

@@ -0,0 +1,42 @@
+import { Value } from "typebox/value";
+import { McpResultPayload } from "./types";
+
+/**
+ * Try to parse a JSON object from a string, returning the first text content.
+ */
+function tryParsePayload(payload: string): string | undefined {
+  const trimmed = payload.trim();
+  if (!trimmed.startsWith("{")) return undefined;
+  try {
+    const data = Value.Parse(McpResultPayload, JSON.parse(trimmed));
+    return data.result.content.find((item) => item.text)?.text;
+  } catch {
+    return undefined;
+  }
+}
+
+/** Parse an MCP response body, handling both plain JSON and SSE streams. */
+export function parseResponse(body: string): string | undefined {
+  const trimmed = body.trim();
+
+  // Try direct JSON parse first
+  if (trimmed) {
+    const direct = tryParsePayload(trimmed);
+    if (direct) {
+      return direct;
+    }
+  }
+
+  // Try SSE lines: "data: {...}"
+  for (const line of body.split("\n")) {
+    if (!line.startsWith("data: ")) {
+      continue;
+    }
+    const data = tryParsePayload(line.slice(6));
+    if (data) {
+      return data;
+    }
+  }
+
+  return undefined;
+}

package/src/lib/types.ts

@@ -0,0 +1,43 @@
+import { StringEnum } from "@earendil-works/pi-ai";
+import type { Static } from "typebox";
+import Type from "typebox";
+
+export const WebsearchParameters = Type.Object({
+  query: Type.String({ description: "Web search query" }),
+  numResults: Type.Optional(
+    Type.Number({
+      description: "Number of search results to return (default: 8)",
+    }),
+  ),
+  livecrawl: Type.Optional(
+    StringEnum(["fallback", "preferred"] as const, {
+      description:
+        "Live crawl mode - 'fallback': use live crawling as backup if cached content unavailable, 'preferred': prioritize live crawling (default: 'fallback')",
+    }),
+  ),
+  type: Type.Optional(
+    StringEnum(["auto", "fast", "deep"] as const, {
+      description:
+        "Search type - 'auto': balanced search, 'fast': quick results, 'deep': comprehensive search (default: 'auto')",
+    }),
+  ),
+  contextMaxCharacters: Type.Optional(
+    Type.Number({
+      description:
+        "Maximum characters for context string optimized for LLMs (default: 10000)",
+    }),
+  ),
+});
+
+export type SearchArgs = Static<typeof WebsearchParameters>;
+
+export const McpResultPayload = Type.Object({
+  result: Type.Object({
+    content: Type.Array(
+      Type.Object({
+        type: Type.String(),
+        text: Type.String(),
+      }),
+    ),
+  }),
+});

package/src/websearch.ts

@@ -1,22 +1,80 @@
-import type { ToolDefinition } from "@earendil-works/pi-coding-agent";
-import Type from "typebox";
+import type { ImageContent, TextContent } from "@earendil-works/pi-ai";
+import {
+  keyText,
+  type Theme,
+  type ToolDefinition,
+} from "@earendil-works/pi-coding-agent";
+import { Container, Spacer, Text } from "@earendil-works/pi-tui";
+import { call } from "./lib/mcp-websearch";
+import { type SearchArgs, WebsearchParameters } from "./lib/types";
+
+const COLLAPSED_PREVIEW_LINES = 7;
+const TIMEOUT_MS = 25_000;
 
 interface WebsearchDetails {
   query: string;
-  numResults: number;
 }
 
-const Parameters = Type.Object({
-  query: Type.String({ description: "Search query" }),
-  numResults: Type.Optional(
-    Type.Number({ description: "Number of results (default: 8)" }),
-  ),
-});
+function isTextContent(c: TextContent | ImageContent): c is TextContent {
+  return c.type === "text";
+}
+
+function renderExpandableResult(
+  details: WebsearchDetails,
+  textContent: string,
+  expanded: boolean,
+  theme: Theme,
+): Container {
+  const container = new Container();
+
+  const title = details.query;
+  container.addChild(
+    new Text(
+      theme.fg("syntaxKeyword", "query: ") + theme.fg("syntaxString", title),
+    ),
+  );
+
+  if (!textContent) {
+    return container;
+  }
+
+  container.addChild(new Spacer(1));
+
+  if (expanded) {
+    container.addChild(new Text(textContent));
+  } else {
+    const lines = textContent
+      .split("\n")
+      .filter(
+        (line, index, arr) =>
+          line.length > 0 || index === 0 || index < arr.length - 1,
+      );
+    const previewLines = lines.slice(0, COLLAPSED_PREVIEW_LINES);
+    const remaining = Math.max(0, lines.length - previewLines.length);
+
+    const preview = previewLines.join("\n");
+    container.addChild(new Text(preview));
+
+    if (remaining > 0) {
+      container.addChild(new Spacer(1));
+      const expandKey = keyText("app.tools.expand") || "Ctrl+O";
+      container.addChild(
+        new Text(
+          theme.fg("muted", `... (${remaining} more lines, `) +
+            theme.fg("dim", expandKey) +
+            theme.fg("muted", " to expand)"),
+        ),
+      );
+    }
+  }
+  return container;
+}
 
 export function createWebsearchTool(): ToolDefinition<
-  typeof Parameters,
+  typeof WebsearchParameters,
   WebsearchDetails
 > {
+  const year = new Date().getFullYear();
   return {
     name: "websearch",
     label: "Web Search",
@@ -31,20 +89,61 @@ Usage notes:
   - Search types when available: 'auto' (balanced), 'fast' (quick results), 'deep' (comprehensive search)
   - Configurable context length for optimal LLM integration`,
     promptSnippet: "Search the web",
-    parameters: Parameters,
-    async execute(_toolCallId, params, _signal, _onUpdate, _ctx) {
-      const query = params.query;
-      const numResults = params.numResults ?? 8;
-
-      return {
-        content: [
-          {
-            type: "text",
-            text: `Search results for: ${query} (${numResults} results requested)`,
-          },
-        ],
-        details: { query, numResults },
+    promptGuidelines: [
+      "Use websearch to find current information, documentation, or answers that require up-to-date web data. Always cite sources from search results.",
+      `The current year is ${year}. You MUST use this year when searching for recent information or current events.\n- Example: If the current year is ${year} and the user asks for "latest AI news", search for "AI news ${year}", NOT "AI news ${year - 1}"`,
+    ],
+    parameters: WebsearchParameters,
+    async execute(_toolCallId, params, signal, _onUpdate, _ctx) {
+      const args: SearchArgs = {
+        query: params.query,
+        type: params.type ?? "auto",
+        numResults: params.numResults ?? 8,
+        livecrawl: params.livecrawl ?? "fallback",
+        contextMaxCharacters: params.contextMaxCharacters,
       };
+
+      try {
+        const result = await call(args, signal, TIMEOUT_MS);
+
+        return {
+          content: [
+            {
+              type: "text",
+              text:
+                result ??
+                "No search results found. Please try a different query.",
+            },
+          ],
+          details: { query: params.query },
+        };
+      } catch (err) {
+        const message = err instanceof Error ? err.message : String(err);
+        throw new Error(`Web search failed: ${message}`);
+      }
+    },
+    renderCall(args, theme, _context) {
+      return new Text(
+        theme.fg("toolTitle", theme.bold("websearch ")) +
+          theme.fg("muted", `"${args.query}"`),
+      );
+    },
+    renderResult(result, options, theme, _context) {
+      if (options.isPartial) {
+        return new Text(theme.fg("warning", "Searching..."));
+      }
+
+      const textContent = result.content
+        .filter(isTextContent)
+        .map((c) => c.text)
+        .join("\n");
+
+      return renderExpandableResult(
+        result.details,
+        textContent,
+        options.expanded,
+        theme,
+      );
     },
   };
 }