mcp server manager

litellm/proxy/_experimental/mcp_server/mcp_server_manager.py

@@ -0,0 +1,125 @@
+"""
+MCP Client Manager
+
+This class is responsible for managing MCP SSE clients.
+
+This is a Proxy 
+"""
+
+import json
+from typing import Any, Dict, List, Optional
+
+from mcp import ClientSession
+from mcp.client.sse import sse_client
+from mcp.types import Tool as MCPTool
+
+from litellm._logging import verbose_logger
+from litellm.types.mcp_server.mcp_server_manager import MCPSSEServer
+
+
+class MCPServerManager:
+    def __init__(self):
+        self.mcp_servers: List[MCPSSEServer] = []
+        """
+        eg.
+        [
+            {
+                "name": "zapier_mcp_server",
+                "url": "https://actions.zapier.com/mcp/sk-ak-2ew3bofIeQIkNoeKIdXrF1Hhhp/sse"
+            },
+            {
+                "name": "google_drive_mcp_server",
+                "url": "https://actions.zapier.com/mcp/sk-ak-2ew3bofIeQIkNoeKIdXrF1Hhhp/sse"
+            }
+        ]
+        """
+
+        self.tool_name_to_mcp_server_name_mapping: Dict[str, str] = {}
+        """
+        {
+            "gmail_send_email": "zapier_mcp_server",
+        }
+        """
+
+    def load_servers_from_config(self, mcp_servers_config: Dict[str, Any]):
+        """
+        Load the MCP Servers from the config
+        """
+        for server_name, server_config in mcp_servers_config.items():
+            self.mcp_servers.append(
+                MCPSSEServer(
+                    name=server_name,
+                    url=server_config["url"],
+                )
+            )
+        verbose_logger.debug(
+            f"Loaded MCP Servers: {json.dumps(self.mcp_servers, indent=4, default=str)}"
+        )
+
+    async def list_tools(self) -> List[MCPTool]:
+        """
+        List all tools available across all MCP Servers.
+
+        Returns:
+            List[MCPTool]: Combined list of tools from all servers
+        """
+        list_tools_result: List[MCPTool] = []
+        verbose_logger.debug("SSE SERVER MANAGER LISTING TOOLS")
+
+        for server in self.mcp_servers:
+            tools = await self._get_tools_from_server(server)
+            list_tools_result.extend(tools)
+
+        return list_tools_result
+
+    async def _get_tools_from_server(self, server: MCPSSEServer) -> List[MCPTool]:
+        """
+        Helper method to get tools from a single MCP server.
+
+        Args:
+            server (MCPSSEServer): The server to query tools from
+
+        Returns:
+            List[MCPTool]: List of tools available on the server
+        """
+        verbose_logger.debug(f"Connecting to url: {server.url}")
+
+        async with sse_client(url=server.url) as (read, write):
+            async with ClientSession(read, write) as session:
+                server.client_session = session
+                await server.client_session.initialize()
+
+                tools_result = await server.client_session.list_tools()
+                verbose_logger.debug(f"Tools from {server.name}: {tools_result}")
+
+                # Update tool to server mapping
+                for tool in tools_result.tools:
+                    self.tool_name_to_mcp_server_name_mapping[tool.name] = server.name
+
+                return tools_result.tools
+
+    async def call_tool(self, name: str, arguments: Dict[str, Any]):
+        """
+        Call a tool with the given name and arguments
+        """
+        mcp_server = self._get_mcp_server_from_tool_name(name)
+        if mcp_server is None:
+            raise ValueError(f"Tool {name} not found")
+        async with sse_client(url=mcp_server.url) as (read, write):
+            async with ClientSession(read, write) as session:
+                mcp_server.client_session = session
+                await mcp_server.client_session.initialize()
+                return await mcp_server.client_session.call_tool(name, arguments)
+
+    def _get_mcp_server_from_tool_name(self, tool_name: str) -> Optional[MCPSSEServer]:
+        """
+        Get the MCP Server from the tool name
+        """
+        if tool_name in self.tool_name_to_mcp_server_name_mapping:
+            for server in self.mcp_servers:
+                if server.name == self.tool_name_to_mcp_server_name_mapping[tool_name]:
+                    return server
+        return None
+
+
+global_mcp_server_manager: MCPServerManager = MCPServerManager()

