diff --git a/litellm/proxy/hooks/mcp_semantic_filter/hook.py b/litellm/proxy/hooks/mcp_semantic_filter/hook.py index 1830961dd06..163e0243379 100644 --- a/litellm/proxy/hooks/mcp_semantic_filter/hook.py +++ b/litellm/proxy/hooks/mcp_semantic_filter/hook.py @@ -51,6 +51,15 @@ def _truncate_csv_at_tool_name_boundary(tool_names_csv: str, max_length: int) -> return head.rsplit(",", 1)[0] +class _OpenAIToolLoader(Protocol): + def __call__( + self, + *, + user_api_key_auth: "UserAPIKeyAuth | None", + mcp_tools_with_litellm_proxy: Sequence[Mapping[str, object]], + ) -> Awaitable[tuple[Sequence[object], Mapping[str, str]]]: ... + + class _PydanticTool(Protocol): def model_dump(self, *, exclude_none: bool) -> Mapping[str, object]: ... @@ -90,15 +99,22 @@ class SemanticToolFilterHook(CustomLogger): 3. Returns only the top-k most relevant tools to the LLM """ - def __init__(self, semantic_filter: "SemanticMCPToolFilter"): + def __init__( + self, + semantic_filter: "SemanticMCPToolFilter", + load_openai_tools: _OpenAIToolLoader | None = None, + ): """ Initialize the hook. Args: semantic_filter: SemanticMCPToolFilter instance + load_openai_tools: loads the MCP tools behind gateway references, defaults to + LiteLLM_Proxy_MCP_Handler.process_mcp_tools_to_openai_format """ super().__init__() self.filter = semantic_filter + self._load_openai_tools = load_openai_tools verbose_proxy_logger.debug( "Initialized SemanticToolFilterHook with filter: enabled=%s, top_k=%s", @@ -139,12 +155,12 @@ class SemanticToolFilterHook(CustomLogger): if not mcp_tools: return [] - # Use single combined method instead of 3 separate calls - # This already handles: fetch -> filter by allowed_tools -> deduplicate -> transform - ( - openai_tools, - _, - ) = await LiteLLM_Proxy_MCP_Handler.process_mcp_tools_to_openai_format( + load_openai_tools: Final = ( + LiteLLM_Proxy_MCP_Handler.process_mcp_tools_to_openai_format + if self._load_openai_tools is None + else self._load_openai_tools + ) + openai_tools, _ = await load_openai_tools( user_api_key_auth=user_api_key_dict, mcp_tools_with_litellm_proxy=mcp_tools ) diff --git a/tests/unit/proxy/mcp_server/test_semantic_tool_filter.py b/tests/unit/proxy/mcp_server/test_semantic_tool_filter.py index 3549e9ec085..0b5cd82efd8 100644 --- a/tests/unit/proxy/mcp_server/test_semantic_tool_filter.py +++ b/tests/unit/proxy/mcp_server/test_semantic_tool_filter.py @@ -2240,12 +2240,15 @@ async def test_semantic_filter_hook_narrows_only_references_the_gateway_serves() assert narrowed == [{**gateway_reference, "allowed_tools": ["srv-tool_1"]}, external_tool] -def test_tool_as_dict_returns_every_expanded_tool_shape_as_a_dict(): +@pytest.mark.asyncio +async def test_expand_mcp_tools_returns_every_expanded_tool_as_a_dict(): from typing import Optional from pydantic import BaseModel, ConfigDict - from litellm.proxy.hooks.mcp_semantic_filter.hook import _tool_as_dict + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + from litellm.proxy.mcp_server.semantic_tool_filter import SemanticMCPToolFilter class _FunctionTool(BaseModel): model_config = ConfigDict(frozen=True) @@ -2264,16 +2267,30 @@ def test_tool_as_dict_returns_every_expanded_tool_shape_as_a_dict(): plain_tool = {"type": "function", "name": "plain_tool", "parameters": {"type": "object"}} unknown_tool = "opaque-tool" - expanded_tools = ( + pipeline_output = [ _FunctionTool(type="function", name="v2_tool"), _LegacyFunctionTool("legacy_tool"), plain_tool, unknown_tool, + ] + auth = UserAPIKeyAuth(api_key="hashed-key") + gateway_reference = {"type": "mcp", "server_url": "litellm_proxy", "require_approval": "never"} + + process_tools = AsyncMock(return_value=(pipeline_output, {})) + hook = SemanticToolFilterHook( + SemanticMCPToolFilter(embedding_model="text-embedding-3-small", litellm_router_instance=Mock()), + load_openai_tools=process_tools, ) - assert [_tool_as_dict(tool) for tool in expanded_tools] == [ + expanded = await hook._expand_mcp_tools([gateway_reference], auth) + + assert expanded == [ {"type": "function", "name": "v2_tool"}, {"type": "function", "name": "legacy_tool"}, plain_tool, unknown_tool, ] + assert process_tools.await_args.kwargs == { + "user_api_key_auth": auth, + "mcp_tools_with_litellm_proxy": [gateway_reference], + }