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:
@@ -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
|
||||
|
||||
@@ -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")
|
||||
|
||||
|
||||
@@ -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);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user