test: add mcp completions test
This commit is contained in:
Yuta Saito
parent
ba43f742ab
commit
1c2942d808
@ -2,203 +2,18 @@
|
||||
|
||||
from typing import (
|
||||
Any,
|
||||
Awaitable,
|
||||
Callable,
|
||||
Dict,
|
||||
Iterable,
|
||||
List,
|
||||
Optional,
|
||||
Union,
|
||||
cast,
|
||||
)
|
||||
|
||||
from litellm.responses.mcp.litellm_proxy_mcp_handler import (
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
)
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
from litellm.types.llms.openai import ToolParam
|
||||
from litellm.types.utils import ModelResponse
|
||||
from litellm.utils import CustomStreamWrapper
|
||||
|
||||
CompletionCallable = Callable[..., Awaitable[Union[ModelResponse, CustomStreamWrapper]]]
|
||||
|
||||
_CHAT_COMPLETION_CALL_ARG_KEYS = [
|
||||
"model",
|
||||
"messages",
|
||||
"functions",
|
||||
"function_call",
|
||||
"timeout",
|
||||
"temperature",
|
||||
"top_p",
|
||||
"n",
|
||||
"stream",
|
||||
"stream_options",
|
||||
"stop",
|
||||
"max_tokens",
|
||||
"max_completion_tokens",
|
||||
"modalities",
|
||||
"prediction",
|
||||
"audio",
|
||||
"presence_penalty",
|
||||
"frequency_penalty",
|
||||
"logit_bias",
|
||||
"user",
|
||||
"response_format",
|
||||
"seed",
|
||||
"tools",
|
||||
"tool_choice",
|
||||
"parallel_tool_calls",
|
||||
"logprobs",
|
||||
"top_logprobs",
|
||||
"deployment_id",
|
||||
"reasoning_effort",
|
||||
"verbosity",
|
||||
"safety_identifier",
|
||||
"service_tier",
|
||||
"base_url",
|
||||
"api_version",
|
||||
"api_key",
|
||||
"model_list",
|
||||
"extra_headers",
|
||||
"thinking",
|
||||
"web_search_options",
|
||||
"shared_session",
|
||||
]
|
||||
|
||||
|
||||
def _build_call_args_from_context(call_context: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""Build kwargs for `acompletion` from the `completion` call context."""
|
||||
|
||||
call_args = {
|
||||
key: call_context.get(key)
|
||||
for key in _CHAT_COMPLETION_CALL_ARG_KEYS
|
||||
if key in call_context
|
||||
}
|
||||
additional_kwargs = dict(call_context.get("kwargs") or {})
|
||||
call_args.update(additional_kwargs)
|
||||
return call_args
|
||||
|
||||
|
||||
async def _call_acompletion_internal(
|
||||
completion_callable: CompletionCallable, **call_args: Any
|
||||
) -> Union[ModelResponse, CustomStreamWrapper]:
|
||||
"""Invoke `acompletion` while skipping MCP interception to avoid recursion."""
|
||||
|
||||
safe_args = dict(call_args)
|
||||
safe_args["_skip_mcp_handler"] = True
|
||||
safe_args.pop("acompletion", None)
|
||||
return await completion_callable(**safe_args)
|
||||
|
||||
|
||||
async def handle_chat_completion_with_mcp(
|
||||
call_context: Dict[str, Any],
|
||||
completion_callable: CompletionCallable,
|
||||
) -> Optional[Union[ModelResponse, CustomStreamWrapper]]:
|
||||
"""Handle MCP-enabled tool execution for chat completion requests."""
|
||||
|
||||
call_args = _build_call_args_from_context(call_context)
|
||||
|
||||
tools = call_args.get("tools")
|
||||
if not tools:
|
||||
return None
|
||||
|
||||
tools_for_mcp = cast(Optional[Iterable[ToolParam]], tools)
|
||||
|
||||
if not LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(
|
||||
tools=tools_for_mcp
|
||||
):
|
||||
return None
|
||||
|
||||
mcp_tools, _ = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools)
|
||||
if not mcp_tools:
|
||||
return None
|
||||
|
||||
base_call_args = dict(call_args)
|
||||
|
||||
user_api_key_auth = call_args.get("user_api_key_auth") or (
|
||||
(call_args.get("metadata", {}) or {}).get("user_api_key_auth")
|
||||
)
|
||||
(
|
||||
deduplicated_mcp_tools,
|
||||
tool_server_map,
|
||||
) = await LiteLLM_Proxy_MCP_Handler._process_mcp_tools_without_openai_transform(
|
||||
user_api_key_auth=user_api_key_auth,
|
||||
mcp_tools_with_litellm_proxy=mcp_tools,
|
||||
)
|
||||
|
||||
openai_tools = LiteLLM_Proxy_MCP_Handler._transform_mcp_tools_to_openai(
|
||||
deduplicated_mcp_tools,
|
||||
target_format="chat",
|
||||
)
|
||||
|
||||
base_call_args["tools"] = openai_tools or None
|
||||
|
||||
should_auto_execute = LiteLLM_Proxy_MCP_Handler._should_auto_execute_tools(
|
||||
mcp_tools_with_litellm_proxy=mcp_tools
|
||||
)
|
||||
|
||||
(
|
||||
mcp_auth_header,
|
||||
mcp_server_auth_headers,
|
||||
oauth2_headers,
|
||||
raw_headers,
|
||||
) = ResponsesAPIRequestUtils.extract_mcp_headers_from_request(
|
||||
secret_fields=base_call_args.get("secret_fields"),
|
||||
tools=tools,
|
||||
)
|
||||
|
||||
if not should_auto_execute:
|
||||
return await _call_acompletion_internal(completion_callable, **base_call_args)
|
||||
|
||||
mock_tool_calls = base_call_args.pop("mock_tool_calls", None)
|
||||
|
||||
initial_call_args = dict(base_call_args)
|
||||
initial_call_args["stream"] = False
|
||||
if mock_tool_calls is not None:
|
||||
initial_call_args["mock_tool_calls"] = mock_tool_calls
|
||||
|
||||
initial_response = await _call_acompletion_internal(
|
||||
completion_callable, **initial_call_args
|
||||
)
|
||||
if not isinstance(initial_response, ModelResponse):
|
||||
return initial_response
|
||||
|
||||
tool_calls = LiteLLM_Proxy_MCP_Handler._extract_tool_calls_from_chat_response(
|
||||
response=initial_response
|
||||
)
|
||||
|
||||
if not tool_calls:
|
||||
if base_call_args.get("stream"):
|
||||
retry_args = dict(base_call_args)
|
||||
retry_args["stream"] = call_args.get("stream")
|
||||
return await _call_acompletion_internal(completion_callable, **retry_args)
|
||||
return initial_response
|
||||
|
||||
tool_results = await LiteLLM_Proxy_MCP_Handler._execute_tool_calls(
|
||||
tool_server_map=tool_server_map,
|
||||
tool_calls=tool_calls,
|
||||
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,
|
||||
)
|
||||
|
||||
if not tool_results:
|
||||
return initial_response
|
||||
|
||||
follow_up_messages = LiteLLM_Proxy_MCP_Handler._create_follow_up_messages_for_chat(
|
||||
original_messages=call_args.get("messages", []),
|
||||
response=initial_response,
|
||||
tool_results=tool_results,
|
||||
)
|
||||
|
||||
follow_up_call_args = dict(base_call_args)
|
||||
follow_up_call_args["messages"] = follow_up_messages
|
||||
follow_up_call_args["stream"] = call_args.get("stream")
|
||||
|
||||
return await _call_acompletion_internal(completion_callable, **follow_up_call_args)
|
||||
|
||||
|
||||
async def acompletion_with_mcp(
|
||||
model: str,
|
||||
|
||||
@ -141,3 +141,174 @@ async def test_acompletion_mcp_respects_manual_approval(monkeypatch):
|
||||
assert isinstance(response, ModelResponse)
|
||||
tool_calls = response.choices[0].message.tool_calls
|
||||
assert tool_calls is not None and len(tool_calls) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_completion_mcp_with_streaming_no_timeout_error(monkeypatch):
|
||||
"""
|
||||
Test that litellm.completion with stream=True and MCP tools does not raise
|
||||
RuntimeError: Timeout context manager should be used inside a task.
|
||||
|
||||
This test ensures that the fix in ba43f742ab86d51b7da63077b85b39d0ac808d30
|
||||
prevents event loop nesting issues when using MCP tools with streaming.
|
||||
|
||||
The fix changes completion() to return a coroutine from acompletion_with_mcp,
|
||||
which acompletion() then awaits, avoiding event loop nesting.
|
||||
"""
|
||||
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
|
||||
|
||||
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):
|
||||
fake_execute.called = True # type: ignore[attr-defined]
|
||||
tool_calls = kwargs.get("tool_calls") or []
|
||||
assert tool_calls, "tool calls should be present during auto execution"
|
||||
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"),
|
||||
}
|
||||
]
|
||||
|
||||
fake_execute.called = False # type: ignore[attr-defined]
|
||||
|
||||
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 a mock streaming response
|
||||
class MockStreamingResponse(CustomStreamWrapper):
|
||||
def __init__(self):
|
||||
self.chunks = [
|
||||
type('Chunk', (), {
|
||||
'choices': [type('Choice', (), {
|
||||
'delta': type('Delta', (), {
|
||||
'content': 'Final'
|
||||
})()
|
||||
})()]
|
||||
})(),
|
||||
type('Chunk', (), {
|
||||
'choices': [type('Choice', (), {
|
||||
'delta': type('Delta', (), {
|
||||
'content': ' answer'
|
||||
})()
|
||||
})()]
|
||||
})(),
|
||||
]
|
||||
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
|
||||
return chunk
|
||||
raise StopIteration
|
||||
|
||||
# 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 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",
|
||||
)
|
||||
# Second call (streaming follow-up)
|
||||
return MockStreamingResponse()
|
||||
|
||||
with patch("litellm.acompletion", side_effect=mock_acompletion):
|
||||
# This should not raise RuntimeError: Timeout context manager should be used inside a task
|
||||
# completion() returns a coroutine when MCP tools are present, which acompletion() awaits
|
||||
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": "{}"},
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
# completion() returns a coroutine when MCP tools are present
|
||||
import asyncio
|
||||
assert asyncio.iscoroutine(response), "completion() should return a coroutine when MCP tools are present"
|
||||
|
||||
# Await the coroutine (this is what acompletion() does internally)
|
||||
# This should not raise RuntimeError: Timeout context manager should be used inside a task
|
||||
result = await response
|
||||
|
||||
# Verify response is a streaming response
|
||||
assert isinstance(result, CustomStreamWrapper) or hasattr(result, '__iter__')
|
||||
|
||||
# Consume the stream to ensure it works
|
||||
chunks = list(result)
|
||||
assert len(chunks) > 0, "Should have received streaming chunks"
|
||||
|
||||
# Verify tool execution was called
|
||||
assert fake_execute.called is True # type: ignore[attr-defined]
|
||||
|
||||
# Verify acompletion was called (should be called by acompletion_with_mcp)
|
||||
assert len(acompletion_calls) >= 1, "acompletion should be called"
|
||||
|
||||
Loading…
Reference in New Issue
Block a user