diff --git a/litellm/proxy/hooks/mcp_semantic_filter/hook.py b/litellm/proxy/hooks/mcp_semantic_filter/hook.py index 6343faaa965..9888baf897e 100644 --- a/litellm/proxy/hooks/mcp_semantic_filter/hook.py +++ b/litellm/proxy/hooks/mcp_semantic_filter/hook.py @@ -123,11 +123,70 @@ async def _expand_mcp_tools( return openai_tools_as_dicts + def _is_mcp_tool(self, tool: object) -> bool: + """ + Check whether *tool* is registered in the MCP semantic router. + + Classification strategy (shape-first, lookup-second): + 1. Chat Completions format dicts are always native. + 2. Responses API function tools are always native. + 3. Everything else is looked up by name in the MCP registry. + """ + if ( + isinstance(tool, dict) + and tool.get("type") == "function" + and isinstance(tool.get("function"), dict) + ): + return False + if ( + isinstance(tool, dict) + and tool.get("type") == "function" + and isinstance(tool.get("name"), str) + ): + return False + name, _ = self.filter._extract_tool_info(tool) + return bool(name) and name in self.filter._tool_map + def _get_metadata_variable_name(self, data: dict) -> str: if "litellm_metadata" in data: return "litellm_metadata" return "metadata" + def _emit_filter_metadata( + self, + data: dict, + mcp_tools: list[object], + filtered_mcp_tools: list[object], + native_tools: list[object], + filtered_tools: list[object], + ) -> None: + """ + Emit response-header metadata when MCP tools were filtered. + + Stats report MCP-only counts so downstream consumers see accurate + semantic filter metrics. Skips metadata entirely for purely-native + requests to avoid spurious headers. + """ + if mcp_tools: + filter_stats = f"{len(mcp_tools)}->{len(filtered_mcp_tools)}" + tool_names_csv = self._get_tool_names_csv(filtered_mcp_tools) + + _metadata_variable_name = self._get_metadata_variable_name(data) + metadata = data.setdefault(_metadata_variable_name, {}) + metadata["litellm_semantic_filter_stats"] = filter_stats + metadata["litellm_semantic_filter_tools"] = tool_names_csv + + verbose_proxy_logger.info( + f"Semantic tool filter: {filter_stats} MCP tools " + f"({len(native_tools)} native preserved, " + f"{len(filtered_tools)} total)" + ) + else: + verbose_proxy_logger.info( + f"Semantic tool filter: all {len(native_tools)} tools " + f"are native, no MCP filtering applied" + ) + async def async_pre_call_hook( self, user_api_key_dict: "UserAPIKeyAuth", @@ -140,53 +199,55 @@ async def async_pre_call_hook( This hook is called before the LLM request is made. It filters the tools list to only include semantically relevant tools. - - Args: - user_api_key_dict: User authentication - cache: Cache instance - data: Request data containing messages and tools - call_type: Type of call (completion, acompletion, etc.) - - Returns: - Modified data dict with filtered tools, or None if no changes """ - # Only filter endpoints that support tools if call_type not in ("completion", "acompletion", "aresponses"): verbose_proxy_logger.debug( f"Skipping semantic filter for call_type={call_type}" ) return None - # Check if tools are present tools = data.get("tools") if not tools: verbose_proxy_logger.debug("No tools in request, skipping semantic filter") return None - original_tool_count = len(tools) - - # Check for MCP references (server_url="litellm_proxy") and expand them + # Expanded MCP tools are in OpenAI nested format which + # filter_tools/_extract_tool_info cannot name-match, so we skip + # semantic filtering and return early. if self._should_expand_mcp_tools(tools): verbose_proxy_logger.debug( "Detected litellm_proxy MCP references, expanding before semantic filtering" ) try: + native_tools_before_expand = [ + t + for t in tools + if not (isinstance(t, dict) and t.get("type") == "mcp") + ] + expanded_tools = await self._expand_mcp_tools(tools, user_api_key_dict) if not expanded_tools: + if native_tools_before_expand: + data["tools"] = native_tools_before_expand + verbose_proxy_logger.warning( + "No MCP tools expanded, preserving " + f"{len(native_tools_before_expand)} native tools" + ) + return data verbose_proxy_logger.warning( "No tools expanded from MCP references" ) return None + data["tools"] = native_tools_before_expand + expanded_tools verbose_proxy_logger.info( - f"Expanded {len(tools)} MCP reference(s) to {len(expanded_tools)} tools" + f"Expanded MCP references to {len(expanded_tools)} tools " + f"({len(native_tools_before_expand)} native preserved), " + f"skipping semantic filter (OpenAI nested format)" ) - - # Update tools for filtering - tools = expanded_tools - original_tool_count = len(tools) + return data except Exception as e: verbose_proxy_logger.error( @@ -194,7 +255,6 @@ async def async_pre_call_hook( ) return None - # Check if messages are present (try both "messages" and "input" for responses API) messages = data.get("messages", []) if not messages: messages = data.get("input", []) @@ -204,13 +264,11 @@ async def async_pre_call_hook( ) return None - # Check if filter is enabled if not self.filter.enabled: verbose_proxy_logger.debug("Semantic filter disabled, skipping") return None try: - # Extract user query from messages user_query = self.filter.extract_user_query(messages) if not user_query: verbose_proxy_logger.debug( @@ -218,33 +276,60 @@ async def async_pre_call_hook( ) return None + native_tools: list[object] = [] + mcp_tools: list[object] = [] + mcp_indices: set[int] = set() + for i, t in enumerate(tools): + if self._is_mcp_tool(t): + mcp_tools.append(t) + mcp_indices.add(i) + else: + native_tools.append(t) + verbose_proxy_logger.debug( - f"Applying semantic filter to {len(tools)} tools " - f"with query: '{user_query[:50]}...'" + f"Applying semantic filter: {len(mcp_tools)} MCP tools, " + f"{len(native_tools)} native tools, " + f"query: '{user_query[:50]}...'" ) - # Filter tools semantically - filtered_tools = await self.filter.filter_tools( - query=user_query, - available_tools=tools, # type: ignore - ) + if mcp_tools: + filtered_mcp_tools = await self.filter.filter_tools( + query=user_query, + available_tools=mcp_tools, # type: ignore + ) + else: + filtered_mcp_tools = [] + + filtered_mcp_names: set[str] = set() + for t in filtered_mcp_tools: + name, _ = self.filter._extract_tool_info(t) + if name: + filtered_mcp_names.add(name) + + filtered_tools: list[object] = [] + for i, t in enumerate(tools): + if i in mcp_indices: + name, _ = self.filter._extract_tool_info(t) + if name in filtered_mcp_names: + filtered_tools.append(t) + else: + filtered_tools.append(t) - # Always update tools and emit header (even if count unchanged) data["tools"] = filtered_tools - # Store filter stats and tool names for response header - filter_stats = f"{original_tool_count}->{len(filtered_tools)}" - tool_names_csv = self._get_tool_names_csv(filtered_tools) - - _metadata_variable_name = self._get_metadata_variable_name(data) - data[_metadata_variable_name][ - "litellm_semantic_filter_stats" - ] = filter_stats - data[_metadata_variable_name][ - "litellm_semantic_filter_tools" - ] = tool_names_csv - - verbose_proxy_logger.info(f"Semantic tool filter: {filter_stats} tools") + try: + self._emit_filter_metadata( + data=data, + mcp_tools=mcp_tools, + filtered_mcp_tools=filtered_mcp_tools, + native_tools=native_tools, + filtered_tools=filtered_tools, + ) + except Exception as e: + verbose_proxy_logger.warning( + f"Failed to emit semantic filter metadata: {e}", + exc_info=True, + ) return data @@ -266,7 +351,7 @@ async def async_post_call_response_headers_hook( from litellm.constants import MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH _metadata_variable_name = self._get_metadata_variable_name(data) - metadata = data[_metadata_variable_name] + metadata = data.get(_metadata_variable_name, {}) filter_stats = metadata.get("litellm_semantic_filter_stats") if not filter_stats: 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..cebc265a148 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 @@ -452,6 +452,430 @@ async def test_semantic_filter_hook_skips_no_tools(): print("✅ Hook correctly skips requests without tools") +@pytest.mark.asyncio +async def test_semantic_filter_hook_preserves_native_tools(): + """ + Regression test: mixed MCP + native tools. + + Given: 5 MCP tools (registered in _tool_map) + 2 native OpenAI-format + function tools (not in _tool_map) + When: The hook filters tools + Then: The native tools must survive unconditionally, and only MCP + tools go through the semantic filter. + """ + from litellm.proxy._experimental.mcp_server.semantic_tool_filter import ( + SemanticMCPToolFilter, + ) + from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + from litellm.types.utils import Embedding, EmbeddingResponse + + mock_router = Mock() + + def mock_embedding_sync(*args, **kwargs): + return EmbeddingResponse( + data=[Embedding(embedding=[0.1] * 1536, index=0, object="embedding")], + model="text-embedding-3-small", + object="list", + usage={"prompt_tokens": 10, "total_tokens": 10}, + ) + + async def mock_embedding_async(*args, **kwargs): + return mock_embedding_sync() + + mock_router.embedding = mock_embedding_sync + mock_router.aembedding = mock_embedding_async + + filter_instance = SemanticMCPToolFilter( + embedding_model="text-embedding-3-small", + litellm_router_instance=mock_router, + top_k=2, + similarity_threshold=0.3, + enabled=True, + ) + + # --- MCP tools (registered in the semantic router) --- + mcp_tools = [ + MCPTool( + name=f"mcp_tool_{i}", + description=f"MCP tool {i}", + inputSchema={"type": "object"}, + ) + for i in range(5) + ] + filter_instance._build_router(mcp_tools) + + # --- Native OpenAI-format function tools (NOT in _tool_map) --- + native_tools = [ + { + "type": "function", + "function": { + "name": "get_current_weather", + "description": "Get the current weather", + "parameters": {"type": "object", "properties": {}}, + }, + }, + { + "type": "function", + "function": { + "name": "search_web", + "description": "Search the web", + "parameters": {"type": "object", "properties": {}}, + }, + }, + ] + + # Combine: MCP tools + native tools + all_tools = list(mcp_tools) + native_tools + + hook = SemanticToolFilterHook(filter_instance) + + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "What is the weather?"}], + "tools": all_tools, + "metadata": {}, + } + + result = await hook.async_pre_call_hook( + user_api_key_dict=Mock(), + cache=Mock(), + data=data, + call_type="completion", + ) + + assert result is not None, "Hook should return modified data" + filtered = result["tools"] + + # Native tools must survive + native_in_result = [ + t for t in filtered if isinstance(t, dict) and t.get("type") == "function" + ] + assert ( + len(native_in_result) == 2 + ), f"Both native tools must survive, got {len(native_in_result)}" + + # MCP tools should be filtered (top_k=2) + mcp_in_result = [t for t in filtered if not isinstance(t, dict)] + assert ( + len(mcp_in_result) <= 2 + ), f"MCP tools should be filtered to top_k=2, got {len(mcp_in_result)}" + + # Total should be native + filtered MCP + assert len(filtered) <= 4, f"Expected at most 4 tools, got {len(filtered)}" + + # Filter stats should be emitted (MCP tools were present) + assert "litellm_semantic_filter_stats" in result["metadata"] + + # Stats should report MCP-only counts, not inflated with native tools + stats = result["metadata"]["litellm_semantic_filter_stats"] + mcp_before, mcp_after = stats.split("->") + assert ( + int(mcp_before) == 5 + ), f"Stats 'from' should be MCP count (5), got {mcp_before}" + assert int(mcp_after) == len( + mcp_in_result + ), f"Stats 'to' should match filtered MCP count, got {mcp_after}" + + print( + f"✅ Hook preserves native tools: {len(all_tools)} -> {len(filtered)} " + f"({len(native_in_result)} native + {len(mcp_in_result)} MCP), " + f"stats={stats}" + ) + + +@pytest.mark.asyncio +async def test_semantic_filter_hook_all_native_tools(): + """ + Regression test: all-native request. + + Given: Only native OpenAI-format function tools (none registered in + the MCP semantic router) + When: The hook processes the request + Then: All tools pass through, and NO spurious semantic filter response + headers are emitted (no litellm_semantic_filter_stats in metadata). + """ + from litellm.proxy._experimental.mcp_server.semantic_tool_filter import ( + SemanticMCPToolFilter, + ) + from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + + mock_router = Mock() + filter_instance = SemanticMCPToolFilter( + embedding_model="text-embedding-3-small", + litellm_router_instance=mock_router, + top_k=3, + similarity_threshold=0.3, + enabled=True, + ) + + # Build router with some MCP tools (so tool_router is not None) + mcp_tools = [ + MCPTool( + name="some_mcp_tool", + description="An MCP tool", + inputSchema={"type": "object"}, + ) + ] + + from litellm.types.utils import Embedding, EmbeddingResponse + + def mock_embedding_sync(*args, **kwargs): + return EmbeddingResponse( + data=[Embedding(embedding=[0.1] * 1536, index=0, object="embedding")], + model="text-embedding-3-small", + object="list", + usage={"prompt_tokens": 10, "total_tokens": 10}, + ) + + async def mock_embedding_async(*args, **kwargs): + return mock_embedding_sync() + + mock_router.embedding = mock_embedding_sync + mock_router.aembedding = mock_embedding_async + + filter_instance._build_router(mcp_tools) + + # --- Only native tools in the request --- + native_tools = [ + { + "type": "function", + "function": { + "name": f"native_func_{i}", + "description": f"Native function {i}", + "parameters": {"type": "object", "properties": {}}, + }, + } + for i in range(3) + ] + + hook = SemanticToolFilterHook(filter_instance) + + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "tools": native_tools, + "metadata": {}, + } + + result = await hook.async_pre_call_hook( + user_api_key_dict=Mock(), + cache=Mock(), + data=data, + call_type="completion", + ) + + assert result is not None, "Hook should return modified data" + filtered = result["tools"] + + # All native tools must pass through + assert ( + len(filtered) == 3 + ), f"All 3 native tools must pass through, got {len(filtered)}" + + # No spurious semantic filter stats (P2 fix) + assert ( + "litellm_semantic_filter_stats" not in result["metadata"] + ), "Should NOT emit semantic filter stats for all-native-tool requests" + + print( + f"✅ Hook passes through all {len(filtered)} native tools, " + f"no spurious filter headers emitted" + ) + + +@pytest.mark.asyncio +async def test_semantic_filter_hook_responses_api_name_collision(): + """ + Regression test: Responses API native tool with MCP-matching name. + + Given: A Responses-API native tool whose top-level ``name`` collides + with an MCP canonical name in ``_tool_map`` + When: The hook classifies tools + Then: The native tool must NOT be sent to the semantic filter, even + though its name matches an MCP canonical. + """ + from litellm.proxy._experimental.mcp_server.semantic_tool_filter import ( + SemanticMCPToolFilter, + ) + from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + from litellm.types.utils import Embedding, EmbeddingResponse + + mock_router = Mock() + + def mock_embedding_sync(*args, **kwargs): + return EmbeddingResponse( + data=[Embedding(embedding=[0.1] * 1536, index=0, object="embedding")], + model="text-embedding-3-small", + object="list", + usage={"prompt_tokens": 10, "total_tokens": 10}, + ) + + async def mock_embedding_async(*args, **kwargs): + return mock_embedding_sync() + + mock_router.embedding = mock_embedding_sync + mock_router.aembedding = mock_embedding_async + + filter_instance = SemanticMCPToolFilter( + embedding_model="text-embedding-3-small", + litellm_router_instance=mock_router, + top_k=2, + similarity_threshold=0.3, + enabled=True, + ) + + # Register an MCP tool with name "github-search" + mcp_tools = [ + MCPTool( + name="github-search", + description="Search GitHub repos", + inputSchema={"type": "object"}, + ) + ] + filter_instance._build_router(mcp_tools) + + # Responses API native tool with SAME name as MCP canonical + responses_api_tool = { + "type": "function", + "name": "github-search", + "description": "Caller-owned search tool", + "parameters": {"type": "object"}, + } + + hook = SemanticToolFilterHook(filter_instance) + + # Verify classification: should be native, not MCP + assert not hook._is_mcp_tool(responses_api_tool), ( + "Responses API tool with type=function + top-level name " + "should be classified as native, not MCP" + ) + + # Full hook test: all-native request should preserve tools + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Search GitHub"}], + "tools": [responses_api_tool], + "metadata": {}, + } + + result = await hook.async_pre_call_hook( + user_api_key_dict=Mock(), + cache=Mock(), + data=data, + call_type="completion", + ) + + # All tools are native → hook returns data with all tools preserved + filtered = (result or data)["tools"] + assert len(filtered) == 1, f"Native tool must survive, got {len(filtered)}" + assert filtered[0]["name"] == "github-search" + + print("✅ Responses API tool with MCP-matching name correctly classified as native") + + +@pytest.mark.asyncio +async def test_semantic_filter_hook_preserves_tool_order(): + """ + Regression test: tool ordering preservation. + + Given: An interleaved request [mcp_tool_A, native_tool, mcp_tool_B] + When: The hook filters tools (all MCP tools survive) + Then: The output order must match the original request order, + NOT native-first. + """ + from litellm.proxy._experimental.mcp_server.semantic_tool_filter import ( + SemanticMCPToolFilter, + ) + from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + from litellm.types.utils import Embedding, EmbeddingResponse + + mock_router = Mock() + + def mock_embedding_sync(*args, **kwargs): + return EmbeddingResponse( + data=[Embedding(embedding=[0.1] * 1536, index=0, object="embedding")], + model="text-embedding-3-small", + object="list", + usage={"prompt_tokens": 10, "total_tokens": 10}, + ) + + async def mock_embedding_async(*args, **kwargs): + return mock_embedding_sync() + + mock_router.embedding = mock_embedding_sync + mock_router.aembedding = mock_embedding_async + + filter_instance = SemanticMCPToolFilter( + embedding_model="text-embedding-3-small", + litellm_router_instance=mock_router, + top_k=5, + similarity_threshold=0.3, + enabled=True, + ) + + # Register MCP tools + mcp_tool_a = MCPTool( + name="github-search", + description="Search GitHub", + inputSchema={"type": "object"}, + ) + mcp_tool_b = MCPTool( + name="github-issue", + description="Create GitHub issue", + inputSchema={"type": "object"}, + ) + filter_instance._build_router([mcp_tool_a, mcp_tool_b]) + + # Mock filter_tools to return both MCP tools (deterministic) + filter_instance.filter_tools = AsyncMock( # type: ignore[method-assign] + return_value=[mcp_tool_a, mcp_tool_b] + ) + + # Native tool (interleaved between MCP tools) + native_tool = { + "type": "function", + "function": { + "name": "weather_lookup", + "description": "Look up weather", + "parameters": {"type": "object", "properties": {}}, + }, + } + + # Original order: [mcp_A, native, mcp_B] + original_tools = [mcp_tool_a, native_tool, mcp_tool_b] + + hook = SemanticToolFilterHook(filter_instance) + + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Search GitHub and check weather"}], + "tools": original_tools, + "metadata": {}, + } + + result = await hook.async_pre_call_hook( + user_api_key_dict=Mock(), + cache=Mock(), + data=data, + call_type="completion", + ) + + assert result is not None, "Hook should return modified data" + filtered = result["tools"] + + # All tools should survive + assert len(filtered) == 3, f"Expected 3 tools, got {len(filtered)}" + + # Order must be preserved: [mcp_A, native, mcp_B] + assert filtered[0] is mcp_tool_a, "First tool should be mcp_tool_a" + assert filtered[1] is native_tool, "Second tool should be native_tool" + assert filtered[2] is mcp_tool_b, "Third tool should be mcp_tool_b" + + print( + "✅ Tool ordering preserved: [mcp_A, native, mcp_B] maintained after filtering" + ) + + class TestGetToolsByNames: """ Regression coverage for SemanticMCPToolFilter._get_tools_by_names @@ -489,9 +913,7 @@ def test_exact_match_unchanged(self): {"name": "send_email", "description": "send mail"}, ] - matched = filter_instance._get_tools_by_names( - ["send_email"], available_tools - ) + matched = filter_instance._get_tools_by_names(["send_email"], available_tools) assert len(matched) == 1 assert matched[0]["name"] == "send_email" @@ -503,9 +925,7 @@ def test_client_prefix_with_underscore_separator(self): client_name = "litellm_" + canonical available_tools = [{"name": client_name, "description": "scrape"}] - matched = filter_instance._get_tools_by_names( - [canonical], available_tools - ) + matched = filter_instance._get_tools_by_names([canonical], available_tools) assert len(matched) == 1 # Must return the incoming tool unchanged so the client-facing @@ -516,13 +936,9 @@ def test_client_prefix_with_dash_separator(self): """Some clients use dash as alias separator; accept that too.""" filter_instance = self._make_filter() canonical = "weather_svc-get_weather" - available_tools = [ - {"name": "mcp-" + canonical, "description": "weather"} - ] + available_tools = [{"name": "mcp-" + canonical, "description": "weather"}] - matched = filter_instance._get_tools_by_names( - [canonical], available_tools - ) + matched = filter_instance._get_tools_by_names([canonical], available_tools) assert len(matched) == 1 assert matched[0]["name"] == "mcp-" + canonical @@ -552,9 +968,7 @@ def test_exact_match_preferred_over_prefixed(self): {"name": "litellm_" + canonical, "description": "wrapped"}, ] - matched = filter_instance._get_tools_by_names( - [canonical], available_tools - ) + matched = filter_instance._get_tools_by_names([canonical], available_tools) assert len(matched) == 1 assert matched[0]["name"] == canonical @@ -567,9 +981,7 @@ def test_same_tool_not_returned_twice(self): separator-anchored suffixes of ``litellm_api-fs-read_file``. """ filter_instance = self._make_filter() - available_tools = [ - {"name": "litellm_api-fs-read_file", "description": "read"} - ] + available_tools = [{"name": "litellm_api-fs-read_file", "description": "read"}] matched = filter_instance._get_tools_by_names( ["fs-read_file", "api-fs-read_file"], available_tools @@ -590,9 +1002,7 @@ def test_suffix_fallback_prefers_shortest_candidate(self): {"name": "my_" + canonical, "description": "plain search"}, ] - matched = filter_instance._get_tools_by_names( - [canonical], available_tools - ) + matched = filter_instance._get_tools_by_names([canonical], available_tools) assert len(matched) == 1 assert matched[0]["name"] == "my_" + canonical