diff --git a/litellm/proxy/_experimental/mcp_server/tool_search.py b/litellm/proxy/_experimental/mcp_server/tool_search.py index fa57a2b3eb21..ab30fc599479 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,13 @@ 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 +_MAX_QUERY_TOKENS = 32 + def coerce_top_k(value: Any, default: int = 5) -> int: try: @@ -21,17 +30,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)))[:_MAX_QUERY_TOKENS] + 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..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 @@ -111,6 +111,66 @@ 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 + + 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: