From e56c060363c741b24579ecc01279c992d2d339d1 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 4 Nov 2026 20:52:44 -0800 Subject: [PATCH 1/9] fix(tool_choice): accept the Responses API's flat allowed_tools layout on chat completions An allowed_tools tool_choice in the flat Responses layout ({"type": "allowed_tools", "mode": ..., "tools": [...]}) answered 400 on litellm.completion, acompletion, Router and the proxy's chat and /v1/responses-to-chat routes. One parser now reads both layouts into the nested Chat layout, and the existing Responses bridge flattens it back for /v1/responses, so either layout reaches OpenAI in the shape each endpoint expects --- .../transformation.py | 4 +- .../allowed_tools_tool_choice.py | 64 ++++++++++++++++--- litellm/utils.py | 37 ++++++----- ...responses_transformation_transformation.py | 4 ++ .../test_allowed_tools_tool_choice.py | 52 +++++++++++++-- .../test_responses_api_bridge_flag.py | 38 +++++++++++ tests/unit/test_main.py | 10 ++- tests/unit/test_utils.py | 15 +++-- 8 files changed, 184 insertions(+), 40 deletions(-) diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 8e6341805db..854d06d0fec 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -36,7 +36,7 @@ from litellm.completion_extras.litellm_responses_transformation.history import ( ) from litellm.integrations.anthropic_cache_control_hook import routed_model_info, supports_openai_prompt_cache_breakpoint from litellm.litellm_core_utils.allowed_tools_tool_choice import ( - parse_chat_allowed_tools_tool_choice, + parse_allowed_tools_tool_choice, responses_allowed_tools_tool_choice, ) from litellm.litellm_core_utils.core_helpers import drops_unsupported_params @@ -802,7 +802,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): if isinstance(choice_type, str) and choice_type in ("auto", "none", "required"): return choice_type if choice_type == "allowed_tools": - parsed: Final = parse_chat_allowed_tools_tool_choice( + parsed: Final = parse_allowed_tools_tool_choice( cast( # cast-ok: widening an untyped dict to object is always sound, the helper validates it with pydantic object, tool_choice ) diff --git a/litellm/litellm_core_utils/allowed_tools_tool_choice.py b/litellm/litellm_core_utils/allowed_tools_tool_choice.py index 5c8f5d83ece..55656977fa7 100644 --- a/litellm/litellm_core_utils/allowed_tools_tool_choice.py +++ b/litellm/litellm_core_utils/allowed_tools_tool_choice.py @@ -1,6 +1,8 @@ from collections.abc import Mapping from typing import Final, Literal +from openai.types.chat.chat_completion_allowed_tool_choice_param import ChatCompletionAllowedToolChoiceParam +from openai.types.chat.chat_completion_allowed_tools_param import ChatCompletionAllowedToolsParam from openai.types.responses.tool_choice_allowed_param import ToolChoiceAllowedParam from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError @@ -11,6 +13,9 @@ ALLOWED_TOOLS_TOOL_CHOICE_PROVIDERS: Final = frozenset({"openai", "azure", "lite CHAT_ALLOWED_TOOLS_LAYOUT: Final = ( '{"type": "allowed_tools", "allowed_tools": {"mode": "auto" | "required", "tools": []}}' ) +RESPONSES_ALLOWED_TOOLS_LAYOUT: Final = ( + '{"type": "allowed_tools", "mode": "auto" | "required", "tools": []}' +) class ChatAllowedTools(BaseModel): @@ -27,7 +32,17 @@ class ChatAllowedToolsToolChoice(BaseModel): allowed_tools: ChatAllowedTools -_CHAT_ALLOWED_TOOLS_TOOL_CHOICE: Final = TypeAdapter(ChatAllowedToolsToolChoice) +class ResponsesAllowedToolsToolChoice(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + type: Literal["allowed_tools"] + mode: Literal["auto", "required"] + tools: tuple[Mapping[str, object], ...] = Field(min_length=1) + + +_ALLOWED_TOOLS_TOOL_CHOICE: Final[TypeAdapter[ChatAllowedToolsToolChoice | ResponsesAllowedToolsToolChoice]] = ( + TypeAdapter(ChatAllowedToolsToolChoice | ResponsesAllowedToolsToolChoice) +) _SELECTOR: Final = TypeAdapter(dict[str, object]) @@ -70,18 +85,51 @@ def raise_if_allowed_tools_tool_choice_unsupported( ) -def parse_chat_allowed_tools_tool_choice(tool_choice: object) -> ChatAllowedToolsToolChoice | None: - try: - return _CHAT_ALLOWED_TOOLS_TOOL_CHOICE.validate_python(tool_choice) - except ValidationError: - return None - - def _has_name(selector: Mapping[str, object]) -> bool: name: Final = selector.get("name") return isinstance(name, str) and name != "" +def _chat_allowed_tool(tool: Mapping[str, object]) -> Mapping[str, object] | None: + tool_type: Final = tool.get("type") + if not isinstance(tool_type, str) or tool_type not in ("function", "custom") or tool_type in tool: + return tool + if not _has_name(tool): + return None + return {**{key: value for key, value in tool.items() if key != "name"}, tool_type: {"name": tool["name"]}} + + +def _chat_layout(choice: ResponsesAllowedToolsToolChoice) -> ChatAllowedToolsToolChoice | None: + nested: Final = tuple(_chat_allowed_tool(tool) for tool in choice.tools) + if any(tool is None for tool in nested): + return None + return ChatAllowedToolsToolChoice( + type="allowed_tools", + allowed_tools=ChatAllowedTools(mode=choice.mode, tools=tuple(tool for tool in nested if tool is not None)), + ) + + +def parse_allowed_tools_tool_choice(tool_choice: object) -> ChatAllowedToolsToolChoice | None: + try: + parsed: Final = _ALLOWED_TOOLS_TOOL_CHOICE.validate_python(tool_choice) + except ValidationError: + return None + match parsed: + case ChatAllowedToolsToolChoice(): + return parsed + case ResponsesAllowedToolsToolChoice(): + return _chat_layout(parsed) + + +def chat_allowed_tools_tool_choice(choice: ChatAllowedToolsToolChoice) -> ChatCompletionAllowedToolChoiceParam: + return ChatCompletionAllowedToolChoiceParam( + type="allowed_tools", + allowed_tools=ChatCompletionAllowedToolsParam( + mode=choice.allowed_tools.mode, tools=[dict(tool) for tool in choice.allowed_tools.tools] + ), + ) + + def _responses_allowed_tool(tool: Mapping[str, object]) -> Mapping[str, object] | None: tool_type: Final = tool.get("type") if tool_type not in ("function", "custom"): diff --git a/litellm/utils.py b/litellm/utils.py index 7f81ece8314..dabc28d4c66 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -86,7 +86,9 @@ from litellm.constants import ( from litellm.litellm_core_utils.agentic_followup_failure import is_agentic_followup_failure from litellm.litellm_core_utils.allowed_tools_tool_choice import ( CHAT_ALLOWED_TOOLS_LAYOUT, - parse_chat_allowed_tools_tool_choice, + RESPONSES_ALLOWED_TOOLS_LAYOUT, + chat_allowed_tools_tool_choice, + parse_allowed_tools_tool_choice, ) from litellm.litellm_core_utils.core_helpers import ( bind_budget_reservation_to_callbacks, @@ -9665,28 +9667,25 @@ def validate_chat_completion_tool_choice( tool_choice_type = tool_choice.get("type") if tool_choice_type in ("auto", "none", "required") and "function" not in tool_choice: return tool_choice_type - if ( - tool_choice_type == "allowed_tools" - and ( - parse_chat_allowed_tools_tool_choice( - cast( # cast-ok: widening an untyped dict to object is always sound, the helper validates it with pydantic - object, tool_choice - ) + if tool_choice_type == "allowed_tools": + allowed_tools: Final = parse_allowed_tools_tool_choice( + cast( # cast-ok: widening an untyped dict to object is always sound, the helper validates it with pydantic + object, tool_choice ) - is None - ) - ): - raise BadRequestError( - message=( - f"Invalid tool choice, tool_choice={tool_choice}. On chat completions an allowed_tools " - f"tool_choice must be {CHAT_ALLOWED_TOOLS_LAYOUT}" - ), - model=model, - llm_provider="", ) + if allowed_tools is None: + raise BadRequestError( + message=( + f"Invalid tool choice, tool_choice={tool_choice}. An allowed_tools tool_choice must be " + f"{CHAT_ALLOWED_TOOLS_LAYOUT} or {RESPONSES_ALLOWED_TOOLS_LAYOUT}" + ), + model=model, + llm_provider="", + ) + return dict(chat_allowed_tools_tool_choice(allowed_tools)) # Standard OpenAI format: {"type": "function", "function": {...}} - if tool_choice_type != "allowed_tools" and (tool_choice_type is None or tool_choice.get("function") is None): + if tool_choice_type is None or tool_choice.get("function") is None: raise BadRequestError( message=f"Invalid tool choice, tool_choice={tool_choice}. Please ensure tool_choice follows the OpenAI spec", model=model, diff --git a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 19aefb430c7..f632b90ee42 100644 --- a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -4558,6 +4558,10 @@ def test_output_item_done_with_stream_map_keeps_empty_delta(): }, {"type": "allowed_tools", "mode": "required", "tools": [{"type": "function", "name": "get_weather"}]}, ), + ( + {"type": "allowed_tools", "mode": "required", "tools": [{"type": "function", "name": "get_weather"}]}, + {"type": "allowed_tools", "mode": "required", "tools": [{"type": "function", "name": "get_weather"}]}, + ), ], ) async def test_acompletion_bridge_normalizes_tool_choice_on_the_wire( diff --git a/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py b/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py index 6ebbdcf73bf..2353308ab98 100644 --- a/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py +++ b/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py @@ -7,7 +7,8 @@ import pytest from litellm.litellm_core_utils.allowed_tools_tool_choice import ( ChatAllowedToolsToolChoice, is_allowed_tools_tool_choice, - parse_chat_allowed_tools_tool_choice, + chat_allowed_tools_tool_choice, + parse_allowed_tools_tool_choice, responses_allowed_tools_tool_choice, ) @@ -26,7 +27,7 @@ _CHAT_LAYOUT: Final[dict[str, object]] = { def _parsed(tool_choice: Mapping[str, object]) -> ChatAllowedToolsToolChoice: - parsed: Final = parse_chat_allowed_tools_tool_choice(tool_choice) + parsed: Final = parse_allowed_tools_tool_choice(tool_choice) assert parsed is not None return parsed @@ -72,18 +73,59 @@ def test_responses_layout_refuses_to_rewrite_a_malformed_or_conflicting_entry(to assert responses_allowed_tools_tool_choice(_parsed(tool_choice)) is None +_RESPONSES_LAYOUT: Final[dict[str, object]] = { + "type": "allowed_tools", + "mode": "required", + "tools": [ + {"type": "function", "name": "get_time"}, + {"type": "custom", "name": "apply_patch", "x-vendor": 1}, + {"type": "mcp", "server_label": "docs"}, + {"type": "function", "function": {"name": "already_nested"}}, + ], +} + + +def test_responses_layout_is_read_into_the_chat_layout(): + assert chat_allowed_tools_tool_choice(_parsed(_RESPONSES_LAYOUT)) == { + "type": "allowed_tools", + "allowed_tools": { + "mode": "required", + "tools": [ + {"type": "function", "function": {"name": "get_time"}}, + {"type": "custom", "x-vendor": 1, "custom": {"name": "apply_patch"}}, + {"type": "mcp", "server_label": "docs"}, + {"type": "function", "function": {"name": "already_nested"}}, + ], + }, + } + + +def test_responses_layout_round_trips_back_to_what_the_caller_sent(): + flat_only: Final = {**_RESPONSES_LAYOUT, "tools": list(_RESPONSES_LAYOUT["tools"])[:3]} + assert responses_allowed_tools_tool_choice(_parsed(flat_only)) == flat_only + + @pytest.mark.parametrize( "tool_choice", [ - {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "function", "name": "get_time"}]}, {"type": "allowed_tools", "allowed_tools": {"mode": "auto", "tools": []}}, {"type": "allowed_tools", "allowed_tools": {"mode": "any", "tools": [{"type": "function"}]}}, {"type": "allowed_tools", "allowed_tools": {"mode": "auto", "tools": [{"type": "function"}]}, "extra": 1}, + {"type": "allowed_tools", "mode": "auto", "tools": []}, + {"type": "allowed_tools", "mode": "none", "tools": [{"type": "function", "name": "get_time"}]}, + {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "function"}]}, + {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "custom", "name": ""}]}, + { + "type": "allowed_tools", + "mode": "auto", + "tools": [{"type": "function", "name": "get_time"}], + "allowed_tools": {"mode": "auto", "tools": [{"type": "function", "name": "get_time"}]}, + }, {"type": "function", "function": {"name": "get_time"}}, ], ) -def test_parse_rejects_anything_but_the_chat_layout(tool_choice: dict[str, object]): - assert parse_chat_allowed_tools_tool_choice(tool_choice) is None +def test_parse_rejects_a_malformed_choice_in_either_layout(tool_choice: dict[str, object]): + assert parse_allowed_tools_tool_choice(tool_choice) is None @pytest.mark.parametrize( diff --git a/tests/unit/responses/test_responses_api_bridge_flag.py b/tests/unit/responses/test_responses_api_bridge_flag.py index 210488f0888..4b7e29d2eb7 100644 --- a/tests/unit/responses/test_responses_api_bridge_flag.py +++ b/tests/unit/responses/test_responses_api_bridge_flag.py @@ -305,6 +305,44 @@ class TestUseResponsesApiBridgeFlag: assert request_body["messages"] == [{"role": "user", "content": "Hello"}] assert response.output[0].content[0].text == "Answer" + def test_bridge_sends_a_responses_layout_allowed_tools_tool_choice_in_the_chat_layout( + self, respx_mock: respx.MockRouter + ): + upstream: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + return_value=httpx.Response( + status_code=200, + json={ + "id": "chatcmpl-123", + "object": "chat.completion", + "created": 1677652288, + "model": "my-custom-model", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "Answer"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 9, "completion_tokens": 1, "total_tokens": 10}, + }, + ) + ) + + litellm.responses( + model="openai/my-custom-model", + input="What time is it?", + use_chat_completions_api=True, + tools=[ + {"type": "function", "name": name, "parameters": {"type": "object", "properties": {}}} + for name in ("get_time", "get_weather") + ], + tool_choice={"type": "allowed_tools", "mode": "required", "tools": [{"type": "function", "name": "get_time"}]}, + api_key="fake-provider-api-key", + num_retries=0, + ) + + assert upstream.call_count == 1 + assert json.loads(upstream.calls[0].request.read())["tool_choice"] == { + "type": "allowed_tools", + "allowed_tools": {"mode": "required", "tools": [{"type": "function", "function": {"name": "get_time"}}]}, + } + def test_bridge_merges_instructions_and_developer_input_for_databricks(self, respx_mock: respx.MockRouter): upstream: Final = respx_mock.post("https://example.databricks.test/serving-endpoints/chat/completions").mock( return_value=httpx.Response( diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 99c6f6c402c..afaab7d8501 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -557,15 +557,21 @@ _ALLOWED_TOOLS_TOOL_CHOICE: Final = { "type": "allowed_tools", "allowed_tools": {"mode": "required", "tools": [{"type": "function", "function": {"name": "get_time"}}]}, } +_RESPONSES_LAYOUT_ALLOWED_TOOLS_TOOL_CHOICE: Final = { + "type": "allowed_tools", + "mode": "required", + "tools": [{"type": "function", "name": "get_time"}], +} _TWO_TOOLS: Final = [ {"type": "function", "function": {"name": "get_time", "parameters": {"type": "object", "properties": {}}}}, {"type": "function", "function": {"name": "get_weather", "parameters": {"type": "object", "properties": {}}}}, ] +@pytest.mark.parametrize("tool_choice", [_ALLOWED_TOOLS_TOOL_CHOICE, _RESPONSES_LAYOUT_ALLOWED_TOOLS_TOOL_CHOICE]) @pytest.mark.parametrize("model", ["openai/gpt-5.5", "azure/my-gpt-deployment"]) def test_completion_forwards_allowed_tools_tool_choice_to_openai_compatible_chat( - respx_mock: respx.MockRouter, openai_api_response: dict[str, object], model: str + respx_mock: respx.MockRouter, openai_api_response: dict[str, object], model: str, tool_choice: dict[str, object] ): mock_route: Final = respx_mock.post(url__regex=r"http://localhost:12346/.*chat/completions.*").mock( return_value=httpx.Response(status_code=200, json=openai_api_response) @@ -575,7 +581,7 @@ def test_completion_forwards_allowed_tools_tool_choice_to_openai_compatible_chat model=model, messages=[{"role": "user", "content": "What time is it?"}], tools=_TWO_TOOLS, - tool_choice=_ALLOWED_TOOLS_TOOL_CHOICE, + tool_choice=tool_choice, api_base="http://localhost:12346/v1", api_key="fake_api_key", api_version="2025-04-01-preview", diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index e24499db071..f242fddc7ba 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -8715,22 +8715,29 @@ def test_validate_tool_choice_invalid_dict_is_a_400(tool_choice): @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @pytest.mark.parametrize("mode", ["auto", "required"]) -def test_validate_tool_choice_allowed_tools_chat_layout_passes_through(mode: str): - tool_choice: Final = { +def test_validate_tool_choice_allowed_tools_either_layout_comes_back_in_the_chat_layout(mode: str): + chat_layout: Final = { "type": "allowed_tools", "allowed_tools": { "mode": mode, "tools": [{"type": "function", "function": {"name": "get_time"}}, {"type": "mcp", "server_label": "docs"}], }, } - assert validate_chat_completion_tool_choice(tool_choice, model=MODEL) is tool_choice + responses_layout: Final = { + "type": "allowed_tools", + "mode": mode, + "tools": [{"type": "function", "name": "get_time"}, {"type": "mcp", "server_label": "docs"}], + } + assert validate_chat_completion_tool_choice(chat_layout, model=MODEL) == chat_layout + assert validate_chat_completion_tool_choice(responses_layout, model=MODEL) == chat_layout @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @pytest.mark.parametrize( "tool_choice", [ - {"type": "allowed_tools", "mode": "required", "tools": [{"type": "function", "name": "get_time"}]}, + {"type": "allowed_tools", "mode": "required", "tools": [{"type": "function"}]}, + {"type": "allowed_tools", "mode": "required", "tools": []}, {"type": "allowed_tools"}, {"type": "allowed_tools", "allowed_tools": {"mode": "required"}}, {"type": "allowed_tools", "allowed_tools": {"mode": "required", "tools": []}}, From 361aecae72f2a9153a12b3998c90515f679896fd Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 5 Nov 2026 10:26:36 -0800 Subject: [PATCH 2/9] fix(tool_choice): drop an allowed_tools choice under drop_params and refuse conflicting flat entries An allowed_tools tool_choice a provider can't enforce, or a malformed one, now drops under drop_params and answers 400 without it, the same rule LiteLLM uses for other unsupported values. That covers providers outside OpenAI and Azure, malformed choices on every provider, and Azure api versions older than 2023-12-01 A flat allowed_tools entry that carries both name and a nested function or custom key answers 400, since the two names can disagree, matching what flattening already refuses --- .../allowed_tools_tool_choice.py | 19 +++++---- litellm/llms/azure/chat/gpt_transformation.py | 14 ++----- litellm/main.py | 15 +++++-- litellm/utils.py | 3 ++ .../test_allowed_tools_tool_choice.py | 5 +++ .../test_azure_chat_gpt_transformation.py | 19 +++++++-- tests/unit/test_main.py | 39 +++++++++++++++++-- tests/unit/test_utils.py | 12 ++++++ 8 files changed, 95 insertions(+), 31 deletions(-) diff --git a/litellm/litellm_core_utils/allowed_tools_tool_choice.py b/litellm/litellm_core_utils/allowed_tools_tool_choice.py index 55656977fa7..276a13fd235 100644 --- a/litellm/litellm_core_utils/allowed_tools_tool_choice.py +++ b/litellm/litellm_core_utils/allowed_tools_tool_choice.py @@ -64,21 +64,22 @@ def is_allowed_tools_tool_choice(tool_choice: object) -> bool: return False -def raise_if_allowed_tools_tool_choice_unsupported( +def drops_unsupported_allowed_tools_tool_choice( tool_choice: object, custom_llm_provider: str, model: str, -) -> None: - """Providers outside the set would silently send the request with no tool restriction at all, so the call stops - here with a 400. ``drop_params`` is deliberately not consulted: dropping a restriction widens what the model may call.""" + drop_params: bool, +) -> bool: if not is_allowed_tools_tool_choice(tool_choice) or custom_llm_provider in ALLOWED_TOOLS_TOOL_CHOICE_PROVIDERS: - return + return False + if drop_params: + return True raise BadRequestError( message=( "tool_choice type 'allowed_tools' is only supported for " f"{', '.join(sorted(ALLOWED_TOOLS_TOOL_CHOICE_PROVIDERS))}, got custom_llm_provider={custom_llm_provider}. " - "drop_params does not drop it, since that would let the model call tools the request excluded. " - "Send only the allowed tools in `tools` with tool_choice='auto' or 'required' instead." + "Send only the allowed tools in `tools` with tool_choice='auto' or 'required', or set drop_params=True " + "to drop the tool_choice, which lets the model call any tool in `tools`." ), model=model, llm_provider=custom_llm_provider, @@ -92,8 +93,10 @@ def _has_name(selector: Mapping[str, object]) -> bool: def _chat_allowed_tool(tool: Mapping[str, object]) -> Mapping[str, object] | None: tool_type: Final = tool.get("type") - if not isinstance(tool_type, str) or tool_type not in ("function", "custom") or tool_type in tool: + if not isinstance(tool_type, str) or tool_type not in ("function", "custom"): return tool + if tool_type in tool: + return None if "name" in tool else tool if not _has_name(tool): return None return {**{key: value for key, value in tool.items() if key != "name"}, tool_type: {"name": tool["name"]}} diff --git a/litellm/llms/azure/chat/gpt_transformation.py b/litellm/llms/azure/chat/gpt_transformation.py index 0fe7caa1e16..e0aceba52f9 100644 --- a/litellm/llms/azure/chat/gpt_transformation.py +++ b/litellm/llms/azure/chat/gpt_transformation.py @@ -1,17 +1,11 @@ from collections.abc import Mapping, Sequence from types import MappingProxyType -from typing import ( - TYPE_CHECKING, - Any, - Final, - cast, # noqa: TID251 # non_default_params reaches the mapper as an untyped dict -) +from typing import TYPE_CHECKING, Any, Final from httpx._models import Headers, Response import litellm from litellm.integrations.anthropic_cache_control_hook import deployment_prompt_cache_breakpoint_support -from litellm.litellm_core_utils.allowed_tools_tool_choice import is_allowed_tools_tool_choice from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.litellm_core_utils.prompt_templates.common_utils import ( chat_completion_prompt_cache_breakpoints_for_wire, @@ -251,12 +245,10 @@ class AzureOpenAIConfig(BaseConfig): or (api_version_year == "2023" and api_version_month < "12") or (api_version_year == "2023" and api_version_month == "12" and api_version_day < "01") ): - if is_allowed_tools_tool_choice( - cast(object, value) # cast-ok: widening an untyped dict value to object is always sound - ) or not (litellm.drop_params is True or (drop_params is not None and drop_params is True)): + if not (litellm.drop_params is True or drop_params is True): raise UnsupportedParamsError( status_code=400, - message=f"""Azure does not support 'tool_choice', for api_version={api_version}. Bump your API version to '2023-12-01-preview' or later. This parameter requires 'api_version="2023-12-01-preview"' or later. drop_params never drops an allowed_tools tool_choice, since that would let the model call tools the request excluded. Azure API Reference: https://learn.microsoft.com/en-us/azure/ai-services/openai/reference#chat-completions""", + message=f"""Azure does not support 'tool_choice', for api_version={api_version}. Bump your API version to '2023-12-01-preview' or later. This parameter requires 'api_version="2023-12-01-preview"' or later. Azure API Reference: https://learn.microsoft.com/en-us/azure/ai-services/openai/reference#chat-completions""", ) elif value == "required" and ( api_version_year == "2024" and api_version_month <= "05" diff --git a/litellm/main.py b/litellm/main.py index 8cb133bf76b..71e149ecc14 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -73,7 +73,7 @@ from litellm.constants import ( from litellm.exceptions import LiteLLMUnknownProvider from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.allowed_tools_tool_choice import ( - raise_if_allowed_tools_tool_choice_unsupported, + drops_unsupported_allowed_tools_tool_choice, ) from litellm.litellm_core_utils.asyncify import asyncify, run_async_function from litellm.litellm_core_utils.audio_utils.utils import ( @@ -5584,7 +5584,13 @@ def completion( messages = validate_and_fix_openai_messages(messages=messages) tools = validate_and_fix_openai_tools(tools=tools) # validate tool_choice - tool_choice = validate_chat_completion_tool_choice(tool_choice=tool_choice, model=model) + tool_choice = validate_chat_completion_tool_choice( + tool_choice=tool_choice, + model=model, + drop_params=drops_unsupported_params( + cast("Mapping[str, object]", kwargs).get("drop_params") # cast-ok: kwargs is the untyped request dict + ), + ) # validate optional params stop = validate_openai_optional_params(stop=stop) thinking = validate_and_fix_thinking_param(thinking=thinking) @@ -5881,12 +5887,13 @@ def completion( if not is_claude_tool_target(custom_llm_provider=custom_llm_provider, model=model): tools = _drop_anthropic_only_tool_keys(tools=tools) - raise_if_allowed_tools_tool_choice_unsupported( + drops_tool_choice: Final = drops_unsupported_allowed_tools_tool_choice( tool_choice=cast( # cast-ok: widening an untyped dict to object is always sound, the helper validates it with pydantic object, tool_choice ), custom_llm_provider=custom_llm_provider, model=model, + drop_params=drop_unsupported_reasoning, ) if provider_specific_header is not None: @@ -6004,7 +6011,7 @@ def completion( "response_format": response_format, "seed": seed, "tools": tools, - "tool_choice": tool_choice, + "tool_choice": None if drops_tool_choice else tool_choice, "max_retries": max_retries, "logprobs": logprobs, "top_logprobs": top_logprobs, diff --git a/litellm/utils.py b/litellm/utils.py index dabc28d4c66..6a960a32a84 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -9655,6 +9655,7 @@ def validate_chat_completion_user_messages(messages: list[AllMessageValues]): def validate_chat_completion_tool_choice( tool_choice: dict | str | None, model: str = "", + drop_params: bool = False, ) -> dict | str | None: """ Confirm the tool choice is passed in the OpenAI format. @@ -9673,6 +9674,8 @@ def validate_chat_completion_tool_choice( object, tool_choice ) ) + if allowed_tools is None and drop_params: + return None if allowed_tools is None: raise BadRequestError( message=( diff --git a/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py b/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py index 2353308ab98..0308ae41abf 100644 --- a/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py +++ b/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py @@ -115,6 +115,11 @@ def test_responses_layout_round_trips_back_to_what_the_caller_sent(): {"type": "allowed_tools", "mode": "none", "tools": [{"type": "function", "name": "get_time"}]}, {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "function"}]}, {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "custom", "name": ""}]}, + { + "type": "allowed_tools", + "mode": "required", + "tools": [{"type": "function", "name": "get_time", "function": {"name": "get_weather"}}], + }, { "type": "allowed_tools", "mode": "auto", diff --git a/tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py b/tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py index f93968eb2e6..e672e42307a 100644 --- a/tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py +++ b/tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py @@ -92,10 +92,21 @@ def test_map_openai_params_drop_params_still_drops_a_named_tool_choice_on_old_ap assert "tool_choice" not in mapped -def test_map_openai_params_drop_params_never_drops_allowed_tools_on_old_api_version(): - with pytest.raises(litellm.UnsupportedParamsError, match="allowed_tools") as exc: - _map_tool_choice_on_pre_tool_choice_api_version( - {"type": "allowed_tools", "allowed_tools": {"mode": "auto", "tools": [{"type": "function"}]}} +def test_map_openai_params_drop_params_drops_allowed_tools_on_old_api_version(): + mapped: Final = _map_tool_choice_on_pre_tool_choice_api_version( + {"type": "allowed_tools", "allowed_tools": {"mode": "auto", "tools": [{"type": "function"}]}} + ) + assert "tool_choice" not in mapped + + +def test_map_openai_params_allowed_tools_on_old_api_version_is_a_400_without_drop_params(): + with pytest.raises(litellm.UnsupportedParamsError, match="2023-12-01-preview") as exc: + AzureOpenAIConfig().map_openai_params( + {"tool_choice": {"type": "allowed_tools", "allowed_tools": {"mode": "auto", "tools": [{"type": "function"}]}}}, + {}, + "gpt-4o", + drop_params=False, + api_version="2023-07-01-preview", ) assert exc.value.status_code == 400 diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index afaab7d8501..42a109fd7e1 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -593,10 +593,10 @@ def test_completion_forwards_allowed_tools_tool_choice_to_openai_compatible_chat assert [tool["function"]["name"] for tool in sent["tools"]] == ["get_time", "get_weather"] -@pytest.mark.parametrize("drop_params", [False, True]) +@pytest.mark.parametrize("tool_choice", [_ALLOWED_TOOLS_TOOL_CHOICE, _RESPONSES_LAYOUT_ALLOWED_TOOLS_TOOL_CHOICE]) @pytest.mark.parametrize("model", ["anthropic/claude-sonnet-4-5", "gemini/gemini-2.5-flash", "bedrock/converse/x"]) def test_completion_allowed_tools_tool_choice_is_a_400_where_the_provider_cannot_enforce_it( - respx_mock: respx.MockRouter, model: str, drop_params: bool + respx_mock: respx.MockRouter, model: str, tool_choice: dict[str, object] ): with pytest.raises( litellm.BadRequestError, match=r"allowed_tools.*only supported for azure, litellm_proxy, openai" @@ -605,8 +605,7 @@ def test_completion_allowed_tools_tool_choice_is_a_400_where_the_provider_cannot model=model, messages=[{"role": "user", "content": "What time is it?"}], tools=_TWO_TOOLS, - tool_choice=_ALLOWED_TOOLS_TOOL_CHOICE, - drop_params=drop_params, + tool_choice=tool_choice, api_key="fake_api_key", aws_access_key_id="fake", aws_secret_access_key="fake", @@ -616,6 +615,38 @@ def test_completion_allowed_tools_tool_choice_is_a_400_where_the_provider_cannot assert not respx_mock.calls +@pytest.mark.parametrize( + ("model", "tool_choice"), + [ + ("hosted_vllm/my-model", _ALLOWED_TOOLS_TOOL_CHOICE), + ("hosted_vllm/my-model", _RESPONSES_LAYOUT_ALLOWED_TOOLS_TOOL_CHOICE), + ("openai/gpt-5.5", {"type": "allowed_tools", "mode": "required", "tools": []}), + ("openai/gpt-5.5", {"type": "allowed_tools", "mode": "required", "tools": [{"type": "function"}]}), + ], +) +def test_completion_drop_params_drops_an_allowed_tools_tool_choice_it_cannot_send( + respx_mock: respx.MockRouter, openai_api_response: dict[str, object], model: str, tool_choice: dict[str, object] +): + mock_route: Final = respx_mock.post(url__regex=r"http://localhost:12346/.*chat/completions.*").mock( + return_value=httpx.Response(status_code=200, json=openai_api_response) + ) + + litellm.completion( + model=model, + messages=[{"role": "user", "content": "What time is it?"}], + tools=_TWO_TOOLS, + tool_choice=tool_choice, + drop_params=True, + api_base="http://localhost:12346/v1", + api_key="fake_api_key", + ) + + assert mock_route.called + sent: Final = json.loads(respx_mock.calls[0].request.content) + assert "tool_choice" not in sent + assert [tool["function"]["name"] for tool in sent["tools"]] == ["get_time", "get_weather"] + + _CLAUDE_TOOL_SEARCH_CALL: Final = { "id": "srvtoolu_01search", "type": "function", diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index f242fddc7ba..42fb73acadc 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -8754,6 +8754,18 @@ def test_validate_tool_choice_allowed_tools_malformed_is_a_400(tool_choice: dict assert '"mode": "auto" | "required"' in str(exc_info.value) +@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") +@pytest.mark.parametrize( + "tool_choice", + [ + {"type": "allowed_tools", "mode": "required", "tools": []}, + {"type": "allowed_tools", "allowed_tools": {"mode": "none", "tools": [{"type": "function"}]}}, + ], +) +def test_validate_tool_choice_allowed_tools_malformed_is_dropped_under_drop_params(tool_choice: dict[str, object]): + assert validate_chat_completion_tool_choice(tool_choice, model=MODEL, drop_params=True) is None + + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @pytest.mark.parametrize("tool_choice", [123, []]) def test_validate_tool_choice_invalid_type_is_a_400(tool_choice): From 5f8f54a6ca0dbd5bfc443b55e4c49e34b6b8176e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 5 Nov 2026 13:56:38 -0800 Subject: [PATCH 3/9] fix(tool-choice): keep the 400 when LiteLLM runs the tools, report auto after a drop Under drop_params an allowed_tools choice a provider can't enforce was dropped, so on MCP gateway tools and in run_tool_loop and arun_tool_loop LiteLLM could run a tool the allow list excluded. Those paths now keep the 400 they answered before the flat layout was accepted, and a malformed choice stays a 400 on MCP gateway tools. /v1/responses served through chat echoed a choice it had dropped, on sync, async and streamed responses. It now reports auto, and a choice the provider received is still echoed as sent. --- .../allowed_tools_tool_choice.py | 22 ++++ litellm/main.py | 13 +- .../handler.py | 21 +++- litellm/responses/main.py | 14 +++ litellm/tool_loop.py | 31 ++++- .../test_allowed_tools_tool_choice.py | 39 ++++++ .../test_anthropic_responses_bridge.py | 113 ++++++++++++++++++ .../mcp/test_aresponses_api_with_mcp.py | 24 ++++ .../mcp/test_chat_completions_handler.py | 38 ++++++ tests/unit/test_main.py | 17 +++ tests/unit/test_tool_loop.py | 74 ++++++++++++ 11 files changed, 398 insertions(+), 8 deletions(-) diff --git a/litellm/litellm_core_utils/allowed_tools_tool_choice.py b/litellm/litellm_core_utils/allowed_tools_tool_choice.py index 276a13fd235..36d110b212e 100644 --- a/litellm/litellm_core_utils/allowed_tools_tool_choice.py +++ b/litellm/litellm_core_utils/allowed_tools_tool_choice.py @@ -124,6 +124,28 @@ def parse_allowed_tools_tool_choice(tool_choice: object) -> ChatAllowedToolsTool return _chat_layout(parsed) +def unenforced_allowed_tools_tool_choice(tool_choice: object, custom_llm_provider: str | None) -> bool: + return is_allowed_tools_tool_choice(tool_choice) and ( + custom_llm_provider not in ALLOWED_TOOLS_TOOL_CHOICE_PROVIDERS + or parse_allowed_tools_tool_choice(tool_choice) is None + ) + + +def refuse_unenforced_allowed_tools_tool_choice( + tool_choice: object, custom_llm_provider: str | None, model: str +) -> None: + if unenforced_allowed_tools_tool_choice(tool_choice, custom_llm_provider): + raise BadRequestError( + message=( + f"This allowed_tools tool_choice can't be enforced on custom_llm_provider={custom_llm_provider}, and " + "LiteLLM runs the tool calls on this request, so drop_params does not drop it. Send only the allowed " + "tools in `tools` with tool_choice='auto' or 'required'." + ), + model=model, + llm_provider=custom_llm_provider or "", + ) + + def chat_allowed_tools_tool_choice(choice: ChatAllowedToolsToolChoice) -> ChatCompletionAllowedToolChoiceParam: return ChatCompletionAllowedToolChoiceParam( type="allowed_tools", diff --git a/litellm/main.py b/litellm/main.py index 71e149ecc14..0b95c3bd15c 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -74,6 +74,7 @@ from litellm.exceptions import LiteLLMUnknownProvider from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.allowed_tools_tool_choice import ( drops_unsupported_allowed_tools_tool_choice, + refuse_unenforced_allowed_tools_tool_choice, ) from litellm.litellm_core_utils.asyncify import asyncify, run_async_function from litellm.litellm_core_utils.audio_utils.utils import ( @@ -5589,7 +5590,8 @@ def completion( model=model, drop_params=drops_unsupported_params( cast("Mapping[str, object]", kwargs).get("drop_params") # cast-ok: kwargs is the untyped request dict - ), + ) + and not may_name_litellm_mcp_gateway(tools), # pyright: ignore[reportUnknownArgumentType] # untyped tools list ) # validate optional params stop = validate_openai_optional_params(stop=stop) @@ -5887,6 +5889,15 @@ def completion( if not is_claude_tool_target(custom_llm_provider=custom_llm_provider, model=model): tools = _drop_anthropic_only_tool_keys(tools=tools) + if skip_mcp_handler: + refuse_unenforced_allowed_tools_tool_choice( + tool_choice=cast( # cast-ok: widening an untyped dict to object is always sound, the helper validates it with pydantic + object, tool_choice + ), + custom_llm_provider=custom_llm_provider, + model=model, + ) + drops_tool_choice: Final = drops_unsupported_allowed_tools_tool_choice( tool_choice=cast( # cast-ok: widening an untyped dict to object is always sound, the helper validates it with pydantic object, tool_choice diff --git a/litellm/responses/litellm_completion_transformation/handler.py b/litellm/responses/litellm_completion_transformation/handler.py index 802ed6bf9c5..153b83c6aab 100644 --- a/litellm/responses/litellm_completion_transformation/handler.py +++ b/litellm/responses/litellm_completion_transformation/handler.py @@ -10,6 +10,7 @@ from typing import Final from pydantic import TypeAdapter import litellm +from litellm.litellm_core_utils.allowed_tools_tool_choice import unenforced_allowed_tools_tool_choice from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.litellm_core_utils.model_param_helper import embeds_thought_signature_in_tool_call_id from litellm.responses.additional_tools import HoistedAdditionalTools, hoist_additional_tools @@ -72,6 +73,17 @@ def _bridged_request( ) +def _with_the_tool_choice_sent( + responses_api_request: ResponsesAPIOptionalRequestParams, custom_llm_provider: str | None +) -> ResponsesAPIOptionalRequestParams: + tool_choice: Final[object] = responses_api_request.get( # pyright: ignore[reportUnknownMemberType] # ToolChoice has an untyped member + "tool_choice" + ) + if not unenforced_allowed_tools_tool_choice(tool_choice, custom_llm_provider): + return responses_api_request + return {**responses_api_request, "tool_choice": None} + + class LiteLLMCompletionTransformationHandler: def response_api_handler( self, @@ -127,13 +139,14 @@ class LiteLLMCompletionTransformationHandler: ), **_chat_verbosity(bridged_request), } + echoed_request: Final = _with_the_tool_choice_sent(bridged_request, custom_llm_provider) if _is_async: return self.async_response_api_handler( previous_response_client_ids, litellm_completion_request=litellm_completion_request, request_input=bridged_input, - responses_api_request=bridged_request, + responses_api_request=echoed_request, **kwargs, ) @@ -146,7 +159,7 @@ class LiteLLMCompletionTransformationHandler: LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( chat_completion_response=litellm_completion_response, request_input=bridged_input, - responses_api_request=bridged_request, + responses_api_request=echoed_request, ) ) @@ -157,7 +170,7 @@ class LiteLLMCompletionTransformationHandler: model=model, litellm_custom_stream_wrapper=litellm_completion_response, request_input=bridged_input, - responses_api_request=bridged_request, + responses_api_request=echoed_request, custom_llm_provider=custom_llm_provider, litellm_metadata=kwargs.get("litellm_metadata", {}), tool_call_provider_fields=not embeds_thought_signature_in_tool_call_id( @@ -209,7 +222,7 @@ class LiteLLMCompletionTransformationHandler: **_chat_verbosity(bridged.responses_api_request), }, request_input=bridged.input, - responses_api_request=bridged.responses_api_request, + responses_api_request=_with_the_tool_choice_sent(bridged.responses_api_request, custom_llm_provider), **request_kwargs, ) diff --git a/litellm/responses/main.py b/litellm/responses/main.py index e9f9eb85d25..5f671f7be08 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -29,6 +29,7 @@ from litellm.integrations.anthropic_cache_control_hook import ( configured_injection_points, routed_model_info, ) +from litellm.litellm_core_utils.allowed_tools_tool_choice import refuse_unenforced_allowed_tools_tool_choice from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.completion_timeout import CompletionTimeout, StreamKind from litellm.litellm_core_utils.core_helpers import drops_unsupported_params, normalize_drop_params @@ -674,6 +675,13 @@ def _bridge_kwargs( return MappingProxyType({key: value for key, value in kwargs.items() if key not in dropped_keys}) +def _refuse_a_choice_the_mcp_tool_run_would_ignore( + skip_mcp_handler: bool, tool_choice: object, custom_llm_provider: str | None, model: str +) -> None: + if skip_mcp_handler: + refuse_unenforced_allowed_tools_tool_choice(tool_choice, custom_llm_provider, model) + + def _raise_untranslatable_responses_fields( failure: _UntranslatableResponsesFields, model: str, custom_llm_provider: str | None ) -> NoReturn: @@ -1762,6 +1770,12 @@ def responses( bridge_kwargs: Final = _bridge_kwargs(kwargs, responses_api_provider_config, allowed_openai_params) if isinstance(bridge_kwargs, _UntranslatableResponsesFields): _raise_untranslatable_responses_fields(bridge_kwargs, model, custom_llm_provider) + _refuse_a_choice_the_mcp_tool_run_would_ignore( + skip_mcp_handler, # pyright: ignore[reportUnknownArgumentType] # popped from untyped kwargs + tool_choice, + custom_llm_provider, + model, + ) return litellm_completion_transformation_handler.response_api_handler( previous_response_client_ids(previous_response_id, kwargs), model=model, diff --git a/litellm/tool_loop.py b/litellm/tool_loop.py index d0b2850bb5b..81d1a8f1b32 100644 --- a/litellm/tool_loop.py +++ b/litellm/tool_loop.py @@ -6,6 +6,10 @@ from typing import Final, TypeAlias, cast from typing_extensions import TypedDict, Unpack from litellm.constants import DEFAULT_TOOL_LOOP_MAX_ROUNDS +from litellm.litellm_core_utils.allowed_tools_tool_choice import ( + is_allowed_tools_tool_choice, + refuse_unenforced_allowed_tools_tool_choice, +) from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionAssistantMessage, @@ -32,11 +36,32 @@ class ToolLoopMaxRoundsExceeded(RuntimeError): super().__init__(f"model still requested tool calls on round {max_rounds} of {max_rounds}") -def _validate_tool_loop_args(max_rounds: int, completion_kwargs: _ToolLoopCompletionKwargs) -> None: +def _text(value: object) -> str | None: + return value if isinstance(value, str) else None + + +def _refuse_unenforced_tool_choice(model: str, completion_kwargs: _ToolLoopCompletionKwargs) -> None: + import litellm + + tool_choice: Final = completion_kwargs.get("tool_choice") + if is_allowed_tools_tool_choice(tool_choice): + refuse_unenforced_allowed_tools_tool_choice( + tool_choice, + litellm.get_llm_provider( + model=model, + custom_llm_provider=_text(completion_kwargs.get("custom_llm_provider")), + api_base=_text(completion_kwargs.get("api_base")), + )[1], + model, + ) + + +def _validate_tool_loop_args(model: str, max_rounds: int, completion_kwargs: _ToolLoopCompletionKwargs) -> None: if max_rounds < 1: raise ValueError(f"max_rounds must be >= 1, got {max_rounds}") if completion_kwargs.get("stream"): raise ValueError("run_tool_loop requires whole responses; stream=True is not supported") + _refuse_unenforced_tool_choice(model, completion_kwargs) def _expect_model_response(response: object) -> ModelResponse: @@ -148,7 +173,7 @@ def run_tool_loop( """ import litellm - _validate_tool_loop_args(max_rounds, completion_kwargs) + _validate_tool_loop_args(model, max_rounds, completion_kwargs) history: tuple[AllMessageValues, ...] = tuple(messages) # rebind-ok: rounds append new turns for round_number in range(1, max_rounds + 1): message = ( @@ -199,7 +224,7 @@ async def arun_tool_loop( """ import litellm - _validate_tool_loop_args(max_rounds, completion_kwargs) + _validate_tool_loop_args(model, max_rounds, completion_kwargs) history: tuple[AllMessageValues, ...] = tuple(messages) # rebind-ok: rounds append new turns for round_number in range(1, max_rounds + 1): message = ( diff --git a/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py b/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py index 0308ae41abf..97d4d6d5152 100644 --- a/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py +++ b/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py @@ -4,11 +4,14 @@ from typing import Final import pytest +from litellm.exceptions import BadRequestError from litellm.litellm_core_utils.allowed_tools_tool_choice import ( ChatAllowedToolsToolChoice, is_allowed_tools_tool_choice, chat_allowed_tools_tool_choice, parse_allowed_tools_tool_choice, + refuse_unenforced_allowed_tools_tool_choice, + unenforced_allowed_tools_tool_choice, responses_allowed_tools_tool_choice, ) @@ -145,3 +148,39 @@ def test_parse_rejects_a_malformed_choice_in_either_layout(tool_choice: dict[str ) def test_is_allowed_tools_tool_choice(tool_choice: Mapping[str, object] | str | None, expected: bool): assert is_allowed_tools_tool_choice(tool_choice) is expected + + +_FLAT_CHOICE: Final = {"type": "allowed_tools", "mode": "required", "tools": [{"type": "function", "name": "get_time"}]} + + +@pytest.mark.parametrize( + ("tool_choice", "custom_llm_provider"), + [ + (_FLAT_CHOICE, "anthropic"), + (_CHAT_LAYOUT, "hosted_vllm"), + ({"type": "allowed_tools", "mode": "required", "tools": []}, "openai"), + ({"type": "allowed_tools", "tools": [{"type": "function", "name": "get_time"}]}, "azure"), + ], +) +def test_a_tool_runner_refuses_a_choice_nothing_would_enforce(tool_choice: object, custom_llm_provider: str): + with pytest.raises(BadRequestError, match="drop_params does not drop it") as exc: + refuse_unenforced_allowed_tools_tool_choice(tool_choice, custom_llm_provider, "the-model") + + assert exc.value.status_code == 400 + + +@pytest.mark.parametrize( + ("tool_choice", "custom_llm_provider"), + [ + (_FLAT_CHOICE, "openai"), + (_CHAT_LAYOUT, "azure"), + (_FLAT_CHOICE, "litellm_proxy"), + ("required", "anthropic"), + ({"type": "function", "function": {"name": "get_time"}}, "anthropic"), + (None, "anthropic"), + ], +) +def test_a_tool_runner_keeps_a_choice_the_provider_enforces_or_any_other_choice( + tool_choice: object, custom_llm_provider: str +): + assert not unenforced_allowed_tools_tool_choice(tool_choice, custom_llm_provider) diff --git a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py index 5ae985ea7e5..66b13b3bbd3 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py +++ b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py @@ -2081,3 +2081,116 @@ async def test_codex_tool_search_follow_up_text_verbosity_follows_provider_param assert dropped == without_text assert allowed == {**without_text, "verbosity": "low"} + + +_FLAT_ALLOWED_WEATHER: Final = { + "type": "allowed_tools", + "mode": "required", + "tools": [{"type": "function", "name": "get_weather"}], +} +_RESPONSES_TOOLS: Final = ( + {"type": "function", "name": "get_weather", "parameters": {"type": "object", "properties": {}}}, + {"type": "function", "name": "get_time", "parameters": {"type": "object", "properties": {}}}, +) + + +def _chat_answer() -> httpx.Response: + return httpx.Response( + 200, + json={ + "id": "chatcmpl-echo", + "object": "chat.completion", + "created": 1739462947, + "model": "the-model", + "choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "done"}}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + ) + + +@pytest.mark.usefixtures(httpx_transport.__name__) +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.asyncio +async def test_responses_reports_auto_for_an_allowed_tools_choice_it_dropped( + respx_mock: respx.MockRouter, use_async: bool +) -> None: + route: Final = respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + return_value=_anthropic_message({"type": "text", "text": "done"}) + ) + call: Final = { + "model": "anthropic/claude-haiku-4-5", + "input": "What is the weather?", + "tools": list(_RESPONSES_TOOLS), + "tool_choice": _FLAT_ALLOWED_WEATHER, + "drop_params": True, + "api_key": "sk-ant-fake", + } + + response: Final = await litellm.aresponses(**call) if use_async else litellm.responses(**call) + + assert "tool_choice" not in json.loads(route.calls[0].request.content) + assert response.tool_choice == "auto" + + +def test_streamed_responses_report_auto_for_an_allowed_tools_choice_it_dropped(respx_mock: respx.MockRouter) -> None: + respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + return_value=httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=_anthropic_sse( + { + "type": "message_start", + "message": { + "id": "msg_dropped_choice", + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5", + "content": [], + "usage": {"input_tokens": 9, "output_tokens": 1}, + }, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "done"}}, + {"type": "content_block_stop", "index": 0}, + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 4}}, + {"type": "message_stop"}, + ), + ) + ) + + events: Final = tuple( + litellm.responses( + model="anthropic/claude-haiku-4-5", + input="What is the weather?", + tools=list(_RESPONSES_TOOLS), + tool_choice=_FLAT_ALLOWED_WEATHER, + drop_params=True, + stream=True, + api_key="sk-ant-fake", + ) + ) + + reported: Final = tuple( + event.response.tool_choice + for event in events + if event.type in (ResponsesAPIStreamEvents.RESPONSE_CREATED, ResponsesAPIStreamEvents.RESPONSE_COMPLETED) + ) + assert reported == ("auto", "auto") + + +def test_responses_through_chat_keeps_reporting_a_choice_the_provider_received(respx_mock: respx.MockRouter) -> None: + route: Final = respx_mock.post("http://localhost:12346/v1/chat/completions").mock(return_value=_chat_answer()) + + response: Final = litellm.responses( + model="openai/gpt-5.5", + input="What is the weather?", + tools=list(_RESPONSES_TOOLS), + tool_choice=_FLAT_ALLOWED_WEATHER, + use_chat_completions_api=True, + api_base="http://localhost:12346/v1", + api_key="sk-fake", + ) + + assert json.loads(route.calls[0].request.content)["tool_choice"]["type"] == "allowed_tools" + assert response.tool_choice["type"] == "allowed_tools" + assert list(response.tool_choice["tools"]) == _FLAT_ALLOWED_WEATHER["tools"] diff --git a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py index 037004ab999..f8eb0378ed1 100644 --- a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py +++ b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py @@ -1696,6 +1696,30 @@ async def _auto_execute( return response +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +async def test_mcp_tools_keep_the_400_for_a_choice_the_chat_bridge_would_drop( + respx_mock: respx.MockRouter, stream: bool +) -> None: + async def call_with_the_choice() -> None: + response: Final = await litellm.aresponses( + model="anthropic/claude-haiku-4-5", + input="Call echo with PING", + tools=[dict(_ECHO_MCP_TOOL)], + tool_choice={"type": "allowed_tools", "mode": "auto", "tools": [{"type": "function", "name": "echo"}]}, + drop_params=True, + stream=stream, + api_key="sk-ant-fake", + ) + if stream: + [event async for event in response] + + with _echo_tool_served("PING"), pytest.raises(litellm.BadRequestError, match="drop_params does not drop it"): + await call_with_the_choice() + + assert not respx_mock.calls + + async def _auto_execute_echo( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, logging_obj: Logging ) -> ResponsesAPIResponse: diff --git a/tests/unit/responses/mcp/test_chat_completions_handler.py b/tests/unit/responses/mcp/test_chat_completions_handler.py index c0d32742e8a..ecef6b44bae 100644 --- a/tests/unit/responses/mcp/test_chat_completions_handler.py +++ b/tests/unit/responses/mcp/test_chat_completions_handler.py @@ -2744,3 +2744,41 @@ def test_stream_auto_execute_read_synchronously_and_dropped_after_its_loop_close assert [raised for raised in unraisable.raised if _names_chat_stream_billing(raised)] == [] assert [str(warning.message) for warning in caught if _names_chat_stream_billing(str(warning.message))] == [] + + +_MCP_ECHO_TOOL: Final = MCPTool.model_validate( + {"name": "echo", "inputSchema": {"type": "object", "properties": {"value": {"type": "string"}}}}, + by_name=False, +) +_FLAT_ALLOWED_ECHO: Final = {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "function", "name": "echo"}]} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("require_approval", ["never", "always"]) +@pytest.mark.parametrize("stream", [False, True]) +async def test_mcp_tools_keep_the_400_for_an_allowed_tools_choice_the_provider_cannot_enforce( + respx_mock: respx.MockRouter, require_approval: str, stream: bool +): + execute: Final = AsyncMock(return_value=[]) + + with ( + patch.object( + LiteLLM_Proxy_MCP_Handler, + "process_mcp_tools_without_openai_transform", + AsyncMock(return_value=([_MCP_ECHO_TOOL], {"echo": "demo"})), + ), + patch.object(LiteLLM_Proxy_MCP_Handler, "execute_tool_calls", execute), + pytest.raises(litellm.BadRequestError, match="drop_params does not drop it"), + ): + await litellm.acompletion( + model="anthropic/claude-haiku-4-5", + messages=[{"role": "user", "content": "Call echo"}], + tools=[{"type": "mcp", "server_url": "litellm_proxy/mcp/demo", "require_approval": require_approval}], + tool_choice=_FLAT_ALLOWED_ECHO, + drop_params=True, + stream=stream, + api_key="sk-ant-fake", + ) + + execute.assert_not_awaited() + assert not respx_mock.calls diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 42a109fd7e1..4f60440ff68 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -647,6 +647,23 @@ def test_completion_drop_params_drops_an_allowed_tools_tool_choice_it_cannot_sen assert [tool["function"]["name"] for tool in sent["tools"]] == ["get_time", "get_weather"] +@pytest.mark.asyncio +async def test_a_malformed_allowed_tools_tool_choice_stays_a_400_on_mcp_gateway_tools(respx_mock: respx.MockRouter): + with pytest.raises(litellm.BadRequestError, match="Invalid tool choice") as exc: + await litellm.acompletion( + model="openai/gpt-5.5", + messages=[{"role": "user", "content": "What time is it?"}], + tools=[{"type": "mcp", "server_url": "litellm_proxy/mcp/demo", "require_approval": "never"}], + tool_choice={"type": "allowed_tools", "mode": "required", "tools": []}, + drop_params=True, + api_base="http://localhost:12346/v1", + api_key="fake_api_key", + ) + + assert exc.value.status_code == 400 + assert not respx_mock.calls + + _CLAUDE_TOOL_SEARCH_CALL: Final = { "id": "srvtoolu_01search", "type": "function", diff --git a/tests/unit/test_tool_loop.py b/tests/unit/test_tool_loop.py index de8f62803cc..22bfd4ccf1d 100644 --- a/tests/unit/test_tool_loop.py +++ b/tests/unit/test_tool_loop.py @@ -248,6 +248,80 @@ def test_stream_rejected_before_any_request(respx_mock: respx.MockRouter) -> Non assert route.call_count == 0 +_FLAT_ALLOWED_TOOLS: Final = {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "function", "name": "a"}]} +_EMPTY_ALLOWED_TOOLS: Final = {"type": "allowed_tools", "mode": "auto", "tools": []} + + +@pytest.mark.parametrize( + ("model", "tool_choice"), + [("anthropic/claude-haiku-4-5", _FLAT_ALLOWED_TOOLS), ("openai/gpt-5-mini", _EMPTY_ALLOWED_TOOLS)], +) +def test_allowed_tools_choice_nothing_enforces_is_refused_under_drop_params( + respx_mock: respx.MockRouter, model: str, tool_choice: dict[str, object] +) -> None: + executed: Final[list[ChatCompletionMessageToolCall]] = [] + + def execute(tool_call: ChatCompletionMessageToolCall) -> ChatCompletionToolMessage: + executed.append(tool_call) + return _tool_result(tool_call) + + with pytest.raises(litellm.BadRequestError, match="drop_params does not drop it"): + litellm.run_tool_loop( + model=model, + messages=[{"role": "user", "content": "hi"}], + tools=WEATHER_TOOLS, + execute_tool=execute, + tool_choice=tool_choice, + drop_params=True, + api_key="sk-test", + ) + + assert not executed + assert not respx_mock.calls + + +@pytest.mark.asyncio +async def test_async_loop_refuses_an_allowed_tools_choice_the_provider_cannot_enforce( + respx_mock: respx.MockRouter, +) -> None: + async def execute(tool_call: ChatCompletionMessageToolCall) -> ChatCompletionToolMessage: + return _tool_result(tool_call) + + with pytest.raises(litellm.BadRequestError, match="drop_params does not drop it"): + await litellm.arun_tool_loop( + model="anthropic/claude-haiku-4-5", + messages=[{"role": "user", "content": "hi"}], + tools=WEATHER_TOOLS, + execute_tool=execute, + tool_choice=_FLAT_ALLOWED_TOOLS, + drop_params=True, + api_key="sk-test", + ) + + assert not respx_mock.calls + + +def test_allowed_tools_choice_the_provider_enforces_reaches_it_every_round(respx_mock: respx.MockRouter) -> None: + respx_mock.post(OPENAI_CHAT_COMPLETIONS_URL).mock( + side_effect=[ + httpx.Response(200, json=_openai_response(None, [_tool_call("call_1", "get_weather", {"city": "Paris"})])), + httpx.Response(200, json=_openai_response("done")), + ] + ) + + answer: Final = litellm.run_tool_loop( + model="openai/gpt-5-mini", + messages=[{"role": "user", "content": "hi"}], + tools=WEATHER_TOOLS, + execute_tool=_tool_result, + tool_choice={"type": "allowed_tools", "mode": "auto", "tools": [{"type": "function", "name": "get_weather"}]}, + api_key="sk-test", + ) + + assert answer == "done" + assert [body["tool_choice"]["type"] for body in _request_bodies(respx_mock)] == ["allowed_tools", "allowed_tools"] + + def test_custom_tool_call_raises_type_error_without_executing(respx_mock: respx.MockRouter) -> None: custom_response: Final = _openai_response( None, From e25a82c3f54ff54178b5ebb98fb1da6a753e1fe2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 5 Nov 2026 18:53:56 -0800 Subject: [PATCH 4/9] fix(tool-choice): refuse an excluded tool right before LiteLLM runs it Checking the provider up front missed a model alias, an old Azure api_version, a litellm_proxy hop, and SDK fallbacks. The tool loop and the MCP gateway now check the allowed_tools list right before they run a tool call, so an excluded call answers 400 and never runs. Under drop_params the choice now drops only when the caller runs every tool, so a provider-run MCP server or built-in keeps the 400. The Azure checks read the request's api_version, which also fixes the /v1/responses echo there. --- .../allowed_tools_tool_choice.py | 89 ++++++++-- litellm/llms/azure/chat/gpt_transformation.py | 7 +- litellm/main.py | 5 +- .../handler.py | 17 +- litellm/responses/main.py | 8 +- .../responses/mcp/chat_completions_handler.py | 2 + .../mcp/litellm_proxy_mcp_handler.py | 7 + .../responses/mcp/mcp_streaming_iterator.py | 1 + litellm/tool_loop.py | 44 ++--- .../test_allowed_tools_tool_choice.py | 69 ++++++++ .../test_anthropic_responses_bridge.py | 50 ++++++ .../mcp/test_aresponses_api_with_mcp.py | 103 ++++++++++-- .../mcp/test_chat_completions_handler.py | 94 ++++++++++- tests/unit/test_main.py | 16 ++ tests/unit/test_tool_loop.py | 155 ++++++++++++++---- 15 files changed, 569 insertions(+), 98 deletions(-) diff --git a/litellm/litellm_core_utils/allowed_tools_tool_choice.py b/litellm/litellm_core_utils/allowed_tools_tool_choice.py index 36d110b212e..ff067b8b8e4 100644 --- a/litellm/litellm_core_utils/allowed_tools_tool_choice.py +++ b/litellm/litellm_core_utils/allowed_tools_tool_choice.py @@ -1,4 +1,4 @@ -from collections.abc import Mapping +from collections.abc import Iterable, Mapping from typing import Final, Literal from openai.types.chat.chat_completion_allowed_tool_choice_param import ChatCompletionAllowedToolChoiceParam @@ -9,6 +9,8 @@ from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError from litellm.exceptions import BadRequestError ALLOWED_TOOLS_TOOL_CHOICE_PROVIDERS: Final = frozenset({"openai", "azure", "litellm_proxy"}) +_FIRST_AZURE_API_VERSION_WITH_TOOL_CHOICE: Final = "2023-12-01" +_TOOL_TYPES_THE_CALLER_RUNS: Final = frozenset({"function", "custom"}) CHAT_ALLOWED_TOOLS_LAYOUT: Final = ( '{"type": "allowed_tools", "allowed_tools": {"mode": "auto" | "required", "tools": []}}' @@ -57,11 +59,19 @@ class _ToolChoiceType(BaseModel): _TOOL_CHOICE_TYPE: Final = TypeAdapter(_ToolChoiceType) -def is_allowed_tools_tool_choice(tool_choice: object) -> bool: +def _declared_type(value: object) -> str | None: try: - return _TOOL_CHOICE_TYPE.validate_python(tool_choice).type == "allowed_tools" + return _TOOL_CHOICE_TYPE.validate_python(value).type except ValidationError: - return False + return None + + +def is_allowed_tools_tool_choice(tool_choice: object) -> bool: + return _declared_type(tool_choice) == "allowed_tools" + + +def caller_runs_every_tool(tools: Iterable[object] | None) -> bool: + return all(_declared_type(tool) in _TOOL_TYPES_THE_CALLER_RUNS for tool in tools or ()) def drops_unsupported_allowed_tools_tool_choice( @@ -69,11 +79,14 @@ def drops_unsupported_allowed_tools_tool_choice( custom_llm_provider: str, model: str, drop_params: bool, + tools: Iterable[object] | None, ) -> bool: if not is_allowed_tools_tool_choice(tool_choice) or custom_llm_provider in ALLOWED_TOOLS_TOOL_CHOICE_PROVIDERS: return False - if drop_params: + if drop_params and caller_runs_every_tool(tools): return True + if drop_params: + raise _a_choice_drop_params_keeps(custom_llm_provider, model) raise BadRequestError( message=( "tool_choice type 'allowed_tools' is only supported for " @@ -124,25 +137,75 @@ def parse_allowed_tools_tool_choice(tool_choice: object) -> ChatAllowedToolsTool return _chat_layout(parsed) -def unenforced_allowed_tools_tool_choice(tool_choice: object, custom_llm_provider: str | None) -> bool: +def azure_api_version_takes_tool_choice(api_version: str) -> bool: + return api_version.count("-") < 2 or api_version >= _FIRST_AZURE_API_VERSION_WITH_TOOL_CHOICE + + +def _azure_request_takes_tool_choice(api_version: str | None) -> bool: + import litellm + from litellm.secret_managers.main import get_secret_str + + return azure_api_version_takes_tool_choice( + api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION") or litellm.AZURE_DEFAULT_API_VERSION + ) + + +def _enforces_allowed_tools(custom_llm_provider: str | None, api_version: str | None) -> bool: + if custom_llm_provider == "azure": + return _azure_request_takes_tool_choice(api_version) + return custom_llm_provider in ALLOWED_TOOLS_TOOL_CHOICE_PROVIDERS + + +def unenforced_allowed_tools_tool_choice( + tool_choice: object, custom_llm_provider: str | None, api_version: str | None = None +) -> bool: return is_allowed_tools_tool_choice(tool_choice) and ( - custom_llm_provider not in ALLOWED_TOOLS_TOOL_CHOICE_PROVIDERS + not _enforces_allowed_tools(custom_llm_provider, api_version) or parse_allowed_tools_tool_choice(tool_choice) is None ) +def _a_choice_drop_params_keeps(custom_llm_provider: str | None, model: str) -> BadRequestError: + return BadRequestError( + message=( + f"This allowed_tools tool_choice can't be enforced on custom_llm_provider={custom_llm_provider}, and the " + "caller does not run every tool on this request, so drop_params does not drop it. Send only the allowed " + "tools in `tools` with tool_choice='auto' or 'required'." + ), + model=model, + llm_provider=custom_llm_provider or "", + ) + + def refuse_unenforced_allowed_tools_tool_choice( - tool_choice: object, custom_llm_provider: str | None, model: str + tool_choice: object, custom_llm_provider: str | None, model: str, api_version: str | None = None ) -> None: - if unenforced_allowed_tools_tool_choice(tool_choice, custom_llm_provider): + if unenforced_allowed_tools_tool_choice(tool_choice, custom_llm_provider, api_version): + raise _a_choice_drop_params_keeps(custom_llm_provider, model) + + +def _allowed_tool_names(tool_choice: object) -> frozenset[object]: + parsed: Final = parse_allowed_tools_tool_choice(tool_choice) + if parsed is None: + return frozenset() + flattened: Final = tuple(_responses_allowed_tool(tool) for tool in parsed.allowed_tools.tools) + return frozenset(tool.get("name") for tool in flattened if tool is not None) + + +def refuse_tool_calls_the_choice_excludes(tool_choice: object, tool_names: Iterable[str], model: str) -> None: + if not is_allowed_tools_tool_choice(tool_choice): + return + allowed: Final = _allowed_tool_names(tool_choice) + excluded: Final = tuple(name for name in tool_names if name not in allowed) + if excluded: raise BadRequestError( message=( - f"This allowed_tools tool_choice can't be enforced on custom_llm_provider={custom_llm_provider}, and " - "LiteLLM runs the tool calls on this request, so drop_params does not drop it. Send only the allowed " - "tools in `tools` with tool_choice='auto' or 'required'." + f"The model called {', '.join(excluded)}, which the allowed_tools tool_choice excludes, because this " + "request did not reach a model that enforces the choice. LiteLLM did not run the call. Send only the " + "allowed tools in `tools` with tool_choice='auto' or 'required'." ), model=model, - llm_provider=custom_llm_provider or "", + llm_provider="", ) diff --git a/litellm/llms/azure/chat/gpt_transformation.py b/litellm/llms/azure/chat/gpt_transformation.py index e0aceba52f9..7ca4ba6c163 100644 --- a/litellm/llms/azure/chat/gpt_transformation.py +++ b/litellm/llms/azure/chat/gpt_transformation.py @@ -6,6 +6,7 @@ from httpx._models import Headers, Response import litellm from litellm.integrations.anthropic_cache_control_hook import deployment_prompt_cache_breakpoint_support +from litellm.litellm_core_utils.allowed_tools_tool_choice import azure_api_version_takes_tool_choice from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.litellm_core_utils.prompt_templates.common_utils import ( chat_completion_prompt_cache_breakpoints_for_wire, @@ -240,11 +241,7 @@ class AzureOpenAIConfig(BaseConfig): if api_version_year is None or api_version_month is None or api_version_day is None: optional_params["tool_choice"] = value else: - if ( - api_version_year < "2023" - or (api_version_year == "2023" and api_version_month < "12") - or (api_version_year == "2023" and api_version_month == "12" and api_version_day < "01") - ): + if not azure_api_version_takes_tool_choice(api_version): if not (litellm.drop_params is True or drop_params is True): raise UnsupportedParamsError( status_code=400, diff --git a/litellm/main.py b/litellm/main.py index 0b95c3bd15c..b45a4ba8812 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -73,6 +73,7 @@ from litellm.constants import ( from litellm.exceptions import LiteLLMUnknownProvider from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.allowed_tools_tool_choice import ( + caller_runs_every_tool, drops_unsupported_allowed_tools_tool_choice, refuse_unenforced_allowed_tools_tool_choice, ) @@ -5591,7 +5592,7 @@ def completion( drop_params=drops_unsupported_params( cast("Mapping[str, object]", kwargs).get("drop_params") # cast-ok: kwargs is the untyped request dict ) - and not may_name_litellm_mcp_gateway(tools), # pyright: ignore[reportUnknownArgumentType] # untyped tools list + and caller_runs_every_tool(tools), # pyright: ignore[reportUnknownArgumentType] # untyped tools list ) # validate optional params stop = validate_openai_optional_params(stop=stop) @@ -5896,6 +5897,7 @@ def completion( ), custom_llm_provider=custom_llm_provider, model=model, + api_version=api_version, ) drops_tool_choice: Final = drops_unsupported_allowed_tools_tool_choice( @@ -5905,6 +5907,7 @@ def completion( custom_llm_provider=custom_llm_provider, model=model, drop_params=drop_unsupported_reasoning, + tools=tools, # pyright: ignore[reportUnknownArgumentType] # untyped tools list ) if provider_specific_header is not None: diff --git a/litellm/responses/litellm_completion_transformation/handler.py b/litellm/responses/litellm_completion_transformation/handler.py index 153b83c6aab..aa5c7028240 100644 --- a/litellm/responses/litellm_completion_transformation/handler.py +++ b/litellm/responses/litellm_completion_transformation/handler.py @@ -74,12 +74,17 @@ def _bridged_request( def _with_the_tool_choice_sent( - responses_api_request: ResponsesAPIOptionalRequestParams, custom_llm_provider: str | None + responses_api_request: ResponsesAPIOptionalRequestParams, + custom_llm_provider: str | None, + request_kwargs: Mapping[str, object], ) -> ResponsesAPIOptionalRequestParams: + api_version: Final = request_kwargs.get("api_version") tool_choice: Final[object] = responses_api_request.get( # pyright: ignore[reportUnknownMemberType] # ToolChoice has an untyped member "tool_choice" ) - if not unenforced_allowed_tools_tool_choice(tool_choice, custom_llm_provider): + if not unenforced_allowed_tools_tool_choice( + tool_choice, custom_llm_provider, api_version if isinstance(api_version, str) else None + ): return responses_api_request return {**responses_api_request, "tool_choice": None} @@ -139,7 +144,9 @@ class LiteLLMCompletionTransformationHandler: ), **_chat_verbosity(bridged_request), } - echoed_request: Final = _with_the_tool_choice_sent(bridged_request, custom_llm_provider) + echoed_request: Final = _with_the_tool_choice_sent( + bridged_request, custom_llm_provider, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs) + ) if _is_async: return self.async_response_api_handler( @@ -222,7 +229,9 @@ class LiteLLMCompletionTransformationHandler: **_chat_verbosity(bridged.responses_api_request), }, request_input=bridged.input, - responses_api_request=_with_the_tool_choice_sent(bridged.responses_api_request, custom_llm_provider), + responses_api_request=_with_the_tool_choice_sent( + bridged.responses_api_request, custom_llm_provider, request_kwargs + ), **request_kwargs, ) diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 5f671f7be08..3dbcc299008 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -478,6 +478,7 @@ async def aresponses_api_with_mcp( tool_server_map=tool_server_map, served_tools=original_mcp_tools, tool_calls=tool_calls, + tool_choice=tool_choice, user_api_key_auth=user_api_key_auth, mcp_auth_header=mcp_auth_header, mcp_server_auth_headers=mcp_server_auth_headers, @@ -676,10 +677,12 @@ def _bridge_kwargs( def _refuse_a_choice_the_mcp_tool_run_would_ignore( - skip_mcp_handler: bool, tool_choice: object, custom_llm_provider: str | None, model: str + skip_mcp_handler: bool, tool_choice: object, custom_llm_provider: str | None, model: str, api_version: object ) -> None: if skip_mcp_handler: - refuse_unenforced_allowed_tools_tool_choice(tool_choice, custom_llm_provider, model) + refuse_unenforced_allowed_tools_tool_choice( + tool_choice, custom_llm_provider, model, api_version if isinstance(api_version, str) else None + ) def _raise_untranslatable_responses_fields( @@ -1775,6 +1778,7 @@ def responses( tool_choice, custom_llm_provider, model, + kwargs.get("api_version"), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # untyped kwargs ) return litellm_completion_transformation_handler.response_api_handler( previous_response_client_ids(previous_response_id, kwargs), diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index f8a52227d1c..c716864d155 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -500,6 +500,7 @@ async def acompletion_with_mcp( tool_server_map=self.tool_server_map, served_tools=deduplicated_mcp_tools, tool_calls=self.tool_calls, + tool_choice=kwargs.get("tool_choice"), 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, @@ -683,6 +684,7 @@ async def acompletion_with_mcp( tool_results: Final = await LiteLLM_Proxy_MCP_Handler.execute_tool_calls( tool_server_map=tool_server_map, tool_calls=tool_calls, + tool_choice=kwargs.get("tool_choice"), served_tools=deduplicated_mcp_tools, user_api_key_auth=user_api_key_auth, mcp_auth_header=mcp_auth_header, diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 1104aee353e..d37ee58eb8f 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -10,6 +10,7 @@ from typing_extensions import Unpack from litellm._logging import verbose_logger from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG +from litellm.litellm_core_utils.allowed_tools_tool_choice import refuse_tool_calls_the_choice_excludes from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy.mcp_server.catalog import catalog_operation, global_manager from litellm.proxy.mcp_server.utils import ( @@ -740,6 +741,7 @@ class LiteLLM_Proxy_MCP_Handler: guardrail_context: Mapping[str, object] | None = None, served_tools: Sequence[MCPTool] | None = None, request_trace_id: str | None = None, + tool_choice: object = None, ) -> list[MCPToolResult]: """Execute tool calls and return results.""" from fastapi import HTTPException @@ -754,6 +756,11 @@ class LiteLLM_Proxy_MCP_Handler: ) from litellm.proxy.proxy_server import proxy_logging_obj + refuse_tool_calls_the_choice_excludes( + tool_choice, + tuple(LiteLLM_Proxy_MCP_Handler.extract_tool_call_details(tool_call)[0] or "" for tool_call in tool_calls), + "", + ) tool_results: Final[list[MCPToolResult]] = [] tool_call_id: str | None = None rules_obj: Final = Rules() diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index 189366feab4..b5e3eee9ecf 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -898,6 +898,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): tool_server_map=self.tool_server_map, served_tools=self.served_tools, tool_calls=tool_calls, + tool_choice=self.original_request_params.get("tool_choice"), 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, diff --git a/litellm/tool_loop.py b/litellm/tool_loop.py index 81d1a8f1b32..1db6b8c6000 100644 --- a/litellm/tool_loop.py +++ b/litellm/tool_loop.py @@ -6,10 +6,7 @@ from typing import Final, TypeAlias, cast from typing_extensions import TypedDict, Unpack from litellm.constants import DEFAULT_TOOL_LOOP_MAX_ROUNDS -from litellm.litellm_core_utils.allowed_tools_tool_choice import ( - is_allowed_tools_tool_choice, - refuse_unenforced_allowed_tools_tool_choice, -) +from litellm.litellm_core_utils.allowed_tools_tool_choice import refuse_tool_calls_the_choice_excludes from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionAssistantMessage, @@ -36,32 +33,11 @@ class ToolLoopMaxRoundsExceeded(RuntimeError): super().__init__(f"model still requested tool calls on round {max_rounds} of {max_rounds}") -def _text(value: object) -> str | None: - return value if isinstance(value, str) else None - - -def _refuse_unenforced_tool_choice(model: str, completion_kwargs: _ToolLoopCompletionKwargs) -> None: - import litellm - - tool_choice: Final = completion_kwargs.get("tool_choice") - if is_allowed_tools_tool_choice(tool_choice): - refuse_unenforced_allowed_tools_tool_choice( - tool_choice, - litellm.get_llm_provider( - model=model, - custom_llm_provider=_text(completion_kwargs.get("custom_llm_provider")), - api_base=_text(completion_kwargs.get("api_base")), - )[1], - model, - ) - - -def _validate_tool_loop_args(model: str, max_rounds: int, completion_kwargs: _ToolLoopCompletionKwargs) -> None: +def _validate_tool_loop_args(max_rounds: int, completion_kwargs: _ToolLoopCompletionKwargs) -> None: if max_rounds < 1: raise ValueError(f"max_rounds must be >= 1, got {max_rounds}") if completion_kwargs.get("stream"): raise ValueError("run_tool_loop requires whole responses; stream=True is not supported") - _refuse_unenforced_tool_choice(model, completion_kwargs) def _expect_model_response(response: object) -> ModelResponse: @@ -139,6 +115,16 @@ def _client_tool_calls( ) +def _refuse_tool_calls_the_choice_excludes( + tool_calls: tuple[ChatCompletionMessageToolCall, ...], model: str, completion_kwargs: _ToolLoopCompletionKwargs +) -> None: + refuse_tool_calls_the_choice_excludes( + completion_kwargs.get("tool_choice"), + tuple(tool_call.function.name or "" for tool_call in tool_calls), + model, + ) + + def _replays_provider_specific_fields( message: Message, tool_calls: tuple[ChatCompletionMessageToolCall, ...], @@ -173,7 +159,7 @@ def run_tool_loop( """ import litellm - _validate_tool_loop_args(model, max_rounds, completion_kwargs) + _validate_tool_loop_args(max_rounds, completion_kwargs) history: tuple[AllMessageValues, ...] = tuple(messages) # rebind-ok: rounds append new turns for round_number in range(1, max_rounds + 1): message = ( @@ -191,6 +177,7 @@ def run_tool_loop( return message.content if round_number == max_rounds: break + _refuse_tool_calls_the_choice_excludes(client_tool_calls, model, completion_kwargs) tool_results = tuple(execute_tool(tool_call) for tool_call in client_tool_calls) history = ( *history, @@ -224,7 +211,7 @@ async def arun_tool_loop( """ import litellm - _validate_tool_loop_args(model, max_rounds, completion_kwargs) + _validate_tool_loop_args(max_rounds, completion_kwargs) history: tuple[AllMessageValues, ...] = tuple(messages) # rebind-ok: rounds append new turns for round_number in range(1, max_rounds + 1): message = ( @@ -242,6 +229,7 @@ async def arun_tool_loop( return message.content if round_number == max_rounds: break + _refuse_tool_calls_the_choice_excludes(client_tool_calls, model, completion_kwargs) tool_results = tuple([await execute_tool(tool_call) for tool_call in client_tool_calls]) history = ( *history, diff --git a/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py b/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py index 97d4d6d5152..698c5dcf2f5 100644 --- a/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py +++ b/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py @@ -4,12 +4,15 @@ from typing import Final import pytest +import litellm from litellm.exceptions import BadRequestError from litellm.litellm_core_utils.allowed_tools_tool_choice import ( ChatAllowedToolsToolChoice, + drops_unsupported_allowed_tools_tool_choice, is_allowed_tools_tool_choice, chat_allowed_tools_tool_choice, parse_allowed_tools_tool_choice, + refuse_tool_calls_the_choice_excludes, refuse_unenforced_allowed_tools_tool_choice, unenforced_allowed_tools_tool_choice, responses_allowed_tools_tool_choice, @@ -184,3 +187,69 @@ def test_a_tool_runner_keeps_a_choice_the_provider_enforces_or_any_other_choice( tool_choice: object, custom_llm_provider: str ): assert not unenforced_allowed_tools_tool_choice(tool_choice, custom_llm_provider) + + +@pytest.mark.parametrize( + ("tool_choice", "tool_names"), + [(_FLAT_CHOICE, ("get_time",)), (_CHAT_LAYOUT, ("get_time", "apply_patch"))], +) +def test_a_tool_call_the_choice_lists_passes(tool_choice: object, tool_names: tuple[str, ...]): + assert refuse_tool_calls_the_choice_excludes(tool_choice, tool_names, "the-model") is None + + +@pytest.mark.parametrize( + "tool_choice", [_FLAT_CHOICE, _CHAT_LAYOUT, {"type": "allowed_tools", "mode": "required", "tools": []}] +) +def test_a_tool_call_the_choice_does_not_list_is_refused(tool_choice: object): + with pytest.raises(BadRequestError, match=r"The model called .*get_weather") as exc: + refuse_tool_calls_the_choice_excludes(tool_choice, ("get_time", "get_weather"), "the-model") + + assert exc.value.status_code == 400 + + +@pytest.mark.parametrize("tool_choice", ["required", None, {"type": "function", "function": {"name": "get_time"}}]) +def test_any_other_choice_leaves_the_tool_calls_alone(tool_choice: object): + assert refuse_tool_calls_the_choice_excludes(tool_choice, ("get_weather",), "the-model") is None + + +@pytest.mark.parametrize( + ("api_version", "unenforced"), + [ + ("2023-07-01-preview", True), + ("2023-11-30", True), + ("2023-12-01-preview", False), + ("2025-04-01-preview", False), + ("v1", False), + ], +) +def test_azure_counts_as_enforcing_only_on_a_version_that_takes_tool_choice(api_version: str, unenforced: bool): + assert unenforced_allowed_tools_tool_choice(_FLAT_CHOICE, "azure", api_version) is unenforced + + +def test_azure_reads_the_same_api_version_fallback_as_the_request(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "api_version", "2023-07-01-preview") + + assert unenforced_allowed_tools_tool_choice(_FLAT_CHOICE, "azure") + + +_FUNCTION_TOOL: Final = {"type": "function", "function": {"name": "get_time"}} + + +@pytest.mark.parametrize( + "tools", + [ + [_FUNCTION_TOOL, {"type": "mcp", "server_label": "docs", "server_url": "https://mcp.example.com"}], + [{"type": "web_search_20250305", "name": "web_search"}], + [{"googleSearch": {}}], + [{"type": ["function"]}], + ["get_time"], + ], +) +def test_drop_params_keeps_the_400_when_the_provider_could_run_a_tool(tools: list[object]): + with pytest.raises(BadRequestError, match="drop_params does not drop it"): + drops_unsupported_allowed_tools_tool_choice(_FLAT_CHOICE, "anthropic", "the-model", True, tools) + + +@pytest.mark.parametrize("tools", [None, [], [_FUNCTION_TOOL, {"type": "custom", "custom": {"name": "patch"}}]]) +def test_drop_params_drops_the_choice_when_the_caller_runs_every_tool(tools: list[object] | None): + assert drops_unsupported_allowed_tools_tool_choice(_FLAT_CHOICE, "anthropic", "the-model", True, tools) diff --git a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py index 66b13b3bbd3..79a655fc913 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py +++ b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py @@ -2132,6 +2132,56 @@ async def test_responses_reports_auto_for_an_allowed_tools_choice_it_dropped( assert response.tool_choice == "auto" +@pytest.mark.usefixtures(httpx_transport.__name__) +@pytest.mark.asyncio +async def test_responses_keep_the_400_when_the_provider_would_run_an_mcp_tool(respx_mock: respx.MockRouter) -> None: + with pytest.raises(litellm.BadRequestError, match="drop_params does not drop it"): + await litellm.aresponses( + model="anthropic/claude-haiku-4-5", + input="What is the weather?", + tools=[{"type": "mcp", "server_label": "docs", "server_url": "https://mcp.example.com/sse"}], + tool_choice=_FLAT_ALLOWED_WEATHER, + drop_params=True, + api_key="sk-ant-fake", + ) + + assert not respx_mock.calls + + +@pytest.mark.usefixtures(httpx_transport.__name__) +def test_responses_report_auto_when_an_old_azure_api_version_dropped_the_choice(respx_mock: respx.MockRouter) -> None: + route: Final = respx_mock.post(url__regex=r"https://azure-bridge\.example\.com/.*chat/completions.*").mock( + return_value=httpx.Response( + 200, + json={ + "id": "chatcmpl-old-azure", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o", + "choices": [ + {"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "done"}} + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + ) + ) + + response: Final = litellm.responses( + model="azure/gpt-4o", + input="What is the weather?", + tools=list(_RESPONSES_TOOLS), + tool_choice=_FLAT_ALLOWED_WEATHER, + drop_params=True, + use_chat_completions_api=True, + api_key="fake", + api_base="https://azure-bridge.example.com", + api_version="2023-07-01-preview", + ) + + assert "tool_choice" not in json.loads(route.calls[0].request.content) + assert response.tool_choice == "auto" + + def test_streamed_responses_report_auto_for_an_allowed_tools_choice_it_dropped(respx_mock: respx.MockRouter) -> None: respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( return_value=httpx.Response( diff --git a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py index f8eb0378ed1..6d6546b1d7e 100644 --- a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py +++ b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py @@ -5,7 +5,7 @@ import logging import sys import warnings from collections.abc import Callable, Iterator, Mapping, Sequence -from contextlib import contextmanager, nullcontext +from contextlib import AbstractContextManager, contextmanager, nullcontext from datetime import datetime from types import MappingProxyType, SimpleNamespace from typing import Any, Final, List, cast @@ -1727,9 +1727,7 @@ async def _auto_execute_echo( return await _auto_execute(logging_obj, _ECHO_MCP_TOOL, "PING") -async def _auto_execute_through_mcp_server( - monkeypatch: pytest.MonkeyPatch, call_tool: AsyncMock, logging_obj: Logging -) -> None: +def _served_by_a_recording_mcp_server(monkeypatch: pytest.MonkeyPatch, call_tool: AsyncMock) -> None: monkeypatch.setitem( sys.modules, "litellm.proxy.proxy_server", SimpleNamespace(proxy_logging_obj=None, prisma_client=None) ) @@ -1743,15 +1741,31 @@ async def _auto_execute_through_mcp_server( get_mcp_server_by_name=MagicMock(return_value=None), ), ) - echo_tool: Final = MCPTool.model_validate( - {"name": "echo", "inputSchema": {"type": "object", "properties": {"value": {"type": "string"}}}}, - by_name=False, - ) - with patch.object( + + +def _echo_tool_listed() -> AbstractContextManager[object]: + return patch.object( LiteLLM_Proxy_MCP_Handler, "process_mcp_tools_without_openai_transform", - AsyncMock(return_value=([echo_tool], {"echo": "demo"})), - ): + AsyncMock( + return_value=( + [ + MCPTool.model_validate( + {"name": "echo", "inputSchema": {"type": "object", "properties": {"value": {"type": "string"}}}}, + by_name=False, + ) + ], + {"echo": "demo"}, + ) + ), + ) + + +async def _auto_execute_through_mcp_server( + monkeypatch: pytest.MonkeyPatch, call_tool: AsyncMock, logging_obj: Logging +) -> None: + _served_by_a_recording_mcp_server(monkeypatch, call_tool) + with _echo_tool_listed(): await litellm.aresponses( model="openai/gpt-5.6", input="Call echo with PING", @@ -1763,6 +1777,73 @@ async def _auto_execute_through_mcp_server( ) +_ALLOWED_ONLY_ANOTHER_TOOL: Final = { + "type": "allowed_tools", + "mode": "auto", + "tools": [{"type": "function", "name": "another_tool"}], +} + + +async def _call_with_a_choice_excluding_echo(stream: bool) -> tuple[object, ...]: + with _echo_tool_listed(): + response: Final = await litellm.aresponses( + model="openai/gpt-5.6", + input="Call echo with PING", + tools=[dict(_ECHO_MCP_TOOL)], + tool_choice=_ALLOWED_ONLY_ANOTHER_TOOL, + api_key="sk-mcp-rounds", + api_base="https://mcp-rounds.example.com/v1", + stream=stream, + ) + return await _drain(response) if stream else (response,) + + +@pytest.mark.asyncio +async def test_an_mcp_tool_the_choice_excludes_never_runs_when_the_model_ignored_the_choice( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + _serve_rounds(respx_mock, monkeypatch, (_ROUND_ONE_BODY,)) + call_tool: Final = AsyncMock() + _served_by_a_recording_mcp_server(monkeypatch, call_tool) + + with pytest.raises(litellm.BadRequestError, match="LiteLLM did not run the call"): + await _call_with_a_choice_excluding_echo(stream=False) + + call_tool.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_an_mcp_tool_the_choice_lists_still_runs(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch): + _serve_rounds(respx_mock, monkeypatch, (_ROUND_ONE_BODY, _ROUND_TWO_BODY)) + call_tool: Final = AsyncMock(return_value=CallToolResult(content=[TextContent(type="text", text="PING")])) + _served_by_a_recording_mcp_server(monkeypatch, call_tool) + + with _echo_tool_listed(): + await litellm.aresponses( + model="openai/gpt-5.6", + input="Call echo with PING", + tools=[dict(_ECHO_MCP_TOOL)], + tool_choice={"type": "allowed_tools", "mode": "auto", "tools": [{"type": "function", "name": "echo"}]}, + api_key="sk-mcp-rounds", + api_base="https://mcp-rounds.example.com/v1", + ) + + assert call_tool.await_count == 1 + + +@pytest.mark.asyncio +async def test_a_streamed_mcp_tool_the_choice_excludes_never_runs_when_the_model_ignored_the_choice( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + _serve_responses(respx_mock, monkeypatch, (_streamed(_ROUND_ONE_BODY),)) + call_tool: Final = AsyncMock() + _served_by_a_recording_mcp_server(monkeypatch, call_tool) + + await _call_with_a_choice_excluding_echo(stream=True) + + call_tool.assert_not_awaited() + + def _rounds_logging_obj(recorder: _RoundRecorder, call_id: str) -> Logging: return Logging( model="openai/gpt-5.6", diff --git a/tests/unit/responses/mcp/test_chat_completions_handler.py b/tests/unit/responses/mcp/test_chat_completions_handler.py index ecef6b44bae..c771b41ff64 100644 --- a/tests/unit/responses/mcp/test_chat_completions_handler.py +++ b/tests/unit/responses/mcp/test_chat_completions_handler.py @@ -7,9 +7,9 @@ import threading import types import warnings from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence -from contextlib import contextmanager +from contextlib import contextmanager, nullcontext from datetime import datetime -from types import MappingProxyType +from types import MappingProxyType, SimpleNamespace from typing import Final, cast from unittest.mock import AsyncMock, MagicMock, patch @@ -2782,3 +2782,93 @@ async def test_mcp_tools_keep_the_400_for_an_allowed_tools_choice_the_provider_c execute.assert_not_awaited() assert not respx_mock.calls + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +async def test_mcp_tools_keep_the_400_on_an_azure_api_version_that_drops_the_choice( + respx_mock: respx.MockRouter, stream: bool +): + execute: Final = AsyncMock(return_value=[]) + + with ( + patch.object( + LiteLLM_Proxy_MCP_Handler, + "process_mcp_tools_without_openai_transform", + AsyncMock(return_value=([_MCP_ECHO_TOOL], {"echo": "demo"})), + ), + patch.object(LiteLLM_Proxy_MCP_Handler, "execute_tool_calls", execute), + pytest.raises(litellm.BadRequestError, match="drop_params does not drop it"), + ): + await litellm.acompletion( + model="azure/gpt-4o", + messages=[{"role": "user", "content": "Call echo"}], + tools=[{"type": "mcp", "server_url": "litellm_proxy/mcp/demo", "require_approval": "never"}], + tool_choice=_FLAT_ALLOWED_ECHO, + drop_params=True, + stream=stream, + api_key="fake", + api_base="https://mcp-chat-rounds.example.com", + api_version="2023-07-01-preview", + ) + + execute.assert_not_awaited() + assert not respx_mock.calls + + +_ALLOWED_ONLY_ANOTHER_TOOL: Final = { + "type": "allowed_tools", + "mode": "auto", + "tools": [{"type": "function", "name": "another_tool"}], +} + + +def _served_by_a_recording_mcp_server(monkeypatch: pytest.MonkeyPatch, call_tool: AsyncMock) -> None: + monkeypatch.setitem( + sys.modules, "litellm.proxy.proxy_server", SimpleNamespace(proxy_logging_obj=None, prisma_client=None) + ) + monkeypatch.setattr( + "litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager", + SimpleNamespace( + catalog=SimpleNamespace(operation=nullcontext), + get_registry=MagicMock(return_value={}), + call_tool=call_tool, + get_mcp_server_from_tool_name=MagicMock(return_value=None), + get_mcp_server_by_name=MagicMock(return_value=None), + ), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +async def test_an_mcp_tool_the_choice_excludes_never_runs_when_the_model_ignored_the_choice( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, stream: bool +): + _serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE) if stream else _ok(_CHAT_ROUND_ONE),)) + call_tool: Final = AsyncMock() + _served_by_a_recording_mcp_server(monkeypatch, call_tool) + + async def call_with_the_choice() -> None: + response: Final = await litellm.acompletion( + model="openai/gpt-5.6", + messages=[{"role": "user", "content": "Call echo with PING"}], + tools=[{"type": "mcp", "server_url": "litellm_proxy/mcp/demo", "require_approval": "never"}], + tool_choice=_ALLOWED_ONLY_ANOTHER_TOOL, + api_key="sk-mcp-chat-rounds", + api_base="https://mcp-chat-rounds.example.com/v1", + stream=stream, + ) + if stream: + await _read_all(response) + + with ( + patch.object( + LiteLLM_Proxy_MCP_Handler, + "process_mcp_tools_without_openai_transform", + AsyncMock(return_value=([_MCP_ECHO_TOOL], {"echo": "demo"})), + ), + pytest.raises(litellm.BadRequestError, match="LiteLLM did not run the call"), + ): + await call_with_the_choice() + + call_tool.assert_not_awaited() diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 4f60440ff68..263cd376a8a 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -664,6 +664,22 @@ async def test_a_malformed_allowed_tools_tool_choice_stays_a_400_on_mcp_gateway_ assert not respx_mock.calls +@pytest.mark.asyncio +async def test_drop_params_keeps_the_400_when_the_provider_would_run_an_mcp_tool(respx_mock: respx.MockRouter): + with pytest.raises(litellm.BadRequestError, match="drop_params does not drop it") as exc: + await litellm.acompletion( + model="anthropic/claude-haiku-4-5", + messages=[{"role": "user", "content": "What time is it?"}], + tools=[{"type": "mcp", "server_label": "docs", "server_url": "https://mcp.example.com/sse"}], + tool_choice={"type": "allowed_tools", "mode": "auto", "tools": [{"type": "function", "name": "get_time"}]}, + drop_params=True, + api_key="sk-ant-fake", + ) + + assert exc.value.status_code == 400 + assert not respx_mock.calls + + _CLAUDE_TOOL_SEARCH_CALL: Final = { "id": "srvtoolu_01search", "type": "function", diff --git a/tests/unit/test_tool_loop.py b/tests/unit/test_tool_loop.py index 22bfd4ccf1d..9b724b644b6 100644 --- a/tests/unit/test_tool_loop.py +++ b/tests/unit/test_tool_loop.py @@ -11,6 +11,7 @@ from litellm.types.llms.openai import ChatCompletionToolMessage from litellm.types.utils import ChatCompletionMessageToolCall OPENAI_CHAT_COMPLETIONS_URL: Final = "https://api.openai.com/v1/chat/completions" +ANTHROPIC_MESSAGES_URL: Final = "https://api.anthropic.com/v1/messages" WEATHER_TOOLS: Final = ( { "type": "function", @@ -248,57 +249,147 @@ def test_stream_rejected_before_any_request(respx_mock: respx.MockRouter) -> Non assert route.call_count == 0 -_FLAT_ALLOWED_TOOLS: Final = {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "function", "name": "a"}]} +_ALLOWED_WEATHER_ONLY: Final = { + "type": "allowed_tools", + "mode": "auto", + "tools": [{"type": "function", "name": "get_weather"}], +} _EMPTY_ALLOWED_TOOLS: Final = {"type": "allowed_tools", "mode": "auto", "tools": []} +_TIME_CALL: Final = _tool_call("call_1", "get_time", {}) -@pytest.mark.parametrize( - ("model", "tool_choice"), - [("anthropic/claude-haiku-4-5", _FLAT_ALLOWED_TOOLS), ("openai/gpt-5-mini", _EMPTY_ALLOWED_TOOLS)], -) -def test_allowed_tools_choice_nothing_enforces_is_refused_under_drop_params( - respx_mock: respx.MockRouter, model: str, tool_choice: dict[str, object] -) -> None: +def _anthropic_message(*content: dict[str, object]) -> httpx.Response: + return httpx.Response( + 200, + json={ + "id": "msg_tool_loop", + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5", + "content": list(content), + "stop_reason": "tool_use" if content[0]["type"] == "tool_use" else "end_turn", + "usage": {"input_tokens": 10, "output_tokens": 5}, + }, + ) + + +def _anthropic_tool_use(name: str) -> httpx.Response: + return _anthropic_message({"type": "tool_use", "id": "toolu_1", "name": name, "input": {}}) + + +def _tools_run_before_the_refusal( + tool_choice: dict[str, object], **completion_kwargs: object +) -> list[ChatCompletionMessageToolCall]: executed: Final[list[ChatCompletionMessageToolCall]] = [] def execute(tool_call: ChatCompletionMessageToolCall) -> ChatCompletionToolMessage: executed.append(tool_call) return _tool_result(tool_call) - with pytest.raises(litellm.BadRequestError, match="drop_params does not drop it"): + with pytest.raises(litellm.BadRequestError, match="LiteLLM did not run the call"): litellm.run_tool_loop( - model=model, messages=[{"role": "user", "content": "hi"}], tools=WEATHER_TOOLS, execute_tool=execute, tool_choice=tool_choice, + api_key="sk-test", + **completion_kwargs, + ) + return executed + + +def test_a_tool_the_dropped_choice_excludes_never_runs(respx_mock: respx.MockRouter) -> None: + respx_mock.post(ANTHROPIC_MESSAGES_URL).mock(return_value=_anthropic_tool_use("get_time")) + + assert not _tools_run_before_the_refusal( + _ALLOWED_WEATHER_ONLY, model="anthropic/claude-haiku-4-5", drop_params=True + ) + assert "tool_choice" not in _request_bodies(respx_mock)[0] + + +def test_a_tool_the_dropped_choice_lists_still_runs(respx_mock: respx.MockRouter) -> None: + respx_mock.post(ANTHROPIC_MESSAGES_URL).mock( + side_effect=[_anthropic_tool_use("get_weather"), _anthropic_message({"type": "text", "text": "done"})] + ) + + assert ( + litellm.run_tool_loop( + model="anthropic/claude-haiku-4-5", + messages=[{"role": "user", "content": "hi"}], + tools=WEATHER_TOOLS, + execute_tool=_tool_result, + tool_choice=_ALLOWED_WEATHER_ONLY, + drop_params=True, + api_key="sk-test", + ) + == "done" + ) + + +@pytest.mark.parametrize( + ("tool_choice", "completion_kwargs"), + [(_ALLOWED_WEATHER_ONLY, {}), (_EMPTY_ALLOWED_TOOLS, {"drop_params": True})], +) +def test_a_tool_the_choice_excludes_never_runs_when_the_model_ignored_the_choice( + respx_mock: respx.MockRouter, tool_choice: dict[str, object], completion_kwargs: dict[str, object] +) -> None: + respx_mock.post(OPENAI_CHAT_COMPLETIONS_URL).mock( + return_value=httpx.Response(200, json=_openai_response(None, [_TIME_CALL])) + ) + + assert not _tools_run_before_the_refusal(tool_choice, model="openai/gpt-5-mini", **completion_kwargs) + + +def test_a_choice_on_a_model_alias_reaches_the_loop( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "model_alias_map", {"the-alias": "openai/gpt-5-mini"}) + respx_mock.post(OPENAI_CHAT_COMPLETIONS_URL).mock( + return_value=httpx.Response(200, json=_openai_response("done")) + ) + + assert ( + litellm.run_tool_loop( + model="the-alias", + messages=[{"role": "user", "content": "hi"}], + tools=WEATHER_TOOLS, + execute_tool=_tool_result, + tool_choice=_ALLOWED_WEATHER_ONLY, + api_key="sk-test", + ) + == "done" + ) + + +@pytest.mark.asyncio +async def test_async_loop_never_runs_a_tool_the_choice_excludes( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "module_level_aclient", AsyncHTTPHandler()) + respx_mock.post(OPENAI_CHAT_COMPLETIONS_URL).mock( + return_value=httpx.Response(200, json=_openai_response(None, [_TIME_CALL])) + ) + executed: Final[list[ChatCompletionMessageToolCall]] = [] + + async def execute(tool_call: ChatCompletionMessageToolCall) -> ChatCompletionToolMessage: + executed.append(tool_call) + return _tool_result(tool_call) + + with pytest.raises(litellm.BadRequestError, match="LiteLLM did not run the call"): + await litellm.arun_tool_loop( + model="openai/gpt-5-mini", + messages=[{"role": "user", "content": "hi"}], + tools=WEATHER_TOOLS, + execute_tool=execute, + tool_choice=_ALLOWED_WEATHER_ONLY, drop_params=True, api_key="sk-test", ) assert not executed - assert not respx_mock.calls - - -@pytest.mark.asyncio -async def test_async_loop_refuses_an_allowed_tools_choice_the_provider_cannot_enforce( - respx_mock: respx.MockRouter, -) -> None: - async def execute(tool_call: ChatCompletionMessageToolCall) -> ChatCompletionToolMessage: - return _tool_result(tool_call) - - with pytest.raises(litellm.BadRequestError, match="drop_params does not drop it"): - await litellm.arun_tool_loop( - model="anthropic/claude-haiku-4-5", - messages=[{"role": "user", "content": "hi"}], - tools=WEATHER_TOOLS, - execute_tool=execute, - tool_choice=_FLAT_ALLOWED_TOOLS, - drop_params=True, - api_key="sk-test", - ) - - assert not respx_mock.calls def test_allowed_tools_choice_the_provider_enforces_reaches_it_every_round(respx_mock: respx.MockRouter) -> None: From be0bb7a57e55712476b903a8244802024416bd96 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 5 Nov 2026 22:38:53 -0800 Subject: [PATCH 5/9] fix(tool-choice): drop the choice under drop_params and hand an excluded MCP call back unrun The drop now applies to every provider that cannot take an allowed_tools tool_choice, whatever the request's tools hold. LiteLLM never runs a call the choice excludes: on MCP paths the whole round goes back to the caller in a 200, the same shape as require_approval, and the tool loop keeps its 400 because it only returns final content. A malformed choice stays a 400 before any request on the tool loop and on chat with MCP gateway tools. --- .../allowed_tools_tool_choice.py | 63 +++---- litellm/main.py | 15 +- litellm/responses/main.py | 22 +-- .../responses/mcp/chat_completions_handler.py | 17 +- .../mcp/litellm_proxy_mcp_handler.py | 15 +- .../responses/mcp/mcp_streaming_iterator.py | 13 +- litellm/tool_loop.py | 12 +- .../test_allowed_tools_tool_choice.py | 44 ++--- .../test_anthropic_responses_bridge.py | 26 +-- .../mcp/test_aresponses_api_with_mcp.py | 82 ++++++--- .../mcp/test_chat_completions_handler.py | 166 ++++++++++-------- tests/unit/test_main.py | 36 ++-- tests/unit/test_tool_loop.py | 23 ++- 13 files changed, 289 insertions(+), 245 deletions(-) diff --git a/litellm/litellm_core_utils/allowed_tools_tool_choice.py b/litellm/litellm_core_utils/allowed_tools_tool_choice.py index ff067b8b8e4..0d836eb8aba 100644 --- a/litellm/litellm_core_utils/allowed_tools_tool_choice.py +++ b/litellm/litellm_core_utils/allowed_tools_tool_choice.py @@ -10,7 +10,6 @@ from litellm.exceptions import BadRequestError ALLOWED_TOOLS_TOOL_CHOICE_PROVIDERS: Final = frozenset({"openai", "azure", "litellm_proxy"}) _FIRST_AZURE_API_VERSION_WITH_TOOL_CHOICE: Final = "2023-12-01" -_TOOL_TYPES_THE_CALLER_RUNS: Final = frozenset({"function", "custom"}) CHAT_ALLOWED_TOOLS_LAYOUT: Final = ( '{"type": "allowed_tools", "allowed_tools": {"mode": "auto" | "required", "tools": []}}' @@ -59,19 +58,11 @@ class _ToolChoiceType(BaseModel): _TOOL_CHOICE_TYPE: Final = TypeAdapter(_ToolChoiceType) -def _declared_type(value: object) -> str | None: - try: - return _TOOL_CHOICE_TYPE.validate_python(value).type - except ValidationError: - return None - - def is_allowed_tools_tool_choice(tool_choice: object) -> bool: - return _declared_type(tool_choice) == "allowed_tools" - - -def caller_runs_every_tool(tools: Iterable[object] | None) -> bool: - return all(_declared_type(tool) in _TOOL_TYPES_THE_CALLER_RUNS for tool in tools or ()) + try: + return _TOOL_CHOICE_TYPE.validate_python(tool_choice).type == "allowed_tools" + except ValidationError: + return False def drops_unsupported_allowed_tools_tool_choice( @@ -79,14 +70,11 @@ def drops_unsupported_allowed_tools_tool_choice( custom_llm_provider: str, model: str, drop_params: bool, - tools: Iterable[object] | None, ) -> bool: if not is_allowed_tools_tool_choice(tool_choice) or custom_llm_provider in ALLOWED_TOOLS_TOOL_CHOICE_PROVIDERS: return False - if drop_params and caller_runs_every_tool(tools): - return True if drop_params: - raise _a_choice_drop_params_keeps(custom_llm_provider, model) + return True raise BadRequestError( message=( "tool_choice type 'allowed_tools' is only supported for " @@ -165,25 +153,6 @@ def unenforced_allowed_tools_tool_choice( ) -def _a_choice_drop_params_keeps(custom_llm_provider: str | None, model: str) -> BadRequestError: - return BadRequestError( - message=( - f"This allowed_tools tool_choice can't be enforced on custom_llm_provider={custom_llm_provider}, and the " - "caller does not run every tool on this request, so drop_params does not drop it. Send only the allowed " - "tools in `tools` with tool_choice='auto' or 'required'." - ), - model=model, - llm_provider=custom_llm_provider or "", - ) - - -def refuse_unenforced_allowed_tools_tool_choice( - tool_choice: object, custom_llm_provider: str | None, model: str, api_version: str | None = None -) -> None: - if unenforced_allowed_tools_tool_choice(tool_choice, custom_llm_provider, api_version): - raise _a_choice_drop_params_keeps(custom_llm_provider, model) - - def _allowed_tool_names(tool_choice: object) -> frozenset[object]: parsed: Final = parse_allowed_tools_tool_choice(tool_choice) if parsed is None: @@ -192,11 +161,27 @@ def _allowed_tool_names(tool_choice: object) -> frozenset[object]: return frozenset(tool.get("name") for tool in flattened if tool is not None) -def refuse_tool_calls_the_choice_excludes(tool_choice: object, tool_names: Iterable[str], model: str) -> None: +def refuse_a_malformed_allowed_tools_tool_choice(tool_choice: object, model: str) -> None: + if is_allowed_tools_tool_choice(tool_choice) and parse_allowed_tools_tool_choice(tool_choice) is None: + raise BadRequestError( + message=( + f"Invalid tool choice, tool_choice={tool_choice}. An allowed_tools tool_choice must be " + f"{CHAT_ALLOWED_TOOLS_LAYOUT} or {RESPONSES_ALLOWED_TOOLS_LAYOUT}" + ), + model=model, + llm_provider="", + ) + + +def tool_calls_the_choice_excludes(tool_choice: object, tool_names: Iterable[str]) -> tuple[str, ...]: if not is_allowed_tools_tool_choice(tool_choice): - return + return () allowed: Final = _allowed_tool_names(tool_choice) - excluded: Final = tuple(name for name in tool_names if name not in allowed) + return tuple(name for name in tool_names if name not in allowed) + + +def refuse_tool_calls_the_choice_excludes(tool_choice: object, tool_names: Iterable[str], model: str) -> None: + excluded: Final = tool_calls_the_choice_excludes(tool_choice, tool_names) if excluded: raise BadRequestError( message=( diff --git a/litellm/main.py b/litellm/main.py index b45a4ba8812..1b884bfeeb8 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -73,9 +73,7 @@ from litellm.constants import ( from litellm.exceptions import LiteLLMUnknownProvider from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.allowed_tools_tool_choice import ( - caller_runs_every_tool, drops_unsupported_allowed_tools_tool_choice, - refuse_unenforced_allowed_tools_tool_choice, ) from litellm.litellm_core_utils.asyncify import asyncify, run_async_function from litellm.litellm_core_utils.audio_utils.utils import ( @@ -5592,7 +5590,7 @@ def completion( drop_params=drops_unsupported_params( cast("Mapping[str, object]", kwargs).get("drop_params") # cast-ok: kwargs is the untyped request dict ) - and caller_runs_every_tool(tools), # pyright: ignore[reportUnknownArgumentType] # untyped tools list + and not may_name_litellm_mcp_gateway(tools), # pyright: ignore[reportUnknownArgumentType] # untyped tools list ) # validate optional params stop = validate_openai_optional_params(stop=stop) @@ -5890,16 +5888,6 @@ def completion( if not is_claude_tool_target(custom_llm_provider=custom_llm_provider, model=model): tools = _drop_anthropic_only_tool_keys(tools=tools) - if skip_mcp_handler: - refuse_unenforced_allowed_tools_tool_choice( - tool_choice=cast( # cast-ok: widening an untyped dict to object is always sound, the helper validates it with pydantic - object, tool_choice - ), - custom_llm_provider=custom_llm_provider, - model=model, - api_version=api_version, - ) - drops_tool_choice: Final = drops_unsupported_allowed_tools_tool_choice( tool_choice=cast( # cast-ok: widening an untyped dict to object is always sound, the helper validates it with pydantic object, tool_choice @@ -5907,7 +5895,6 @@ def completion( custom_llm_provider=custom_llm_provider, model=model, drop_params=drop_unsupported_reasoning, - tools=tools, # pyright: ignore[reportUnknownArgumentType] # untyped tools list ) if provider_specific_header is not None: diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 3dbcc299008..8f98f686b7e 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -29,7 +29,6 @@ from litellm.integrations.anthropic_cache_control_hook import ( configured_injection_points, routed_model_info, ) -from litellm.litellm_core_utils.allowed_tools_tool_choice import refuse_unenforced_allowed_tools_tool_choice from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.completion_timeout import CompletionTimeout, StreamKind from litellm.litellm_core_utils.core_helpers import drops_unsupported_params, normalize_drop_params @@ -457,7 +456,9 @@ async def aresponses_api_with_mcp( ######################################################### if should_auto_execute and isinstance(response, ResponsesAPIResponse): with first_round_billed_on_failure(logging_obj, response): - tool_calls: Final = LiteLLM_Proxy_MCP_Handler.extract_tool_calls_from_response(response=response) + tool_calls: Final = LiteLLM_Proxy_MCP_Handler.tool_calls_to_run( + LiteLLM_Proxy_MCP_Handler.extract_tool_calls_from_response(response=response), tool_choice + ) if tool_calls: user_api_key_auth = kwargs.get("litellm_metadata", {}).get("user_api_key_auth") @@ -478,7 +479,6 @@ async def aresponses_api_with_mcp( tool_server_map=tool_server_map, served_tools=original_mcp_tools, tool_calls=tool_calls, - tool_choice=tool_choice, user_api_key_auth=user_api_key_auth, mcp_auth_header=mcp_auth_header, mcp_server_auth_headers=mcp_server_auth_headers, @@ -676,15 +676,6 @@ def _bridge_kwargs( return MappingProxyType({key: value for key, value in kwargs.items() if key not in dropped_keys}) -def _refuse_a_choice_the_mcp_tool_run_would_ignore( - skip_mcp_handler: bool, tool_choice: object, custom_llm_provider: str | None, model: str, api_version: object -) -> None: - if skip_mcp_handler: - refuse_unenforced_allowed_tools_tool_choice( - tool_choice, custom_llm_provider, model, api_version if isinstance(api_version, str) else None - ) - - def _raise_untranslatable_responses_fields( failure: _UntranslatableResponsesFields, model: str, custom_llm_provider: str | None ) -> NoReturn: @@ -1773,13 +1764,6 @@ def responses( bridge_kwargs: Final = _bridge_kwargs(kwargs, responses_api_provider_config, allowed_openai_params) if isinstance(bridge_kwargs, _UntranslatableResponsesFields): _raise_untranslatable_responses_fields(bridge_kwargs, model, custom_llm_provider) - _refuse_a_choice_the_mcp_tool_run_would_ignore( - skip_mcp_handler, # pyright: ignore[reportUnknownArgumentType] # popped from untyped kwargs - tool_choice, - custom_llm_provider, - model, - kwargs.get("api_version"), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # untyped kwargs - ) return litellm_completion_transformation_handler.response_api_handler( previous_response_client_ids(previous_response_id, kwargs), model=model, diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index c716864d155..14050b014bc 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -2,6 +2,7 @@ import asyncio import logging +from collections.abc import Sequence from types import MappingProxyType from typing import TYPE_CHECKING, Final, cast @@ -38,7 +39,7 @@ class _MCPCompletionKwargs(TypedDict, total=False, extra_items=object): def _add_mcp_metadata_to_response( response: ModelResponse | CustomStreamWrapper, openai_tools: list | None, - tool_calls: list | None = None, + tool_calls: Sequence[object] | None = None, tool_results: list | None = None, ) -> None: """ @@ -253,7 +254,7 @@ async def acompletion_with_mcp( self.base_call_args = base_call_args self.request_tags = request_tags self.collected_chunks: list[ModelResponseStream] = [] - self.tool_calls: list | None = None + self.tool_calls: Sequence[object] | None = None self.tool_results: list | None = None self.complete_response: ModelResponse | None = None self.stream_exhausted = False @@ -490,8 +491,9 @@ async def acompletion_with_mcp( 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 + self.tool_calls = LiteLLM_Proxy_MCP_Handler.tool_calls_to_run( + LiteLLM_Proxy_MCP_Handler.extract_tool_calls_from_chat_response(response=complete_response), + kwargs.get("tool_choice"), ) if self.tool_calls: @@ -500,7 +502,6 @@ async def acompletion_with_mcp( tool_server_map=self.tool_server_map, served_tools=deduplicated_mcp_tools, tool_calls=self.tool_calls, - tool_choice=kwargs.get("tool_choice"), 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, @@ -671,7 +672,10 @@ async def acompletion_with_mcp( if not isinstance(initial_response, ModelResponse): return initial_response - tool_calls: Final = LiteLLM_Proxy_MCP_Handler.extract_tool_calls_from_chat_response(response=initial_response) + tool_calls: Final = LiteLLM_Proxy_MCP_Handler.tool_calls_to_run( + LiteLLM_Proxy_MCP_Handler.extract_tool_calls_from_chat_response(response=initial_response), + kwargs.get("tool_choice"), + ) if not tool_calls: _add_mcp_metadata_to_response( @@ -684,7 +688,6 @@ async def acompletion_with_mcp( tool_results: Final = await LiteLLM_Proxy_MCP_Handler.execute_tool_calls( tool_server_map=tool_server_map, tool_calls=tool_calls, - tool_choice=kwargs.get("tool_choice"), served_tools=deduplicated_mcp_tools, user_api_key_auth=user_api_key_auth, mcp_auth_header=mcp_auth_header, diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index d37ee58eb8f..6cda545259b 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -10,7 +10,7 @@ from typing_extensions import Unpack from litellm._logging import verbose_logger from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG -from litellm.litellm_core_utils.allowed_tools_tool_choice import refuse_tool_calls_the_choice_excludes +from litellm.litellm_core_utils.allowed_tools_tool_choice import tool_calls_the_choice_excludes from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy.mcp_server.catalog import catalog_operation, global_manager from litellm.proxy.mcp_server.utils import ( @@ -609,6 +609,13 @@ class LiteLLM_Proxy_MCP_Handler: _should_auto_execute_tools = should_auto_execute_tools + @staticmethod + def tool_calls_to_run(tool_calls: Sequence[object], tool_choice: object) -> Sequence[object]: + names: Final = tuple( + LiteLLM_Proxy_MCP_Handler.extract_tool_call_details(tool_call)[0] or "" for tool_call in tool_calls + ) + return () if tool_calls_the_choice_excludes(tool_choice, names) else tool_calls + @staticmethod def extract_tool_calls_from_response(response: ResponsesAPIResponse) -> list[object]: """Extract tool calls from the response output.""" @@ -741,7 +748,6 @@ class LiteLLM_Proxy_MCP_Handler: guardrail_context: Mapping[str, object] | None = None, served_tools: Sequence[MCPTool] | None = None, request_trace_id: str | None = None, - tool_choice: object = None, ) -> list[MCPToolResult]: """Execute tool calls and return results.""" from fastapi import HTTPException @@ -756,11 +762,6 @@ class LiteLLM_Proxy_MCP_Handler: ) from litellm.proxy.proxy_server import proxy_logging_obj - refuse_tool_calls_the_choice_excludes( - tool_choice, - tuple(LiteLLM_Proxy_MCP_Handler.extract_tool_call_details(tool_call)[0] or "" for tool_call in tool_calls), - "", - ) tool_results: Final[list[MCPToolResult]] = [] tool_call_id: str | None = None rules_obj: Final = Rules() diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index b5e3eee9ecf..44b07b15f2a 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -838,10 +838,14 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): try: # Extract tool calls from the response - if self.collected_response is not None: - tool_calls = LiteLLM_Proxy_MCP_Handler.extract_tool_calls_from_response(self.collected_response) - else: - tool_calls = [] + tool_calls: Final = ( + LiteLLM_Proxy_MCP_Handler.tool_calls_to_run( + LiteLLM_Proxy_MCP_Handler.extract_tool_calls_from_response(self.collected_response), + self.original_request_params.get("tool_choice"), + ) + if self.collected_response is not None + else () + ) if not tool_calls: return self.tool_call_round += 1 @@ -898,7 +902,6 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): tool_server_map=self.tool_server_map, served_tools=self.served_tools, tool_calls=tool_calls, - tool_choice=self.original_request_params.get("tool_choice"), 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, diff --git a/litellm/tool_loop.py b/litellm/tool_loop.py index 1db6b8c6000..9fc9af3a1c3 100644 --- a/litellm/tool_loop.py +++ b/litellm/tool_loop.py @@ -6,7 +6,10 @@ from typing import Final, TypeAlias, cast from typing_extensions import TypedDict, Unpack from litellm.constants import DEFAULT_TOOL_LOOP_MAX_ROUNDS -from litellm.litellm_core_utils.allowed_tools_tool_choice import refuse_tool_calls_the_choice_excludes +from litellm.litellm_core_utils.allowed_tools_tool_choice import ( + refuse_a_malformed_allowed_tools_tool_choice, + refuse_tool_calls_the_choice_excludes, +) from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionAssistantMessage, @@ -33,11 +36,12 @@ class ToolLoopMaxRoundsExceeded(RuntimeError): super().__init__(f"model still requested tool calls on round {max_rounds} of {max_rounds}") -def _validate_tool_loop_args(max_rounds: int, completion_kwargs: _ToolLoopCompletionKwargs) -> None: +def _validate_tool_loop_args(max_rounds: int, model: str, completion_kwargs: _ToolLoopCompletionKwargs) -> None: if max_rounds < 1: raise ValueError(f"max_rounds must be >= 1, got {max_rounds}") if completion_kwargs.get("stream"): raise ValueError("run_tool_loop requires whole responses; stream=True is not supported") + refuse_a_malformed_allowed_tools_tool_choice(completion_kwargs.get("tool_choice"), model) def _expect_model_response(response: object) -> ModelResponse: @@ -159,7 +163,7 @@ def run_tool_loop( """ import litellm - _validate_tool_loop_args(max_rounds, completion_kwargs) + _validate_tool_loop_args(max_rounds, model, completion_kwargs) history: tuple[AllMessageValues, ...] = tuple(messages) # rebind-ok: rounds append new turns for round_number in range(1, max_rounds + 1): message = ( @@ -211,7 +215,7 @@ async def arun_tool_loop( """ import litellm - _validate_tool_loop_args(max_rounds, completion_kwargs) + _validate_tool_loop_args(max_rounds, model, completion_kwargs) history: tuple[AllMessageValues, ...] = tuple(messages) # rebind-ok: rounds append new turns for round_number in range(1, max_rounds + 1): message = ( diff --git a/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py b/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py index 698c5dcf2f5..a32ebe3e9b5 100644 --- a/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py +++ b/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py @@ -12,8 +12,8 @@ from litellm.litellm_core_utils.allowed_tools_tool_choice import ( is_allowed_tools_tool_choice, chat_allowed_tools_tool_choice, parse_allowed_tools_tool_choice, + refuse_a_malformed_allowed_tools_tool_choice, refuse_tool_calls_the_choice_excludes, - refuse_unenforced_allowed_tools_tool_choice, unenforced_allowed_tools_tool_choice, responses_allowed_tools_tool_choice, ) @@ -165,11 +165,8 @@ _FLAT_CHOICE: Final = {"type": "allowed_tools", "mode": "required", "tools": [{" ({"type": "allowed_tools", "tools": [{"type": "function", "name": "get_time"}]}, "azure"), ], ) -def test_a_tool_runner_refuses_a_choice_nothing_would_enforce(tool_choice: object, custom_llm_provider: str): - with pytest.raises(BadRequestError, match="drop_params does not drop it") as exc: - refuse_unenforced_allowed_tools_tool_choice(tool_choice, custom_llm_provider, "the-model") - - assert exc.value.status_code == 400 +def test_a_choice_nothing_would_enforce_counts_as_unenforced(tool_choice: object, custom_llm_provider: str): + assert unenforced_allowed_tools_tool_choice(tool_choice, custom_llm_provider) @pytest.mark.parametrize( @@ -183,7 +180,7 @@ def test_a_tool_runner_refuses_a_choice_nothing_would_enforce(tool_choice: objec (None, "anthropic"), ], ) -def test_a_tool_runner_keeps_a_choice_the_provider_enforces_or_any_other_choice( +def test_a_choice_the_provider_enforces_or_any_other_choice_counts_as_enforced( tool_choice: object, custom_llm_provider: str ): assert not unenforced_allowed_tools_tool_choice(tool_choice, custom_llm_provider) @@ -232,24 +229,31 @@ def test_azure_reads_the_same_api_version_fallback_as_the_request(monkeypatch: p assert unenforced_allowed_tools_tool_choice(_FLAT_CHOICE, "azure") -_FUNCTION_TOOL: Final = {"type": "function", "function": {"name": "get_time"}} +def test_drop_params_drops_a_choice_the_provider_cannot_take(): + assert drops_unsupported_allowed_tools_tool_choice(_FLAT_CHOICE, "anthropic", "the-model", True) + + +def test_a_choice_the_provider_cannot_take_answers_400_without_drop_params(): + with pytest.raises(BadRequestError, match="only supported for") as exc: + drops_unsupported_allowed_tools_tool_choice(_FLAT_CHOICE, "anthropic", "the-model", False) + + assert exc.value.status_code == 400 @pytest.mark.parametrize( - "tools", + "tool_choice", [ - [_FUNCTION_TOOL, {"type": "mcp", "server_label": "docs", "server_url": "https://mcp.example.com"}], - [{"type": "web_search_20250305", "name": "web_search"}], - [{"googleSearch": {}}], - [{"type": ["function"]}], - ["get_time"], + {"type": "allowed_tools", "mode": "required", "tools": []}, + {"type": "allowed_tools", "tools": [{"type": "function", "name": "get_time"}]}, ], ) -def test_drop_params_keeps_the_400_when_the_provider_could_run_a_tool(tools: list[object]): - with pytest.raises(BadRequestError, match="drop_params does not drop it"): - drops_unsupported_allowed_tools_tool_choice(_FLAT_CHOICE, "anthropic", "the-model", True, tools) +def test_a_malformed_choice_is_refused(tool_choice: object): + with pytest.raises(BadRequestError, match="Invalid tool choice") as exc: + refuse_a_malformed_allowed_tools_tool_choice(tool_choice, "the-model") + + assert exc.value.status_code == 400 -@pytest.mark.parametrize("tools", [None, [], [_FUNCTION_TOOL, {"type": "custom", "custom": {"name": "patch"}}]]) -def test_drop_params_drops_the_choice_when_the_caller_runs_every_tool(tools: list[object] | None): - assert drops_unsupported_allowed_tools_tool_choice(_FLAT_CHOICE, "anthropic", "the-model", True, tools) +@pytest.mark.parametrize("tool_choice", [_FLAT_CHOICE, _CHAT_LAYOUT, "required", None]) +def test_a_valid_choice_or_any_other_choice_is_not_called_malformed(tool_choice: object): + assert refuse_a_malformed_allowed_tools_tool_choice(tool_choice, "the-model") is None diff --git a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py index 79a655fc913..983271c0b49 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py +++ b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py @@ -2134,18 +2134,22 @@ async def test_responses_reports_auto_for_an_allowed_tools_choice_it_dropped( @pytest.mark.usefixtures(httpx_transport.__name__) @pytest.mark.asyncio -async def test_responses_keep_the_400_when_the_provider_would_run_an_mcp_tool(respx_mock: respx.MockRouter) -> None: - with pytest.raises(litellm.BadRequestError, match="drop_params does not drop it"): - await litellm.aresponses( - model="anthropic/claude-haiku-4-5", - input="What is the weather?", - tools=[{"type": "mcp", "server_label": "docs", "server_url": "https://mcp.example.com/sse"}], - tool_choice=_FLAT_ALLOWED_WEATHER, - drop_params=True, - api_key="sk-ant-fake", - ) +async def test_responses_drop_the_choice_beside_a_tool_the_provider_runs(respx_mock: respx.MockRouter) -> None: + route: Final = respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + return_value=_anthropic_message({"type": "text", "text": "done"}) + ) - assert not respx_mock.calls + response: Final = await litellm.aresponses( + model="anthropic/claude-haiku-4-5", + input="What is the weather?", + tools=[{"type": "mcp", "server_label": "docs", "server_url": "https://mcp.example.com/sse"}], + tool_choice=_FLAT_ALLOWED_WEATHER, + drop_params=True, + api_key="sk-ant-fake", + ) + + assert "tool_choice" not in json.loads(route.calls[0].request.content) + assert response.tool_choice == "auto" @pytest.mark.usefixtures(httpx_transport.__name__) diff --git a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py index 6d6546b1d7e..28db003f4c1 100644 --- a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py +++ b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py @@ -31,6 +31,7 @@ from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Ha from litellm.responses.mcp.mcp_streaming_iterator import MCPEnhancedStreamingIterator from litellm.types.llms.openai import ( OpenAIMcpServerTool, + ResponseCompletedEvent, ResponsesAPIResponse, ResponsesAPIStreamingResponse, ToolParam, @@ -1696,30 +1697,6 @@ async def _auto_execute( return response -@pytest.mark.asyncio -@pytest.mark.parametrize("stream", [False, True]) -async def test_mcp_tools_keep_the_400_for_a_choice_the_chat_bridge_would_drop( - respx_mock: respx.MockRouter, stream: bool -) -> None: - async def call_with_the_choice() -> None: - response: Final = await litellm.aresponses( - model="anthropic/claude-haiku-4-5", - input="Call echo with PING", - tools=[dict(_ECHO_MCP_TOOL)], - tool_choice={"type": "allowed_tools", "mode": "auto", "tools": [{"type": "function", "name": "echo"}]}, - drop_params=True, - stream=stream, - api_key="sk-ant-fake", - ) - if stream: - [event async for event in response] - - with _echo_tool_served("PING"), pytest.raises(litellm.BadRequestError, match="drop_params does not drop it"): - await call_with_the_choice() - - assert not respx_mock.calls - - async def _auto_execute_echo( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, logging_obj: Logging ) -> ResponsesAPIResponse: @@ -1806,10 +1783,57 @@ async def test_an_mcp_tool_the_choice_excludes_never_runs_when_the_model_ignored call_tool: Final = AsyncMock() _served_by_a_recording_mcp_server(monkeypatch, call_tool) - with pytest.raises(litellm.BadRequestError, match="LiteLLM did not run the call"): - await _call_with_a_choice_excluding_echo(stream=False) + (response,) = await _call_with_a_choice_excluding_echo(stream=False) call_tool.assert_not_awaited() + assert isinstance(response, ResponsesAPIResponse) + assert response.status == "completed" + assert _output_types(response.output) == ("function_call",) + + +def _anthropic_call_to_echo() -> httpx.Response: + return httpx.Response( + 200, + json={ + "id": "msg_echo", + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5", + "content": [{"type": "tool_use", "id": "toolu_echo", "name": "echo", "input": {"value": "PING"}}], + "stop_reason": "tool_use", + "usage": {"input_tokens": 10, "output_tokens": 5}, + }, + ) + + +@pytest.mark.asyncio +async def test_a_dropped_choice_hands_back_an_mcp_call_it_excludes( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + litellm.in_memory_llm_clients_cache.flush_cache() + route: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock( + return_value=_anthropic_call_to_echo() + ) + call_tool: Final = AsyncMock() + _served_by_a_recording_mcp_server(monkeypatch, call_tool) + + with _echo_tool_listed(): + response: Final = await litellm.aresponses( + model="anthropic/claude-haiku-4-5", + input="Call echo with PING", + tools=[dict(_ECHO_MCP_TOOL)], + tool_choice=_ALLOWED_ONLY_ANOTHER_TOOL, + drop_params=True, + api_key="sk-ant-fake", + ) + + call_tool.assert_not_awaited() + assert "tool_choice" not in json.loads(route.calls[0].request.content) + assert isinstance(response, ResponsesAPIResponse) + assert response.tool_choice == "auto" + assert _output_types(response.output) == ("function_call",) @pytest.mark.asyncio @@ -1839,9 +1863,13 @@ async def test_a_streamed_mcp_tool_the_choice_excludes_never_runs_when_the_model call_tool: Final = AsyncMock() _served_by_a_recording_mcp_server(monkeypatch, call_tool) - await _call_with_a_choice_excluding_echo(stream=True) + events: Final = await _call_with_a_choice_excluding_echo(stream=True) call_tool.assert_not_awaited() + last: Final = events[-1] + assert isinstance(last, ResponseCompletedEvent) + assert _output_types(last.response.output) == ("function_call",) + assert not any("mcp_call" in str(getattr(event, "type", "")) for event in events) def _rounds_logging_obj(recorder: _RoundRecorder, call_id: str) -> Logging: diff --git a/tests/unit/responses/mcp/test_chat_completions_handler.py b/tests/unit/responses/mcp/test_chat_completions_handler.py index c771b41ff64..bf5c7c0846a 100644 --- a/tests/unit/responses/mcp/test_chat_completions_handler.py +++ b/tests/unit/responses/mcp/test_chat_completions_handler.py @@ -17,6 +17,7 @@ import httpx import pytest import respx from httpx import Response +from mcp.types import CallToolResult, TextContent from mcp.types import Tool as MCPTool from starlette.responses import StreamingResponse from starlette.types import Message @@ -2753,69 +2754,6 @@ _MCP_ECHO_TOOL: Final = MCPTool.model_validate( _FLAT_ALLOWED_ECHO: Final = {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "function", "name": "echo"}]} -@pytest.mark.asyncio -@pytest.mark.parametrize("require_approval", ["never", "always"]) -@pytest.mark.parametrize("stream", [False, True]) -async def test_mcp_tools_keep_the_400_for_an_allowed_tools_choice_the_provider_cannot_enforce( - respx_mock: respx.MockRouter, require_approval: str, stream: bool -): - execute: Final = AsyncMock(return_value=[]) - - with ( - patch.object( - LiteLLM_Proxy_MCP_Handler, - "process_mcp_tools_without_openai_transform", - AsyncMock(return_value=([_MCP_ECHO_TOOL], {"echo": "demo"})), - ), - patch.object(LiteLLM_Proxy_MCP_Handler, "execute_tool_calls", execute), - pytest.raises(litellm.BadRequestError, match="drop_params does not drop it"), - ): - await litellm.acompletion( - model="anthropic/claude-haiku-4-5", - messages=[{"role": "user", "content": "Call echo"}], - tools=[{"type": "mcp", "server_url": "litellm_proxy/mcp/demo", "require_approval": require_approval}], - tool_choice=_FLAT_ALLOWED_ECHO, - drop_params=True, - stream=stream, - api_key="sk-ant-fake", - ) - - execute.assert_not_awaited() - assert not respx_mock.calls - - -@pytest.mark.asyncio -@pytest.mark.parametrize("stream", [False, True]) -async def test_mcp_tools_keep_the_400_on_an_azure_api_version_that_drops_the_choice( - respx_mock: respx.MockRouter, stream: bool -): - execute: Final = AsyncMock(return_value=[]) - - with ( - patch.object( - LiteLLM_Proxy_MCP_Handler, - "process_mcp_tools_without_openai_transform", - AsyncMock(return_value=([_MCP_ECHO_TOOL], {"echo": "demo"})), - ), - patch.object(LiteLLM_Proxy_MCP_Handler, "execute_tool_calls", execute), - pytest.raises(litellm.BadRequestError, match="drop_params does not drop it"), - ): - await litellm.acompletion( - model="azure/gpt-4o", - messages=[{"role": "user", "content": "Call echo"}], - tools=[{"type": "mcp", "server_url": "litellm_proxy/mcp/demo", "require_approval": "never"}], - tool_choice=_FLAT_ALLOWED_ECHO, - drop_params=True, - stream=stream, - api_key="fake", - api_base="https://mcp-chat-rounds.example.com", - api_version="2023-07-01-preview", - ) - - execute.assert_not_awaited() - assert not respx_mock.calls - - _ALLOWED_ONLY_ANOTHER_TOOL: Final = { "type": "allowed_tools", "mode": "auto", @@ -2848,7 +2786,11 @@ async def test_an_mcp_tool_the_choice_excludes_never_runs_when_the_model_ignored call_tool: Final = AsyncMock() _served_by_a_recording_mcp_server(monkeypatch, call_tool) - async def call_with_the_choice() -> None: + with patch.object( + LiteLLM_Proxy_MCP_Handler, + "process_mcp_tools_without_openai_transform", + AsyncMock(return_value=([_MCP_ECHO_TOOL], {"echo": "demo"})), + ): response: Final = await litellm.acompletion( model="openai/gpt-5.6", messages=[{"role": "user", "content": "Call echo with PING"}], @@ -2858,17 +2800,93 @@ async def test_an_mcp_tool_the_choice_excludes_never_runs_when_the_model_ignored api_base="https://mcp-chat-rounds.example.com/v1", stream=stream, ) - if stream: - await _read_all(response) + chunks: Final = await _read_all(response) if stream else () - with ( - patch.object( - LiteLLM_Proxy_MCP_Handler, - "process_mcp_tools_without_openai_transform", - AsyncMock(return_value=([_MCP_ECHO_TOOL], {"echo": "demo"})), - ), - pytest.raises(litellm.BadRequestError, match="LiteLLM did not run the call"), + handed_back: Final = ( + tuple(call.function.name for chunk in chunks for call in chunk.choices[0].delta.tool_calls or ()) + if stream + else tuple(call.function.name for call in response.choices[0].message.tool_calls) + ) + assert handed_back == ("echo",) + call_tool.assert_not_awaited() + + +def _anthropic_call_to_echo() -> Response: + return Response( + 200, + json={ + "id": "msg_echo", + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5", + "content": [{"type": "tool_use", "id": "toolu_echo", "name": "echo", "input": {"value": "PING"}}], + "stop_reason": "tool_use", + "usage": {"input_tokens": 10, "output_tokens": 5}, + }, + ) + + +async def _call_echo_on_anthropic_under_drop_params( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, rounds: Sequence[Response], tool_choice: object +) -> tuple[ModelResponse, respx.Route]: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + litellm.in_memory_llm_clients_cache.flush_cache() + route: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock(side_effect=list(rounds)) + with patch.object( + LiteLLM_Proxy_MCP_Handler, + "process_mcp_tools_without_openai_transform", + AsyncMock(return_value=([_MCP_ECHO_TOOL], {"echo": "demo"})), ): - await call_with_the_choice() + response: Final = await litellm.acompletion( + model="anthropic/claude-haiku-4-5", + messages=[{"role": "user", "content": "Call echo with PING"}], + tools=[{"type": "mcp", "server_url": "litellm_proxy/mcp/demo", "require_approval": "never"}], + tool_choice=tool_choice, + drop_params=True, + api_key="sk-ant-fake", + ) + return response, route + + +@pytest.mark.asyncio +async def test_a_dropped_choice_hands_back_an_mcp_call_it_excludes( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + call_tool: Final = AsyncMock() + _served_by_a_recording_mcp_server(monkeypatch, call_tool) + + response, route = await _call_echo_on_anthropic_under_drop_params( + respx_mock, monkeypatch, (_anthropic_call_to_echo(),), _ALLOWED_ONLY_ANOTHER_TOOL + ) call_tool.assert_not_awaited() + assert "tool_choice" not in json.loads(route.calls[0].request.content) + assert [call.function.name for call in response.choices[0].message.tool_calls] == ["echo"] + + +@pytest.mark.asyncio +async def test_a_dropped_choice_still_runs_an_mcp_call_it_lists( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + call_tool: Final = AsyncMock(return_value=CallToolResult(content=[TextContent(type="text", text="PING")])) + _served_by_a_recording_mcp_server(monkeypatch, call_tool) + answer: Final = Response( + 200, + json={ + "id": "msg_answer", + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5", + "content": [{"type": "text", "text": "PING back"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 10, "output_tokens": 5}, + }, + ) + + response, _route = await _call_echo_on_anthropic_under_drop_params( + respx_mock, monkeypatch, (_anthropic_call_to_echo(), answer), _FLAT_ALLOWED_ECHO + ) + + assert call_tool.await_count == 1 + assert response.choices[0].message.content == "PING back" diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 263cd376a8a..eac7ac8ed60 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -664,20 +664,32 @@ async def test_a_malformed_allowed_tools_tool_choice_stays_a_400_on_mcp_gateway_ assert not respx_mock.calls -@pytest.mark.asyncio -async def test_drop_params_keeps_the_400_when_the_provider_would_run_an_mcp_tool(respx_mock: respx.MockRouter): - with pytest.raises(litellm.BadRequestError, match="drop_params does not drop it") as exc: - await litellm.acompletion( - model="anthropic/claude-haiku-4-5", - messages=[{"role": "user", "content": "What time is it?"}], - tools=[{"type": "mcp", "server_label": "docs", "server_url": "https://mcp.example.com/sse"}], - tool_choice={"type": "allowed_tools", "mode": "auto", "tools": [{"type": "function", "name": "get_time"}]}, - drop_params=True, - api_key="sk-ant-fake", +def test_drop_params_drops_the_choice_beside_a_tool_the_provider_runs(respx_mock: respx.MockRouter): + route: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock( + return_value=httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5", + "content": [{"type": "text", "text": "done"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, ) + ) - assert exc.value.status_code == 400 - assert not respx_mock.calls + litellm.completion( + model="anthropic/claude-haiku-4-5", + messages=[{"role": "user", "content": "What time is it?"}], + tools=[{"type": "mcp", "server_label": "docs", "server_url": "https://mcp.example.com/sse"}], + tool_choice={"type": "allowed_tools", "mode": "auto", "tools": [{"type": "function", "name": "get_time"}]}, + drop_params=True, + api_key="sk-ant-fake", + ) + + assert "tool_choice" not in json.loads(route.calls[0].request.content) _CLAUDE_TOOL_SEARCH_CALL: Final = { diff --git a/tests/unit/test_tool_loop.py b/tests/unit/test_tool_loop.py index 9b724b644b6..9d0539c80d2 100644 --- a/tests/unit/test_tool_loop.py +++ b/tests/unit/test_tool_loop.py @@ -326,18 +326,29 @@ def test_a_tool_the_dropped_choice_lists_still_runs(respx_mock: respx.MockRouter ) -@pytest.mark.parametrize( - ("tool_choice", "completion_kwargs"), - [(_ALLOWED_WEATHER_ONLY, {}), (_EMPTY_ALLOWED_TOOLS, {"drop_params": True})], -) def test_a_tool_the_choice_excludes_never_runs_when_the_model_ignored_the_choice( - respx_mock: respx.MockRouter, tool_choice: dict[str, object], completion_kwargs: dict[str, object] + respx_mock: respx.MockRouter, ) -> None: respx_mock.post(OPENAI_CHAT_COMPLETIONS_URL).mock( return_value=httpx.Response(200, json=_openai_response(None, [_TIME_CALL])) ) - assert not _tools_run_before_the_refusal(tool_choice, model="openai/gpt-5-mini", **completion_kwargs) + assert not _tools_run_before_the_refusal(_ALLOWED_WEATHER_ONLY, model="openai/gpt-5-mini") + + +def test_a_malformed_choice_is_refused_before_any_request(respx_mock: respx.MockRouter) -> None: + with pytest.raises(litellm.BadRequestError, match="Invalid tool choice"): + litellm.run_tool_loop( + model="openai/gpt-5-mini", + messages=[{"role": "user", "content": "hi"}], + tools=WEATHER_TOOLS, + execute_tool=_tool_result, + tool_choice=_EMPTY_ALLOWED_TOOLS, + drop_params=True, + api_key="sk-test", + ) + + assert not respx_mock.calls def test_a_choice_on_a_model_alias_reaches_the_loop( From 2ac020c070c7fdd47ec912f5ae334923ec85356c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 6 Nov 2026 02:09:06 -0800 Subject: [PATCH 6/9] fix(tool-choice): refuse an allowed_tools entry whose name is not a string An mcp or other non-function entry with a list or dict name passed the parser and crashed when the allow list was hashed. The parser now treats it as malformed on both layouts, so it answers the same 400 as an empty allow list, or drops under drop_params. --- .../allowed_tools_tool_choice.py | 8 ++++++-- .../test_allowed_tools_tool_choice.py | 16 ++++++++++++++++ .../mcp/test_aresponses_api_with_mcp.py | 11 +++++++++-- tests/unit/test_main.py | 14 ++++++++++++-- tests/unit/test_tool_loop.py | 10 ++++++++-- 5 files changed, 51 insertions(+), 8 deletions(-) diff --git a/litellm/litellm_core_utils/allowed_tools_tool_choice.py b/litellm/litellm_core_utils/allowed_tools_tool_choice.py index 0d836eb8aba..7ee53dcddab 100644 --- a/litellm/litellm_core_utils/allowed_tools_tool_choice.py +++ b/litellm/litellm_core_utils/allowed_tools_tool_choice.py @@ -92,6 +92,10 @@ def _has_name(selector: Mapping[str, object]) -> bool: return isinstance(name, str) and name != "" +def _names_on_other_tool_types_are_strings(tools: Iterable[Mapping[str, object]]) -> bool: + return all(tool.get("type") in ("function", "custom") or "name" not in tool or _has_name(tool) for tool in tools) + + def _chat_allowed_tool(tool: Mapping[str, object]) -> Mapping[str, object] | None: tool_type: Final = tool.get("type") if not isinstance(tool_type, str) or tool_type not in ("function", "custom"): @@ -120,9 +124,9 @@ def parse_allowed_tools_tool_choice(tool_choice: object) -> ChatAllowedToolsTool return None match parsed: case ChatAllowedToolsToolChoice(): - return parsed + return parsed if _names_on_other_tool_types_are_strings(parsed.allowed_tools.tools) else None case ResponsesAllowedToolsToolChoice(): - return _chat_layout(parsed) + return _chat_layout(parsed) if _names_on_other_tool_types_are_strings(parsed.tools) else None def azure_api_version_takes_tool_choice(api_version: str) -> bool: diff --git a/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py b/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py index a32ebe3e9b5..8b124818d6a 100644 --- a/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py +++ b/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py @@ -14,10 +14,20 @@ from litellm.litellm_core_utils.allowed_tools_tool_choice import ( parse_allowed_tools_tool_choice, refuse_a_malformed_allowed_tools_tool_choice, refuse_tool_calls_the_choice_excludes, + tool_calls_the_choice_excludes, unenforced_allowed_tools_tool_choice, responses_allowed_tools_tool_choice, ) +_CHOICES_WITH_A_NAME_THAT_IS_NOT_A_STRING: Final = ( + {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "mcp", "name": ["get_time"]}]}, + {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "mcp", "name": {"first": "get_time"}}]}, + {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "mcp", "name": 7}]}, + {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "mcp", "name": ""}]}, + {"type": "allowed_tools", "allowed_tools": {"mode": "auto", "tools": [{"type": "mcp", "name": ["get_time"]}]}}, + {"type": "allowed_tools", "allowed_tools": {"mode": "auto", "tools": [{"type": "mcp", "name": 7}]}}, +) + _CHAT_LAYOUT: Final[dict[str, object]] = { "type": "allowed_tools", "allowed_tools": { @@ -245,6 +255,7 @@ def test_a_choice_the_provider_cannot_take_answers_400_without_drop_params(): [ {"type": "allowed_tools", "mode": "required", "tools": []}, {"type": "allowed_tools", "tools": [{"type": "function", "name": "get_time"}]}, + *_CHOICES_WITH_A_NAME_THAT_IS_NOT_A_STRING, ], ) def test_a_malformed_choice_is_refused(tool_choice: object): @@ -254,6 +265,11 @@ def test_a_malformed_choice_is_refused(tool_choice: object): assert exc.value.status_code == 400 +@pytest.mark.parametrize("tool_choice", _CHOICES_WITH_A_NAME_THAT_IS_NOT_A_STRING) +def test_a_name_that_is_not_a_string_lists_no_tool(tool_choice: object): + assert tool_calls_the_choice_excludes(tool_choice, ("get_time",)) == ("get_time",) + + @pytest.mark.parametrize("tool_choice", [_FLAT_CHOICE, _CHAT_LAYOUT, "required", None]) def test_a_valid_choice_or_any_other_choice_is_not_called_malformed(tool_choice: object): assert refuse_a_malformed_allowed_tools_tool_choice(tool_choice, "the-model") is None diff --git a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py index 28db003f4c1..01b223f5e26 100644 --- a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py +++ b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py @@ -1807,8 +1807,15 @@ def _anthropic_call_to_echo() -> httpx.Response: @pytest.mark.asyncio +@pytest.mark.parametrize( + "tool_choice", + [ + _ALLOWED_ONLY_ANOTHER_TOOL, + {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "mcp", "name": ["echo"]}]}, + ], +) async def test_a_dropped_choice_hands_back_an_mcp_call_it_excludes( - respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch + tool_choice: dict[str, object], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") @@ -1824,7 +1831,7 @@ async def test_a_dropped_choice_hands_back_an_mcp_call_it_excludes( model="anthropic/claude-haiku-4-5", input="Call echo with PING", tools=[dict(_ECHO_MCP_TOOL)], - tool_choice=_ALLOWED_ONLY_ANOTHER_TOOL, + tool_choice=tool_choice, drop_params=True, api_key="sk-ant-fake", ) diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index eac7ac8ed60..7313a1600f9 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -648,13 +648,23 @@ def test_completion_drop_params_drops_an_allowed_tools_tool_choice_it_cannot_sen @pytest.mark.asyncio -async def test_a_malformed_allowed_tools_tool_choice_stays_a_400_on_mcp_gateway_tools(respx_mock: respx.MockRouter): +@pytest.mark.parametrize( + "tool_choice", + [ + {"type": "allowed_tools", "mode": "required", "tools": []}, + {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "mcp", "name": ["get_time"]}]}, + {"type": "allowed_tools", "allowed_tools": {"mode": "auto", "tools": [{"type": "mcp", "name": ["get_time"]}]}}, + ], +) +async def test_a_malformed_allowed_tools_tool_choice_stays_a_400_on_mcp_gateway_tools( + tool_choice: dict[str, object], respx_mock: respx.MockRouter +): with pytest.raises(litellm.BadRequestError, match="Invalid tool choice") as exc: await litellm.acompletion( model="openai/gpt-5.5", messages=[{"role": "user", "content": "What time is it?"}], tools=[{"type": "mcp", "server_url": "litellm_proxy/mcp/demo", "require_approval": "never"}], - tool_choice={"type": "allowed_tools", "mode": "required", "tools": []}, + tool_choice=tool_choice, drop_params=True, api_base="http://localhost:12346/v1", api_key="fake_api_key", diff --git a/tests/unit/test_tool_loop.py b/tests/unit/test_tool_loop.py index 9d0539c80d2..bf5c00ef53f 100644 --- a/tests/unit/test_tool_loop.py +++ b/tests/unit/test_tool_loop.py @@ -336,14 +336,20 @@ def test_a_tool_the_choice_excludes_never_runs_when_the_model_ignored_the_choice assert not _tools_run_before_the_refusal(_ALLOWED_WEATHER_ONLY, model="openai/gpt-5-mini") -def test_a_malformed_choice_is_refused_before_any_request(respx_mock: respx.MockRouter) -> None: +@pytest.mark.parametrize( + "tool_choice", + [_EMPTY_ALLOWED_TOOLS, {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "mcp", "name": ["get_time"]}]}], +) +def test_a_malformed_choice_is_refused_before_any_request( + tool_choice: dict[str, object], respx_mock: respx.MockRouter +) -> None: with pytest.raises(litellm.BadRequestError, match="Invalid tool choice"): litellm.run_tool_loop( model="openai/gpt-5-mini", messages=[{"role": "user", "content": "hi"}], tools=WEATHER_TOOLS, execute_tool=_tool_result, - tool_choice=_EMPTY_ALLOWED_TOOLS, + tool_choice=tool_choice, drop_params=True, api_key="sk-test", ) From 939034d8a4bc3a9f7d337471ca97738df933e300 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 6 Nov 2026 06:44:51 -0800 Subject: [PATCH 7/9] fix(tool-choice): honor server-level and null-name mcp entries and stop hooks running excluded calls --- .../allowed_tools_tool_choice.py | 46 +++- .../chat_completion_agentic_loop.py | 124 ++++++++--- .../litellm_core_utils/get_litellm_params.py | 3 +- litellm/main.py | 18 +- litellm/responses/main.py | 5 +- .../responses/mcp/chat_completions_handler.py | 4 + .../mcp/litellm_proxy_mcp_handler.py | 15 +- .../responses/mcp/mcp_streaming_iterator.py | 2 + litellm/responses/mcp/tool_precheck.py | 15 ++ litellm/types/litellm_params.py | 1 + .../test_allowed_tools_tool_choice.py | 43 ++++ .../test_chat_completion_agentic_loop.py | 196 +++++++++++++++++- .../mcp/test_chat_completions_handler.py | 152 ++++++++------ tests/unit/test_tool_loop.py | 1 - tests/unit/types/test_litellm_params.py | 1 + 15 files changed, 520 insertions(+), 106 deletions(-) diff --git a/litellm/litellm_core_utils/allowed_tools_tool_choice.py b/litellm/litellm_core_utils/allowed_tools_tool_choice.py index 7ee53dcddab..fd0ba2c8e4f 100644 --- a/litellm/litellm_core_utils/allowed_tools_tool_choice.py +++ b/litellm/litellm_core_utils/allowed_tools_tool_choice.py @@ -1,4 +1,5 @@ from collections.abc import Iterable, Mapping +from types import MappingProxyType from typing import Final, Literal from openai.types.chat.chat_completion_allowed_tool_choice_param import ChatCompletionAllowedToolChoiceParam @@ -92,8 +93,21 @@ def _has_name(selector: Mapping[str, object]) -> bool: return isinstance(name, str) and name != "" -def _names_on_other_tool_types_are_strings(tools: Iterable[Mapping[str, object]]) -> bool: - return all(tool.get("type") in ("function", "custom") or "name" not in tool or _has_name(tool) for tool in tools) +def _names_a_tool_or_nothing(tool: Mapping[str, object]) -> bool: + tool_type: Final = tool.get("type") + if not isinstance(tool_type, str) or tool_type not in ("function", "custom"): + return tool.get("name") is None or _has_name(tool) + match tool.get(tool_type): + case {"name": str()}: + return True + case {"name": object()}: + return False + case _: + return True + + +def _names_are_strings(tools: Iterable[Mapping[str, object]]) -> bool: + return all(map(_names_a_tool_or_nothing, tools)) def _chat_allowed_tool(tool: Mapping[str, object]) -> Mapping[str, object] | None: @@ -124,9 +138,9 @@ def parse_allowed_tools_tool_choice(tool_choice: object) -> ChatAllowedToolsTool return None match parsed: case ChatAllowedToolsToolChoice(): - return parsed if _names_on_other_tool_types_are_strings(parsed.allowed_tools.tools) else None + return parsed if _names_are_strings(parsed.allowed_tools.tools) else None case ResponsesAllowedToolsToolChoice(): - return _chat_layout(parsed) if _names_on_other_tool_types_are_strings(parsed.tools) else None + return _chat_layout(parsed) if _names_are_strings(parsed.tools) else None def azure_api_version_takes_tool_choice(api_version: str) -> bool: @@ -177,11 +191,31 @@ def refuse_a_malformed_allowed_tools_tool_choice(tool_choice: object, model: str ) -def tool_calls_the_choice_excludes(tool_choice: object, tool_names: Iterable[str]) -> tuple[str, ...]: +def _whole_server_label(tool: Mapping[str, object]) -> str | None: + label: Final = tool.get("server_label") + return label if tool.get("type") == "mcp" and tool.get("name") is None and isinstance(label, str) else None + + +def _whole_server_labels(tool_choice: object) -> frozenset[str]: + parsed: Final = parse_allowed_tools_tool_choice(tool_choice) + labels: Final = () if parsed is None else tuple(map(_whole_server_label, parsed.allowed_tools.tools)) + return frozenset(label for label in labels if label is not None) + + +def tool_calls_the_choice_excludes( + tool_choice: object, + tool_names: Iterable[str], + server_labels_by_tool: Mapping[str, frozenset[str]] = MappingProxyType({}), +) -> tuple[str, ...]: if not is_allowed_tools_tool_choice(tool_choice): return () allowed: Final = _allowed_tool_names(tool_choice) - return tuple(name for name in tool_names if name not in allowed) + whole_servers: Final = _whole_server_labels(tool_choice) + return tuple( + name + for name in tool_names + if name not in allowed and whole_servers.isdisjoint(server_labels_by_tool.get(name, frozenset())) + ) def refuse_tool_calls_the_choice_excludes(tool_choice: object, tool_names: Iterable[str], model: str) -> None: diff --git a/litellm/litellm_core_utils/chat_completion_agentic_loop.py b/litellm/litellm_core_utils/chat_completion_agentic_loop.py index ce98d2cbb01..42a8b3013ee 100644 --- a/litellm/litellm_core_utils/chat_completion_agentic_loop.py +++ b/litellm/litellm_core_utils/chat_completion_agentic_loop.py @@ -18,7 +18,12 @@ from litellm.litellm_core_utils.agentic_loop_settings import ( DEFAULT_MAX_AGENTIC_LOOPS, validated_max_agentic_loops, ) +from litellm.litellm_core_utils.allowed_tools_tool_choice import ( + is_allowed_tools_tool_choice, + tool_calls_the_choice_excludes, +) from litellm.litellm_core_utils.core_helpers import qualify_agentic_followup_model +from litellm.litellm_core_utils.get_litellm_params import AGENTIC_LOOP_TOOL_CHOICE_KEY from litellm.litellm_core_utils.hidden_params import get_hidden_params, served_from_cache from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import is_str_keyed_dict, object_items @@ -34,7 +39,11 @@ from litellm.types.integrations.custom_logger import ( is_interception_internal_key, stashed_stream_options, ) -from litellm.types.utils import ModelResponse +from litellm.types.utils import ( + ChatCompletionMessageCustomToolCall, + ChatCompletionMessageToolCall, + ModelResponse, +) from litellm.utils import CustomStreamWrapper _DispatchResult: Final = TypeVar("_DispatchResult") @@ -181,6 +190,32 @@ def _end_refused_turn( return finalized if replay is None else replay +def _tool_call_name(call: ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall) -> str: + match call: + case ChatCompletionMessageCustomToolCall(): + return call.custom.name + case ChatCompletionMessageToolCall(): + return call.function.name or "" + + +def _called_tool_names(response: ModelResponse) -> tuple[str, ...]: + calls: Final = chain.from_iterable(choice.message.tool_calls or () for choice in response.choices) + return tuple(map(_tool_call_name, calls)) + + +def _holds_a_call_the_choice_excludes(response: ModelResponse, kwargs: Mapping[str, object]) -> bool: + return bool(tool_calls_the_choice_excludes(kwargs.get(AGENTIC_LOOP_TOOL_CHOICE_KEY), _called_tool_names(response))) + + +def _with_caller_tool_choice( + kwargs: Mapping[str, object], optional_params: Mapping[str, object] +) -> Mapping[str, object]: + tool_choice: Final = optional_params.get("tool_choice") + if AGENTIC_LOOP_TOOL_CHOICE_KEY in kwargs or not is_allowed_tools_tool_choice(tool_choice): + return kwargs + return {**kwargs, AGENTIC_LOOP_TOOL_CHOICE_KEY: tool_choice} + + def _request_tools(optional_params: Mapping[str, object]) -> list[dict[str, object]]: return [tool for tool in object_items(optional_params.get("tools")) if is_str_keyed_dict(tool)] @@ -238,6 +273,34 @@ async def _refused_chat_turn( return finalize_refused_chat_turn(response, claimed) +def _fingerprint_of_a_turn_to_run( + *, + response: ModelResponse, + request_kwargs: Mapping[str, object], + tool_calls: object, + fingerprints: Sequence[str], + depth: int, + max_loops: int, + model: str, + stream: bool, +) -> str | None: + if _holds_a_call_the_choice_excludes(response, request_kwargs): + return None + try: + return _check_agentic_loop_safety( + tool_calls=tool_calls, + fingerprints=fingerprints, + depth=depth, + max_loops=max_loops, + model=model, + ) + except AgenticLoopSafetyError as e: + if stream: + raise + verbose_logger.warning("LiteLLM.AgenticLoopRefused: ending turn [rail=%s]", type(e).__name__) + return None + + def _with_agentic_loop_metadata(kwargs_for_followup: Mapping[str, object]) -> Mapping[str, object]: metadata: Final = kwargs_for_followup.get("litellm_metadata") return MappingProxyType( @@ -371,7 +434,8 @@ async def maybe_run_chat_completion_agentic_loop( stream: bool, ) -> ModelResponse | CustomStreamWrapper | None: callbacks: Final = _custom_loggers(logging_obj) - depth, max_loops, fingerprints = _agentic_loop_settings(kwargs) + request_kwargs: Final = _with_caller_tool_choice(kwargs, optional_params) + depth, max_loops, fingerprints = _agentic_loop_settings(request_kwargs) tools: Final = _request_tools(optional_params) for callback in callbacks: @@ -379,7 +443,7 @@ async def maybe_run_chat_completion_agentic_loop( continue hook_kwargs = { - **kwargs, + **request_kwargs, "_agentic_loop_api_surface": CHAT_COMPLETION_AGENTIC_SURFACE, "custom_llm_provider": custom_llm_provider, } @@ -403,18 +467,17 @@ async def maybe_run_chat_completion_agentic_loop( if not should_run: continue - try: - fingerprint = _check_agentic_loop_safety( - tool_calls=tool_calls, - fingerprints=fingerprints, - depth=depth, - max_loops=max_loops, - model=model, - ) - except AgenticLoopSafetyError as e: - if stream: - raise - verbose_logger.warning("LiteLLM.AgenticLoopRefused: ending turn [rail=%s]", type(e).__name__) + fingerprint = _fingerprint_of_a_turn_to_run( + response=response, + request_kwargs=request_kwargs, + tool_calls=tool_calls, + fingerprints=fingerprints, + depth=depth, + max_loops=max_loops, + model=model, + stream=stream, + ) + if fingerprint is None: return _end_refused_turn( await _refused_chat_turn( callbacks, @@ -425,9 +488,9 @@ async def maybe_run_chat_completion_agentic_loop( messages=messages, tools=tools, custom_llm_provider=custom_llm_provider, - kwargs=kwargs, + kwargs=request_kwargs, ), - kwargs=kwargs, + kwargs=request_kwargs, depth=depth, model=model, custom_llm_provider=custom_llm_provider, @@ -471,7 +534,7 @@ async def maybe_run_chat_completion_agentic_loop( callback=callback, model=model, optional_params=optional_params, - kwargs=kwargs, + kwargs=dict(request_kwargs), logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, depth=depth, @@ -482,7 +545,7 @@ async def maybe_run_chat_completion_agentic_loop( return _converted_stream_replay( response, - kwargs=kwargs, + kwargs=request_kwargs, depth=depth, model=model, custom_llm_provider=custom_llm_provider, @@ -526,18 +589,17 @@ async def run_chat_completion_agentic_hooks( if not should_run: continue - try: - fingerprint = _check_agentic_loop_safety( - tool_calls=tool_calls, - fingerprints=fingerprints, - depth=depth, - max_loops=max_loops, - model=model, - ) - except AgenticLoopSafetyError as e: - if stream: - raise - verbose_logger.warning("LiteLLM.AgenticLoopRefused: ending turn [rail=%s]", type(e).__name__) + fingerprint = _fingerprint_of_a_turn_to_run( + response=response, + request_kwargs=kwargs, + tool_calls=tool_calls, + fingerprints=fingerprints, + depth=depth, + max_loops=max_loops, + model=model, + stream=stream, + ) + if fingerprint is None: return await _refused_chat_turn( callbacks, callback, diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index d0ac5c00513..cd90706a26f 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -33,8 +33,9 @@ AWS_CREDENTIAL_KWARGS_KEYS: Final = frozenset( ) PROVIDER_AFFINITY_HEADER_KWARG_KEY: Final = "provider_affinity_header" +AGENTIC_LOOP_TOOL_CHOICE_KEY: Final = "_agentic_loop_tool_choice" AGENTIC_LOOP_CONTROL_KWARGS_KEYS: Final = frozenset( - {"_agentic_loop_depth", "_agentic_loop_fingerprints", "max_agentic_loops"} + {"_agentic_loop_depth", "_agentic_loop_fingerprints", "max_agentic_loops", AGENTIC_LOOP_TOOL_CHOICE_KEY} ) OAUTH_TOKEN_EXCHANGE_KWARGS_KEYS: Final = frozenset( { diff --git a/litellm/main.py b/litellm/main.py index 1b884bfeeb8..77220aba5ac 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -74,6 +74,7 @@ from litellm.exceptions import LiteLLMUnknownProvider from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.allowed_tools_tool_choice import ( drops_unsupported_allowed_tools_tool_choice, + is_allowed_tools_tool_choice, ) from litellm.litellm_core_utils.asyncify import asyncify, run_async_function from litellm.litellm_core_utils.audio_utils.utils import ( @@ -88,6 +89,7 @@ from litellm.litellm_core_utils.completion_timeout import CompletionTimeout, Str from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.get_litellm_params import ( AGENTIC_LOOP_CONTROL_KWARGS_KEYS, + AGENTIC_LOOP_TOOL_CHOICE_KEY, AWS_CREDENTIAL_KWARGS_KEYS, OAUTH_TOKEN_EXCHANGE_KWARGS_KEYS, PROVIDER_AFFINITY_HEADER_KWARG_KEY, @@ -5584,6 +5586,9 @@ def completion( messages = validate_and_fix_openai_messages(messages=messages) tools = validate_and_fix_openai_tools(tools=tools) # validate tool_choice + requested_tool_choice: Final = cast( # cast-ok: widening an untyped dict to object is always sound + object, tool_choice + ) tool_choice = validate_chat_completion_tool_choice( tool_choice=tool_choice, model=model, @@ -6201,7 +6206,18 @@ def completion( if key in kwargs }, ) - litellm_params: Final = with_control_options(requested_litellm_params, control_options) + litellm_params: Final = with_control_options( + { + **requested_litellm_params, + **( + {AGENTIC_LOOP_TOOL_CHOICE_KEY: requested_tool_choice} + if AGENTIC_LOOP_TOOL_CHOICE_KEY not in kwargs + and is_allowed_tools_tool_choice(requested_tool_choice) + else {} + ), + }, + control_options, + ) affinity_api_base: Final = ( str(client.base_url) if isinstance(client, openai.OpenAI | openai.AsyncOpenAI) diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 8f98f686b7e..f70f9758cb5 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -457,7 +457,10 @@ async def aresponses_api_with_mcp( if should_auto_execute and isinstance(response, ResponsesAPIResponse): with first_round_billed_on_failure(logging_obj, response): tool_calls: Final = LiteLLM_Proxy_MCP_Handler.tool_calls_to_run( - LiteLLM_Proxy_MCP_Handler.extract_tool_calls_from_response(response=response), tool_choice + LiteLLM_Proxy_MCP_Handler.extract_tool_calls_from_response(response=response), + tool_choice, + tool_server_map, + mcp_tools_with_litellm_proxy, ) if tool_calls: diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index 14050b014bc..f622362dce4 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -494,6 +494,8 @@ async def acompletion_with_mcp( self.tool_calls = LiteLLM_Proxy_MCP_Handler.tool_calls_to_run( LiteLLM_Proxy_MCP_Handler.extract_tool_calls_from_chat_response(response=complete_response), kwargs.get("tool_choice"), + self.tool_server_map, + mcp_tools_with_litellm_proxy, ) if self.tool_calls: @@ -675,6 +677,8 @@ async def acompletion_with_mcp( tool_calls: Final = LiteLLM_Proxy_MCP_Handler.tool_calls_to_run( LiteLLM_Proxy_MCP_Handler.extract_tool_calls_from_chat_response(response=initial_response), kwargs.get("tool_choice"), + tool_server_map, + mcp_tools_with_litellm_proxy, ) if not tool_calls: diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 6cda545259b..d1188e2ecbc 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -1,6 +1,7 @@ import traceback from collections.abc import Awaitable, Callable, Collection, Iterable, Mapping, Sequence from datetime import datetime +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypedDict, overload from openai.types.chat import ChatCompletionToolParam @@ -26,6 +27,7 @@ from litellm.responses.mcp.tool_precheck import ( names_gateway_explicitly, names_whole_gateway, proxy_path_mcp_name, + server_labels_serving, ) from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator from litellm.types.llms.openai import ( @@ -610,11 +612,20 @@ class LiteLLM_Proxy_MCP_Handler: _should_auto_execute_tools = should_auto_execute_tools @staticmethod - def tool_calls_to_run(tool_calls: Sequence[object], tool_choice: object) -> Sequence[object]: + def tool_calls_to_run( + tool_calls: Sequence[object], + tool_choice: object, + tool_server_map: Mapping[str, str], + mcp_tools: Iterable[object], + ) -> Sequence[object]: names: Final = tuple( LiteLLM_Proxy_MCP_Handler.extract_tool_call_details(tool_call)[0] or "" for tool_call in tool_calls ) - return () if tool_calls_the_choice_excludes(tool_choice, names) else tool_calls + references: Final = tuple(mcp_tools) + labels_by_tool: Final = MappingProxyType( + {name: server_labels_serving(references, tool_server_map.get(name)) for name in names} + ) + return () if tool_calls_the_choice_excludes(tool_choice, names, labels_by_tool) else tool_calls @staticmethod def extract_tool_calls_from_response(response: ResponsesAPIResponse) -> list[object]: diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index 44b07b15f2a..2f366837b3b 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -842,6 +842,8 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): LiteLLM_Proxy_MCP_Handler.tool_calls_to_run( LiteLLM_Proxy_MCP_Handler.extract_tool_calls_from_response(self.collected_response), self.original_request_params.get("tool_choice"), + self.tool_server_map, + self.mcp_tools_with_litellm_proxy, ) if self.collected_response is not None else () diff --git a/litellm/responses/mcp/tool_precheck.py b/litellm/responses/mcp/tool_precheck.py index a67d81e70ae..cbd8a36d049 100644 --- a/litellm/responses/mcp/tool_precheck.py +++ b/litellm/responses/mcp/tool_precheck.py @@ -29,6 +29,21 @@ def names_whole_gateway(tool: object) -> bool: ) +def _server_label_serving(tool: object, server_name: str) -> str | None: + match tool: + case {"type": "mcp", "server_label": str() as label} if ( + names_whole_gateway(tool) or _mcp_server_url(tool) == f"{LITELLM_PROXY_MCP_SERVER_URL_PREFIX}{server_name}" + ): + return label + case _: + return None + + +def server_labels_serving(mcp_tools: Iterable[object], server_name: str | None) -> frozenset[str]: + labels: Final = () if server_name is None else tuple(_server_label_serving(tool, server_name) for tool in mcp_tools) + return frozenset(label for label in labels if label is not None) + + def proxy_path_mcp_name(tool: object) -> str | None: server_url: Final = _mcp_server_url(tool) match: Final = None if server_url is None else _PROXY_MCP_PATH_RE.match(server_url) diff --git a/litellm/types/litellm_params.py b/litellm/types/litellm_params.py index 9f63d27f5f3..d65a9a31e61 100644 --- a/litellm/types/litellm_params.py +++ b/litellm/types/litellm_params.py @@ -361,6 +361,7 @@ class AgenticLoopState: api_surface: Literal["chat_completions", "responses"] | None = field( default=None, metadata=wire("_agentic_loop_api_surface") ) + tool_choice: Mapping[str, object] | None = field(default=None, metadata=wire("_agentic_loop_tool_choice")) code_interpreter_active: bool | None = field(default=None, metadata=wire("_code_interpreter_interception_active")) code_interpreter_sandbox_key: str | None = field( default=None, metadata=wire("_code_interpreter_interception_sandbox_key") diff --git a/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py b/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py index 8b124818d6a..7c19e0b673c 100644 --- a/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py +++ b/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py @@ -26,8 +26,15 @@ _CHOICES_WITH_A_NAME_THAT_IS_NOT_A_STRING: Final = ( {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "mcp", "name": ""}]}, {"type": "allowed_tools", "allowed_tools": {"mode": "auto", "tools": [{"type": "mcp", "name": ["get_time"]}]}}, {"type": "allowed_tools", "allowed_tools": {"mode": "auto", "tools": [{"type": "mcp", "name": 7}]}}, + { + "type": "allowed_tools", + "allowed_tools": {"mode": "auto", "tools": [{"type": "function", "function": {"name": ["get_time"]}}]}, + }, + {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "function", "function": {"name": 7}}]}, ) +_WHOLE_DOCS_SERVER: Final = {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "mcp", "server_label": "docs"}]} + _CHAT_LAYOUT: Final[dict[str, object]] = { "type": "allowed_tools", "allowed_tools": { @@ -273,3 +280,39 @@ def test_a_name_that_is_not_a_string_lists_no_tool(tool_choice: object): @pytest.mark.parametrize("tool_choice", [_FLAT_CHOICE, _CHAT_LAYOUT, "required", None]) def test_a_valid_choice_or_any_other_choice_is_not_called_malformed(tool_choice: object): assert refuse_a_malformed_allowed_tools_tool_choice(tool_choice, "the-model") is None + + +@pytest.mark.parametrize( + "tool_choice", + [ + {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "mcp", "server_label": "docs", "name": None}]}, + { + "type": "allowed_tools", + "allowed_tools": {"mode": "auto", "tools": [{"type": "mcp", "server_label": "docs", "name": None}]}, + }, + ], +) +def test_an_mcp_entry_with_a_null_name_is_a_valid_choice(tool_choice: object): + assert parse_allowed_tools_tool_choice(tool_choice) is not None + assert refuse_a_malformed_allowed_tools_tool_choice(tool_choice, "the-model") is None + + +def test_an_mcp_entry_without_a_name_allows_every_call_its_server_serves(): + assert tool_calls_the_choice_excludes(_WHOLE_DOCS_SERVER, ("search",), {"search": frozenset({"docs"})}) == () + + +@pytest.mark.parametrize( + ("tool_choice", "server_labels"), + [ + (_WHOLE_DOCS_SERVER, frozenset({"other"})), + (_WHOLE_DOCS_SERVER, frozenset()), + ( + {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "mcp", "server_label": "docs", "name": "read"}]}, + frozenset({"docs"}), + ), + ], +) +def test_an_mcp_entry_allows_no_call_outside_its_server_or_its_name( + tool_choice: object, server_labels: frozenset[str] +): + assert tool_calls_the_choice_excludes(tool_choice, ("search",), {"search": server_labels}) == ("search",) diff --git a/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py b/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py index d26da679246..20bd02811fc 100644 --- a/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py +++ b/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py @@ -22,7 +22,7 @@ they stay out of the body even without it. import datetime import json -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence from typing import Any, Dict, Final, List, Optional, Tuple from unittest.mock import AsyncMock, MagicMock, patch @@ -38,6 +38,7 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.chat_completion_agentic_loop import ( _execute_chat_completion_agentic_plan, maybe_run_chat_completion_agentic_loop, + run_chat_completion_agentic_hooks, ) from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper @@ -67,6 +68,7 @@ _INTERNAL_CONTROL_FIELDS = ( "_code_interpreter_interception_active", "_code_interpreter_interception_sandbox_key", "_code_interpreter_interception_converted_stream", + "_agentic_loop_tool_choice", "litellm_metadata", ) @@ -132,6 +134,27 @@ def _tool_call_model_response() -> ModelResponse: ) +def _custom_tool_call_model_response() -> ModelResponse: + return ModelResponse( + choices=[ + Choices( + finish_reason="tool_calls", + message=Message( + role="assistant", + content=None, + tool_calls=[ + { + "id": "call_abc", + "type": "custom", + "custom": {"name": "litellm_code_execution", "input": "print(6*7)"}, + } + ], + ), + ) + ] + ) + + def _plain_model_response(content: str = "The answer is 42") -> ModelResponse: return ModelResponse( choices=[ @@ -1090,3 +1113,174 @@ class TestAgenticFollowUpKeepsTheProviderPrefix: response: Final = await self._run_followup("openai/gpt-4o", "openrouter", patched_model="openai/gpt-4o") assert response._hidden_params["custom_llm_provider"] == "openrouter" + + +_ALLOWS_ONLY_GET_WEATHER: Final = { + "type": "allowed_tools", + "mode": "auto", + "tools": [{"type": "function", "name": "get_weather"}], +} +_ALLOWS_THE_HOOKS_CALL: Final = { + "type": "allowed_tools", + "mode": "auto", + "tools": [{"type": "function", "name": "litellm_code_execution"}], +} + + +async def _dispatch_under_the_choice(tool_choice: object, acompletion_mock: AsyncMock) -> Any: + gate: Final = _GateOnlyLogger( + plan=AgenticLoopPlan( + run_agentic_loop=True, request_patch=AgenticLoopRequestPatch(messages=_patched_messages()) + ), + tool_calls=_REFUSED_GATE_TOOL_CALLS, + ) + with patch.object(litellm, "acompletion", acompletion_mock): + return await maybe_run_chat_completion_agentic_loop( + response=_tool_call_model_response(), + model="gpt-4o-mini", + messages=[{"role": "user", "content": "what is 6*7?"}], + optional_params={"tool_choice": tool_choice}, + kwargs={}, + logging_obj=_LoggingWithCallbacks((gate,)), + custom_llm_provider="openai", + stream=False, + ) + + +@pytest.mark.asyncio +async def test_an_intercepted_call_the_choice_excludes_ends_the_turn_without_running(restore_callbacks): + acompletion_mock: Final = AsyncMock() + + result: Final = await _dispatch_under_the_choice(_ALLOWS_ONLY_GET_WEATHER, acompletion_mock) + + acompletion_mock.assert_not_awaited() + assert isinstance(result, ModelResponse) + assert not result.choices[0].message.tool_calls + assert result.choices[0].finish_reason == "stop" + + +@pytest.mark.asyncio +async def test_an_intercepted_call_the_choice_lists_still_runs(restore_callbacks): + followup: Final = _plain_model_response("done") + acompletion_mock: Final = AsyncMock(return_value=followup) + + result: Final = await _dispatch_under_the_choice(_ALLOWS_THE_HOOKS_CALL, acompletion_mock) + + acompletion_mock.assert_awaited_once() + assert result is followup + + +class _ChatHookLogger(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.runs = 0 + + async def async_should_run_chat_completion_agentic_loop( + self, + response: object, + model: str, + messages: list[dict], + tools: list[dict] | None, + stream: bool, + custom_llm_provider: str, + kwargs: dict, + ) -> tuple[bool, dict]: + return True, {"tool_calls": [{"id": "call_abc", "name": "litellm_code_execution"}]} + + async def async_run_chat_completion_agentic_loop( + self, + tools: dict, + model: str, + messages: list[dict], + response: object, + optional_params: dict, + logging_obj: Any, + stream: bool, + kwargs: dict, + ) -> object: + self.runs += 1 + return _plain_model_response("ran the hook") + + +def _logging_with(hook: CustomLogger) -> Logging: + return Logging( + model="gpt-4o-mini", + messages=[], + stream=False, + call_type="acompletion", + start_time=datetime.datetime.now(), + litellm_call_id="call-test", + function_id="function-test", + dynamic_success_callbacks=[hook], + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model_response", [_tool_call_model_response, _custom_tool_call_model_response]) +async def test_a_chat_hook_never_runs_a_call_the_choice_excludes(model_response: Callable[[], ModelResponse]): + hook: Final = _ChatHookLogger() + + result: Final = await run_chat_completion_agentic_hooks( + response=model_response(), + model="gpt-4o-mini", + messages=[{"role": "user", "content": "what is 6*7?"}], + optional_params={}, + kwargs={"_agentic_loop_tool_choice": _ALLOWS_ONLY_GET_WEATHER}, + logging_obj=_logging_with(hook), + custom_llm_provider="openai", + stream=False, + ) + + assert hook.runs == 0 + assert isinstance(result, ModelResponse) + assert not result.choices[0].message.tool_calls + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model_response", "tool_type"), + [(_tool_call_model_response, "function"), (_custom_tool_call_model_response, "custom")], +) +async def test_a_chat_hook_runs_a_call_the_choice_lists(model_response: Callable[[], ModelResponse], tool_type: str): + hook: Final = _ChatHookLogger() + + result: Final = await run_chat_completion_agentic_hooks( + response=model_response(), + model="gpt-4o-mini", + messages=[{"role": "user", "content": "what is 6*7?"}], + optional_params={}, + kwargs={ + "_agentic_loop_tool_choice": { + "type": "allowed_tools", + "mode": "auto", + "tools": [{"type": tool_type, "name": "litellm_code_execution"}], + } + }, + logging_obj=_logging_with(hook), + custom_llm_provider="openai", + stream=False, + ) + + assert hook.runs == 1 + assert isinstance(result, ModelResponse) + assert result.choices[0].message.content == "ran the hook" + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("tool_choice", "hook_runs"), [(_ALLOWS_ONLY_GET_WEATHER, 0), (_ALLOWS_THE_HOOKS_CALL, 1)]) +async def test_acompletion_hands_the_callers_choice_to_the_chat_hooks(tool_choice: dict[str, object], hook_runs: int): + hook: Final = _ChatHookLogger() + mock_client: Final = MagicMock() + mock_client.post = AsyncMock(return_value=_raw_response_for(_tool_call_model_response())) + + response: Final = await litellm.acompletion( + model="openai/gpt-4o-mini", + messages=[{"role": "user", "content": "what is 6*7?"}], + tools=[{"type": "function", "function": {"name": "get_weather", "parameters": {"type": "object"}}}], + tool_choice=tool_choice, + api_key="sk-test", + client=mock_client, + success_callback=[hook], + ) + + assert (hook.runs, mock_client.post.await_count, response.choices[0].message.tool_calls) == (hook_runs, 1, None) diff --git a/tests/unit/responses/mcp/test_chat_completions_handler.py b/tests/unit/responses/mcp/test_chat_completions_handler.py index bf5c7c0846a..ea013184a15 100644 --- a/tests/unit/responses/mcp/test_chat_completions_handler.py +++ b/tests/unit/responses/mcp/test_chat_completions_handler.py @@ -26,6 +26,7 @@ import litellm from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.proxy.mcp_server.faults.list_outcomes import AggregateToolListing from litellm.responses.mcp import chat_completions_handler from litellm.responses.mcp.chat_completions_handler import ( acompletion_with_mcp, @@ -34,6 +35,7 @@ from litellm.responses.mcp.litellm_proxy_mcp_handler import ( LiteLLM_Proxy_MCP_Handler, ) from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.types.mcp_server.mcp_server_manager import MCPServer from litellm.types.utils import ModelResponse, ModelResponseStream, StandardLoggingPayload, Usage from litellm.utils import CustomStreamWrapper @@ -2761,7 +2763,21 @@ _ALLOWED_ONLY_ANOTHER_TOOL: Final = { } +_DEMO_MCP_SERVER: Final = MCPServer( + server_id="demo", + name="demo", + server_name="demo", + transport="http", + url="https://demo.example/mcp", + auth_type="none", +) + + def _served_by_a_recording_mcp_server(monkeypatch: pytest.MonkeyPatch, call_tool: AsyncMock) -> None: + monkeypatch.setattr( + "litellm.proxy.mcp_server.server._get_tools_from_mcp_servers", + AsyncMock(return_value=AggregateToolListing(tools=[_MCP_ECHO_TOOL], outcomes={})), + ) monkeypatch.setitem( sys.modules, "litellm.proxy.proxy_server", SimpleNamespace(proxy_logging_obj=None, prisma_client=None) ) @@ -2773,6 +2789,8 @@ def _served_by_a_recording_mcp_server(monkeypatch: pytest.MonkeyPatch, call_tool call_tool=call_tool, get_mcp_server_from_tool_name=MagicMock(return_value=None), get_mcp_server_by_name=MagicMock(return_value=None), + get_allowed_mcp_servers=AsyncMock(return_value=[_DEMO_MCP_SERVER.server_id]), + get_mcp_servers_from_ids=MagicMock(return_value=[_DEMO_MCP_SERVER]), ), ) @@ -2786,21 +2804,16 @@ async def test_an_mcp_tool_the_choice_excludes_never_runs_when_the_model_ignored call_tool: Final = AsyncMock() _served_by_a_recording_mcp_server(monkeypatch, call_tool) - with patch.object( - LiteLLM_Proxy_MCP_Handler, - "process_mcp_tools_without_openai_transform", - AsyncMock(return_value=([_MCP_ECHO_TOOL], {"echo": "demo"})), - ): - response: Final = await litellm.acompletion( - model="openai/gpt-5.6", - messages=[{"role": "user", "content": "Call echo with PING"}], - tools=[{"type": "mcp", "server_url": "litellm_proxy/mcp/demo", "require_approval": "never"}], - tool_choice=_ALLOWED_ONLY_ANOTHER_TOOL, - api_key="sk-mcp-chat-rounds", - api_base="https://mcp-chat-rounds.example.com/v1", - stream=stream, - ) - chunks: Final = await _read_all(response) if stream else () + response: Final = await litellm.acompletion( + model="openai/gpt-5.6", + messages=[{"role": "user", "content": "Call echo with PING"}], + tools=[{"type": "mcp", "server_url": "litellm_proxy/mcp/demo", "require_approval": "never"}], + tool_choice=_ALLOWED_ONLY_ANOTHER_TOOL, + api_key="sk-mcp-chat-rounds", + api_base="https://mcp-chat-rounds.example.com/v1", + stream=stream, + ) + chunks: Final = await _read_all(response) if stream else () handed_back: Final = ( tuple(call.function.name for chunk in chunks for call in chunk.choices[0].delta.tool_calls or ()) @@ -2826,52 +2839,8 @@ def _anthropic_call_to_echo() -> Response: ) -async def _call_echo_on_anthropic_under_drop_params( - respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, rounds: Sequence[Response], tool_choice: object -) -> tuple[ModelResponse, respx.Route]: - monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") - litellm.in_memory_llm_clients_cache.flush_cache() - route: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock(side_effect=list(rounds)) - with patch.object( - LiteLLM_Proxy_MCP_Handler, - "process_mcp_tools_without_openai_transform", - AsyncMock(return_value=([_MCP_ECHO_TOOL], {"echo": "demo"})), - ): - response: Final = await litellm.acompletion( - model="anthropic/claude-haiku-4-5", - messages=[{"role": "user", "content": "Call echo with PING"}], - tools=[{"type": "mcp", "server_url": "litellm_proxy/mcp/demo", "require_approval": "never"}], - tool_choice=tool_choice, - drop_params=True, - api_key="sk-ant-fake", - ) - return response, route - - -@pytest.mark.asyncio -async def test_a_dropped_choice_hands_back_an_mcp_call_it_excludes( - respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch -): - call_tool: Final = AsyncMock() - _served_by_a_recording_mcp_server(monkeypatch, call_tool) - - response, route = await _call_echo_on_anthropic_under_drop_params( - respx_mock, monkeypatch, (_anthropic_call_to_echo(),), _ALLOWED_ONLY_ANOTHER_TOOL - ) - - call_tool.assert_not_awaited() - assert "tool_choice" not in json.loads(route.calls[0].request.content) - assert [call.function.name for call in response.choices[0].message.tool_calls] == ["echo"] - - -@pytest.mark.asyncio -async def test_a_dropped_choice_still_runs_an_mcp_call_it_lists( - respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch -): - call_tool: Final = AsyncMock(return_value=CallToolResult(content=[TextContent(type="text", text="PING")])) - _served_by_a_recording_mcp_server(monkeypatch, call_tool) - answer: Final = Response( +def _anthropic_answer() -> Response: + return Response( 200, json={ "id": "msg_answer", @@ -2884,8 +2853,67 @@ async def test_a_dropped_choice_still_runs_an_mcp_call_it_lists( }, ) + +_WHOLE_DEMO_SERVER: Final = {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "mcp", "server_label": "demo"}]} +_WHOLE_OTHER_SERVER: Final = {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "mcp", "server_label": "other"}]} + + +async def _call_echo_on_anthropic_under_drop_params( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + rounds: Sequence[Response], + tool_choice: object, + server_url: str = "litellm_proxy/mcp/demo", +) -> tuple[ModelResponse, respx.Route]: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + litellm.in_memory_llm_clients_cache.flush_cache() + route: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock(side_effect=list(rounds)) + response: Final = await litellm.acompletion( + model="anthropic/claude-haiku-4-5", + messages=[{"role": "user", "content": "Call echo with PING"}], + tools=[{"type": "mcp", "server_label": "demo", "server_url": server_url, "require_approval": "never"}], + tool_choice=tool_choice, + drop_params=True, + api_key="sk-ant-fake", + ) + return response, route + + +@pytest.mark.asyncio +@pytest.mark.parametrize("tool_choice", [_ALLOWED_ONLY_ANOTHER_TOOL, _WHOLE_OTHER_SERVER]) +async def test_a_dropped_choice_hands_back_an_mcp_call_it_excludes( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, tool_choice: object +): + call_tool: Final = AsyncMock() + _served_by_a_recording_mcp_server(monkeypatch, call_tool) + + response, route = await _call_echo_on_anthropic_under_drop_params( + respx_mock, monkeypatch, (_anthropic_call_to_echo(),), tool_choice + ) + + call_tool.assert_not_awaited() + assert "tool_choice" not in json.loads(route.calls[0].request.content) + assert [call.function.name for call in response.choices[0].message.tool_calls] == ["echo"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("tool_choice", "server_url"), + [ + (_FLAT_ALLOWED_ECHO, "litellm_proxy/mcp/demo"), + (_WHOLE_DEMO_SERVER, "litellm_proxy/mcp/demo"), + (_WHOLE_DEMO_SERVER, "litellm_proxy"), + ], +) +async def test_a_dropped_choice_still_runs_an_mcp_call_it_lists( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, tool_choice: object, server_url: str +): + call_tool: Final = AsyncMock(return_value=CallToolResult(content=[TextContent(type="text", text="PING")])) + _served_by_a_recording_mcp_server(monkeypatch, call_tool) + response, _route = await _call_echo_on_anthropic_under_drop_params( - respx_mock, monkeypatch, (_anthropic_call_to_echo(), answer), _FLAT_ALLOWED_ECHO + respx_mock, monkeypatch, (_anthropic_call_to_echo(), _anthropic_answer()), tool_choice, server_url ) assert call_tool.await_count == 1 diff --git a/tests/unit/test_tool_loop.py b/tests/unit/test_tool_loop.py index bf5c00ef53f..e679f9817c1 100644 --- a/tests/unit/test_tool_loop.py +++ b/tests/unit/test_tool_loop.py @@ -663,7 +663,6 @@ async def test_arun_tool_loop_two_rounds(respx_mock: respx.MockRouter, monkeypat assert messages == messages_snapshot -ANTHROPIC_MESSAGES_URL: Final = "https://api.anthropic.com/v1/messages" TOOL_SEARCH_CALL: Final = { "type": "server_tool_use", "id": "srvtoolu_01Search", diff --git a/tests/unit/types/test_litellm_params.py b/tests/unit/types/test_litellm_params.py index d71791b6f11..3caac93d2cd 100644 --- a/tests/unit/types/test_litellm_params.py +++ b/tests/unit/types/test_litellm_params.py @@ -237,6 +237,7 @@ AGENTIC_LOOP_STATE_NAMES: Final = ( "_agentic_loop_depth", "_agentic_loop_fingerprints", "_agentic_loop_api_surface", + "_agentic_loop_tool_choice", "_code_interpreter_interception_active", "_code_interpreter_interception_sandbox_key", "_code_interpreter_interception_session_scoped", From 4755727276c84a50107c3d5aef9d398e7d0d4603 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 6 Nov 2026 07:49:49 -0800 Subject: [PATCH 8/9] fix(tool-choice): match named mcp entries by server label and let built-in entries allow their interception call --- .../code_interpreter_interception/handler.py | 17 ++++- .../websearch_interception/handler.py | 31 ++++++++-- .../allowed_tools_tool_choice.py | 58 ++++++++++++++--- .../test_handler.py | 34 ++++++++++ .../test_websearch_responses.py | 34 ++++++++++ .../test_allowed_tools_tool_choice.py | 62 +++++++++++++++++++ 6 files changed, 220 insertions(+), 16 deletions(-) diff --git a/litellm/integrations/code_interpreter_interception/handler.py b/litellm/integrations/code_interpreter_interception/handler.py index b035732afa6..dd36eee978d 100644 --- a/litellm/integrations/code_interpreter_interception/handler.py +++ b/litellm/integrations/code_interpreter_interception/handler.py @@ -18,6 +18,7 @@ from pydantic import ValidationError import litellm from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.allowed_tools_tool_choice import allowed_tools_tool_choice_with_builtin_as_function from litellm.llms.base_llm.sandbox.transformation import ( CodeExecutionResult, ContainerHandle, @@ -175,6 +176,10 @@ def _sandbox_unavailable(step: str, error: Exception, output: str = _SANDBOX_UNA return _SandboxUnavailable(output=output) +def _is_code_interpreter_tool(tool: Mapping[str, object]) -> bool: + return tool.get("type") == "code_interpreter" + + def _output_item_type(item: object) -> object: if isinstance(item, dict): item_mapping: Final[dict[str, object]] = item @@ -343,8 +348,8 @@ class CodeInterpreterInterceptionLogger(CustomLogger): (function_tool if isinstance(tool, dict) and tool.get("type") == "code_interpreter" else tool) for tool in tools ] - if self._tool_choice_targets_code_interpreter(kwargs.get("tool_choice")): - kwargs["tool_choice"] = self._get_function_tool_choice(call_type=call_type) + if "tool_choice" in kwargs: + kwargs["tool_choice"] = self._intercepted_tool_choice(kwargs["tool_choice"], call_type) return kwargs @staticmethod @@ -428,6 +433,14 @@ class CodeInterpreterInterceptionLogger(CustomLogger): or (isinstance(function, dict) and function.get("name") == LITELLM_CODE_EXECUTION_TOOL_NAME) ) + @classmethod + def _intercepted_tool_choice(cls, tool_choice: object, call_type: CallTypes | None) -> object: + if cls._tool_choice_targets_code_interpreter(tool_choice): + return cls._get_function_tool_choice(call_type=call_type) + return allowed_tools_tool_choice_with_builtin_as_function( + tool_choice, _is_code_interpreter_tool, LITELLM_CODE_EXECUTION_TOOL_NAME + ) + def _resolve_provider(self, kwargs: dict[str, object]) -> str | None: provider: Final = kwargs.get("custom_llm_provider") if isinstance(provider, str) and provider: diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index d41647fe5f1..632864d05c9 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -9,7 +9,7 @@ server-side using litellm router's search tools. import asyncio import math import uuid -from collections.abc import AsyncIterator, Mapping, Sequence +from collections.abc import AsyncIterator, Callable, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, TypeVar, cast @@ -38,6 +38,7 @@ from litellm.integrations.websearch_interception.transformation import ( from litellm.litellm_core_utils.agentic_loop_settings import ( validated_max_agentic_loops, ) +from litellm.litellm_core_utils.allowed_tools_tool_choice import allowed_tools_tool_choice_with_builtin_as_function from litellm.litellm_core_utils.core_helpers import qualify_provider_stripped_model from litellm.llms.base_llm.search.transformation import SearchResponse from litellm.types.integrations.custom_logger import ( @@ -263,6 +264,18 @@ class _ParentRequestCorrelation: _NO_ACOMPLETION_NAMED: Final[_AcompletionNamedParams] = {} +def _tool_choice_naming_litellm_web_search( + kwargs: Mapping[str, object], is_search_tool: Callable[[Mapping[str, object]], bool] +) -> Mapping[str, object]: + if "tool_choice" not in kwargs: + return MappingProxyType({}) + return { + "tool_choice": allowed_tools_tool_choice_with_builtin_as_function( + kwargs["tool_choice"], is_search_tool, LITELLM_WEB_SEARCH_TOOL_NAME + ) + } + + class WebSearchInterceptionLogger(CustomLogger): """ CustomLogger that intercepts WebSearch tool calls for models that don't @@ -514,13 +527,17 @@ class WebSearchInterceptionLogger(CustomLogger): # Keep other tools as-is converted_tools.append(tool) - kwargs["tools"] = converted_tools + converted_kwargs: Final[dict[str, object]] = { + **kwargs, + "tools": converted_tools, + **_tool_choice_naming_litellm_web_search(kwargs, lambda tool: is_search_tool(dict(tool))), + } if kwargs.get("stream"): verbose_logger.debug("WebSearchInterception: deployment hook converting stream=True to stream=False") - return as_converted_stream(kwargs, WEBSEARCH_INTERCEPTION_PREFIX) + return as_converted_stream(converted_kwargs, WEBSEARCH_INTERCEPTION_PREFIX) - return kwargs + return converted_kwargs def _convert_responses_tools( self, kwargs: Mapping[str, object], tools: Sequence[dict[str, object]] @@ -535,7 +552,11 @@ class WebSearchInterceptionLogger(CustomLogger): get_litellm_web_search_tool_responses() if is_web_search_tool_responses(tool) else tool for tool in tools ] - converted_kwargs: Final = {**kwargs, "tools": converted_tools} + converted_kwargs: Final = { + **kwargs, + "tools": converted_tools, + **_tool_choice_naming_litellm_web_search(kwargs, is_web_search_tool_responses), + } if kwargs.get("stream"): verbose_logger.debug("WebSearchInterception: deployment hook converting stream=True to stream=False") diff --git a/litellm/litellm_core_utils/allowed_tools_tool_choice.py b/litellm/litellm_core_utils/allowed_tools_tool_choice.py index fd0ba2c8e4f..ae256580b93 100644 --- a/litellm/litellm_core_utils/allowed_tools_tool_choice.py +++ b/litellm/litellm_core_utils/allowed_tools_tool_choice.py @@ -1,4 +1,4 @@ -from collections.abc import Iterable, Mapping +from collections.abc import Callable, Iterable, Mapping from types import MappingProxyType from typing import Final, Literal @@ -176,7 +176,7 @@ def _allowed_tool_names(tool_choice: object) -> frozenset[object]: if parsed is None: return frozenset() flattened: Final = tuple(_responses_allowed_tool(tool) for tool in parsed.allowed_tools.tools) - return frozenset(tool.get("name") for tool in flattened if tool is not None) + return frozenset(tool.get("name") for tool in flattened if tool is not None and _server_entry(tool) is None) def refuse_a_malformed_allowed_tools_tool_choice(tool_choice: object, model: str) -> None: @@ -191,15 +191,22 @@ def refuse_a_malformed_allowed_tools_tool_choice(tool_choice: object, model: str ) -def _whole_server_label(tool: Mapping[str, object]) -> str | None: +def _server_entry(tool: Mapping[str, object]) -> tuple[str, str | None] | None: label: Final = tool.get("server_label") - return label if tool.get("type") == "mcp" and tool.get("name") is None and isinstance(label, str) else None + name: Final = tool.get("name") + if tool.get("type") != "mcp" or not isinstance(label, str): + return None + return label, name if isinstance(name, str) else None -def _whole_server_labels(tool_choice: object) -> frozenset[str]: +def _server_entries(tool_choice: object) -> frozenset[tuple[str, str | None]]: parsed: Final = parse_allowed_tools_tool_choice(tool_choice) - labels: Final = () if parsed is None else tuple(map(_whole_server_label, parsed.allowed_tools.tools)) - return frozenset(label for label in labels if label is not None) + entries: Final = () if parsed is None else tuple(map(_server_entry, parsed.allowed_tools.tools)) + return frozenset(entry for entry in entries if entry is not None) + + +def _a_server_entry_allows(entries: frozenset[tuple[str, str | None]], name: str, labels: frozenset[str]) -> bool: + return any(label in labels and entry_name in (None, name) for label, entry_name in entries) def tool_calls_the_choice_excludes( @@ -210,11 +217,12 @@ def tool_calls_the_choice_excludes( if not is_allowed_tools_tool_choice(tool_choice): return () allowed: Final = _allowed_tool_names(tool_choice) - whole_servers: Final = _whole_server_labels(tool_choice) + server_entries: Final = _server_entries(tool_choice) return tuple( name for name in tool_names - if name not in allowed and whole_servers.isdisjoint(server_labels_by_tool.get(name, frozenset())) + if name not in allowed + and not _a_server_entry_allows(server_entries, name, server_labels_by_tool.get(name, frozenset())) ) @@ -232,6 +240,38 @@ def refuse_tool_calls_the_choice_excludes(tool_choice: object, tool_names: Itera ) +def _with_builtin_as( + tool: Mapping[str, object], is_builtin: Callable[[Mapping[str, object]], bool], function_entry: Mapping[str, object] +) -> Mapping[str, object]: + return function_entry if tool.get("type") not in ("function", "custom", "mcp") and is_builtin(tool) else tool + + +def allowed_tools_tool_choice_with_builtin_as_function( + tool_choice: object, is_builtin: Callable[[Mapping[str, object]], bool], function_name: str +) -> object: + try: + parsed: Final = _ALLOWED_TOOLS_TOOL_CHOICE.validate_python(tool_choice) + except ValidationError: + return tool_choice + match parsed: + case ChatAllowedToolsToolChoice(): + nested: Final = {"type": "function", "function": {"name": function_name}} + return { + "type": "allowed_tools", + "allowed_tools": { + "mode": parsed.allowed_tools.mode, + "tools": [dict(_with_builtin_as(tool, is_builtin, nested)) for tool in parsed.allowed_tools.tools], + }, + } + case ResponsesAllowedToolsToolChoice(): + flat: Final = {"type": "function", "name": function_name} + return { + "type": "allowed_tools", + "mode": parsed.mode, + "tools": [dict(_with_builtin_as(tool, is_builtin, flat)) for tool in parsed.tools], + } + + def chat_allowed_tools_tool_choice(choice: ChatAllowedToolsToolChoice) -> ChatCompletionAllowedToolChoiceParam: return ChatCompletionAllowedToolChoiceParam( type="allowed_tools", diff --git a/tests/unit/integrations/code_interpreter_interception/test_handler.py b/tests/unit/integrations/code_interpreter_interception/test_handler.py index f4f6a127fc7..5fde826ee23 100644 --- a/tests/unit/integrations/code_interpreter_interception/test_handler.py +++ b/tests/unit/integrations/code_interpreter_interception/test_handler.py @@ -12,6 +12,7 @@ from typing import Final import httpx import pytest +from pydantic import TypeAdapter import respx import litellm @@ -27,6 +28,7 @@ from litellm.integrations.code_interpreter_interception.handler import ( LITELLM_CODE_EXECUTION_TOOL_NAME, CodeInterpreterInterceptionLogger, ) +from litellm.litellm_core_utils.allowed_tools_tool_choice import tool_calls_the_choice_excludes from litellm.llms.base_llm.sandbox.transformation import CodeExecutionResult from litellm.types.integrations.custom_logger import ( CHAT_COMPLETION_AGENTIC_SURFACE, @@ -1203,6 +1205,38 @@ async def test_extract_tool_calls_reads_object_attributes(): # --------------------------------------------------------------------------- +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("call_type", "tool_choice"), + [ + (CallTypes.aresponses, {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "code_interpreter"}]}), + ( + CallTypes.acompletion, + {"type": "allowed_tools", "allowed_tools": {"mode": "auto", "tools": [{"type": "code_interpreter"}]}}, + ), + ], +) +async def test_pre_call_lets_an_allowed_tools_choice_listing_code_interpreter_allow_the_run( + call_type: CallTypes, tool_choice: dict[str, object] +): + logger: Final = CodeInterpreterInterceptionLogger() + + result: Final = TypeAdapter(dict[str, object]).validate_python( + await logger.async_pre_call_deployment_hook( + { + "tools": [{"type": "code_interpreter", "container": {"type": "auto"}}], + "tool_choice": tool_choice, + "custom_llm_provider": "openai", + }, + call_type, + ) + ) + + assert tool_calls_the_choice_excludes(result["tool_choice"], (LITELLM_CODE_EXECUTION_TOOL_NAME, "get_weather")) == ( + "get_weather", + ) + + @pytest.mark.asyncio async def test_pre_call_uses_session_id_from_metadata_as_sandbox_key(): """When session_id is in request metadata, it becomes the sandbox key so the diff --git a/tests/unit/integrations/websearch_interception/test_websearch_responses.py b/tests/unit/integrations/websearch_interception/test_websearch_responses.py index 6ca812a5e40..0ad758cf53f 100644 --- a/tests/unit/integrations/websearch_interception/test_websearch_responses.py +++ b/tests/unit/integrations/websearch_interception/test_websearch_responses.py @@ -21,6 +21,7 @@ from litellm.integrations.websearch_interception.handler import ( from litellm.types.integrations.custom_logger import ( RESPONSES_AGENTIC_SURFACE, ) +from litellm.litellm_core_utils.allowed_tools_tool_choice import tool_calls_the_choice_excludes from litellm.litellm_core_utils.litellm_logging import Logging from litellm.types.utils import CallTypes, LlmProviders @@ -390,6 +391,39 @@ async def test_deployment_hook_converts_native_responses_web_search_tool(): assert tool["parameters"]["required"] == ["query"] +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("call_type", "tool_choice"), + [ + (CallTypes.aresponses, {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "web_search"}]}), + (CallTypes.acompletion, {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "web_search"}]}), + ( + CallTypes.acompletion, + {"type": "allowed_tools", "allowed_tools": {"mode": "required", "tools": [{"type": "web_search"}]}}, + ), + ], +) +async def test_deployment_hook_lets_an_allowed_tools_choice_listing_web_search_allow_the_search( + call_type: CallTypes, tool_choice: dict[str, object] +): + logger: Final = WebSearchInterceptionLogger(enabled_providers=[LlmProviders.OPENAI]) + + result: Final = await logger.async_pre_call_deployment_hook( + kwargs={ + "model": "gpt-4o", + "custom_llm_provider": "openai", + "tools": [{"type": "web_search"}, {"type": "function", "name": "get_weather"}], + "tool_choice": tool_choice, + }, + call_type=call_type, + ) + + assert result is not None + assert tool_calls_the_choice_excludes(result["tool_choice"], ("litellm_web_search", "get_weather")) == ( + "get_weather", + ) + + @pytest.mark.asyncio async def test_deployment_hook_responses_returns_none_without_web_search(): """No web search tool in a responses request -> deployment hook makes no change.""" diff --git a/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py b/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py index 7c19e0b673c..5d7fbbb2228 100644 --- a/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py +++ b/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py @@ -8,6 +8,7 @@ import litellm from litellm.exceptions import BadRequestError from litellm.litellm_core_utils.allowed_tools_tool_choice import ( ChatAllowedToolsToolChoice, + allowed_tools_tool_choice_with_builtin_as_function, drops_unsupported_allowed_tools_tool_choice, is_allowed_tools_tool_choice, chat_allowed_tools_tool_choice, @@ -297,6 +298,17 @@ def test_an_mcp_entry_with_a_null_name_is_a_valid_choice(tool_choice: object): assert refuse_a_malformed_allowed_tools_tool_choice(tool_choice, "the-model") is None +_DOCS_SEARCH: Final = { + "type": "allowed_tools", + "mode": "auto", + "tools": [{"type": "mcp", "server_label": "docs", "name": "search"}], +} + + +def test_a_named_mcp_entry_allows_its_call_on_a_server_its_label_serves(): + assert tool_calls_the_choice_excludes(_DOCS_SEARCH, ("search",), {"search": frozenset({"docs", "other"})}) == () + + def test_an_mcp_entry_without_a_name_allows_every_call_its_server_serves(): assert tool_calls_the_choice_excludes(_WHOLE_DOCS_SERVER, ("search",), {"search": frozenset({"docs"})}) == () @@ -310,9 +322,59 @@ def test_an_mcp_entry_without_a_name_allows_every_call_its_server_serves(): {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "mcp", "server_label": "docs", "name": "read"}]}, frozenset({"docs"}), ), + (_DOCS_SEARCH, frozenset({"other"})), + (_DOCS_SEARCH, frozenset()), ], ) def test_an_mcp_entry_allows_no_call_outside_its_server_or_its_name( tool_choice: object, server_labels: frozenset[str] ): assert tool_calls_the_choice_excludes(tool_choice, ("search",), {"search": server_labels}) == ("search",) + + +def _is_web_search(tool: Mapping[str, object]) -> bool: + return tool.get("type") == "web_search" + + +@pytest.mark.parametrize( + "tool_choice", + [ + {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "web_search"}, {"type": "function", "name": "get_time"}]}, + { + "type": "allowed_tools", + "allowed_tools": { + "mode": "auto", + "tools": [{"type": "web_search"}, {"type": "function", "function": {"name": "get_time"}}], + }, + }, + ], +) +def test_a_builtin_entry_allows_the_function_that_replaces_it_in_the_same_layout(tool_choice: dict[str, object]): + rewritten: Final = allowed_tools_tool_choice_with_builtin_as_function( + tool_choice, _is_web_search, "litellm_web_search" + ) + + assert tool_calls_the_choice_excludes(tool_choice, ("litellm_web_search", "get_time")) == ("litellm_web_search",) + assert tool_calls_the_choice_excludes(rewritten, ("litellm_web_search", "get_time", "get_weather")) == ( + "get_weather", + ) + assert isinstance(rewritten, dict) and rewritten.keys() == tool_choice.keys() + + +def test_a_function_or_mcp_entry_is_never_taken_for_a_builtin(): + tool_choice: Final = { + "type": "allowed_tools", + "mode": "auto", + "tools": [{"type": "function", "name": "web_search"}, {"type": "mcp", "server_label": "docs", "name": "web_search"}], + } + + rewritten: Final = allowed_tools_tool_choice_with_builtin_as_function( + tool_choice, lambda tool: tool.get("name") == "web_search", "litellm_web_search" + ) + + assert parse_allowed_tools_tool_choice(rewritten) == parse_allowed_tools_tool_choice(tool_choice) + + +@pytest.mark.parametrize("tool_choice", ["auto", None, {"type": "allowed_tools", "mode": "auto", "tools": []}]) +def test_any_other_or_malformed_choice_is_left_as_sent(tool_choice: object): + assert allowed_tools_tool_choice_with_builtin_as_function(tool_choice, _is_web_search, "litellm_web_search") is tool_choice From cd076fb13abea196d91e5804fa4b8394949f5e6d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 6 Nov 2026 08:19:23 -0800 Subject: [PATCH 9/9] fix(tool-choice): leave a malformed allowed_tools choice unrewritten in the interception hooks --- .../litellm_core_utils/allowed_tools_tool_choice.py | 5 ++--- .../test_allowed_tools_tool_choice.py | 12 +++++++++++- 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/litellm/litellm_core_utils/allowed_tools_tool_choice.py b/litellm/litellm_core_utils/allowed_tools_tool_choice.py index ae256580b93..1d254aeb559 100644 --- a/litellm/litellm_core_utils/allowed_tools_tool_choice.py +++ b/litellm/litellm_core_utils/allowed_tools_tool_choice.py @@ -249,10 +249,9 @@ def _with_builtin_as( def allowed_tools_tool_choice_with_builtin_as_function( tool_choice: object, is_builtin: Callable[[Mapping[str, object]], bool], function_name: str ) -> object: - try: - parsed: Final = _ALLOWED_TOOLS_TOOL_CHOICE.validate_python(tool_choice) - except ValidationError: + if parse_allowed_tools_tool_choice(tool_choice) is None: return tool_choice + parsed: Final = _ALLOWED_TOOLS_TOOL_CHOICE.validate_python(tool_choice) match parsed: case ChatAllowedToolsToolChoice(): nested: Final = {"type": "function", "function": {"name": function_name}} diff --git a/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py b/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py index 5d7fbbb2228..5a87f0e6434 100644 --- a/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py +++ b/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py @@ -375,6 +375,16 @@ def test_a_function_or_mcp_entry_is_never_taken_for_a_builtin(): assert parse_allowed_tools_tool_choice(rewritten) == parse_allowed_tools_tool_choice(tool_choice) -@pytest.mark.parametrize("tool_choice", ["auto", None, {"type": "allowed_tools", "mode": "auto", "tools": []}]) +@pytest.mark.parametrize( + "tool_choice", + [ + "auto", + None, + {"type": "allowed_tools", "mode": "auto", "tools": []}, + {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "web_search", "name": []}]}, + {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "web_search"}, {"type": "function", "name": 3}]}, + {"type": "allowed_tools", "allowed_tools": {"mode": "auto", "tools": [{"type": "web_search", "name": {}}]}}, + ], +) def test_any_other_or_malformed_choice_is_left_as_sent(tool_choice: object): assert allowed_tools_tool_choice_with_builtin_as_function(tool_choice, _is_web_search, "litellm_web_search") is tool_choice