diff --git a/tests/mcp_tests/test_semantic_tool_filter_e2e.py b/tests/mcp_tests/test_semantic_tool_filter_e2e.py new file mode 100644 index 00000000000..930d005ee46 --- /dev/null +++ b/tests/mcp_tests/test_semantic_tool_filter_e2e.py @@ -0,0 +1,74 @@ +""" +End-to-end test for MCP Semantic Tool Filtering +""" +import asyncio +import os +import sys +from unittest.mock import Mock + +import pytest + +sys.path.insert(0, os.path.abspath("../..")) +import json + +from mcp.types import Tool as MCPTool + + +@pytest.mark.asyncio +async def test_e2e_semantic_filter(): + """E2E: Load router/filter and verify hook filters tools.""" + from litellm import Router + from litellm.proxy.hooks.semantic_tool_filter_hook import SemanticToolFilterHook + 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"}, + }] + ) + + filter_instance = SemanticMCPToolFilter( + embedding_model="text-embedding-3-small", + litellm_router_instance=router, + top_k=3, + enabled=True, + ) + + hook = SemanticToolFilterHook(filter_instance) + + # Create 10 tools + 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"}), + ] + + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Send an email"}], + "tools": tools, + } + + # Call hook + result = await hook.async_pre_call_hook( + user_api_key_dict=Mock(), + cache=Mock(), + data=data, + call_type="completion", + ) + + print("final result tools: ", json.dumps(result["tools"], indent=4)) + + # Single assertion: hook filtered tools + assert result and len(result["tools"]) < len(tools)