litellm/proxy/_experimental/mcp_server/server.py

@@ -36,6 +36,7 @@ if MCP_AVAILABLE:
     from mcp.types import TextContent as MCPTextContent
     from mcp.types import Tool as MCPTool
 
+    from .mcp_server_manager import global_mcp_server_manager
     from .sse_transport import SseServerTransport
     from .tool_registry import global_mcp_tool_registry
 
@@ -67,24 +68,64 @@ if MCP_AVAILABLE:
                     inputSchema=tool.input_schema,
                 )
             )
-
+        verbose_logger.debug(
+            "GLOBAL MCP TOOLS: %s", global_mcp_tool_registry.list_tools()
+        )
+        sse_tools: List[MCPTool] = await global_mcp_server_manager.list_tools()
+        verbose_logger.debug("SSE TOOLS: %s", sse_tools)
+        if sse_tools is not None:
+            tools.extend(sse_tools)
         return tools
 
     @server.call_tool()
-    async def handle_call_tool(
+    async def mcp_server_tool_call(
         name: str, arguments: Dict[str, Any] | None
     ) -> List[Union[MCPTextContent, MCPImageContent, MCPEmbeddedResource]]:
         """
         Call a specific tool with the provided arguments
+
+        Args:
+            name (str): Name of the tool to call
+            arguments (Dict[str, Any] | None): Arguments to pass to the tool
+
+        Returns:
+            List[Union[MCPTextContent, MCPImageContent, MCPEmbeddedResource]]: Tool execution results
+
+        Raises:
+            HTTPException: If tool not found or arguments missing
         """
-        tool = global_mcp_tool_registry.get_tool(name)
-        if not tool:
-            raise HTTPException(status_code=404, detail=f"Tool '{name}' not found")
+        # Validate arguments
         if arguments is None:
             raise HTTPException(
                 status_code=400, detail="Request arguments are required"
             )
 
+        # Try managed server tool first
+        if name in global_mcp_server_manager.tool_name_to_mcp_server_name_mapping:
+            return await _handle_managed_mcp_tool(name, arguments)
+
+        # Fall back to local tool registry
+        return await _handle_local_mcp_tool(name, arguments)
+
+    async def _handle_managed_mcp_tool(
+        name: str, arguments: Dict[str, Any]
+    ) -> List[Union[MCPTextContent, MCPImageContent, MCPEmbeddedResource]]:
+        """Handle tool execution for managed server tools"""
+        call_tool_result = await global_mcp_server_manager.call_tool(
+            name=name,
+            arguments=arguments,
+        )
+        verbose_logger.debug("CALL TOOL RESULT: %s", call_tool_result)
+        return call_tool_result.content
+
+    async def _handle_local_mcp_tool(
+        name: str, arguments: Dict[str, Any]
+    ) -> List[Union[MCPTextContent, MCPImageContent, MCPEmbeddedResource]]:
+        """Handle tool execution for local registry tools"""
+        tool = global_mcp_tool_registry.get_tool(name)
+        if not tool:
+            raise HTTPException(status_code=404, detail=f"Tool '{name}' not found")
+
         try:
             result = tool.handler(**arguments)
             return [MCPTextContent(text=str(result), type="text")]

litellm/types/mcp_server/mcp_server_manager.py

@@ -0,0 +1,11 @@
+from typing import Optional
+
+from mcp import ClientSession
+from pydantic import BaseModel, ConfigDict
+
+
+class MCPSSEServer(BaseModel):
+    name: str
+    url: str
+    client_session: Optional[ClientSession] = None
+    model_config = ConfigDict(arbitrary_types_allowed=True)