dev_sandbox/.pi/agent/extensions/web-search.ts

/**
 * Web Search Extension (Tavily API)
 *
 * Registers a `web_search` tool the LLM can call for current/up-to-date info.
 *
 * Setup:
 *   1. Get an API key at https://tavily.com/ (email signup, free 1k queries/month)
 *   2. export TAVILY_API_KEY=tvly-...
 */

import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { StringEnum } from "@earendil-works/pi-ai";
import { Type } from "typebox";

type TavilyResult = {
	title: string;
	url: string;
	content: string;
	score?: number;
	published_date?: string;
};

type TavilyResponse = {
	query: string;
	answer?: string;
	results: TavilyResult[];
};

export default function webSearchExtension(pi: ExtensionAPI) {
	pi.registerTool({
		name: "web_search",
		label: "Web Search",
		description:
			"Search the web via Tavily and return a synthesized answer plus result titles, URLs, and snippets. " +
			"Use for current events, docs lookups, or any fact that may be newer than training data.",
		promptSnippet: "Search the web (Tavily) for up-to-date information",
		promptGuidelines: [
			"Use web_search when the user asks about current events, recent releases, or facts that may be outdated in training data.",
			"After web_search returns results, call fetch_url on the most relevant URLs if the snippets are insufficient. Prefer fetch_url over raw curl/wget in bash — it handles JS walls and anti-bot pages automatically.",
		],
		parameters: Type.Object({
			query: Type.String({ description: "Search query" }),
			max_results: Type.Optional(
				Type.Integer({
					description: "Max results to return (1-20, default 5)",
					minimum: 1,
					maximum: 20,
				}),
			),
			search_depth: Type.Optional(
				StringEnum(["basic", "advanced"] as const, {
					description:
						"'basic' is faster/cheaper; 'advanced' does deeper crawling for harder queries. Default 'basic'.",
				}),
			),
			topic: Type.Optional(
				StringEnum(["general", "news"] as const, {
					description: "Use 'news' for recent-events queries. Default 'general'.",
				}),
			),
			include_answer: Type.Optional(
				Type.Boolean({
					description: "Ask Tavily to return a short synthesized answer. Default true.",
				}),
			),
		}),
		async execute(_toolCallId, params, signal) {
			const apiKey = process.env.TAVILY_API_KEY;
			if (!apiKey) {
				return {
					content: [
						{
							type: "text",
							text:
								"TAVILY_API_KEY is not set. Get a key at https://tavily.com/ " +
								"and `export TAVILY_API_KEY=...` before launching pi.",
						},
					],
					isError: true,
					details: {},
				};
			}

			const body = {
				query: params.query,
				max_results: params.max_results ?? 5,
				search_depth: params.search_depth ?? "basic",
				topic: params.topic ?? "general",
				include_answer: params.include_answer ?? true,
			};

			let resp: Response;
			try {
				resp = await fetch("https://api.tavily.com/search", {
					method: "POST",
					headers: {
						"Content-Type": "application/json",
						Authorization: `Bearer ${apiKey}`,
					},
					body: JSON.stringify(body),
					signal,
				});
			} catch (err) {
				const msg = err instanceof Error ? err.message : String(err);
				return {
					content: [{ type: "text", text: `Network error calling Tavily: ${msg}` }],
					isError: true,
					details: {},
				};
			}

			if (!resp.ok) {
				const errBody = await resp.text().catch(() => "");
				return {
					content: [
						{
							type: "text",
							text: `Tavily search failed: ${resp.status} ${resp.statusText}\n${errBody}`,
						},
					],
					isError: true,
					details: { status: resp.status },
				};
			}

			const data = (await resp.json()) as TavilyResponse;
			const results = data.results ?? [];

			if (results.length === 0 && !data.answer) {
				return {
					content: [{ type: "text", text: `No results for: ${params.query}` }],
					details: { query: params.query, results: [] },
				};
			}

			const parts: string[] = [];
			if (data.answer) {
				parts.push(`Answer: ${data.answer}`);
			}
			if (results.length > 0) {
				const formatted = results
					.map(
						(r, i) =>
							`${i + 1}. ${r.title}\n   ${r.url}\n   ${r.content}${
								r.published_date ? ` (${r.published_date})` : ""
							}`,
					)
					.join("\n\n");
				parts.push(`Results:\n${formatted}`);
			}

			return {
				content: [{ type: "text", text: parts.join("\n\n") }],
				details: {
					query: params.query,
					answer: data.answer,
					results,
				},
			};
		},
	});
}