mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
test: Fail MCP streaming test when LiteLLM logs errors during follow-up calls
This commit is contained in:
Yuta Saito
parent
b887159b13
commit
1ff5b3ed28
1 changed files with 96 additions and 69 deletions
|
|
@ -1,3 +1,4 @@
|
|||
import logging
|
||||
import os
|
||||
import sys
|
||||
import pytest
|
||||
|
|
@ -667,7 +668,9 @@ async def test_streaming_mcp_events_validation():
|
|||
pytest.param("anthropic/claude-4-5-haiku", id="anthropic"),
|
||||
],
|
||||
)
|
||||
async def test_streaming_responses_api_with_mcp_tools(model: str):
|
||||
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"
|
||||
|
||||
|
|
@ -700,75 +703,99 @@ async def test_streaming_responses_api_with_mcp_tools(model: str):
|
|||
]
|
||||
|
||||
# Only mock the MCP-specific operations, let LLM responses be real
|
||||
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):
|
||||
"""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=[
|
||||
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):
|
||||
"""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,
|
||||
{
|
||||
"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!")
|
||||
"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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue