mirror of
https://bitbucket.org/siakitem/my-pi.git
synced 2026-10-26 10:45:22 +00:00
feat(pi-web-search): 统一 web_search/web_fetch 并由主模型路由搜索厂商
7 个按厂商命名的搜索工具收敛为 2 个统一工具，schema 与提示词不再暴露 provider 概念：web_search 由当前主模型按内置规则在 Tavily/Exa/Keenable 间路由（会话内 LRU 缓存，失败按可用性降级），web_fetch 固定走 Exa 并 静默降级。三家均改为纯 fetch 直连——Exa 直接对 mcp.exa.ai 发 JSON-RPC 并缓存会话，彻底移除 pi-mcp-adapter、@tavily/pi-extension 与 @keenable/pi-search 依赖。 Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
云服务部-叶林立
co-authored by
Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
parent
6fe6cb29ae
commit
ee1e6999d7
@@ -0,0 +1,178 @@
|
||||
import assert from "node:assert/strict";
|
||||
import test from "node:test";
|
||||
|
||||
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
|
||||
|
||||
import webSearchExtension, {
|
||||
createFetchRunners,
|
||||
createProviderRunners,
|
||||
dispatchFetch,
|
||||
dispatchSearch,
|
||||
normalizeUrls,
|
||||
toExaParams,
|
||||
toKeenableParams,
|
||||
toTavilyParams,
|
||||
type WebSearchParams,
|
||||
} from "../index.ts";
|
||||
|
||||
const PARAMS: WebSearchParams = { query: "release notes" };
|
||||
|
||||
function ctxAnswering(answer: string, calls: { count: number } = { count: 0 }) {
|
||||
return {
|
||||
model: { provider: "test", id: "main" },
|
||||
modelRegistry: {
|
||||
complete: async () => {
|
||||
calls.count += 1;
|
||||
return { content: [{ type: "text", text: answer }] };
|
||||
},
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
function withEnv<T>(key: string, value: string | undefined, fn: () => Promise<T> | T): Promise<T> | T {
|
||||
const previous = process.env[key];
|
||||
if (value === undefined) delete process.env[key];
|
||||
else process.env[key] = value;
|
||||
const restore = () => {
|
||||
if (previous === undefined) delete process.env[key];
|
||||
else process.env[key] = previous;
|
||||
};
|
||||
try {
|
||||
const out = fn();
|
||||
if (out instanceof Promise) return out.finally(restore);
|
||||
restore();
|
||||
return out;
|
||||
} catch (error) {
|
||||
restore();
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
test("extension registers only web_search and web_fetch, with no provider in the schema", () => {
|
||||
const tools = new Map<string, { parameters?: { properties?: Record<string, unknown> }; description?: string }>();
|
||||
const pi = {
|
||||
registerTool(definition: { name: string }) {
|
||||
tools.set(definition.name, definition as never);
|
||||
},
|
||||
on() {},
|
||||
} as unknown as ExtensionAPI;
|
||||
|
||||
webSearchExtension(pi);
|
||||
assert.deepEqual([...tools.keys()].sort(), ["web_fetch", "web_search"]);
|
||||
const schemaText = JSON.stringify([...tools.values()].map((t) => t.parameters));
|
||||
assert.doesNotMatch(schemaText, /tavily|exa|keenable|"provider"/iu);
|
||||
});
|
||||
|
||||
test("dispatchSearch runs the model-picked provider first and records details", async () => {
|
||||
await withEnv("TAVILY_API_KEY", "k", async () => {
|
||||
const ran: string[] = [];
|
||||
const runners = createProviderRunners({
|
||||
exa: async () => {
|
||||
ran.push("exa");
|
||||
return "exa results";
|
||||
},
|
||||
});
|
||||
const result = await dispatchSearch(PARAMS, ctxAnswering("exa"), new Map(), undefined, runners);
|
||||
assert.deepEqual(ran, ["exa"]);
|
||||
assert.equal(result.content[0]!.type === "text" && result.content[0]!.text, "exa results");
|
||||
assert.deepEqual(result.details, { provider: "exa", routedBy: "model" });
|
||||
});
|
||||
});
|
||||
|
||||
test("dispatchSearch falls back across remaining providers when the pick fails", async () => {
|
||||
await withEnv("TAVILY_API_KEY", "k", async () => {
|
||||
const ran: string[] = [];
|
||||
const runners = createProviderRunners({
|
||||
tavily: async () => {
|
||||
ran.push("tavily");
|
||||
throw new Error("tavily down");
|
||||
},
|
||||
exa: async () => {
|
||||
ran.push("exa");
|
||||
return "exa results";
|
||||
},
|
||||
});
|
||||
const result = await dispatchSearch(PARAMS, ctxAnswering("tavily"), new Map(), undefined, runners);
|
||||
assert.deepEqual(ran, ["tavily", "exa"]);
|
||||
assert.equal((result.details as Record<string, unknown>).provider, "exa");
|
||||
assert.equal((result.details as Record<string, unknown>).routedBy, "provider-fallback");
|
||||
});
|
||||
});
|
||||
|
||||
test("dispatchSearch reports an error result when every provider fails", async () => {
|
||||
await withEnv("TAVILY_API_KEY", "k", async () => {
|
||||
const runners = createProviderRunners({
|
||||
tavily: async () => { throw new Error("t"); },
|
||||
exa: async () => { throw new Error("e"); },
|
||||
keenable: async () => { throw new Error("k"); },
|
||||
});
|
||||
const result = await dispatchSearch(PARAMS, ctxAnswering("exa"), new Map(), undefined, runners);
|
||||
assert.equal((result.details as Record<string, unknown>).error, true);
|
||||
assert.match(result.content[0]!.type === "text" ? result.content[0]!.text : "", /failed on every available provider/);
|
||||
});
|
||||
});
|
||||
|
||||
test("normalizeUrls accepts a string or array, trims, and deduplicates", () => {
|
||||
assert.deepEqual(normalizeUrls(" https://a.example "), ["https://a.example"]);
|
||||
assert.deepEqual(normalizeUrls(["https://a", " https://a ", "https://b"]), ["https://a", "https://b"]);
|
||||
});
|
||||
|
||||
test("dispatchFetch prefers exa and degrades to keenable without a model call", async () => {
|
||||
await withEnv("TAVILY_API_KEY", undefined, async () => {
|
||||
const fetchers = createFetchRunners({
|
||||
exa: async (urls) => {
|
||||
throw new Error(`exa down for ${urls.length}`);
|
||||
},
|
||||
keenable: async (urls) => `keenable pages: ${urls.join(",")}`,
|
||||
});
|
||||
const result = await dispatchFetch({ urls: ["https://a.example", "https://b.example"] }, undefined, fetchers);
|
||||
assert.equal((result.details as Record<string, unknown>).provider, "keenable");
|
||||
assert.match(result.content[0]!.type === "text" ? result.content[0]!.text : "", /keenable pages: https:\/\/a\.example,https:\/\/b\.example/);
|
||||
});
|
||||
});
|
||||
|
||||
test("dispatchFetch validates input and caps the URL count", async () => {
|
||||
const empty = await dispatchFetch({ urls: [] }, undefined);
|
||||
assert.equal((empty.details as Record<string, unknown>).error, true);
|
||||
|
||||
const tooMany = await dispatchFetch({ urls: Array.from({ length: 11 }, (_, i) => `https://e${i}.example`) }, undefined);
|
||||
assert.equal((tooMany.details as Record<string, unknown>).error, true);
|
||||
assert.match(tooMany.content[0]!.type === "text" ? tooMany.content[0]!.text : "", /at most 10/);
|
||||
});
|
||||
|
||||
test("unified params map onto each provider's native arguments", () => {
|
||||
const params: WebSearchParams = {
|
||||
query: "q",
|
||||
max_results: 4,
|
||||
site: "a.dev",
|
||||
published_after: "2025-01-01",
|
||||
published_before: "2025-06-01",
|
||||
depth: "advanced",
|
||||
};
|
||||
assert.deepEqual(toTavilyParams(params), {
|
||||
query: "q",
|
||||
max_results: 4,
|
||||
search_depth: "advanced",
|
||||
include_domains: ["a.dev"],
|
||||
start_date: "2025-01-01",
|
||||
end_date: "2025-06-01",
|
||||
});
|
||||
assert.deepEqual(toExaParams(params), {
|
||||
query: "q",
|
||||
numResults: 4,
|
||||
includeDomains: ["a.dev"],
|
||||
startPublishedDate: "2025-01-01",
|
||||
endPublishedDate: "2025-06-01",
|
||||
});
|
||||
assert.deepEqual(toKeenableParams(params), {
|
||||
query: "q",
|
||||
site: "a.dev",
|
||||
published_after: "2025-01-01",
|
||||
published_before: "2025-06-01",
|
||||
});
|
||||
|
||||
const defaults = toTavilyParams({ query: "q" });
|
||||
assert.equal(defaults.max_results, 8);
|
||||
assert.equal(defaults.search_depth, "basic");
|
||||
assert.equal(defaults.include_domains, undefined);
|
||||
});
|
||||
@@ -0,0 +1,140 @@
|
||||
import assert from "node:assert/strict";
|
||||
import test from "node:test";
|
||||
|
||||
import { ExaError, exaFetchPages, exaSearch, parseRpcPayload, resetExaSession } from "../clients/exa.ts";
|
||||
|
||||
function sse(payload: Record<string, unknown>): string {
|
||||
return `event: message\ndata: ${JSON.stringify(payload)}\n\n`;
|
||||
}
|
||||
|
||||
interface RecordedCall {
|
||||
method?: string;
|
||||
params?: Record<string, unknown>;
|
||||
headers: Record<string, string>;
|
||||
}
|
||||
|
||||
/**
|
||||
* A fetch stub that speaks the Exa endpoint protocol: initialize answers with
|
||||
* the scripted session id and the request's own id, notifications get a bare
|
||||
* 202, and tools/call delegates to the supplied handler (default 404).
|
||||
*/
|
||||
function exaStub(options: {
|
||||
sessionIds?: string[];
|
||||
onToolCall?: (args: Record<string, unknown>, id: number, session: string | undefined) => Response;
|
||||
} = {}) {
|
||||
const calls: RecordedCall[] = [];
|
||||
let sessionIndex = 0;
|
||||
const fetchImpl = async (_url: string | URL, init?: RequestInit): Promise<Response> => {
|
||||
const headers = (init?.headers ?? {}) as Record<string, string>;
|
||||
const body = JSON.parse(String(init?.body)) as { id?: number; method?: string; params?: Record<string, unknown> };
|
||||
calls.push({ method: body.method, params: body.params, headers });
|
||||
if (body.method === "initialize") {
|
||||
const sid = options.sessionIds?.[Math.min(sessionIndex++, (options.sessionIds?.length ?? 1) - 1)] ?? `sid-${sessionIndex}`;
|
||||
return new Response(
|
||||
sse({ jsonrpc: "2.0", id: body.id, result: { protocolVersion: "2025-03-26", capabilities: {}, serverInfo: { name: "exa" } } }),
|
||||
{ status: 200, headers: { "mcp-session-id": sid, "content-type": "text/event-stream" } },
|
||||
);
|
||||
}
|
||||
if (body.method === "notifications/initialized") return new Response("", { status: 202 });
|
||||
if (body.method === "tools/call") {
|
||||
return options.onToolCall?.(body.params ?? {}, body.id ?? 0, headers["mcp-session-id"]) ?? new Response("Not Found", { status: 404 });
|
||||
}
|
||||
return new Response("unexpected", { status: 500 });
|
||||
};
|
||||
return { fetchImpl, calls };
|
||||
}
|
||||
|
||||
function okToolResult(id: number, text: string): Response {
|
||||
return new Response(
|
||||
sse({ jsonrpc: "2.0", id, result: { content: [{ type: "text", text }] } }),
|
||||
{ status: 200, headers: { "content-type": "text/event-stream" } },
|
||||
);
|
||||
}
|
||||
|
||||
test("parseRpcPayload reads JSON bodies and SSE data lines, matching the request id", () => {
|
||||
const direct = JSON.stringify({ jsonrpc: "2.0", id: 7, result: { ok: true } });
|
||||
assert.deepEqual(parseRpcPayload(direct, 7).result, { ok: true });
|
||||
|
||||
const body =
|
||||
`event: message\ndata: {"jsonrpc":"2.0","id":1,"result":{"a":1}}\n\n` +
|
||||
`event: message\ndata: {"jsonrpc":"2.0","id":2,"result":{"b":2}}\n\n`;
|
||||
assert.deepEqual(parseRpcPayload(body, 2).result, { b: 2 });
|
||||
|
||||
assert.throws(() => parseRpcPayload("not json at all", 1), ExaError);
|
||||
assert.throws(() => parseRpcPayload(sse({ jsonrpc: "2.0", id: 99, result: {} }), 1), ExaError);
|
||||
});
|
||||
|
||||
test("exaSearch initializes a session, notifies, then posts tools/call with the session id", async () => {
|
||||
resetExaSession();
|
||||
const { fetchImpl, calls } = exaStub({
|
||||
sessionIds: ["sid-A"],
|
||||
onToolCall: (_args, id) => okToolResult(id, "search results"),
|
||||
});
|
||||
|
||||
const text = await exaSearch({ query: "docs", numResults: 3, includeDomains: ["python.org"] }, undefined, fetchImpl);
|
||||
assert.equal(text, "search results");
|
||||
assert.deepEqual(calls.map((c) => c.method), ["initialize", "notifications/initialized", "tools/call"]);
|
||||
|
||||
const call = calls[2]!;
|
||||
assert.equal(call.headers["mcp-session-id"], "sid-A");
|
||||
assert.equal((call.params as { name: string }).name, "web_search_advanced_exa");
|
||||
const args = (call.params as { arguments: Record<string, unknown> }).arguments;
|
||||
assert.equal(args.numResults, 3);
|
||||
assert.equal(args.type, "auto");
|
||||
assert.deepEqual(args.includeDomains, ["python.org"]);
|
||||
assert.ok((args.textMaxCharacters as number) > 0);
|
||||
});
|
||||
|
||||
test("exaSearch reuses the cached session across calls", async () => {
|
||||
resetExaSession();
|
||||
const { fetchImpl, calls } = exaStub({
|
||||
sessionIds: ["sid-1"],
|
||||
onToolCall: (args, id) => okToolResult(id, `result for ${(args as { name: string }).name}`),
|
||||
});
|
||||
|
||||
await exaSearch({ query: "a" }, undefined, fetchImpl);
|
||||
const text = await exaFetchPages(["https://example.com"], 500, undefined, fetchImpl);
|
||||
assert.match(text, /web_fetch_exa/);
|
||||
assert.equal(calls.length, 4);
|
||||
assert.equal((calls[3]!.params as { name: string }).name, "web_fetch_exa");
|
||||
assert.equal(calls[3]!.headers["mcp-session-id"], "sid-1");
|
||||
assert.equal((calls[3]!.params as { arguments: { maxCharacters: number } }).arguments.maxCharacters, 500);
|
||||
});
|
||||
|
||||
test("a 404 from tools/call drops the session, re-initializes, and retries once", async () => {
|
||||
resetExaSession();
|
||||
let toolCalls = 0;
|
||||
const { fetchImpl, calls } = exaStub({
|
||||
sessionIds: ["stale-sid", "fresh-sid"],
|
||||
onToolCall: (_args, id) => {
|
||||
toolCalls += 1;
|
||||
return toolCalls === 1 ? new Response("Not Found", { status: 404 }) : okToolResult(id, "after re-init");
|
||||
},
|
||||
});
|
||||
|
||||
const text = await exaSearch({ query: "retry me" }, undefined, fetchImpl);
|
||||
assert.equal(text, "after re-init");
|
||||
assert.deepEqual(calls.map((c) => c.method), [
|
||||
"initialize",
|
||||
"notifications/initialized",
|
||||
"tools/call",
|
||||
"initialize",
|
||||
"notifications/initialized",
|
||||
"tools/call",
|
||||
]);
|
||||
assert.equal(calls[5]!.headers["mcp-session-id"], "fresh-sid");
|
||||
});
|
||||
|
||||
test("isError tool results and HTTP failures raise ExaError", async () => {
|
||||
resetExaSession();
|
||||
const isError = exaStub({
|
||||
onToolCall: (_args, id) =>
|
||||
new Response(sse({ jsonrpc: "2.0", id, result: { isError: true, content: [{ type: "text", text: "bad args" }] } }), { status: 200 }),
|
||||
});
|
||||
await assert.rejects(exaFetchPages(["https://example.com"], undefined, undefined, isError.fetchImpl), /bad args/);
|
||||
|
||||
resetExaSession();
|
||||
const httpFail = exaStub({ onToolCall: () => new Response("boom", { status: 503 }) });
|
||||
// Session init succeeds; tools/call hits a non-404 HTTP error.
|
||||
await assert.rejects(exaSearch({ query: "x" }, undefined, httpFail.fetchImpl), /HTTP 503/);
|
||||
});
|
||||
@@ -0,0 +1,109 @@
|
||||
import assert from "node:assert/strict";
|
||||
import test from "node:test";
|
||||
|
||||
import {
|
||||
availableProviders,
|
||||
buildRouterPrompt,
|
||||
parseProviderChoice,
|
||||
routeSearch,
|
||||
ROUTE_CACHE_LIMIT,
|
||||
type RouterContext,
|
||||
type SearchProvider,
|
||||
} from "../router.ts";
|
||||
|
||||
const ALL: SearchProvider[] = ["tavily", "exa", "keenable"];
|
||||
const ENV_ALL_KEYS = { TAVILY_API_KEY: "test-key" };
|
||||
|
||||
function ctxAnswering(answer: string, calls: { count: number } = { count: 0 }): RouterContext {
|
||||
return {
|
||||
model: { provider: "test", id: "main" },
|
||||
modelRegistry: {
|
||||
complete: async () => {
|
||||
calls.count += 1;
|
||||
return { content: [{ type: "text", text: answer }] };
|
||||
},
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
test("availability: tavily needs a key, exa and keenable are keyless", () => {
|
||||
assert.deepEqual(availableProviders({}), ["exa", "keenable"]);
|
||||
assert.deepEqual(availableProviders(ENV_ALL_KEYS), ["tavily", "exa", "keenable"]);
|
||||
});
|
||||
|
||||
test("router prompt lists only candidate rules plus the query and caller hints", () => {
|
||||
const prompt = buildRouterPrompt(
|
||||
{ query: "python 3.14 release notes", site: "python.org", published_after: "2025-10-01" },
|
||||
ALL,
|
||||
);
|
||||
assert.match(prompt, /- tavily: /);
|
||||
assert.match(prompt, /- exa: /);
|
||||
assert.match(prompt, /- keenable: /);
|
||||
assert.match(prompt, /Query: python 3\.14 release notes/);
|
||||
assert.match(prompt, /site filter.*python\.org|site: python\.org/);
|
||||
assert.match(prompt, /Provider:$/);
|
||||
|
||||
const subset = buildRouterPrompt({ query: "hi" }, ["exa", "keenable"]);
|
||||
assert.doesNotMatch(subset, /- tavily/);
|
||||
assert.match(subset, /exactly one word: exa, keenable/);
|
||||
});
|
||||
|
||||
test("parseProviderChoice takes the first available candidate token", () => {
|
||||
assert.equal(parseProviderChoice("tavily", ALL), "tavily");
|
||||
assert.equal(parseProviderChoice("I would pick Exa here.", ALL), "exa");
|
||||
assert.equal(parseProviderChoice("keenable!", ALL), "keenable");
|
||||
assert.equal(parseProviderChoice("google maybe", ALL), undefined);
|
||||
// Model picked an unavailable provider but mentioned a valid one later.
|
||||
assert.equal(parseProviderChoice("tavily or keenable", ["exa", "keenable"]), "keenable");
|
||||
assert.equal(parseProviderChoice("tavily", ["exa", "keenable"]), undefined);
|
||||
});
|
||||
|
||||
test("routeSearch honors the model pick and caches identical requests", async () => {
|
||||
const calls = { count: 0 };
|
||||
const ctx = ctxAnswering("exa", calls);
|
||||
const cache = new Map<string, SearchProvider>();
|
||||
const request = { query: "official MCP spec" };
|
||||
|
||||
const first = await routeSearch(request, ctx, cache, undefined, ENV_ALL_KEYS);
|
||||
assert.deepEqual(first, { provider: "exa", routedBy: "model" });
|
||||
|
||||
const second = await routeSearch(request, ctx, cache, undefined, ENV_ALL_KEYS);
|
||||
assert.deepEqual(second, { provider: "exa", routedBy: "cache" });
|
||||
assert.equal(calls.count, 1);
|
||||
});
|
||||
|
||||
test("routeSearch falls back to the first available provider on unroutable or failed model calls", async () => {
|
||||
const cache = new Map<string, SearchProvider>();
|
||||
const unparseable = await routeSearch({ query: "x" }, ctxAnswering("no idea"), cache, undefined, ENV_ALL_KEYS);
|
||||
assert.deepEqual(unparseable, { provider: "tavily", routedBy: "fallback" });
|
||||
|
||||
const throwing: RouterContext = {
|
||||
model: {},
|
||||
modelRegistry: { complete: async () => { throw new Error("boom"); } },
|
||||
};
|
||||
const failed = await routeSearch({ query: "x" }, throwing, cache, undefined, ENV_ALL_KEYS);
|
||||
assert.deepEqual(failed, { provider: "tavily", routedBy: "fallback" });
|
||||
});
|
||||
|
||||
test("routeSearch without a tavily key excludes it from candidates", async () => {
|
||||
const cache = new Map<string, SearchProvider>();
|
||||
const decision = await routeSearch({ query: "x" }, ctxAnswering("tavily"), cache, undefined, {});
|
||||
// "tavily" is not a candidate → unroutable → first available fallback.
|
||||
assert.deepEqual(decision, { provider: "exa", routedBy: "fallback" });
|
||||
});
|
||||
|
||||
test("routeSearch works without a model", async () => {
|
||||
const cache = new Map<string, SearchProvider>();
|
||||
const decision = await routeSearch({ query: "x" }, {}, cache, undefined, ENV_ALL_KEYS);
|
||||
assert.deepEqual(decision, { provider: "tavily", routedBy: "fallback" });
|
||||
});
|
||||
|
||||
test("route cache evicts the oldest entry at the limit", async () => {
|
||||
const cache = new Map<string, SearchProvider>();
|
||||
for (let i = 0; i < ROUTE_CACHE_LIMIT; i++) cache.set(`dummy-${i}`, "exa");
|
||||
const decision = await routeSearch({ query: "fresh query" }, ctxAnswering("keenable"), cache, undefined, ENV_ALL_KEYS);
|
||||
assert.equal(decision.provider, "keenable");
|
||||
assert.equal(cache.size, ROUTE_CACHE_LIMIT);
|
||||
assert.equal(cache.has("dummy-0"), false);
|
||||
assert.equal(cache.has("dummy-1"), true);
|
||||
});
|
||||
Reference in New Issue
Block a user