From b08cf974ce00d986389dfad2cd4149ad2f3ca784 Mon Sep 17 00:00:00 2001 From: Sushit-prog Date: Wed, 26 Aug 2026 06:39:58 +0530 Subject: [PATCH 1/4] fix(vertex_ai): prevent Gemini from misinterpreting $ref in tool results --- .../prompt_templates/factory.py | 21 +++- ...llm_core_utils_prompt_templates_factory.py | 103 ++++++++++++++++++ .../test_vertex_ai_gemini_transformation.py | 55 ++++++++++ 3 files changed, 178 insertions(+), 1 deletion(-) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 86cfbf70255..88e0205b838 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -1324,6 +1324,25 @@ def convert_to_gemini_tool_call_invoke( raise Exception(f"Unable to convert openai tool calls={message} to gemini tool calls. Received error={e}") +def _contains_json_schema_ref(payload: object) -> bool: + """ + Gemini interprets {"$ref": ...} objects anywhere inside function_response.response + as references to named parts in function_response.parts and rejects the request when + no matching part exists. Tool results carrying JSON Schema data must therefore not be + sent structurally. https://github.com/BerriAI/litellm/issues/38223 + """ + stack: list[Any] = [payload] + while stack: + node = stack.pop() + if isinstance(node, dict): + if "$ref" in node: + return True + stack.extend(node.values()) + elif isinstance(node, list): + stack.extend(node) + return False + + def convert_to_gemini_tool_call_result( message: ChatCompletionToolMessage | ChatCompletionFunctionMessage, last_message_with_tool_calls: dict | None, @@ -1469,7 +1488,7 @@ def convert_to_gemini_tool_call_result( if content_str.strip().startswith("{") or content_str.strip().startswith("["): # Try to parse as JSON (for Computer Use structured responses) parsed: Final = json.loads(content_str) - if isinstance(parsed, dict): + if isinstance(parsed, dict) and not _contains_json_schema_ref(parsed): response_data = parsed # Use the parsed JSON directly else: response_data = {"content": content_str} 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 6265779b90d..fc75a292e5f 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 @@ -905,6 +905,109 @@ def test_convert_gemini_tool_call_result_with_data_url_extra_params(): ), f"expected clean 'image/png', got '{inline_parts[0]['mime_type']}'" +def _gemini_tool_result_fixture( + content: str, +) -> tuple[ChatCompletionToolMessage, dict]: + message = ChatCompletionToolMessage( + role="tool", + tool_call_id="call_schema", + content=content, + ) + last_message_with_tool_calls = { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_schema", + "type": "function", + "index": 0, + "function": {"name": "inspect_schema", "arguments": "{}"}, + } + ], + } + return message, last_message_with_tool_calls + + +def test_convert_gemini_tool_result_wraps_output_containing_json_schema_ref(): + """ + Gemini treats {"$ref": ...} objects inside function_response.response as + references to named parts in function_response.parts, so schema-bearing + tool output must be delivered opaquely instead of structurally. + Fixes: https://github.com/BerriAI/litellm/issues/38223 + """ + schema_json = ( + '{"$defs":{"humanScalar":{"type":"string"}},' + '"type":"object","properties":{"value":{"$ref":"#/$defs/humanScalar"}}}' + ) + message, last_message_with_tool_calls = _gemini_tool_result_fixture(schema_json) + + result = convert_to_gemini_tool_call_result( + message=message, + last_message_with_tool_calls=last_message_with_tool_calls, + ) + + function_response = result["function_response"] + assert function_response["name"] == "inspect_schema" + assert function_response["response"] == {"content": schema_json} + + +def test_convert_gemini_tool_result_wraps_deeply_nested_json_schema_ref(): + """ + A $ref key at any depth must trigger opaque delivery. + Fixes: https://github.com/BerriAI/litellm/issues/38223 + """ + nested_json = '{"results":[{"schemas":[{"node":{"$ref":"#/$defs/deep"}}]}]}' + message, last_message_with_tool_calls = _gemini_tool_result_fixture(nested_json) + + result = convert_to_gemini_tool_call_result( + message=message, + last_message_with_tool_calls=last_message_with_tool_calls, + ) + + function_response = result["function_response"] + assert function_response["response"] == {"content": nested_json} + + +@pytest.mark.parametrize( + "schema_json", + [ + '{"$defs":{"Foo":{"type":"string"}},"type":"object","properties":{"bar":{"type":"string"}}}', + '{"definitions":{"Foo":{"type":"string"}},"type":"object","properties":{"bar":{"type":"string"}}}', + '{"status":"ok","count":2,"tags":["a","b"]}', + ], +) +def test_convert_gemini_tool_result_preserves_structured_passthrough_without_refs( + schema_json: str, +): + """ + Tool output carrying $defs / definitions / plain data but no $ref keys + keeps the existing structured function-response behavior unchanged. + Fixes: https://github.com/BerriAI/litellm/issues/38223 + """ + message, last_message_with_tool_calls = _gemini_tool_result_fixture(schema_json) + + result = convert_to_gemini_tool_call_result( + message=message, + last_message_with_tool_calls=last_message_with_tool_calls, + ) + + function_response = result["function_response"] + assert function_response["name"] == "inspect_schema" + assert function_response["response"] == json.loads(schema_json) + + +def test_convert_gemini_tool_result_malformed_json_still_wrapped_as_content(): + """Unparseable tool output keeps taking the existing content-wrapping path.""" + message, last_message_with_tool_calls = _gemini_tool_result_fixture("{not valid json") + + result = convert_to_gemini_tool_call_result( + message=message, + last_message_with_tool_calls=last_message_with_tool_calls, + ) + + assert result["function_response"]["response"] == {"content": "{not valid json"} + + def test_bedrock_tools_unpack_defs(): """ Test that the unpack_defs method handles nested $ref inside anyOf items correctly diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py index 8c1de12e7d9..eba8210ee09 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py @@ -412,6 +412,61 @@ def test_empty_content_handling(): assert contents[0]["parts"][0]["text"] == "" +def test_transform_request_body_wraps_tool_result_with_json_schema_ref(): + """ + Regression: Gemini treats {"$ref": ...} objects inside functionResponse.response + as references to named parts in function_response.parts and rejects the request. + The shared request-body transformation must deliver schema-bearing tool output + opaquely instead of structurally. + Fixes: https://github.com/BerriAI/litellm/issues/38223 + """ + schema_json = ( + '{"$defs":{"humanScalar":{"type":"string"}},' + '"type":"object","properties":{"value":{"$ref":"#/$defs/humanScalar"}}}' + ) + messages = [ + {"role": "user", "content": "Inspect this schema."}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_schema", + "type": "function", + "function": {"name": "inspect_schema", "arguments": "{}"}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_schema", "content": schema_json}, + {"role": "user", "content": "Acknowledge the result."}, + ] + + result = _transform_request_body( + messages=messages, + model="gemini-2.0-flash", + optional_params={}, + custom_llm_provider="vertex_ai", + litellm_params={}, + cached_content=None, + ) + + parts = [part for content in result["contents"] for part in content.get("parts", [])] + function_responses = [ + part[key] for part in parts for key in ("function_response", "functionResponse") if key in part + ] + assert len(function_responses) == 1 + assert function_responses[0]["response"] == {"content": schema_json} + + stack: list[object] = list(parts) + while stack: + node = stack.pop() + if isinstance(node, dict): + assert "$ref" not in node + stack.extend(node.values()) + elif isinstance(node, list): + stack.extend(node) + + def test_thought_signature_extraction_from_response(): """Test that thought signatures are extracted from Gemini response parts and stored in provider_specific_fields""" from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( From 4c9b6ae16a0ef44ad2633928a4ef85996b58febc Mon Sep 17 00:00:00 2001 From: Sushit-prog Date: Wed, 26 Aug 2026 07:07:37 +0530 Subject: [PATCH 2/4] fix(vertex_ai): use object instead of Any in schema-ref traversal --- litellm/litellm_core_utils/prompt_templates/factory.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 88e0205b838..2a8de054043 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -1331,7 +1331,7 @@ def _contains_json_schema_ref(payload: object) -> bool: no matching part exists. Tool results carrying JSON Schema data must therefore not be sent structurally. https://github.com/BerriAI/litellm/issues/38223 """ - stack: list[Any] = [payload] + stack: list[object] = [payload] while stack: node = stack.pop() if isinstance(node, dict): From 3c427735e1971725c38a8658427b18e185e9b988 Mon Sep 17 00:00:00 2001 From: Sushit-prog Date: Wed, 26 Aug 2026 07:54:24 +0530 Subject: [PATCH 3/4] fix(vertex_ai): use immutable tuple traversal to satisfy type-discipline gate --- .../prompt_templates/factory.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 2a8de054043..7b3822bafdc 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -7,6 +7,7 @@ import re import xml.etree.ElementTree as ET from collections.abc import Iterator, Mapping, Sequence from enum import Enum +from itertools import chain from typing import Any, Final, TypedDict, cast, overload from jinja2.sandbox import ImmutableSandboxedEnvironment @@ -1331,15 +1332,13 @@ def _contains_json_schema_ref(payload: object) -> bool: no matching part exists. Tool results carrying JSON Schema data must therefore not be sent structurally. https://github.com/BerriAI/litellm/issues/38223 """ - stack: list[object] = [payload] - while stack: - node = stack.pop() - if isinstance(node, dict): - if "$ref" in node: - return True - stack.extend(node.values()) - elif isinstance(node, list): - stack.extend(node) + level: tuple[object, ...] = (payload,) + while level: + if any(isinstance(x, dict) and "$ref" in x for x in level): + return True + level = tuple( + chain.from_iterable(v.values() if isinstance(v, dict) else v for v in level if isinstance(v, (dict, list))) + ) return False From 56e5bc001790b4f1151a5d883af1722f478ba9b8 Mon Sep 17 00:00:00 2001 From: Sushit-prog Date: Wed, 26 Aug 2026 08:17:42 +0530 Subject: [PATCH 4/4] test(vertex_ai): add direct contract test for _contains_json_schema_ref --- ...tellm_core_utils_prompt_templates_factory.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) 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 fc75a292e5f..0e854923c46 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 @@ -1008,6 +1008,23 @@ def test_convert_gemini_tool_result_malformed_json_still_wrapped_as_content(): assert result["function_response"]["response"] == {"content": "{not valid json"} +def test_contains_json_schema_ref_contract(): + """ + Direct contract pin for the helper behind #38223: returns True exactly when + some dict at any depth carries a $ref key, across mixed dict/list/scalar + shapes, and False for the same shape without one. + """ + from litellm.litellm_core_utils.prompt_templates.factory import ( + _contains_json_schema_ref, + ) + + ref_bearing = {"a": [{"$ref": "#/$defs/x"}, "plain", 3], "b": None} + assert _contains_json_schema_ref(ref_bearing) is True + + clean_twin = {"a": [{"refs": "#/$defs/x"}, "plain", 3], "b": None} + assert _contains_json_schema_ref(clean_twin) is False + + def test_bedrock_tools_unpack_defs(): """ Test that the unpack_defs method handles nested $ref inside anyOf items correctly