diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index d1313ddfe..651551286 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -42,6 +42,7 @@ - Fixed `pi update` naming pnpm for installs that only live under `PNPM_HOME`, for example a wrapper under `$PNPM_HOME/global-nub/`. pi said "not managed by a global pnpm install"; it now names the package manager only when that manager owns the install ([#7756](https://github.com/earendil-works/pi/issues/7756)) - Fixed image paste, copy, and opening the browser for sign-in failing silently on WSL when the working directory is deep. Windows programs such as `powershell.exe` and `explorer.exe` cannot start from a directory whose `\\wsl.localhost\...` path exceeds 260 characters; pi now starts these helpers from `/` ([#11592](https://github.com/earendil-works/pi/issues/11592)) - Fixed the external editor failing silently, for example when the editor command is wrong or a Windows editor cannot start from a deep WSL directory. pi now shows a status message such as `External editor exited with code 1; prompt unchanged` ([#11592](https://github.com/earendil-works/pi/issues/11592)) +- Fixed `tool_search` and codemode `searchTools()` ranking a tool below tools that only mention its name in their descriptions, for example `peekaboo app` ranking `mcp__peekaboo__app` 14th. A query word in a tool's name now scores more than any number of mentions in its description, a tool whose whole name is in the query gets a large bonus, and long server instructions no longer lower the scores of that server's tools ([#11331](https://github.com/earendil-works/pi/issues/11331)) ## [1.4.0] - 2026-10-21 diff --git a/packages/coding-agent/src/extensions/tool-search/tool.ts b/packages/coding-agent/src/extensions/tool-search/tool.ts index b9fa4ade0..0d51696eb 100644 --- a/packages/coding-agent/src/extensions/tool-search/tool.ts +++ b/packages/coding-agent/src/extensions/tool-search/tool.ts @@ -20,10 +20,13 @@ import type { export const TOOL_SEARCH_TOOL_NAME = "tool_search"; export const DEFAULT_TOOL_SEARCH_LIMIT = 8; -/** A tool as the ranker sees it: its name and the text built by {@link createToolSearchDocument}. */ +/** A tool as the ranker sees it: its name and the texts built by {@link createToolSearchDocument}. */ export interface ToolSearchDocument { name: string; + /** Description, schema descriptions, and property names. */ text: string; + /** Name, description, and instructions of the tool's namespace. */ + namespaceText?: string; } export interface ToolSearchMatch { @@ -101,18 +104,19 @@ function schemaText(schema: unknown, parts: string[]): void { } /** - * Search text of a tool: the name, the name with `_` - * as spaces, the description, schema descriptions and property names, and the namespace with its - * description and instructions. + * Search text of a tool: the description and the schema descriptions and property names, and + * separately the namespace name, description, and instructions. The ranker reads the name from `name`. */ export function createToolSearchDocument( tool: Pick, namespace?: ToolNamespace, ): ToolSearchDocument { - const parts = [tool.name, tool.name.replaceAll("_", " "), tool.description]; + const parts = [tool.description]; schemaText(tool.parameters, parts); - if (namespace) parts.push(namespace.name, namespace.description ?? "", namespace.instructions ?? ""); - return { name: tool.name, text: parts.filter((part) => part.trim()).join(" ") }; + const namespaceText = namespace + ? [namespace.name, namespace.description ?? "", namespace.instructions ?? ""].join(" ") + : ""; + return { name: tool.name, text: parts.filter((part) => part.trim()).join(" "), namespaceText }; } /** Lowercase with characters invalid in script identifiers as `_`, so `dev-radius` matches `dev_radius`. */ @@ -124,18 +128,63 @@ function normalizeName(name: string): string { } /** - * Whether `query` names the tool: its full name or a suffix after a `__` separator, so `thread_read` - * and `app__thread_read` both name `mcp__app__thread_read`. + * Whether the whole query is the tool's name: its full name or a suffix after a `__` separator, so + * `thread_read` and `app__thread_read` both name `mcp__app__thread_read`. */ -function isNameQuery(query: string, name: string): boolean { +function isExactNameQuery(query: string, name: string): boolean { const normalized = normalizeName(name); return normalized === query || normalized.endsWith(`__${query}`); } +/** A document split into the fields the ranker weighs differently. */ +interface IndexedDocument { + name: string; + /** Terms of the name, for example `mcp`, `app`, `thread`, `read` for `mcp__app__thread_read`. */ + nameTerms: ReadonlySet; + /** Terms of the last `__` segment of the name, the part models say: `thread`, `read`. */ + localNameTerms: readonly string[]; + textCounts: ReadonlyMap; + textLength: number; + namespaceTerms: ReadonlySet; +} + +function indexDocument(document: ToolSearchDocument): IndexedDocument { + const textCounts = new Map(); + for (const term of tokenize(document.text)) textCounts.set(term, (textCounts.get(term) ?? 0) + 1); + const segments = document.name.split("__"); + return { + name: document.name, + nameTerms: new Set(tokenize(document.name)), + localNameTerms: [...new Set(tokenize(segments[segments.length - 1]))], + textCounts, + textLength: [...textCounts.values()].reduce((sum, count) => sum + count, 0), + namespaceTerms: new Set(tokenize(document.namespaceText ?? "")), + }; +} + /** - * Okapi BM25 with the usual parameters. A query that names a tool ranks that tool first, because - * length normalization can otherwise rank a short description sharing the name's words above it. - * Ties keep document order. + * A query term in the name counts this many times its IDF. The text part of a term is at most + * `k1 + 1` (2.2) times its IDF and about 1 for a single mention, so a name match beats any number of + * mentions in the description. + */ +const NAME_WEIGHT = 3; + +/** + * A query term in the namespace name, description, or instructions counts this many times its IDF. They + * are shared by every tool of the server, so they help find the server but not the tool in it. + */ +const NAMESPACE_WEIGHT = 0.5; + +/** + * Okapi BM25 over the description, with the name and the namespace as separate fields: + * - a query term in the name scores more than any number of mentions in the description; + * - a tool whose whole local name is in the query (`app` in `peekaboo app` for + * `mcp__peekaboo__app`) gets a bonus worth one more match of every query term, so it mostly wins + * over tools that only mention the words; + * - a query that is exactly a tool's name ranks that tool first. + * + * The description alone is length-normalized, so a long name or long namespace instructions do + * not lower a tool's score. Ties keep document order. */ export class Bm25Ranker implements ToolRanker { private readonly k1: number; @@ -147,35 +196,40 @@ export class Bm25Ranker implements ToolRanker { } rank(query: string, documents: readonly ToolSearchDocument[], limit: number): ToolSearchMatch[] { - const queryTerms = [...new Set(tokenize(query))]; - if (queryTerms.length === 0 || documents.length === 0 || limit <= 0) return []; - const termCounts = documents.map((document) => { - const counts = new Map(); - for (const term of tokenize(document.text)) counts.set(term, (counts.get(term) ?? 0) + 1); - return counts; - }); - const lengths = termCounts.map((counts) => [...counts.values()].reduce((sum, count) => sum + count, 0)); - const averageLength = lengths.reduce((sum, length) => sum + length, 0) / documents.length || 1; + const queryTerms = new Set(tokenize(query)); + if (queryTerms.size === 0 || documents.length === 0 || limit <= 0) return []; + const indexed = documents.map(indexDocument); + const averageLength = indexed.reduce((sum, document) => sum + document.textLength, 0) / indexed.length || 1; const idf = new Map( - queryTerms.map((term) => { - const frequency = termCounts.filter((counts) => counts.has(term)).length; - return [term, Math.log(1 + (documents.length - frequency + 0.5) / (frequency + 0.5))] as const; + [...queryTerms].map((term) => { + const frequency = indexed.filter( + (document) => + document.nameTerms.has(term) || document.textCounts.has(term) || document.namespaceTerms.has(term), + ).length; + return [term, Math.log(1 + (indexed.length - frequency + 0.5) / (frequency + 0.5))] as const; }), ); - const nameQuery = normalizeName(query); - const matches: (ToolSearchMatch & { named: boolean })[] = []; - documents.forEach((document, index) => { + const fullNameBonus = [...idf.values()].reduce((sum, value) => sum + value, 0); + const exactQuery = normalizeName(query); + const matches: (ToolSearchMatch & { exact: boolean })[] = []; + for (const document of indexed) { + const norm = this.k1 * (1 - this.b + (this.b * document.textLength) / averageLength); let score = 0; for (const term of queryTerms) { - const count = termCounts[index].get(term); - if (!count) continue; - const norm = this.k1 * (1 - this.b + (this.b * lengths[index]) / averageLength); - score += (idf.get(term) ?? 0) * ((count * (this.k1 + 1)) / (count + norm)); + const count = document.textCounts.get(term) ?? 0; + const weight = + (document.nameTerms.has(term) ? NAME_WEIGHT : 0) + + (document.namespaceTerms.has(term) ? NAMESPACE_WEIGHT : 0) + + (count * (this.k1 + 1)) / (count + norm); + score += (idf.get(term) ?? 0) * weight; } - if (score > 0) matches.push({ name: document.name, score, named: isNameQuery(nameQuery, document.name) }); - }); + if (score === 0) continue; + const local = document.localNameTerms; + if (local.length > 0 && local.every((term) => queryTerms.has(term))) score += fullNameBonus; + matches.push({ name: document.name, score, exact: isExactNameQuery(exactQuery, document.name) }); + } return matches - .sort((a, b) => Number(b.named) - Number(a.named) || b.score - a.score) + .sort((a, b) => Number(b.exact) - Number(a.exact) || b.score - a.score) .slice(0, limit) .map(({ name, score }) => ({ name, score })); } diff --git a/packages/coding-agent/test/tool-search.test.ts b/packages/coding-agent/test/tool-search.test.ts index fd55b5c56..da5df3433 100644 --- a/packages/coding-agent/test/tool-search.test.ts +++ b/packages/coding-agent/test/tool-search.test.ts @@ -76,6 +76,48 @@ describe("Bm25Ranker", () => { expect(ranker.rank("read", named, 1).map((match) => match.name)).toEqual(["mcp__app__queue_read"]); }); + // https://github.com/earendil-works/pi/issues/11331 + it("ranks a name match above mentions in descriptions", () => { + const documents = [ + tool("mcp__ci__status", "Deploy status: last deploy, deploy log, deploy errors."), + tool( + "mcp__ci__deploy_release", + "Start a release of the current branch to an environment, run the migrations, switch traffic over, and wait until the new version is live and healthy.", + ), + ].map((entry) => createToolSearchDocument(entry)); + expect(new Bm25Ranker().rank("deploy", documents, 1).map((match) => match.name)).toEqual([ + "mcp__ci__deploy_release", + ]); + }); + + // https://github.com/earendil-works/pi/issues/11331 + it("ranks a tool first when the query contains its whole name", () => { + const peekaboo = { name: "mcp__peekaboo", instructions: "Peekaboo controls macOS apps by app name." }; + const documents = [ + tool("mcp__peekaboo__see", "Capture the app window of the peekaboo app and describe the app."), + tool("mcp__peekaboo__menu", "Click a menu item of an app."), + tool( + "mcp__peekaboo__app", + "Launch, quit, relaunch, focus, hide, or switch applications. Targets an application by name, bundle identifier, or process ID, and can wait until it is ready.", + ), + tool("mcp__app__queue_read", "Read a queued message of this thread."), + tool( + "mcp__dev-app__thread_read", + "Read durable state and a paginated timeline from any thread. The default view returns user messages, assistant messages and plans.", + ), + ].map((entry) => + createToolSearchDocument(entry, entry.name.startsWith("mcp__peekaboo__") ? peekaboo : undefined), + ); + const ranker = new Bm25Ranker(); + const first = (query: string) => ranker.rank(query, documents, 1).map((match) => match.name); + expect(first("peekaboo app")).toEqual(["mcp__peekaboo__app"]); + expect(first("open the peekaboo app tool")).toEqual(["mcp__peekaboo__app"]); + expect(first("describe app")).toEqual(["mcp__peekaboo__app"]); + expect(first("read thread_read state")).toEqual(["mcp__dev-app__thread_read"]); + // Without the whole name, mentions in descriptions decide. + expect(first("describe window")).toEqual(["mcp__peekaboo__see"]); + }); + it("includes the namespace in the search text", () => { const document = createToolSearchDocument(tool("mcp__x__run", "Run it."), { name: "mcp__x",