From a1ec0dd0e2963bec022a2e0784013eea6b5abd34 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Sat, 29 Mar 2025 12:52:46 -0700
Subject: [PATCH] add testing mcp server

---
 litellm/proxy/proxy_config.yaml             | 14 +++--
 tests/mcp_tests/test_mcp_server.py          | 37 ++++++++++++
 tests/pass_through_tests/test_mcp_routes.py | 66 ++++++++++++++++++---
 3 files changed, 105 insertions(+), 12 deletions(-)
 create mode 100644 tests/mcp_tests/test_mcp_server.py

diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml
index 106003f996..9da28ff51f 100644
--- a/litellm/proxy/proxy_config.yaml
+++ b/litellm/proxy/proxy_config.yaml
@@ -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",
+  },
+  }
+
+}
\ No newline at end of file
diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py
new file mode 100644
index 0000000000..17abd1d6a5
--- /dev/null
+++ b/tests/mcp_tests/test_mcp_server.py
@@ -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 
+
+"""
diff --git a/tests/pass_through_tests/test_mcp_routes.py b/tests/pass_through_tests/test_mcp_routes.py
index 687efe6195..d496a430af 100644
--- a/tests/pass_through_tests/test_mcp_routes.py
+++ b/tests/pass_through_tests/test_mcp_routes.py
@@ -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