diff --git a/litellm/proxy/_lazy_openapi_snapshot.py b/litellm/proxy/_lazy_openapi_snapshot.py index c63ff8d0733..a9d5396cc29 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.py +++ b/litellm/proxy/_lazy_openapi_snapshot.py @@ -86,7 +86,6 @@ def generate_snapshot() -> Dict[str, Dict]: import importlib from fastapi.openapi.utils import get_openapi - from litellm.proxy._lazy_features import LAZY_FEATURES from litellm.proxy.proxy_server import app, ensure_unique_openapi_operation_ids @@ -133,7 +132,6 @@ def generate_snapshot() -> Dict[str, Dict]: } return fragments - if __name__ == "__main__": fragments = generate_snapshot() SNAPSHOT_FILE.write_text(json.dumps(fragments, indent=2, sort_keys=True) + "\n") diff --git a/litellm/proxy/hooks/mcp_semantic_filter/hook.py b/litellm/proxy/hooks/mcp_semantic_filter/hook.py index e05e346e3f9..645e0eb84aa 100644 --- a/litellm/proxy/hooks/mcp_semantic_filter/hook.py +++ b/litellm/proxy/hooks/mcp_semantic_filter/hook.py @@ -5,6 +5,7 @@ Pre-call hook that filters MCP tools semantically before LLM inference. Reduces context window size and improves tool selection accuracy. """ +from collections import Counter from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union from litellm._logging import verbose_proxy_logger @@ -177,6 +178,30 @@ class SemanticToolFilterHook(CustomLogger): return native_tools, mcp_tools + def _merge_filtered_tools_preserving_request_order( + self, + tools: List[Any], + filtered_mcp_tools: List[Any], + ) -> List[Any]: + filtered_mcp_tool_counts = Counter( + tool_name + for tool in filtered_mcp_tools + if (tool_name := self._get_tool_name(tool)) + ) + filtered_tools: List[Any] = [] + + for tool in tools: + if not self._is_mcp_router_tool(tool): + filtered_tools.append(tool) + continue + + tool_name = self._get_tool_name(tool) + if filtered_mcp_tool_counts[tool_name] > 0: + filtered_tools.append(tool) + filtered_mcp_tool_counts[tool_name] -= 1 + + return filtered_tools + async def async_pre_call_hook( self, user_api_key_dict: "UserAPIKeyAuth", @@ -272,7 +297,7 @@ class SemanticToolFilterHook(CustomLogger): f"with query: '{user_query[:50]}...'" ) - native_tools, mcp_tools = self._partition_mcp_router_tools(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" @@ -284,7 +309,9 @@ class SemanticToolFilterHook(CustomLogger): query=user_query, available_tools=mcp_tools, # type: ignore ) - filtered_tools = native_tools + filtered_mcp_tools + filtered_tools = self._merge_filtered_tools_preserving_request_order( + tools, filtered_mcp_tools + ) # Always update tools and emit header (even if count unchanged) data["tools"] = filtered_tools 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 2fdd95e65fa..49dda138977 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 @@ -558,6 +558,65 @@ async def test_semantic_filter_hook_preserves_native_openai_tools(): ) +@pytest.mark.asyncio +async def test_semantic_filter_hook_preserves_remaining_tool_order(): + """ + Filtering MCP router tools should not reorder caller-owned native tools + around the MCP tools that survive semantic filtering. + """ + 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_search_tool = {"name": "github-search", "description": "Search GitHub"} + mcp_issue_tool = {"name": "github-issue", "description": "Create GitHub issue"} + filter_instance._tool_map = { + mcp_search_tool["name"]: mcp_search_tool, + mcp_issue_tool["name"]: mcp_issue_tool, + } + filter_instance.filter_tools = AsyncMock( # type: ignore[method-assign] + return_value=[mcp_issue_tool, mcp_search_tool] + ) + + 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": "Search GitHub"}], + "tools": [mcp_search_tool, native_tool, mcp_issue_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 + assert result["tools"] == [mcp_search_tool, native_tool, mcp_issue_tool] + assert ( + result["metadata"]["litellm_semantic_filter_tools"] + == "github-search,weather_lookup,github-issue" + ) + + @pytest.mark.asyncio async def test_semantic_filter_hook_skips_all_native_openai_tools(): """