mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-04 02:31:27 +00:00
fix(mcp): suppress intermediate tool-call turn in auto-exec streaming
When the proxy auto-executes MCP tools (require_approval: never) with stream=true, the client received the model's first turn verbatim: raw tool_call deltas plus a mid-stream finish_reason "tool_calls", followed by the follow-up answer and a second finish_reason. OpenAI-compatible clients treat the first finish_reason as end-of-message, so chat UIs rendered an empty reply and discarded the real answer. MCPStreamingIterator now holds back tool-call delta chunks and, once the follow-up stream is created, drops them along with the intermediate finish_reason chunk. The client sees a single assistant message ending with one terminal finish_reason. Content deltas from the first turn still stream immediately, and if tool execution produces no follow-up the held chunks are flushed so the stream still terminates as before. mcp_list_tools metadata now lands on the first chunk actually yielded to the client, and mcp_tool_calls / mcp_call_results move to the terminal chunk of the follow-up stream since their previous carrier chunk is no longer surfaced. Fixes #31910
This commit is contained in:
fangkangmi
parent
c370503091
commit
8e8e96d6a8
3 changed files with 445 additions and 221 deletions
|
|
@ -245,6 +245,9 @@ async def acompletion_with_mcp(
|
|||
self.follow_up_stream = None
|
||||
self.follow_up_iterator = None
|
||||
self.follow_up_exhausted = False
|
||||
self.held_tool_call_chunks: list[ModelResponseStream] = []
|
||||
self.pending_chunks: list[ModelResponseStream] = []
|
||||
self.any_chunk_yielded = False
|
||||
|
||||
async def __aiter__(self):
|
||||
return self
|
||||
|
|
@ -319,9 +322,42 @@ async def acompletion_with_mcp(
|
|||
"Error draining inner MCP stream after final chunk; spend logging may be incomplete"
|
||||
)
|
||||
|
||||
def _is_final_chunk(self, chunk: ModelResponseStream) -> bool:
|
||||
return bool(
|
||||
hasattr(chunk, "choices")
|
||||
and chunk.choices
|
||||
and hasattr(chunk.choices[0], "finish_reason")
|
||||
and chunk.choices[0].finish_reason is not None
|
||||
)
|
||||
|
||||
def _chunk_has_tool_call_delta(self, chunk: ModelResponseStream) -> bool:
|
||||
choices = getattr(chunk, "choices", None) or []
|
||||
return any(getattr(getattr(choice, "delta", None), "tool_calls", None) for choice in choices)
|
||||
|
||||
def _yield_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream:
|
||||
if not self.any_chunk_yielded:
|
||||
self.any_chunk_yielded = True
|
||||
chunk = self._add_mcp_list_tools_to_chunk(chunk)
|
||||
return chunk
|
||||
|
||||
async def _finish_initial_turn(self):
|
||||
await self._process_tool_calls()
|
||||
if self.tool_results and self.complete_response:
|
||||
await self._prepare_follow_up_call()
|
||||
# Drain inner stream so CustomStreamWrapper fires its
|
||||
# end-of-stream handler (dispatch_success_handlers →
|
||||
# _ProxyDBLogger → LiteLLM_SpendLogs). The CSW may
|
||||
# yield one usage chunk before raising StopAsyncIteration.
|
||||
await self._drain_inner_stream()
|
||||
if self.follow_up_stream is not None:
|
||||
self.held_tool_call_chunks = []
|
||||
|
||||
async def __anext__(self):
|
||||
if self.pending_chunks:
|
||||
return self._yield_chunk(self.pending_chunks.pop(0))
|
||||
|
||||
# Phase 1: Collect and yield initial stream chunks
|
||||
if not self.stream_exhausted:
|
||||
while not self.stream_exhausted:
|
||||
# Get the iterator from the stream wrapper
|
||||
if not hasattr(self, "_stream_iterator"):
|
||||
self._stream_iterator = self.stream_wrapper.__aiter__()
|
||||
|
|
@ -333,46 +369,35 @@ async def acompletion_with_mcp(
|
|||
|
||||
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:
|
||||
self.stream_exhausted = True
|
||||
await self._process_tool_calls()
|
||||
chunk = self._add_mcp_tool_metadata_to_final_chunk(chunk)
|
||||
if self.tool_results and self.complete_response:
|
||||
await self._prepare_follow_up_call()
|
||||
# Drain inner stream so CustomStreamWrapper fires its
|
||||
# end-of-stream handler (dispatch_success_handlers →
|
||||
# _ProxyDBLogger → LiteLLM_SpendLogs). The CSW may
|
||||
# yield one usage chunk before raising StopAsyncIteration.
|
||||
await self._drain_inner_stream()
|
||||
|
||||
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()
|
||||
await self._drain_inner_stream()
|
||||
return final_chunk
|
||||
if not self.collected_chunks:
|
||||
break
|
||||
await self._finish_initial_turn()
|
||||
if self.follow_up_stream is not None:
|
||||
break
|
||||
final_chunk = self._add_mcp_tool_metadata_to_final_chunk(self.collected_chunks[-1])
|
||||
self.pending_chunks = self.held_tool_call_chunks + [final_chunk]
|
||||
self.held_tool_call_chunks = []
|
||||
return self._yield_chunk(self.pending_chunks.pop(0))
|
||||
|
||||
self.collected_chunks.append(chunk)
|
||||
|
||||
if self._is_final_chunk(chunk):
|
||||
self.stream_exhausted = True
|
||||
await self._finish_initial_turn()
|
||||
if self.follow_up_stream is not None:
|
||||
break
|
||||
chunk = self._add_mcp_tool_metadata_to_final_chunk(chunk)
|
||||
self.pending_chunks = self.held_tool_call_chunks + [chunk]
|
||||
self.held_tool_call_chunks = []
|
||||
return self._yield_chunk(self.pending_chunks.pop(0))
|
||||
|
||||
if self._chunk_has_tool_call_delta(chunk):
|
||||
self.held_tool_call_chunks.append(chunk)
|
||||
continue
|
||||
|
||||
return self._yield_chunk(chunk)
|
||||
|
||||
# Phase 2: Yield follow-up stream chunks if available
|
||||
if self.follow_up_stream and not self.follow_up_exhausted:
|
||||
|
|
@ -387,7 +412,9 @@ async def acompletion_with_mcp(
|
|||
from litellm._logging import verbose_logger
|
||||
|
||||
verbose_logger.debug(f"Follow-up chunk yielded: {chunk}")
|
||||
return chunk
|
||||
if self._is_final_chunk(chunk):
|
||||
chunk = self._add_mcp_tool_metadata_to_final_chunk(chunk)
|
||||
return self._yield_chunk(chunk)
|
||||
except StopAsyncIteration:
|
||||
self.follow_up_exhausted = True
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
|
|||
|
|
@ -464,9 +464,9 @@ async def test_completion_mcp_with_streaming_no_timeout_error(monkeypatch):
|
|||
async def test_mcp_metadata_in_streaming_final_chunk(monkeypatch):
|
||||
"""
|
||||
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
|
||||
- mcp_list_tools should be in the first chunk yielded to the client
|
||||
- mcp_tool_calls and mcp_call_results should be in the final chunk of the follow-up response
|
||||
- The intermediate tool-call turn must not leak into the client stream (issue #31910)
|
||||
"""
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
|
@ -734,95 +734,59 @@ async def test_mcp_metadata_in_streaming_final_chunk(monkeypatch):
|
|||
all_chunks = executor.submit(consume_stream).result()
|
||||
assert len(all_chunks) > 0, "Should have received streaming chunks"
|
||||
|
||||
# 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 == "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]
|
||||
finish_reasons = [
|
||||
chunk.choices[0].finish_reason
|
||||
for chunk in all_chunks
|
||||
if chunk.choices and chunk.choices[0].finish_reason is not None
|
||||
]
|
||||
assert finish_reasons == ["stop"], (
|
||||
f"Intermediate tool-call turn must not leak; expected a single stop. Got: {finish_reasons}"
|
||||
)
|
||||
assert all(
|
||||
not getattr(chunk.choices[0].delta, "tool_calls", None)
|
||||
for chunk in all_chunks
|
||||
if chunk.choices and chunk.choices[0].delta
|
||||
), "tool_call deltas must not leak into the client stream"
|
||||
|
||||
first_chunk = all_chunks[0]
|
||||
first_provider_fields = getattr(
|
||||
first_chunk.choices[0].delta, "provider_specific_fields", None
|
||||
)
|
||||
assert (
|
||||
initial_final_chunk is not None
|
||||
), "Should have a final chunk from initial response"
|
||||
first_provider_fields is not None
|
||||
), "First chunk should have provider_specific_fields"
|
||||
assert (
|
||||
"mcp_list_tools" in first_provider_fields
|
||||
), "First chunk should have mcp_list_tools"
|
||||
|
||||
# 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"
|
||||
final_chunk = all_chunks[-1]
|
||||
assert final_chunk.choices[0].finish_reason == "stop"
|
||||
final_provider_fields = getattr(
|
||||
final_chunk.choices[0].delta, "provider_specific_fields", None
|
||||
)
|
||||
assert (
|
||||
final_provider_fields is not None
|
||||
), "Final chunk should have provider_specific_fields"
|
||||
assert "mcp_tool_calls" in final_provider_fields, "Should have mcp_tool_calls"
|
||||
assert (
|
||||
"mcp_call_results" in final_provider_fields
|
||||
), "Should have mcp_call_results"
|
||||
|
||||
# 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_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"
|
||||
content = "".join(
|
||||
chunk.choices[0].delta.content or ""
|
||||
for chunk in all_chunks
|
||||
if chunk.choices and chunk.choices[0].delta
|
||||
)
|
||||
assert content == "Hello world!"
|
||||
|
||||
|
||||
@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
|
||||
- mcp_list_tools should appear in the first chunk yielded to the client
|
||||
- mcp_tool_calls and mcp_call_results should appear in the terminal chunk of the stream
|
||||
- The client stream must contain exactly one terminal finish_reason (issue #31910)
|
||||
"""
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
|
@ -1069,66 +1033,38 @@ async def test_mcp_streaming_metadata_ordering(monkeypatch):
|
|||
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
|
||||
finish_reasons = [
|
||||
chunk.choices[0].finish_reason
|
||||
for chunk in all_chunks
|
||||
if chunk.choices and chunk.choices[0].finish_reason is not None
|
||||
]
|
||||
assert finish_reasons == ["stop"], (
|
||||
f"Client stream must contain exactly one terminal finish_reason. Got: {finish_reasons}"
|
||||
)
|
||||
|
||||
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"
|
||||
first_provider_fields = getattr(
|
||||
all_chunks[0].choices[0].delta, "provider_specific_fields", None
|
||||
)
|
||||
assert (
|
||||
tool_calls_finish_reason_seen
|
||||
), "Should have seen tool_calls finish_reason"
|
||||
assert follow_up_content_seen, "Should have seen follow-up content"
|
||||
first_provider_fields is not None and "mcp_list_tools" in first_provider_fields
|
||||
), "mcp_list_tools should be in the first chunk"
|
||||
|
||||
terminal_chunk = all_chunks[-1]
|
||||
assert terminal_chunk.choices[0].finish_reason == "stop"
|
||||
terminal_provider_fields = getattr(
|
||||
terminal_chunk.choices[0].delta, "provider_specific_fields", None
|
||||
)
|
||||
assert terminal_provider_fields is not None
|
||||
assert (
|
||||
"mcp_tool_calls" in terminal_provider_fields
|
||||
), "mcp_tool_calls should be in the terminal chunk"
|
||||
assert (
|
||||
"mcp_call_results" in terminal_provider_fields
|
||||
), "mcp_call_results should be in the terminal chunk"
|
||||
|
||||
content = "".join(
|
||||
chunk.choices[0].delta.content or ""
|
||||
for chunk in all_chunks
|
||||
if chunk.choices and chunk.choices[0].delta
|
||||
)
|
||||
assert content == "Hello world!", "Follow-up answer content must reach the client"
|
||||
|
|
|
|||
|
|
@ -742,8 +742,8 @@ async def test_acompletion_with_mcp_streaming_initial_call_is_streaming(monkeypa
|
|||
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
|
||||
- mcp_list_tools should be in the first chunk yielded to the client
|
||||
- mcp_tool_calls and mcp_call_results should be in the final chunk of the follow-up response
|
||||
"""
|
||||
from litellm.utils import CustomStreamWrapper
|
||||
from litellm.types.utils import (
|
||||
|
|
@ -971,35 +971,16 @@ async def test_acompletion_with_mcp_streaming_metadata_in_correct_chunks(monkeyp
|
|||
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
|
||||
finish_reasons = [
|
||||
chunk.choices[0].finish_reason
|
||||
for chunk in all_chunks
|
||||
if chunk.choices and chunk.choices[0].finish_reason is not None
|
||||
]
|
||||
assert finish_reasons == ["stop"], f"Client stream must end with a single stop. Got: {finish_reasons}"
|
||||
|
||||
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 == "tool_calls"
|
||||
):
|
||||
initial_final_chunk = chunk
|
||||
|
||||
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"
|
||||
|
||||
# Verify mcp_list_tools is in the first chunk
|
||||
assert (
|
||||
hasattr(first_chunk, "choices") and first_chunk.choices
|
||||
), "First chunk must have choices"
|
||||
first_choice = first_chunk.choices[0]
|
||||
assert (
|
||||
hasattr(first_choice, "delta") and first_choice.delta
|
||||
), "First choice must have delta"
|
||||
first_chunk = all_chunks[0]
|
||||
first_provider_fields = getattr(
|
||||
first_choice.delta, "provider_specific_fields", None
|
||||
first_chunk.choices[0].delta, "provider_specific_fields", None
|
||||
)
|
||||
assert (
|
||||
first_provider_fields is not None
|
||||
|
|
@ -1008,16 +989,10 @@ async def test_acompletion_with_mcp_streaming_metadata_in_correct_chunks(monkeyp
|
|||
"mcp_list_tools" in first_provider_fields
|
||||
), "First chunk should have mcp_list_tools"
|
||||
|
||||
# Verify mcp_tool_calls and mcp_call_results are in the final chunk of initial response
|
||||
assert (
|
||||
hasattr(initial_final_chunk, "choices") and initial_final_chunk.choices
|
||||
), "Final chunk must have choices"
|
||||
final_choice = initial_final_chunk.choices[0]
|
||||
assert (
|
||||
hasattr(final_choice, "delta") and final_choice.delta
|
||||
), "Final choice must have delta"
|
||||
final_chunk = all_chunks[-1]
|
||||
assert final_chunk.choices[0].finish_reason == "stop"
|
||||
final_provider_fields = getattr(
|
||||
final_choice.delta, "provider_specific_fields", None
|
||||
final_chunk.choices[0].delta, "provider_specific_fields", None
|
||||
)
|
||||
assert (
|
||||
final_provider_fields is not None
|
||||
|
|
@ -1344,3 +1319,289 @@ async def test_acompletion_with_mcp_streaming_drains_inner_stream_after_exhausti
|
|||
|
||||
assert len(all_chunks) == 3
|
||||
assert initial_stream.drained_after_exhaustion is True
|
||||
|
||||
|
||||
def _create_mcp_stream_chunk(content, finish_reason=None, tool_calls=None):
|
||||
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
|
||||
|
||||
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,
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def _make_mock_stream_class(stream_chunks):
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from litellm.utils import CustomStreamWrapper
|
||||
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.model_call_details = {}
|
||||
|
||||
class MockStream(CustomStreamWrapper):
|
||||
def __init__(self):
|
||||
super().__init__(
|
||||
completion_stream=None,
|
||||
model="test-model",
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
self.chunks = stream_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
|
||||
|
||||
return MockStream
|
||||
|
||||
|
||||
def _patch_mcp_auto_exec_scaffolding(monkeypatch, tools, openai_tools, tool_calls, tool_results):
|
||||
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)),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acompletion_with_mcp_streaming_suppresses_intermediate_tool_call_turn(monkeypatch):
|
||||
"""
|
||||
Regression test for https://github.com/BerriAI/litellm/issues/31910:
|
||||
when the proxy auto-executes MCP tools with stream=True, the intermediate
|
||||
tool-call turn (raw tool_call deltas + finish_reason "tool_calls") must not
|
||||
leak into the client stream. The client should see a single assistant
|
||||
message ending with exactly one terminal finish_reason.
|
||||
"""
|
||||
from litellm.types.utils import 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"}]
|
||||
|
||||
initial_chunks = [
|
||||
_create_mcp_stream_chunk("Let me check. "),
|
||||
_create_mcp_stream_chunk(
|
||||
None,
|
||||
tool_calls=[
|
||||
ChatCompletionDeltaToolCall(
|
||||
id="call-1",
|
||||
type="function",
|
||||
function=Function(name="local_search", arguments="{}"),
|
||||
index=0,
|
||||
)
|
||||
],
|
||||
),
|
||||
_create_mcp_stream_chunk("", finish_reason="tool_calls"),
|
||||
]
|
||||
follow_up_chunks = [
|
||||
_create_mcp_stream_chunk("Hello"),
|
||||
_create_mcp_stream_chunk(" world", finish_reason="stop"),
|
||||
]
|
||||
|
||||
InitialStream = _make_mock_stream_class(initial_chunks)
|
||||
FollowUpStream = _make_mock_stream_class(follow_up_chunks)
|
||||
|
||||
async def mock_acompletion(**kwargs):
|
||||
messages = kwargs.get("messages", [])
|
||||
is_follow_up = any(
|
||||
isinstance(msg, dict) and msg.get("role") == "tool" for msg in messages
|
||||
)
|
||||
return FollowUpStream() if is_follow_up else InitialStream()
|
||||
|
||||
mock_acompletion_func = AsyncMock(side_effect=mock_acompletion)
|
||||
_patch_mcp_auto_exec_scaffolding(monkeypatch, tools, openai_tools, tool_calls, tool_results)
|
||||
|
||||
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="gpt-4o-mini",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
tools=tools,
|
||||
stream=True,
|
||||
)
|
||||
|
||||
all_chunks = []
|
||||
async for chunk in result:
|
||||
all_chunks.append(chunk)
|
||||
|
||||
assert all(
|
||||
not getattr(chunk.choices[0].delta, "tool_calls", None) for chunk in all_chunks if chunk.choices
|
||||
), f"tool_call deltas must not leak into the client stream. Got: {all_chunks}"
|
||||
|
||||
finish_reasons = [
|
||||
chunk.choices[0].finish_reason
|
||||
for chunk in all_chunks
|
||||
if chunk.choices and chunk.choices[0].finish_reason is not None
|
||||
]
|
||||
assert finish_reasons == ["stop"], f"Expected a single terminal stop. Got: {finish_reasons}"
|
||||
assert all_chunks[-1].choices[0].finish_reason == "stop"
|
||||
|
||||
content = "".join(
|
||||
chunk.choices[0].delta.content or "" for chunk in all_chunks if chunk.choices and chunk.choices[0].delta
|
||||
)
|
||||
assert content == "Let me check. Hello world"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acompletion_with_mcp_streaming_flushes_tool_call_turn_when_no_follow_up(monkeypatch):
|
||||
"""
|
||||
When tool execution produces no results (so no follow-up stream is created),
|
||||
the held tool-call chunks and the finish_reason "tool_calls" chunk must be
|
||||
flushed so the client still receives a complete, terminated stream.
|
||||
"""
|
||||
from litellm.types.utils import 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": "{}"},
|
||||
}
|
||||
]
|
||||
|
||||
initial_chunks = [
|
||||
_create_mcp_stream_chunk(
|
||||
None,
|
||||
tool_calls=[
|
||||
ChatCompletionDeltaToolCall(
|
||||
id="call-1",
|
||||
type="function",
|
||||
function=Function(name="local_search", arguments="{}"),
|
||||
index=0,
|
||||
)
|
||||
],
|
||||
),
|
||||
_create_mcp_stream_chunk("", finish_reason="tool_calls"),
|
||||
]
|
||||
|
||||
InitialStream = _make_mock_stream_class(initial_chunks)
|
||||
|
||||
async def mock_acompletion(**kwargs):
|
||||
return InitialStream()
|
||||
|
||||
mock_acompletion_func = AsyncMock(side_effect=mock_acompletion)
|
||||
_patch_mcp_auto_exec_scaffolding(monkeypatch, tools, openai_tools, tool_calls, tool_results=[])
|
||||
|
||||
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="gpt-4o-mini",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
tools=tools,
|
||||
stream=True,
|
||||
)
|
||||
|
||||
all_chunks = []
|
||||
async for chunk in result:
|
||||
all_chunks.append(chunk)
|
||||
|
||||
assert any(
|
||||
getattr(chunk.choices[0].delta, "tool_calls", None) for chunk in all_chunks if chunk.choices
|
||||
), "Held tool-call chunks must be flushed when no follow-up stream is created"
|
||||
assert all_chunks[-1].choices[0].finish_reason == "tool_calls"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue