diff --git a/litellm/proxy/hooks/mcp_semantic_filter/hook.py b/litellm/proxy/hooks/mcp_semantic_filter/hook.py index 1830961dd06..b7b65d6f0db 100644 --- a/litellm/proxy/hooks/mcp_semantic_filter/hook.py +++ b/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 ) 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 65e59fab0da..afbc439cb99 100644 --- a/tests/unit/proxy/mcp_server/test_semantic_tool_filter.py +++ b/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"},