From 4d97818f98ec64c3c6075819c740661c996ba7c9 Mon Sep 17 00:00:00 2001 From: giulio-leone Date: Wed, 4 Mar 2026 06:01:42 +0100 Subject: [PATCH] fix(tools): gracefully repair truncated JSON in tool call arguments --- .../prompt_templates/common_utils.py | 95 ++++++++++-- .../prompt_templates/factory.py | 10 +- tests/llm_translation/test_prompt_factory.py | 144 ++++++++++++++++++ 3 files changed, 237 insertions(+), 12 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 125f2585a3..d59b8d8871 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -20,6 +20,7 @@ from typing import ( cast, ) +from litellm import verbose_logger from litellm.router_utils.batch_utils import InMemoryFile from litellm.types.llms.openai import ( AllMessageValues, @@ -1278,16 +1279,76 @@ def extract_images_from_message(message: AllMessageValues) -> List[str]: return images +def _attempt_json_repair(s: str) -> Optional[Any]: + """ + Attempt to repair truncated JSON produced by LLM tool calls. + + Handles the most common truncation patterns where the model generates + valid JSON that is cut short (missing closing brackets/braces). + + Returns the parsed value on success, or None if repair fails. + """ + import json + + stripped = s.rstrip() + if not stripped: + return None + + # Track the stack of unmatched openers to respect nesting order + opener_stack: list = [] + in_string = False + escape_next = False + + for ch in stripped: + if escape_next: + escape_next = False + continue + if ch == "\\": + if in_string: + escape_next = True + continue + if ch == '"': + in_string = not in_string + continue + if in_string: + continue + if ch == "{": + opener_stack.append("}") + elif ch == "[": + opener_stack.append("]") + elif ch in ("}", "]"): + if opener_stack and opener_stack[-1] == ch: + opener_stack.pop() + + if not opener_stack: + return None + + # Remove trailing comma before we close brackets + candidate = stripped.rstrip(",") + + # Close in reverse order of opening (respects nesting) + candidate += "".join(reversed(opener_stack)) + + try: + return json.loads(candidate) + except json.JSONDecodeError: + pass + + return None + + def parse_tool_call_arguments( arguments: Optional[str], tool_name: Optional[str] = None, context: Optional[str] = None, -) -> Dict[str, Any]: +) -> Any: """ Parse tool call arguments from a JSON string. - This function handles malformed JSON gracefully by raising a ValueError - with context about what failed and what the problematic input was. + 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. Args: arguments: The JSON string containing tool arguments, or None. @@ -1295,19 +1356,34 @@ def parse_tool_call_arguments( context: Optional context string (e.g., "Anthropic Messages API"). Returns: - Parsed arguments as a dictionary. Returns empty dict if arguments is None or empty. + 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. Raises: - ValueError: If the arguments string is not valid JSON. + ValueError: If the arguments string is not valid JSON and cannot be repaired. """ import json - if not arguments: + if not arguments or not arguments.strip(): return {} try: return json.loads(arguments) - except json.JSONDecodeError as e: + except json.JSONDecodeError as original_error: + repaired = _attempt_json_repair(arguments) + if repaired is not None: + verbose_logger.warning( + "Repaired truncated tool call arguments for tool '%s' (%s). " + "Original (%d chars): %.200s%s", + tool_name or "", + context or "unknown context", + len(arguments), + arguments, + "..." if len(arguments) > 200 else "", + ) + return repaired + error_parts = ["Failed to parse tool call arguments"] if tool_name: @@ -1316,10 +1392,11 @@ def parse_tool_call_arguments( error_parts.append(f"({context})") error_message = ( - " ".join(error_parts) + f". Error: {str(e)}. Arguments: {arguments}" + " ".join(error_parts) + + f". Error: {str(original_error)}. Arguments: {arguments}" ) - raise ValueError(error_message) from e + raise ValueError(error_message) from original_error def split_concatenated_json_objects(raw: str) -> List[Dict[str, Any]]: diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 796223ff8e..a694cec7d6 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -1035,9 +1035,13 @@ def convert_to_anthropic_tool_invoke_xml(tool_calls: list) -> str: parsed_args = parse_tool_call_arguments( tool_arguments, tool_name=tool_name, context="Anthropic XML tool invoke" ) - parameters = "".join( - f"<{param}>{val}\n" for param, val in parsed_args.items() - ) + if isinstance(parsed_args, dict): + parameters = "".join( + f"<{param}>{val}\n" + for param, val in parsed_args.items() + ) + else: + parameters = f"{parsed_args}\n" invokes += ( "\n" f"{tool_name}\n" diff --git a/tests/llm_translation/test_prompt_factory.py b/tests/llm_translation/test_prompt_factory.py index 9f902f2bd8..0e7ed28e1a 100644 --- a/tests/llm_translation/test_prompt_factory.py +++ b/tests/llm_translation/test_prompt_factory.py @@ -1457,3 +1457,147 @@ def test_convert_to_anthropic_tool_invoke_malformed_json(): error_msg = str(exc_info.value) assert "bad_tool" in error_msg assert '{"truncated' in error_msg + + +# ============ _attempt_json_repair Tests ============ +# Tests for the JSON repair utility that fixes truncated tool call arguments + + +def test_attempt_json_repair_missing_closing_brace(): + """Repair JSON truncated with a missing closing brace (issue #22312).""" + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + _attempt_json_repair, + ) + + truncated = '{"command": ["bash","-lc","find /x/repos -name \'messages.py\' -type f"]' + result = _attempt_json_repair(truncated) + assert result is not None + assert result["command"] == ["bash", "-lc", "find /x/repos -name 'messages.py' -type f"] + + +def test_attempt_json_repair_missing_bracket_and_brace(): + """Repair JSON truncated with both missing ] and }.""" + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + _attempt_json_repair, + ) + + truncated = '{"items": [1, 2, 3' + result = _attempt_json_repair(truncated) + assert result is not None + assert result["items"] == [1, 2, 3] + + +def test_attempt_json_repair_trailing_comma(): + """Repair JSON with a trailing comma before missing close.""" + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + _attempt_json_repair, + ) + + truncated = '{"a": 1, "b": 2,' + result = _attempt_json_repair(truncated) + assert result is not None + assert result == {"a": 1, "b": 2} + + +def test_attempt_json_repair_returns_none_for_unterminated_string(): + """Cannot repair an unterminated string — returns None.""" + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + _attempt_json_repair, + ) + + assert _attempt_json_repair('{"key": "incomplete value') is None + + +def test_attempt_json_repair_returns_none_for_valid_json(): + """Valid JSON has no unmatched brackets — returns None (no repair needed).""" + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + _attempt_json_repair, + ) + + assert _attempt_json_repair('{"key": "value"}') is None + + +def test_attempt_json_repair_returns_none_for_empty(): + """Empty / whitespace input returns None.""" + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + _attempt_json_repair, + ) + + assert _attempt_json_repair("") is None + assert _attempt_json_repair(" ") is None + + +def test_attempt_json_repair_interleaved_nesting(): + """Repair JSON with interleaved {} and [] nesting.""" + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + _attempt_json_repair, + ) + + # {"a": [{"b": 2 needs }]} not ]}} + truncated = '{"a": [{"b": 2' + result = _attempt_json_repair(truncated) + assert result is not None + assert result == {"a": [{"b": 2}]} + + +def test_attempt_json_repair_deeply_nested(): + """Repair deeply nested truncated JSON.""" + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + _attempt_json_repair, + ) + + truncated = '{"x": {"y": [1, {"z": [2, 3' + result = _attempt_json_repair(truncated) + assert result is not None + assert result == {"x": {"y": [1, {"z": [2, 3]}]}} + + +def test_parse_tool_call_arguments_whitespace_only(): + """Whitespace-only input returns empty dict.""" + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + parse_tool_call_arguments, + ) + + assert parse_tool_call_arguments(" ") == {} + assert parse_tool_call_arguments("\n") == {} + + +def test_parse_tool_call_arguments_non_object_json(): + """Non-object JSON (list, string, number) is returned as-is (no wrapping).""" + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + parse_tool_call_arguments, + ) + + result = parse_tool_call_arguments('[1, 2, 3]') + assert result == [1, 2, 3] + + +def test_parse_tool_call_arguments_repairs_truncated_json(): + """parse_tool_call_arguments should repair truncated JSON instead of raising.""" + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + parse_tool_call_arguments, + ) + + truncated = '{"command": ["bash","-lc","find /x -type f"]' + result = parse_tool_call_arguments( + truncated, tool_name="shell", context="Anthropic tool invoke" + ) + assert result == {"command": ["bash", "-lc", "find /x -type f"]} + + +def test_parse_tool_call_arguments_still_raises_for_unrepairable(): + """parse_tool_call_arguments raises ValueError when repair also fails.""" + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + parse_tool_call_arguments, + ) + + with pytest.raises(ValueError) as exc_info: + parse_tool_call_arguments( + '{"key": "unterminated', + tool_name="test_tool", + context="test context", + ) + + error_msg = str(exc_info.value) + assert "test_tool" in error_msg + assert "test context" in error_msg