Merge branch 'main' into litellm_contributor_fix_mar_21

This commit is contained in:
Ishaan Jaff
2025-03-21 21:07:20 -07:00
committed by GitHub
14 changed files with 674 additions and 81 deletions
+56 -2
View File
@@ -680,6 +680,50 @@ jobs:
paths:
- llm_translation_coverage.xml
- llm_translation_coverage
mcp_testing:
docker:
- image: cimg/python:3.11
auth:
username: ${DOCKERHUB_USERNAME}
password: ${DOCKERHUB_PASSWORD}
working_directory: ~/project
steps:
- checkout
- run:
name: Install Dependencies
command: |
python -m pip install --upgrade pip
python -m pip install -r requirements.txt
pip install "pytest==7.3.1"
pip install "pytest-retry==1.6.3"
pip install "pytest-cov==5.0.0"
pip install "pytest-asyncio==0.21.1"
pip install "respx==0.21.1"
pip install "pydantic==2.7.2"
pip install "mcp==1.4.1"
# Run pytest and generate JUnit XML report
- run:
name: Run tests
command: |
pwd
ls
python -m pytest -vv tests/mcp_tests --cov=litellm --cov-report=xml -x -s -v --junitxml=test-results/junit.xml --durations=5
no_output_timeout: 120m
- run:
name: Rename the coverage files
command: |
mv coverage.xml mcp_coverage.xml
mv .coverage mcp_coverage
# Store test results
- store_test_results:
path: test-results
- persist_to_workspace:
root: .
paths:
- mcp_coverage.xml
- mcp_coverage
llm_responses_api_testing:
docker:
- image: cimg/python:3.11
@@ -744,6 +788,8 @@ jobs:
pip install "pytest-asyncio==0.21.1"
pip install "respx==0.21.1"
pip install "hypercorn==0.17.3"
pip install "pydantic==2.7.2"
pip install "mcp==1.4.1"
# Run pytest and generate JUnit XML report
- run:
name: Run tests
@@ -1353,7 +1399,7 @@ jobs:
command: |
pwd
ls
python -m pytest -s -vv tests/*.py -x --junitxml=test-results/junit.xml --durations=5 --ignore=tests/otel_tests --ignore=tests/pass_through_tests --ignore=tests/proxy_admin_ui_tests --ignore=tests/load_tests --ignore=tests/llm_translation --ignore=tests/llm_responses_api_testing --ignore=tests/image_gen_tests --ignore=tests/pass_through_unit_tests
python -m pytest -s -vv tests/*.py -x --junitxml=test-results/junit.xml --durations=5 --ignore=tests/otel_tests --ignore=tests/pass_through_tests --ignore=tests/proxy_admin_ui_tests --ignore=tests/load_tests --ignore=tests/llm_translation --ignore=tests/llm_responses_api_testing --ignore=tests/mcp_tests --ignore=tests/image_gen_tests --ignore=tests/pass_through_unit_tests
no_output_timeout: 120m
# Store test results
@@ -2112,7 +2158,7 @@ jobs:
python -m venv venv
. venv/bin/activate
pip install coverage
coverage combine llm_translation_coverage llm_responses_api_coverage logging_coverage litellm_router_coverage local_testing_coverage litellm_assistants_api_coverage auth_ui_unit_tests_coverage langfuse_coverage caching_coverage litellm_proxy_unit_tests_coverage image_gen_coverage pass_through_unit_tests_coverage batches_coverage litellm_proxy_security_tests_coverage
coverage combine llm_translation_coverage llm_responses_api_coverage mcp_coverage logging_coverage litellm_router_coverage local_testing_coverage litellm_assistants_api_coverage auth_ui_unit_tests_coverage langfuse_coverage caching_coverage litellm_proxy_unit_tests_coverage image_gen_coverage pass_through_unit_tests_coverage batches_coverage litellm_proxy_security_tests_coverage
coverage xml
- codecov/upload:
file: ./coverage.xml
@@ -2473,6 +2519,12 @@ workflows:
only:
- main
- /litellm_.*/
- mcp_testing:
filters:
branches:
only:
- main
- /litellm_.*/
- llm_responses_api_testing:
filters:
branches:
@@ -2518,6 +2570,7 @@ workflows:
- upload-coverage:
requires:
- llm_translation_testing
- mcp_testing
- llm_responses_api_testing
- litellm_mapped_tests
- batches_testing
@@ -2577,6 +2630,7 @@ workflows:
- load_testing
- test_bad_database_url
- llm_translation_testing
- mcp_testing
- llm_responses_api_testing
- litellm_mapped_tests
- batches_testing
+232 -74
View File
@@ -1,114 +1,272 @@
import Tabs from '@theme/Tabs';
import TabItem from '@theme/TabItem';
import Image from '@theme/IdealImage';
# /mcp Model Context Protocol [Beta]
# /mcp [BETA] - Model Context Protocol
Use Model Context Protocol with LiteLLM
<Image
img={require('../img/litellm_mcp.png')}
style={{width: '100%', display: 'block', margin: '2rem auto'}}
/>
<p style={{textAlign: 'left', color: '#666'}}>
LiteLLM MCP Architecture: Use MCP tools with all LiteLLM supported models
</p>
Use Model Context Protocol with LiteLLM.
## Overview
LiteLLM acts as a MCP bridge to utilize MCP tools with all LiteLLM supported models. LiteLLM offers the following features for using MCP
- **List** Available MCP Tools: OpenAI clients can view all available MCP tools
- `litellm.experimental_mcp_client.load_mcp_tools` to list all available MCP tools
- **Call** MCP Tools: OpenAI clients can call MCP tools
- `litellm.experimental_mcp_client.call_openai_tool` to call an OpenAI tool on an MCP server
LiteLLM acts as a MCP bridge to utilize **MCP tools** with **all LiteLLM supported models**. LiteLLM offers a client that exposes a tools method for retrieving tools from a MCP server.
## Usage
### 1. List Available MCP Tools
In this example we'll use `litellm.experimental_mcp_client.load_mcp_tools` to list all available MCP tools on any MCP server. This method can be used in two ways:
- `format="mcp"` - (default) Return MCP tools
- Returns: `mcp.types.Tool`
- `format="openai"` - Return MCP tools converted to OpenAI API compatible tools. Allows using with OpenAI endpoints.
- Returns: `openai.types.chat.ChatCompletionToolParam`
<Tabs>
<TabItem value="sdk" label="LiteLLM Python SDK">
```python
import asyncio
```python title="MCP Client List Tools" showLineNumbers
# Create server parameters for stdio connection
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
import os
import litellm
from litellm import experimental_create_mcp_client
from litellm.mcp_stdio import experimental_stdio_mcp_transport
from litellm import experimental_mcp_client
async def main():
client_one = None
try:
# Initialize an MCP client to connect to a `stdio` MCP server:
transport = experimental_stdio_mcp_transport(
command='node',
args=['src/stdio/dist/server.js']
)
client_one = await experimental_create_mcp_client(
transport=transport
)
server_params = StdioServerParameters(
command="python3",
# Make sure to update to the full absolute path to your mcp_server.py file
args=["./mcp_server.py"],
)
tools = await client_one.list_tools(format="openai")
response = await litellm.completion(
async with stdio_client(server_params) as (read, write):
async with ClientSession(read, write) as session:
# Initialize the connection
await session.initialize()
# Get tools
tools = await experimental_mcp_client.load_mcp_tools(session=session, format="openai")
print("MCP TOOLS: ", tools)
messages = [{"role": "user", "content": "what's (3 + 5)"}]
llm_response = await litellm.acompletion(
model="gpt-4o",
api_key=os.getenv("OPENAI_API_KEY"),
messages=messages,
tools=tools,
messages=[
{
"role": "user",
"content": "Find products under $100"
}
]
)
print(response.text)
except Exception as error:
print(error)
finally:
await asyncio.gather(
client_one.close() if client_one else asyncio.sleep(0),
)
if __name__ == "__main__":
asyncio.run(main())
print("LLM RESPONSE: ", json.dumps(llm_response, indent=4, default=str))
```
</TabItem>
<TabItem value="proxy" label="LiteLLM Proxy Server">
```python
import asyncio
<TabItem value="openai" label="OpenAI SDK + LiteLLM Proxy">
In this example we'll walk through how you can use the OpenAI SDK pointed to the LiteLLM proxy to call MCP tools. The key difference here is we use the OpenAI SDK to make the LLM API request
```python title="MCP Client List Tools" showLineNumbers
# Create server parameters for stdio connection
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
import os
from openai import OpenAI
from litellm import experimental_create_mcp_client
from litellm.mcp_stdio import experimental_stdio_mcp_transport
from litellm import experimental_mcp_client
async def main():
client_one = None
server_params = StdioServerParameters(
command="python3",
# Make sure to update to the full absolute path to your mcp_server.py file
args=["./mcp_server.py"],
)
try:
# Initialize an MCP client to connect to a `stdio` MCP server:
transport = experimental_stdio_mcp_transport(
command='node',
args=['src/stdio/dist/server.js']
)
client_one = await experimental_create_mcp_client(
transport=transport
)
async with stdio_client(server_params) as (read, write):
async with ClientSession(read, write) as session:
# Initialize the connection
await session.initialize()
# Get tools from MCP client
tools = await client_one.list_tools(format="openai")
# Use OpenAI client connected to LiteLLM Proxy Server
# Get tools using litellm mcp client
tools = await experimental_mcp_client.load_mcp_tools(session=session, format="openai")
print("MCP TOOLS: ", tools)
# Use OpenAI SDK pointed to LiteLLM proxy
client = OpenAI(
api_key="sk-1234",
base_url="http://0.0.0.0:4000"
api_key="your-api-key", # Your LiteLLM proxy API key
base_url="http://localhost:4000" # Your LiteLLM proxy URL
)
response = client.chat.completions.create(
messages = [{"role": "user", "content": "what's (3 + 5)"}]
llm_response = client.chat.completions.create(
model="gpt-4",
messages=messages,
tools=tools
)
print("LLM RESPONSE: ", llm_response)
```
</TabItem>
</Tabs>
### 2. List and Call MCP Tools
In this example we'll use
- `litellm.experimental_mcp_client.load_mcp_tools` to list all available MCP tools on any MCP server
- `litellm.experimental_mcp_client.call_openai_tool` to call an OpenAI tool on an MCP server
The first llm response returns a list of OpenAI tools. We take the first tool call from the LLM response and pass it to `litellm.experimental_mcp_client.call_openai_tool` to call the tool on the MCP server.
#### How `litellm.experimental_mcp_client.call_openai_tool` works
- Accepts an OpenAI Tool Call from the LLM response
- Converts the OpenAI Tool Call to an MCP Tool
- Calls the MCP Tool on the MCP server
- Returns the result of the MCP Tool call
<Tabs>
<TabItem value="sdk" label="LiteLLM Python SDK">
```python title="MCP Client List and Call Tools" showLineNumbers
# Create server parameters for stdio connection
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
import os
import litellm
from litellm import experimental_mcp_client
server_params = StdioServerParameters(
command="python3",
# Make sure to update to the full absolute path to your mcp_server.py file
args=["./mcp_server.py"],
)
async with stdio_client(server_params) as (read, write):
async with ClientSession(read, write) as session:
# Initialize the connection
await session.initialize()
# Get tools
tools = await experimental_mcp_client.load_mcp_tools(session=session, format="openai")
print("MCP TOOLS: ", tools)
messages = [{"role": "user", "content": "what's (3 + 5)"}]
llm_response = await litellm.acompletion(
model="gpt-4o",
api_key=os.getenv("OPENAI_API_KEY"),
messages=messages,
tools=tools,
messages=[
{
"role": "user",
"content": "Find products under $100"
}
]
)
print("LLM RESPONSE: ", json.dumps(llm_response, indent=4, default=str))
openai_tool = llm_response["choices"][0]["message"]["tool_calls"][0]
# Call the tool using MCP client
call_result = await experimental_mcp_client.call_openai_tool(
session=session,
openai_tool=openai_tool,
)
print("MCP TOOL CALL RESULT: ", call_result)
# send the tool result to the LLM
messages.append(llm_response["choices"][0]["message"])
messages.append(
{
"role": "tool",
"content": str(call_result.content[0].text),
"tool_call_id": openai_tool["id"],
}
)
print("final messages with tool result: ", messages)
llm_response = await litellm.acompletion(
model="gpt-4o",
api_key=os.getenv("OPENAI_API_KEY"),
messages=messages,
tools=tools,
)
print(
"FINAL LLM RESPONSE: ", json.dumps(llm_response, indent=4, default=str)
)
```
</TabItem>
<TabItem value="proxy" label="OpenAI SDK + LiteLLM Proxy">
In this example we'll walk through how you can use the OpenAI SDK pointed to the LiteLLM proxy to call MCP tools. The key difference here is we use the OpenAI SDK to make the LLM API request
```python title="MCP Client with OpenAI SDK" showLineNumbers
# Create server parameters for stdio connection
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
import os
from openai import OpenAI
from litellm import experimental_mcp_client
server_params = StdioServerParameters(
command="python3",
# Make sure to update to the full absolute path to your mcp_server.py file
args=["./mcp_server.py"],
)
async with stdio_client(server_params) as (read, write):
async with ClientSession(read, write) as session:
# Initialize the connection
await session.initialize()
# Get tools using litellm mcp client
tools = await experimental_mcp_client.load_mcp_tools(session=session, format="openai")
print("MCP TOOLS: ", tools)
# Use OpenAI SDK pointed to LiteLLM proxy
client = OpenAI(
api_key="your-api-key", # Your LiteLLM proxy API key
base_url="http://localhost:8000" # Your LiteLLM proxy URL
)
print(response.choices[0].message.content)
except Exception as error:
print(error)
finally:
await asyncio.gather(
client_one.close() if client_one else asyncio.sleep(0),
messages = [{"role": "user", "content": "what's (3 + 5)"}]
llm_response = client.chat.completions.create(
model="gpt-4",
messages=messages,
tools=tools
)
print("LLM RESPONSE: ", llm_response)
if __name__ == "__main__":
asyncio.run(main())
# Get the first tool call
tool_call = llm_response.choices[0].message.tool_calls[0]
# Call the tool using MCP client
call_result = await experimental_mcp_client.call_openai_tool(
session=session,
openai_tool=tool_call.model_dump(),
)
print("MCP TOOL CALL RESULT: ", call_result)
# Send the tool result back to the LLM
messages.append(llm_response.choices[0].message.model_dump())
messages.append({
"role": "tool",
"content": str(call_result.content[0].text),
"tool_call_id": tool_call.id,
})
final_response = client.chat.completions.create(
model="gpt-4",
messages=messages,
tools=tools
)
print("FINAL RESPONSE: ", final_response)
```
</TabItem>
Binary file not shown.

After

Width:  |  Height:  |  Size: 112 KiB

+1 -1
View File
@@ -2,7 +2,7 @@
import warnings
warnings.filterwarnings("ignore", message=".*conflict with protected namespace.*")
### INIT VARIABLES #########
### INIT VARIABLES ##########
import threading
import os
from typing import Callable, List, Optional, Dict, Union, Any, Literal, get_args
@@ -0,0 +1,6 @@
# LiteLLM MCP Client
LiteLLM MCP Client is a client that allows you to use MCP tools with LiteLLM.
@@ -0,0 +1,3 @@
from .tools import call_openai_tool, load_mcp_tools
__all__ = ["load_mcp_tools", "call_openai_tool"]
+109
View File
@@ -0,0 +1,109 @@
import json
from typing import List, Literal, Union
from mcp import ClientSession
from mcp.types import CallToolRequestParams as MCPCallToolRequestParams
from mcp.types import CallToolResult as MCPCallToolResult
from mcp.types import Tool as MCPTool
from openai.types.chat import ChatCompletionToolParam
from openai.types.shared_params.function_definition import FunctionDefinition
from litellm.types.utils import ChatCompletionMessageToolCall
########################################################
# List MCP Tool functions
########################################################
def transform_mcp_tool_to_openai_tool(mcp_tool: MCPTool) -> ChatCompletionToolParam:
"""Convert an MCP tool to an OpenAI tool."""
return ChatCompletionToolParam(
type="function",
function=FunctionDefinition(
name=mcp_tool.name,
description=mcp_tool.description or "",
parameters=mcp_tool.inputSchema,
strict=False,
),
)
async def load_mcp_tools(
session: ClientSession, format: Literal["mcp", "openai"] = "mcp"
) -> Union[List[MCPTool], List[ChatCompletionToolParam]]:
"""
Load all available MCP tools
Args:
session: The MCP session to use
format: The format to convert the tools to
By default, the tools are returned in MCP format.
If format is set to "openai", the tools are converted to OpenAI API compatible tools.
"""
tools = await session.list_tools()
if format == "openai":
return [
transform_mcp_tool_to_openai_tool(mcp_tool=tool) for tool in tools.tools
]
return tools.tools
########################################################
# Call MCP Tool functions
########################################################
async def call_mcp_tool(
session: ClientSession,
call_tool_request_params: MCPCallToolRequestParams,
) -> MCPCallToolResult:
"""Call an MCP tool."""
tool_result = await session.call_tool(
name=call_tool_request_params.name,
arguments=call_tool_request_params.arguments,
)
return tool_result
def _get_function_arguments(function: FunctionDefinition) -> dict:
"""Helper to safely get and parse function arguments."""
arguments = function.get("arguments", {})
if isinstance(arguments, str):
try:
arguments = json.loads(arguments)
except json.JSONDecodeError:
arguments = {}
return arguments if isinstance(arguments, dict) else {}
def _transform_openai_tool_call_to_mcp_tool_call_request(
openai_tool: ChatCompletionMessageToolCall,
) -> MCPCallToolRequestParams:
"""Convert an OpenAI ChatCompletionMessageToolCall to an MCP CallToolRequestParams."""
function = openai_tool["function"]
return MCPCallToolRequestParams(
name=function["name"],
arguments=_get_function_arguments(function),
)
async def call_openai_tool(
session: ClientSession,
openai_tool: ChatCompletionMessageToolCall,
) -> MCPCallToolResult:
"""
Call an OpenAI tool using MCP client.
Args:
session: The MCP session to use
openai_tool: The OpenAI tool to call. You can get this from the `choices[0].message.tool_calls[0]` of the response from the OpenAI API.
Returns:
The result of the MCP tool call.
"""
mcp_tool_call_request_params = _transform_openai_tool_call_to_mcp_tool_call_request(
openai_tool=openai_tool,
)
return await call_mcp_tool(
session=session,
call_tool_request_params=mcp_tool_call_request_params,
)
@@ -1435,7 +1435,7 @@
"input_cost_per_token_batches": 0.0000375,
"output_cost_per_token_batches": 0.000075,
"cache_read_input_token_cost": 0.0000375,
"litellm_provider": "openai",
"litellm_provider": "azure",
"mode": "chat",
"supports_function_calling": true,
"supports_parallel_function_calling": true,
+1 -1
View File
@@ -1435,7 +1435,7 @@
"input_cost_per_token_batches": 0.0000375,
"output_cost_per_token_batches": 0.000075,
"cache_read_input_token_cost": 0.0000375,
"litellm_provider": "openai",
"litellm_provider": "azure",
"mode": "chat",
"supports_function_calling": true,
"supports_parallel_function_calling": true,
+2 -2
View File
@@ -1,6 +1,6 @@
[tool.poetry]
name = "litellm"
version = "1.63.12"
version = "1.63.14"
description = "Library to easily interface with LLM API providers"
authors = ["BerriAI"]
license = "MIT"
@@ -100,7 +100,7 @@ requires = ["poetry-core", "wheel"]
build-backend = "poetry.core.masonry.api"
[tool.commitizen]
version = "1.63.12"
version = "1.63.14"
version_files = [
"pyproject.toml:^version"
]
@@ -0,0 +1,157 @@
import json
import os
import sys
from unittest.mock import AsyncMock, MagicMock
import pytest
sys.path.insert(
0, os.path.abspath("../../..")
) # Adds the parent directory to the system path
from mcp.types import (
CallToolRequestParams,
CallToolResult,
ListToolsResult,
TextContent,
)
from mcp.types import Tool as MCPTool
from litellm.experimental_mcp_client.tools import (
_get_function_arguments,
_transform_openai_tool_call_to_mcp_tool_call_request,
call_mcp_tool,
call_openai_tool,
load_mcp_tools,
transform_mcp_tool_to_openai_tool,
)
@pytest.fixture
def mock_mcp_tool():
return MCPTool(
name="test_tool",
description="A test tool",
inputSchema={"type": "object", "properties": {"test": {"type": "string"}}},
)
@pytest.fixture
def mock_session():
session = MagicMock()
session.list_tools = AsyncMock()
session.call_tool = AsyncMock()
return session
@pytest.fixture
def mock_list_tools_result():
return ListToolsResult(
tools=[
MCPTool(
name="test_tool",
description="A test tool",
inputSchema={
"type": "object",
"properties": {"test": {"type": "string"}},
},
)
]
)
@pytest.fixture
def mock_mcp_tool_call_result():
return CallToolResult(content=[TextContent(type="text", text="test_output")])
def test_transform_mcp_tool_to_openai_tool(mock_mcp_tool):
openai_tool = transform_mcp_tool_to_openai_tool(mock_mcp_tool)
assert openai_tool["type"] == "function"
assert openai_tool["function"]["name"] == "test_tool"
assert openai_tool["function"]["description"] == "A test tool"
assert openai_tool["function"]["parameters"] == {
"type": "object",
"properties": {"test": {"type": "string"}},
}
def test_transform_openai_tool_call_to_mcp_tool_call_request(mock_mcp_tool):
openai_tool = {
"function": {"name": "test_tool", "arguments": json.dumps({"test": "value"})}
}
mcp_tool_call_request = _transform_openai_tool_call_to_mcp_tool_call_request(
openai_tool
)
assert mcp_tool_call_request.name == "test_tool"
assert mcp_tool_call_request.arguments == {"test": "value"}
@pytest.mark.asyncio()
async def test_load_mcp_tools_mcp_format(mock_session, mock_list_tools_result):
mock_session.list_tools.return_value = mock_list_tools_result
result = await load_mcp_tools(mock_session, format="mcp")
assert len(result) == 1
assert isinstance(result[0], MCPTool)
assert result[0].name == "test_tool"
mock_session.list_tools.assert_called_once()
@pytest.mark.asyncio()
async def test_load_mcp_tools_openai_format(mock_session, mock_list_tools_result):
mock_session.list_tools.return_value = mock_list_tools_result
result = await load_mcp_tools(mock_session, format="openai")
assert len(result) == 1
assert result[0]["type"] == "function"
assert result[0]["function"]["name"] == "test_tool"
mock_session.list_tools.assert_called_once()
def test_get_function_arguments():
# Test with string arguments
function = {"arguments": '{"test": "value"}'}
result = _get_function_arguments(function)
assert result == {"test": "value"}
# Test with dict arguments
function = {"arguments": {"test": "value"}}
result = _get_function_arguments(function)
assert result == {"test": "value"}
# Test with invalid JSON string
function = {"arguments": "invalid json"}
result = _get_function_arguments(function)
assert result == {}
# Test with no arguments
function = {}
result = _get_function_arguments(function)
assert result == {}
@pytest.mark.asyncio()
async def test_call_openai_tool(mock_session, mock_mcp_tool_call_result):
mock_session.call_tool.return_value = mock_mcp_tool_call_result
openai_tool = {
"function": {"name": "test_tool", "arguments": json.dumps({"test": "value"})}
}
result = await call_openai_tool(mock_session, openai_tool)
print("result of call_openai_tool", result)
assert result.content[0].text == "test_output"
mock_session.call_tool.assert_called_once_with(
name="test_tool", arguments={"test": "value"}
)
@pytest.mark.asyncio()
async def test_call_mcp_tool(mock_session, mock_mcp_tool_call_result):
mock_session.call_tool.return_value = mock_mcp_tool_call_result
request_params = CallToolRequestParams(
name="test_tool", arguments={"test": "value"}
)
result = await call_mcp_tool(mock_session, request_params)
print("call_mcp_tool result", result)
assert result.content[0].text == "test_output"
mock_session.call_tool.assert_called_once_with(
name="test_tool", arguments={"test": "value"}
)
+20
View File
@@ -0,0 +1,20 @@
# math_server.py
from mcp.server.fastmcp import FastMCP
mcp = FastMCP("Math")
@mcp.tool()
def add(a: int, b: int) -> int:
"""Add two numbers"""
return a + b
@mcp.tool()
def multiply(a: int, b: int) -> int:
"""Multiply two numbers"""
return a * b
if __name__ == "__main__":
mcp.run(transport="stdio")
@@ -0,0 +1,86 @@
# 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 mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
import os
from litellm import experimental_mcp_client
import litellm
import pytest
import json
@pytest.mark.asyncio
async def test_mcp_agent():
local_server_path = "./mcp_server.py"
ci_cd_server_path = "tests/mcp_tests/mcp_server.py"
server_params = StdioServerParameters(
command="python3",
# Make sure to update to the full absolute path to your math_server.py file
args=[ci_cd_server_path],
)
async with stdio_client(server_params) as (read, write):
async with ClientSession(read, write) as session:
# Initialize the connection
await session.initialize()
# Get tools
tools = await experimental_mcp_client.load_mcp_tools(
session=session, format="openai"
)
print("MCP TOOLS: ", tools)
# Create and run the agent
messages = [{"role": "user", "content": "what's (3 + 5)"}]
llm_response = await litellm.acompletion(
model="gpt-4o",
api_key=os.getenv("OPENAI_API_KEY"),
messages=messages,
tools=tools,
tool_choice="required",
)
print("LLM RESPONSE: ", json.dumps(llm_response, indent=4, default=str))
# Add assertions to verify the response
assert llm_response["choices"][0]["message"]["tool_calls"] is not None
assert (
llm_response["choices"][0]["message"]["tool_calls"][0]["function"][
"name"
]
== "add"
)
openai_tool = llm_response["choices"][0]["message"]["tool_calls"][0]
# Call the tool using MCP client
call_result = await experimental_mcp_client.call_openai_tool(
session=session,
openai_tool=openai_tool,
)
print("CALL RESULT: ", call_result)
# send the tool result to the LLM
messages.append(llm_response["choices"][0]["message"])
messages.append(
{
"role": "tool",
"content": str(call_result.content[0].text),
"tool_call_id": openai_tool["id"],
}
)
print("final messages: ", messages)
llm_response = await litellm.acompletion(
model="gpt-4o",
api_key=os.getenv("OPENAI_API_KEY"),
messages=messages,
tools=tools,
)
print(
"FINAL LLM RESPONSE: ", json.dumps(llm_response, indent=4, default=str)
)