/** * 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. */ 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 = { 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; Exa and Keenable both work keyless. */ export function availableProviders(env: Record = process.env): SearchProvider[] { const providers: SearchProvider[] = []; if (env.TAVILY_API_KEY?.trim()) 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; }; } export async function routeSearch( request: SearchRequest, ctx: RouterContext, cache: Map, signal?: AbortSignal, env: Record = process.env, ): Promise { const candidates = availableProviders(env); 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" }; } }