Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 43 additions & 1 deletion src/websearch/native.ts
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ export type NativeAuthResult =

export interface NativeModelRegistry {
getApiKeyAndHeaders(model: NativeModelInfo): Promise<NativeAuthResult>;
getAvailable?(): NativeModelInfo[];
}

interface NativeProviderMapping {
Expand Down Expand Up @@ -74,6 +75,7 @@ function buildEndpointUrl(baseUrl: string, resource: string): string {
export async function buildNativeEntry(
model: NativeModelInfo | undefined,
modelRegistry: NativeModelRegistry | undefined,
id = "native",
): Promise<SearchProviderEntry | null> {
if (!model || !modelRegistry) return null;

Expand All @@ -86,11 +88,51 @@ export async function buildNativeEntry(
if (!auth.ok || !auth.apiKey) return null;

return {
id: "native",
id,
provider: mapping.provider,
apiKey: auth.apiKey,
baseUrl,
model: model.id,
priority: -1,
};
}

function nativeEntryKey(entry: SearchProviderEntry): string {
return `${entry.provider}:${entry.baseUrl ?? ""}:${entry.model ?? ""}`;
}

function stableIdPart(value: string): string {
return (
value
.toLowerCase()
.replace(/[^a-z0-9]+/g, "-")
.replace(/^-+|-+$/g, "") || "model"
);
}

function discoveredNativeEntryId(entry: SearchProviderEntry): string {
return `native-${entry.provider}-${stableIdPart(entry.model ?? "model")}`;
}

export async function buildNativeEntries(
model: NativeModelInfo | undefined,
modelRegistry: NativeModelRegistry | undefined,
): Promise<SearchProviderEntry[]> {
if (!modelRegistry) return [];

const entries: SearchProviderEntry[] = [];
const activeEntry = await buildNativeEntry(model, modelRegistry);
if (activeEntry) entries.push(activeEntry);

const seen = new Set(entries.map(nativeEntryKey));
const availableModels = modelRegistry.getAvailable?.() ?? [];
for (const availableModel of availableModels) {
const entry = await buildNativeEntry(availableModel, modelRegistry, "native-discovered");
if (!entry) continue;
const key = nativeEntryKey(entry);
if (seen.has(key)) continue;
seen.add(key);
entries.push({ ...entry, id: discoveredNativeEntryId(entry) });
}
return entries;
}
6 changes: 3 additions & 3 deletions src/websearch/tool.ts
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import { defineTool } from "@mariozechner/pi-coding-agent";
import { Type } from "typebox";

import { buildNativeEntry, type NativeModelInfo, type NativeModelRegistry } from "./native.js";
import { buildNativeEntries, type NativeModelInfo, type NativeModelRegistry } from "./native.js";
import { renderSearchCall, renderSearchResult } from "./renderers.js";
import { createSearchRoutingState, formatSearchText, performSearch, type SearchRoutingState } from "./search.js";
import type {
Expand Down Expand Up @@ -36,8 +36,8 @@ interface WebSearchToolContext {

async function configWithNativeRoute(config: WebsearchConfig, ctx?: WebSearchToolContext): Promise<WebsearchConfig> {
if (!config.auto) return config;
const nativeEntry = await buildNativeEntry(ctx?.model, ctx?.modelRegistry);
return nativeEntry ? { ...config, providers: [nativeEntry, ...config.providers] } : config;
const nativeEntries = await buildNativeEntries(ctx?.model, ctx?.modelRegistry);
return nativeEntries.length > 0 ? { ...config, providers: [...nativeEntries, ...config.providers] } : config;
}

function providerLabel(provider: SearchProviderEntry): string {
Expand Down
86 changes: 85 additions & 1 deletion test/tool.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -27,8 +27,11 @@ function config(auto: boolean, maxResults?: number): WebsearchConfig {
};
}

function context(model: NativeModelInfo) {
function context(model: NativeModelInfo, availableModels: NativeModelInfo[] = []) {
const modelRegistry: NativeModelRegistry = {
getAvailable() {
return availableModels;
},
async getApiKeyAndHeaders() {
return { ok: true, apiKey: "native-test" };
},
Expand Down Expand Up @@ -269,6 +272,87 @@ describe("web_search tool definition", () => {
expect("provider" in details).toBe(false);
});

it("#given auto enabled and registry has available search-capable models #when executing #then prepends discovered providers", async () => {
// given
const requestedUrls: string[] = [];
vi.stubGlobal("fetch", async (input: string | URL | Request): Promise<Response> => {
const url = String(input);
requestedUrls.push(url);
if (url.includes("messages")) {
return jsonResponse({
content: [
{
type: "web_search_tool_result",
content: [{ title: "Claude", url: "https://claude.example.com", page_age: "1 day ago" }],
},
],
});
}
return jsonResponse({ results: [{ title: "Manual", url: "https://manual.example.com", text: "manual" }] });
});
const tool = withNativeExecutionContext(
createWebSearchTool(() => ({ ok: true, config: config(true), source: "test" })),
);
const availableModels: NativeModelInfo[] = [
{ provider: "anthropic", id: "claude-sonnet-4", baseUrl: "https://anthropic.gateway.example.com/v1" },
{ provider: "openai", id: "gpt-5.5", baseUrl: "https://openai.gateway.example.com/v1" },
];

// when
const result = await tool.execute(
"tool-call",
{ query: "discovered route" },
undefined,
undefined,
context({ provider: "openai", id: "gpt-3.5", baseUrl: "https://gateway.example.com/v1" }, availableModels),
);

// then
const details = result.details as SearchDetails;
expect(requestedUrls).toEqual(["https://anthropic.gateway.example.com/v1/messages"]);
expect(details.provider).toBe("anthropic");
expect(details.entryId).toBe("native-anthropic-claude-sonnet-4");
expect(details.attempts?.map((attempt) => attempt.entryId)).toEqual(["native-anthropic-claude-sonnet-4"]);
});

it("#given available models share provider #when discovering entries #then IDs stay unique", async () => {
// given
const requestedUrls: string[] = [];
vi.stubGlobal("fetch", async (input: string | URL | Request): Promise<Response> => {
requestedUrls.push(String(input));
return jsonResponse({ results: [] });
});
const tool = withNativeExecutionContext(
createWebSearchTool(() => ({ ok: true, config: config(true), source: "test" })),
);
const availableModels: NativeModelInfo[] = [
{ provider: "openai", id: "gpt-4.1", baseUrl: "https://openai.gateway.example.com/v1" },
{ provider: "openai", id: "gpt-5.5", baseUrl: "https://openai.gateway.example.com/v1" },
];

// when
const result = await tool.execute(
"tool-call",
{ query: "duplicate provider ids" },
undefined,
undefined,
context({ provider: "custom", id: "not-search", baseUrl: "https://gateway.example.com/v1" }, availableModels),
);

// then
const details = result.details as SearchDetails;
expect(requestedUrls).toEqual([
"https://openai.gateway.example.com/v1/responses",
"https://openai.gateway.example.com/v1/responses",
"https://gateway.example.com/exa",
]);
expect(details.attempts?.map((attempt) => attempt.entryId)).toEqual([
"native-openai-gpt-4-1",
"native-openai-gpt-5-5",
"manual",
]);
});

it("#given auto enabled and unsupported active model #when executing #then does not prepend native provider", async () => {
// given
const requestedUrls: string[] = [];
Expand Down
Loading