refactor(mcp): let the semantic filter take its MCP tool loader as a parameter

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
         )
 

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],
+    }