Merge pull request #19623 from BerriAI/litellm_fix_completions_mcp_output_ordering

[fix] completions mcp output ordering
This commit is contained in:
YutaSaito
2026-01-23 15:56:02 +09:00
committed by GitHub
7 changed files with 1633 additions and 193 deletions
@@ -1571,6 +1571,50 @@ class CustomStreamWrapper:
)
return chunk
def _add_mcp_list_tools_to_first_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream:
"""
Add mcp_list_tools from _hidden_params to the first chunk's delta.provider_specific_fields.
This method checks if MCP metadata with mcp_list_tools is stored in _hidden_params
and adds it to the first chunk's delta.provider_specific_fields.
"""
try:
# Check if MCP metadata should be added to first chunk
if not hasattr(self, "_hidden_params") or not self._hidden_params:
return chunk
mcp_metadata = self._hidden_params.get("mcp_metadata")
if not mcp_metadata or not isinstance(mcp_metadata, dict):
return chunk
# Only add mcp_list_tools to first chunk (not tool_calls or tool_results)
mcp_list_tools = mcp_metadata.get("mcp_list_tools")
if not mcp_list_tools:
return chunk
# Add mcp_list_tools to delta.provider_specific_fields
if hasattr(chunk, "choices") and chunk.choices:
for choice in chunk.choices:
if isinstance(choice, StreamingChoices) and hasattr(choice, "delta") and choice.delta:
# Get existing provider_specific_fields or create new dict
provider_fields = (
getattr(choice.delta, "provider_specific_fields", None) or {}
)
# Add only mcp_list_tools to first chunk
provider_fields["mcp_list_tools"] = mcp_list_tools
# Set the provider_specific_fields
setattr(choice.delta, "provider_specific_fields", provider_fields)
except Exception as e:
from litellm._logging import verbose_logger
verbose_logger.exception(
f"Error adding MCP list tools to first chunk: {str(e)}"
)
return chunk
def _add_mcp_metadata_to_final_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream:
"""
Add MCP metadata from _hidden_params to the final chunk's delta.provider_specific_fields.
@@ -1727,6 +1771,12 @@ class CustomStreamWrapper:
)
# HANDLE STREAM OPTIONS
self.chunks.append(response)
# Add mcp_list_tools to first chunk if present
if not self.sent_first_chunk:
response = self._add_mcp_list_tools_to_first_chunk(response)
self.sent_first_chunk = True
if hasattr(
response, "usage"
): # remove usage from chunk, only send on final chunk
@@ -1894,6 +1944,11 @@ class CustomStreamWrapper:
input=self.response_uptil_now, model=self.model
)
self.chunks.append(processed_chunk)
# Add mcp_list_tools to first chunk if present
if not self.sent_first_chunk:
processed_chunk = self._add_mcp_list_tools_to_first_chunk(processed_chunk)
self.sent_first_chunk = True
if hasattr(
processed_chunk, "usage"
): # remove usage from chunk, only send on final chunk
+360 -13
View File
@@ -5,6 +5,7 @@ from typing import (
List,
Optional,
Union,
cast,
)
from litellm.responses.mcp.litellm_proxy_mcp_handler import (
@@ -78,7 +79,7 @@ def _add_mcp_metadata_to_response(
setattr(message, "provider_specific_fields", provider_fields)
async def acompletion_with_mcp(
async def acompletion_with_mcp( # noqa: PLR0915
model: str,
messages: List,
tools: Optional[List] = None,
@@ -174,10 +175,367 @@ async def acompletion_with_mcp(
)
return response
# For auto-execute: disable streaming for initial call
# For auto-execute: handle streaming vs non-streaming differently
stream = kwargs.get("stream", False)
mock_tool_calls = base_call_args.pop("mock_tool_calls", None)
if stream:
# Streaming mode: make initial call with streaming, collect chunks, detect tool calls
initial_call_args = dict(base_call_args)
initial_call_args["stream"] = True
if mock_tool_calls is not None:
initial_call_args["mock_tool_calls"] = mock_tool_calls
# Make initial streaming call
initial_stream = await litellm_acompletion(**initial_call_args)
if not isinstance(initial_stream, CustomStreamWrapper):
# Not a stream, return as-is
if isinstance(initial_stream, ModelResponse):
_add_mcp_metadata_to_response(
response=initial_stream,
openai_tools=openai_tools,
)
return initial_stream
# Create a custom async generator that collects chunks and handles tool execution
from litellm.main import stream_chunk_builder
from litellm.types.utils import ModelResponseStream
class MCPStreamingIterator:
"""Custom iterator that collects chunks, detects tool calls, and adds MCP metadata to final chunk."""
def __init__(self, stream_wrapper, messages, tool_server_map, user_api_key_auth,
mcp_auth_header, mcp_server_auth_headers, oauth2_headers, raw_headers,
litellm_call_id, litellm_trace_id, openai_tools, base_call_args):
self.stream_wrapper = stream_wrapper
self.messages = messages
self.tool_server_map = tool_server_map
self.user_api_key_auth = user_api_key_auth
self.mcp_auth_header = mcp_auth_header
self.mcp_server_auth_headers = mcp_server_auth_headers
self.oauth2_headers = oauth2_headers
self.raw_headers = raw_headers
self.litellm_call_id = litellm_call_id
self.litellm_trace_id = litellm_trace_id
self.openai_tools = openai_tools
self.base_call_args = base_call_args
self.collected_chunks: List[ModelResponseStream] = []
self.tool_calls: Optional[List] = None
self.tool_results: Optional[List] = None
self.complete_response: Optional[ModelResponse] = None
self.stream_exhausted = False
self.tool_execution_done = False
self.follow_up_stream = None
self.follow_up_iterator = None
self.follow_up_exhausted = False
async def __aiter__(self):
return self
def _add_mcp_list_tools_to_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream:
"""Add mcp_list_tools to the first chunk."""
from litellm.types.utils import StreamingChoices, add_provider_specific_fields
if not self.openai_tools:
return chunk
if hasattr(chunk, "choices") and chunk.choices:
for choice in chunk.choices:
if isinstance(choice, StreamingChoices) and hasattr(choice, "delta") and choice.delta:
# Get existing provider_specific_fields or create new dict
existing_fields = getattr(choice.delta, "provider_specific_fields", None) or {}
provider_fields = dict(existing_fields) # Create a copy to avoid mutating the original
# Add only mcp_list_tools to first chunk
provider_fields["mcp_list_tools"] = self.openai_tools
# Use add_provider_specific_fields to ensure proper setting
# This function handles Pydantic model attribute setting correctly
add_provider_specific_fields(choice.delta, provider_fields)
return chunk
def _add_mcp_tool_metadata_to_final_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream:
"""Add mcp_tool_calls and mcp_call_results to the final chunk."""
from litellm.types.utils import StreamingChoices, add_provider_specific_fields
if hasattr(chunk, "choices") and chunk.choices:
for choice in chunk.choices:
if isinstance(choice, StreamingChoices) and hasattr(choice, "delta") and choice.delta:
# Get existing provider_specific_fields or create new dict
# Access the attribute directly to handle Pydantic model attributes correctly
existing_fields = {}
if hasattr(choice.delta, "provider_specific_fields"):
attr_value = getattr(choice.delta, "provider_specific_fields", None)
if attr_value is not None:
# Create a copy to avoid mutating the original
existing_fields = dict(attr_value) if isinstance(attr_value, dict) else {}
provider_fields = existing_fields
# Add tool_calls and tool_results if available
if self.tool_calls:
provider_fields["mcp_tool_calls"] = self.tool_calls
if self.tool_results:
provider_fields["mcp_call_results"] = self.tool_results
# Use add_provider_specific_fields to ensure proper setting
# This function handles Pydantic model attribute setting correctly
add_provider_specific_fields(choice.delta, provider_fields)
return chunk
async def __anext__(self):
# Phase 1: Collect and yield initial stream chunks
if not self.stream_exhausted:
# Get the iterator from the stream wrapper
if not hasattr(self, '_stream_iterator'):
self._stream_iterator = self.stream_wrapper.__aiter__()
# Add mcp_list_tools to the first chunk (available from the start)
_add_mcp_metadata_to_response(
response=self.stream_wrapper,
openai_tools=self.openai_tools,
)
try:
chunk = await self._stream_iterator.__anext__()
self.collected_chunks.append(chunk)
# Add mcp_list_tools to the first chunk
if len(self.collected_chunks) == 1:
chunk = self._add_mcp_list_tools_to_chunk(chunk)
# Check if this is the final chunk (has finish_reason)
is_final = (
hasattr(chunk, "choices")
and chunk.choices
and hasattr(chunk.choices[0], "finish_reason")
and chunk.choices[0].finish_reason is not None
)
if is_final:
# This is the final chunk, mark stream as exhausted
self.stream_exhausted = True
# Process tool calls after we've collected all chunks
await self._process_tool_calls()
# Apply MCP metadata (tool_calls and tool_results) to final chunk
chunk = self._add_mcp_tool_metadata_to_final_chunk(chunk)
# If we have tool results, prepare follow-up call immediately
if self.tool_results and self.complete_response:
await self._prepare_follow_up_call()
return chunk
except StopAsyncIteration:
self.stream_exhausted = True
# Process tool calls after stream is exhausted
await self._process_tool_calls()
# If we have chunks, yield the final one with metadata
if self.collected_chunks:
final_chunk = self.collected_chunks[-1]
final_chunk = self._add_mcp_tool_metadata_to_final_chunk(final_chunk)
# If we have tool results, prepare follow-up call
if self.tool_results and self.complete_response:
await self._prepare_follow_up_call()
return final_chunk
# Phase 2: Yield follow-up stream chunks if available
if self.follow_up_stream and not self.follow_up_exhausted:
if not self.follow_up_iterator:
self.follow_up_iterator = self.follow_up_stream.__aiter__()
from litellm._logging import verbose_logger
verbose_logger.debug("Follow-up stream iterator created")
try:
chunk = await self.follow_up_iterator.__anext__()
from litellm._logging import verbose_logger
verbose_logger.debug(f"Follow-up chunk yielded: {chunk}")
return chunk
except StopAsyncIteration:
self.follow_up_exhausted = True
from litellm._logging import verbose_logger
verbose_logger.debug("Follow-up stream exhausted")
# After follow-up stream is exhausted, check if we need to raise StopAsyncIteration
raise StopAsyncIteration
# If we're here and follow_up_stream is None but we expected it, log a warning
if self.stream_exhausted and self.tool_results and self.complete_response and self.follow_up_stream is None:
from litellm._logging import verbose_logger
verbose_logger.warning(
"Follow-up stream was not created despite having tool results"
)
raise StopAsyncIteration
async def _process_tool_calls(self):
"""Process tool calls after streaming completes."""
if self.tool_execution_done:
return
self.tool_execution_done = True
if not self.collected_chunks:
return
# Build complete response from chunks
complete_response = stream_chunk_builder(
chunks=self.collected_chunks,
messages=self.messages,
)
if isinstance(complete_response, ModelResponse):
self.complete_response = complete_response
# Extract tool calls from complete response
self.tool_calls = LiteLLM_Proxy_MCP_Handler._extract_tool_calls_from_chat_response(
response=complete_response
)
if self.tool_calls:
# Execute tool calls
self.tool_results = await LiteLLM_Proxy_MCP_Handler._execute_tool_calls(
tool_server_map=self.tool_server_map,
tool_calls=self.tool_calls,
user_api_key_auth=self.user_api_key_auth,
mcp_auth_header=self.mcp_auth_header,
mcp_server_auth_headers=self.mcp_server_auth_headers,
oauth2_headers=self.oauth2_headers,
raw_headers=self.raw_headers,
litellm_call_id=self.litellm_call_id,
litellm_trace_id=self.litellm_trace_id,
)
async def _prepare_follow_up_call(self):
"""Prepare and initiate follow-up call with tool results."""
if self.follow_up_stream is not None:
return # Already prepared
if not self.tool_results or not self.complete_response:
return
# Create follow-up messages with tool results
follow_up_messages = LiteLLM_Proxy_MCP_Handler._create_follow_up_messages_for_chat(
original_messages=self.messages,
response=self.complete_response,
tool_results=self.tool_results,
)
# Make follow-up call with streaming
follow_up_call_args = dict(self.base_call_args)
follow_up_call_args["messages"] = follow_up_messages
follow_up_call_args["stream"] = True
# Ensure follow-up call doesn't trigger MCP handler again
follow_up_call_args["_skip_mcp_handler"] = True
# Import litellm here to ensure we get the patched version
# This ensures the patch works correctly in tests
import litellm
follow_up_response = await litellm.acompletion(**follow_up_call_args)
# Ensure follow-up response is a CustomStreamWrapper
if isinstance(follow_up_response, CustomStreamWrapper):
self.follow_up_stream = follow_up_response
from litellm._logging import verbose_logger
verbose_logger.debug("Follow-up stream created successfully")
else:
# Unexpected response type - log and set to None
from litellm._logging import verbose_logger
verbose_logger.warning(
f"Follow-up response is not a CustomStreamWrapper: {type(follow_up_response)}"
)
self.follow_up_stream = None
# Create the custom iterator
iterator = MCPStreamingIterator(
stream_wrapper=initial_stream,
messages=messages,
tool_server_map=tool_server_map,
user_api_key_auth=user_api_key_auth,
mcp_auth_header=mcp_auth_header,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
litellm_call_id=kwargs.get("litellm_call_id"),
litellm_trace_id=kwargs.get("litellm_trace_id"),
openai_tools=openai_tools,
base_call_args=base_call_args,
)
# Create a wrapper class that delegates to our custom iterator
# We'll use a simple approach: just replace the __aiter__ method
class MCPStreamWrapper(CustomStreamWrapper):
def __init__(self, original_wrapper, custom_iterator):
# Initialize with the same parameters as original wrapper
super().__init__(
completion_stream=None,
model=getattr(original_wrapper, "model", "unknown"),
logging_obj=getattr(original_wrapper, "logging_obj", None),
custom_llm_provider=getattr(original_wrapper, "custom_llm_provider", None),
stream_options=getattr(original_wrapper, "stream_options", None),
make_call=getattr(original_wrapper, "make_call", None),
_response_headers=getattr(original_wrapper, "_response_headers", None),
)
self._original_wrapper = original_wrapper
self._custom_iterator = custom_iterator
# Copy important attributes from original wrapper
if hasattr(original_wrapper, "_hidden_params"):
self._hidden_params = original_wrapper._hidden_params
# For synchronous iteration, we need to run the async iterator
self._sync_iterator = None
self._sync_loop = None
def __aiter__(self):
return self._custom_iterator
def __iter__(self):
# For synchronous iteration, create a sync wrapper
if self._sync_iterator is None:
import asyncio
try:
self._sync_loop = asyncio.get_event_loop()
except RuntimeError:
self._sync_loop = asyncio.new_event_loop()
asyncio.set_event_loop(self._sync_loop)
self._sync_iterator = _SyncIteratorWrapper(self._custom_iterator, self._sync_loop)
return self._sync_iterator
def __next__(self):
# Delegate to sync iterator
if self._sync_iterator is None:
self.__iter__()
return next(self._sync_iterator)
def __getattr__(self, name):
# Delegate all other attributes to original wrapper
return getattr(self._original_wrapper, name)
# Helper class to wrap async iterator for sync iteration
class _SyncIteratorWrapper:
def __init__(self, async_iterator, loop):
self._async_iterator = async_iterator
self._loop = loop
self._iterator = None
def __iter__(self):
return self
def __next__(self):
if self._iterator is None:
# __aiter__ might be async, so we need to await it
aiter_result = self._async_iterator.__aiter__()
if hasattr(aiter_result, '__await__'):
# It's a coroutine, await it
self._iterator = self._loop.run_until_complete(aiter_result)
else:
# It's already an iterator
self._iterator = aiter_result
try:
return self._loop.run_until_complete(self._iterator.__anext__())
except StopAsyncIteration:
raise StopIteration
return cast(CustomStreamWrapper, MCPStreamWrapper(initial_stream, iterator))
# Non-streaming mode: use existing logic
initial_call_args = dict(base_call_args)
initial_call_args["stream"] = False
if mock_tool_calls is not None:
@@ -195,17 +553,6 @@ async def acompletion_with_mcp(
)
if not tool_calls:
# No tool calls, return response or retry with streaming if needed
if stream:
retry_args = dict(base_call_args)
retry_args["stream"] = stream
response = await litellm_acompletion(**retry_args)
if isinstance(response, (ModelResponse, CustomStreamWrapper)):
_add_mcp_metadata_to_response(
response=response,
openai_tools=openai_tools,
)
return response
_add_mcp_metadata_to_response(
response=initial_response,
openai_tools=openai_tools,
@@ -775,6 +775,12 @@ class LiteLLM_Proxy_MCP_Handler:
first_choice, "message", None
):
message_to_append = first_choice.message.model_dump(exclude_none=True)
# Ensure tool_calls have arguments field (required by OpenAI API)
if message_to_append.get("tool_calls"):
for tool_call in message_to_append["tool_calls"]:
if isinstance(tool_call, dict) and "function" in tool_call:
if "arguments" not in tool_call["function"]:
tool_call["function"]["arguments"] = "{}"
except Exception:
verbose_logger.exception("Failed to convert assistant message for MCP flow")
+628 -89
View File
@@ -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"
@@ -607,12 +607,16 @@ const ChatUI: React.FC<ChatUIProps> = ({
console.log("ChatUI: Received MCP event:", event);
setMCPEvents((prev) => {
// Check if this is a duplicate event (same item_id and type)
const isDuplicate = prev.some(
(existingEvent) =>
existingEvent.item_id === event.item_id &&
existingEvent.type === event.type &&
existingEvent.sequence_number === event.sequence_number,
);
// Only check for duplicates if item_id is defined (for mcp_list_tools, item_id is "mcp_list_tools")
const isDuplicate = event.item_id
? prev.some(
(existingEvent) =>
existingEvent.item_id === event.item_id &&
existingEvent.type === event.type &&
(existingEvent.sequence_number === event.sequence_number ||
(existingEvent.sequence_number === undefined && event.sequence_number === undefined)),
)
: false;
if (isDuplicate) {
console.log("ChatUI: Duplicate MCP event, skipping");
@@ -59,12 +59,13 @@ export async function makeOpenAIChatCompletionRequest(
let fullResponseContent = "";
let fullReasoningContent = "";
// Track MCP metadata from final chunk
// Track MCP metadata cumulatively across chunks
let mcpMetadata: {
mcp_list_tools?: any[];
mcp_tool_calls?: any[];
mcp_call_results?: any[];
} | null = null;
} = {};
let mcpListToolsProcessed = false;
// Build tools array
const tools: any[] = [];
@@ -167,16 +168,48 @@ export async function makeOpenAIChatCompletionRequest(
onSearchResults(delta.provider_specific_fields.search_results);
}
// Check for MCP metadata in provider_specific_fields (typically in final chunk)
// Check for MCP metadata in provider_specific_fields
if (delta && delta.provider_specific_fields) {
const providerFields = delta.provider_specific_fields;
// Merge MCP metadata cumulatively (don't overwrite)
if (providerFields.mcp_list_tools && !mcpMetadata.mcp_list_tools) {
mcpMetadata.mcp_list_tools = providerFields.mcp_list_tools;
// Process mcp_list_tools immediately when found (typically in first chunk)
if (onMCPEvent && !mcpListToolsProcessed) {
mcpListToolsProcessed = true;
const toolsEvent: MCPEvent = {
type: "response.output_item.done",
item_id: "mcp_list_tools", // Add item_id to prevent duplicate detection issues
item: {
type: "mcp_list_tools",
tools: providerFields.mcp_list_tools.map((tool: any) => ({
name: tool.function?.name || tool.name || "",
description: tool.function?.description || tool.description || "",
input_schema: tool.function?.parameters || tool.input_schema || {},
})),
},
timestamp: Date.now(),
};
onMCPEvent(toolsEvent);
console.log("MCP list_tools event sent:", toolsEvent);
}
}
if (providerFields.mcp_tool_calls) {
mcpMetadata.mcp_tool_calls = providerFields.mcp_tool_calls;
}
if (providerFields.mcp_call_results) {
mcpMetadata.mcp_call_results = providerFields.mcp_call_results;
}
if (providerFields.mcp_list_tools || providerFields.mcp_tool_calls || providerFields.mcp_call_results) {
mcpMetadata = {
mcp_list_tools: providerFields.mcp_list_tools,
mcp_tool_calls: providerFields.mcp_tool_calls,
mcp_call_results: providerFields.mcp_call_results,
};
console.log("MCP metadata found in chunk:", mcpMetadata);
console.log("MCP metadata found in chunk:", {
mcp_list_tools: providerFields.mcp_list_tools ? "present" : "absent",
mcp_tool_calls: providerFields.mcp_tool_calls ? "present" : "absent",
mcp_call_results: providerFields.mcp_call_results ? "present" : "absent",
});
}
}
@@ -204,35 +237,19 @@ export async function makeOpenAIChatCompletionRequest(
}
}
// Process MCP metadata from final chunk and convert to MCPEvent format
if (mcpMetadata && onMCPEvent) {
// Convert mcp_list_tools to MCPEvent
if (mcpMetadata.mcp_list_tools && mcpMetadata.mcp_list_tools.length > 0) {
const toolsEvent: MCPEvent = {
type: "response.output_item.done",
item: {
type: "mcp_list_tools",
tools: mcpMetadata.mcp_list_tools.map((tool: any) => ({
name: tool.function?.name || tool.name || "",
description: tool.function?.description || tool.description || "",
input_schema: tool.function?.parameters || tool.input_schema || {},
})),
},
timestamp: Date.now(),
};
onMCPEvent(toolsEvent);
}
// Process remaining MCP metadata (mcp_tool_calls and mcp_call_results) after stream completes
// Note: mcp_list_tools is already processed when found in the first chunk
if (onMCPEvent && (mcpMetadata.mcp_tool_calls || mcpMetadata.mcp_call_results)) {
// Convert mcp_tool_calls and mcp_call_results to MCPEvent[]
if (mcpMetadata?.mcp_tool_calls && mcpMetadata?.mcp_tool_calls.length > 0) {
mcpMetadata?.mcp_tool_calls.forEach((toolCall: any, index: number) => {
if (mcpMetadata.mcp_tool_calls && mcpMetadata.mcp_tool_calls.length > 0) {
mcpMetadata.mcp_tool_calls.forEach((toolCall: any, index: number) => {
const functionName = toolCall.function?.name || toolCall.name || "";
const functionArgs = toolCall.function?.arguments || toolCall.arguments || "{}";
// Find corresponding result
const result = mcpMetadata?.mcp_call_results?.find(
const result = mcpMetadata.mcp_call_results?.find(
(r: any) => r.tool_call_id === toolCall.id || r.tool_call_id === toolCall.call_id
) || mcpMetadata?.mcp_call_results?.[index];
) || mcpMetadata.mcp_call_results?.[index];
const callEvent: MCPEvent = {
type: "response.output_item.done",
@@ -246,6 +263,7 @@ export async function makeOpenAIChatCompletionRequest(
timestamp: Date.now(),
};
onMCPEvent(callEvent);
console.log("MCP call event sent:", callEvent);
});
}
}