pi-config/extensions/tavily-search/index.ts

/**
 * @jmcombs/pi-tavily-search — Real-time web search for the Pi coding agent.
 *
 * Registers a `tavily_search` tool that the LLM can call to perform a Tavily
 * web search. The Tavily API key is resolved from (in order):
 *   1. `AuthStorage` under the "tavily" key (`~/.pi/agent/auth.json`)
 *   2. The `TAVILY_API_KEY` environment variable
 *
 * See README.md for configuration details and recommended secret-storage
 * patterns (env var, plain auth.json, or shell-resolved 1Password / Keychain).
 */

import { type ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { readFileSync, existsSync } from "fs";
import { join } from "path";
import { Type, type Static } from "typebox";

function getTavilyKey(): string | undefined {
  // 1. Try env var
  const envKey = process.env.TAVILY_API_KEY;
  if (envKey) return envKey;
  // 2. Try auth.json (works with both old and new pi)
  try {
    const authPath = join(process.env.HOME || "", ".pi/agent/auth.json");
    if (existsSync(authPath)) {
      const auth = JSON.parse(readFileSync(authPath, "utf-8"));
      const tavily = auth["tavily"];
      if (tavily?.type === "api_key") return tavily.key;
    }
  } catch {}
  return undefined;
}

const TAVILY_SEARCH_ENDPOINT = "https://api.tavily.com/search";

// ── Tool parameter schema ──────────────────────────────────────────────

const tavilySearchSchema = Type.Object({
  query: Type.String({
    description: "The search query to perform.",
    minLength: 1,
  }),
});

export type TavilySearchInput = Static<typeof tavilySearchSchema>;

// ── Tavily API response types ──────────────────────────────────────────
//
// Documented at https://docs.tavily.com/documentation/api-reference/endpoint/search
// We model only the fields we actually consume; unknown fields pass through
// untouched in the `details.raw` field returned by the tool.

interface TavilySearchResult {
  title: string;
  url: string;
  content: string;
  score?: number;
  raw_content?: string | null;
}

interface TavilySearchResponse {
  query?: string;
  answer?: string;
  results?: TavilySearchResult[];
}

// ── Helpers ────────────────────────────────────────────────────────────

const MISSING_KEY_MESSAGE = [
  "Error: No Tavily API key configured.",
  "",
  "Configure one of the following:",
  "  • Environment variable: export TAVILY_API_KEY=<your-key>",
  '  • ~/.pi/agent/auth.json: { "tavily": { "type": "api_key", "key": "<your-key>" } }',
  '  • Shell-resolved: { "tavily": { "type": "api_key", "key": "!security find-generic-password -ws tavily" } }',
  '  • 1Password: { "tavily": { "type": "api_key", "key": "!op read \'op://Personal/tavily/credential\'" } }',
].join("\n");

function formatResults(data: TavilySearchResponse, query: string): string {
  const results = data.results ?? [];
  if (results.length === 0) {
    return `No search results found for "${query}".`;
  }

  const formatted = results
    .map((r) => `Title: ${r.title}\nURL: ${r.url}\nContent: ${r.content}\n`)
    .join("\n---\n");

  const answer = data.answer ? `Answer: ${data.answer}\n\n` : "";
  return `${answer}Search results for "${query}":\n\n${formatted}`;
}

// ── Extension factory ──────────────────────────────────────────────────

export default function (pi: ExtensionAPI): void {

  pi.registerTool({
    name: "tavily_search",
    label: "Tavily Web Search",
    description:
      "Performs a web search using the Tavily API to get real-time information from the internet.",
    parameters: tavilySearchSchema,
    async execute(_toolCallId, params, signal, _onUpdate, _ctx) {
      const apiKey = getTavilyKey();
      if (!apiKey) {
        return {
          content: [{ type: "text", text: MISSING_KEY_MESSAGE }],
          details: { error: "missing_api_key" },
          isError: true,
        };
      }

      try {
        const response = await fetch(TAVILY_SEARCH_ENDPOINT, {
          method: "POST",
          headers: { "Content-Type": "application/json" },
          body: JSON.stringify({
            api_key: apiKey,
            query: params.query,
            search_depth: "advanced",
            max_results: 5,
          }),
          signal,
        });

        if (!response.ok) {
          const errorText = await response.text();
          return {
            content: [
              {
                type: "text",
                text: `Tavily API error: ${String(response.status)} ${response.statusText}\n${errorText}`,
              },
            ],
            details: { status: response.status, body: errorText },
            isError: true,
          };
        }

        const data = (await response.json()) as TavilySearchResponse;
        return {
          content: [{ type: "text", text: formatResults(data, params.query) }],
          details: { raw: data },
        };
      } catch (error) {
        const message = error instanceof Error ? error.message : String(error);
        return {
          content: [{ type: "text", text: `Error performing Tavily search: ${message}` }],
          details: { error: message },
          isError: true,
        };
      }
    },
  });
}