Skip to content
Open
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
2 changes: 2 additions & 0 deletions knip.json
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,8 @@
"!utils/WorkspacePathResolver.ts!",
"!test-utils/**!",
"!utils/versionToken.ts!",
"!services/tools/Bm25Ranker.ts!",
"!services/tools/types.ts!",
"!core/tools/helpers/toolResultFormatting.ts!"
],
"ignoreDependencies": [
Expand Down
105 changes: 105 additions & 0 deletions src/services/tools/Bm25Ranker.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,105 @@
import { Ranker, ToolDoc } from "./types"

type IndexedDocument = {
item: ToolDoc
termFrequencies: Map<string, number>
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[] = []

Check warning on line 14 in src/services/tools/Bm25Ranker.ts

View workflow job for this annotation

GitHub Actions / mutation-diff

Mutation test advisory

src/services/tools/Bm25Ranker.ts:14: Survived ArrayDeclaration mutant (replacement: ["Stryker was here"]). See the job summary for the complete list and resolution guidance.
private documentFrequency = new Map<string, number>()
private averageDocumentLength = 0

rank(query: string, items: ToolDoc[], k: number): ToolDoc[] {
const queryTerms = tokenize(query)

if (queryTerms.length === 0 || items.length === 0 || k <= 0) {

Check warning on line 21 in src/services/tools/Bm25Ranker.ts

View workflow job for this annotation

GitHub Actions / mutation-diff

Mutation test advisory

src/services/tools/Bm25Ranker.ts:21: 8 mutation test gaps; example: Survived ConditionalExpression mutant (replacement: false). See the job summary for the complete list and resolution guidance.
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)

Check warning on line 34 in src/services/tools/Bm25Ranker.ts

View workflow job for this annotation

GitHub Actions / mutation-diff

Mutation test advisory

src/services/tools/Bm25Ranker.ts:34: Survived ArithmeticOperator mutant (replacement: left.document.index + right.document.index). See the job summary for the complete list and resolution guidance.

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<string, number>()
this.documents = items.map((item, index) => {
const termFrequencies = countTerms(item)

for (const term of termFrequencies.keys()) {

Check warning on line 49 in src/services/tools/Bm25Ranker.ts

View workflow job for this annotation

GitHub Actions / mutation-diff

Mutation test advisory

src/services/tools/Bm25Ranker.ts:49: Survived BlockStatement mutant (replacement: {}). See the job summary for the complete list and resolution guidance.
this.documentFrequency.set(term, (this.documentFrequency.get(term) ?? 0) + 1)
}

return {
item,
termFrequencies,
length: Array.from(termFrequencies.values()).reduce((total, count) => total + count, 0),

Check warning on line 56 in src/services/tools/Bm25Ranker.ts

View workflow job for this annotation

GitHub Actions / mutation-diff

Mutation test advisory

src/services/tools/Bm25Ranker.ts:56: Survived ArithmeticOperator mutant (replacement: total - count). See the job summary for the complete list and resolution guidance.
index,
}
})
this.averageDocumentLength =
this.documents.reduce((total, document) => total + document.length, 0) / this.documents.length

Check warning on line 61 in src/services/tools/Bm25Ranker.ts

View workflow job for this annotation

GitHub Actions / mutation-diff

Mutation test advisory

src/services/tools/Bm25Ranker.ts:61: Survived ArithmeticOperator mutant (replacement: this.documents.reduce((total, document) => total + document.length, 0) * this.documents.length). See the job summary for the complete list and resolution guidance.
}

private score(document: IndexedDocument, queryTerms: Set<string>): 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

Check warning on line 75 in src/services/tools/Bm25Ranker.ts

View workflow job for this annotation

GitHub Actions / mutation-diff

Mutation test advisory

src/services/tools/Bm25Ranker.ts:75: Survived LogicalOperator mutant (replacement: this.documentFrequency.get(term) && 0). See the job summary for the complete list and resolution guidance.
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
.normalize("NFC")
.toLowerCase()
.split(/[^\p{L}\p{M}\p{N}]+/u)
.filter(Boolean)
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Resolved
Comment thread
edelauna marked this conversation as resolved.
Resolved

function countTerms(item: ToolDoc): Map<string, number> {
const frequencies = new Map<string, number>()
const text = `${item.serverName} ${item.toolName} ${item.description}`

for (const term of tokenize(text)) {
frequencies.set(term, (frequencies.get(term) ?? 0) + 1)
}

return frequencies
}
96 changes: 96 additions & 0 deletions src/services/tools/__tests__/Bm25Ranker.spec.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,96 @@
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("postgres", "query", "Send data and query a Postgres database"),
tool("jira", "createIssue", "Send messages and create Jira issues"),
tool("slack", "postMessage", "Send a Slack message"),
]

expect(new Bm25Ranker().rank("send slack message", items, 3)).toEqual([items[2], items[1], items[0]])
})

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("splits underscore-separated tool names into separate terms", () => {
const matching = tool("github", "create_issue", "Open a ticket")
const nonMatching = tool("calendar", "list_events", "Show upcoming events")

expect(new Bm25Ranker().rank("create issue", [nonMatching, matching], 10)).toEqual([matching])
})

it("preserves Unicode letters and numbers in tokens", () => {
const russian = tool("погода-сервер", "прогноз", "Прогноз погоды")
const arabic = tool("أدوات", "مقياس", "يدعم ٣ بروتوكولات")

expect(new Bm25Ranker().rank("погода", [arabic, russian], 10)).toEqual([russian])
expect(new Bm25Ranker().rank("٣", [russian, arabic], 10)).toEqual([arabic])
})

it("matches composed and decomposed forms of the same accented character", () => {
const decomposed = tool("café", "menu", "Liste du café")
const composed = tool("café", "menu", "Menu")
const other = tool("calendar", "events", "Show events")

expect(new Bm25Ranker().rank("café", [other, decomposed], 10)).toEqual([decomposed])
expect(new Bm25Ranker().rank("café", [other, composed], 10)).toEqual([composed])
})

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]])
})
})
9 changes: 9 additions & 0 deletions src/services/tools/types.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
export type ToolDoc = {
serverName: string
toolName: string
description: string
}

export interface Ranker {
rank(query: string, items: ToolDoc[], k: number): ToolDoc[]
}
Loading