phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-27 19:54:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix mcp type imports

Browse source
This commit is contained in:
Ishaan Jaff 2025-03-24 11:01:29 -07:00
parent f133bb07d1
commit d4b53d3fa9
1 changed files with 10 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

17
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -3,14 +3,17 @@ LiteLLM MCP Server Routes
""" """
import asyncio import asyncio
from typing import Any, Dict from typing import Any, Dict, List, Union
import mcp.types as types
from anyio import BrokenResourceError from anyio import BrokenResourceError
from fastapi import APIRouter, HTTPException, Request from fastapi import APIRouter, HTTPException, Request
from fastapi.responses import StreamingResponse from fastapi.responses import StreamingResponse
from mcp.server import NotificationOptions, Server from mcp.server import NotificationOptions, Server
from mcp.server.models import InitializationOptions from mcp.server.models import InitializationOptions
from mcp.types import EmbeddedResource as MCPEmbeddedResource
from mcp.types import ImageContent as MCPImageContent
from mcp.types import TextContent as MCPTextContent
from mcp.types import Tool as MCPTool
from pydantic import ValidationError from pydantic import ValidationError
from litellm._logging import verbose_logger from litellm._logging import verbose_logger
@ -36,14 +39,14 @@ sse = SseServerTransport("/mcp/sse/messages")
@server.list_tools() @server.list_tools()
async def list_tools() -> list[types.Tool]: async def list_tools() -> list[MCPTool]:
""" """
List all available tools List all available tools
""" """
tools = [] tools = []
for tool in global_mcp_tool_registry.list_tools(): for tool in global_mcp_tool_registry.list_tools():
tools.append( tools.append(
types.Tool( MCPTool(
name=tool.name, name=tool.name,
description=tool.description, description=tool.description,
inputSchema=tool.input_schema, inputSchema=tool.input_schema,
@ -56,7 +59,7 @@ async def list_tools() -> list[types.Tool]:
@server.call_tool() @server.call_tool()
async def handle_call_tool( async def handle_call_tool(
name: str, arguments: Dict[str, Any] | None name: str, arguments: Dict[str, Any] | None
) -> list[types.TextContent | types.ImageContent | types.EmbeddedResource]: ) -> List[Union[MCPTextContent, MCPImageContent, MCPEmbeddedResource]]:
""" """
Call a specific tool with the provided arguments Call a specific tool with the provided arguments
""" """
@ -68,9 +71,9 @@ async def handle_call_tool(
try: try:
result = tool.handler(**arguments) result = tool.handler(**arguments)
return [types.TextContent(text=str(result), type="text")] return [MCPTextContent(text=str(result), type="text")]
except Exception as e: except Exception as e:
return [types.TextContent(text=f"Error: {str(e)}", type="text")] return [MCPTextContent(text=f"Error: {str(e)}", type="text")]
@router.get("/", response_class=StreamingResponse) @router.get("/", response_class=StreamingResponse)