diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 3565c4468c..8e83aea69a 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1099,7 +1099,7 @@ class Logging(LiteLLMLoggingBaseClass): standard_built_in_tools_params=self.standard_built_in_tools_params, ) ) - elif isinstance(result, dict): # pass-through endpoints + elif isinstance(result, dict) or isinstance(result, list): ## STANDARDIZED LOGGING PAYLOAD self.model_call_details["standard_logging_object"] = ( get_standard_logging_object_payload( diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index a2b378fe09..c73eabd629 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -6,6 +6,7 @@ This class is responsible for managing MCP SSE clients. This is a Proxy """ +import asyncio import json from typing import Any, Dict, List, Optional @@ -60,6 +61,8 @@ class MCPServerManager: f"Loaded MCP Servers: {json.dumps(self.mcp_servers, indent=4, default=str)}" ) + self.initialize_tool_name_to_mcp_server_name_mapping() + async def list_tools(self) -> List[MCPTool]: """ List all tools available across all MCP Servers. @@ -102,6 +105,29 @@ class MCPServerManager: return tools_result.tools + def initialize_tool_name_to_mcp_server_name_mapping(self): + """ + On startup, initialize the tool name to MCP server name mapping + """ + try: + if asyncio.get_running_loop(): + asyncio.create_task( + self._initialize_tool_name_to_mcp_server_name_mapping() + ) + except RuntimeError as e: # no running event loop + verbose_logger.exception( + f"No running event loop - skipping tool name to MCP server name mapping initialization: {str(e)}" + ) + + async def _initialize_tool_name_to_mcp_server_name_mapping(self): + """ + Call list_tools for each server and update the tool name to MCP server name mapping + """ + for server in self.mcp_servers: + tools = await self._get_tools_from_server(server) + for tool in tools: + self.tool_name_to_mcp_server_name_mapping[tool.name] = server.name + async def call_tool(self, name: str, arguments: Dict[str, Any]): """ Call a tool with the given name and arguments diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 3e94576efb..32c33904a3 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -3,7 +3,7 @@ LiteLLM MCP Server Routes """ import asyncio -from typing import Any, Dict, List, Union +from typing import Any, Dict, List, Optional, Union from anyio import BrokenResourceError from fastapi import APIRouter, Depends, HTTPException, Request @@ -11,10 +11,12 @@ from fastapi.responses import StreamingResponse from pydantic import ValidationError from litellm._logging import verbose_logger +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.types.mcp_server.mcp_server_manager import ( ListMCPToolsRestAPIResponseObject, ) +from litellm.utils import client # Check if MCP is available # "mcp" requires python 3.10 or higher, but several litellm users use python 3.8 @@ -110,8 +112,9 @@ if MCP_AVAILABLE: ) return response + @client async def call_mcp_tool( - name: str, arguments: Dict[str, Any] | None + name: str, arguments: Optional[Dict[str, Any]] = None, **kwargs: Any ) -> List[Union[MCPTextContent, MCPImageContent, MCPEmbeddedResource]]: """ Call a specific tool with the provided arguments @@ -221,14 +224,23 @@ if MCP_AVAILABLE: return list_tools_result @router.post("/tools/call", dependencies=[Depends(user_api_key_auth)]) - async def call_tool_rest_api(request: Request): + async def call_tool_rest_api( + request: Request, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + ): + """ + REST API to call a specific MCP tool with the provided arguments + """ + from litellm.proxy.proxy_server import add_litellm_data_to_request, proxy_config + data = await request.json() - name = data.get("name") - arguments = data.get("arguments") - return await call_mcp_tool( - name=name, - arguments=arguments, + data = await add_litellm_data_to_request( + data=data, + request=request, + user_api_key_dict=user_api_key_dict, + proxy_config=proxy_config, ) + return await call_mcp_tool(**data) options = InitializationOptions( server_name="litellm-mcp-server",