mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-13 06:23:32 +00:00
add testing mcp server
This commit is contained in:
@@ -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",
|
||||
},
|
||||
}
|
||||
|
||||
}
|
||||
@@ -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
|
||||
|
||||
"""
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user