diff --git a/src/services/tools/Bm25Ranker.ts b/src/services/tools/Bm25Ranker.ts new file mode 100644 index 0000000000..c90f5a874d --- /dev/null +++ b/src/services/tools/Bm25Ranker.ts @@ -0,0 +1,101 @@ +import { Ranker, ToolDoc } from "./types" + +type IndexedDocument = { + item: ToolDoc + termFrequencies: Map + length: number + index: number +} + +export class Bm25Ranker implements Ranker { + private readonly k1 = 1.5 + private readonly b = 0.75 + private indexedItems: ToolDoc[] | undefined + private documents: IndexedDocument[] = [] + private documentFrequency = new Map() + private averageDocumentLength = 0 + + rank(query: string, items: ToolDoc[], k: number): ToolDoc[] { + const queryTerms = tokenize(query) + + if (queryTerms.length === 0 || items.length === 0 || k <= 0) { + return [] + } + + this.ensureIndex(items) + + const scores = this.documents + .map((document) => ({ + document, + score: this.score(document, new Set(queryTerms)), + })) + .filter(({ score }) => score > 0) + + scores.sort((left, right) => right.score - left.score || left.document.index - right.document.index) + + return scores.slice(0, k).map(({ document }) => document.item) + } + + private ensureIndex(items: ToolDoc[]): void { + if (this.indexedItems === items) { + return + } + + this.indexedItems = items + this.documentFrequency = new Map() + this.documents = items.map((item, index) => { + const termFrequencies = countTerms(item) + + for (const term of termFrequencies.keys()) { + this.documentFrequency.set(term, (this.documentFrequency.get(term) ?? 0) + 1) + } + + return { + item, + termFrequencies, + length: Array.from(termFrequencies.values()).reduce((total, count) => total + count, 0), + index, + } + }) + this.averageDocumentLength = + this.documents.reduce((total, document) => total + document.length, 0) / this.documents.length + } + + private score(document: IndexedDocument, queryTerms: Set): number { + let score = 0 + const documentCount = this.documents.length + + for (const term of queryTerms) { + const termFrequency = document.termFrequencies.get(term) + + if (!termFrequency) { + continue + } + + const documentFrequency = this.documentFrequency.get(term) ?? 0 + const inverseDocumentFrequency = Math.log( + 1 + (documentCount - documentFrequency + 0.5) / (documentFrequency + 0.5), + ) + const normalization = this.k1 * (1 - this.b + this.b * (document.length / this.averageDocumentLength)) + + score += inverseDocumentFrequency * ((termFrequency * (this.k1 + 1)) / (termFrequency + normalization)) + } + + return score + } +} + +function tokenize(value: string): string[] { + return value.toLowerCase().split(/\W+/).filter(Boolean) +} + +function countTerms(item: ToolDoc): Map { + const frequencies = new Map() + const text = `${item.serverName} ${item.toolName} ${item.description}` + + for (const term of tokenize(text)) { + frequencies.set(term, (frequencies.get(term) ?? 0) + 1) + } + + return frequencies +} diff --git a/src/services/tools/__tests__/Bm25Ranker.spec.ts b/src/services/tools/__tests__/Bm25Ranker.spec.ts new file mode 100644 index 0000000000..f93ffb9727 --- /dev/null +++ b/src/services/tools/__tests__/Bm25Ranker.spec.ts @@ -0,0 +1,72 @@ +import { describe, expect, it } from "vitest" +import { Bm25Ranker } from "../Bm25Ranker" +import { ToolDoc } from "../types" + +const tool = (serverName: string, toolName: string, description: string): ToolDoc => ({ + serverName, + toolName, + description, +}) + +describe("Bm25Ranker", () => { + it("ranks matching MCP tools by BM25 relevance", () => { + const items = [ + tool("slack", "postMessage", "Send a Slack message"), + tool("jira", "createIssue", "Send messages and create Jira issues"), + tool("postgres", "query", "Send data and query a Postgres database"), + ] + + expect(new Bm25Ranker().rank("send slack message", items, 3)).toEqual([items[0], items[1], items[2]]) + }) + + it("returns an empty result for empty queries, empty items, and no matches", () => { + const ranker = new Bm25Ranker() + const items = [tool("slack", "postMessage", "Send a message")] + + expect(ranker.rank(" ", items, 10)).toEqual([]) + expect(ranker.rank("message", [], 10)).toEqual([]) + expect(ranker.rank("unrelated", items, 10)).toEqual([]) + }) + + it("truncates results to the requested top-k", () => { + const items = [ + tool("one", "search", "Search records"), + tool("two", "search", "Search records"), + tool("three", "search", "Search records"), + ] + + expect(new Bm25Ranker().rank("search", items, 2)).toEqual(items.slice(0, 2)) + expect(new Bm25Ranker().rank("search", items, 0)).toEqual([]) + }) + + it("preserves input order for equal scores", () => { + const items = [tool("first", "lookup", "Find a record"), tool("second", "lookup", "Find a record")] + + expect(new Bm25Ranker().rank("record", items, 10)).toEqual(items) + }) + + it("tokenizes case-insensitively around punctuation", () => { + const matching = tool("Slack-Server", "POST.Message", "Send a message") + const nonMatching = tool("calendar", "createEvent", "Create a calendar event") + + expect(new Bm25Ranker().rank("SLACK post message", [nonMatching, matching], 10)).toEqual([matching]) + }) + + it("handles tools with empty descriptions", () => { + const item = tool("GitHub", "listRepos", "") + + expect(new Bm25Ranker().rank("github repos", [item], 10)).toEqual([item]) + }) + + it("rebuilds only when the items array reference changes", () => { + const items = [tool("slack", "send", "Send a message")] + const ranker = new Bm25Ranker() + + ranker.rank("send", items, 10) + items.push(tool("jira", "create", "Create an issue")) + expect(ranker.rank("issue", items, 10)).toEqual([]) + + const newItems = items.slice() + expect(ranker.rank("issue", newItems, 10)).toEqual([newItems[1]]) + }) +}) diff --git a/src/services/tools/types.ts b/src/services/tools/types.ts new file mode 100644 index 0000000000..86170c8aa8 --- /dev/null +++ b/src/services/tools/types.ts @@ -0,0 +1,9 @@ +export type ToolDoc = { + serverName: string + toolName: string + description: string +} + +export interface Ranker { + rank(query: string, items: ToolDoc[], k: number): ToolDoc[] +}