mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
Address MCP chat review follow-ups
Expose the non-streaming MCP auto-execution round limit as litellm.max_mcp_auto_execution_rounds with an env-backed default. Keep MCP metadata consistent by only aggregating tool calls after execution returns results for that round. Add coverage for no-tool-call responses, empty execution results, and the configurable round-limit path.
This commit is contained in:
nvr
parent
6df35fa1a6
commit
8d43894b25
3 changed files with 350 additions and 4 deletions
|
|
@ -213,6 +213,9 @@ token: Optional[
|
|||
] = None # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648
|
||||
telemetry = True
|
||||
max_tokens: int = DEFAULT_MAX_TOKENS # OpenAI Defaults
|
||||
max_mcp_auto_execution_rounds: int = int(
|
||||
os.getenv("LITELLM_MAX_MCP_AUTO_EXECUTION_ROUNDS", 6)
|
||||
)
|
||||
drop_params = bool(os.getenv("LITELLM_DROP_PARAMS", False))
|
||||
modify_params = bool(os.getenv("LITELLM_MODIFY_PARAMS", False))
|
||||
use_chat_completions_url_for_anthropic_messages: bool = bool(
|
||||
|
|
|
|||
|
|
@ -15,7 +15,21 @@ from litellm.responses.utils import ResponsesAPIRequestUtils
|
|||
from litellm.types.utils import ModelResponse
|
||||
from litellm.utils import CustomStreamWrapper
|
||||
|
||||
MAX_CHAT_MCP_AUTO_EXECUTION_ROUNDS = 6
|
||||
DEFAULT_CHAT_MCP_AUTO_EXECUTION_ROUNDS = 6
|
||||
|
||||
|
||||
def _get_max_chat_mcp_auto_execution_rounds() -> int:
|
||||
import litellm
|
||||
|
||||
configured_rounds = getattr(
|
||||
litellm,
|
||||
"max_mcp_auto_execution_rounds",
|
||||
DEFAULT_CHAT_MCP_AUTO_EXECUTION_ROUNDS,
|
||||
)
|
||||
try:
|
||||
return max(1, int(configured_rounds))
|
||||
except (TypeError, ValueError):
|
||||
return DEFAULT_CHAT_MCP_AUTO_EXECUTION_ROUNDS
|
||||
|
||||
|
||||
def _add_mcp_metadata_to_response(
|
||||
|
|
@ -615,7 +629,7 @@ async def acompletion_with_mcp( # noqa: PLR0915
|
|||
aggregated_tool_calls: List[Any] = []
|
||||
aggregated_tool_results: List[Any] = []
|
||||
|
||||
for _round_idx in range(MAX_CHAT_MCP_AUTO_EXECUTION_ROUNDS):
|
||||
for _round_idx in range(_get_max_chat_mcp_auto_execution_rounds()):
|
||||
if not isinstance(response, ModelResponse):
|
||||
return response
|
||||
|
||||
|
|
@ -632,8 +646,6 @@ async def acompletion_with_mcp( # noqa: PLR0915
|
|||
)
|
||||
return response
|
||||
|
||||
aggregated_tool_calls.extend(tool_calls)
|
||||
|
||||
tool_results = await LiteLLM_Proxy_MCP_Handler._execute_tool_calls(
|
||||
tool_server_map=tool_server_map,
|
||||
tool_calls=tool_calls,
|
||||
|
|
@ -655,6 +667,7 @@ async def acompletion_with_mcp( # noqa: PLR0915
|
|||
)
|
||||
return response
|
||||
|
||||
aggregated_tool_calls.extend(tool_calls)
|
||||
aggregated_tool_results.extend(tool_results)
|
||||
current_messages = LiteLLM_Proxy_MCP_Handler._create_follow_up_messages_for_chat(
|
||||
original_messages=current_messages,
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
import pytest
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import litellm
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
from litellm.responses.mcp import chat_completions_handler
|
||||
|
|
@ -531,6 +532,335 @@ async def test_acompletion_with_mcp_non_streaming_auto_exec_supports_multiple_ro
|
|||
assert provider_fields["mcp_call_results"][1]["tool_call_id"] == "call-2"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acompletion_with_mcp_non_streaming_returns_response_when_no_tool_calls(
|
||||
monkeypatch,
|
||||
):
|
||||
tools = [
|
||||
{
|
||||
"type": "mcp",
|
||||
"server_url": "litellm_proxy/mcp/local",
|
||||
"require_approval": "never",
|
||||
}
|
||||
]
|
||||
openai_tools = [{"type": "function", "function": {"name": "local_search"}}]
|
||||
response_without_tool_calls = ModelResponse(
|
||||
id="resp-1",
|
||||
model="test-model",
|
||||
created=1234567890,
|
||||
object="chat.completion",
|
||||
choices=[
|
||||
{
|
||||
"index": 0,
|
||||
"finish_reason": "stop",
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": "final answer",
|
||||
},
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
mock_acompletion = AsyncMock(return_value=response_without_tool_calls)
|
||||
|
||||
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(
|
||||
ResponsesAPIRequestUtils,
|
||||
"extract_mcp_headers_from_request",
|
||||
staticmethod(lambda **_: (None, None, None, None)),
|
||||
)
|
||||
|
||||
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=False,
|
||||
tool_choice="required",
|
||||
)
|
||||
|
||||
assert isinstance(result, ModelResponse)
|
||||
provider_fields = result.choices[0].message.provider_specific_fields
|
||||
assert provider_fields is not None
|
||||
assert provider_fields["mcp_list_tools"] == openai_tools
|
||||
assert "mcp_tool_calls" not in provider_fields
|
||||
assert "mcp_call_results" not in provider_fields
|
||||
mock_acompletion.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acompletion_with_mcp_non_streaming_omits_unexecuted_tool_calls_from_metadata(
|
||||
monkeypatch,
|
||||
):
|
||||
tools = [
|
||||
{
|
||||
"type": "mcp",
|
||||
"server_url": "litellm_proxy/mcp/local",
|
||||
"require_approval": "never",
|
||||
}
|
||||
]
|
||||
openai_tools = [{"type": "function", "function": {"name": "local_search"}}]
|
||||
tool_call = {
|
||||
"id": "call-1",
|
||||
"type": "function",
|
||||
"function": {"name": "local_search", "arguments": "{}"},
|
||||
}
|
||||
response_with_tool_call = ModelResponse(
|
||||
id="resp-1",
|
||||
model="test-model",
|
||||
created=1234567890,
|
||||
object="chat.completion",
|
||||
choices=[
|
||||
{
|
||||
"index": 0,
|
||||
"finish_reason": "tool_calls",
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"tool_calls": [tool_call],
|
||||
},
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
mock_acompletion = AsyncMock(return_value=response_with_tool_call)
|
||||
|
||||
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),
|
||||
)
|
||||
|
||||
async def mock_execute(**_):
|
||||
return []
|
||||
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"_execute_tool_calls",
|
||||
mock_execute,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
ResponsesAPIRequestUtils,
|
||||
"extract_mcp_headers_from_request",
|
||||
staticmethod(lambda **_: (None, None, None, None)),
|
||||
)
|
||||
|
||||
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=False,
|
||||
tool_choice="required",
|
||||
)
|
||||
|
||||
assert isinstance(result, ModelResponse)
|
||||
provider_fields = result.choices[0].message.provider_specific_fields
|
||||
assert provider_fields is not None
|
||||
assert provider_fields["mcp_list_tools"] == openai_tools
|
||||
assert "mcp_tool_calls" not in provider_fields
|
||||
assert "mcp_call_results" not in provider_fields
|
||||
mock_acompletion.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acompletion_with_mcp_non_streaming_uses_configured_round_limit(
|
||||
monkeypatch,
|
||||
):
|
||||
tools = [
|
||||
{
|
||||
"type": "mcp",
|
||||
"server_url": "litellm_proxy/mcp/local",
|
||||
"require_approval": "never",
|
||||
}
|
||||
]
|
||||
openai_tools = [{"type": "function", "function": {"name": "local_search"}}]
|
||||
first_tool_call = {
|
||||
"id": "call-1",
|
||||
"type": "function",
|
||||
"function": {"name": "local_search", "arguments": "{}"},
|
||||
}
|
||||
second_tool_call = {
|
||||
"id": "call-2",
|
||||
"type": "function",
|
||||
"function": {"name": "local_search", "arguments": '{"phase":"second"}'},
|
||||
}
|
||||
first_response = ModelResponse(
|
||||
id="resp-1",
|
||||
model="test-model",
|
||||
created=1234567890,
|
||||
object="chat.completion",
|
||||
choices=[
|
||||
{
|
||||
"index": 0,
|
||||
"finish_reason": "tool_calls",
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"tool_calls": [first_tool_call],
|
||||
},
|
||||
}
|
||||
],
|
||||
)
|
||||
second_response = ModelResponse(
|
||||
id="resp-2",
|
||||
model="test-model",
|
||||
created=1234567891,
|
||||
object="chat.completion",
|
||||
choices=[
|
||||
{
|
||||
"index": 0,
|
||||
"finish_reason": "tool_calls",
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"tool_calls": [second_tool_call],
|
||||
},
|
||||
}
|
||||
],
|
||||
)
|
||||
mock_acompletion = AsyncMock(side_effect=[first_response, second_response])
|
||||
|
||||
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),
|
||||
)
|
||||
|
||||
async def mock_execute(**kwargs):
|
||||
tool_id = kwargs["tool_calls"][0]["id"]
|
||||
return [
|
||||
{
|
||||
"tool_call_id": tool_id,
|
||||
"name": "local_search",
|
||||
"result": f"result-for-{tool_id}",
|
||||
}
|
||||
]
|
||||
|
||||
monkeypatch.setattr(
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
"_execute_tool_calls",
|
||||
mock_execute,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
ResponsesAPIRequestUtils,
|
||||
"extract_mcp_headers_from_request",
|
||||
staticmethod(lambda **_: (None, None, None, None)),
|
||||
)
|
||||
monkeypatch.setattr(litellm, "max_mcp_auto_execution_rounds", 1)
|
||||
|
||||
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=False,
|
||||
tool_choice="required",
|
||||
)
|
||||
|
||||
assert isinstance(result, ModelResponse)
|
||||
assert result.choices[0].message.tool_calls is not None
|
||||
assert result.choices[0].message.tool_calls[0].id == "call-2"
|
||||
assert mock_acompletion.await_count == 2
|
||||
provider_fields = result.choices[0].message.provider_specific_fields
|
||||
assert provider_fields is not None
|
||||
assert len(provider_fields["mcp_tool_calls"]) == 1
|
||||
assert len(provider_fields["mcp_call_results"]) == 1
|
||||
assert provider_fields["mcp_tool_calls"][0]["id"] == "call-1"
|
||||
assert provider_fields["mcp_call_results"][0]["tool_call_id"] == "call-1"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acompletion_with_mcp_adds_metadata_to_streaming(monkeypatch):
|
||||
"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue