Skip to content

feat(mcp): rank mcp_tool_search results with BM25 scoring - #33949

Open
chasecheese wants to merge 4 commits into
BerriAI:litellm_oss_daily_2026_07_17from
chasecheese:litellm_mcp_tool_search_bm25_scoring
Open

feat(mcp): rank mcp_tool_search results with BM25 scoring#33949
chasecheese wants to merge 4 commits into
BerriAI:litellm_oss_daily_2026_07_17from
chasecheese:litellm_mcp_tool_search_bm25_scoring

Conversation

@chasecheese

@chasecheese chasecheese commented Jul 20, 2026

Copy link
Copy Markdown

Relevant issues

N/A

Linear ticket

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

  • I have added meaningful tests
  • My PR passes all CI/CD checks (e.g., lint, format, unit tests)
  • My PR's scope is as isolated as possible; it only solves 1 specific problem
  • I have received a Greptile Confidence Score of at least 4/5 before requesting a maintainer review (Greptile reviews automatically once the PR is opened; only comment @greptileai to 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_tools at 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 function mcp_tool_search ranks 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 changes

query: 'add'
  1ebf2a78a9 (base): ['address-book-lookup', 'address-book-update', 'math-add']
  3aecb33ec8 (PR):   ['math-add', 'address-book-lookup', 'address-book-update']
query: 'add numbers'
  1ebf2a78a9 (base): ['math-add', 'address-book-lookup', 'address-book-update', 'math-multiply']
  3aecb33ec8 (PR):   ['math-add', 'math-multiply', 'address-book-lookup', 'address-book-update']
query: 'list issues'
  1ebf2a78a9 (base): ['jira-list_issues', 'github-list_repos', 'slack-list_channels', 'files-list', 'calendar-list_events']
  3aecb33ec8 (PR):   ['jira-list_issues', 'files-list', 'calendar-list_events', 'github-list_repos', 'slack-list_channels']
query: 'create issue github'
  1ebf2a78a9 (base): ['github-create_issue', 'jira-create_issue', 'github-list_repos', 'github-search_code', 'jira-list_issues']
  3aecb33ec8 (PR):   ['github-create_issue', 'jira-create_issue', 'github-search_code', 'notion-create_page', 'github-list_repos']

To reproduce end to end on a live proxy, use the setup from #31777 (a key with mcp_tool_search_enabled and an MCP server) and call POST /mcp-rest/tools/call with {"name": "mcp_tool_search", "arguments": {"query": "add"}}; the response shape is identical to before, with the ranking shown above

Type

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 for add matches address by 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 deterministic

This 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 keeps channel matching channels while math-ookup for add. 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 requests
This also fixes a small bug where top_k <= 0 fed a negative value into a Python slice and silently dropped tools from the end of the results; it now returns an empty list

Queries 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.py pass 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 positive top_k returns empty, repeated query tokens are not double counted, and query tokens past the 32 token cap are ignored

Final Attestation

  • The tests check the right things, including the edge cases, and regressions in the respective real-world customer uer this PR

@CLAassistant

CLAassistant commented Jul 20, 2026

Copy link
Copy Markdown

CLA assistant check
All committers have signed the CLA.

@greptile-apps

greptile-apps Bot commented Jul 20, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR replaces the original substring-count scoring in mcp_tool_search with an on-the-fly BM25 ranking computed against the per-request, permission-filtered tool catalog. It also fixes a latent bug where top_k ≤ 0 produced a silent incorrect result via Python's negative-slice semantics.

  • BM25 implementation: tokenization splits on any non-alphanumeric character; exact token hits score 1.0, prefix hits 0.3, name tokens weight 3× description tokens; IDF is derived from document frequency in the live catalog, so rare terms outrank common ones without any persistent index or new dependency.
  • Security bound: queries are capped at 32 unique tokens before scoring begins, keeping per-request CPU work proportional to the catalog size rather than attacker-controlled input length.
  • Test coverage: 8 new unit tests pin the new ranking properties; all 45 existing tests pass without modification.

Confidence Score: 5/5

Safe 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.

Important Files Changed

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

Comment thread litellm/proxy/_experimental/mcp_server/tool_search.py
@veria-ai

veria-ai Bot commented Jul 20, 2026

Copy link
Copy Markdown
Contributor

PR overview

All previously flagged issues have been addressed. No open security concerns remain on this pull request.

Security review

No open security issues remain on this pull request.

Fixed/addressed: 1 · PR risk: 0/10

@chasecheese chasecheese reopened this Jul 20, 2026
@chasecheese
chasecheese changed the base branch from litellm_internal_staging to litellm_oss_daily_2026_07_17 July 20, 2026 04:26
@codecov

codecov Bot commented Jul 20, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

@chasecheese

Copy link
Copy Markdown
Author

@greptileai

@codspeed-hq

codspeed-hq Bot commented Jul 20, 2026

Copy link
Copy Markdown
Contributor

Merging this PR will not alter performance

✅ 31 untouched benchmarks


Comparing chasecheese:litellm_mcp_tool_search_bm25_scoring (3aecb33) with litellm_oss_daily_2026_07_17 (1ebf2a7)

Open in CodSpeed

@chasecheese
chasecheese requested a review from a team July 20, 2026 05:06
@chasecheese

Copy link
Copy Markdown
Author

osv-scan also fails on the base branch head (1ebf2a7); this PR changes no dependencies. The advisory needs a separate dependency bump PR

@chasecheese

Copy link
Copy Markdown
Author

@greptileai

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants