mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-15 22:22:12 +00:00
Merge pull request #19623 from BerriAI/litellm_fix_completions_mcp_output_ordering
[fix] completions mcp output ordering
This commit is contained in:
YutaSaito
@@ -3,6 +3,7 @@ from unittest.mock import AsyncMock, patch
|
||||
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
from litellm.responses.mcp import chat_completions_handler
|
||||
from litellm.responses.mcp.chat_completions_handler import (
|
||||
acompletion_with_mcp,
|
||||
)
|
||||
@@ -91,24 +92,116 @@ async def test_acompletion_with_mcp_without_auto_execution_calls_model(monkeypat
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acompletion_with_mcp_auto_exec_performs_follow_up(monkeypatch):
|
||||
from litellm.utils import CustomStreamWrapper
|
||||
from litellm.types.utils import ModelResponseStream, StreamingChoices, Delta, ChatCompletionDeltaToolCall, Function
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
tools = [{"type": "function", "function": {"name": "tool"}}]
|
||||
initial_response = ModelResponse(
|
||||
id="1",
|
||||
model="test",
|
||||
choices=[],
|
||||
created=0,
|
||||
object="chat.completion",
|
||||
)
|
||||
follow_up_response = ModelResponse(
|
||||
id="2",
|
||||
model="test",
|
||||
choices=[],
|
||||
created=0,
|
||||
object="chat.completion",
|
||||
)
|
||||
mock_acompletion = AsyncMock(
|
||||
side_effect=[initial_response, follow_up_response]
|
||||
)
|
||||
|
||||
# Create mock streaming chunks for initial response
|
||||
def create_chunk(content, finish_reason=None, tool_calls=None):
|
||||
return ModelResponseStream(
|
||||
id="test-stream",
|
||||
model="test",
|
||||
created=1234567890,
|
||||
object="chat.completion.chunk",
|
||||
choices=[
|
||||
StreamingChoices(
|
||||
index=0,
|
||||
delta=Delta(
|
||||
content=content,
|
||||
role="assistant",
|
||||
tool_calls=tool_calls,
|
||||
),
|
||||
finish_reason=finish_reason,
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
initial_chunks = [
|
||||
create_chunk(
|
||||
"",
|
||||
finish_reason="tool_calls",
|
||||
tool_calls=[
|
||||
ChatCompletionDeltaToolCall(
|
||||
id="call-1",
|
||||
type="function",
|
||||
function=Function(name="tool", arguments="{}"),
|
||||
index=0,
|
||||
)
|
||||
],
|
||||
),
|
||||
]
|
||||
|
||||
follow_up_chunks = [
|
||||
create_chunk("Hello"),
|
||||
create_chunk(" world", finish_reason="stop"),
|
||||
]
|
||||
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.model_call_details = {}
|
||||
|
||||
class InitialStreamingResponse(CustomStreamWrapper):
|
||||
def __init__(self):
|
||||
super().__init__(
|
||||
completion_stream=None,
|
||||
model="test",
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
self.chunks = initial_chunks
|
||||
self._index = 0
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
if self._index < len(self.chunks):
|
||||
chunk = self.chunks[self._index]
|
||||
self._index += 1
|
||||
return chunk
|
||||
raise StopAsyncIteration
|
||||
|
||||
class FollowUpStreamingResponse(CustomStreamWrapper):
|
||||
def __init__(self):
|
||||
super().__init__(
|
||||
completion_stream=None,
|
||||
model="test",
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
self.chunks = follow_up_chunks
|
||||
self._index = 0
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
if self._index < len(self.chunks):
|
||||
chunk = self.chunks[self._index]
|
||||
self._index += 1
|
||||
return chunk
|
||||
raise StopAsyncIteration
|
||||
|
||||
async def mock_acompletion(**kwargs):
|
||||
if kwargs.get("stream", False):
|
||||
messages = kwargs.get("messages", [])
|
||||
is_follow_up = any(
|
||||
msg.get("role") == "tool" or (isinstance(msg, dict) and "tool_call_id" in str(msg))
|
||||
for msg in messages
|
||||
)
|
||||
if is_follow_up:
|
||||
return FollowUpStreamingResponse()
|
||||
else:
|
||||
return InitialStreamingResponse()
|
||||
# Non-streaming should not happen
|
||||
return ModelResponse(
|
||||
id="1",
|
||||
model="test",
|
||||
choices=[],
|
||||
created=0,
|
||||
object="chat.completion",
|
||||
)
|
||||
|
||||
mock_acompletion_func = AsyncMock(side_effect=mock_acompletion)
|
||||
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
@@ -141,10 +234,10 @@ async def test_acompletion_with_mcp_auto_exec_performs_follow_up(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"_extract_tool_calls_from_chat_response",
|
||||
staticmethod(lambda **_: ["call"]),
|
||||
staticmethod(lambda **_: [{"id": "call-1", "type": "function", "function": {"name": "tool", "arguments": "{}"}}]),
|
||||
)
|
||||
async def mock_execute(**_):
|
||||
return ["result"]
|
||||
return [{"tool_call_id": "call-1", "result": "executed"}]
|
||||
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
@@ -154,7 +247,11 @@ async def test_acompletion_with_mcp_auto_exec_performs_follow_up(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"_create_follow_up_messages_for_chat",
|
||||
staticmethod(lambda **_: ["follow-up"]),
|
||||
staticmethod(lambda **_: [
|
||||
{"role": "user", "content": "hello"},
|
||||
{"role": "assistant", "tool_calls": [{"id": "call-1", "type": "function", "function": {"name": "tool", "arguments": "{}"}}]},
|
||||
{"role": "tool", "tool_call_id": "call-1", "name": "tool", "content": "executed"}
|
||||
]),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
ResponsesAPIRequestUtils,
|
||||
@@ -162,21 +259,40 @@ async def test_acompletion_with_mcp_auto_exec_performs_follow_up(monkeypatch):
|
||||
staticmethod(lambda **_: (None, None, None, None)),
|
||||
)
|
||||
|
||||
with patch("litellm.acompletion", mock_acompletion):
|
||||
# Patch litellm.acompletion at module level to catch function-level imports
|
||||
with patch("litellm.acompletion", mock_acompletion_func), \
|
||||
patch.object(chat_completions_handler, "litellm_acompletion", mock_acompletion_func, create=True):
|
||||
result = await acompletion_with_mcp(
|
||||
model="test-model",
|
||||
messages=["msg"],
|
||||
model="gpt-4o-mini",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
tools=tools,
|
||||
stream=True,
|
||||
)
|
||||
|
||||
assert result is follow_up_response
|
||||
assert mock_acompletion.await_count == 2
|
||||
first_call = mock_acompletion.await_args_list[0].kwargs
|
||||
second_call = mock_acompletion.await_args_list[1].kwargs
|
||||
assert first_call["stream"] is False
|
||||
assert second_call["messages"] == ["follow-up"]
|
||||
assert second_call["stream"] is True
|
||||
# Consume the stream to trigger the iterator and follow-up call
|
||||
# The initial stream has one chunk with finish_reason="tool_calls"
|
||||
# which will trigger tool execution and follow-up call
|
||||
chunks = []
|
||||
async for chunk in result:
|
||||
chunks.append(chunk)
|
||||
# After consuming the initial chunk, the follow-up call should be made
|
||||
# Break after first chunk since that's when follow-up is triggered
|
||||
break
|
||||
|
||||
# With new implementation, first call should be streaming
|
||||
assert mock_acompletion_func.await_count >= 2
|
||||
first_call = mock_acompletion_func.await_args_list[0].kwargs
|
||||
# First call should be streaming in new implementation
|
||||
assert first_call["stream"] is True
|
||||
# Find the follow-up call (should have tool role messages)
|
||||
follow_up_call = None
|
||||
for call in mock_acompletion_func.await_args_list:
|
||||
messages = call.kwargs.get("messages", [])
|
||||
if messages and any(msg.get("role") == "tool" for msg in messages if isinstance(msg, dict)):
|
||||
follow_up_call = call.kwargs
|
||||
break
|
||||
assert follow_up_call is not None, "Should have a follow-up call"
|
||||
assert follow_up_call["stream"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -191,7 +307,7 @@ async def test_acompletion_with_mcp_adds_metadata_to_streaming(monkeypatch):
|
||||
|
||||
tools = [{"type": "mcp", "server_url": "litellm_proxy/mcp/local"}]
|
||||
openai_tools = [{"type": "function", "function": {"name": "local_search"}}]
|
||||
tool_calls = [{"id": "call-1", "type": "function", "function": {"name": "local_search"}}]
|
||||
tool_calls = [{"id": "call-1", "type": "function", "function": {"name": "local_search", "arguments": "{}"}}]
|
||||
tool_results = [{"tool_call_id": "call-1", "result": "executed"}]
|
||||
|
||||
# Create mock streaming chunks
|
||||
@@ -234,19 +350,21 @@ async def test_acompletion_with_mcp_adds_metadata_to_streaming(monkeypatch):
|
||||
self._index = 0
|
||||
self.sent_last_chunk = False
|
||||
|
||||
def __iter__(self):
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
def __next__(self):
|
||||
async def __anext__(self):
|
||||
if self._index < len(self.chunks):
|
||||
chunk = self.chunks[self._index]
|
||||
self._index += 1
|
||||
if self._index == len(self.chunks):
|
||||
self.sent_last_chunk = True
|
||||
# Call the method that adds MCP metadata to final chunk
|
||||
chunk = self._add_mcp_metadata_to_final_chunk(chunk)
|
||||
# Add mcp_list_tools to first chunk if present
|
||||
if not self.sent_first_chunk:
|
||||
chunk = self._add_mcp_list_tools_to_first_chunk(chunk)
|
||||
self.sent_first_chunk = True
|
||||
return chunk
|
||||
raise StopIteration
|
||||
raise StopAsyncIteration
|
||||
|
||||
mock_acompletion = AsyncMock(return_value=MockStreamingResponse())
|
||||
|
||||
@@ -286,7 +404,7 @@ async def test_acompletion_with_mcp_adds_metadata_to_streaming(monkeypatch):
|
||||
|
||||
with patch("litellm.acompletion", mock_acompletion):
|
||||
result = await acompletion_with_mcp(
|
||||
model="test-model",
|
||||
model="gpt-4o-mini",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
tools=tools,
|
||||
stream=True,
|
||||
@@ -302,30 +420,383 @@ async def test_acompletion_with_mcp_adds_metadata_to_streaming(monkeypatch):
|
||||
assert "mcp_list_tools" in mcp_metadata
|
||||
assert mcp_metadata["mcp_list_tools"] == openai_tools
|
||||
|
||||
# Consume the stream and check final chunk
|
||||
all_chunks = list(result)
|
||||
# Consume the stream and check chunks
|
||||
all_chunks = []
|
||||
async for chunk in result:
|
||||
all_chunks.append(chunk)
|
||||
assert len(all_chunks) > 0
|
||||
|
||||
# Find the final chunk (with finish_reason)
|
||||
final_chunk = None
|
||||
# Verify mcp_list_tools is in the first chunk
|
||||
first_chunk = all_chunks[0] if all_chunks else None
|
||||
assert first_chunk is not None, "Should have a first chunk"
|
||||
if hasattr(first_chunk, "choices") and first_chunk.choices:
|
||||
choice = first_chunk.choices[0]
|
||||
if hasattr(choice, "delta") and choice.delta:
|
||||
provider_fields = getattr(choice.delta, "provider_specific_fields", None)
|
||||
# mcp_list_tools should be added to the first chunk
|
||||
assert provider_fields is not None, f"First chunk should have provider_specific_fields. Delta: {choice.delta}"
|
||||
assert "mcp_list_tools" in provider_fields, f"First chunk should have mcp_list_tools. Fields: {provider_fields}"
|
||||
assert provider_fields["mcp_list_tools"] == openai_tools
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acompletion_with_mcp_streaming_initial_call_is_streaming(monkeypatch):
|
||||
"""
|
||||
Test that acompletion_with_mcp makes the initial LLM call with streaming=True
|
||||
when stream=True is requested, instead of making a non-streaming call first.
|
||||
"""
|
||||
from litellm.utils import CustomStreamWrapper
|
||||
from litellm.types.utils import ModelResponseStream, StreamingChoices, Delta
|
||||
|
||||
tools = [{"type": "mcp", "server_url": "litellm_proxy/mcp/local"}]
|
||||
openai_tools = [{"type": "function", "function": {"name": "local_search"}}]
|
||||
|
||||
# Create mock streaming chunks
|
||||
def create_chunk(content, finish_reason=None):
|
||||
return ModelResponseStream(
|
||||
id="test-stream",
|
||||
model="test-model",
|
||||
created=1234567890,
|
||||
object="chat.completion.chunk",
|
||||
choices=[
|
||||
StreamingChoices(
|
||||
index=0,
|
||||
delta=Delta(
|
||||
content=content,
|
||||
role="assistant",
|
||||
),
|
||||
finish_reason=finish_reason,
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
chunks = [
|
||||
create_chunk("", finish_reason="tool_calls"), # Final chunk with tool_calls
|
||||
]
|
||||
|
||||
# Create a proper CustomStreamWrapper
|
||||
from unittest.mock import MagicMock
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.model_call_details = {}
|
||||
|
||||
class MockStreamingResponse(CustomStreamWrapper):
|
||||
def __init__(self):
|
||||
super().__init__(
|
||||
completion_stream=None,
|
||||
model="test-model",
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
self.chunks = chunks
|
||||
self._index = 0
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
if self._index < len(self.chunks):
|
||||
chunk = self.chunks[self._index]
|
||||
self._index += 1
|
||||
return chunk
|
||||
raise StopAsyncIteration
|
||||
|
||||
mock_acompletion = AsyncMock(return_value=MockStreamingResponse())
|
||||
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"_should_use_litellm_mcp_gateway",
|
||||
staticmethod(lambda tools: True),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"_parse_mcp_tools",
|
||||
staticmethod(lambda tools: (tools, [])),
|
||||
)
|
||||
async def mock_process(**_):
|
||||
return (tools, {"local_search": "local"})
|
||||
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"_process_mcp_tools_without_openai_transform",
|
||||
mock_process,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"_transform_mcp_tools_to_openai",
|
||||
staticmethod(lambda *_, **__: openai_tools),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"_should_auto_execute_tools",
|
||||
staticmethod(lambda **_: True),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"_extract_tool_calls_from_chat_response",
|
||||
staticmethod(lambda **_: [{"id": "call-1", "type": "function", "function": {"name": "local_search", "arguments": "{}"}}]),
|
||||
)
|
||||
async def mock_execute(**_):
|
||||
return [{"tool_call_id": "call-1", "result": "executed"}]
|
||||
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"_execute_tool_calls",
|
||||
mock_execute,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"_create_follow_up_messages_for_chat",
|
||||
staticmethod(lambda **_: [
|
||||
{"role": "user", "content": "hello"},
|
||||
{"role": "assistant", "tool_calls": [{"id": "call-1", "type": "function", "function": {"name": "local_search", "arguments": "{}"}}]},
|
||||
{"role": "tool", "tool_call_id": "call-1", "name": "local_search", "content": "executed"}
|
||||
]),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
ResponsesAPIRequestUtils,
|
||||
"extract_mcp_headers_from_request",
|
||||
staticmethod(lambda **_: (None, None, None, None)),
|
||||
)
|
||||
|
||||
# Patch litellm.acompletion at module level to catch function-level imports
|
||||
with patch("litellm.acompletion", mock_acompletion), \
|
||||
patch.object(chat_completions_handler, "litellm_acompletion", mock_acompletion, create=True):
|
||||
result = await acompletion_with_mcp(
|
||||
model="gpt-4o-mini",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
tools=tools,
|
||||
stream=True,
|
||||
)
|
||||
|
||||
# Verify result is CustomStreamWrapper
|
||||
assert isinstance(result, CustomStreamWrapper)
|
||||
|
||||
# Verify that the first call was made with stream=True
|
||||
assert mock_acompletion.await_count >= 1
|
||||
first_call = mock_acompletion.await_args_list[0].kwargs
|
||||
assert first_call["stream"] is True, "First call should be streaming with new implementation"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acompletion_with_mcp_streaming_metadata_in_correct_chunks(monkeypatch):
|
||||
"""
|
||||
Test that MCP metadata is added to the correct chunks:
|
||||
- mcp_list_tools should be in the first chunk
|
||||
- mcp_tool_calls and mcp_call_results should be in the final chunk of initial response
|
||||
"""
|
||||
from litellm.utils import CustomStreamWrapper
|
||||
from litellm.types.utils import ModelResponseStream, StreamingChoices, Delta, ChatCompletionDeltaToolCall, Function
|
||||
|
||||
tools = [{"type": "mcp", "server_url": "litellm_proxy/mcp/local"}]
|
||||
openai_tools = [{"type": "function", "function": {"name": "local_search"}}]
|
||||
tool_calls = [{"id": "call-1", "type": "function", "function": {"name": "local_search", "arguments": "{}"}}]
|
||||
tool_results = [{"tool_call_id": "call-1", "result": "executed"}]
|
||||
|
||||
# Create mock streaming chunks
|
||||
def create_chunk(content, finish_reason=None, tool_calls=None):
|
||||
return ModelResponseStream(
|
||||
id="test-stream",
|
||||
model="test-model",
|
||||
created=1234567890,
|
||||
object="chat.completion.chunk",
|
||||
choices=[
|
||||
StreamingChoices(
|
||||
index=0,
|
||||
delta=Delta(
|
||||
content=content,
|
||||
role="assistant",
|
||||
tool_calls=tool_calls,
|
||||
),
|
||||
finish_reason=finish_reason,
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
initial_chunks = [
|
||||
create_chunk(
|
||||
"",
|
||||
finish_reason="tool_calls",
|
||||
tool_calls=[
|
||||
ChatCompletionDeltaToolCall(
|
||||
id="call-1",
|
||||
type="function",
|
||||
function=Function(name="local_search", arguments="{}"),
|
||||
index=0,
|
||||
)
|
||||
],
|
||||
), # Final chunk with tool_calls
|
||||
]
|
||||
|
||||
follow_up_chunks = [
|
||||
create_chunk("Hello"),
|
||||
create_chunk(" world", finish_reason="stop"),
|
||||
]
|
||||
|
||||
# Create a proper CustomStreamWrapper
|
||||
from unittest.mock import MagicMock
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.model_call_details = {}
|
||||
|
||||
class InitialStreamingResponse(CustomStreamWrapper):
|
||||
def __init__(self):
|
||||
super().__init__(
|
||||
completion_stream=None,
|
||||
model="test-model",
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
self.chunks = initial_chunks
|
||||
self._index = 0
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
if self._index < len(self.chunks):
|
||||
chunk = self.chunks[self._index]
|
||||
self._index += 1
|
||||
return chunk
|
||||
raise StopAsyncIteration
|
||||
|
||||
class FollowUpStreamingResponse(CustomStreamWrapper):
|
||||
def __init__(self):
|
||||
super().__init__(
|
||||
completion_stream=None,
|
||||
model="test-model",
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
self.chunks = follow_up_chunks
|
||||
self._index = 0
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
if self._index < len(self.chunks):
|
||||
chunk = self.chunks[self._index]
|
||||
self._index += 1
|
||||
return chunk
|
||||
raise StopAsyncIteration
|
||||
|
||||
acompletion_calls = []
|
||||
|
||||
async def mock_acompletion(**kwargs):
|
||||
acompletion_calls.append(kwargs)
|
||||
if kwargs.get("stream", False):
|
||||
messages = kwargs.get("messages", [])
|
||||
is_follow_up = any(
|
||||
msg.get("role") == "tool" or (isinstance(msg, dict) and "tool_call_id" in str(msg))
|
||||
for msg in messages
|
||||
)
|
||||
if is_follow_up:
|
||||
return FollowUpStreamingResponse()
|
||||
else:
|
||||
return InitialStreamingResponse()
|
||||
pytest.fail("Non-streaming call should not happen with new implementation")
|
||||
|
||||
mock_acompletion_func = AsyncMock(side_effect=mock_acompletion)
|
||||
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"_should_use_litellm_mcp_gateway",
|
||||
staticmethod(lambda tools: True),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"_parse_mcp_tools",
|
||||
staticmethod(lambda tools: (tools, [])),
|
||||
)
|
||||
async def mock_process(**_):
|
||||
return (tools, {"local_search": "local"})
|
||||
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"_process_mcp_tools_without_openai_transform",
|
||||
mock_process,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"_transform_mcp_tools_to_openai",
|
||||
staticmethod(lambda *_, **__: openai_tools),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"_should_auto_execute_tools",
|
||||
staticmethod(lambda **_: True),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"_extract_tool_calls_from_chat_response",
|
||||
staticmethod(lambda **_: tool_calls),
|
||||
)
|
||||
async def mock_execute(**_):
|
||||
return tool_results
|
||||
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"_execute_tool_calls",
|
||||
mock_execute,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"_create_follow_up_messages_for_chat",
|
||||
staticmethod(lambda **_: [
|
||||
{"role": "user", "content": "hello"},
|
||||
{"role": "assistant", "tool_calls": [{"id": "call-1", "type": "function", "function": {"name": "local_search", "arguments": "{}"}}]},
|
||||
{"role": "tool", "tool_call_id": "call-1", "name": "local_search", "content": "executed"}
|
||||
]),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
ResponsesAPIRequestUtils,
|
||||
"extract_mcp_headers_from_request",
|
||||
staticmethod(lambda **_: (None, None, None, None)),
|
||||
)
|
||||
|
||||
# Patch litellm.acompletion at module level to catch function-level imports
|
||||
with patch("litellm.acompletion", mock_acompletion_func), \
|
||||
patch.object(chat_completions_handler, "litellm_acompletion", side_effect=mock_acompletion, create=True):
|
||||
result = await acompletion_with_mcp(
|
||||
model="gpt-4o-mini",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
tools=tools,
|
||||
stream=True,
|
||||
)
|
||||
|
||||
# Verify result is CustomStreamWrapper
|
||||
assert isinstance(result, CustomStreamWrapper)
|
||||
|
||||
# Consume the stream and verify metadata placement
|
||||
all_chunks = []
|
||||
async for chunk in result:
|
||||
all_chunks.append(chunk)
|
||||
assert len(all_chunks) > 0
|
||||
|
||||
# Find first chunk and final chunk from initial response
|
||||
# mcp_list_tools is added to the first chunk (all_chunks[0])
|
||||
first_chunk = all_chunks[0] if all_chunks else None
|
||||
initial_final_chunk = None
|
||||
|
||||
for chunk in all_chunks:
|
||||
if hasattr(chunk, "choices") and chunk.choices:
|
||||
choice = chunk.choices[0]
|
||||
if hasattr(choice, "finish_reason") and choice.finish_reason:
|
||||
final_chunk = chunk
|
||||
break
|
||||
if hasattr(choice, "finish_reason") and choice.finish_reason == "tool_calls":
|
||||
initial_final_chunk = chunk
|
||||
|
||||
# If no chunk with finish_reason, use the last chunk
|
||||
if final_chunk is None and all_chunks:
|
||||
final_chunk = all_chunks[-1]
|
||||
assert first_chunk is not None, "Should have a first chunk"
|
||||
assert initial_final_chunk is not None, "Should have a final chunk from initial response"
|
||||
|
||||
assert final_chunk is not None, "Should have a final chunk"
|
||||
# print(first_chunk)
|
||||
# Verify mcp_list_tools is in the first chunk
|
||||
if hasattr(first_chunk, "choices") and first_chunk.choices:
|
||||
choice = first_chunk.choices[0]
|
||||
if hasattr(choice, "delta") and choice.delta:
|
||||
provider_fields = getattr(choice.delta, "provider_specific_fields", None)
|
||||
assert provider_fields is not None, "First chunk should have provider_specific_fields"
|
||||
assert "mcp_list_tools" in provider_fields, "First chunk should have mcp_list_tools"
|
||||
|
||||
# Verify MCP metadata is in the final chunk's delta.provider_specific_fields
|
||||
if hasattr(final_chunk, "choices") and final_chunk.choices:
|
||||
choice = final_chunk.choices[0]
|
||||
# Verify mcp_tool_calls and mcp_call_results are in the final chunk of initial response
|
||||
if hasattr(initial_final_chunk, "choices") and initial_final_chunk.choices:
|
||||
choice = initial_final_chunk.choices[0]
|
||||
if hasattr(choice, "delta") and choice.delta:
|
||||
provider_fields = getattr(choice.delta, "provider_specific_fields", None)
|
||||
assert provider_fields is not None, "Final chunk should have provider_specific_fields"
|
||||
assert "mcp_list_tools" in provider_fields, "Should have mcp_list_tools"
|
||||
assert provider_fields["mcp_list_tools"] == openai_tools
|
||||
assert "mcp_tool_calls" in provider_fields, "Should have mcp_tool_calls"
|
||||
assert "mcp_call_results" in provider_fields, "Should have mcp_call_results"
|
||||
|
||||
Reference in New Issue
Block a user