my-pi/pi-web-search/index.ts

/**
 * pi-web-search: one unified web_search + web_fetch surface.
 *
 * The main model never sees provider names — web_search routes internally via
 * the session's current model (see router.ts), and web_fetch always reads
 * pages through Exa with graceful degradation. Provider clients are plain
 * fetch calls (clients/*) — no MCP adapter, no upstream extension wrappers.
 */

import type { AgentToolResult, ExtensionAPI } from "@earendil-works/pi-coding-agent";

import { getSearchApiKey, type SearchKeyResolver } from "./auth.ts";
import { exaFetchPages, exaSearch } from "./clients/exa.ts";
import { keenableFetch, keenableSearch } from "./clients/keenable.ts";
import { tavilyAvailable, tavilyExtract, tavilySearch } from "./clients/tavily.ts";
import {
	availableProviders,
	routeSearch,
	type SearchProvider,
	type RouterContext,
} from "./router.ts";

const DEFAULT_MAX_RESULTS = 8;
const MAX_RESULTS_LIMIT = 20;
const MAX_FETCH_URLS = 10;
const DEFAULT_FETCH_MAX_CHARACTERS = 20000;
const FETCH_MAX_CHARACTERS_LIMIT = 100000;

function clampInt(value: number | undefined, fallback: number, min: number, max: number): number {
	if (typeof value !== "number" || !Number.isFinite(value)) return fallback;
	return Math.min(Math.max(Math.floor(value), min), max);
}

export interface WebSearchParams {
	query: string;
	max_results?: number;
	site?: string;
	published_after?: string;
	published_before?: string;
	depth?: "basic" | "advanced";
}

export interface WebFetchParams {
	urls: string | string[];
	max_characters?: number;
}

export type ProviderRunner = (params: WebSearchParams, signal?: AbortSignal) => Promise<string>;
export type FetchRunner = (urls: string[], maxCharacters: number, signal?: AbortSignal) => Promise<string>;

function formatSearchPayload(payload: unknown): string {
	return JSON.stringify(payload, null, 2);
}

export function toTavilyParams(params: WebSearchParams) {
	return {
		query: params.query,
		max_results: clampInt(params.max_results, DEFAULT_MAX_RESULTS, 1, MAX_RESULTS_LIMIT),
		search_depth: (params.depth ?? "basic") as "basic" | "advanced",
		include_domains: params.site ? [params.site] : undefined,
		start_date: params.published_after,
		end_date: params.published_before,
	};
}

export function toExaParams(params: WebSearchParams) {
	return {
		query: params.query,
		numResults: clampInt(params.max_results, DEFAULT_MAX_RESULTS, 1, MAX_RESULTS_LIMIT),
		includeDomains: params.site ? [params.site] : undefined,
		startPublishedDate: params.published_after,
		endPublishedDate: params.published_before,
	};
}

export function toKeenableParams(params: WebSearchParams) {
	return {
		query: params.query,
		site: params.site,
		published_after: params.published_after,
		published_before: params.published_before,
	};
}

/** Default per-provider parameter mapping from the unified schema. */
export function createProviderRunners(
	overrides: Partial<Record<SearchProvider, ProviderRunner>> = {},
): Record<SearchProvider, ProviderRunner> {
	return {
		tavily:
			overrides.tavily ??
			(async (params, signal) => formatSearchPayload(await tavilySearch(toTavilyParams(params), signal))),
		exa: overrides.exa ?? ((params, signal) => exaSearch(toExaParams(params), signal)),
		keenable:
			overrides.keenable ??
			(async (params, signal) => formatSearchPayload(await keenableSearch(toKeenableParams(params), signal))),
	};
}

function errorText(error: unknown): string {
	return error instanceof Error ? error.message : String(error);
}

function toResult(text: string, details: Record<string, unknown>): AgentToolResult {
	return { content: [{ type: "text", text }], details };
}

/**
 * Route once via the main model, then try the picked provider first and fall
 * back across the remaining available providers on transport errors.
 */
export async function dispatchSearch(
	params: WebSearchParams,
	ctx: RouterContext,
	routeCache: Map<string, SearchProvider>,
	signal?: AbortSignal,
	runners: Record<SearchProvider, ProviderRunner> = createProviderRunners(),
	resolve: SearchKeyResolver = getSearchApiKey,
): Promise<AgentToolResult> {
	const decision = await routeSearch(params, ctx, routeCache, signal, resolve);
	const order = [decision.provider, ...availableProviders(resolve).filter((p) => p !== decision.provider)];
	let lastError: unknown;
	for (const provider of order) {
		try {
			const text = await runners[provider](params, signal);
			return toResult(text, {
				provider,
				routedBy: provider === decision.provider ? decision.routedBy : "provider-fallback",
			});
		} catch (error) {
			lastError = error;
		}
	}
	return toResult(`web_search failed on every available provider: ${errorText(lastError)}`, { error: true });
}

export function normalizeUrls(urls: string | string[]): string[] {
	const list = (Array.isArray(urls) ? urls : [urls]).map((url) => url.trim()).filter(Boolean);
	return [...new Set(list)];
}

