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

/**
 * Provider router for web_search. The decision is delegated to the session's
 * current main model (ctx.model via ctx.modelRegistry.complete) so no extra
 * model configuration is needed; deterministic rules only handle availability
 * and failure fallbacks, never the routing itself.
 */

import { getSearchApiKey, type SearchKeyResolver } from "./auth.ts";

export type SearchProvider = "tavily" | "exa" | "keenable";

export interface SearchRequest {
	query: string;
	site?: string;
	published_after?: string;
	published_before?: string;
	depth?: string;
}

export type RoutedBy = "model" | "cache" | "fallback";

export interface RouteDecision {
	provider: SearchProvider;
	routedBy: RoutedBy;
}

const ROUTER_TIMEOUT_MS = 15_000;
export const ROUTE_CACHE_LIMIT = 32;

/** Order used only when the model cannot decide or no candidates remain. */
export const PROVIDER_FALLBACK_ORDER: SearchProvider[] = ["tavily", "exa", "keenable"];

const PROVIDER_RULES: Record<SearchProvider, string> = {
	tavily:
		"Broad discovery. Best when keywords are uncertain, for news and recent or trending topics, " +
		"and when many candidate sources are useful. Results can mix official pages, blogs, forums, " +
		"and secondary articles; returned content is naturally long.",
	exa:
		"Precision sources. Best for official or API documentation, release notes and version " +
		"information, academic papers and research material, and fact confirmation that needs " +
		"high-quality sources.",
	keenable:
		"Focused and compact. Best for Chinese-language pages (policies, government notices), " +
		"site-restricted search, publication-date filtering, and tight control over returned " +
		"content length.",
};

/** Tavily requires a key in auth.json; Exa and Keenable both work keyless. */
export function availableProviders(resolve: SearchKeyResolver = getSearchApiKey): SearchProvider[] {
	const providers: SearchProvider[] = [];
	if (resolve("tavily")) providers.push("tavily");
	providers.push("exa", "keenable");
	return providers;
}

export function buildRouterPrompt(request: SearchRequest, candidates: SearchProvider[]): string {
	const rules = candidates.map((provider) => `- ${provider}: ${PROVIDER_RULES[provider]}`).join("\n");
	const lines = [
		"Pick the single best web search provider for this query.",
		`Answer with exactly one word: ${candidates.join(", ")}.`,
		"",
		rules,
		"",
		`Query: ${request.query}`,
	];
	if (request.site) lines.push(`The caller asked to restrict results to the site: ${request.site}`);
	if (request.published_after || request.published_before) {
		lines.push(
			`The caller asked for a publication-date filter: ${request.published_after ?? "…"} to ${request.published_before ?? "…"}`,
		);
	}
	if (request.depth === "advanced") lines.push("The caller asked for a deep, high-precision search.");
	lines.push("Provider:");
	return lines.join("\n");
}

/** First candidate-matching provider token in the model's reply, else undefined. */
export function parseProviderChoice(text: string, candidates: SearchProvider[]): SearchProvider | undefined {
	const matches = text.toLowerCase().match(/\b(tavily|exa|keenable)\b/gu) ?? [];
	for (const match of matches) {
		const pick = match as SearchProvider;
		if (candidates.includes(pick)) return pick;
	}
	return undefined;
}

export function routeCacheKey(request: SearchRequest): string {
	return [
		request.query.trim().toLowerCase(),
		request.site ?? "",
		request.published_after ?? "",
		request.published_before ?? "",
		request.depth ?? "",
	].join(" ");
}

interface ModelMessage {
	stopReason?: string;
	errorMessage?: string;
	content?: Array<{ type?: string; text?: string }>;
}

/** Minimal structural view of the extension context needed for routing. */
export interface RouterContext {
	model?: unknown;
	modelRegistry?: {
		complete?: (model: unknown, request: unknown, options: unknown) => Promise<ModelMessage>;
	};
}

export async function routeSearch(
	request: SearchRequest,
	ctx: RouterContext,
	cache: Map<string, SearchProvider>,
	signal?: AbortSignal,
	resolve: SearchKeyResolver = getSearchApiKey,
): Promise<RouteDecision> {
	const candidates = availableProviders(resolve);
	const fallback = PROVIDER_FALLBACK_ORDER.find((provider) => candidates.includes(provider)) ?? candidates[0];
	if (!fallback || candidates.length <= 1) {
		return { provider: fallback ?? "keenable", routedBy: "fallback" };
	}

	const key = routeCacheKey(request);
	const cached = cache.get(key);
	if (cached !== undefined && candidates.includes(cached)) {
		return { provider: cached, routedBy: "cache" };
	}

	const complete = ctx.modelRegistry?.complete;
	if (!ctx.model || typeof complete !== "function") {
		return { provider: fallback, routedBy: "fallback" };
	}

	try {
		const timeout = AbortSignal.timeout(ROUTER_TIMEOUT_MS);
		const message = await complete.call(ctx.modelRegistry, ctx.model, {
			messages: [
				{
					role: "user",
					content: [{ type: "text", text: buildRouterPrompt(request, candidates) }],
					timestamp: Date.now(),
				},
			],
		}, {
			signal: signal ? AbortSignal.any([signal, timeout]) : timeout,
		});
		if (message.stopReason === "error" || message.stopReason === "aborted") {
			throw new Error(message.errorMessage ?? `routing stopped: ${message.stopReason}`);
		}
		const text = (message.content ?? [])
			.flatMap((part) => (part.type === "text" && typeof part.text === "string" ? [part.text] : []))
			.join("\n");
		const pick = parseProviderChoice(text, candidates);
		if (!pick) return { provider: fallback, routedBy: "fallback" };
		if (!cache.has(key) && cache.size >= ROUTE_CACHE_LIMIT) {
			const oldest = cache.keys().next().value;
			if (oldest !== undefined) cache.delete(oldest);
		}
		cache.set(key, pick);
		return { provider: pick, routedBy: "model" };
	} catch {
		return { provider: fallback, routedBy: "fallback" };
	}
}