add testing mcp server

This commit is contained in:
Ishaan Jaff
2025-03-29 12:52:46 -07:00
parent 6aa660b3f5
commit a1ec0dd0e2
3 changed files with 105 additions and 12 deletions
+10 -4
View File
@@ -1,7 +1,13 @@
model_list:
- model_name: fake-openai-endpoint
- model_name: gpt-4o
litellm_params:
model: openai/fake
api_key: fake-key
api_base: https://exampleopenaiendpoint-production.up.railway.app/
model: openai/gpt-4o
mcp_servers:
{
"Zapier MCP": {
"url": "os.environ/ZAPIER_MCP_SERVER_URL",
},
}
}
+37
View File
@@ -0,0 +1,37 @@
# Create server parameters for stdio connection
import os
import sys
import pytest
sys.path.insert(
0, os.path.abspath("../../..")
) # Adds the parent directory to the system path
from litellm.proxy._experimental.mcp_server.mcp_client_manager import (
MCPServerManager,
MCPSSEServer,
)
MCP_SERVERS = [
MCPSSEServer(name="zapier_mcp_server", url=os.environ.get("ZAPIER_MCP_SERVER_URL")),
]
mcp_server_manager = MCPServerManager(mcp_servers=MCP_SERVERS)
@pytest.mark.asyncio
async def test_mcp_server_manager():
tools = await mcp_server_manager.list_tools()
print("TOOLS FROM MCP SERVER MANAGER== ", tools)
result = await mcp_server_manager.call_tool(
name="gmail_send_email", arguments={"body": "Test"}
)
print("RESULT FROM CALLING TOOL FROM MCP SERVER MANAGER== ", result)
"""
TODO test with multiple MCP servers and calling a specific
"""
+58 -8
View File
@@ -2,15 +2,24 @@
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 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_to_mcp_tool_call_request,
)
import json
async def main():
model = ChatOpenAI(model="gpt-4o", api_key="sk-12")
@pytest.mark.asyncio
async def test_mcp_routes():
model = ChatOpenAI(
model="gpt-4o", 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:
@@ -25,11 +34,52 @@ async def main():
print("Tools loaded")
print(tools)
# # Create and run the agent
# agent = create_react_agent(model, tools)
# agent_response = await agent.ainvoke({"messages": "what's (3 + 5) x 12?"})
# Create and run the agent
agent = create_react_agent(model, tools)
agent_response = await agent.ainvoke({"messages": "Send an "})
print(agent_response)
# Run the async function
if __name__ == "__main__":
asyncio.run(main())
@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):
async with ClientSession(read, write) as session:
await session.initialize()
MCP_TOOLS = await session.list_tools()
print("MCP TOOLS from litellm proxy: ", MCP_TOOLS)
messages = [
{
"role": "user",
"content": "send an email about litellm supporting MCP and send it to krrish@berri.ai",
}
]
llm_response = await openai_client.chat.completions.create(
model="gpt-4o",
messages=messages,
tools=[
transform_mcp_tool_to_openai_tool(tool) for tool in MCP_TOOLS.tools
],
tool_choice="required",
)
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]
# Call the tool using MCP client
mcp_tool_call_request = (
_transform_openai_tool_call_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