Added LiteLLM to the stack

Development/litellm/tests/test_litellm/experimental_mcp_client/test_mcp_client.py

@@ -0,0 +1,79 @@
+import sys
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+# Add the parent directory to the path so we can import litellm
+sys.path.insert(0, '../../../')
+
+from litellm.experimental_mcp_client.client import MCPClient
+from litellm.types.mcp import MCPStdioConfig, MCPTransport
+
+
+class TestMCPClient:
+    """Test MCP Client stdio functionality"""
+
+    def test_mcp_client_stdio_init(self):
+        """Test MCPClient initialization with stdio config"""
+        stdio_config = MCPStdioConfig(
+            command="python",
+            args=["-m", "my_mcp_server"],
+            env={"DEBUG": "1"}
+        )
+        
+        client = MCPClient(
+            transport_type=MCPTransport.stdio,
+            stdio_config=stdio_config
+        )
+        
+        assert client.transport_type == MCPTransport.stdio
+        assert client.stdio_config == stdio_config
+        assert client.stdio_config["command"] == "python"
+        assert client.stdio_config["args"] == ["-m", "my_mcp_server"]
+
+    @pytest.mark.asyncio
+    async def test_mcp_client_stdio_connect_error(self):
+        """Test MCP client stdio connection error handling"""
+        # Test missing stdio_config
+        client = MCPClient(transport_type=MCPTransport.stdio)
+        
+        with pytest.raises(ValueError, match="stdio_config is required for stdio transport"):
+            await client.connect()
+
+    @pytest.mark.asyncio
+    @patch('litellm.experimental_mcp_client.client.stdio_client')
+    @patch('litellm.experimental_mcp_client.client.ClientSession')
+    async def test_mcp_client_stdio_connect_success(self, mock_session, mock_stdio_client):
+        """Test successful stdio connection"""
+        # Setup mocks
+        mock_transport = (MagicMock(), MagicMock())
+        mock_stdio_client.return_value.__aenter__ = AsyncMock(return_value=mock_transport)
+        
+        mock_session_instance = MagicMock()
+        mock_session_instance.__aenter__ = AsyncMock(return_value=mock_session_instance)
+        mock_session_instance.initialize = AsyncMock()
+        mock_session.return_value = mock_session_instance
+        
+        stdio_config = MCPStdioConfig(
+            command="python",
+            args=["-m", "my_mcp_server"],
+            env={"DEBUG": "1"}
+        )
+        
+        client = MCPClient(
+            transport_type=MCPTransport.stdio,
+            stdio_config=stdio_config
+        )
+        
+        await client.connect()
+        
+        # Verify stdio_client was called with correct parameters
+        mock_stdio_client.assert_called_once()
+        call_args = mock_stdio_client.call_args[0][0]
+        assert call_args.command == "python"
+        assert call_args.args == ["-m", "my_mcp_server"]
+        assert call_args.env == {"DEBUG": "1"}
+
+
+if __name__ == "__main__":
+    pytest.main([__file__]) 

Development/litellm/tests/test_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,
+    call_mcp_tool,
+    call_openai_tool,
+    load_mcp_tools,
+    transform_mcp_tool_to_openai_tool,
+    transform_openai_tool_call_request_to_mcp_tool_call_request,
+)
+
+
+@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 testtransform_openai_tool_call_request_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_request_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"}
+    )