diff --git a/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py b/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py new file mode 100644 index 00000000000..ee1a447f73b --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/semantic_tool_filter.py @@ -0,0 +1,254 @@ +""" +Semantic MCP Tool Filtering using semantic-router + +This module provides semantic filtering for MCP tools to reduce context window size +and improve tool selection accuracy. It leverages the existing semantic-router library +and LiteLLMRouterEncoder to provide efficient tool filtering based on user queries. +""" +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union + +from litellm._logging import verbose_logger + +if TYPE_CHECKING: + from mcp.types import Tool as MCPTool + from semantic_router.routers import SemanticRouter + from semantic_router.schema import RouteChoice + + from litellm.router import Router + + +class SemanticMCPToolFilter: + """ + Filters MCP tools using semantic-router library. + + Converts MCP tools to semantic-router Routes and uses SemanticRouter + to find the most relevant tools for a given user query. + """ + + def __init__( + self, + embedding_model: str, + litellm_router_instance: "Router", + top_k: int = 10, + similarity_threshold: float = 0.3, + enabled: bool = True, + ): + """ + Initialize the semantic tool filter. + + Args: + embedding_model: Model to use for generating embeddings (e.g., "text-embedding-3-small") + litellm_router_instance: Router instance for embedding generation + top_k: Maximum number of tools to return + similarity_threshold: Minimum similarity score for filtering + enabled: Whether filtering is enabled + """ + self.enabled = enabled + self.top_k = top_k + self.similarity_threshold = similarity_threshold + self.embedding_model = embedding_model + self.router_instance = litellm_router_instance + self.tool_router: Optional["SemanticRouter"] = None + self._tool_map: Dict[str, "MCPTool"] = {} # name -> tool + + verbose_logger.debug( + f"Initialized SemanticMCPToolFilter: enabled={enabled}, " + f"top_k={top_k}, threshold={similarity_threshold}, " + f"model={embedding_model}" + ) + + def _mcp_tools_to_routes(self, tools: List["MCPTool"]) -> List: + """ + Convert MCP tools to semantic-router Routes. + + Args: + tools: List of MCP tools + + Returns: + List of Route objects + """ + from semantic_router.routers.base import Route + + routes = [] + self._tool_map = {} + + for tool in tools: + self._tool_map[tool.name] = tool + + # Use tool description as both description and utterance + description = tool.description or tool.name + utterances = [description] if description else [] + + routes.append( + Route( + name=tool.name, + description=description, + utterances=utterances, + score_threshold=self.similarity_threshold, + ) + ) + + verbose_logger.debug(f"Converted {len(tools)} MCP tools to Routes") + return routes + + def rebuild_router(self, tools: List["MCPTool"]) -> None: + """ + Rebuild semantic router with updated tools. + + This should be called whenever the tool list changes (server add/update/remove). + + Args: + tools: Updated list of all available MCP tools + """ + from semantic_router.routers import SemanticRouter + + from litellm.router_strategy.auto_router.litellm_encoder import ( + LiteLLMRouterEncoder, + ) + + if not tools: + self.tool_router = None + verbose_logger.debug("No tools provided, semantic router set to None") + return + + try: + routes = self._mcp_tools_to_routes(tools) + + self.tool_router = SemanticRouter( + routes=routes, + encoder=LiteLLMRouterEncoder( + litellm_router_instance=self.router_instance, + model_name=self.embedding_model, + score_threshold=self.similarity_threshold, + ), + auto_sync="local", # Build index immediately + ) + + verbose_logger.info( + f"Rebuilt semantic router with {len(routes)} tool routes" + ) + + except Exception as e: + verbose_logger.error(f"Failed to rebuild semantic router: {e}") + self.tool_router = None + raise + + async def filter_tools( + self, + query: str, + available_tools: List["MCPTool"], + top_k: Optional[int] = None, + ) -> List["MCPTool"]: + """ + Filter tools semantically based on query. + + Args: + query: User query to match against tools + available_tools: Full list of available tools + top_k: Override default top_k (optional) + + Returns: + Filtered and ordered list of tools (up to top_k) + """ + if not self.enabled or not available_tools: + return available_tools + + if not query or not query.strip(): + verbose_logger.debug("Empty query, returning all tools") + return available_tools + + top_k = top_k or self.top_k + + try: + # Rebuild router if needed (first time or tools changed) + if self.tool_router is None: + verbose_logger.debug("Router not initialized, rebuilding...") + self.rebuild_router(available_tools) + + if self.tool_router is None: + verbose_logger.warning("Router rebuild failed, returning all tools") + return available_tools + + # Query semantic router + from semantic_router.schema import RouteChoice + + verbose_logger.debug(f"Querying semantic router with: '{query[:50]}...'") + matches = self.tool_router(text=query) + + if not matches: + verbose_logger.warning( + f"No tools matched query. Returning all {len(available_tools)} tools." + ) + return available_tools + + # Extract matched tool names + matched_names: List[str] = [] + if isinstance(matches, RouteChoice): + matched_names = [matches.name] + elif isinstance(matches, list): + # semantic-router returns list of RouteChoice, take top_k + matched_names = [m.name for m in matches[:top_k] if hasattr(m, 'name')] + + if not matched_names: + verbose_logger.warning("No matched tool names extracted, returning all tools") + return available_tools + + # Filter available tools by matched names (preserve order from semantic router) + matched_name_set = set(matched_names) + filtered = [ + tool for tool in available_tools + if tool.name in matched_name_set + ] + + # Reorder based on semantic router's ordering + ordered_filtered = [] + for name in matched_names: + for tool in filtered: + if tool.name == name: + ordered_filtered.append(tool) + break + + verbose_logger.info( + f"Semantic filter: {len(available_tools)} -> {len(ordered_filtered)} tools. " + f"Top matches: {matched_names[:3]}" + ) + + return ordered_filtered if ordered_filtered else available_tools + + except Exception as e: + verbose_logger.error( + f"Semantic tool filter failed: {e}. Returning all tools.", + exc_info=True + ) + return available_tools + + def extract_user_query(self, messages: List[Dict[str, Any]]) -> str: + """ + Extract user query from messages. + + Args: + messages: List of message dictionaries + + Returns: + Extracted query string + """ + # Get the last user message + for msg in reversed(messages): + if msg.get("role") == "user": + content = msg.get("content", "") + + # Handle string content + if isinstance(content, str): + return content + + # Handle content blocks (list) + elif isinstance(content, list): + texts = [] + for block in content: + if isinstance(block, dict) and block.get("type") == "text": + texts.append(block.get("text", "")) + elif isinstance(block, str): + texts.append(block) + return " ".join(texts) + + return ""