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

* test(mcp): check the tool dict conversion without patching the MCP handler

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

* refactor(mcp): type the semantic filter's default tool loader

Annotate the chosen loader as _OpenAIToolLoader and narrow its return to
FunctionToolParam, so a default that returns raw MCP tools fails the type
check. Drop the unused patch import.

litellm/proxy/hooks/mcp_semantic_filter/hook.py

@@ -24,6 +24,8 @@ from litellm.proxy.mcp_server.semantic_tool_filter import (
 )
 
 if TYPE_CHECKING:
+    from openai.types.responses.function_tool_param import FunctionToolParam
+
     from litellm.caching.caching import DualCache
     from litellm.proxy._types import UserAPIKeyAuth
     from litellm.proxy.mcp_server.semantic_tool_filter import (
@@ -51,6 +53,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["FunctionToolParam"], Mapping[str, str]]]: ...
+
+
 class _PydanticTool(Protocol):
     def model_dump(self, *, exclude_none: bool) -> Mapping[str, object]: ...
 
@@ -90,15 +101,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 +157,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[_OpenAIToolLoader] = (
+            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

@@ -7,7 +7,7 @@ an ordered set of top K tools based on semantic similarity.
 
 import asyncio
 import sys
-from unittest.mock import AsyncMock, Mock, patch
+from unittest.mock import AsyncMock, Mock
 
 import pytest
 
@@ -2249,7 +2249,6 @@ async def test_expand_mcp_tools_returns_every_expanded_tool_as_a_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
-    from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Handler
 
     class _FunctionTool(BaseModel):
         model_config = ConfigDict(frozen=True)
@@ -2277,15 +2276,13 @@ async def test_expand_mcp_tools_returns_every_expanded_tool_as_a_dict():
     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())
+        SemanticMCPToolFilter(embedding_model="text-embedding-3-small", litellm_router_instance=Mock()),
+        load_openai_tools=process_tools,
     )
-    with patch.object(
-        LiteLLM_Proxy_MCP_Handler,
-        "process_mcp_tools_to_openai_format",
-        new=AsyncMock(return_value=(pipeline_output, {})),
-    ) as process_tools:
-        expanded = await hook._expand_mcp_tools([gateway_reference], auth)
+
+    expanded = await hook._expand_mcp_tools([gateway_reference], auth)
 
     assert expanded == [
         {"type": "function", "name": "v2_tool"},