From 130c1660fd1258ca0853dd4382f30a90cb56fb77 Mon Sep 17 00:00:00 2001 From: Genmin Date: Thu, 30 Apr 2026 20:10:26 -0700 Subject: [PATCH] fix mcp semantic filter native tool passthrough --- .../proxy/hooks/mcp_semantic_filter/hook.py | 75 ++++++-- .../mcp_server/test_semantic_tool_filter.py | 182 ++++++++++++++++-- 2 files changed, 233 insertions(+), 24 deletions(-) diff --git a/litellm/proxy/hooks/mcp_semantic_filter/hook.py b/litellm/proxy/hooks/mcp_semantic_filter/hook.py index 6343faaa965..e05e346e3f9 100644 --- a/litellm/proxy/hooks/mcp_semantic_filter/hook.py +++ b/litellm/proxy/hooks/mcp_semantic_filter/hook.py @@ -77,10 +77,10 @@ class SemanticToolFilterHook(CustomLogger): ) # Parse to separate MCP tools from other tools - mcp_tools, _ = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools) + mcp_tools, other_tools = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools) if not mcp_tools: - return [] + return other_tools # Use single combined method instead of 3 separate calls # This already handles: fetch -> filter by allowed_tools -> deduplicate -> transform @@ -118,16 +118,65 @@ class SemanticToolFilterHook(CustomLogger): openai_tools_as_dicts.append(tool) verbose_proxy_logger.debug( - f"Expanded {len(mcp_tools)} MCP reference(s) to {len(openai_tools_as_dicts)} tools (all as dicts)" + f"Expanded {len(mcp_tools)} MCP reference(s) to {len(openai_tools_as_dicts)} " + f"tools and preserved {len(other_tools)} non-MCP tool(s)" ) - return openai_tools_as_dicts + return other_tools + openai_tools_as_dicts def _get_metadata_variable_name(self, data: dict) -> str: if "litellm_metadata" in data: return "litellm_metadata" return "metadata" + @staticmethod + def _get_tool_name(tool: Any) -> str: + if isinstance(tool, dict): + function = tool.get("function") + if isinstance(function, dict) and isinstance(function.get("name"), str): + return function["name"] + + name = tool.get("name") + if isinstance(name, str): + return name + + return "" + + name = getattr(tool, "name", "") + return name if isinstance(name, str) else "" + + def _is_mcp_router_tool(self, tool: Any) -> bool: + tool_name = self._get_tool_name(tool) + if not tool_name: + return False + + tool_map = getattr(self.filter, "_tool_map", {}) + if tool_name in tool_map: + return True + + name_matches_canonical = getattr(self.filter, "_name_matches_canonical", None) + if name_matches_canonical is None: + return False + + return any( + name_matches_canonical(tool_name, canonical_name) + for canonical_name in tool_map + ) + + def _partition_mcp_router_tools( + self, tools: List[Any] + ) -> tuple[List[Any], List[Any]]: + native_tools: List[Any] = [] + mcp_tools: List[Any] = [] + + for tool in tools: + if self._is_mcp_router_tool(tool): + mcp_tools.append(tool) + else: + native_tools.append(tool) + + return native_tools, mcp_tools + async def async_pre_call_hook( self, user_api_key_dict: "UserAPIKeyAuth", @@ -223,11 +272,19 @@ class SemanticToolFilterHook(CustomLogger): f"with query: '{user_query[:50]}...'" ) + native_tools, mcp_tools = self._partition_mcp_router_tools(tools) + if not mcp_tools: + verbose_proxy_logger.debug( + "No MCP router tools in request, skipping semantic filter" + ) + return None + # Filter tools semantically - filtered_tools = await self.filter.filter_tools( + filtered_mcp_tools = await self.filter.filter_tools( query=user_query, - available_tools=tools, # type: ignore + available_tools=mcp_tools, # type: ignore ) + filtered_tools = native_tools + filtered_mcp_tools # Always update tools and emit header (even if count unchanged) data["tools"] = filtered_tools @@ -294,11 +351,7 @@ class SemanticToolFilterHook(CustomLogger): tool_names = [] for tool in tools: - name = ( - tool.get("name", "") - if isinstance(tool, dict) - else getattr(tool, "name", "") - ) + name = self._get_tool_name(tool) if name: tool_names.append(name) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py index 2558df8533b..2fdd95e65fa 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py @@ -5,10 +5,9 @@ Tests the core filtering logic that takes a long list of tools and returns an ordered set of top K tools based on semantic similarity. """ -import asyncio import os import sys -from unittest.mock import AsyncMock, Mock, patch +from unittest.mock import AsyncMock, Mock import pytest @@ -137,12 +136,6 @@ async def test_semantic_filter_basic_filtering(): tool, "description" ), "Filtered result should be MCPTool with description" - filtered_names = [t.name for t in filtered] - print( - f"✅ Successfully filtered {len(tools)} tools down to top {len(filtered)}: {filtered_names}" - ) - print(f" Filter respects top_k parameter correctly") - @pytest.mark.asyncio async def test_semantic_filter_top_k_limiting(): @@ -206,8 +199,6 @@ async def test_semantic_filter_top_k_limiting(): # Should return at most 5 tools assert len(filtered) <= 5, f"Expected at most 5 tools, got {len(filtered)}" - print(f"Returned {len(filtered)} tools out of {len(tools)} (top_k=5)") - @pytest.mark.asyncio async def test_semantic_filter_disabled(): @@ -403,8 +394,6 @@ async def test_semantic_filter_hook_triggers_on_completion(): tools ), f"Hook should filter tools, got {len(result['tools'])}/{len(tools)}" - print(f"✅ Hook triggered correctly: {len(tools)} -> {len(result['tools'])} tools") - @pytest.mark.asyncio async def test_semantic_filter_hook_skips_no_tools(): @@ -449,7 +438,174 @@ async def test_semantic_filter_hook_skips_no_tools(): # Should return None (no modification) assert result is None, "Hook should skip requests without tools" - print("✅ Hook correctly skips requests without tools") + + +@pytest.mark.asyncio +async def test_expand_mcp_tools_preserves_native_tools(monkeypatch): + """ + MCP-reference expansion should not discard native tools that arrived in + the same request. + """ + from litellm.proxy._experimental.mcp_server.semantic_tool_filter import ( + SemanticMCPToolFilter, + ) + from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + from litellm.responses.mcp.litellm_proxy_mcp_handler import ( + LiteLLM_Proxy_MCP_Handler, + ) + + async def fake_process_mcp_tools_to_openai_format(*args, **kwargs): + return ([{"name": "github-search", "description": "Search GitHub"}], None) + + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_process_mcp_tools_to_openai_format", + fake_process_mcp_tools_to_openai_format, + ) + + filter_instance = SemanticMCPToolFilter( + embedding_model="text-embedding-3-small", + litellm_router_instance=Mock(), + top_k=3, + similarity_threshold=0.3, + enabled=True, + ) + hook = SemanticToolFilterHook(filter_instance) + native_tool = { + "type": "function", + "function": { + "name": "weather_lookup", + "description": "Look up weather for a city", + "parameters": {"type": "object"}, + }, + } + mcp_reference = {"type": "mcp", "server_url": "litellm_proxy/mcp/github"} + + expanded = await hook._expand_mcp_tools( + [mcp_reference, native_tool], + user_api_key_dict=Mock(), + ) + + assert expanded == [ + native_tool, + {"name": "github-search", "description": "Search GitHub"}, + ] + + +@pytest.mark.asyncio +async def test_semantic_filter_hook_preserves_native_openai_tools(): + """ + The MCP semantic filter must not drop native tools owned by the caller. + + Native tools are not present in the MCP router map, so they should bypass + MCP semantic filtering and be merged back into the request unchanged. + """ + from litellm.proxy._experimental.mcp_server.semantic_tool_filter import ( + SemanticMCPToolFilter, + ) + from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + + filter_instance = SemanticMCPToolFilter( + embedding_model="text-embedding-3-small", + litellm_router_instance=Mock(), + top_k=3, + similarity_threshold=0.3, + enabled=True, + ) + mcp_tool = {"name": "github-search", "description": "Search GitHub"} + filter_instance._tool_map = {mcp_tool["name"]: mcp_tool} + filter_instance.filter_tools = AsyncMock(return_value=[mcp_tool]) # type: ignore[method-assign] + + native_tool = { + "type": "function", + "function": { + "name": "weather_lookup", + "description": "Look up weather for a city", + "parameters": {"type": "object"}, + }, + } + responses_native_tool = { + "type": "function", + "name": "calculator", + "description": "Evaluate an expression", + "parameters": {"type": "object"}, + } + hook = SemanticToolFilterHook(filter_instance) + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Check weather and search GitHub"}], + "tools": [native_tool, responses_native_tool, mcp_tool], + "metadata": {}, + } + + result = await hook.async_pre_call_hook( + user_api_key_dict=Mock(), + cache=Mock(), + data=data, + call_type="completion", + ) + + assert result is data + filter_instance.filter_tools.assert_awaited_once_with( + query="Check weather and search GitHub", + available_tools=[mcp_tool], + ) + assert result["tools"] == [native_tool, responses_native_tool, mcp_tool] + assert result["metadata"]["litellm_semantic_filter_stats"] == "3->3" + assert ( + result["metadata"]["litellm_semantic_filter_tools"] + == "weather_lookup,calculator,github-search" + ) + + +@pytest.mark.asyncio +async def test_semantic_filter_hook_skips_all_native_openai_tools(): + """ + If a request contains only caller-owned native tools, the MCP semantic + filter should leave the request untouched and avoid emitting filter stats. + """ + from litellm.proxy._experimental.mcp_server.semantic_tool_filter import ( + SemanticMCPToolFilter, + ) + from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + + filter_instance = SemanticMCPToolFilter( + embedding_model="text-embedding-3-small", + litellm_router_instance=Mock(), + top_k=3, + similarity_threshold=0.3, + enabled=True, + ) + filter_instance._tool_map = {"github-search": {"name": "github-search"}} + filter_instance.filter_tools = AsyncMock(return_value=[]) # type: ignore[method-assign] + + native_tool = { + "type": "function", + "function": { + "name": "weather_lookup", + "description": "Look up weather for a city", + "parameters": {"type": "object"}, + }, + } + hook = SemanticToolFilterHook(filter_instance) + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Check weather"}], + "tools": [native_tool], + "metadata": {}, + } + + result = await hook.async_pre_call_hook( + user_api_key_dict=Mock(), + cache=Mock(), + data=data, + call_type="completion", + ) + + assert result is None + filter_instance.filter_tools.assert_not_awaited() + assert data["tools"] == [native_tool] + assert "litellm_semantic_filter_stats" not in data["metadata"] class TestGetToolsByNames: