diff --git a/tests/mcp_tests/test_semantic_tool_filter_e2e.py b/tests/mcp_tests/test_semantic_tool_filter_e2e.py index b7d9e4a0784..c619f0f6d46 100644 --- a/tests/mcp_tests/test_semantic_tool_filter_e2e.py +++ b/tests/mcp_tests/test_semantic_tool_filter_e2e.py @@ -1,6 +1,7 @@ """ End-to-end test for MCP Semantic Tool Filtering """ + import asyncio import os import sys @@ -15,6 +16,7 @@ from mcp.types import Tool as MCPTool # Check if semantic-router is available try: import semantic_router + SEMANTIC_ROUTER_AVAILABLE = True except ImportError: SEMANTIC_ROUTER_AVAILABLE = False @@ -26,8 +28,7 @@ except ImportError: reason="semantic-router not installed. Install the `litellm[semantic-router]` extra." ) @pytest.mark.skipif( - not os.environ.get("OPENAI_API_KEY"), - reason="OPENAI_API_KEY not set in environment" + not os.environ.get("OPENAI_API_KEY"), reason="OPENAI_API_KEY not set in environment" ) async def test_e2e_semantic_filter(): """E2E: Load router/filter and verify hook filters tools.""" @@ -36,48 +37,88 @@ async def test_e2e_semantic_filter(): from litellm.proxy._experimental.mcp_server.semantic_tool_filter import ( SemanticMCPToolFilter, ) - + # Create router and filter router = Router( - model_list=[{ - "model_name": "text-embedding-3-small", - "litellm_params": {"model": "openai/text-embedding-3-small"}, - }] + model_list=[ + { + "model_name": "text-embedding-3-small", + "litellm_params": {"model": "openai/text-embedding-3-small"}, + } + ] ) - + filter_instance = SemanticMCPToolFilter( embedding_model="text-embedding-3-small", litellm_router_instance=router, top_k=3, enabled=True, ) - - # Create 10 tools + + # Create 10 tools (MCP-prefixed names so the hook classifies them as MCP) tools = [ - MCPTool(name="gmail_send", description="Send an email via Gmail", inputSchema={"type": "object"}), - MCPTool(name="calendar_create", description="Create a calendar event", inputSchema={"type": "object"}), - MCPTool(name="file_upload", description="Upload a file", inputSchema={"type": "object"}), - MCPTool(name="web_search", description="Search the web", inputSchema={"type": "object"}), - MCPTool(name="slack_send", description="Send Slack message", inputSchema={"type": "object"}), - MCPTool(name="doc_read", description="Read document", inputSchema={"type": "object"}), - MCPTool(name="db_query", description="Query database", inputSchema={"type": "object"}), - MCPTool(name="api_call", description="Make API call", inputSchema={"type": "object"}), - MCPTool(name="task_create", description="Create task", inputSchema={"type": "object"}), - MCPTool(name="note_add", description="Add note", inputSchema={"type": "object"}), + MCPTool( + name="gmail-send_email", + description="Send an email via Gmail", + inputSchema={"type": "object"}, + ), + MCPTool( + name="calendar-create_event", + description="Create a calendar event", + inputSchema={"type": "object"}, + ), + MCPTool( + name="files-upload", + description="Upload a file", + inputSchema={"type": "object"}, + ), + MCPTool( + name="web-search", + description="Search the web", + inputSchema={"type": "object"}, + ), + MCPTool( + name="slack-send_message", + description="Send Slack message", + inputSchema={"type": "object"}, + ), + MCPTool( + name="docs-read", + description="Read document", + inputSchema={"type": "object"}, + ), + MCPTool( + name="db-query", + description="Query database", + inputSchema={"type": "object"}, + ), + MCPTool( + name="api-call", description="Make API call", inputSchema={"type": "object"} + ), + MCPTool( + name="tasks-create", + description="Create task", + inputSchema={"type": "object"}, + ), + MCPTool( + name="notes-add", description="Add note", inputSchema={"type": "object"} + ), ] - + # Build router with test tools filter_instance._build_router(tools) - + hook = SemanticToolFilterHook(filter_instance) - + data = { "model": "gpt-4", - "messages": [{"role": "user", "content": "Send an email and create a calendar event"}], + "messages": [ + {"role": "user", "content": "Send an email and create a calendar event"} + ], "tools": tools, "metadata": {}, # Initialize metadata dict for hook to store filter stats } - + # Call hook result = await hook.async_pre_call_hook( user_api_key_dict=Mock(), @@ -87,7 +128,11 @@ async def test_e2e_semantic_filter(): ) # Single assertion: hook filtered tools - assert result and len(result["tools"]) < len(tools), f"Expected filtered tools, got {len(result['tools'])} tools (original: {len(tools)})" - - print(f"✅ E2E test passed: Filtering reduced tools from {len(tools)} to {len(result['tools'])}") + assert result and len(result["tools"]) < len( + tools + ), f"Expected filtered tools, got {len(result['tools'])} tools (original: {len(tools)})" + + print( + f"✅ E2E test passed: Filtering reduced tools from {len(tools)} to {len(result['tools'])}" + ) print(f" Filtered tools: {[t.name for t in result['tools']]}")