diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index c73eabd629d..9eb44f55977 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -8,11 +8,17 @@ This is a Proxy import asyncio import json -from typing import Any, Dict, List, Optional +from typing import TYPE_CHECKING, Any, Dict, List, Optional -from mcp import ClientSession -from mcp.client.sse import sse_client -from mcp.types import Tool as MCPTool +if TYPE_CHECKING: + from mcp import ClientSession + from mcp.client.sse import sse_client + from mcp.types import Tool as MCPTool +else: + # Provide fallback types for runtime incase `mcp` is not installed + ClientSession = None + MCPTool = object + sse_client = None from litellm._logging import verbose_logger from litellm.types.mcp_server.mcp_server_manager import MCPInfo, MCPSSEServer @@ -93,10 +99,9 @@ class MCPServerManager: 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() + await session.initialize() - tools_result = await server.client_session.list_tools() + tools_result = await session.list_tools() verbose_logger.debug(f"Tools from {server.name}: {tools_result}") # Update tool to server mapping @@ -137,9 +142,8 @@ class MCPServerManager: 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) + await session.initialize() + return await session.call_tool(name, arguments) def _get_mcp_server_from_tool_name(self, tool_name: str) -> Optional[MCPSSEServer]: """ diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 981926df3cb..9752423e7e6 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -1,10 +1,16 @@ -from typing import Any, Dict, List, Optional +from typing import TYPE_CHECKING, Optional -from mcp import ClientSession -from mcp.types import Tool as MCPTool from pydantic import BaseModel, ConfigDict from typing_extensions import TypedDict +if TYPE_CHECKING: + from mcp import ClientSession + from mcp.types import Tool as MCPTool +else: + # Provide fallback types for runtime incase `mcp` is not installed + ClientSession = None + MCPTool = object + class MCPInfo(TypedDict, total=False): server_name: str @@ -14,7 +20,6 @@ class MCPInfo(TypedDict, total=False): class MCPSSEServer(BaseModel): name: str url: str - client_session: Optional[ClientSession] = None mcp_info: Optional[MCPInfo] = None model_config = ConfigDict(arbitrary_types_allowed=True)