1701 lines
60 KiB
Python
1701 lines
60 KiB
Python
import logging
|
|
import os
|
|
import sys
|
|
import pytest
|
|
from typing import List, Any, cast
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
sys.path.insert(0, os.path.abspath("../../.."))
|
|
|
|
# Import required modules
|
|
import litellm
|
|
from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Handler
|
|
from litellm.types.llms.openai import (
|
|
ResponsesAPIResponse,
|
|
ResponsesAPIStreamingResponse,
|
|
OpenAIMcpServerTool,
|
|
ToolParam,
|
|
)
|
|
|
|
|
|
class MockUserAPIKeyAuth:
|
|
"""Mock UserAPIKeyAuth for testing"""
|
|
|
|
def __init__(self):
|
|
self.api_key = "test_key"
|
|
self.user_id = "test_user"
|
|
self.team_id = "test_team"
|
|
self.user_email = "test@example.com"
|
|
self.max_budget = 100.0
|
|
self.spend = 0.0
|
|
self.models = []
|
|
self.aliases = {}
|
|
self.config = {}
|
|
self.permissions = {}
|
|
self.metadata = {}
|
|
self.object_permission_id = "test_permission_id"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mcp_helper_methods():
|
|
"""Test the core MCP helper methods in LiteLLM_Proxy_MCP_Handler"""
|
|
|
|
# Test _should_use_litellm_mcp_gateway
|
|
mcp_tools: List[Any] = [
|
|
{"type": "mcp", "server_url": "litellm_proxy", "require_approval": "never"}
|
|
]
|
|
|
|
other_tools: List[Any] = [
|
|
{
|
|
"type": "function",
|
|
"name": "get_weather",
|
|
"description": "Get weather info",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {"location": {"type": "string"}},
|
|
},
|
|
}
|
|
]
|
|
|
|
# Should return True for MCP tools with litellm_proxy
|
|
assert LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(mcp_tools) == True
|
|
|
|
# Should return False for other tools
|
|
assert (
|
|
LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(other_tools) == False
|
|
)
|
|
|
|
# Should return False for None
|
|
assert LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(None) == False
|
|
|
|
# Test _parse_mcp_tools
|
|
mixed_tools = mcp_tools + other_tools
|
|
mcp_parsed, other_parsed = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(mixed_tools)
|
|
|
|
assert len(mcp_parsed) == 1
|
|
assert len(other_parsed) == 1
|
|
assert mcp_parsed[0]["type"] == "mcp"
|
|
assert other_parsed[0]["type"] == "function"
|
|
|
|
# Test _should_auto_execute_tools
|
|
mcp_tools_never = [{"require_approval": "never"}]
|
|
mcp_tools_always = [{"require_approval": "always"}]
|
|
|
|
assert LiteLLM_Proxy_MCP_Handler._should_auto_execute_tools(mcp_tools_never) == True
|
|
assert (
|
|
LiteLLM_Proxy_MCP_Handler._should_auto_execute_tools(mcp_tools_always) == False
|
|
)
|
|
|
|
print("✓ MCP helper methods test passed!")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mcp_output_elements_addition():
|
|
"""Test adding MCP output elements to response"""
|
|
|
|
# Create a mock response
|
|
mock_response = ResponsesAPIResponse(
|
|
**{ # type: ignore
|
|
"id": "test_response_id",
|
|
"object": "response",
|
|
"created_at": 1234567890,
|
|
"status": "completed",
|
|
"error": None,
|
|
"incomplete_details": None,
|
|
"instructions": None,
|
|
"max_output_tokens": None,
|
|
"model": "gpt-4o",
|
|
"output": [
|
|
{
|
|
"type": "message",
|
|
"id": "msg_123",
|
|
"status": "completed",
|
|
"role": "assistant",
|
|
"content": [
|
|
{
|
|
"type": "output_text",
|
|
"text": "Hello, world!",
|
|
"annotations": [],
|
|
}
|
|
],
|
|
}
|
|
],
|
|
"parallel_tool_calls": True,
|
|
"previous_response_id": None,
|
|
"reasoning": {"effort": None, "summary": None},
|
|
"store": True,
|
|
"temperature": 1.0,
|
|
"text": {"format": {"type": "text"}},
|
|
"tool_choice": "auto",
|
|
"tools": [],
|
|
"top_p": 1.0,
|
|
"truncation": "disabled",
|
|
"usage": {
|
|
"input_tokens": 10,
|
|
"input_tokens_details": {"cached_tokens": 0},
|
|
"output_tokens": 5,
|
|
"output_tokens_details": {"reasoning_tokens": 0},
|
|
"total_tokens": 15,
|
|
},
|
|
"user": None,
|
|
"metadata": {},
|
|
}
|
|
)
|
|
|
|
# Mock MCP tools and tool results
|
|
mock_mcp_tools = [
|
|
{
|
|
"name": "test_tool",
|
|
"description": "A test tool",
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"properties": {"query": {"type": "string"}},
|
|
},
|
|
}
|
|
]
|
|
|
|
mock_tool_results = [
|
|
{"tool_call_id": "call_123", "result": "Tool executed successfully"}
|
|
]
|
|
|
|
# Test adding output elements
|
|
updated_response = LiteLLM_Proxy_MCP_Handler._add_mcp_output_elements_to_response(
|
|
response=mock_response,
|
|
mcp_tools_fetched=mock_mcp_tools,
|
|
tool_results=mock_tool_results,
|
|
)
|
|
|
|
# Verify output elements were added
|
|
assert len(updated_response.output) == 3 # Original + 2 new elements
|
|
|
|
# Check that MCP tools output was added - handle both dict and object cases
|
|
mcp_tools_output = updated_response.output[1]
|
|
if hasattr(mcp_tools_output, "type"):
|
|
# Handle as object with attributes
|
|
output_obj = cast(Any, mcp_tools_output)
|
|
assert output_obj.type == "mcp_tools_fetched"
|
|
assert output_obj.role == "system"
|
|
assert output_obj.status == "completed"
|
|
elif isinstance(mcp_tools_output, dict):
|
|
# Handle as dictionary
|
|
assert mcp_tools_output["type"] == "mcp_tools_fetched"
|
|
assert mcp_tools_output["role"] == "system"
|
|
assert mcp_tools_output["status"] == "completed"
|
|
|
|
# Check that tool results output was added
|
|
tool_results_output = updated_response.output[2]
|
|
if hasattr(tool_results_output, "type"):
|
|
# Handle as object with attributes
|
|
output_obj = cast(Any, tool_results_output)
|
|
assert output_obj.type == "tool_execution_results"
|
|
assert output_obj.role == "system"
|
|
assert output_obj.status == "completed"
|
|
elif isinstance(tool_results_output, dict):
|
|
# Handle as dictionary
|
|
assert tool_results_output["type"] == "tool_execution_results"
|
|
assert tool_results_output["role"] == "system"
|
|
assert tool_results_output["status"] == "completed"
|
|
|
|
print("✓ MCP output elements addition test passed!")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aresponses_api_with_mcp_mock_integration():
|
|
"""
|
|
Test the core MCP integration logic without complex external mocking.
|
|
This focuses on verifying the MCP tool parsing and handling works correctly.
|
|
"""
|
|
# Define MCP tools with litellm_proxy server_url and require_approval="never"
|
|
mcp_tools: List[OpenAIMcpServerTool] = [
|
|
{
|
|
"type": "mcp",
|
|
"server_url": "litellm_proxy",
|
|
"require_approval": "never",
|
|
"server_label": "test_server",
|
|
}
|
|
]
|
|
|
|
# Test the helper methods that the integration relies on
|
|
from litellm.responses.mcp.litellm_proxy_mcp_handler import (
|
|
LiteLLM_Proxy_MCP_Handler,
|
|
)
|
|
|
|
# Test 1: Verify MCP tools are detected correctly
|
|
should_use_mcp = LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(
|
|
cast(Any, mcp_tools)
|
|
)
|
|
assert (
|
|
should_use_mcp == True
|
|
), "Should detect MCP tools with litellm_proxy server_url"
|
|
|
|
# Test 2: Verify auto-execution detection works
|
|
should_auto_execute = LiteLLM_Proxy_MCP_Handler._should_auto_execute_tools(
|
|
cast(Any, mcp_tools)
|
|
)
|
|
assert (
|
|
should_auto_execute == True
|
|
), "Should auto-execute tools with require_approval='never'"
|
|
|
|
# Test 3: Verify tool parsing works correctly
|
|
mcp_parsed, other_parsed = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(
|
|
cast(Any, mcp_tools)
|
|
)
|
|
assert len(mcp_parsed) == 1, "Should parse one MCP tool"
|
|
assert len(other_parsed) == 0, "Should have no other tools"
|
|
assert mcp_parsed[0]["type"] == "mcp", "Parsed tool should be MCP type"
|
|
assert mcp_parsed[0]["server_url"] == "litellm_proxy", "Should preserve server_url"
|
|
assert (
|
|
mcp_parsed[0].get("require_approval") == "never"
|
|
), "Should preserve require_approval"
|
|
|
|
# Test 4: Test with mixed tools
|
|
mixed_tools = mcp_tools + [
|
|
{"type": "function", "name": "test_function", "parameters": {"type": "object"}}
|
|
]
|
|
|
|
mcp_parsed, other_parsed = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(
|
|
cast(Any, mixed_tools)
|
|
)
|
|
assert len(mcp_parsed) == 1, "Should parse one MCP tool from mixed list"
|
|
assert len(other_parsed) == 1, "Should have one other tool from mixed list"
|
|
|
|
print("✓ MCP integration core logic test completed successfully!")
|
|
print(f"MCP tools detected: {should_use_mcp}")
|
|
print(f"Auto-execute enabled: {should_auto_execute}")
|
|
print(f"MCP tools parsed: {len(mcp_parsed)}")
|
|
print(f"Other tools parsed: {len(other_parsed)}")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aresponses_api_with_mcp_passes_mcp_server_auth_headers_to_process_tools():
|
|
"""
|
|
Test that MCP auth headers from secret_fields (e.g. x-mcp-linear_config-authorization)
|
|
are passed to _process_mcp_tools_without_openai_transform when using the responses API.
|
|
"""
|
|
from litellm.responses.main import aresponses_api_with_mcp
|
|
|
|
captured_process_kwargs = {}
|
|
|
|
async def mock_process(**kwargs):
|
|
captured_process_kwargs.update(kwargs)
|
|
return ([], {})
|
|
|
|
mock_response = ResponsesAPIResponse(
|
|
**{
|
|
"id": "resp_test",
|
|
"object": "response",
|
|
"created_at": 1234567890,
|
|
"status": "completed",
|
|
"error": None,
|
|
"incomplete_details": None,
|
|
"instructions": None,
|
|
"max_output_tokens": None,
|
|
"model": "gpt-4o",
|
|
"output": [
|
|
{
|
|
"type": "message",
|
|
"id": "msg_1",
|
|
"status": "completed",
|
|
"role": "assistant",
|
|
"content": [],
|
|
}
|
|
],
|
|
"parallel_tool_calls": True,
|
|
"previous_response_id": None,
|
|
"reasoning": {"effort": None, "summary": None},
|
|
"store": True,
|
|
"temperature": 1.0,
|
|
"text": {"format": {"type": "text"}},
|
|
"tool_choice": "auto",
|
|
"tools": [],
|
|
"top_p": 1.0,
|
|
"truncation": "disabled",
|
|
"usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
|
|
"user": None,
|
|
"metadata": {},
|
|
}
|
|
)
|
|
|
|
mcp_tools = [{"type": "mcp", "server_url": "litellm_proxy"}]
|
|
secret_fields = {
|
|
"raw_headers": {"x-mcp-linear_config-authorization": "Bearer linear-token"},
|
|
}
|
|
|
|
with (
|
|
patch.object(
|
|
LiteLLM_Proxy_MCP_Handler,
|
|
"_process_mcp_tools_without_openai_transform",
|
|
mock_process,
|
|
),
|
|
patch(
|
|
"litellm.responses.main.aresponses",
|
|
new_callable=AsyncMock,
|
|
return_value=mock_response,
|
|
),
|
|
):
|
|
await aresponses_api_with_mcp(
|
|
input=[{"role": "user", "type": "message", "content": "hi"}],
|
|
model="gpt-4o",
|
|
tools=mcp_tools,
|
|
secret_fields=secret_fields,
|
|
)
|
|
|
|
assert "mcp_server_auth_headers" in captured_process_kwargs
|
|
mcp_server_auth_headers = captured_process_kwargs["mcp_server_auth_headers"]
|
|
assert mcp_server_auth_headers is not None
|
|
assert "linear_config" in mcp_server_auth_headers
|
|
assert (
|
|
mcp_server_auth_headers["linear_config"]["Authorization"]
|
|
== "Bearer linear-token"
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mcp_allowed_tools_filtering():
|
|
"""
|
|
Test the allowed_tools filtering functionality for MCP tools.
|
|
This test verifies that when allowed_tools is specified in MCP tool config,
|
|
only the allowed tools are passed to the LLM.
|
|
"""
|
|
from litellm.responses.mcp.litellm_proxy_mcp_handler import (
|
|
LiteLLM_Proxy_MCP_Handler,
|
|
)
|
|
|
|
# Mock MCP tools returned from the server (simulating all available tools)
|
|
mock_mcp_tools_from_server = [
|
|
# Mock MCP tool object with name attribute
|
|
type(
|
|
"MCPTool",
|
|
(),
|
|
{
|
|
"name": "search_tiktoken_documentation",
|
|
"description": "Search tiktoken documentation",
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"properties": {"query": {"type": "string"}},
|
|
},
|
|
},
|
|
)(),
|
|
type(
|
|
"MCPTool",
|
|
(),
|
|
{
|
|
"name": "fetch_tiktoken_documentation",
|
|
"description": "Fetch tiktoken documentation",
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"properties": {"path": {"type": "string"}},
|
|
},
|
|
},
|
|
)(),
|
|
type(
|
|
"MCPTool",
|
|
(),
|
|
{
|
|
"name": "list_tiktoken_functions",
|
|
"description": "List tiktoken functions",
|
|
"inputSchema": {"type": "object", "properties": {}},
|
|
},
|
|
)(),
|
|
type(
|
|
"MCPTool",
|
|
(),
|
|
{
|
|
"name": "get_tiktoken_examples",
|
|
"description": "Get tiktoken examples",
|
|
"inputSchema": {"type": "object", "properties": {}},
|
|
},
|
|
)(),
|
|
]
|
|
|
|
allowed_mcp_servers = ["gitmcp"]
|
|
|
|
# Test Case 1: MCP tool config with allowed_tools specified
|
|
mcp_tool_config_with_allowed_tools = [
|
|
{
|
|
"type": "mcp",
|
|
"server_label": "gitmcp",
|
|
"server_url": "https://gitmcp.io/openai/tiktoken",
|
|
"allowed_tools": [
|
|
"search_tiktoken_documentation",
|
|
"fetch_tiktoken_documentation",
|
|
],
|
|
"require_approval": "never",
|
|
}
|
|
]
|
|
|
|
# Filter tools using the helper function
|
|
filtered_tools = LiteLLM_Proxy_MCP_Handler._filter_mcp_tools_by_allowed_tools(
|
|
mcp_tools=mock_mcp_tools_from_server,
|
|
mcp_tools_with_litellm_proxy=cast(
|
|
List[ToolParam], mcp_tool_config_with_allowed_tools
|
|
),
|
|
)
|
|
|
|
# Should only return the 2 allowed tools
|
|
assert (
|
|
len(filtered_tools) == 2
|
|
), f"Expected 2 filtered tools, got {len(filtered_tools)}"
|
|
|
|
# Check that only allowed tools are included
|
|
filtered_tool_names = [tool.name for tool in filtered_tools]
|
|
expected_allowed_tools = [
|
|
"search_tiktoken_documentation",
|
|
"fetch_tiktoken_documentation",
|
|
]
|
|
|
|
assert set(filtered_tool_names) == set(
|
|
expected_allowed_tools
|
|
), f"Expected tools {expected_allowed_tools}, got {filtered_tool_names}"
|
|
|
|
# Verify excluded tools are not present
|
|
excluded_tools = ["list_tiktoken_functions", "get_tiktoken_examples"]
|
|
for excluded_tool in excluded_tools:
|
|
assert (
|
|
excluded_tool not in filtered_tool_names
|
|
), f"Tool {excluded_tool} should have been filtered out"
|
|
|
|
print("✓ Test Case 1: allowed_tools filtering works correctly")
|
|
|
|
# Test Case 2: MCP tool config without allowed_tools (should return all tools)
|
|
mcp_tool_config_without_allowed_tools = [
|
|
{
|
|
"type": "mcp",
|
|
"server_label": "gitmcp",
|
|
"server_url": "https://gitmcp.io/openai/tiktoken",
|
|
"require_approval": "never",
|
|
}
|
|
]
|
|
|
|
filtered_tools_all = LiteLLM_Proxy_MCP_Handler._filter_mcp_tools_by_allowed_tools(
|
|
mcp_tools=mock_mcp_tools_from_server,
|
|
mcp_tools_with_litellm_proxy=cast(
|
|
List[ToolParam], mcp_tool_config_without_allowed_tools
|
|
),
|
|
)
|
|
|
|
# Should return all 4 tools when no allowed_tools specified
|
|
assert (
|
|
len(filtered_tools_all) == 4
|
|
), f"Expected 4 tools when no allowed_tools specified, got {len(filtered_tools_all)}"
|
|
|
|
print("✓ Test Case 2: no allowed_tools returns all tools")
|
|
|
|
# Test Case 3: Test deduplication of duplicate tools
|
|
mock_mcp_tools_with_duplicates = [
|
|
# First instance of duplicate tool
|
|
type(
|
|
"MCPTool",
|
|
(),
|
|
{
|
|
"name": "GitMCP-fetch_litellm_documentation",
|
|
"description": "Fetch entire documentation file from GitHub repository: BerriAI/litellm. Useful for general questions. Always call this tool first if asked about BerriAI/litellm.",
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"properties": {},
|
|
"additionalProperties": False,
|
|
},
|
|
},
|
|
)(),
|
|
# Second instance of duplicate tool (should be filtered out)
|
|
type(
|
|
"MCPTool",
|
|
(),
|
|
{
|
|
"name": "GitMCP-fetch_litellm_documentation",
|
|
"description": "Fetch entire documentation file from GitHub repository: BerriAI/litellm. Useful for general questions. Always call this tool first if asked about BerriAI/litellm.",
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"properties": {},
|
|
"additionalProperties": False,
|
|
},
|
|
},
|
|
)(),
|
|
# Other unique tools
|
|
type(
|
|
"MCPTool",
|
|
(),
|
|
{
|
|
"name": "GitMCP-search_litellm_documentation",
|
|
"description": "Semantically search within the fetched documentation from GitHub repository: BerriAI/litellm. Useful for specific queries.",
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"properties": {"query": {"type": "string"}},
|
|
"required": ["query"],
|
|
"additionalProperties": False,
|
|
},
|
|
},
|
|
)(),
|
|
]
|
|
|
|
mcp_tool_config_with_duplicates = [
|
|
{
|
|
"type": "mcp",
|
|
"server_label": "litellm",
|
|
"server_url": "litellm_proxy/mcp",
|
|
"require_approval": "never",
|
|
"allowed_tools": ["GitMCP-fetch_litellm_documentation"],
|
|
}
|
|
]
|
|
|
|
# First filter by allowed tools
|
|
filtered_tools_with_duplicates = (
|
|
LiteLLM_Proxy_MCP_Handler._filter_mcp_tools_by_allowed_tools(
|
|
mcp_tools=mock_mcp_tools_with_duplicates,
|
|
mcp_tools_with_litellm_proxy=cast(
|
|
List[ToolParam], mcp_tool_config_with_duplicates
|
|
),
|
|
)
|
|
)
|
|
|
|
# Then deduplicate the filtered tools
|
|
filtered_tools_deduplicated, _ = LiteLLM_Proxy_MCP_Handler._deduplicate_mcp_tools(
|
|
filtered_tools_with_duplicates, []
|
|
)
|
|
|
|
# Should only return 1 tool (the duplicate should be removed)
|
|
assert (
|
|
len(filtered_tools_deduplicated) == 1
|
|
), f"Expected 1 tool after deduplication, got {len(filtered_tools_deduplicated)}"
|
|
|
|
# Check that the correct tool is present
|
|
assert (
|
|
filtered_tools_deduplicated[0].name == "GitMCP-fetch_litellm_documentation"
|
|
), f"Expected GitMCP-fetch_litellm_documentation, got {filtered_tools_deduplicated[0].name}"
|
|
|
|
print("✓ Test Case 3: duplicate tools are properly deduplicated")
|
|
|
|
# Test Case 3b: Test standalone deduplication method
|
|
standalone_deduplicated, _ = LiteLLM_Proxy_MCP_Handler._deduplicate_mcp_tools(
|
|
mock_mcp_tools_with_duplicates, allowed_mcp_servers
|
|
)
|
|
|
|
# Should return 2 unique tools (GitMCP-fetch_litellm_documentation and GitMCP-search_litellm_documentation)
|
|
assert (
|
|
len(standalone_deduplicated) == 2
|
|
), f"Expected 2 unique tools after standalone deduplication, got {len(standalone_deduplicated)}"
|
|
|
|
unique_tool_names = [tool.name for tool in standalone_deduplicated]
|
|
expected_unique_names = [
|
|
"GitMCP-fetch_litellm_documentation",
|
|
"GitMCP-search_litellm_documentation",
|
|
]
|
|
assert set(unique_tool_names) == set(
|
|
expected_unique_names
|
|
), f"Expected {expected_unique_names}, got {unique_tool_names}"
|
|
|
|
print("✓ Test Case 3b: standalone deduplication method works correctly")
|
|
|
|
# Test Case 4: Multiple MCP tool configs with different allowed_tools
|
|
multiple_mcp_configs = [
|
|
{
|
|
"type": "mcp",
|
|
"server_label": "gitmcp1",
|
|
"server_url": "https://gitmcp.io/openai/tiktoken",
|
|
"allowed_tools": ["search_tiktoken_documentation"],
|
|
"require_approval": "never",
|
|
},
|
|
{
|
|
"type": "mcp",
|
|
"server_label": "gitmcp2",
|
|
"server_url": "https://gitmcp.io/openai/tiktoken",
|
|
"allowed_tools": ["fetch_tiktoken_documentation", "get_tiktoken_examples"],
|
|
"require_approval": "never",
|
|
},
|
|
]
|
|
|
|
filtered_tools_multiple = (
|
|
LiteLLM_Proxy_MCP_Handler._filter_mcp_tools_by_allowed_tools(
|
|
mcp_tools=mock_mcp_tools_from_server,
|
|
mcp_tools_with_litellm_proxy=cast(List[ToolParam], multiple_mcp_configs),
|
|
)
|
|
)
|
|
|
|
# Should return union of all allowed tools (3 unique tools)
|
|
assert (
|
|
len(filtered_tools_multiple) == 3
|
|
), f"Expected 3 tools from multiple configs, got {len(filtered_tools_multiple)}"
|
|
|
|
filtered_multiple_names = [tool.name for tool in filtered_tools_multiple]
|
|
expected_multiple_tools = [
|
|
"search_tiktoken_documentation",
|
|
"fetch_tiktoken_documentation",
|
|
"get_tiktoken_examples",
|
|
]
|
|
|
|
assert set(filtered_multiple_names) == set(
|
|
expected_multiple_tools
|
|
), f"Expected tools {expected_multiple_tools}, got {filtered_multiple_names}"
|
|
|
|
print(
|
|
"✓ Test Case 3: multiple MCP configs with different allowed_tools works correctly"
|
|
)
|
|
|
|
# Test Case 4: Empty allowed_tools list (should return no tools)
|
|
mcp_config_empty_allowed = [
|
|
{
|
|
"type": "mcp",
|
|
"server_label": "gitmcp",
|
|
"server_url": "https://gitmcp.io/openai/tiktoken",
|
|
"allowed_tools": [],
|
|
"require_approval": "never",
|
|
}
|
|
]
|
|
|
|
filtered_tools_empty = LiteLLM_Proxy_MCP_Handler._filter_mcp_tools_by_allowed_tools(
|
|
mcp_tools=mock_mcp_tools_from_server,
|
|
mcp_tools_with_litellm_proxy=cast(List[ToolParam], mcp_config_empty_allowed),
|
|
)
|
|
|
|
# Should return all tools when allowed_tools is empty list (no filtering)
|
|
assert (
|
|
len(filtered_tools_empty) == 4
|
|
), f"Expected 4 tools when allowed_tools is empty list, got {len(filtered_tools_empty)}"
|
|
|
|
print("✓ Test Case 4: empty allowed_tools list returns all tools")
|
|
|
|
print("✓ MCP allowed_tools filtering test completed successfully!")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_streaming_mcp_events_validation():
|
|
"""
|
|
Test that MCP streaming events are properly emitted when using streaming with MCP tools.
|
|
|
|
This test validates:
|
|
1. MCP discovery events are emitted first
|
|
2. Regular streaming response events follow
|
|
3. Tool execution events are emitted when tools are auto-executed
|
|
"""
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
from litellm.types.llms.openai import ResponsesAPIStreamEvents
|
|
|
|
# Mock MCP tools that would be returned from the manager
|
|
mock_mcp_tools = [
|
|
type(
|
|
"MCPTool",
|
|
(),
|
|
{
|
|
"name": "search_repo",
|
|
"description": "Search BerriAI/litellm repository for information",
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"properties": {
|
|
"query": {"type": "string", "description": "Search query"}
|
|
},
|
|
"required": ["query"],
|
|
},
|
|
},
|
|
)(),
|
|
type(
|
|
"MCPTool",
|
|
(),
|
|
{
|
|
"name": "get_repo_info",
|
|
"description": "Get repository information",
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"properties": {
|
|
"repo_name": {
|
|
"type": "string",
|
|
"description": "Repository name",
|
|
}
|
|
},
|
|
"required": ["repo_name"],
|
|
},
|
|
},
|
|
)(),
|
|
]
|
|
|
|
# Build fake streaming chunks that the inner aresponses() call would yield
|
|
fake_response_obj = MagicMock(spec=ResponsesAPIResponse)
|
|
fake_response_obj.id = "resp_fake_123"
|
|
fake_response_obj.output = []
|
|
|
|
fake_created_chunk = MagicMock(spec=ResponsesAPIStreamingResponse)
|
|
fake_created_chunk.type = ResponsesAPIStreamEvents.RESPONSE_CREATED
|
|
fake_created_chunk.response = fake_response_obj
|
|
|
|
fake_in_progress_chunk = MagicMock(spec=ResponsesAPIStreamingResponse)
|
|
fake_in_progress_chunk.type = ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS
|
|
fake_in_progress_chunk.response = fake_response_obj
|
|
|
|
fake_output_item_added_chunk = MagicMock(spec=ResponsesAPIStreamingResponse)
|
|
fake_output_item_added_chunk.type = ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED
|
|
fake_output_item_added_chunk.response = fake_response_obj
|
|
|
|
fake_completed_chunk = MagicMock(spec=ResponsesAPIStreamingResponse)
|
|
fake_completed_chunk.type = ResponsesAPIStreamEvents.RESPONSE_COMPLETED
|
|
fake_completed_chunk.response = fake_response_obj
|
|
|
|
# Create a fake async iterator for the inner LLM streaming call
|
|
class FakeAsyncIterator:
|
|
def __init__(self, chunks):
|
|
self._chunks = list(chunks)
|
|
self._index = 0
|
|
|
|
def __aiter__(self):
|
|
return self
|
|
|
|
async def __anext__(self):
|
|
if self._index >= len(self._chunks):
|
|
raise StopAsyncIteration
|
|
chunk = self._chunks[self._index]
|
|
self._index += 1
|
|
return chunk
|
|
|
|
fake_stream = FakeAsyncIterator(
|
|
[
|
|
fake_created_chunk,
|
|
fake_in_progress_chunk,
|
|
fake_output_item_added_chunk,
|
|
fake_completed_chunk,
|
|
]
|
|
)
|
|
|
|
# Mock the MCP operations and the inner aresponses call
|
|
with (
|
|
patch.object(
|
|
LiteLLM_Proxy_MCP_Handler,
|
|
"_get_mcp_tools_from_manager",
|
|
new_callable=AsyncMock,
|
|
) as mock_get_tools,
|
|
patch.object(
|
|
LiteLLM_Proxy_MCP_Handler,
|
|
"_execute_tool_calls",
|
|
new_callable=AsyncMock,
|
|
) as mock_execute_tools,
|
|
patch(
|
|
"litellm.responses.main.aresponses",
|
|
new_callable=AsyncMock,
|
|
return_value=fake_stream,
|
|
),
|
|
):
|
|
# Setup MCP mocks
|
|
mock_get_tools.return_value = (mock_mcp_tools, ["test_server"])
|
|
|
|
async def mock_execute_tool_calls_side_effect(
|
|
tool_server_map, tool_calls, user_api_key_auth, **kwargs
|
|
):
|
|
"""Mock tool execution with realistic results"""
|
|
results = []
|
|
for tool_call in tool_calls:
|
|
call_id = None
|
|
if isinstance(tool_call, dict):
|
|
call_id = tool_call.get("call_id") or tool_call.get("id")
|
|
elif hasattr(tool_call, "call_id"):
|
|
call_id = tool_call.call_id
|
|
elif hasattr(tool_call, "id"):
|
|
call_id = tool_call.id
|
|
|
|
if call_id:
|
|
results.append(
|
|
{
|
|
"tool_call_id": call_id,
|
|
"result": "LiteLLM is a unified interface for 100+ LLMs.",
|
|
}
|
|
)
|
|
return results
|
|
|
|
mock_execute_tools.side_effect = mock_execute_tool_calls_side_effect
|
|
|
|
# Configure MCP tool with streaming and auto-execution
|
|
mcp_tool_config = {
|
|
"type": "mcp",
|
|
"server_url": "litellm_proxy/mcp/test_server",
|
|
"require_approval": "never", # This enables auto-execution
|
|
}
|
|
|
|
# Make streaming request with MCP tools
|
|
response = await litellm.aresponses(
|
|
model="gpt-4o-mini",
|
|
tools=[mcp_tool_config],
|
|
tool_choice="required",
|
|
input=[
|
|
{
|
|
"role": "user",
|
|
"type": "message",
|
|
"content": "What is LiteLLM? Give me a brief overview.",
|
|
}
|
|
],
|
|
stream=True,
|
|
)
|
|
|
|
assert hasattr(
|
|
response, "__aiter__"
|
|
), "Response should be async iterable for streaming"
|
|
|
|
# Collect all streaming events
|
|
events = []
|
|
event_types = []
|
|
mcp_discovery_events = []
|
|
regular_events = []
|
|
|
|
async for chunk in response:
|
|
events.append(chunk)
|
|
event_type = getattr(chunk, "type", "unknown")
|
|
event_types.append(event_type)
|
|
|
|
# Categorize events
|
|
if event_type in [
|
|
ResponsesAPIStreamEvents.MCP_LIST_TOOLS_IN_PROGRESS,
|
|
ResponsesAPIStreamEvents.MCP_LIST_TOOLS_COMPLETED,
|
|
]:
|
|
mcp_discovery_events.append(chunk)
|
|
else:
|
|
regular_events.append(chunk)
|
|
|
|
# Validate that we got streaming events
|
|
assert len(events) > 0, "Should have received at least some streaming events"
|
|
|
|
# Validate MCP discovery events were emitted
|
|
assert (
|
|
len(mcp_discovery_events) > 0
|
|
), "Should have received MCP discovery events"
|
|
|
|
# Check that discovery events come before regular content events
|
|
first_discovery_idx = next(
|
|
i
|
|
for i, e in enumerate(events)
|
|
if getattr(e, "type", None)
|
|
in [
|
|
ResponsesAPIStreamEvents.MCP_LIST_TOOLS_IN_PROGRESS,
|
|
ResponsesAPIStreamEvents.MCP_LIST_TOOLS_COMPLETED,
|
|
]
|
|
)
|
|
# The output_item.added event triggers the transition to MCP discovery,
|
|
# so discovery events should appear after it in the stream
|
|
assert (
|
|
first_discovery_idx > 0
|
|
), "MCP discovery events should follow the initial output_item.added event"
|
|
|
|
# Verify MCP mocks were called
|
|
assert mock_get_tools.called, "MCP tools should have been fetched"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"model",
|
|
[
|
|
pytest.param("gpt-4o-mini", id="openai"),
|
|
pytest.param("claude-haiku-4-5", id="anthropic"),
|
|
],
|
|
)
|
|
async def test_streaming_responses_api_with_mcp_tools(
|
|
model: str, caplog: pytest.LogCaptureFixture
|
|
):
|
|
"""
|
|
Test the streaming responses API with MCP tools when using server_url="litellm_proxy"
|
|
|
|
Under the hood the follow occurs
|
|
|
|
- MCP: responses called litellm MCP manager.list_tools (MOCKED)
|
|
- Request 1: Made to model under test with fetched tools (REAL LLM CALL)
|
|
- MCP: Execute tool call from request 1 and returns result (MOCKED)
|
|
- Request 2: Made to model under test with fetched tools and tool results (REAL LLM CALL)
|
|
|
|
Return the user the result of request 2
|
|
"""
|
|
# Skip test if API keys are not set for the respective models
|
|
if ("claude" in model.lower() or "anthropic" in model.lower()) and not os.getenv(
|
|
"ANTHROPIC_API_KEY"
|
|
):
|
|
pytest.skip("ANTHROPIC_API_KEY not set, skipping anthropic model test")
|
|
if ("gpt" in model.lower() or "openai" in model.lower()) and not os.getenv(
|
|
"OPENAI_API_KEY"
|
|
):
|
|
pytest.skip("OPENAI_API_KEY not set, skipping openai model test")
|
|
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
print("🧪 Testing basic streaming with MCP tools...")
|
|
|
|
# Mock MCP tools that would be returned from the manager
|
|
mock_mcp_tools = [
|
|
type(
|
|
"MCPTool",
|
|
(),
|
|
{
|
|
"name": "search_repo",
|
|
"description": "Search BerriAI/litellm repository for information",
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"properties": {
|
|
"query": {"type": "string", "description": "Search query"}
|
|
},
|
|
"required": ["query"],
|
|
},
|
|
},
|
|
)()
|
|
]
|
|
|
|
# Only mock the MCP-specific operations, let LLM responses be real
|
|
with caplog.at_level(logging.ERROR):
|
|
with (
|
|
patch.object(
|
|
LiteLLM_Proxy_MCP_Handler,
|
|
"_get_mcp_tools_from_manager",
|
|
new_callable=AsyncMock,
|
|
) as mock_get_tools,
|
|
patch.object(
|
|
LiteLLM_Proxy_MCP_Handler,
|
|
"_execute_tool_calls",
|
|
new_callable=AsyncMock,
|
|
) as mock_execute_tools,
|
|
):
|
|
# Setup MCP mocks only
|
|
mock_get_tools.return_value = (mock_mcp_tools, ["litellm_proxy"])
|
|
|
|
# Create a dynamic mock that will match the actual tool call ID from the LLM response
|
|
def mock_execute_tool_calls_side_effect(
|
|
tool_calls, user_api_key_auth, **kwargs
|
|
):
|
|
"""Mock function that returns results matching the actual tool call IDs from the LLM"""
|
|
results = []
|
|
for tool_call in tool_calls:
|
|
# Extract call_id from the tool call
|
|
call_id = None
|
|
if isinstance(tool_call, dict):
|
|
call_id = tool_call.get("call_id") or tool_call.get("id")
|
|
elif hasattr(tool_call, "call_id"):
|
|
call_id = tool_call.call_id
|
|
elif hasattr(tool_call, "id"):
|
|
call_id = tool_call.id
|
|
|
|
if call_id:
|
|
results.append(
|
|
{
|
|
"tool_call_id": call_id,
|
|
"result": "LiteLLM is a unified interface for 100+ LLMs that translates inputs to provider-specific completion endpoints and provides consistent OpenAI-format output.",
|
|
}
|
|
)
|
|
return results
|
|
|
|
mock_execute_tools.side_effect = mock_execute_tool_calls_side_effect
|
|
|
|
# Make the actual call - LLM responses will be real
|
|
mcp_tool_config = cast(
|
|
Any,
|
|
{
|
|
"type": "mcp",
|
|
"server_url": "litellm_proxy",
|
|
"require_approval": "never",
|
|
},
|
|
)
|
|
response = await litellm.aresponses(
|
|
model=model,
|
|
tools=[mcp_tool_config],
|
|
tool_choice="required",
|
|
input=[
|
|
{
|
|
"role": "user",
|
|
"type": "message",
|
|
"content": "give me a TLDR of what BerriAI/litellm is about",
|
|
}
|
|
],
|
|
stream=True,
|
|
)
|
|
|
|
print(f"📋 Response type: {type(response)}")
|
|
assert hasattr(
|
|
response, "__aiter__"
|
|
), "Response should be an async streaming response"
|
|
|
|
# Collect streaming chunks
|
|
chunks = []
|
|
async for chunk in response:
|
|
chunks.append(chunk)
|
|
print(f"📦 Chunk type: {getattr(chunk, 'type', 'unknown')}")
|
|
|
|
print(f"📊 Total chunks received: {len(chunks)}")
|
|
|
|
# Verify MCP mocks were called (may be called multiple times in streaming)
|
|
assert (
|
|
mock_get_tools.call_count >= 1
|
|
), f"Expected MCP tools to be fetched at least once, got {mock_get_tools.call_count}"
|
|
print(f"MCP tools fetched: {len(mock_mcp_tools)}")
|
|
|
|
# Verify we got a response
|
|
assert response is not None
|
|
assert len(chunks) > 0, "Should have received streaming chunks"
|
|
|
|
print("Basic streaming responses API with MCP tools test passed!")
|
|
|
|
lite_errors = [
|
|
record
|
|
for record in caplog.records
|
|
if record.levelno >= logging.ERROR
|
|
and ("LiteLLM" in record.name or "LiteLLM" in record.getMessage())
|
|
]
|
|
assert not lite_errors, "Unexpected LiteLLM errors: " + ", ".join(
|
|
record.getMessage() for record in lite_errors
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mcp_parameter_preparation_helpers():
|
|
"""
|
|
Test the new parameter preparation helper methods for clean MCP handling.
|
|
|
|
Tests:
|
|
1. _prepare_initial_call_params - handles stream disabling for auto-execute
|
|
2. _prepare_follow_up_call_params - restores stream and removes tool_choice
|
|
3. _build_request_params - clean parameter merging
|
|
"""
|
|
from litellm.responses.mcp.litellm_proxy_mcp_handler import (
|
|
LiteLLM_Proxy_MCP_Handler,
|
|
)
|
|
|
|
print("🧪 Testing MCP parameter preparation helpers...")
|
|
|
|
# Test _prepare_initial_call_params
|
|
base_call_params = {
|
|
"stream": True,
|
|
"temperature": 0.7,
|
|
"tool_choice": "required",
|
|
"max_output_tokens": 1000,
|
|
}
|
|
|
|
# Test Case 1: Auto-execute scenario (should disable streaming)
|
|
initial_params_auto = LiteLLM_Proxy_MCP_Handler._prepare_initial_call_params(
|
|
call_params=base_call_params, should_auto_execute=True
|
|
)
|
|
|
|
assert (
|
|
initial_params_auto["stream"] == False
|
|
), "Stream should be disabled for auto-execute"
|
|
assert initial_params_auto["temperature"] == 0.7, "Other params should be preserved"
|
|
assert (
|
|
initial_params_auto["tool_choice"] == "required"
|
|
), "tool_choice should be preserved for initial call"
|
|
assert base_call_params["stream"] == True, "Original params should not be mutated"
|
|
|
|
print("✅ _prepare_initial_call_params (auto-execute) works correctly")
|
|
|
|
# Test Case 2: No auto-execute scenario (should preserve streaming)
|
|
initial_params_no_auto = LiteLLM_Proxy_MCP_Handler._prepare_initial_call_params(
|
|
call_params=base_call_params, should_auto_execute=False
|
|
)
|
|
|
|
assert (
|
|
initial_params_no_auto["stream"] == True
|
|
), "Stream should be preserved when not auto-executing"
|
|
assert (
|
|
initial_params_no_auto["temperature"] == 0.7
|
|
), "Other params should be preserved"
|
|
|
|
print("✅ _prepare_initial_call_params (no auto-execute) works correctly")
|
|
|
|
# Test _prepare_follow_up_call_params
|
|
follow_up_params = LiteLLM_Proxy_MCP_Handler._prepare_follow_up_call_params(
|
|
call_params=base_call_params, original_stream_setting=True
|
|
)
|
|
|
|
assert (
|
|
follow_up_params["stream"] == True
|
|
), "Stream should be restored to original setting"
|
|
assert (
|
|
"tool_choice" not in follow_up_params
|
|
), "tool_choice should be removed for follow-up call"
|
|
assert follow_up_params["temperature"] == 0.7, "Other params should be preserved"
|
|
assert (
|
|
base_call_params["tool_choice"] == "required"
|
|
), "Original params should not be mutated"
|
|
|
|
print("✅ _prepare_follow_up_call_params works correctly")
|
|
|
|
# Test _build_request_params
|
|
input_data = [{"role": "user", "content": "test", "type": "message"}]
|
|
model = "gpt-4o-mini"
|
|
tools = [{"type": "function", "name": "test_tool"}]
|
|
call_params = {"stream": True, "temperature": 0.8}
|
|
previous_response_id = "resp_123"
|
|
extra_kwargs = {"custom_param": "test_value"}
|
|
|
|
request_params = LiteLLM_Proxy_MCP_Handler._build_request_params(
|
|
input=input_data,
|
|
model=model,
|
|
all_tools=tools,
|
|
call_params=call_params,
|
|
previous_response_id=previous_response_id,
|
|
**extra_kwargs,
|
|
)
|
|
|
|
# Verify core parameters
|
|
assert request_params["input"] == input_data, "Input should be included"
|
|
assert request_params["model"] == model, "Model should be included"
|
|
assert request_params["tools"] == tools, "Tools should be included"
|
|
assert (
|
|
request_params["previous_response_id"] == previous_response_id
|
|
), "Previous response ID should be included"
|
|
|
|
# Verify call_params are merged
|
|
assert request_params["stream"] == True, "call_params should be merged"
|
|
assert request_params["temperature"] == 0.8, "call_params should be merged"
|
|
|
|
# Verify extra kwargs are merged
|
|
assert (
|
|
request_params["custom_param"] == "test_value"
|
|
), "Extra kwargs should be merged"
|
|
|
|
print("✅ _build_request_params works correctly")
|
|
|
|
# Test _build_request_params with None previous_response_id
|
|
request_params_no_prev = LiteLLM_Proxy_MCP_Handler._build_request_params(
|
|
input=input_data,
|
|
model=model,
|
|
all_tools=tools,
|
|
call_params=call_params,
|
|
previous_response_id=None,
|
|
)
|
|
|
|
assert (
|
|
"previous_response_id" not in request_params_no_prev
|
|
), "None previous_response_id should not be included"
|
|
|
|
print("✅ _build_request_params handles None previous_response_id correctly")
|
|
|
|
print("🎉 All MCP parameter preparation helper tests passed!")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mcp_tool_execution_events_creation():
|
|
"""
|
|
Test the _create_tool_execution_events helper method for generating streaming events.
|
|
"""
|
|
from litellm.responses.mcp.litellm_proxy_mcp_handler import (
|
|
LiteLLM_Proxy_MCP_Handler,
|
|
)
|
|
|
|
print("Testing MCP tool execution events creation...")
|
|
|
|
# Mock tool calls (simulating what comes from LLM response in function_call format)
|
|
mock_tool_calls = [
|
|
{
|
|
"id": "call_abc123",
|
|
"name": "search_repo",
|
|
"arguments": '{"query": "LiteLLM overview"}',
|
|
"type": "function_call",
|
|
},
|
|
{
|
|
"id": "call_def456",
|
|
"name": "get_repo_info",
|
|
"arguments": '{"repo_name": "BerriAI/litellm"}',
|
|
"type": "function_call",
|
|
},
|
|
]
|
|
|
|
# Mock tool results (simulating what comes from tool execution)
|
|
mock_tool_results = [
|
|
{
|
|
"tool_call_id": "call_abc123",
|
|
"result": "LiteLLM is a unified interface for 100+ LLMs",
|
|
},
|
|
{
|
|
"tool_call_id": "call_def456",
|
|
"result": "Repository: BerriAI/litellm - Python library for LLM integration",
|
|
},
|
|
]
|
|
|
|
# Create tool execution events
|
|
execution_events = LiteLLM_Proxy_MCP_Handler._create_tool_execution_events(
|
|
tool_calls=mock_tool_calls, tool_results=mock_tool_results
|
|
)
|
|
|
|
# Verify events were created
|
|
assert len(execution_events) > 0, "Should create tool execution events"
|
|
print(f"Created {len(execution_events)} tool execution events")
|
|
|
|
# Verify events have proper structure
|
|
for event in execution_events:
|
|
assert hasattr(event, "type"), "Event should have type attribute"
|
|
event_type = str(event.type)
|
|
assert (
|
|
"mcp_call" in event_type.lower() or "output_item" in event_type.lower()
|
|
), f"Event should be MCP-related: {event_type}"
|
|
|
|
# Check for sequence numbers
|
|
if hasattr(event, "sequence_number"):
|
|
assert isinstance(
|
|
event.sequence_number, int
|
|
), "Sequence number should be integer"
|
|
assert event.sequence_number > 0, "Sequence number should be positive"
|
|
|
|
print("Tool execution events have proper structure")
|
|
|
|
# Test with empty inputs
|
|
empty_events = LiteLLM_Proxy_MCP_Handler._create_tool_execution_events(
|
|
tool_calls=[], tool_results=[]
|
|
)
|
|
|
|
assert len(empty_events) == 0, "Should create no events for empty inputs"
|
|
print("Handles empty inputs correctly")
|
|
|
|
print("MCP tool execution events creation test passed!")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_no_duplicate_mcp_tools_in_streaming_e2e():
|
|
"""
|
|
End-to-end test to validate that MCP tools are not duplicated when using streaming.
|
|
|
|
This test protects against the bug where:
|
|
1. Parent function (aresponses_api_with_mcp) processed MCP tools once
|
|
2. Streaming iterator processed MCP tools again, causing duplicates
|
|
|
|
The test mocks the MCP manager response but validates the actual tools
|
|
sent to the LLM to ensure no duplication occurs.
|
|
"""
|
|
from unittest.mock import AsyncMock, patch, call
|
|
from litellm.responses.mcp.litellm_proxy_mcp_handler import (
|
|
LiteLLM_Proxy_MCP_Handler,
|
|
)
|
|
|
|
print("Testing no duplicate MCP tools in streaming E2E...")
|
|
|
|
# Mock MCP tools that would be returned from the manager
|
|
mock_mcp_tools = [
|
|
type(
|
|
"MCPTool",
|
|
(),
|
|
{
|
|
"name": "search_docs",
|
|
"description": "Search documentation for information",
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"properties": {
|
|
"query": {"type": "string", "description": "Search query"}
|
|
},
|
|
"required": ["query"],
|
|
},
|
|
},
|
|
)(),
|
|
type(
|
|
"MCPTool",
|
|
(),
|
|
{
|
|
"name": "get_file_content",
|
|
"description": "Get content of a specific file",
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"properties": {
|
|
"file_path": {"type": "string", "description": "Path to file"}
|
|
},
|
|
"required": ["file_path"],
|
|
},
|
|
},
|
|
)(),
|
|
]
|
|
|
|
# Track all calls to the underlying LLM to detect duplicates
|
|
llm_call_tools = []
|
|
|
|
async def capture_llm_tools(**kwargs):
|
|
"""Capture the tools parameter from LLM calls"""
|
|
tools = kwargs.get("tools", [])
|
|
llm_call_tools.append(tools)
|
|
|
|
# Return a minimal mock async streaming response
|
|
class MockStreamingResponse:
|
|
async def __aiter__(self):
|
|
yield type(
|
|
"MockChunk", (), {"type": "response.completed", "output": []}
|
|
)()
|
|
|
|
return MockStreamingResponse()
|
|
|
|
# Mock both the MCP manager and the underlying LLM call
|
|
with (
|
|
patch.object(
|
|
LiteLLM_Proxy_MCP_Handler,
|
|
"_get_mcp_tools_from_manager",
|
|
new_callable=AsyncMock,
|
|
) as mock_get_tools,
|
|
patch("litellm.aresponses", side_effect=capture_llm_tools) as mock_aresponses,
|
|
):
|
|
|
|
# Setup MCP mock to return our test tools
|
|
mock_get_tools.return_value = mock_mcp_tools
|
|
|
|
# Configure MCP tool for streaming
|
|
mcp_tool_config = {
|
|
"type": "mcp",
|
|
"server_url": "litellm_proxy/mcp/test_server",
|
|
"require_approval": "always", # Disable auto-execution to focus on tool duplication
|
|
}
|
|
|
|
print("Making streaming request with MCP tools...")
|
|
|
|
# Make streaming request with MCP tools
|
|
try:
|
|
response = await litellm.aresponses(
|
|
model="gpt-4o-mini",
|
|
tools=[mcp_tool_config],
|
|
input=[
|
|
{
|
|
"role": "user",
|
|
"type": "message",
|
|
"content": "Search the documentation for information about authentication.",
|
|
}
|
|
],
|
|
stream=True,
|
|
)
|
|
|
|
# Consume the streaming response
|
|
chunks = []
|
|
async for chunk in response:
|
|
chunks.append(chunk)
|
|
|
|
except Exception as e:
|
|
print(f"Request failed (expected for test): {e}")
|
|
# Continue with validation even if request fails
|
|
|
|
# Validate underlying LLM was called (this proves our mocking works)
|
|
assert len(llm_call_tools) > 0, "LLM should have been called at least once"
|
|
print(f"LLM called {len(llm_call_tools)} time(s)")
|
|
|
|
# If MCP tools were processed, validate they were fetched exactly once
|
|
# (This protects against duplicate fetching)
|
|
if mock_get_tools.call_count > 0:
|
|
assert (
|
|
mock_get_tools.call_count == 1
|
|
), f"MCP tools should be fetched exactly once, got {mock_get_tools.call_count} calls"
|
|
print(f"MCP tools fetched exactly once: {mock_get_tools.call_count}")
|
|
else:
|
|
print(
|
|
"MCP tools not fetched (likely due to test mocking - this is OK for validation)"
|
|
)
|
|
|
|
# Analyze tools sent to LLM for duplicates
|
|
for call_idx, tools_in_call in enumerate(llm_call_tools):
|
|
print(f"LLM Call {call_idx + 1}: {len(tools_in_call)} tools")
|
|
|
|
if tools_in_call:
|
|
# Extract tool names to check for duplicates
|
|
tool_names = []
|
|
for tool in tools_in_call:
|
|
if isinstance(tool, dict):
|
|
tool_name = tool.get("function", {}).get("name") or tool.get(
|
|
"name"
|
|
)
|
|
else:
|
|
tool_name = getattr(tool, "name", str(tool))
|
|
|
|
if tool_name:
|
|
tool_names.append(tool_name)
|
|
|
|
print(f" Tool names: {tool_names}")
|
|
|
|
# Check for duplicate tool names
|
|
unique_tool_names = set(tool_names)
|
|
duplicates = [name for name in tool_names if tool_names.count(name) > 1]
|
|
|
|
assert (
|
|
len(duplicates) == 0
|
|
), f"Found duplicate tools in LLM call {call_idx + 1}: {duplicates}"
|
|
assert len(tool_names) == len(
|
|
unique_tool_names
|
|
), f"Tool names should be unique in call {call_idx + 1}"
|
|
|
|
print(f" No duplicate tools found in call {call_idx + 1}")
|
|
|
|
# Validate that MCP tools were properly transformed to OpenAI format
|
|
openai_format_tools = [
|
|
tool
|
|
for tool in tools_in_call
|
|
if isinstance(tool, dict) and "function" in tool
|
|
]
|
|
if openai_format_tools:
|
|
print(f" Found {len(openai_format_tools)} OpenAI-format tools")
|
|
|
|
# Verify tools have proper OpenAI structure
|
|
for tool in openai_format_tools:
|
|
assert "type" in tool, "Tool should have 'type' field"
|
|
assert (
|
|
tool["type"] == "function"
|
|
), "Tool type should be 'function'"
|
|
assert "function" in tool, "Tool should have 'function' field"
|
|
assert "name" in tool["function"], "Function should have 'name'"
|
|
assert (
|
|
"description" in tool["function"]
|
|
), "Function should have 'description'"
|
|
assert (
|
|
"parameters" in tool["function"]
|
|
), "Function should have 'parameters'"
|
|
|
|
print(f" All tools have proper OpenAI format")
|
|
|
|
# The key validation: ensure no duplicate fetching occurred
|
|
# This is the main protection against the bug we fixed
|
|
if mock_get_tools.call_count > 1:
|
|
print(
|
|
f"ERROR: Duplicate MCP fetching detected! Called {mock_get_tools.call_count} times"
|
|
)
|
|
assert (
|
|
False
|
|
), f"MCP tools should be fetched exactly once, but were fetched {mock_get_tools.call_count} times"
|
|
|
|
# Additional validation: ensure no duplicate tools in any LLM call
|
|
total_duplicates_found = 0
|
|
for call_idx, tools_in_call in enumerate(llm_call_tools):
|
|
if tools_in_call:
|
|
tool_names = []
|
|
for tool in tools_in_call:
|
|
if isinstance(tool, dict):
|
|
tool_name = tool.get("function", {}).get("name") or tool.get(
|
|
"name"
|
|
)
|
|
if tool_name:
|
|
tool_names.append(tool_name)
|
|
|
|
duplicates = [name for name in tool_names if tool_names.count(name) > 1]
|
|
if duplicates:
|
|
total_duplicates_found += len(set(duplicates))
|
|
print(
|
|
f"ERROR: Duplicate tools in call {call_idx + 1}: {set(duplicates)}"
|
|
)
|
|
|
|
if total_duplicates_found > 0:
|
|
assert (
|
|
False
|
|
), f"Found {total_duplicates_found} duplicate tools across all LLM calls"
|
|
|
|
print("No duplicate MCP tools E2E test passed!")
|
|
print(f"Summary:")
|
|
print(f" - MCP manager called: {mock_get_tools.call_count} time(s)")
|
|
print(f" - LLM called: {len(llm_call_tools)} time(s)")
|
|
print(
|
|
f" - Unique tools per call: {[len(set(getattr(t.get('function', {}), 'name', 'unknown') if isinstance(t, dict) else str(t) for t in tools)) for tools in llm_call_tools]}"
|
|
)
|
|
print(f" - No duplicate tools detected")
|
|
|
|
return {
|
|
"mcp_manager_calls": mock_get_tools.call_count,
|
|
"llm_calls": len(llm_call_tools),
|
|
"tools_per_call": [len(tools) for tools in llm_call_tools],
|
|
"duplicate_tools_found": False,
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("model", ["gpt-4o-mini"])
|
|
async def test_streaming_mcp_event_order_and_response_id_consistency(
|
|
model: str, caplog: pytest.LogCaptureFixture
|
|
):
|
|
"""
|
|
Test that:
|
|
1. Streaming events are emitted in correct order (response.created, response.in_progress, response.output_item.added before MCP events)
|
|
2. All response lifecycle events share the same response ID within a cycle
|
|
"""
|
|
if ("gpt" in model.lower() or "openai" in model.lower()) and not os.getenv(
|
|
"OPENAI_API_KEY"
|
|
):
|
|
pytest.skip("OPENAI_API_KEY not set, skipping openai model test")
|
|
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
mock_mcp_tools = [
|
|
type(
|
|
"MCPTool",
|
|
(),
|
|
{
|
|
"name": "get_weather",
|
|
"description": "Get weather for a city",
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"properties": {
|
|
"city": {"type": "string", "description": "City name"}
|
|
},
|
|
"required": ["city"],
|
|
},
|
|
},
|
|
)()
|
|
]
|
|
|
|
with caplog.at_level(logging.ERROR):
|
|
with (
|
|
patch.object(
|
|
LiteLLM_Proxy_MCP_Handler,
|
|
"_get_mcp_tools_from_manager",
|
|
new_callable=AsyncMock,
|
|
) as mock_get_tools,
|
|
patch.object(
|
|
LiteLLM_Proxy_MCP_Handler,
|
|
"_execute_tool_calls",
|
|
new_callable=AsyncMock,
|
|
) as mock_execute_tools,
|
|
):
|
|
mock_get_tools.return_value = (mock_mcp_tools, ["litellm_proxy"])
|
|
|
|
def mock_execute_side_effect(tool_calls, user_api_key_auth, **kwargs):
|
|
results = []
|
|
for tool_call in tool_calls:
|
|
call_id = None
|
|
if isinstance(tool_call, dict):
|
|
call_id = tool_call.get("call_id") or tool_call.get("id")
|
|
elif hasattr(tool_call, "call_id"):
|
|
call_id = tool_call.call_id
|
|
elif hasattr(tool_call, "id"):
|
|
call_id = tool_call.id
|
|
if call_id:
|
|
results.append(
|
|
{
|
|
"tool_call_id": call_id,
|
|
"result": "Sunny, 72°F",
|
|
}
|
|
)
|
|
return results
|
|
|
|
mock_execute_tools.side_effect = mock_execute_side_effect
|
|
|
|
mcp_tool_config = cast(
|
|
Any,
|
|
{
|
|
"type": "mcp",
|
|
"server_url": "litellm_proxy",
|
|
"require_approval": "never",
|
|
},
|
|
)
|
|
|
|
response = await litellm.aresponses(
|
|
model=model,
|
|
tools=[mcp_tool_config],
|
|
input=[
|
|
{
|
|
"role": "user",
|
|
"type": "message",
|
|
"content": "What's the weather in San Francisco?",
|
|
}
|
|
],
|
|
stream=True,
|
|
)
|
|
|
|
events = []
|
|
async for chunk in response:
|
|
events.append(chunk)
|
|
|
|
assert len(events) > 0, "Should receive streaming events"
|
|
|
|
created_idx = next(
|
|
(
|
|
i
|
|
for i, e in enumerate(events)
|
|
if getattr(e, "type", None) == "response.created"
|
|
),
|
|
None,
|
|
)
|
|
in_progress_idx = next(
|
|
(
|
|
i
|
|
for i, e in enumerate(events)
|
|
if getattr(e, "type", None) == "response.in_progress"
|
|
),
|
|
None,
|
|
)
|
|
output_item_added_idx = next(
|
|
(
|
|
i
|
|
for i, e in enumerate(events)
|
|
if getattr(e, "type", None) == "response.output_item.added"
|
|
),
|
|
None,
|
|
)
|
|
mcp_in_progress_idx = next(
|
|
(
|
|
i
|
|
for i, e in enumerate(events)
|
|
if "mcp_list_tools.in_progress" in str(getattr(e, "type", ""))
|
|
),
|
|
None,
|
|
)
|
|
completed_idx = next(
|
|
(
|
|
i
|
|
for i, e in enumerate(events)
|
|
if getattr(e, "type", None) == "response.completed"
|
|
),
|
|
None,
|
|
)
|
|
|
|
assert created_idx is not None, "response.created event should be present"
|
|
assert (
|
|
in_progress_idx is not None
|
|
), "response.in_progress event should be present"
|
|
assert (
|
|
output_item_added_idx is not None
|
|
), "response.output_item.added event should be present"
|
|
|
|
assert (
|
|
created_idx < in_progress_idx
|
|
), "response.created should come before response.in_progress"
|
|
assert (
|
|
in_progress_idx < output_item_added_idx
|
|
), "response.in_progress should come before response.output_item.added"
|
|
|
|
if mcp_in_progress_idx is not None:
|
|
assert (
|
|
output_item_added_idx < mcp_in_progress_idx
|
|
), "response.output_item.added should come before response.mcp_list_tools.in_progress"
|
|
|
|
response_ids = []
|
|
for i, event in enumerate(events):
|
|
event_type = getattr(event, "type", None)
|
|
if hasattr(event, "response"):
|
|
response_obj = getattr(event, "response", None)
|
|
if response_obj and hasattr(response_obj, "id"):
|
|
event_type_value = (
|
|
event_type.value
|
|
if hasattr(event_type, "value")
|
|
else str(event_type)
|
|
)
|
|
if any(
|
|
x in event_type_value
|
|
for x in [
|
|
"response.created",
|
|
"response.in_progress",
|
|
"response.completed",
|
|
]
|
|
):
|
|
response_ids.append((i, event_type_value, response_obj.id))
|
|
|
|
assert (
|
|
len(response_ids) >= 2
|
|
), f"Should have at least 2 response lifecycle events. Found {len(response_ids)}"
|
|
|
|
cycles = []
|
|
current_cycle = []
|
|
current_id = None
|
|
|
|
for idx, event_type, resp_id in response_ids:
|
|
if current_id is None or resp_id == current_id:
|
|
current_cycle.append((idx, event_type, resp_id))
|
|
current_id = resp_id
|
|
else:
|
|
if current_cycle:
|
|
cycles.append(current_cycle)
|
|
current_cycle = [(idx, event_type, resp_id)]
|
|
current_id = resp_id
|
|
if current_cycle:
|
|
cycles.append(current_cycle)
|
|
|
|
for cycle_num, cycle in enumerate(cycles):
|
|
cycle_ids = set(resp_id for _, _, resp_id in cycle)
|
|
assert (
|
|
len(cycle_ids) == 1
|
|
), f"Cycle {cycle_num + 1} should have consistent response ID. Found {len(cycle_ids)} unique IDs"
|
|
|
|
assert (
|
|
completed_idx is not None
|
|
), "response.completed event should be present"
|
|
|
|
lite_errors = [
|
|
record
|
|
for record in caplog.records
|
|
if record.levelno >= logging.ERROR
|
|
and ("LiteLLM" in record.name or "LiteLLM" in record.getMessage())
|
|
]
|
|
assert not lite_errors, "Unexpected LiteLLM errors: " + ", ".join(
|
|
record.getMessage() for record in lite_errors
|
|
)
|