diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 18eb19ff56c..294fa2324d3 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -2203,6 +2203,14 @@ def _attempt_json_repair(s: str) -> object | None: return None +# Upper bound on how many argument objects a single tool-call string may +# recover into. Concatenated recovery fans one provider tool call out into +# multiple proxy-side actions (e.g. guardrail retrievals), so an unbounded +# split would let a crafted model response amplify one call into arbitrarily +# many authenticated downstream requests. +MAX_RECOVERED_ARGUMENT_OBJECTS: Final = 8 + + def parse_tool_call_arguments( arguments: str | None, tool_name: str | None = None, @@ -2254,7 +2262,17 @@ def parse_tool_call_arguments( return repaired if allow_concatenated: - split_arguments: Final = split_concatenated_json_objects(arguments) + # strict=True: refuse partial recovery -- recovering the complete + # prefix while silently dropping a malformed tail would execute an + # incomplete tool sequence. + split_arguments: Final = split_concatenated_json_objects(arguments, strict=True) + if len(split_arguments) > MAX_RECOVERED_ARGUMENT_OBJECTS: + raise ValueError( + f"Failed to parse tool call arguments for tool '{tool_name or ''}' " + f"({context or 'unknown context'}): recovered {len(split_arguments)} " + f"concatenated argument objects, exceeding the per-call limit of " + f"{MAX_RECOVERED_ARGUMENT_OBJECTS}. Arguments: {arguments}" + ) from original_error if split_arguments: verbose_logger.warning( "Recovered %d concatenated tool call argument object(s) for tool '%s' " @@ -2280,7 +2298,7 @@ def parse_tool_call_arguments( raise ValueError(error_message) from original_error -def split_concatenated_json_objects(raw: str) -> list[dict[str, object]]: +def split_concatenated_json_objects(raw: str, strict: bool = False) -> list[dict[str, object]]: """ Split a string that contains one or more concatenated JSON objects into a list of parsed dicts. @@ -2301,6 +2319,10 @@ def split_concatenated_json_objects(raw: str) -> list[dict[str, object]]: with a warning, rather than raising. Callers treat an empty result as ``input={}`` so the conversation can continue instead of hard-failing. + Pass ``strict=True`` to reject partial recovery entirely: an unparseable + tail then yields an empty list, so callers that would execute the + recovered objects never run an incomplete sequence. + Returns ------- list[dict] @@ -2336,6 +2358,8 @@ def split_concatenated_json_objects(raw: str) -> list[dict[str, object]]: idx, e, ) + if strict: + return [] break if isinstance(obj, dict): results.append(obj) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index b364c751f82..75e04371be3 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -10,6 +10,7 @@ from enum import Enum from typing import Any, Final, TypedDict, cast, overload from jinja2.sandbox import ImmutableSandboxedEnvironment +from typing_extensions import ReadOnly import litellm import litellm.types @@ -5348,9 +5349,9 @@ def get_attribute_or_key(tool_or_function, attribute, default=None): class NormalizedToolCall(TypedDict): - id: str | None - name: str | None - arguments: Mapping[str, object] + id: ReadOnly[str | None] + name: ReadOnly[str | None] + arguments: ReadOnly[Mapping[str, object]] def _parse_tool_call_arguments( @@ -5402,78 +5403,85 @@ def _normalized_tool_calls( ) +def _chat_tool_calls_of_choice(choice: object) -> tuple[object, ...]: + message: Final = get_attribute_or_key(choice, "message", None) + choice_tool_calls: Final = get_attribute_or_key(message, "tool_calls", None) if message else None + if isinstance(choice_tool_calls, list): + return tuple(choice_tool_calls) + return () + + +def _normalized_from_chat_tool_call(tool_call: object) -> tuple[NormalizedToolCall, ...]: + function: Final = get_attribute_or_key(tool_call, "function", None) + if function is None: + return () + name: Final = get_attribute_or_key(function, "name") + return _normalized_tool_calls( + get_attribute_or_key(tool_call, "id"), + name, + _parse_tool_call_arguments( + get_attribute_or_key(function, "arguments", "{}"), + tool_name=name, + context="chat completions", + ), + ) + + def _tool_calls_from_chat_completion_response( response: object, include_all_choices: bool = False -) -> list[NormalizedToolCall]: +) -> tuple[NormalizedToolCall, ...]: choices: Final = get_attribute_or_key(response, "choices", None) if not (isinstance(choices, list) and choices): - return [] - tool_calls: Final[list[object]] = [] - for choice in choices if include_all_choices else choices[:1]: - message = get_attribute_or_key(choice, "message", None) - choice_tool_calls = get_attribute_or_key(message, "tool_calls", None) if message else None - if isinstance(choice_tool_calls, list): - tool_calls.extend(choice_tool_calls) - result: Final[list[NormalizedToolCall]] = [] - for tc in tool_calls: - fn = get_attribute_or_key(tc, "function", None) - if fn is None: - continue - name = get_attribute_or_key(fn, "name") - result.extend( - _normalized_tool_calls( - get_attribute_or_key(tc, "id"), - name, - _parse_tool_call_arguments( - get_attribute_or_key(fn, "arguments", "{}"), - tool_name=name, - context="chat completions", - ), - ) - ) - return result + return () + selected_choices: Final = choices if include_all_choices else choices[:1] + return tuple( + normalized + for choice in selected_choices + for tool_call in _chat_tool_calls_of_choice(choice) + for normalized in _normalized_from_chat_tool_call(tool_call) + ) -def _tool_calls_from_responses_api_response(response: object) -> list[NormalizedToolCall]: +def _normalized_from_responses_item(item: object) -> tuple[NormalizedToolCall, ...]: + if get_attribute_or_key(item, "type") != "function_call": + return () + name: Final = get_attribute_or_key(item, "name") + return _normalized_tool_calls( + get_attribute_or_key(item, "call_id") or get_attribute_or_key(item, "id"), + name, + _parse_tool_call_arguments( + get_attribute_or_key(item, "arguments", "{}"), + tool_name=name, + context="responses API", + ), + ) + + +def _tool_calls_from_responses_api_response(response: object) -> tuple[NormalizedToolCall, ...]: output: Final = get_attribute_or_key(response, "output", None) if not isinstance(output, list): - return [] - result: Final[list[NormalizedToolCall]] = [] - for item in output: - if get_attribute_or_key(item, "type") != "function_call": - continue - name = get_attribute_or_key(item, "name") - result.extend( - _normalized_tool_calls( - get_attribute_or_key(item, "call_id") or get_attribute_or_key(item, "id"), - name, - _parse_tool_call_arguments( - get_attribute_or_key(item, "arguments", "{}"), - tool_name=name, - context="responses API", - ), - ) - ) - return result + return () + return tuple(normalized for item in output for normalized in _normalized_from_responses_item(item)) -def _tool_calls_from_anthropic_messages_response(response: object) -> list[NormalizedToolCall]: +def _normalized_from_anthropic_block(block: object) -> tuple[NormalizedToolCall, ...]: + if get_attribute_or_key(block, "type") != "tool_use": + return () + raw_input: Final = get_attribute_or_key(block, "input", {}) + return ( + NormalizedToolCall( + id=get_attribute_or_key(block, "id"), + name=get_attribute_or_key(block, "name"), + arguments=raw_input if isinstance(raw_input, dict) else {}, + ), + ) + + +def _tool_calls_from_anthropic_messages_response(response: object) -> tuple[NormalizedToolCall, ...]: content: Final = get_attribute_or_key(response, "content", None) if not isinstance(content, list): - return [] - result: Final[list[NormalizedToolCall]] = [] - for block in content: - if get_attribute_or_key(block, "type") != "tool_use": - continue - raw_input = get_attribute_or_key(block, "input", {}) - result.append( - NormalizedToolCall( - id=get_attribute_or_key(block, "id"), - name=get_attribute_or_key(block, "name"), - arguments=raw_input if isinstance(raw_input, dict) else {}, - ) - ) - return result + return () + return tuple(normalized for block in content for normalized in _normalized_from_anthropic_block(block)) def get_tool_calls_from_response(response: object, include_all_choices: bool = False) -> list[NormalizedToolCall]: @@ -5494,17 +5502,15 @@ def get_tool_calls_from_response(response: object, include_all_choices: bool = F Callers that only care about a specific tool should filter the result by ``name`` themselves -- this returns every tool call found. """ - chat_tool_calls = _tool_calls_from_chat_completion_response(response, include_all_choices=include_all_choices) + chat_tool_calls: Final = _tool_calls_from_chat_completion_response( + response, include_all_choices=include_all_choices + ) if chat_tool_calls: - return chat_tool_calls - for extractor in ( - _tool_calls_from_responses_api_response, - _tool_calls_from_anthropic_messages_response, - ): - tool_calls = extractor(response) - if tool_calls: - return tool_calls - return [] + return list(chat_tool_calls) + responses_tool_calls: Final = _tool_calls_from_responses_api_response(response) + if responses_tool_calls: + return list(responses_tool_calls) + return list(_tool_calls_from_anthropic_messages_response(response)) def has_tool_with_name(tools: object, tool_name: str) -> bool: diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index 63b26be0f26..8e6049dd6ab 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -211,16 +211,47 @@ def test_parse_tool_call_arguments_concatenated_objects_disabled_by_default(): parse_tool_call_arguments(raw, tool_name="weather", context="chat completions") -def test_parse_tool_call_arguments_concatenated_partial_tail(): - """Only complete objects before a truncated tail are recovered.""" +def test_parse_tool_call_arguments_concatenated_partial_tail_rejected(): + """ + A truncated tail rejects the whole recovery instead of executing an + incomplete tool sequence (only the complete prefix would have run). + """ raw = '{"city": "Paris"}{"units": "celsius"}{"forecast":' + with pytest.raises(ValueError, match="Failed to parse tool call arguments"): + parse_tool_call_arguments( + raw, + tool_name="weather", + context="chat completions", + allow_concatenated=True, + ) + + +def test_parse_tool_call_arguments_concatenated_over_limit_rejected(): + """ + Recovery fans one provider tool call out into many proxy-side actions, so + it is capped: more objects than the per-call limit raises instead of + amplifying the call. + """ + raw = "".join(f'{{"n": {index}}}' for index in range(9)) + with pytest.raises(ValueError, match="exceeding the per-call limit"): + parse_tool_call_arguments( + raw, + tool_name="weather", + context="chat completions", + allow_concatenated=True, + ) + + +def test_parse_tool_call_arguments_concatenated_at_limit_accepted(): + """Exactly the per-call limit of objects still recovers.""" + raw = "".join(f'{{"n": {index}}}' for index in range(8)) result = parse_tool_call_arguments( raw, tool_name="weather", context="chat completions", allow_concatenated=True, ) - assert result == [{"city": "Paris"}, {"units": "celsius"}] + assert result == [{"n": index} for index in range(8)] def test_split_concatenated_json_single_object(): @@ -297,6 +328,17 @@ def test_split_concatenated_json_salvages_prefix_before_truncated_tail(): assert result == [{"a": 1}, {"b": 2}] +def test_split_concatenated_json_strict_rejects_truncated_tail(): + """ + strict=True discards the whole result when a tail cannot be parsed, so + callers that execute the recovered objects never run a partial sequence. + """ + raw = '{"a": 1}{"b": 2}{"c":' + assert split_concatenated_json_objects(raw, strict=True) == [] + # Default mode keeps the graceful salvage behavior (Bedrock replay path). + assert split_concatenated_json_objects(raw) == [{"a": 1}, {"b": 2}] + + # --------------------------------------------------------------------------- # Regression tests for non-OpenAI file content blocks. # diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index 48f424d58c0..1762c78f47e 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -3458,6 +3458,72 @@ def test_get_tool_calls_from_response_splits_concatenated_responses_arguments(): ] +def test_get_tool_calls_from_response_rejects_partial_concatenated_tail(): + """ + A truncated tail must not execute an incomplete tool sequence: the whole + recovery is rejected and the single original call degrades to {}. + """ + from litellm.litellm_core_utils.prompt_templates.factory import ( + get_tool_calls_from_response, + ) + + response: Final = { + "choices": [ + { + "message": { + "tool_calls": [ + { + "id": "call_1", + "function": { + "name": "search", + "arguments": '{"query": "first"}{"query":', + }, + } + ] + } + } + ] + } + + tool_calls: Final = get_tool_calls_from_response(response) + + assert tool_calls == [{"id": "call_1", "name": "search", "arguments": {}}] + + +def test_get_tool_calls_from_response_caps_concatenated_expansion(): + """ + One provider tool call must not amplify into arbitrarily many proxy-side + calls: recovering more objects than the per-call limit is rejected. + """ + from litellm.litellm_core_utils.prompt_templates.factory import ( + get_tool_calls_from_response, + ) + + response: Final = { + "choices": [ + { + "message": { + "tool_calls": [ + { + "id": "call_1", + "function": { + "name": "search", + "arguments": "".join( + f'{{"query": "q{index}"}}' for index in range(9) + ), + }, + } + ] + } + } + ] + } + + tool_calls: Final = get_tool_calls_from_response(response) + + assert tool_calls == [{"id": "call_1", "name": "search", "arguments": {}}] + + def test_get_tool_calls_from_response_non_object_arguments_returns_empty(): from litellm.litellm_core_utils.prompt_templates.factory import ( get_tool_calls_from_response,