feat(mcp): rank mcp_tool_search results with BM25 scoring - #33949
feat(mcp): rank mcp_tool_search results with BM25 scoring#33949chasecheese wants to merge 4 commits into
Conversation
Greptile SummaryThis PR replaces the original substring-count scoring in
Confidence Score: 5/5Safe to merge — the change is self-contained within the experimental MCP tool-search module, introduces no new dependencies, and the wire contract is unchanged. The BM25 math is correct, all edge cases (empty query, zero-length docs, non-positive top_k, single-doc catalog, negative IDF) are handled, the query-token cap bounds per-request work, and 8 new regression tests accompany 45 unchanged existing ones. No existing test assertions were weakened. No files require special attention.
|
| Filename | Overview |
|---|---|
| litellm/proxy/_experimental/mcp_server/tool_search.py | Replaces substring-count scoring with a correct BM25 implementation; the math is sound, edge cases (empty query, top_k ≤ 0, zero-length docs, single-doc catalog) are all handled, and the query-token cap of 32 properly bounds per-request work. |
| tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py | Eight new regression tests added alongside 45 unchanged existing tests; each new test pins a distinct behavioral guarantee (exact-vs-prefix ranking, name vs description weight, IDF effect, tie determinism, token boundary handling, top_k edge cases, deduplication, and the 32-token cap) without weakening any prior assertion. |
Reviews (4): Last reviewed commit: "Revert "fix(deps): bump mcp to 1.28.1 to..." | Re-trigger Greptile
PR overviewAll previously flagged issues have been addressed. No open security concerns remain on this pull request. Security reviewNo open security issues remain on this pull request. Fixed/addressed: 1 · PR risk: 0/10 |
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
This reverts commit 8f9478c.
|
osv-scan also fails on the base branch head (1ebf2a7); this PR changes no dependencies. The advisory needs a separate dependency bump PR |
Relevant issues
N/A
Linear ticket
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
@greptileaito re-request a review after pushing changes)Delays in PR merge?
If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).
Screenshots / Proof of Fix
Ranking comparison of
search_toolsat the base commit (1ebf2a7) vs this PR (3aecb33), over the same 15 tool catalog spanning address book, math, GitHub, Slack, Jira, files, email, calendar and Notion tools. This is the exact functionmcp_tool_searchranks the permission filtered catalog with; the wire contract of the virtual tool (arguments, response shape) is unchanged, so only the ordering and cutoff of results changesTo reproduce end to end on a live proxy, use the setup from #31777 (a key with
mcp_tool_search_enabledand an MCP server) and callPOST /mcp-rest/tools/callwith{"name": "mcp_tool_search", "arguments": {"query": "add"}}; the response shape is identical to before, with the ranking shown aboveType
New Feature
Changes
mcp_tool_search(added in #31777) ranked tools by counting how many query tokens appear as substrings of the tool's name plus description. That breaks down in three ways: a query foraddmatchesaddressby substring, every hit counts the same no matter how common the word is across the catalog, and tied scores fall back to server enumeration order, so results are not deterministicThis PR replaces that scoring with BM25 computed per request over the same permission filtered catalog. Tokenization splits on anything outside
[a-z0-9], so hyphens, underscores and punctuation are word boundaries. A query token scores a full hit on an exact toand a discounted hit (0.3) when it is a prefix of a catalog token, which keepschannelmatchingchannelswhilemath-ookupforadd. Name tokens weigh 3x description tokens. IDF is computed on the fly from the per request catalog, sorare terms outrank common ones without any index, cache or new dependency. Ties break by tool name, making results deterministic across requestsThis also fixes a small bug where
top_k <= 0fed a negative value into a Python slice and silently dropped tools from the end of the results; it now returns an empty listQueries are capped at the first 32 unique tokens, keeping the per request scoring work bounded by the catalog size rather than by attacker controlled query length
The existing 45 tests in
test_mcp_tool_search.pypass unchanged. Eight new regression tests pin the new behavior: exact word beats substring false positive, name hits outrank description hits, rare tokens outrank common ones, ties are deterministic, hyphen and underscore are token boundaries, non positivetop_kreturns empty, repeated query tokens are not double counted, and query tokens past the 32 token cap are ignoredFinal Attestation