mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-11 11:04:32 +00:00
Merge pull request #19623 from BerriAI/litellm_fix_completions_mcp_output_ordering
[fix] completions mcp output ordering
This commit is contained in:
YutaSaito
@@ -206,8 +206,18 @@ async def test_completion_mcp_with_streaming_no_timeout_error(monkeypatch):
|
||||
)
|
||||
|
||||
# Create a mock streaming response
|
||||
from unittest.mock import MagicMock, AsyncMock
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.model_call_details = {}
|
||||
logging_obj.async_failure_handler = AsyncMock()
|
||||
|
||||
class MockStreamingResponse(CustomStreamWrapper):
|
||||
def __init__(self):
|
||||
super().__init__(
|
||||
completion_stream=None,
|
||||
model="gpt-4o-mini",
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
self.chunks = [
|
||||
type('Chunk', (), {
|
||||
'choices': [type('Choice', (), {
|
||||
@@ -233,39 +243,150 @@ async def test_completion_mcp_with_streaming_no_timeout_error(monkeypatch):
|
||||
if self._index < len(self.chunks):
|
||||
chunk = self.chunks[self._index]
|
||||
self._index += 1
|
||||
# Add mcp_list_tools to first chunk if present
|
||||
if not self.sent_first_chunk:
|
||||
chunk = self._add_mcp_list_tools_to_first_chunk(chunk)
|
||||
self.sent_first_chunk = True
|
||||
return chunk
|
||||
raise StopIteration
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
if self._index < len(self.chunks):
|
||||
chunk = self.chunks[self._index]
|
||||
self._index += 1
|
||||
# Add mcp_list_tools to first chunk if present
|
||||
if not self.sent_first_chunk:
|
||||
chunk = self._add_mcp_list_tools_to_first_chunk(chunk)
|
||||
self.sent_first_chunk = True
|
||||
return chunk
|
||||
raise StopAsyncIteration
|
||||
|
||||
# Track calls to acompletion
|
||||
acompletion_calls = []
|
||||
|
||||
# Create mock streaming response for initial call
|
||||
from unittest.mock import MagicMock, AsyncMock
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.model_call_details = {}
|
||||
logging_obj.async_failure_handler = AsyncMock()
|
||||
|
||||
from litellm.types.utils import (
|
||||
ModelResponseStream,
|
||||
StreamingChoices,
|
||||
Delta,
|
||||
ChatCompletionDeltaToolCall,
|
||||
Function,
|
||||
)
|
||||
|
||||
# Create initial streaming chunks with tool_calls
|
||||
# Add tool_calls to the final chunk so stream_chunk_builder can extract them
|
||||
tool_calls = [
|
||||
ChatCompletionDeltaToolCall(
|
||||
id="call-1",
|
||||
type="function",
|
||||
function=Function(name="local_search", arguments="{}"),
|
||||
index=0,
|
||||
)
|
||||
]
|
||||
|
||||
initial_chunks = [
|
||||
ModelResponseStream(
|
||||
id="test-1",
|
||||
model="gpt-4o-mini",
|
||||
created=1234567890,
|
||||
object="chat.completion.chunk",
|
||||
choices=[
|
||||
StreamingChoices(
|
||||
index=0,
|
||||
delta=Delta(
|
||||
content="",
|
||||
role="assistant",
|
||||
tool_calls=tool_calls,
|
||||
),
|
||||
finish_reason="tool_calls",
|
||||
)
|
||||
],
|
||||
)
|
||||
]
|
||||
|
||||
class InitialStreamingResponse(CustomStreamWrapper):
|
||||
def __init__(self):
|
||||
super().__init__(
|
||||
completion_stream=None,
|
||||
model="gpt-4o-mini",
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
self.chunks = initial_chunks
|
||||
self._index = 0
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
||||
def __next__(self):
|
||||
if self._index < len(self.chunks):
|
||||
chunk = self.chunks[self._index]
|
||||
self._index += 1
|
||||
# Add mcp_list_tools to first chunk if present
|
||||
if not self.sent_first_chunk:
|
||||
chunk = self._add_mcp_list_tools_to_first_chunk(chunk)
|
||||
self.sent_first_chunk = True
|
||||
return chunk
|
||||
raise StopIteration
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
if self._index < len(self.chunks):
|
||||
chunk = self.chunks[self._index]
|
||||
self._index += 1
|
||||
# Add mcp_list_tools to first chunk if present
|
||||
if not self.sent_first_chunk:
|
||||
chunk = self._add_mcp_list_tools_to_first_chunk(chunk)
|
||||
self.sent_first_chunk = True
|
||||
return chunk
|
||||
raise StopAsyncIteration
|
||||
|
||||
async def mock_acompletion(**kwargs):
|
||||
acompletion_calls.append(kwargs)
|
||||
# First call (non-streaming for tool extraction)
|
||||
if not kwargs.get("stream", False):
|
||||
# Return a ModelResponse with tool_calls using dict format
|
||||
return ModelResponse(
|
||||
id="test-1",
|
||||
model="gpt-4o-mini",
|
||||
choices=[{
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"tool_calls": [{
|
||||
"id": "call-1",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "local_search",
|
||||
"arguments": "{}"
|
||||
}
|
||||
}]
|
||||
},
|
||||
"finish_reason": "tool_calls"
|
||||
}],
|
||||
created=0,
|
||||
object="chat.completion",
|
||||
# With new implementation, first call should be streaming
|
||||
if kwargs.get("stream", False):
|
||||
# Check if this is the follow-up call
|
||||
messages = kwargs.get("messages", [])
|
||||
is_follow_up = any(
|
||||
msg.get("role") == "tool" or (isinstance(msg, dict) and "tool_call_id" in str(msg))
|
||||
for msg in messages
|
||||
)
|
||||
# Second call (streaming follow-up)
|
||||
return MockStreamingResponse()
|
||||
if is_follow_up:
|
||||
# Follow-up call (streaming)
|
||||
return MockStreamingResponse()
|
||||
else:
|
||||
# Initial call (streaming)
|
||||
return InitialStreamingResponse()
|
||||
# Non-streaming call should not happen with new implementation, but handle it
|
||||
return ModelResponse(
|
||||
id="test-1",
|
||||
model="gpt-4o-mini",
|
||||
choices=[{
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"tool_calls": [{
|
||||
"id": "call-1",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "local_search",
|
||||
"arguments": "{}"
|
||||
}
|
||||
}]
|
||||
},
|
||||
"finish_reason": "tool_calls"
|
||||
}],
|
||||
created=0,
|
||||
object="chat.completion",
|
||||
)
|
||||
|
||||
with patch("litellm.acompletion", side_effect=mock_acompletion):
|
||||
# This should not raise RuntimeError: Timeout context manager should be used inside a task
|
||||
@@ -303,8 +424,12 @@ async def test_completion_mcp_with_streaming_no_timeout_error(monkeypatch):
|
||||
# Verify response is a streaming response
|
||||
assert isinstance(result, CustomStreamWrapper) or hasattr(result, '__iter__')
|
||||
|
||||
# Consume the stream to ensure it works
|
||||
chunks = list(result)
|
||||
# Consume the stream to ensure it works (run in separate thread to avoid event loop conflict)
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
def consume_stream():
|
||||
return list(result)
|
||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
||||
chunks = executor.submit(consume_stream).result()
|
||||
assert len(chunks) > 0, "Should have received streaming chunks"
|
||||
|
||||
# Verify tool execution was called
|
||||
@@ -317,8 +442,10 @@ async def test_completion_mcp_with_streaming_no_timeout_error(monkeypatch):
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_metadata_in_streaming_final_chunk(monkeypatch):
|
||||
"""
|
||||
Test that MCP metadata is added to the final streaming chunk's
|
||||
delta.provider_specific_fields when using MCP tools with streaming.
|
||||
Test that MCP metadata is added correctly to streaming chunks:
|
||||
- mcp_list_tools should be in the first chunk
|
||||
- mcp_tool_calls and mcp_call_results should be in the final chunk of initial response
|
||||
- Follow-up response should be streamed after initial response
|
||||
"""
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
@@ -328,7 +455,13 @@ async def test_mcp_metadata_in_streaming_final_chunk(monkeypatch):
|
||||
)
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
from litellm.utils import CustomStreamWrapper
|
||||
from litellm.types.utils import ModelResponseStream, StreamingChoices, Delta
|
||||
from litellm.types.utils import (
|
||||
ModelResponseStream,
|
||||
StreamingChoices,
|
||||
Delta,
|
||||
ChatCompletionDeltaToolCall,
|
||||
Function,
|
||||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
|
||||
dummy_tool = SimpleNamespace(
|
||||
@@ -369,7 +502,13 @@ async def test_mcp_metadata_in_streaming_final_chunk(monkeypatch):
|
||||
)
|
||||
|
||||
# Create mock streaming chunks
|
||||
def create_chunk(content, finish_reason=None):
|
||||
def create_chunk(content, finish_reason=None, tool_calls=None):
|
||||
delta = Delta(
|
||||
content=content,
|
||||
role="assistant",
|
||||
)
|
||||
if tool_calls:
|
||||
delta.tool_calls = tool_calls
|
||||
return ModelResponseStream(
|
||||
id="test-stream",
|
||||
model="gpt-4o-mini",
|
||||
@@ -378,36 +517,48 @@ async def test_mcp_metadata_in_streaming_final_chunk(monkeypatch):
|
||||
choices=[
|
||||
StreamingChoices(
|
||||
index=0,
|
||||
delta=Delta(
|
||||
content=content,
|
||||
role="assistant",
|
||||
),
|
||||
delta=delta,
|
||||
finish_reason=finish_reason,
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
chunks = [
|
||||
# Create initial streaming chunks with tool_calls
|
||||
# Add tool_calls to the final chunk so stream_chunk_builder can extract them
|
||||
tool_calls = [
|
||||
ChatCompletionDeltaToolCall(
|
||||
id="call-1",
|
||||
type="function",
|
||||
function=Function(name="local_search", arguments="{}"),
|
||||
index=0,
|
||||
)
|
||||
]
|
||||
initial_chunks = [
|
||||
create_chunk("", finish_reason="tool_calls", tool_calls=tool_calls), # Final chunk with tool_calls
|
||||
]
|
||||
|
||||
# Create follow-up streaming chunks
|
||||
follow_up_chunks = [
|
||||
create_chunk("Hello"),
|
||||
create_chunk(" world"),
|
||||
create_chunk("!", finish_reason="stop"), # Final chunk
|
||||
]
|
||||
|
||||
# Create a proper CustomStreamWrapper with logging_obj
|
||||
from unittest.mock import MagicMock
|
||||
from unittest.mock import MagicMock, AsyncMock
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.model_call_details = {}
|
||||
logging_obj.async_failure_handler = AsyncMock()
|
||||
|
||||
class MockStreamingResponse(CustomStreamWrapper):
|
||||
class InitialStreamingResponse(CustomStreamWrapper):
|
||||
def __init__(self):
|
||||
super().__init__(
|
||||
completion_stream=None,
|
||||
model="gpt-4o-mini",
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
self.chunks = chunks
|
||||
self.chunks = initial_chunks
|
||||
self._index = 0
|
||||
self.sent_last_chunk = False
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
@@ -416,42 +567,106 @@ async def test_mcp_metadata_in_streaming_final_chunk(monkeypatch):
|
||||
if self._index < len(self.chunks):
|
||||
chunk = self.chunks[self._index]
|
||||
self._index += 1
|
||||
if self._index == len(self.chunks):
|
||||
self.sent_last_chunk = True
|
||||
# Call the method that adds MCP metadata to final chunk
|
||||
chunk = self._add_mcp_metadata_to_final_chunk(chunk)
|
||||
# Add mcp_list_tools to first chunk if present
|
||||
if not self.sent_first_chunk:
|
||||
chunk = self._add_mcp_list_tools_to_first_chunk(chunk)
|
||||
self.sent_first_chunk = True
|
||||
return chunk
|
||||
raise StopIteration
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
if self._index < len(self.chunks):
|
||||
chunk = self.chunks[self._index]
|
||||
self._index += 1
|
||||
# Add mcp_list_tools to first chunk if present
|
||||
if not self.sent_first_chunk:
|
||||
chunk = self._add_mcp_list_tools_to_first_chunk(chunk)
|
||||
self.sent_first_chunk = True
|
||||
return chunk
|
||||
raise StopAsyncIteration
|
||||
|
||||
class FollowUpStreamingResponse(CustomStreamWrapper):
|
||||
def __init__(self):
|
||||
super().__init__(
|
||||
completion_stream=None,
|
||||
model="gpt-4o-mini",
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
self.chunks = follow_up_chunks
|
||||
self._index = 0
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
||||
def __next__(self):
|
||||
if self._index < len(self.chunks):
|
||||
chunk = self.chunks[self._index]
|
||||
self._index += 1
|
||||
# Add mcp_list_tools to first chunk if present
|
||||
if not self.sent_first_chunk:
|
||||
chunk = self._add_mcp_list_tools_to_first_chunk(chunk)
|
||||
self.sent_first_chunk = True
|
||||
return chunk
|
||||
raise StopIteration
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
if self._index < len(self.chunks):
|
||||
chunk = self.chunks[self._index]
|
||||
self._index += 1
|
||||
# Add mcp_list_tools to first chunk if present
|
||||
if not self.sent_first_chunk:
|
||||
chunk = self._add_mcp_list_tools_to_first_chunk(chunk)
|
||||
self.sent_first_chunk = True
|
||||
return chunk
|
||||
raise StopAsyncIteration
|
||||
|
||||
# Track calls to acompletion
|
||||
acompletion_calls = []
|
||||
|
||||
async def mock_acompletion(**kwargs):
|
||||
acompletion_calls.append(kwargs)
|
||||
# First call (non-streaming for tool extraction)
|
||||
if not kwargs.get("stream", False):
|
||||
return ModelResponse(
|
||||
id="test-1",
|
||||
model="gpt-4o-mini",
|
||||
choices=[{
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"tool_calls": [{
|
||||
"id": "call-1",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "local_search",
|
||||
"arguments": "{}"
|
||||
}
|
||||
}]
|
||||
},
|
||||
"finish_reason": "tool_calls"
|
||||
}],
|
||||
created=0,
|
||||
object="chat.completion",
|
||||
# With new implementation, first call should be streaming
|
||||
if kwargs.get("stream", False):
|
||||
# Check if this is the follow-up call (has tool results in messages)
|
||||
messages = kwargs.get("messages", [])
|
||||
is_follow_up = any(
|
||||
msg.get("role") == "tool" or (isinstance(msg, dict) and "tool_call_id" in str(msg))
|
||||
for msg in messages
|
||||
)
|
||||
# Second call (streaming follow-up)
|
||||
return MockStreamingResponse()
|
||||
|
||||
if is_follow_up:
|
||||
# Follow-up call - return follow-up chunks
|
||||
return FollowUpStreamingResponse()
|
||||
else:
|
||||
# Initial streaming call - return chunks with tool_calls
|
||||
return InitialStreamingResponse()
|
||||
# Non-streaming call should not happen with new implementation, but handle it
|
||||
return ModelResponse(
|
||||
id="test-1",
|
||||
model="gpt-4o-mini",
|
||||
choices=[{
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"tool_calls": [{
|
||||
"id": "call-1",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "local_search",
|
||||
"arguments": "{}"
|
||||
}
|
||||
}]
|
||||
},
|
||||
"finish_reason": "tool_calls"
|
||||
}],
|
||||
created=0,
|
||||
object="chat.completion",
|
||||
)
|
||||
|
||||
with patch("litellm.acompletion", side_effect=mock_acompletion):
|
||||
response = litellm.completion(
|
||||
@@ -482,39 +697,363 @@ async def test_mcp_metadata_in_streaming_final_chunk(monkeypatch):
|
||||
|
||||
assert isinstance(result, CustomStreamWrapper)
|
||||
|
||||
# Verify _hidden_params contains mcp_metadata
|
||||
assert hasattr(result, "_hidden_params")
|
||||
assert "mcp_metadata" in result._hidden_params
|
||||
mcp_metadata = result._hidden_params["mcp_metadata"]
|
||||
assert "mcp_list_tools" in mcp_metadata
|
||||
assert "mcp_tool_calls" in mcp_metadata
|
||||
assert "mcp_call_results" in mcp_metadata
|
||||
# Consume the stream and check chunks (run in separate thread to avoid event loop conflict)
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
def consume_stream():
|
||||
return list(result)
|
||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
||||
all_chunks = executor.submit(consume_stream).result()
|
||||
assert len(all_chunks) > 0, "Should have received streaming chunks"
|
||||
|
||||
# Consume the stream and check final chunk
|
||||
all_chunks = list(result)
|
||||
assert len(all_chunks) > 0
|
||||
|
||||
# Find the final chunk (with finish_reason)
|
||||
final_chunk = None
|
||||
# Find chunks from initial response (with tool_calls finish_reason)
|
||||
initial_chunks_list = []
|
||||
follow_up_chunks_list = []
|
||||
for chunk in all_chunks:
|
||||
if hasattr(chunk, "choices") and chunk.choices:
|
||||
choice = chunk.choices[0]
|
||||
if hasattr(choice, "finish_reason") and choice.finish_reason:
|
||||
final_chunk = chunk
|
||||
if hasattr(choice, "finish_reason") and choice.finish_reason == "tool_calls":
|
||||
initial_chunks_list.append(chunk)
|
||||
elif hasattr(choice, "finish_reason") and choice.finish_reason == "stop":
|
||||
follow_up_chunks_list.append(chunk)
|
||||
elif not hasattr(choice, "finish_reason") or choice.finish_reason is None:
|
||||
# Chunks without finish_reason could be from either stream
|
||||
# Check if we've seen tool_calls yet
|
||||
if initial_chunks_list:
|
||||
follow_up_chunks_list.append(chunk)
|
||||
else:
|
||||
initial_chunks_list.append(chunk)
|
||||
|
||||
# Verify initial response chunks
|
||||
assert len(initial_chunks_list) > 0, "Should have initial response chunks"
|
||||
|
||||
# Find the final chunk from initial response (with tool_calls finish_reason)
|
||||
initial_final_chunk = None
|
||||
for chunk in initial_chunks_list:
|
||||
if hasattr(chunk, "choices") and chunk.choices:
|
||||
choice = chunk.choices[0]
|
||||
if hasattr(choice, "finish_reason") and choice.finish_reason == "tool_calls":
|
||||
initial_final_chunk = chunk
|
||||
break
|
||||
|
||||
if initial_final_chunk is None and initial_chunks_list:
|
||||
initial_final_chunk = initial_chunks_list[-1]
|
||||
|
||||
# If no chunk with finish_reason, use the last chunk
|
||||
if final_chunk is None and all_chunks:
|
||||
final_chunk = all_chunks[-1]
|
||||
assert initial_final_chunk is not None, "Should have a final chunk from initial response"
|
||||
|
||||
assert final_chunk is not None, "Should have a final chunk"
|
||||
# Verify mcp_list_tools is in the first chunk of initial response
|
||||
first_chunk = initial_chunks_list[0] if initial_chunks_list else None
|
||||
assert first_chunk is not None, "Should have a first chunk"
|
||||
if hasattr(first_chunk, "choices") and first_chunk.choices:
|
||||
choice = first_chunk.choices[0]
|
||||
if hasattr(choice, "delta") and choice.delta:
|
||||
provider_fields = getattr(choice.delta, "provider_specific_fields", None)
|
||||
assert provider_fields is not None, "First chunk should have provider_specific_fields"
|
||||
assert "mcp_list_tools" in provider_fields, "First chunk should have mcp_list_tools"
|
||||
|
||||
# Verify MCP metadata is in the final chunk's delta.provider_specific_fields
|
||||
if hasattr(final_chunk, "choices") and final_chunk.choices:
|
||||
choice = final_chunk.choices[0]
|
||||
# Verify mcp_tool_calls and mcp_call_results are in the final chunk of initial response
|
||||
if hasattr(initial_final_chunk, "choices") and initial_final_chunk.choices:
|
||||
choice = initial_final_chunk.choices[0]
|
||||
if hasattr(choice, "delta") and choice.delta:
|
||||
provider_fields = getattr(choice.delta, "provider_specific_fields", None)
|
||||
assert provider_fields is not None, "Final chunk should have provider_specific_fields"
|
||||
assert "mcp_list_tools" in provider_fields, "Should have mcp_list_tools"
|
||||
assert "mcp_tool_calls" in provider_fields, "Should have mcp_tool_calls"
|
||||
assert "mcp_call_results" in provider_fields, "Should have mcp_call_results"
|
||||
|
||||
# Verify follow-up response chunks are present
|
||||
assert len(follow_up_chunks_list) > 0, "Should have follow-up response chunks"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_streaming_metadata_ordering(monkeypatch):
|
||||
"""
|
||||
Test that MCP metadata appears in the correct order:
|
||||
- mcp_list_tools should appear in the first chunk (before tool_calls)
|
||||
- mcp_tool_calls and mcp_call_results should appear in the final chunk of initial response
|
||||
- Follow-up response should be streamed after initial response completes
|
||||
"""
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
from litellm.responses.mcp.litellm_proxy_mcp_handler import (
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
)
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
from litellm.utils import CustomStreamWrapper
|
||||
from litellm.types.utils import (
|
||||
ModelResponseStream,
|
||||
StreamingChoices,
|
||||
Delta,
|
||||
ChatCompletionDeltaToolCall,
|
||||
Function,
|
||||
)
|
||||
|
||||
dummy_tool = SimpleNamespace(
|
||||
name="local_search",
|
||||
description="search",
|
||||
inputSchema={"type": "object", "properties": {}},
|
||||
)
|
||||
|
||||
async def fake_process(user_api_key_auth, mcp_tools_with_litellm_proxy):
|
||||
return [dummy_tool], {"local_search": "local"}
|
||||
|
||||
async def fake_execute(**kwargs):
|
||||
tool_calls = kwargs.get("tool_calls") or []
|
||||
call_entry = tool_calls[0]
|
||||
call_id = call_entry.get("id") or call_entry.get("call_id") or "call"
|
||||
return [
|
||||
{
|
||||
"tool_call_id": call_id,
|
||||
"result": "executed",
|
||||
"name": call_entry.get("name", "local_search"),
|
||||
}
|
||||
]
|
||||
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"_process_mcp_tools_without_openai_transform",
|
||||
fake_process,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"_execute_tool_calls",
|
||||
fake_execute,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
ResponsesAPIRequestUtils,
|
||||
"extract_mcp_headers_from_request",
|
||||
staticmethod(lambda secret_fields, tools: (None, None, None, None)),
|
||||
)
|
||||
|
||||
# Create mock streaming chunks
|
||||
def create_chunk(content, finish_reason=None, tool_calls=None):
|
||||
delta = Delta(
|
||||
content=content,
|
||||
role="assistant",
|
||||
)
|
||||
if tool_calls:
|
||||
delta.tool_calls = tool_calls
|
||||
return ModelResponseStream(
|
||||
id="test-stream",
|
||||
model="gpt-4o-mini",
|
||||
created=1234567890,
|
||||
object="chat.completion.chunk",
|
||||
choices=[
|
||||
StreamingChoices(
|
||||
index=0,
|
||||
delta=delta,
|
||||
finish_reason=finish_reason,
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
# Create initial streaming chunks with tool_calls
|
||||
# Add tool_calls to the final chunk so stream_chunk_builder can extract them
|
||||
tool_calls = [
|
||||
ChatCompletionDeltaToolCall(
|
||||
id="call-1",
|
||||
type="function",
|
||||
function=Function(name="local_search", arguments="{}"),
|
||||
index=0,
|
||||
)
|
||||
]
|
||||
initial_chunks = [
|
||||
create_chunk("", finish_reason="tool_calls", tool_calls=tool_calls), # Final chunk with tool_calls
|
||||
]
|
||||
|
||||
# Create follow-up streaming chunks
|
||||
follow_up_chunks = [
|
||||
create_chunk("Hello"),
|
||||
create_chunk(" world"),
|
||||
create_chunk("!", finish_reason="stop"), # Final chunk
|
||||
]
|
||||
|
||||
# Create a proper CustomStreamWrapper with logging_obj
|
||||
from unittest.mock import MagicMock, AsyncMock
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.model_call_details = {}
|
||||
logging_obj.async_failure_handler = AsyncMock()
|
||||
|
||||
class InitialStreamingResponse(CustomStreamWrapper):
|
||||
def __init__(self):
|
||||
super().__init__(
|
||||
completion_stream=None,
|
||||
model="gpt-4o-mini",
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
self.chunks = initial_chunks
|
||||
self._index = 0
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
||||
def __next__(self):
|
||||
if self._index < len(self.chunks):
|
||||
chunk = self.chunks[self._index]
|
||||
self._index += 1
|
||||
# Add mcp_list_tools to first chunk if present
|
||||
if not self.sent_first_chunk:
|
||||
chunk = self._add_mcp_list_tools_to_first_chunk(chunk)
|
||||
self.sent_first_chunk = True
|
||||
return chunk
|
||||
raise StopIteration
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
if self._index < len(self.chunks):
|
||||
chunk = self.chunks[self._index]
|
||||
self._index += 1
|
||||
# Add mcp_list_tools to first chunk if present
|
||||
if not self.sent_first_chunk:
|
||||
chunk = self._add_mcp_list_tools_to_first_chunk(chunk)
|
||||
self.sent_first_chunk = True
|
||||
return chunk
|
||||
raise StopAsyncIteration
|
||||
|
||||
class FollowUpStreamingResponse(CustomStreamWrapper):
|
||||
def __init__(self):
|
||||
super().__init__(
|
||||
completion_stream=None,
|
||||
model="gpt-4o-mini",
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
self.chunks = follow_up_chunks
|
||||
self._index = 0
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
||||
def __next__(self):
|
||||
if self._index < len(self.chunks):
|
||||
chunk = self.chunks[self._index]
|
||||
self._index += 1
|
||||
# Add mcp_list_tools to first chunk if present
|
||||
if not self.sent_first_chunk:
|
||||
chunk = self._add_mcp_list_tools_to_first_chunk(chunk)
|
||||
self.sent_first_chunk = True
|
||||
return chunk
|
||||
raise StopIteration
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
if self._index < len(self.chunks):
|
||||
chunk = self.chunks[self._index]
|
||||
self._index += 1
|
||||
# Add mcp_list_tools to first chunk if present
|
||||
if not self.sent_first_chunk:
|
||||
chunk = self._add_mcp_list_tools_to_first_chunk(chunk)
|
||||
self.sent_first_chunk = True
|
||||
return chunk
|
||||
raise StopAsyncIteration
|
||||
|
||||
# Track calls to acompletion
|
||||
acompletion_calls = []
|
||||
|
||||
async def mock_acompletion(**kwargs):
|
||||
acompletion_calls.append(kwargs)
|
||||
# With new implementation, first call should be streaming
|
||||
if kwargs.get("stream", False):
|
||||
# Check if this is the follow-up call (has tool results in messages)
|
||||
messages = kwargs.get("messages", [])
|
||||
is_follow_up = any(
|
||||
msg.get("role") == "tool" or (isinstance(msg, dict) and "tool_call_id" in str(msg))
|
||||
for msg in messages
|
||||
)
|
||||
|
||||
if is_follow_up:
|
||||
# Follow-up call - return follow-up chunks
|
||||
return FollowUpStreamingResponse()
|
||||
else:
|
||||
# Initial streaming call - return chunks with tool_calls
|
||||
return InitialStreamingResponse()
|
||||
# Non-streaming call should not happen with new implementation
|
||||
pytest.fail("Non-streaming call should not happen with new implementation")
|
||||
|
||||
with patch("litellm.acompletion", side_effect=mock_acompletion):
|
||||
response = litellm.completion(
|
||||
model="gpt-4o-mini",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
tools=[
|
||||
{
|
||||
"type": "mcp",
|
||||
"server_url": "litellm_proxy/mcp/local",
|
||||
"server_label": "local",
|
||||
"require_approval": "never",
|
||||
}
|
||||
],
|
||||
stream=True,
|
||||
mock_response="Final answer",
|
||||
mock_tool_calls=[
|
||||
{
|
||||
"id": "call-1",
|
||||
"type": "function",
|
||||
"function": {"name": "local_search", "arguments": "{}"},
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
import asyncio
|
||||
assert asyncio.iscoroutine(response)
|
||||
result = await response
|
||||
|
||||
assert isinstance(result, CustomStreamWrapper)
|
||||
|
||||
# Consume the stream and verify order (run in separate thread to avoid event loop conflict)
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
def consume_stream():
|
||||
return list(result)
|
||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
||||
all_chunks = executor.submit(consume_stream).result()
|
||||
assert len(all_chunks) > 0, "Should have received streaming chunks"
|
||||
|
||||
# Track when we see each type of metadata
|
||||
mcp_list_tools_seen = False
|
||||
mcp_tool_calls_seen = False
|
||||
mcp_call_results_seen = False
|
||||
tool_calls_finish_reason_seen = False
|
||||
follow_up_content_seen = False
|
||||
|
||||
for i, chunk in enumerate(all_chunks):
|
||||
if hasattr(chunk, "choices") and chunk.choices:
|
||||
choice = chunk.choices[0]
|
||||
if hasattr(choice, "delta") and choice.delta:
|
||||
provider_fields = getattr(choice.delta, "provider_specific_fields", None)
|
||||
if provider_fields:
|
||||
if "mcp_list_tools" in provider_fields:
|
||||
mcp_list_tools_seen = True
|
||||
# mcp_list_tools should appear before tool_calls finish_reason
|
||||
assert not tool_calls_finish_reason_seen, \
|
||||
"mcp_list_tools should appear before tool_calls finish_reason"
|
||||
if "mcp_tool_calls" in provider_fields:
|
||||
mcp_tool_calls_seen = True
|
||||
if "mcp_call_results" in provider_fields:
|
||||
mcp_call_results_seen = True
|
||||
|
||||
if hasattr(choice, "finish_reason") and choice.finish_reason == "tool_calls":
|
||||
tool_calls_finish_reason_seen = True
|
||||
# mcp_tool_calls and mcp_call_results should be in the same chunk as tool_calls finish_reason
|
||||
if hasattr(choice, "delta") and choice.delta:
|
||||
provider_fields = getattr(choice.delta, "provider_specific_fields", None)
|
||||
assert provider_fields is not None
|
||||
assert "mcp_tool_calls" in provider_fields, \
|
||||
"mcp_tool_calls should be in the chunk with tool_calls finish_reason"
|
||||
assert "mcp_call_results" in provider_fields, \
|
||||
"mcp_call_results should be in the chunk with tool_calls finish_reason"
|
||||
|
||||
if hasattr(choice, "delta") and choice.delta and choice.delta.content:
|
||||
content = choice.delta.content
|
||||
if content and ("Hello" in content or "world" in content or "!" in content):
|
||||
follow_up_content_seen = True
|
||||
# Follow-up content should appear after tool_calls finish_reason
|
||||
assert tool_calls_finish_reason_seen, \
|
||||
"Follow-up content should appear after tool_calls finish_reason"
|
||||
|
||||
# Verify all metadata was seen
|
||||
assert mcp_list_tools_seen, "Should have seen mcp_list_tools"
|
||||
assert mcp_tool_calls_seen, "Should have seen mcp_tool_calls"
|
||||
assert mcp_call_results_seen, "Should have seen mcp_call_results"
|
||||
assert tool_calls_finish_reason_seen, "Should have seen tool_calls finish_reason"
|
||||
assert follow_up_content_seen, "Should have seen follow-up content"
|
||||
|
||||
@@ -3,6 +3,7 @@ from unittest.mock import AsyncMock, patch
|
||||
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
from litellm.responses.mcp import chat_completions_handler
|
||||
from litellm.responses.mcp.chat_completions_handler import (
|
||||
acompletion_with_mcp,
|
||||
)
|
||||
@@ -91,24 +92,116 @@ async def test_acompletion_with_mcp_without_auto_execution_calls_model(monkeypat
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acompletion_with_mcp_auto_exec_performs_follow_up(monkeypatch):
|
||||
from litellm.utils import CustomStreamWrapper
|
||||
from litellm.types.utils import ModelResponseStream, StreamingChoices, Delta, ChatCompletionDeltaToolCall, Function
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
tools = [{"type": "function", "function": {"name": "tool"}}]
|
||||
initial_response = ModelResponse(
|
||||
id="1",
|
||||
model="test",
|
||||
choices=[],
|
||||
created=0,
|
||||
object="chat.completion",
|
||||
)
|
||||
follow_up_response = ModelResponse(
|
||||
id="2",
|
||||
model="test",
|
||||
choices=[],
|
||||
created=0,
|
||||
object="chat.completion",
|
||||
)
|
||||
mock_acompletion = AsyncMock(
|
||||
side_effect=[initial_response, follow_up_response]
|
||||
)
|
||||
|
||||
# Create mock streaming chunks for initial response
|
||||
def create_chunk(content, finish_reason=None, tool_calls=None):
|
||||
return ModelResponseStream(
|
||||
id="test-stream",
|
||||
model="test",
|
||||
created=1234567890,
|
||||
object="chat.completion.chunk",
|
||||
choices=[
|
||||
StreamingChoices(
|
||||
index=0,
|
||||
delta=Delta(
|
||||
content=content,
|
||||
role="assistant",
|
||||
tool_calls=tool_calls,
|
||||
),
|
||||
finish_reason=finish_reason,
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
initial_chunks = [
|
||||
create_chunk(
|
||||
"",
|
||||
finish_reason="tool_calls",
|
||||
tool_calls=[
|
||||
ChatCompletionDeltaToolCall(
|
||||
id="call-1",
|
||||
type="function",
|
||||
function=Function(name="tool", arguments="{}"),
|
||||
index=0,
|
||||
)
|
||||
],
|
||||
),
|
||||
]
|
||||
|
||||
follow_up_chunks = [
|
||||
create_chunk("Hello"),
|
||||
create_chunk(" world", finish_reason="stop"),
|
||||
]
|
||||
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.model_call_details = {}
|
||||
|
||||
class InitialStreamingResponse(CustomStreamWrapper):
|
||||
def __init__(self):
|
||||
super().__init__(
|
||||
completion_stream=None,
|
||||
model="test",
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
self.chunks = initial_chunks
|
||||
self._index = 0
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
if self._index < len(self.chunks):
|
||||
chunk = self.chunks[self._index]
|
||||
self._index += 1
|
||||
return chunk
|
||||
raise StopAsyncIteration
|
||||
|
||||
class FollowUpStreamingResponse(CustomStreamWrapper):
|
||||
def __init__(self):
|
||||
super().__init__(
|
||||
completion_stream=None,
|
||||
model="test",
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
self.chunks = follow_up_chunks
|
||||
self._index = 0
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
if self._index < len(self.chunks):
|
||||
chunk = self.chunks[self._index]
|
||||
self._index += 1
|
||||
return chunk
|
||||
raise StopAsyncIteration
|
||||
|
||||
async def mock_acompletion(**kwargs):
|
||||
if kwargs.get("stream", False):
|
||||
messages = kwargs.get("messages", [])
|
||||
is_follow_up = any(
|
||||
msg.get("role") == "tool" or (isinstance(msg, dict) and "tool_call_id" in str(msg))
|
||||
for msg in messages
|
||||
)
|
||||
if is_follow_up:
|
||||
return FollowUpStreamingResponse()
|
||||
else:
|
||||
return InitialStreamingResponse()
|
||||
# Non-streaming should not happen
|
||||
return ModelResponse(
|
||||
id="1",
|
||||
model="test",
|
||||
choices=[],
|
||||
created=0,
|
||||
object="chat.completion",
|
||||
)
|
||||
|
||||
mock_acompletion_func = AsyncMock(side_effect=mock_acompletion)
|
||||
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
@@ -141,10 +234,10 @@ async def test_acompletion_with_mcp_auto_exec_performs_follow_up(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"_extract_tool_calls_from_chat_response",
|
||||
staticmethod(lambda **_: ["call"]),
|
||||
staticmethod(lambda **_: [{"id": "call-1", "type": "function", "function": {"name": "tool", "arguments": "{}"}}]),
|
||||
)
|
||||
async def mock_execute(**_):
|
||||
return ["result"]
|
||||
return [{"tool_call_id": "call-1", "result": "executed"}]
|
||||
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
@@ -154,7 +247,11 @@ async def test_acompletion_with_mcp_auto_exec_performs_follow_up(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"_create_follow_up_messages_for_chat",
|
||||
staticmethod(lambda **_: ["follow-up"]),
|
||||
staticmethod(lambda **_: [
|
||||
{"role": "user", "content": "hello"},
|
||||
{"role": "assistant", "tool_calls": [{"id": "call-1", "type": "function", "function": {"name": "tool", "arguments": "{}"}}]},
|
||||
{"role": "tool", "tool_call_id": "call-1", "name": "tool", "content": "executed"}
|
||||
]),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
ResponsesAPIRequestUtils,
|
||||
@@ -162,21 +259,40 @@ async def test_acompletion_with_mcp_auto_exec_performs_follow_up(monkeypatch):
|
||||
staticmethod(lambda **_: (None, None, None, None)),
|
||||
)
|
||||
|
||||
with patch("litellm.acompletion", mock_acompletion):
|
||||
# Patch litellm.acompletion at module level to catch function-level imports
|
||||
with patch("litellm.acompletion", mock_acompletion_func), \
|
||||
patch.object(chat_completions_handler, "litellm_acompletion", mock_acompletion_func, create=True):
|
||||
result = await acompletion_with_mcp(
|
||||
model="test-model",
|
||||
messages=["msg"],
|
||||
model="gpt-4o-mini",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
tools=tools,
|
||||
stream=True,
|
||||
)
|
||||
|
||||
assert result is follow_up_response
|
||||
assert mock_acompletion.await_count == 2
|
||||
first_call = mock_acompletion.await_args_list[0].kwargs
|
||||
second_call = mock_acompletion.await_args_list[1].kwargs
|
||||
assert first_call["stream"] is False
|
||||
assert second_call["messages"] == ["follow-up"]
|
||||
assert second_call["stream"] is True
|
||||
# Consume the stream to trigger the iterator and follow-up call
|
||||
# The initial stream has one chunk with finish_reason="tool_calls"
|
||||
# which will trigger tool execution and follow-up call
|
||||
chunks = []
|
||||
async for chunk in result:
|
||||
chunks.append(chunk)
|
||||
# After consuming the initial chunk, the follow-up call should be made
|
||||
# Break after first chunk since that's when follow-up is triggered
|
||||
break
|
||||
|
||||
# With new implementation, first call should be streaming
|
||||
assert mock_acompletion_func.await_count >= 2
|
||||
first_call = mock_acompletion_func.await_args_list[0].kwargs
|
||||
# First call should be streaming in new implementation
|
||||
assert first_call["stream"] is True
|
||||
# Find the follow-up call (should have tool role messages)
|
||||
follow_up_call = None
|
||||
for call in mock_acompletion_func.await_args_list:
|
||||
messages = call.kwargs.get("messages", [])
|
||||
if messages and any(msg.get("role") == "tool" for msg in messages if isinstance(msg, dict)):
|
||||
follow_up_call = call.kwargs
|
||||
break
|
||||
assert follow_up_call is not None, "Should have a follow-up call"
|
||||
assert follow_up_call["stream"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -191,7 +307,7 @@ async def test_acompletion_with_mcp_adds_metadata_to_streaming(monkeypatch):
|
||||
|
||||
tools = [{"type": "mcp", "server_url": "litellm_proxy/mcp/local"}]
|
||||
openai_tools = [{"type": "function", "function": {"name": "local_search"}}]
|
||||
tool_calls = [{"id": "call-1", "type": "function", "function": {"name": "local_search"}}]
|
||||
tool_calls = [{"id": "call-1", "type": "function", "function": {"name": "local_search", "arguments": "{}"}}]
|
||||
tool_results = [{"tool_call_id": "call-1", "result": "executed"}]
|
||||
|
||||
# Create mock streaming chunks
|
||||
@@ -234,19 +350,21 @@ async def test_acompletion_with_mcp_adds_metadata_to_streaming(monkeypatch):
|
||||
self._index = 0
|
||||
self.sent_last_chunk = False
|
||||
|
||||
def __iter__(self):
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
def __next__(self):
|
||||
async def __anext__(self):
|
||||
if self._index < len(self.chunks):
|
||||
chunk = self.chunks[self._index]
|
||||
self._index += 1
|
||||
if self._index == len(self.chunks):
|
||||
self.sent_last_chunk = True
|
||||
# Call the method that adds MCP metadata to final chunk
|
||||
chunk = self._add_mcp_metadata_to_final_chunk(chunk)
|
||||
# Add mcp_list_tools to first chunk if present
|
||||
if not self.sent_first_chunk:
|
||||
chunk = self._add_mcp_list_tools_to_first_chunk(chunk)
|
||||
self.sent_first_chunk = True
|
||||
return chunk
|
||||
raise StopIteration
|
||||
raise StopAsyncIteration
|
||||
|
||||
mock_acompletion = AsyncMock(return_value=MockStreamingResponse())
|
||||
|
||||
@@ -286,7 +404,7 @@ async def test_acompletion_with_mcp_adds_metadata_to_streaming(monkeypatch):
|
||||
|
||||
with patch("litellm.acompletion", mock_acompletion):
|
||||
result = await acompletion_with_mcp(
|
||||
model="test-model",
|
||||
model="gpt-4o-mini",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
tools=tools,
|
||||
stream=True,
|
||||
@@ -302,30 +420,383 @@ async def test_acompletion_with_mcp_adds_metadata_to_streaming(monkeypatch):
|
||||
assert "mcp_list_tools" in mcp_metadata
|
||||
assert mcp_metadata["mcp_list_tools"] == openai_tools
|
||||
|
||||
# Consume the stream and check final chunk
|
||||
all_chunks = list(result)
|
||||
# Consume the stream and check chunks
|
||||
all_chunks = []
|
||||
async for chunk in result:
|
||||
all_chunks.append(chunk)
|
||||
assert len(all_chunks) > 0
|
||||
|
||||
# Find the final chunk (with finish_reason)
|
||||
final_chunk = None
|
||||
# Verify mcp_list_tools is in the first chunk
|
||||
first_chunk = all_chunks[0] if all_chunks else None
|
||||
assert first_chunk is not None, "Should have a first chunk"
|
||||
if hasattr(first_chunk, "choices") and first_chunk.choices:
|
||||
choice = first_chunk.choices[0]
|
||||
if hasattr(choice, "delta") and choice.delta:
|
||||
provider_fields = getattr(choice.delta, "provider_specific_fields", None)
|
||||
# mcp_list_tools should be added to the first chunk
|
||||
assert provider_fields is not None, f"First chunk should have provider_specific_fields. Delta: {choice.delta}"
|
||||
assert "mcp_list_tools" in provider_fields, f"First chunk should have mcp_list_tools. Fields: {provider_fields}"
|
||||
assert provider_fields["mcp_list_tools"] == openai_tools
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acompletion_with_mcp_streaming_initial_call_is_streaming(monkeypatch):
|
||||
"""
|
||||
Test that acompletion_with_mcp makes the initial LLM call with streaming=True
|
||||
when stream=True is requested, instead of making a non-streaming call first.
|
||||
"""
|
||||
from litellm.utils import CustomStreamWrapper
|
||||
from litellm.types.utils import ModelResponseStream, StreamingChoices, Delta
|
||||
|
||||
tools = [{"type": "mcp", "server_url": "litellm_proxy/mcp/local"}]
|
||||
openai_tools = [{"type": "function", "function": {"name": "local_search"}}]
|
||||
|
||||
# Create mock streaming chunks
|
||||
def create_chunk(content, finish_reason=None):
|
||||
return ModelResponseStream(
|
||||
id="test-stream",
|
||||
model="test-model",
|
||||
created=1234567890,
|
||||
object="chat.completion.chunk",
|
||||
choices=[
|
||||
StreamingChoices(
|
||||
index=0,
|
||||
delta=Delta(
|
||||
content=content,
|
||||
role="assistant",
|
||||
),
|
||||
finish_reason=finish_reason,
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
chunks = [
|
||||
create_chunk("", finish_reason="tool_calls"), # Final chunk with tool_calls
|
||||
]
|
||||
|
||||
# Create a proper CustomStreamWrapper
|
||||
from unittest.mock import MagicMock
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.model_call_details = {}
|
||||
|
||||
class MockStreamingResponse(CustomStreamWrapper):
|
||||
def __init__(self):
|
||||
super().__init__(
|
||||
completion_stream=None,
|
||||
model="test-model",
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
self.chunks = chunks
|
||||
self._index = 0
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
if self._index < len(self.chunks):
|
||||
chunk = self.chunks[self._index]
|
||||
self._index += 1
|
||||
return chunk
|
||||
raise StopAsyncIteration
|
||||
|
||||
mock_acompletion = AsyncMock(return_value=MockStreamingResponse())
|
||||
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"_should_use_litellm_mcp_gateway",
|
||||
staticmethod(lambda tools: True),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"_parse_mcp_tools",
|
||||
staticmethod(lambda tools: (tools, [])),
|
||||
)
|
||||
async def mock_process(**_):
|
||||
return (tools, {"local_search": "local"})
|
||||
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"_process_mcp_tools_without_openai_transform",
|
||||
mock_process,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"_transform_mcp_tools_to_openai",
|
||||
staticmethod(lambda *_, **__: openai_tools),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"_should_auto_execute_tools",
|
||||
staticmethod(lambda **_: True),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"_extract_tool_calls_from_chat_response",
|
||||
staticmethod(lambda **_: [{"id": "call-1", "type": "function", "function": {"name": "local_search", "arguments": "{}"}}]),
|
||||
)
|
||||
async def mock_execute(**_):
|
||||
return [{"tool_call_id": "call-1", "result": "executed"}]
|
||||
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"_execute_tool_calls",
|
||||
mock_execute,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"_create_follow_up_messages_for_chat",
|
||||
staticmethod(lambda **_: [
|
||||
{"role": "user", "content": "hello"},
|
||||
{"role": "assistant", "tool_calls": [{"id": "call-1", "type": "function", "function": {"name": "local_search", "arguments": "{}"}}]},
|
||||
{"role": "tool", "tool_call_id": "call-1", "name": "local_search", "content": "executed"}
|
||||
]),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
ResponsesAPIRequestUtils,
|
||||
"extract_mcp_headers_from_request",
|
||||
staticmethod(lambda **_: (None, None, None, None)),
|
||||
)
|
||||
|
||||
# Patch litellm.acompletion at module level to catch function-level imports
|
||||
with patch("litellm.acompletion", mock_acompletion), \
|
||||
patch.object(chat_completions_handler, "litellm_acompletion", mock_acompletion, create=True):
|
||||
result = await acompletion_with_mcp(
|
||||
model="gpt-4o-mini",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
tools=tools,
|
||||
stream=True,
|
||||
)
|
||||
|
||||
# Verify result is CustomStreamWrapper
|
||||
assert isinstance(result, CustomStreamWrapper)
|
||||
|
||||
# Verify that the first call was made with stream=True
|
||||
assert mock_acompletion.await_count >= 1
|
||||
first_call = mock_acompletion.await_args_list[0].kwargs
|
||||
assert first_call["stream"] is True, "First call should be streaming with new implementation"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acompletion_with_mcp_streaming_metadata_in_correct_chunks(monkeypatch):
|
||||
"""
|
||||
Test that MCP metadata is added to the correct chunks:
|
||||
- mcp_list_tools should be in the first chunk
|
||||
- mcp_tool_calls and mcp_call_results should be in the final chunk of initial response
|
||||
"""
|
||||
from litellm.utils import CustomStreamWrapper
|
||||
from litellm.types.utils import ModelResponseStream, StreamingChoices, Delta, ChatCompletionDeltaToolCall, Function
|
||||
|
||||
tools = [{"type": "mcp", "server_url": "litellm_proxy/mcp/local"}]
|
||||
openai_tools = [{"type": "function", "function": {"name": "local_search"}}]
|
||||
tool_calls = [{"id": "call-1", "type": "function", "function": {"name": "local_search", "arguments": "{}"}}]
|
||||
tool_results = [{"tool_call_id": "call-1", "result": "executed"}]
|
||||
|
||||
# Create mock streaming chunks
|
||||
def create_chunk(content, finish_reason=None, tool_calls=None):
|
||||
return ModelResponseStream(
|
||||
id="test-stream",
|
||||
model="test-model",
|
||||
created=1234567890,
|
||||
object="chat.completion.chunk",
|
||||
choices=[
|
||||
StreamingChoices(
|
||||
index=0,
|
||||
delta=Delta(
|
||||
content=content,
|
||||
role="assistant",
|
||||
tool_calls=tool_calls,
|
||||
),
|
||||
finish_reason=finish_reason,
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
initial_chunks = [
|
||||
create_chunk(
|
||||
"",
|
||||
finish_reason="tool_calls",
|
||||
tool_calls=[
|
||||
ChatCompletionDeltaToolCall(
|
||||
id="call-1",
|
||||
type="function",
|
||||
function=Function(name="local_search", arguments="{}"),
|
||||
index=0,
|
||||
)
|
||||
],
|
||||
), # Final chunk with tool_calls
|
||||
]
|
||||
|
||||
follow_up_chunks = [
|
||||
create_chunk("Hello"),
|
||||
create_chunk(" world", finish_reason="stop"),
|
||||
]
|
||||
|
||||
# Create a proper CustomStreamWrapper
|
||||
from unittest.mock import MagicMock
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.model_call_details = {}
|
||||
|
||||
class InitialStreamingResponse(CustomStreamWrapper):
|
||||
def __init__(self):
|
||||
super().__init__(
|
||||
completion_stream=None,
|
||||
model="test-model",
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
self.chunks = initial_chunks
|
||||
self._index = 0
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
if self._index < len(self.chunks):
|
||||
chunk = self.chunks[self._index]
|
||||
self._index += 1
|
||||
return chunk
|
||||
raise StopAsyncIteration
|
||||
|
||||
class FollowUpStreamingResponse(CustomStreamWrapper):
|
||||
def __init__(self):
|
||||
super().__init__(
|
||||
completion_stream=None,
|
||||
model="test-model",
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
self.chunks = follow_up_chunks
|
||||
self._index = 0
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
if self._index < len(self.chunks):
|
||||
chunk = self.chunks[self._index]
|
||||
self._index += 1
|
||||
return chunk
|
||||
raise StopAsyncIteration
|
||||
|
||||
acompletion_calls = []
|
||||
|
||||
async def mock_acompletion(**kwargs):
|
||||
acompletion_calls.append(kwargs)
|
||||
if kwargs.get("stream", False):
|
||||
messages = kwargs.get("messages", [])
|
||||
is_follow_up = any(
|
||||
msg.get("role") == "tool" or (isinstance(msg, dict) and "tool_call_id" in str(msg))
|
||||
for msg in messages
|
||||
)
|
||||
if is_follow_up:
|
||||
return FollowUpStreamingResponse()
|
||||
else:
|
||||
return InitialStreamingResponse()
|
||||
pytest.fail("Non-streaming call should not happen with new implementation")
|
||||
|
||||
mock_acompletion_func = AsyncMock(side_effect=mock_acompletion)
|
||||
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"_should_use_litellm_mcp_gateway",
|
||||
staticmethod(lambda tools: True),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"_parse_mcp_tools",
|
||||
staticmethod(lambda tools: (tools, [])),
|
||||
)
|
||||
async def mock_process(**_):
|
||||
return (tools, {"local_search": "local"})
|
||||
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"_process_mcp_tools_without_openai_transform",
|
||||
mock_process,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"_transform_mcp_tools_to_openai",
|
||||
staticmethod(lambda *_, **__: openai_tools),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"_should_auto_execute_tools",
|
||||
staticmethod(lambda **_: True),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"_extract_tool_calls_from_chat_response",
|
||||
staticmethod(lambda **_: tool_calls),
|
||||
)
|
||||
async def mock_execute(**_):
|
||||
return tool_results
|
||||
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"_execute_tool_calls",
|
||||
mock_execute,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"_create_follow_up_messages_for_chat",
|
||||
staticmethod(lambda **_: [
|
||||
{"role": "user", "content": "hello"},
|
||||
{"role": "assistant", "tool_calls": [{"id": "call-1", "type": "function", "function": {"name": "local_search", "arguments": "{}"}}]},
|
||||
{"role": "tool", "tool_call_id": "call-1", "name": "local_search", "content": "executed"}
|
||||
]),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
ResponsesAPIRequestUtils,
|
||||
"extract_mcp_headers_from_request",
|
||||
staticmethod(lambda **_: (None, None, None, None)),
|
||||
)
|
||||
|
||||
# Patch litellm.acompletion at module level to catch function-level imports
|
||||
with patch("litellm.acompletion", mock_acompletion_func), \
|
||||
patch.object(chat_completions_handler, "litellm_acompletion", side_effect=mock_acompletion, create=True):
|
||||
result = await acompletion_with_mcp(
|
||||
model="gpt-4o-mini",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
tools=tools,
|
||||
stream=True,
|
||||
)
|
||||
|
||||
# Verify result is CustomStreamWrapper
|
||||
assert isinstance(result, CustomStreamWrapper)
|
||||
|
||||
# Consume the stream and verify metadata placement
|
||||
all_chunks = []
|
||||
async for chunk in result:
|
||||
all_chunks.append(chunk)
|
||||
assert len(all_chunks) > 0
|
||||
|
||||
# Find first chunk and final chunk from initial response
|
||||
# mcp_list_tools is added to the first chunk (all_chunks[0])
|
||||
first_chunk = all_chunks[0] if all_chunks else None
|
||||
initial_final_chunk = None
|
||||
|
||||
for chunk in all_chunks:
|
||||
if hasattr(chunk, "choices") and chunk.choices:
|
||||
choice = chunk.choices[0]
|
||||
if hasattr(choice, "finish_reason") and choice.finish_reason:
|
||||
final_chunk = chunk
|
||||
break
|
||||
if hasattr(choice, "finish_reason") and choice.finish_reason == "tool_calls":
|
||||
initial_final_chunk = chunk
|
||||
|
||||
# If no chunk with finish_reason, use the last chunk
|
||||
if final_chunk is None and all_chunks:
|
||||
final_chunk = all_chunks[-1]
|
||||
assert first_chunk is not None, "Should have a first chunk"
|
||||
assert initial_final_chunk is not None, "Should have a final chunk from initial response"
|
||||
|
||||
assert final_chunk is not None, "Should have a final chunk"
|
||||
# print(first_chunk)
|
||||
# Verify mcp_list_tools is in the first chunk
|
||||
if hasattr(first_chunk, "choices") and first_chunk.choices:
|
||||
choice = first_chunk.choices[0]
|
||||
if hasattr(choice, "delta") and choice.delta:
|
||||
provider_fields = getattr(choice.delta, "provider_specific_fields", None)
|
||||
assert provider_fields is not None, "First chunk should have provider_specific_fields"
|
||||
assert "mcp_list_tools" in provider_fields, "First chunk should have mcp_list_tools"
|
||||
|
||||
# Verify MCP metadata is in the final chunk's delta.provider_specific_fields
|
||||
if hasattr(final_chunk, "choices") and final_chunk.choices:
|
||||
choice = final_chunk.choices[0]
|
||||
# Verify mcp_tool_calls and mcp_call_results are in the final chunk of initial response
|
||||
if hasattr(initial_final_chunk, "choices") and initial_final_chunk.choices:
|
||||
choice = initial_final_chunk.choices[0]
|
||||
if hasattr(choice, "delta") and choice.delta:
|
||||
provider_fields = getattr(choice.delta, "provider_specific_fields", None)
|
||||
assert provider_fields is not None, "Final chunk should have provider_specific_fields"
|
||||
assert "mcp_list_tools" in provider_fields, "Should have mcp_list_tools"
|
||||
assert provider_fields["mcp_list_tools"] == openai_tools
|
||||
assert "mcp_tool_calls" in provider_fields, "Should have mcp_tool_calls"
|
||||
assert "mcp_call_results" in provider_fields, "Should have mcp_call_results"
|
||||
|
||||
Reference in New Issue
Block a user