From 28c80048ed47d53befd819f22b5ac1bb112ec3fd Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Mon, 2 Nov 2026 18:51:58 -0800 Subject: [PATCH] fix(tool_choice): accept OpenAI allowed_tools on chat completions and flatten it for the Responses bridge (#49043) * fix(tool_choice): accept OpenAI allowed_tools on chat completions and flatten it for the Responses bridge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test: type the allowed_tools tool_choice fixtures Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test: raw regex for the allowed_tools provider gate match Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(tool_choice): move the allowed_tools provider gate out of completion() so pyright can still analyze it Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(tool_choice): keep the allowed_tools helpers typed so the basedpyright gate stays level Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(tool_choice): type the allowed_tools validation boundary with Any instead of casts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(tool_choice): widen tool_choice to object at the allowed_tools helpers with a cast-ok reason Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(tool_choice): drop allowed_tools test docstrings, type the new tests, and revert unrelated reformatting Also reword the allowed_tools 400 so it no longer blames the flat Responses layout for every malformed choice --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../transformation.py | 22 +++- .../allowed_tools_tool_choice.py | 111 ++++++++++++++++++ litellm/llms/azure/chat/gpt_transformation.py | 16 ++- litellm/main.py | 11 ++ litellm/utils.py | 25 +++- ...responses_transformation_transformation.py | 43 +++++++ .../test_allowed_tools_tool_choice.py | 100 ++++++++++++++++ .../test_azure_chat_gpt_transformation.py | 36 ++++++ tests/unit/test_main.py | 57 +++++++++ tests/unit/test_utils.py | 34 ++++++ 10 files changed, 448 insertions(+), 7 deletions(-) create mode 100644 litellm/litellm_core_utils/allowed_tools_tool_choice.py create mode 100644 tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index c9d9cc25049..0a500173c5b 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -16,6 +16,7 @@ from openai.types.responses.response_input_param import ( ResponseCustomToolCallOutputParam, ResponseCustomToolCallParam, ) +from openai.types.responses.tool_choice_allowed_param import ToolChoiceAllowedParam from openai.types.responses.tool_choice_custom_param import ToolChoiceCustomParam from openai.types.responses.tool_choice_function_param import ToolChoiceFunctionParam from openai.types.responses.tool_param import FunctionToolParam @@ -34,6 +35,10 @@ from litellm.completion_extras.litellm_responses_transformation.history import ( requested_phase, ) from litellm.integrations.anthropic_cache_control_hook import supports_openai_prompt_cache_breakpoint +from litellm.litellm_core_utils.allowed_tools_tool_choice import ( + parse_chat_allowed_tools_tool_choice, + responses_allowed_tools_tool_choice, +) from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.litellm_core_utils.hidden_params import get_hidden_params, get_or_create_hidden_params from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -783,13 +788,28 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): def _normalize_tool_choice_for_responses_api( self, tool_choice: _ToolChoiceT - ) -> _ToolChoiceT | ToolChoiceFunctionParam | ToolChoiceCustomParam | Literal["auto", "none", "required"]: + ) -> ( + _ToolChoiceT + | ToolChoiceFunctionParam + | ToolChoiceCustomParam + | ToolChoiceAllowedParam + | Literal["auto", "none", "required"] + ): """Chat tool_choice nests the name under function/custom; Responses API expects top-level name.""" if not isinstance(tool_choice, dict): return tool_choice choice_type: Final = tool_choice.get("type") 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( + cast( # cast-ok: widening an untyped dict to object is always sound, the helper validates it with pydantic + object, tool_choice + ) + ) + flat: Final = responses_allowed_tools_tool_choice(parsed) if parsed is not None else None + if flat is not None: + return flat if choice_type not in ("function", "custom"): return tool_choice if isinstance(tool_choice.get("name"), str) and tool_choice.get("name"): diff --git a/litellm/litellm_core_utils/allowed_tools_tool_choice.py b/litellm/litellm_core_utils/allowed_tools_tool_choice.py new file mode 100644 index 00000000000..5c8f5d83ece --- /dev/null +++ b/litellm/litellm_core_utils/allowed_tools_tool_choice.py @@ -0,0 +1,111 @@ +from collections.abc import Mapping +from typing import Final, Literal + +from openai.types.responses.tool_choice_allowed_param import ToolChoiceAllowedParam +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError + +from litellm.exceptions import BadRequestError + +ALLOWED_TOOLS_TOOL_CHOICE_PROVIDERS: Final = frozenset({"openai", "azure", "litellm_proxy"}) + +CHAT_ALLOWED_TOOLS_LAYOUT: Final = ( + '{"type": "allowed_tools", "allowed_tools": {"mode": "auto" | "required", "tools": []}}' +) + + +class ChatAllowedTools(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + mode: Literal["auto", "required"] + tools: tuple[Mapping[str, object], ...] = Field(min_length=1) + + +class ChatAllowedToolsToolChoice(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + type: Literal["allowed_tools"] + allowed_tools: ChatAllowedTools + + +_CHAT_ALLOWED_TOOLS_TOOL_CHOICE: Final = TypeAdapter(ChatAllowedToolsToolChoice) + + +_SELECTOR: Final = TypeAdapter(dict[str, object]) + + +class _ToolChoiceType(BaseModel): + model_config = ConfigDict(frozen=True) + + type: str + + +_TOOL_CHOICE_TYPE: Final = TypeAdapter(_ToolChoiceType) + + +def is_allowed_tools_tool_choice(tool_choice: object) -> bool: + try: + return _TOOL_CHOICE_TYPE.validate_python(tool_choice).type == "allowed_tools" + except ValidationError: + return False + + +def raise_if_allowed_tools_tool_choice_unsupported( + 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.""" + if not is_allowed_tools_tool_choice(tool_choice) or custom_llm_provider in ALLOWED_TOOLS_TOOL_CHOICE_PROVIDERS: + return + 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." + ), + model=model, + llm_provider=custom_llm_provider, + ) + + +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 _responses_allowed_tool(tool: Mapping[str, object]) -> Mapping[str, object] | None: + tool_type: Final = tool.get("type") + if tool_type not in ("function", "custom"): + return tool + if tool_type not in tool: + return tool if _has_name(tool) else None + try: + nested: Final = _SELECTOR.validate_python(tool[tool_type]) + except ValidationError: + return None + outer: Final = {key: value for key, value in tool.items() if key != tool_type} + if not _has_name(nested) or outer.keys() & nested.keys(): + return None + return {**outer, **nested} + + +def responses_allowed_tools_tool_choice(choice: ChatAllowedToolsToolChoice) -> ToolChoiceAllowedParam | None: + """The flat Responses API layout of a chat ``allowed_tools`` choice, or ``None`` when an entry is malformed or + carries conflicting keys, so the caller sends the original rather than a rewrite that could mean something looser.""" + flattened: Final = tuple(_responses_allowed_tool(tool) for tool in choice.allowed_tools.tools) + if any(tool is None for tool in flattened): + return None + return ToolChoiceAllowedParam( + type="allowed_tools", + mode=choice.allowed_tools.mode, + tools=[dict(tool) for tool in flattened if tool is not None], + ) diff --git a/litellm/llms/azure/chat/gpt_transformation.py b/litellm/llms/azure/chat/gpt_transformation.py index fdb012756fc..0fe7caa1e16 100644 --- a/litellm/llms/azure/chat/gpt_transformation.py +++ b/litellm/llms/azure/chat/gpt_transformation.py @@ -1,11 +1,17 @@ from collections.abc import Mapping, Sequence from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final +from typing import ( + TYPE_CHECKING, + Any, + Final, + cast, # noqa: TID251 # non_default_params reaches the mapper as an untyped dict +) 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, @@ -245,12 +251,12 @@ 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 litellm.drop_params is True or (drop_params is not None and drop_params is True): - pass - else: + 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)): 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. 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. 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""", ) elif value == "required" and ( api_version_year == "2024" and api_version_month <= "05" diff --git a/litellm/main.py b/litellm/main.py index 78767b03eb7..38dd9f90be9 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -70,6 +70,9 @@ 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, +) from litellm.litellm_core_utils.asyncify import asyncify, run_async_function from litellm.litellm_core_utils.audio_utils.utils import ( calculate_request_duration, @@ -5956,6 +5959,14 @@ 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( + 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, + ) + if provider_specific_header is not None: headers.update( ProviderSpecificHeaderUtils.get_provider_specific_headers( diff --git a/litellm/utils.py b/litellm/utils.py index 55423a42260..5a0fc18dfdc 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -83,6 +83,10 @@ from litellm.constants import ( PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO, TOOL_CHOICE_OBJECT_TOKEN_COUNT, ) +from litellm.litellm_core_utils.allowed_tools_tool_choice import ( + CHAT_ALLOWED_TOOLS_LAYOUT, + parse_chat_allowed_tools_tool_choice, +) from litellm.litellm_core_utils.core_helpers import ( bind_budget_reservation_to_callbacks, normalize_drop_params, @@ -9470,9 +9474,28 @@ 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 + ) + ) + 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="", + ) # Standard OpenAI format: {"type": "function", "function": {...}} - if tool_choice.get("type") is None or tool_choice.get("function") is None: + if tool_choice_type != "allowed_tools" and (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 558350a790a..9e40e128952 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 @@ -3064,6 +3064,38 @@ def test_map_optional_params_tool_choice_chat_nested_to_responses_api(): {"type": "custom", "name": "ApplyPatch"}, ), ({"type": "custom"}, {"type": "custom"}), + ( + { + "type": "allowed_tools", + "allowed_tools": { + "mode": "required", + "tools": [ + {"type": "function", "function": {"name": "get_time"}}, + {"type": "custom", "custom": {"name": "apply_patch"}}, + {"type": "mcp", "server_label": "docs"}, + {"type": "function", "name": "already_flat"}, + ], + }, + }, + { + "type": "allowed_tools", + "mode": "required", + "tools": [ + {"type": "function", "name": "get_time"}, + {"type": "custom", "name": "apply_patch"}, + {"type": "mcp", "server_label": "docs"}, + {"type": "function", "name": "already_flat"}, + ], + }, + ), + ( + {"type": "allowed_tools", "allowed_tools": {"mode": "auto", "tools": [{"type": "function"}]}}, + {"type": "allowed_tools", "allowed_tools": {"mode": "auto", "tools": [{"type": "function"}]}}, + ), + ( + {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "function", "name": "get_time"}]}, + {"type": "allowed_tools", "mode": "auto", "tools": [{"type": "function", "name": "get_time"}]}, + ), ], ) def test_normalize_tool_choice_for_responses_api(tool_choice, expected): @@ -4105,6 +4137,16 @@ def test_output_item_done_with_stream_map_keeps_empty_delta(): ({"type": "required"}, "required"), ("auto", "auto"), ({"type": "function", "function": {"name": "get_weather"}}, {"type": "function", "name": "get_weather"}), + ( + { + "type": "allowed_tools", + "allowed_tools": { + "mode": "required", + "tools": [{"type": "function", "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( @@ -4177,6 +4219,7 @@ async def test_acompletion_bridge_normalizes_tool_choice_on_the_wire( post_kwargs = mock_post.call_args.kwargs request_body = post_kwargs["json"] if "json" in post_kwargs else json.loads(post_kwargs["data"]) assert request_body["tool_choice"] == expected_wire_tool_choice + assert [tool["name"] for tool in request_body["tools"]] == ["get_weather"] def _make_incomplete_responses_api_response( 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 new file mode 100644 index 00000000000..6ebbdcf73bf --- /dev/null +++ b/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py @@ -0,0 +1,100 @@ +import copy +from collections.abc import Mapping +from typing import Final + +import pytest + +from litellm.litellm_core_utils.allowed_tools_tool_choice import ( + ChatAllowedToolsToolChoice, + is_allowed_tools_tool_choice, + parse_chat_allowed_tools_tool_choice, + responses_allowed_tools_tool_choice, +) + +_CHAT_LAYOUT: Final[dict[str, object]] = { + "type": "allowed_tools", + "allowed_tools": { + "mode": "auto", + "tools": [ + {"type": "function", "function": {"name": "get_time"}}, + {"type": "custom", "custom": {"name": "apply_patch", "x-vendor": 1}, "x-trace": "a"}, + {"type": "mcp", "server_label": "docs", "server_url": "https://docs.example"}, + {"type": "function", "name": "already_flat"}, + ], + }, +} + + +def _parsed(tool_choice: Mapping[str, object]) -> ChatAllowedToolsToolChoice: + parsed: Final = parse_chat_allowed_tools_tool_choice(tool_choice) + assert parsed is not None + return parsed + + +def test_responses_layout_lifts_mode_and_tools_and_flattens_names_in_order(): + assert responses_allowed_tools_tool_choice(_parsed(_CHAT_LAYOUT)) == { + "type": "allowed_tools", + "mode": "auto", + "tools": [ + {"type": "function", "name": "get_time"}, + {"type": "custom", "name": "apply_patch", "x-vendor": 1, "x-trace": "a"}, + {"type": "mcp", "server_label": "docs", "server_url": "https://docs.example"}, + {"type": "function", "name": "already_flat"}, + ], + } + + +def test_responses_layout_leaves_the_caller_dict_untouched(): + original: Final = copy.deepcopy(_CHAT_LAYOUT) + flat: Final = responses_allowed_tools_tool_choice(_parsed(_CHAT_LAYOUT)) + assert flat is not None + next(iter(flat["tools"]))["name"] = "mutated" + assert _CHAT_LAYOUT == original + + +@pytest.mark.parametrize( + "tool", + [ + {"type": "function"}, + {"type": "function", "name": ""}, + {"type": "function", "function": {}}, + {"type": "function", "function": "get_time"}, + {"type": "custom", "custom": {"name": ""}}, + {"type": "function", "name": "get_time", "function": {"name": "get_weather"}}, + {"type": "function", "function": {"name": "get_time", "type": "custom"}}, + ], +) +def test_responses_layout_refuses_to_rewrite_a_malformed_or_conflicting_entry(tool: dict[str, object]): + tool_choice: Final[dict[str, object]] = { + "type": "allowed_tools", + "allowed_tools": {"mode": "required", "tools": [tool]}, + } + assert responses_allowed_tools_tool_choice(_parsed(tool_choice)) is None + + +@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": "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 + + +@pytest.mark.parametrize( + ("tool_choice", "expected"), + [ + (_CHAT_LAYOUT, True), + ({"type": "allowed_tools"}, True), + ({"type": "function"}, False), + ("auto", False), + (None, False), + ], +) +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 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 1f1ff6a5748..dd29ee83bc6 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 @@ -75,6 +75,42 @@ def test_map_openai_params_with_preview_api_version(): ) +def _map_tool_choice_on_pre_tool_choice_api_version(tool_choice: dict[str, object]) -> dict[str, object]: + return AzureOpenAIConfig().map_openai_params( + {"tool_choice": tool_choice, "tools": [{"type": "function", "function": {"name": "get_time"}}]}, + {}, + "gpt-4o", + drop_params=True, + api_version="2023-07-01-preview", + ) + + +def test_map_openai_params_drop_params_still_drops_a_named_tool_choice_on_old_api_version(): + mapped: Final = _map_tool_choice_on_pre_tool_choice_api_version( + {"type": "function", "function": {"name": "get_time"}} + ) + 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"}]}} + ) + assert exc.value.status_code == 400 + + +def test_map_openai_params_forwards_allowed_tools_on_current_api_version(): + tool_choice: Final[dict[str, object]] = { + "type": "allowed_tools", + "allowed_tools": {"mode": "required", "tools": [{"type": "function", "function": {"name": "get_time"}}]}, + } + mapped: Final = AzureOpenAIConfig().map_openai_params( + {"tool_choice": tool_choice}, {}, "gpt-4o", drop_params=True, api_version="2025-04-01-preview" + ) + assert mapped["tool_choice"] == tool_choice + + def test_transform_request_hoists_tool_message_image(): """Azure builds its request via convert_to_azure_openai_messages without the OpenAIGPTConfig._transform_messages pipeline, so transform_request must hoist diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 85864b2024e..95c5958c748 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -553,6 +553,63 @@ def test_completion_strips_eager_input_streaming_before_openai(respx_mock: respx assert sent_tool["function"]["name"] == "write_file" +_ALLOWED_TOOLS_TOOL_CHOICE: Final = { + "type": "allowed_tools", + "allowed_tools": {"mode": "required", "tools": [{"type": "function", "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("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 +): + 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=_ALLOWED_TOOLS_TOOL_CHOICE, + api_base="http://localhost:12346/v1", + api_key="fake_api_key", + api_version="2025-04-01-preview", + ) + + assert mock_route.called + sent: Final = json.loads(respx_mock.calls[0].request.content) + assert sent["tool_choice"] == _ALLOWED_TOOLS_TOOL_CHOICE + assert [tool["function"]["name"] for tool in sent["tools"]] == ["get_time", "get_weather"] + + +@pytest.mark.parametrize("drop_params", [False, True]) +@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 +): + with pytest.raises( + litellm.BadRequestError, match=r"allowed_tools.*only supported for azure, litellm_proxy, openai" + ) as exc: + litellm.completion( + model=model, + messages=[{"role": "user", "content": "What time is it?"}], + tools=_TWO_TOOLS, + tool_choice=_ALLOWED_TOOLS_TOOL_CHOICE, + drop_params=drop_params, + api_key="fake_api_key", + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + 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_utils.py b/tests/unit/test_utils.py index a9b56568afb..19b3f317927 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -8651,6 +8651,40 @@ def test_validate_tool_choice_invalid_dict_is_a_400(tool_choice): assert exc_info.value.model == MODEL +@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 = { + "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 + + +@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"}, + {"type": "allowed_tools", "allowed_tools": {"mode": "required"}}, + {"type": "allowed_tools", "allowed_tools": {"mode": "required", "tools": []}}, + {"type": "allowed_tools", "allowed_tools": {"mode": "none", "tools": [{"type": "function"}]}}, + {"type": "allowed_tools", "allowed_tools": {"tools": [{"type": "function", "function": {"name": "a"}}]}}, + {"type": "allowed_tools", "allowed_tools": {"mode": "auto", "tools": "get_time"}}, + ], +) +def test_validate_tool_choice_allowed_tools_malformed_is_a_400(tool_choice: dict[str, object]): + with pytest.raises(litellm.BadRequestError, match="allowed_tools") as exc_info: + validate_chat_completion_tool_choice(tool_choice, model=MODEL) + assert exc_info.value.status_code == 400 + assert exc_info.value.model == MODEL + assert '"mode": "auto" | "required"' in str(exc_info.value) + + @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):