diff --git a/litellm/experimental_mcp_client/tools.py b/litellm/experimental_mcp_client/tools.py index 592b0767fd..cdc26af4b7 100644 --- a/litellm/experimental_mcp_client/tools.py +++ b/litellm/experimental_mcp_client/tools.py @@ -1,5 +1,5 @@ import json -from typing import List, Literal, Union +from typing import Dict, List, Literal, Union from mcp import ClientSession from mcp.types import CallToolRequestParams as MCPCallToolRequestParams @@ -77,7 +77,7 @@ def _get_function_arguments(function: FunctionDefinition) -> dict: def transform_openai_tool_call_request_to_mcp_tool_call_request( - openai_tool: ChatCompletionMessageToolCall, + openai_tool: Union[ChatCompletionMessageToolCall, Dict], ) -> MCPCallToolRequestParams: """Convert an OpenAI ChatCompletionMessageToolCall to an MCP CallToolRequestParams.""" function = openai_tool["function"] diff --git a/tests/pass_through_tests/test_mcp_routes.py b/tests/pass_through_tests/test_mcp_routes.py index 3435d53585..f7f0a0d17f 100644 --- a/tests/pass_through_tests/test_mcp_routes.py +++ b/tests/pass_through_tests/test_mcp_routes.py @@ -1,85 +1,60 @@ -# Create server parameters for stdio connection import asyncio -import os - -import pytest -from langchain_mcp_adapters.tools import load_mcp_tools -from langchain_openai import ChatOpenAI -from langgraph.prebuilt import create_react_agent +from openai import AsyncOpenAI +from openai.types.chat import ChatCompletionUserMessageParam from mcp import ClientSession from mcp.client.sse import sse_client from litellm.experimental_mcp_client.tools import ( transform_mcp_tool_to_openai_tool, transform_openai_tool_call_request_to_mcp_tool_call_request, ) -import json -@pytest.mark.asyncio -async def test_mcp_routes(): - model = ChatOpenAI( - model="gpt-4o", api_key="sk-1234", base_url="http://localhost:4000" - ) +async def main(): + # Initialize clients + client = AsyncOpenAI(api_key="sk-1234", base_url="http://localhost:4000") - async with sse_client(url="http://localhost:4000/mcp/") as (read, write): - async with ClientSession(read, write) as session: - # Initialize the connection - print("Initializing session") - await session.initialize() - print("Session initialized") - - # Get tools - print("Loading tools") - tools = await load_mcp_tools(session) - print("Tools loaded") - print(tools) - - # Create and run the agent - agent = create_react_agent(model, tools) - agent_response = await agent.ainvoke({"messages": "Send an "}) - print(agent_response) - - -@pytest.mark.asyncio -async def test_mcp_routes_with_vertex_ai(): - # Create and run the agent - from openai import AsyncOpenAI - - openai_client = AsyncOpenAI(api_key="sk-1234", base_url="http://localhost:4000") - async with sse_client(url="http://localhost:4000/mcp/") as (read, write): + # Connect to MCP + async with sse_client("http://localhost:4000/mcp/") as (read, write): async with ClientSession(read, write) as session: await session.initialize() - MCP_TOOLS = await session.list_tools() + mcp_tools = await session.list_tools() + print("List of MCP tools for MCP server:", mcp_tools.tools) - print("MCP TOOLS from litellm proxy: ", MCP_TOOLS) + # Create message messages = [ - { - "role": "user", - "content": "send an email about litellm supporting MCP and send it to krrish@berri.ai", - } + ChatCompletionUserMessageParam( + content="Send an email about LiteLLM supporting MCP", role="user" + ) ] - llm_response = await openai_client.chat.completions.create( + + # Request with tools + response = await client.chat.completions.create( model="gpt-4o", messages=messages, tools=[ - transform_mcp_tool_to_openai_tool(tool) for tool in MCP_TOOLS.tools + transform_mcp_tool_to_openai_tool(tool) for tool in mcp_tools.tools ], - tool_choice="required", + tool_choice="auto", ) - print("LLM RESPONSE: ", json.dumps(llm_response, indent=4, default=str)) - # Add assertions to verify the response - openai_tool = llm_response.choices[0].message.tool_calls[0] + # Handle tool call + if response.choices[0].message.tool_calls: + tool_call = response.choices[0].message.tool_calls[0] + if tool_call: + # Convert format + mcp_call = ( + transform_openai_tool_call_request_to_mcp_tool_call_request( + openai_tool=tool_call.model_dump() + ) + ) - # Call the tool using MCP client - mcp_tool_call_request = ( - transform_openai_tool_call_request_to_mcp_tool_call_request( - openai_tool.model_dump() - ) - ) - call_result = await session.call_tool( - name=mcp_tool_call_request.name, - arguments=mcp_tool_call_request.arguments, - ) - print("CALL RESULT: ", call_result) - pass + # Execute tool + result = await session.call_tool( + name=mcp_call.name, arguments=mcp_call.arguments + ) + + print("Result:", result) + + +# Run it +asyncio.run(main())