Merge 0343c96558 into 0c98afa780

litellm/litellm_core_utils/prompt_templates/common_utils.py

@@ -2314,18 +2314,28 @@ 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,
     context: str | None = None,
+    allow_concatenated: bool = False,
 ) -> Any:
     """
     Parse tool call arguments from a JSON string.
 
     When the JSON is malformed (e.g. truncated by the model), this function
-    attempts a lightweight repair (closing unmatched brackets/braces) before
-    raising an error.  A warning is logged whenever repair succeeds so that
-    callers are aware the arguments were not perfectly formed.
+    attempts a lightweight repair (closing unmatched brackets/braces).  If
+    ``allow_concatenated`` is true, it also attempts to split concatenated
+    JSON objects.  A warning is logged whenever repair or splitting succeeds
+    so that callers are aware the arguments were not perfectly formed.
 
     Args:
         arguments: The JSON string containing tool arguments, or None.
@@ -2334,8 +2344,10 @@ def parse_tool_call_arguments(
 
     Returns:
         Parsed arguments (usually a dict, but may be any JSON-deserializable
-        type such as list, str, int, float, or None).  Returns empty dict if
-        arguments is None or empty.
+        type such as list, str, int, float, or None).  When
+        ``allow_concatenated`` is true, concatenated JSON objects are
+        returned as a list of dicts.  Returns empty dict if arguments is
+        None or empty.
 
     Raises:
         ValueError: If the arguments string is not valid JSON and cannot be repaired.
@@ -2360,6 +2372,31 @@ def parse_tool_call_arguments(
             )
             return repaired
 
+        if allow_concatenated:
+            # 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 '<unknown>'}' "
+                    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' "
+                    "(%s). Original (%d chars): %.200s%s",
+                    len(split_arguments),
+                    tool_name or "<unknown>",
+                    context or "unknown context",
+                    len(arguments),
+                    arguments,
+                    "..." if len(arguments) > 200 else "",
+                )
+                return split_arguments
+
         error_parts: Final = ["Failed to parse tool call arguments"]
 
         if tool_name:
@@ -2372,7 +2409,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.
@@ -2390,10 +2427,13 @@ def split_concatenated_json_objects(raw: str) -> list[dict[str, object]]:
     The walk degrades gracefully: if the string is malformed or truncated
     (e.g. a stream that ended mid-tool-call), whatever complete objects were
     parsed before the bad tail are returned and the remainder is discarded
-    with a warning, rather than raising.  The sole caller
-    (``_convert_to_bedrock_tool_call_invoke``) treats an empty result as
+    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]
@@ -2429,6 +2469,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)

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,12 +5349,14 @@ def get_attribute_or_key(tool_or_function, attribute, default=None):
 
 
 class NormalizedToolCall(TypedDict):
-    id: str | None
-    name: str | None
-    arguments: dict[str, object]
+    id: ReadOnly[str | None]
+    name: ReadOnly[str | None]
+    arguments: ReadOnly[Mapping[str, object]]
 
 
-def _parse_tool_call_arguments(raw: object, tool_name: str | None, context: str) -> dict[str, object]:
+def _parse_tool_call_arguments(
+    raw: object, tool_name: str | None, context: str
+) -> Mapping[str, object] | Sequence[Mapping[str, object]]:
     # Anthropic's tool_use blocks already carry a parsed dict in "input";
     # chat completions and the Responses API carry a JSON string that may be
     # truncated by the model, so route those through the repair-aware parser.
@@ -5363,89 +5366,141 @@ def _parse_tool_call_arguments(raw: object, tool_name: str | None, context: str)
         return {}
     normalized_raw: Final = "{}" if raw == REDACTED_BY_LITELLM else raw
     from litellm.litellm_core_utils.prompt_templates.common_utils import (
+        MAX_RECOVERED_ARGUMENT_OBJECTS,
         parse_tool_call_arguments,
     )
 
     try:
-        parsed: Final = parse_tool_call_arguments(normalized_raw, tool_name=tool_name, context=context)
+        direct: Final = json.loads(normalized_raw)
+    except json.JSONDecodeError:
+        pass  # malformed JSON: fall through to repair / concatenated recovery
+    else:
+        # Valid JSON keeps the historical object-only contract: a non-object
+        # root (e.g. a JSON array of objects) degrades to {} rather than
+        # becoming an uncapped expansion vector.
+        return direct if isinstance(direct, dict) else {}
+
+    try:
+        parsed: Final = parse_tool_call_arguments(
+            normalized_raw,
+            tool_name=tool_name,
+            context=context,
+            allow_concatenated=True,
+        )
     except ValueError as e:
         verbose_logger.warning("Failed to parse tool call arguments: %s", e)
         return {}
-    return parsed if isinstance(parsed, dict) else {}
+    if isinstance(parsed, dict):
+        return parsed
+    # Only malformed-input recovery (repair or concatenated split) can yield a
+    # sequence here; hold it to the same per-call bound as concatenated
+    # recovery so expansion is never an amplification vector.
+    if (
+        isinstance(parsed, list)
+        and parsed
+        and all(isinstance(item, dict) for item in parsed)
+        and len(parsed) <= MAX_RECOVERED_ARGUMENT_OBJECTS
+    ):
+        return parsed
+    return {}
+
+
+def _normalized_tool_calls(
+    call_id: str | None,
+    name: str | None,
+    parsed_arguments: Mapping[str, object] | Sequence[Mapping[str, object]],
+) -> tuple[NormalizedToolCall, ...]:
+    if isinstance(parsed_arguments, Mapping):
+        return (NormalizedToolCall(id=call_id, name=name, arguments=parsed_arguments),)
+    return tuple(
+        NormalizedToolCall(
+            id=call_id if argument_index == 0 or call_id is None else f"{call_id}_{argument_index}",
+            name=name,
+            arguments=arguments,
+        )
+        for argument_index, arguments in enumerate(parsed_arguments)
+    )
+
+
+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.append(
-            NormalizedToolCall(
-                id=get_attribute_or_key(tc, "id"),
-                name=name,
-                arguments=_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.append(
-            NormalizedToolCall(
-                id=get_attribute_or_key(item, "call_id") or get_attribute_or_key(item, "id"),
-                name=name,
-                arguments=_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]:
@@ -5466,17 +5521,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:

tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py

@@ -20,6 +20,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
     handle_any_messages_to_chat_completion_str_messages_conversion,
     hoist_images_from_tool_messages,
     is_encrypted_reasoning_block,
+    parse_tool_call_arguments,
     responses_reasoning_items_from_thinking_blocks,
     split_concatenated_json_objects,
     strip_encrypted_reasoning_from_messages,
@@ -194,6 +195,68 @@ def test_convert_prefix_message_to_non_prefix_messages():
 # ── split_concatenated_json_objects tests ──
 
 
+def test_parse_tool_call_arguments_concatenated_objects_opt_in():
+    """Concatenated JSON objects are split only when explicitly enabled."""
+    raw = '{"city": "Paris"}{"units": "celsius"}'
+    result = parse_tool_call_arguments(
+        raw,
+        tool_name="weather",
+        context="chat completions",
+        allow_concatenated=True,
+    )
+    assert result == [{"city": "Paris"}, {"units": "celsius"}]
+
+
+def test_parse_tool_call_arguments_concatenated_objects_disabled_by_default():
+    """Existing parser callers keep the original strict behavior."""
+    raw = '{"city": "Paris"}{"units": "celsius"}'
+    with pytest.raises(ValueError, match="Extra data"):
+        parse_tool_call_arguments(raw, tool_name="weather", context="chat completions")
+
+
+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 == [{"n": index} for index in range(8)]
+
+
 def test_split_concatenated_json_single_object():
     """A single valid JSON object is returned as a one-element list."""
     result = split_concatenated_json_objects('{"location": "Boston"}')
@@ -268,6 +331,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.
 #

tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py

@@ -3401,6 +3401,306 @@ def test_get_tool_calls_from_response_warns_for_malformed_arguments(caplog):
     assert "Failed to parse tool call arguments" in caplog.text
 
 
+def test_get_tool_calls_from_response_splits_concatenated_arguments():
+    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": "second"}',
+                            },
+                        }
+                    ]
+                }
+            }
+        ]
+    }
+
+    tool_calls: Final = get_tool_calls_from_response(response)
+
+    assert tool_calls == [
+        {"id": "call_1", "name": "search", "arguments": {"query": "first"}},
+        {"id": "call_1_1", "name": "search", "arguments": {"query": "second"}},
+    ]
+
+
+def test_get_tool_calls_from_response_splits_concatenated_responses_arguments():
+    from litellm.litellm_core_utils.prompt_templates.factory import (
+        get_tool_calls_from_response,
+    )
+
+    response: Final = {
+        "choices": None,
+        "output": [
+            {
+                "type": "function_call",
+                "id": "fc_1",
+                "call_id": "call_1",
+                "name": "search",
+                "arguments": '{"query": "first"}{"query": "second"}',
+            }
+        ],
+    }
+
+    tool_calls: Final = get_tool_calls_from_response(response)
+
+    assert tool_calls == [
+        {"id": "call_1", "name": "search", "arguments": {"query": "first"}},
+        {"id": "call_1_1", "name": "search", "arguments": {"query": "second"}},
+    ]
+
+
+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_valid_json_array_not_expanded():
+    """
+    A valid JSON array of objects is not concatenated recovery: it keeps the
+    historical object-only semantics (degrades to {}) and never expands, so
+    it cannot bypass the per-call expansion cap.
+    """
+    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": json.dumps(
+                                    [{"query": f"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_valid_json_array_of_two_not_expanded():
+    """Even a small valid JSON array degrades to {} (object-only contract)."""
+    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": json.dumps([{"query": "a"}, {"query": "b"}]),
+                            },
+                        }
+                    ]
+                }
+            }
+        ]
+    }
+
+    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,
+    )
+
+    response: Final = {
+        "choices": [
+            {
+                "message": {
+                    "tool_calls": [
+                        {
+                            "id": "call_1",
+                            "function": {
+                                "name": "search",
+                                "arguments": "[1, 2, 3]",
+                            },
+                        }
+                    ]
+                }
+            }
+        ]
+    }
+
+    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_accepts_dict_arguments():
+    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"},
+                            },
+                        }
+                    ]
+                }
+            }
+        ]
+    }
+
+    tool_calls: Final = get_tool_calls_from_response(response)
+
+    assert tool_calls == [{"id": "call_1", "name": "search", "arguments": {"query": "first"}}]
+
+
+def test_get_tool_calls_from_response_ignores_non_string_arguments():
+    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": 123},
+                        }
+                    ]
+                }
+            }
+        ]
+    }
+
+    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_ignores_malformed_tool_call_entries():
+    from litellm.litellm_core_utils.prompt_templates.factory import (
+        get_tool_calls_from_response,
+    )
+
+    response: Final = {
+        "choices": [
+            {
+                "message": {
+                    "tool_calls": [{"id": "call_1"}],
+                }
+            }
+        ],
+        "output": None,
+    }
+
+    assert get_tool_calls_from_response(response) == []
+
+
+def test_get_tool_calls_from_response_ignores_non_function_output_items():
+    from litellm.litellm_core_utils.prompt_templates.factory import (
+        get_tool_calls_from_response,
+    )
+
+    response: Final = {
+        "choices": None,
+        "output": [{"type": "message", "id": "msg_1"}],
+    }
+
+    assert get_tool_calls_from_response(response) == []
+
+
 def test_group_tool_exchanges_pairs_assistant_with_its_tool_rows():
     from litellm.litellm_core_utils.prompt_templates.factory import group_tool_exchanges