diff --git a/.circleci/config.yml b/.circleci/config.yml index a1a7e69c40..aa18f0c5ad 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -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 diff --git a/docs/my-website/docs/mcp.md b/docs/my-website/docs/mcp.md index fbdeaa191f..6ebbd7a1fa 100644 --- a/docs/my-website/docs/mcp.md +++ b/docs/my-website/docs/mcp.md @@ -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 + + + +

+ LiteLLM MCP Architecture: Use MCP tools with all LiteLLM supported models +

-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` + -```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)) ``` - -```python -import asyncio + + +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) +``` + + + + +### 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 + + + + +```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) + ) +``` + + + + +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) ``` diff --git a/docs/my-website/img/litellm_mcp.png b/docs/my-website/img/litellm_mcp.png new file mode 100644 index 0000000000..cef822eeb2 Binary files /dev/null and b/docs/my-website/img/litellm_mcp.png differ diff --git a/litellm/__init__.py b/litellm/__init__.py index 762a058c7e..25da650440 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -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 diff --git a/litellm/experimental_mcp_client/Readme.md b/litellm/experimental_mcp_client/Readme.md new file mode 100644 index 0000000000..4fbd624369 --- /dev/null +++ b/litellm/experimental_mcp_client/Readme.md @@ -0,0 +1,6 @@ +# LiteLLM MCP Client + +LiteLLM MCP Client is a client that allows you to use MCP tools with LiteLLM. + + + diff --git a/litellm/experimental_mcp_client/__init__.py b/litellm/experimental_mcp_client/__init__.py new file mode 100644 index 0000000000..7110d5375e --- /dev/null +++ b/litellm/experimental_mcp_client/__init__.py @@ -0,0 +1,3 @@ +from .tools import call_openai_tool, load_mcp_tools + +__all__ = ["load_mcp_tools", "call_openai_tool"] diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/litellm/experimental_mcp_client/tools.py b/litellm/experimental_mcp_client/tools.py new file mode 100644 index 0000000000..f4ebbf4af4 --- /dev/null +++ b/litellm/experimental_mcp_client/tools.py @@ -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, + ) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 1d7b8794b5..c5e0329013 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -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, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 1d7b8794b5..c5e0329013 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -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, diff --git a/pyproject.toml b/pyproject.toml index 38d5687800..000adb98f6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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" ] diff --git a/tests/litellm/experimental_mcp_client/test_tools.py b/tests/litellm/experimental_mcp_client/test_tools.py new file mode 100644 index 0000000000..7089d83217 --- /dev/null +++ b/tests/litellm/experimental_mcp_client/test_tools.py @@ -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"} + ) diff --git a/tests/mcp_tests/mcp_server.py b/tests/mcp_tests/mcp_server.py new file mode 100644 index 0000000000..99a67edd02 --- /dev/null +++ b/tests/mcp_tests/mcp_server.py @@ -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") diff --git a/tests/mcp_tests/test_mcp_litellm_client.py b/tests/mcp_tests/test_mcp_litellm_client.py new file mode 100644 index 0000000000..0f8fb7994a --- /dev/null +++ b/tests/mcp_tests/test_mcp_litellm_client.py @@ -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) + )