/** web_fetch always prefers Exa; Keenable and Tavily are silent degradation paths. */
export function createFetchRunners(
	overrides: Partial<Record<SearchProvider, FetchRunner>> = {},
): Record<SearchProvider, FetchRunner> {
	return {
		exa: overrides.exa ?? ((urls, maxCharacters, signal) => exaFetchPages(urls, maxCharacters, signal)),
		keenable:
			overrides.keenable ??
			(async (urls, _maxCharacters, signal) => {
				const pages = await Promise.all(urls.map((url) => keenableFetch(url, signal)));
				return pages
					.map((page) => `# ${page.title ?? page.url ?? "page"}\n${page.content ?? "(no content)"}`)
					.join("\n\n---\n\n");
			}),
		tavily:
			overrides.tavily ??
			(async (urls, _maxCharacters, signal) => {
				const { results, failed } = await tavilyExtract(urls, signal);
				const body = results
					.map((page) => `# ${page.url ?? "page"}\n${page.content ?? "(no content)"}`)
					.join("\n\n---\n\n");
				return failed.length ? `${body}\n\nFailed URLs: ${failed.join(", ")}` : body;
			}),
	};
}

export async function dispatchFetch(
	params: WebFetchParams,
	signal?: AbortSignal,
	fetchers: Record<SearchProvider, FetchRunner> = createFetchRunners(),
	resolve: SearchKeyResolver = getSearchApiKey,
): Promise<AgentToolResult> {
	const urls = normalizeUrls(params.urls);
	if (!urls.length) {
		return toResult("web_fetch needs at least one http(s) URL.", { error: true });
	}
	if (urls.length > MAX_FETCH_URLS) {
		return toResult(`web_fetch accepts at most ${MAX_FETCH_URLS} URLs per call; received ${urls.length}.`, { error: true });
	}
	const maxCharacters = clampInt(params.max_characters, DEFAULT_FETCH_MAX_CHARACTERS, 1, FETCH_MAX_CHARACTERS_LIMIT);
	const order: SearchProvider[] = ["exa", "keenable"];
	if (tavilyAvailable(resolve)) order.push("tavily");

	let lastError: unknown;
	for (const provider of order) {
		try {
			const text = await fetchers[provider](urls, maxCharacters, signal);
			return toResult(text, { provider });
		} catch (error) {
			lastError = error;
		}
	}
	return toResult(`web_fetch failed on every available provider: ${errorText(lastError)}`, { error: true });
}

export default function webSearchExtension(pi: ExtensionAPI): void {
	const routeCache = new Map<string, SearchProvider>();

	pi.registerTool({
		name: "web_search",
		label: "Web Search",
		description:
			"Search the web for current or citable information and return a compact list of candidate " +
			"pages (title, url, summary). Provider selection is automatic. Follow up with web_fetch to " +
			"read the most relevant pages.",
		promptSnippet: "Search the web for current or citable information.",
		promptGuidelines: [
			"Use web_search when the task needs current, real-world, or citable web information.",
			"Keep queries focused; split multi-topic requests into separate web_search calls.",
			"Use site and published_* filters only when they materially narrow the results.",
			"After web_search, read the most relevant pages with web_fetch before relying on a claim.",
		],
		parameters: {
			type: "object",
			properties: {
				query: {
					type: "string",
					description: "The search query — a natural-language description of the ideal page, not just keywords.",
				},
				max_results: {
					type: "number",
					minimum: 1,
					maximum: MAX_RESULTS_LIMIT,
					description: `Maximum number of results to return (default ${DEFAULT_MAX_RESULTS}, max ${MAX_RESULTS_LIMIT}).`,
				},
				site: {
					type: "string",
					description: "Restrict results to a single domain, e.g. 'arxiv.org'.",
				},
				published_after: {
					type: "string",
					description: "Only pages published on/after this date (YYYY-MM-DD).",
				},
				published_before: {
					type: "string",
					description: "Only pages published on/before this date (YYYY-MM-DD).",
				},
				depth: {
					type: "string",
					enum: ["basic", "advanced"],
					description: "'basic' (default) balanced search; 'advanced' for higher precision at higher cost.",
				},
			},
			required: ["query"],
			additionalProperties: false,
		},
		async execute(_toolCallId, params, signal, _onUpdate, ctx) {
			return dispatchSearch(params as WebSearchParams, ctx as unknown as RouterContext, routeCache, signal);
		},
	});

	pi.registerTool({
		name: "web_fetch",
		label: "Web Fetch",
		description:
			"Fetch the full content of one or more http(s) URLs as clean markdown. Use after web_search " +
			"to read selected pages, or directly when the user provides a page to read.",
		promptSnippet: "Read the full content of known URLs.",
		promptGuidelines: [
			"Use web_fetch to read the full content of a URL found via web_search or provided by the user.",
			"Batch multiple URLs into a single web_fetch call when possible.",
		],
		parameters: {
			type: "object",
			properties: {
				urls: {
					anyOf: [
						{ type: "string", description: "A single http(s) URL to read." },
						{
							type: "array",
							items: { type: "string" },
							maxItems: MAX_FETCH_URLS,
							description: `A list of http(s) URLs to read (max ${MAX_FETCH_URLS}).`,
						},
					],
				},
				max_characters: {
					type: "number",
					minimum: 1,
					maximum: FETCH_MAX_CHARACTERS_LIMIT,
					description: `Maximum characters to extract per page (default ${DEFAULT_FETCH_MAX_CHARACTERS}, max ${FETCH_MAX_CHARACTERS_LIMIT}).`,
				},
			},
			required: ["urls"],
			additionalProperties: false,
		},
		async execute(_toolCallId, params, signal) {
			return dispatchFetch(params as WebFetchParams, signal);
		},
	});
}