From 5d19cb6ee28d79572bd089d58582ffcd0009d3b7 Mon Sep 17 00:00:00 2001 From: Hashim1999164 Date: Thu, 16 Jul 2026 01:28:59 +0500 Subject: [PATCH] fix(mcp): make mcp_tool_search default top_k configurable Allow operators to set a global litellm_settings default or per-key object_permission override instead of always defaulting to 5. Fixes #33440 --- .../migration.sql | 2 + .../litellm_proxy_extras/schema.prisma | 1 + litellm/models/object_permission.py | 1 + .../mcp_server/rest_endpoints.py | 11 +++- .../proxy/_experimental/mcp_server/server.py | 12 +++- .../_experimental/mcp_server/tool_search.py | 48 +++++++++++++-- litellm/proxy/_types.py | 1 + litellm/proxy/schema.prisma | 1 + litellm/types/object_permission.py | 1 + schema.prisma | 1 + .../mcp_server/test_mcp_tool_search.py | 60 +++++++++++++++++++ 11 files changed, 129 insertions(+), 10 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260716000000_add_mcp_tool_search_top_k/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260716000000_add_mcp_tool_search_top_k/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260716000000_add_mcp_tool_search_top_k/migration.sql new file mode 100644 index 00000000000..94240b245ea --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260716000000_add_mcp_tool_search_top_k/migration.sql @@ -0,0 +1,2 @@ +-- AlterTable +ALTER TABLE "LiteLLM_ObjectPermissionTable" ADD COLUMN "mcp_tool_search_top_k" INTEGER; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index a23cecc3911..faa66fae557 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -280,6 +280,7 @@ model LiteLLM_ObjectPermissionTable { mcp_toolsets String[] @default([]) // Toolset IDs granted to this key/team/user search_tools String[] @default([]) // search_tool_name values this key/team/user may call mcp_tool_search_enabled Boolean? + mcp_tool_search_top_k Int? teams LiteLLM_TeamTable[] projects LiteLLM_ProjectTable[] verification_tokens LiteLLM_VerificationToken[] diff --git a/litellm/models/object_permission.py b/litellm/models/object_permission.py index 3052a2af459..ec9b6671e8b 100644 --- a/litellm/models/object_permission.py +++ b/litellm/models/object_permission.py @@ -25,3 +25,4 @@ class LiteLLM_ObjectPermissionTable(LiteLLMPydanticObjectBase): blocked_tools: Optional[List[str]] = [] search_tools: Optional[List[str]] = [] mcp_tool_search_enabled: Optional[bool] = None + mcp_tool_search_top_k: Optional[int] = None diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 111fde86ea0..0511c574801 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -139,9 +139,11 @@ if MCP_AVAILABLE: ) from litellm.proxy._experimental.mcp_server.tool_search import ( MCP_TOOL_SEARCH_TOOL_NAME, - coerce_top_k, + get_mcp_tool_search_default_top_k, + get_virtual_tool_definitions, handle_mcp_tool_call, handle_mcp_tool_search, + resolve_mcp_tool_search_top_k, ) from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.proxy_server import general_settings, proxy_config, proxy_logging_obj @@ -162,7 +164,7 @@ if MCP_AVAILABLE: if tool_name == MCP_TOOL_SEARCH_TOOL_NAME: return await handle_mcp_tool_search( query=tool_arguments.get("query", ""), - top_k=coerce_top_k(tool_arguments.get("top_k", 5)), + top_k=resolve_mcp_tool_search_top_k(tool_arguments.get("top_k"), user_api_key_dict), user_api_key_dict=user_api_key_dict, client_ip=rest_client_ip, mcp_auth_header=virtual_mcp_auth_header, @@ -733,11 +735,14 @@ if MCP_AVAILABLE: ) ): from litellm.proxy._experimental.mcp_server.tool_search import ( + get_mcp_tool_search_default_top_k, get_virtual_tool_definitions, ) return { - "tools": get_virtual_tool_definitions(), + "tools": get_virtual_tool_definitions( + default_top_k=get_mcp_tool_search_default_top_k(user_api_key_dict) + ), "error": None, "message": "Successfully retrieved tools", } diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 68a61b85175..0432af9ff97 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -702,10 +702,16 @@ if MCP_AVAILABLE: from mcp.types import Tool from litellm.proxy._experimental.mcp_server.tool_search import ( + get_mcp_tool_search_default_top_k, get_virtual_tool_definitions, ) - return [Tool(**d) for d in get_virtual_tool_definitions()] + return [ + Tool(**d) + for d in get_virtual_tool_definitions( + default_top_k=get_mcp_tool_search_default_top_k(user_api_key_auth) + ) + ] # Get mcp_servers from context variable verbose_logger.debug("MCP list_tools - Calling _list_mcp_tools") @@ -821,9 +827,9 @@ if MCP_AVAILABLE: from litellm.proxy._experimental.mcp_server.tool_search import ( MCP_TOOL_CALL_TOOL_NAME, MCP_TOOL_SEARCH_TOOL_NAME, - coerce_top_k, handle_mcp_tool_call, handle_mcp_tool_search, + resolve_mcp_tool_search_top_k, ) if name not in (MCP_TOOL_SEARCH_TOOL_NAME, MCP_TOOL_CALL_TOOL_NAME): @@ -848,7 +854,7 @@ if MCP_AVAILABLE: if name == MCP_TOOL_SEARCH_TOOL_NAME: return await handle_mcp_tool_search( query=args.get("query", ""), - top_k=coerce_top_k(args.get("top_k", 5)), + top_k=resolve_mcp_tool_search_top_k(args.get("top_k"), user_api_key_auth), user_api_key_dict=user_api_key_auth, client_ip=client_ip, mcp_servers=mcp_servers, diff --git a/litellm/proxy/_experimental/mcp_server/tool_search.py b/litellm/proxy/_experimental/mcp_server/tool_search.py index fa57a2b3eb2..dd59f9f6860 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_search.py +++ b/litellm/proxy/_experimental/mcp_server/tool_search.py @@ -12,16 +12,54 @@ if TYPE_CHECKING: MCP_TOOL_SEARCH_TOOL_NAME: str = "mcp_tool_search" MCP_TOOL_CALL_TOOL_NAME: str = "mcp_tool_call" +DEFAULT_MCP_TOOL_SEARCH_TOP_K: int = 5 -def coerce_top_k(value: Any, default: int = 5) -> int: +def _get_litellm_settings() -> dict[str, Any]: + try: + from litellm.proxy.proxy_server import proxy_config + + return proxy_config.get_config_state().get("litellm_settings") or {} + except Exception: + return {} + + +def get_mcp_tool_search_default_top_k( + user_api_key_dict: Optional["UserAPIKeyAuth"] = None, +) -> int: + """Resolve the default top_k for mcp_tool_search (per-key, then global, then 5).""" + if user_api_key_dict is not None: + object_permission = getattr(user_api_key_dict, "object_permission", None) + if object_permission is not None: + key_top_k = getattr(object_permission, "mcp_tool_search_top_k", None) + if key_top_k is not None: + return coerce_top_k(key_top_k, default=DEFAULT_MCP_TOOL_SEARCH_TOP_K) + + global_top_k = _get_litellm_settings().get("mcp_tool_search_default_top_k") + if global_top_k is not None: + return coerce_top_k(global_top_k, default=DEFAULT_MCP_TOOL_SEARCH_TOP_K) + + return DEFAULT_MCP_TOOL_SEARCH_TOP_K + + +def resolve_mcp_tool_search_top_k( + explicit_top_k: Any, + user_api_key_dict: Optional["UserAPIKeyAuth"] = None, +) -> int: + default_top_k = get_mcp_tool_search_default_top_k(user_api_key_dict) + if explicit_top_k is None: + return default_top_k + return coerce_top_k(explicit_top_k, default=default_top_k) + + +def coerce_top_k(value: Any, default: int = DEFAULT_MCP_TOOL_SEARCH_TOP_K) -> int: try: return int(value) except (TypeError, ValueError): return default -def search_tools(query: str, tools: list[dict[str, Any]], top_k: int = 5) -> list[dict[str, Any]]: +def search_tools(query: str, tools: list[dict[str, Any]], top_k: int = DEFAULT_MCP_TOOL_SEARCH_TOP_K) -> list[dict[str, Any]]: if not query: return [] tokens = query.lower().split() @@ -34,7 +72,9 @@ def search_tools(query: str, tools: list[dict[str, Any]], top_k: int = 5) -> lis return [tool for _, tool in sorted(scored, key=lambda x: x[0], reverse=True)[:top_k]] -def get_virtual_tool_definitions() -> list[dict[str, Any]]: +def get_virtual_tool_definitions( + default_top_k: int = DEFAULT_MCP_TOOL_SEARCH_TOP_K, +) -> list[dict[str, Any]]: return [ { "name": MCP_TOOL_SEARCH_TOOL_NAME, @@ -49,7 +89,7 @@ def get_virtual_tool_definitions() -> list[dict[str, Any]]: "top_k": { "type": "integer", "description": "Maximum number of results to return.", - "default": 5, + "default": default_top_k, }, }, "required": ["query"], diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 5e3ea4b7dcb..dd8150cba35 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1009,6 +1009,7 @@ class LiteLLM_ObjectPermissionBase(LiteLLMPydanticObjectBase): models: Optional[List[str]] = None search_tools: Optional[List[str]] = None mcp_tool_search_enabled: Optional[bool] = None + mcp_tool_search_top_k: Optional[int] = None from litellm.types.object_permission import ( # noqa: E402 diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index a23cecc3911..faa66fae557 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -280,6 +280,7 @@ model LiteLLM_ObjectPermissionTable { mcp_toolsets String[] @default([]) // Toolset IDs granted to this key/team/user search_tools String[] @default([]) // search_tool_name values this key/team/user may call mcp_tool_search_enabled Boolean? + mcp_tool_search_top_k Int? teams LiteLLM_TeamTable[] projects LiteLLM_ProjectTable[] verification_tokens LiteLLM_VerificationToken[] diff --git a/litellm/types/object_permission.py b/litellm/types/object_permission.py index d0458173fbf..024e141781d 100644 --- a/litellm/types/object_permission.py +++ b/litellm/types/object_permission.py @@ -25,3 +25,4 @@ class ObjectPermissionDict(TypedDict, total=False): models: Optional[list[str]] search_tools: Optional[list[str]] mcp_tool_search_enabled: Optional[bool] + mcp_tool_search_top_k: Optional[int] diff --git a/schema.prisma b/schema.prisma index a23cecc3911..faa66fae557 100644 --- a/schema.prisma +++ b/schema.prisma @@ -280,6 +280,7 @@ model LiteLLM_ObjectPermissionTable { mcp_toolsets String[] @default([]) // Toolset IDs granted to this key/team/user search_tools String[] @default([]) // search_tool_name values this key/team/user may call mcp_tool_search_enabled Boolean? + mcp_tool_search_top_k Int? teams LiteLLM_TeamTable[] projects LiteLLM_ProjectTable[] verification_tokens LiteLLM_VerificationToken[] 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 b8f0b205831..e860365954e 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 @@ -19,8 +19,11 @@ from litellm.models.object_permission import LiteLLM_ObjectPermissionTable from litellm.proxy._experimental.mcp_server.tool_search import ( MCP_TOOL_CALL_TOOL_NAME, MCP_TOOL_SEARCH_TOOL_NAME, + DEFAULT_MCP_TOOL_SEARCH_TOP_K, coerce_top_k, + get_mcp_tool_search_default_top_k, get_virtual_tool_definitions, + resolve_mcp_tool_search_top_k, search_tools, ) from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth @@ -72,6 +75,58 @@ class TestCoerceTopK: assert coerce_top_k("nope", default=10) == 10 +class TestMcpToolSearchDefaultTopK: + def test_builtin_default(self) -> None: + assert get_mcp_tool_search_default_top_k() == DEFAULT_MCP_TOOL_SEARCH_TOP_K + + def test_per_key_override(self) -> None: + uak = UserAPIKeyAuth( + api_key="k", + object_permission=_make_perm(mcp_tool_search_top_k=10), + ) + assert get_mcp_tool_search_default_top_k(uak) == 10 + + def test_global_litellm_settings_override(self, monkeypatch: pytest.MonkeyPatch) -> None: + mock_config = MagicMock() + mock_config.get_config_state.return_value = { + "litellm_settings": {"mcp_tool_search_default_top_k": 12} + } + monkeypatch.setattr( + "litellm.proxy.proxy_server.proxy_config", + mock_config, + ) + assert get_mcp_tool_search_default_top_k() == 12 + + def test_per_key_beats_global(self, monkeypatch: pytest.MonkeyPatch) -> None: + mock_config = MagicMock() + mock_config.get_config_state.return_value = { + "litellm_settings": {"mcp_tool_search_default_top_k": 12} + } + monkeypatch.setattr( + "litellm.proxy.proxy_server.proxy_config", + mock_config, + ) + uak = UserAPIKeyAuth( + api_key="k", + object_permission=_make_perm(mcp_tool_search_top_k=8), + ) + assert get_mcp_tool_search_default_top_k(uak) == 8 + + def test_resolve_uses_explicit_top_k(self) -> None: + uak = UserAPIKeyAuth( + api_key="k", + object_permission=_make_perm(mcp_tool_search_top_k=10), + ) + assert resolve_mcp_tool_search_top_k(3, uak) == 3 + + def test_resolve_uses_default_when_omitted(self) -> None: + uak = UserAPIKeyAuth( + api_key="k", + object_permission=_make_perm(mcp_tool_search_top_k=10), + ) + assert resolve_mcp_tool_search_top_k(None, uak) == 10 + + class TestSearchTools: def test_returns_matching_tools(self) -> None: results = search_tools("github issue", SAMPLE_TOOLS) @@ -139,6 +194,11 @@ class TestGetVirtualToolDefinitions: assert "arguments" in props assert "tool_name" in call_tool["inputSchema"]["required"] + def test_mcp_tool_search_schema_top_k_default(self) -> None: + tools = get_virtual_tool_definitions(default_top_k=10) + search_tool = next(t for t in tools if t["name"] == MCP_TOOL_SEARCH_TOOL_NAME) + assert search_tool["inputSchema"]["properties"]["top_k"]["default"] == 10 + def test_all_tools_have_description(self) -> None: for tool in get_virtual_tool_definitions(): assert tool.get("description"), f"{tool['name']} missing description"