Merge 56e5bc0017 into c8635ecc67

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
@@ -1324,6 +1325,23 @@ 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
+    """
+    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
+
+
 def convert_to_gemini_tool_call_result(
     message: ChatCompletionToolMessage | ChatCompletionFunctionMessage,
     last_message_with_tool_calls: dict | None,
@@ -1469,7 +1487,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}

tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py

@@ -905,6 +905,126 @@ 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_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

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 (