scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101

fix(core): harden concatenated tool argument recovery

Browse source 
- Reject partial recovery: a malformed tail now fails the whole parse
  instead of executing an incomplete tool sequence (strict split).
- Cap recovered objects per call (8) to prevent tool-call amplification.
- Rewrite extractors as one-shot tuple constructions per repo style
  (LIT001/LIT002) and mark NormalizedToolCall fields ReadOnly (LIT012).
This commit is contained in:
ggbond 2026-09-11 10:50:15 +08:00
parent 7704317022
commit 3d4e3549fe
4 changed files with 216 additions and 78 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

28
litellm/litellm_core_utils/prompt_templates/common_utils.py
View file

@ -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 '<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' "
@ -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)

152
litellm/litellm_core_utils/prompt_templates/factory.py
View file

@ -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:

48
tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py
View file

@ -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.
#

66
tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py
View file

@ -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,