From 10486dd17ebed082cf8833f451b3341b2c25f1be Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 29 Mar 2025 18:40:58 -0700 Subject: [PATCH] fix listing mcp tools --- .../proxy/_experimental/mcp_server/server.py | 21 ++++++++++++++----- .../types/mcp_server/mcp_server_manager.py | 17 --------------- 2 files changed, 16 insertions(+), 22 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 603370711f..fe1eccb048 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -8,16 +8,14 @@ from typing import Any, Dict, List, Optional, Union from anyio import BrokenResourceError from fastapi import APIRouter, Depends, HTTPException, Request from fastapi.responses import StreamingResponse -from pydantic import ValidationError +from pydantic import ConfigDict, ValidationError from litellm._logging import verbose_logger from litellm.constants import MCP_TOOL_NAME_PREFIX from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj 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.types.mcp_server.mcp_server_manager import MCPInfo from litellm.types.utils import StandardLoggingMCPToolCall from litellm.utils import client @@ -49,6 +47,19 @@ if MCP_AVAILABLE: from .sse_transport import SseServerTransport from .tool_registry import global_mcp_tool_registry + ###################################################### + ############ MCP Tools List REST API Response Object # + # Defined here because we don't want to add `mcp` as a + # required dependency for `litellm` pip package + ###################################################### + class ListMCPToolsRestAPIResponseObject(MCPTool): + """ + Object returned by the /tools/list REST API route. + """ + + mcp_info: Optional[MCPInfo] = None + model_config = ConfigDict(arbitrary_types_allowed=True) + ######################################################## ############ Initialize the MCP Server ################# ######################################################## @@ -224,7 +235,7 @@ if MCP_AVAILABLE: ############ MCP Server REST API Routes ################# ######################################################## @router.get("/tools/list", dependencies=[Depends(user_api_key_auth)]) - async def list_tool_rest_api(): + async def list_tool_rest_api() -> List[ListMCPToolsRestAPIResponseObject]: """ List all available tools with information about the server they belong to. diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 9752423e7e..aecd11aa1a 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -3,14 +3,6 @@ from typing import TYPE_CHECKING, Optional 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 @@ -22,12 +14,3 @@ class MCPSSEServer(BaseModel): url: str mcp_info: Optional[MCPInfo] = None model_config = ConfigDict(arbitrary_types_allowed=True) - - -class ListMCPToolsRestAPIResponseObject(MCPTool): - """ - Object returned by the /tools/list REST API route. - """ - - mcp_info: Optional[MCPInfo] = None - model_config = ConfigDict(arbitrary_types_allowed=True)