From 3aecb33ec82dfa48c58f6157bb6415adcf5241ff Mon Sep 17 00:00:00 2001 From: chasecheese Date: Mon, 20 Jul 2026 11:02:24 +0800 Subject: [PATCH 1/4] feat(mcp): rank mcp_tool_search results with BM25 scoring --- .../_experimental/mcp_server/tool_search.py | 52 +++++++++++++++--- .../mcp_server/test_mcp_tool_search.py | 54 +++++++++++++++++++ 2 files changed, 99 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/tool_search.py b/litellm/proxy/_experimental/mcp_server/tool_search.py index fa57a2b3eb21..095df503c8fb 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_search.py +++ b/litellm/proxy/_experimental/mcp_server/tool_search.py @@ -1,6 +1,8 @@ from __future__ import annotations import json +import math +import re from datetime import datetime from typing import TYPE_CHECKING, Any, Optional @@ -13,6 +15,12 @@ MCP_TOOL_SEARCH_TOOL_NAME: str = "mcp_tool_search" MCP_TOOL_CALL_TOOL_NAME: str = "mcp_tool_call" +_TOKEN_RE = re.compile(r"[a-z0-9]+") +_BM25_K1 = 1.5 +_BM25_B = 0.75 +_NAME_WEIGHT = 3.0 +_PREFIX_WEIGHT = 0.3 + def coerce_top_k(value: Any, default: int = 5) -> int: try: @@ -21,17 +29,47 @@ def coerce_top_k(value: Any, default: int = 5) -> int: return default +def _tokenize(text: str) -> tuple[str, ...]: + return tuple(_TOKEN_RE.findall(text.lower())) + + +def _field_tf(query_token: str, field_tokens: tuple[str, ...]) -> float: + return sum( + 1.0 if token == query_token else _PREFIX_WEIGHT if token.startswith(query_token) else 0.0 + for token in field_tokens + ) + + def search_tools(query: str, tools: list[dict[str, Any]], top_k: int = 5) -> list[dict[str, Any]]: - if not query: + query_tokens = tuple(dict.fromkeys(_tokenize(query))) + if not query_tokens or not tools or top_k <= 0: return [] - tokens = query.lower().split() + docs = tuple( + (tool, _tokenize(str(tool.get("name", ""))), _tokenize(str(tool.get("description") or ""))) for tool in tools + ) + tf_rows = tuple( + tuple(_NAME_WEIGHT * _field_tf(token, name_tokens) + _field_tf(token, desc_tokens) for token in query_tokens) + for _, name_tokens, desc_tokens in docs + ) + doc_lengths = tuple(_NAME_WEIGHT * len(name_tokens) + len(desc_tokens) for _, name_tokens, desc_tokens in docs) + avg_doc_length = (sum(doc_lengths) / len(doc_lengths)) or 1.0 + doc_count = len(docs) + idfs = tuple( + math.log(1.0 + (doc_count - df + 0.5) / (df + 0.5)) + for df in (sum(1 for row in tf_rows if row[i] > 0.0) for i in range(len(query_tokens))) + ) - def _score(tool: dict[str, Any]) -> int: - haystack = (tool.get("name", "") + " " + tool.get("description", "")).lower() - return sum(1 for t in tokens if t in haystack) + def _bm25(row: tuple[float, ...], doc_length: float) -> float: + norm = _BM25_K1 * (1.0 - _BM25_B + _BM25_B * doc_length / avg_doc_length) + return sum(idf * tf * (_BM25_K1 + 1.0) / (tf + norm) for idf, tf in zip(idfs, row) if tf > 0.0) - scored = ((s, tool) for tool in tools if (s := _score(tool)) > 0) - return [tool for _, tool in sorted(scored, key=lambda x: x[0], reverse=True)[:top_k]] + scored = tuple( + (score, tool) + for (tool, _, _), row, doc_length in zip(docs, tf_rows, doc_lengths) + if (score := _bm25(row, doc_length)) > 0.0 + ) + ranked = sorted(scored, key=lambda item: (-item[0], str(item[1].get("name", "")))) + return [tool for _, tool in ranked[:top_k]] def get_virtual_tool_definitions() -> list[dict[str, Any]]: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py index b8f0b205831d..5e7145e59dd7 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py @@ -111,6 +111,60 @@ def test_result_tools_have_full_schema(self) -> None: assert "description" in tool assert "inputSchema" in tool + def test_exact_word_outranks_prefix_false_positive(self) -> None: + tools = _make_tools( + [ + ("address-lookup", "Look up a contact in the address book"), + ("math-add", "Add two numbers together"), + ] + ) + results = search_tools("add", tools) + assert results[0]["name"] == "math-add" + + def test_name_match_outranks_description_match(self) -> None: + tools = _make_tools( + [ + ("files-list", "Search the file index"), + ("code-search", "Find matches in the repository"), + ] + ) + results = search_tools("search", tools) + assert results[0]["name"] == "code-search" + + def test_rare_token_outranks_common_token(self) -> None: + tools = _make_tools( + [ + ("files-list", "List all files"), + ("repos-list", "List repositories"), + ("channels-list", "List channels"), + ("email-send", "Send an email message"), + ] + ) + results = search_tools("list email", tools) + assert results[0]["name"] == "email-send" + + def test_tie_break_is_deterministic_by_name(self) -> None: + specs = [ + ("b-tool", "Fetch the weather"), + ("a-tool", "Fetch the weather"), + ] + forward = [t["name"] for t in search_tools("weather", _make_tools(specs))] + reverse = [t["name"] for t in search_tools("weather", _make_tools(specs[::-1]))] + assert forward == reverse == ["a-tool", "b-tool"] + + def test_hyphen_and_underscore_are_token_boundaries(self) -> None: + results = search_tools("create_issue", SAMPLE_TOOLS) + assert results[0]["name"] == "github-create_issue" + + def test_non_positive_top_k_returns_empty(self) -> None: + assert search_tools("github", SAMPLE_TOOLS, top_k=0) == [] + assert search_tools("github", SAMPLE_TOOLS, top_k=-1) == [] + + def test_repeated_query_token_not_double_counted(self) -> None: + single = [t["name"] for t in search_tools("github issue", SAMPLE_TOOLS)] + repeated = [t["name"] for t in search_tools("github github issue", SAMPLE_TOOLS)] + assert single == repeated + class TestGetVirtualToolDefinitions: def test_returns_two_tools(self) -> None: From 7bd49475ce853670902c987a942a21f3fb0cba73 Mon Sep 17 00:00:00 2001 From: chasecheese Date: Mon, 20 Jul 2026 12:36:04 +0800 Subject: [PATCH 2/4] fix(mcp): cap mcp_tool_search query tokens to bound scoring cost --- litellm/proxy/_experimental/mcp_server/tool_search.py | 3 ++- .../proxy/_experimental/mcp_server/test_mcp_tool_search.py | 6 ++++++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_experimental/mcp_server/tool_search.py b/litellm/proxy/_experimental/mcp_server/tool_search.py index 095df503c8fb..ab30fc599479 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_search.py +++ b/litellm/proxy/_experimental/mcp_server/tool_search.py @@ -20,6 +20,7 @@ _BM25_B = 0.75 _NAME_WEIGHT = 3.0 _PREFIX_WEIGHT = 0.3 +_MAX_QUERY_TOKENS = 32 def coerce_top_k(value: Any, default: int = 5) -> int: @@ -41,7 +42,7 @@ def _field_tf(query_token: str, field_tokens: tuple[str, ...]) -> float: def search_tools(query: str, tools: list[dict[str, Any]], top_k: int = 5) -> list[dict[str, Any]]: - query_tokens = tuple(dict.fromkeys(_tokenize(query))) + query_tokens = tuple(dict.fromkeys(_tokenize(query)))[:_MAX_QUERY_TOKENS] if not query_tokens or not tools or top_k <= 0: return [] docs = tuple( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py index 5e7145e59dd7..cd0831438e38 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py @@ -165,6 +165,12 @@ def test_repeated_query_token_not_double_counted(self) -> None: repeated = [t["name"] for t in search_tools("github github issue", SAMPLE_TOOLS)] assert single == repeated + def test_query_tokens_capped_at_32(self) -> None: + filler = " ".join(f"junk{i}" for i in range(32)) + assert search_tools(f"{filler} github", SAMPLE_TOOLS) == [] + capped = [t["name"] for t in search_tools(f"github {filler}", SAMPLE_TOOLS)] + assert "github-create_issue" in capped + class TestGetVirtualToolDefinitions: def test_returns_two_tools(self) -> None: From 8f9478cc49ed6a84fc93add91f37826c38fb3e1c Mon Sep 17 00:00:00 2001 From: chasecheese Date: Mon, 20 Jul 2026 13:06:54 +0800 Subject: [PATCH 3/4] fix(deps): bump mcp to 1.28.1 to clear osv-scan advisories --- uv.lock | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/uv.lock b/uv.lock index 8dfc4bd5fcc0..b1a04c99bc07 100644 --- a/uv.lock +++ b/uv.lock @@ -10,7 +10,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-07-13T21:03:01.672393Z" +exclude-newer = "2026-07-17T04:53:41.725625623Z" exclude-newer-span = "P3D" [manifest] @@ -4431,7 +4431,7 @@ wheels = [ [[package]] name = "mcp" -version = "1.26.0" +version = "1.28.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anyio" }, @@ -4449,9 +4449,9 @@ dependencies = [ { name = "typing-inspection" }, { name = "uvicorn", marker = "sys_platform != 'emscripten'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/fc/6d/62e76bbb8144d6ed86e202b5edd8a4cb631e7c8130f3f4893c3f90262b10/mcp-1.26.0.tar.gz", hash = "sha256:db6e2ef491eecc1a0d93711a76f28dec2e05999f93afd48795da1c1137142c66", size = 608005, upload-time = "2026-01-24T19:40:32.468Z" } +sdist = { url = "https://files.pythonhosted.org/packages/6e/77/9450b8f251a13affb6281997d0523c4615f8a8b35d0b21ff30db3a5aac9d/mcp-1.28.1.tar.gz", hash = "sha256:d51e36a5f5644faea4f85ea649bfffa6bc6c26770d42798ad6a3de3d2ba69683", size = 638501, upload-time = "2026-06-26T12:57:29.093Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/fd/d9/eaa1f80170d2b7c5ba23f3b59f766f3a0bb41155fbc32a69adfa1adaaef9/mcp-1.26.0-py3-none-any.whl", hash = "sha256:904a21c33c25aa98ddbeb47273033c435e595bbacfdb177f4bd87f6dceebe1ca", size = 233615, upload-time = "2026-01-24T19:40:30.652Z" }, + { url = "https://files.pythonhosted.org/packages/e2/5e/d118fce19f87a2e7d8101c35c8ae0ec289098a4df0ff244cec23e415aca0/mcp-1.28.1-py3-none-any.whl", hash = "sha256:2726bca5e7193f61c5dde8b12500a6de2d9acf6d1a1c0be9e8c2e706437991df", size = 222620, upload-time = "2026-06-26T12:57:27.218Z" }, ] [[package]] From 80eba1a18967b98ec828976be9e8d2126d00b28b Mon Sep 17 00:00:00 2001 From: chasecheese Date: Mon, 20 Jul 2026 13:31:10 +0800 Subject: [PATCH 4/4] Revert "fix(deps): bump mcp to 1.28.1 to clear osv-scan advisories" This reverts commit 8f9478cc49ed6a84fc93add91f37826c38fb3e1c. --- uv.lock | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/uv.lock b/uv.lock index b1a04c99bc07..8dfc4bd5fcc0 100644 --- a/uv.lock +++ b/uv.lock @@ -10,7 +10,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-07-17T04:53:41.725625623Z" +exclude-newer = "2026-07-13T21:03:01.672393Z" exclude-newer-span = "P3D" [manifest] @@ -4431,7 +4431,7 @@ wheels = [ [[package]] name = "mcp" -version = "1.28.1" +version = "1.26.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anyio" }, @@ -4449,9 +4449,9 @@ dependencies = [ { name = "typing-inspection" }, { name = "uvicorn", marker = "sys_platform != 'emscripten'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/6e/77/9450b8f251a13affb6281997d0523c4615f8a8b35d0b21ff30db3a5aac9d/mcp-1.28.1.tar.gz", hash = "sha256:d51e36a5f5644faea4f85ea649bfffa6bc6c26770d42798ad6a3de3d2ba69683", size = 638501, upload-time = "2026-06-26T12:57:29.093Z" } +sdist = { url = "https://files.pythonhosted.org/packages/fc/6d/62e76bbb8144d6ed86e202b5edd8a4cb631e7c8130f3f4893c3f90262b10/mcp-1.26.0.tar.gz", hash = "sha256:db6e2ef491eecc1a0d93711a76f28dec2e05999f93afd48795da1c1137142c66", size = 608005, upload-time = "2026-01-24T19:40:32.468Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/e2/5e/d118fce19f87a2e7d8101c35c8ae0ec289098a4df0ff244cec23e415aca0/mcp-1.28.1-py3-none-any.whl", hash = "sha256:2726bca5e7193f61c5dde8b12500a6de2d9acf6d1a1c0be9e8c2e706437991df", size = 222620, upload-time = "2026-06-26T12:57:27.218Z" }, + { url = "https://files.pythonhosted.org/packages/fd/d9/eaa1f80170d2b7c5ba23f3b59f766f3a0bb41155fbc32a69adfa1adaaef9/mcp-1.26.0-py3-none-any.whl", hash = "sha256:904a21c33c25aa98ddbeb47273033c435e595bbacfdb177f4bd87f6dceebe1ca", size = 233615, upload-time = "2026-01-24T19:40:30.652Z" }, ] [[package]]