From 03f45fee70d9deed15ce68f707bc69591c489ce8 Mon Sep 17 00:00:00 2001 From: Sumedh Aerram Date: Mon, 2 Nov 2026 13:51:24 -0800 Subject: [PATCH] fix(models): send structured output for non-BaseModel schemas in LiteLlm _to_litellm_response_format only handled dicts and BaseModels, so list, int, Literal and Annotated output schemas got no response_format, and Annotated aliases raised a TypeError. Convert them with schema_to_json_schema. For strict json_schema targets, nest non-object roots under a "response" property, keeping $defs at the root, and unwrap the value from the reply, including the final streaming aggregate. Unwrap only when ADK wrapped. Map schemas and unconvertible types keep the previous warning. Fall back to response_json_schema when response_schema is unset. The Gemini path is unchanged. Fixes #7946 --- src/google/adk/models/lite_llm.py | 150 +++++++++-- tests/unittests/models/test_litellm.py | 334 +++++++++++++++++++++++++ 2 files changed, 457 insertions(+), 27 deletions(-) diff --git a/src/google/adk/models/lite_llm.py b/src/google/adk/models/lite_llm.py index e06eb32248b..8633e2f0ec1 100644 --- a/src/google/adk/models/lite_llm.py +++ b/src/google/adk/models/lite_llm.py @@ -31,6 +31,7 @@ from typing import cast from typing import Dict from typing import Generator +from typing import get_origin from typing import Iterable from typing import List from typing import Literal @@ -55,6 +56,7 @@ from pydantic import BaseModel from pydantic import Field from pydantic import PrivateAttr +from pydantic import PydanticUserError from typing_extensions import NotRequired from typing_extensions import override from typing_extensions import Required @@ -63,6 +65,7 @@ from ..utils import streaming_utils from ..utils._google_client_headers import merge_tracking_headers from ..utils._schema_utils import lowercase_schema_types +from ..utils._schema_utils import schema_to_json_schema from ..utils.model_name_utils import is_gemini_model from ._capabilities import LlmCapabilities from .base_llm import BaseLlm @@ -3814,21 +3817,13 @@ def _enforce_strict_openai_schema(schema: dict[str, Any]) -> None: _enforce_strict_openai_schema(schema["items"]) -def _to_litellm_response_format( - response_schema: types.SchemaUnion, - model: str, -) -> dict[str, Any] | None: - """Converts ADK response schema objects into LiteLLM-compatible payloads. +_OUTPUT_SCHEMA_VALUE_KEY = "response" - Args: - response_schema: The response schema to convert. - model: The model string to determine the appropriate format. Gemini models - use 'response_schema' key, while OpenAI-compatible models use - 'json_schema' key. - Returns: - A dictionary with the appropriate response format for LiteLLM. - """ +def _get_litellm_output_schema( + response_schema: types.SchemaUnion, +) -> tuple[dict[str, Any], str]: + """Returns the JSON schema and its name before provider-specific wrapping.""" schema_name = "response" if isinstance(response_schema, dict): @@ -3837,7 +3832,7 @@ def _to_litellm_response_format( isinstance(schema_type, str) and schema_type.lower() in _LITELLM_STRUCTURED_TYPES ): - return response_schema + return response_schema, schema_name schema_dict = copy.deepcopy(response_schema) if "title" in schema_dict: schema_name = str(schema_dict["title"]) @@ -3859,30 +3854,90 @@ def _to_litellm_response_format( else: schema_dict = response_schema.__class__.model_json_schema() schema_name = response_schema.__class__.__name__ + elif get_origin(response_schema) is not None or isinstance( + response_schema, type + ): + schema_dict = schema_to_json_schema(response_schema) elif hasattr(response_schema, "model_dump"): schema_dict = copy.deepcopy( response_schema.model_dump(exclude_none=True, mode="json") ) schema_name = response_schema.__class__.__name__ else: + schema_dict = schema_to_json_schema(response_schema) + + return schema_dict, schema_name + + +def _to_litellm_response_format( + response_schema: types.SchemaUnion, + model: str, +) -> dict[str, Any] | None: + """Converts ADK response schema objects into LiteLLM-compatible payloads. + + Args: + response_schema: The response schema to convert. + model: The model string to determine the appropriate format. Gemini models + use 'response_schema' key, while OpenAI-compatible models use + 'json_schema' key. + + Returns: + A dictionary with the appropriate response format for LiteLLM. + """ + return _prepare_litellm_response_format(response_schema, model)[0] + + +def _prepare_litellm_response_format( + response_schema: types.SchemaUnion, + model: str, +) -> tuple[dict[str, Any] | None, bool]: + """Builds the provider format and records whether its value was nested.""" + try: + schema_dict, schema_name = _get_litellm_output_schema(response_schema) + except (PydanticUserError, TypeError): logger.warning( "Unsupported response_schema type %s for LiteLLM structured outputs.", type(response_schema), ) - return None + return None, False + if ( + isinstance(schema_dict.get("type"), str) + and schema_dict["type"].lower() in _LITELLM_STRUCTURED_TYPES + ): + return schema_dict, False # Gemini models use a special response format with 'response_schema' key if _is_litellm_gemini_model(model): return { "type": "json_object", "response_schema": schema_dict, - } + }, False - # OpenAI-compatible format (default) per LiteLLM docs: - # https://docs.litellm.ai/docs/completion/json_mode - if isinstance(schema_dict, dict): - lowercase_schema_types(schema_dict) - _enforce_strict_openai_schema(schema_dict) + # OpenAI strict output requires an object root. Keep references relative to + # the document root when nesting arrays, scalars, or root references. + lowercase_schema_types(schema_dict) + if ( + schema_dict.get("type") == "object" + and "properties" not in schema_dict + and isinstance(schema_dict.get("additionalProperties"), dict) + ): + logger.warning( + "Unsupported response_schema type %s for LiteLLM structured outputs.", + type(response_schema), + ) + return None, False + nests_output = schema_dict.get("type") != "object" + if nests_output: + defs = schema_dict.pop("$defs", None) + schema_dict = { + "type": "object", + "properties": {_OUTPUT_SCHEMA_VALUE_KEY: schema_dict}, + "required": [_OUTPUT_SCHEMA_VALUE_KEY], + "additionalProperties": False, + } + if defs: + schema_dict["$defs"] = defs + _enforce_strict_openai_schema(schema_dict) return { "type": "json_schema", @@ -3891,7 +3946,28 @@ def _to_litellm_response_format( "strict": True, "schema": schema_dict, }, - } + }, nests_output + + +def _unwrap_litellm_output(response: LlmResponse) -> None: + """Unwraps a complete structured reply, preserving malformed/error replies.""" + if not response.content or not response.content.parts: + return + parts = [ + part for part in response.content.parts if part.text and not part.thought + ] + if not parts: + return + try: + value = json.loads("".join(part.text or "" for part in parts)) + except json.JSONDecodeError: + return + if isinstance(value, dict) and _OUTPUT_SCHEMA_VALUE_KEY in value: + parts[0].text = json.dumps( + value[_OUTPUT_SCHEMA_VALUE_KEY], ensure_ascii=False + ) + for part in parts[1:]: + response.content.parts.remove(part) _ANTHROPIC_MIN_THINKING_BUDGET_TOKENS = 1024 @@ -4030,11 +4106,14 @@ async def _get_completion_inputs( # 3. Handle response format response_format: dict[str, Any] | None = None - if llm_request.config and llm_request.config.response_schema: - response_format = _to_litellm_response_format( - llm_request.config.response_schema, - model=model, - ) + if llm_request.config: + response_schema = llm_request.config.response_schema + if response_schema is None: + response_schema = llm_request.config.response_json_schema + if response_schema is not None: + response_format = _to_litellm_response_format( + response_schema, model=model + ) # 4. Extract generation parameters generation_params: dict[str, Any] | None = None @@ -4708,6 +4787,15 @@ async def generate_content_async( ), ) ) + nests_output = False + if response_format and llm_request.config: + response_schema = llm_request.config.response_schema + if response_schema is None: + response_schema = llm_request.config.response_json_schema + if response_schema is not None: + _, nests_output = _prepare_litellm_response_format( + response_schema, effective_model + ) normalized_messages = _normalize_ollama_chat_messages( messages, model=effective_model, @@ -4727,6 +4815,9 @@ async def generate_content_async( "response_format": response_format, } completion_args.update(self._additional_args) + nests_output = ( + nests_output and completion_args["response_format"] == response_format + ) # A caller who named their own injection points at construction has said # more about their provider than the app-level config can, so leave those @@ -5164,6 +5255,7 @@ def _has_held_back_tool_calls() -> bool: elif isinstance(chunk, TextChunk): if chunk.text: text_parts.append(chunk.text) + # Partial text is raw wrapped JSON; only the final aggregate is unwrapped. yield _message_to_generate_content_response( ChatCompletionAssistantMessage( role="assistant", @@ -5311,6 +5403,8 @@ def _has_held_back_tool_calls() -> bool: usage_metadata = None if grounding_metadata: aggregated_llm_response.grounding_metadata = grounding_metadata + if nests_output: + _unwrap_litellm_output(aggregated_llm_response) yield aggregated_llm_response if aggregated_llm_response_with_tool_call: @@ -5341,6 +5435,8 @@ def _has_held_back_tool_calls() -> bool: for p in llm_response.content.parts if not p.thought or (has_tool_calls and p.thought_signature) ] + if nests_output: + _unwrap_litellm_output(llm_response) yield llm_response @classmethod diff --git a/tests/unittests/models/test_litellm.py b/tests/unittests/models/test_litellm.py index 34c17bc9c0b..2a5febcd5a8 100644 --- a/tests/unittests/models/test_litellm.py +++ b/tests/unittests/models/test_litellm.py @@ -19,6 +19,10 @@ import os import sys import tempfile +from typing import Annotated +from typing import Callable +from typing import ForwardRef +from typing import Literal import unittest from unittest.mock import ANY from unittest.mock import AsyncMock @@ -62,6 +66,7 @@ from google.adk.models.lite_llm import _parse_deepseek_tool_calls_from_text from google.adk.models.lite_llm import _parse_tool_call_arguments from google.adk.models.lite_llm import _parse_tool_calls_from_text +from google.adk.models.lite_llm import _prepare_litellm_response_format from google.adk.models.lite_llm import _redact_file_uri_for_log from google.adk.models.lite_llm import _redirect_litellm_loggers_to_stdout from google.adk.models.lite_llm import _safe_json_serialize @@ -14023,3 +14028,332 @@ def test_strict_json_does_not_log_warning(self, caplog): with caplog.at_level(logging.WARNING, logger="google_adk"): _parse_tool_call_arguments('{"a": 1}', function_name="test_fn") assert "Repaired" not in caplog.text + + +@pytest.mark.parametrize("model", ["openai/gpt-4o", "gemini/gemini-2.5-flash"]) +@pytest.mark.parametrize( + "schema, expected", + [ + ( + list[_StructuredOutput], + {"type": "array", "items": {"$ref": "#/$defs/_StructuredOutput"}}, + ), + (int, {"type": "integer"}), + ( + dict[str, int], + {"type": "object", "additionalProperties": {"type": "integer"}}, + ), + (Literal["yes", "no"], {"type": "string", "enum": ["yes", "no"]}), + ( + {"type": "array", "items": {"type": "integer"}}, + {"type": "array", "items": {"type": "integer"}}, + ), + ], +) +def test_non_basemodel_output_schema_response_format( + schema, expected, model, caplog +): + """All supported output types reach LiteLLM, with object roots for strict mode.""" + formatted = _to_litellm_response_format(schema, model=model) + + if schema == dict[str, int] and not model.startswith("gemini/"): + assert formatted is None + assert "Unsupported response_schema type" in caplog.text + return + if model.startswith("gemini/"): + assert formatted["type"] == "json_object" + actual = dict(formatted["response_schema"]) + else: + assert formatted["type"] == "json_schema" + assert formatted["json_schema"]["strict"] is True + actual = dict(formatted["json_schema"]["schema"]) + defs = actual.pop("$defs", None) + if schema == list[_StructuredOutput]: + assert "_StructuredOutput" in defs + if not model.startswith("gemini/") and expected["type"] != "object": + assert actual == { + "type": "object", + "properties": {"response": expected}, + "required": ["response"], + "additionalProperties": False, + } + else: + assert actual == expected + + +@pytest.mark.parametrize("model", ["openai/gpt-4o", "gemini/gemini-2.5-flash"]) +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize( + "schema, wire_text, expected", + [ + (int, '{"response": 7}', "7"), + (int, '{"response": 0}', "0"), + ( + list[_StructuredOutput], + '{"response": [{"value": 2}]}', + '[{"value": 2}]', + ), + (Literal["yes", "no"], '{"response": "yes"}', '"yes"'), + ( + {"type": "array", "items": {"type": "integer"}}, + '{"response": []}', + "[]", + ), + (_StructuredOutput, '{"value": 7}', '{"value": 7}'), + ( + {"type": "object", "properties": {"response": {"type": "integer"}}}, + '{"response": 7}', + '{"response": 7}', + ), + (dict[str, int], '{"response": 7}', '{"response": 7}'), + (int, '{"response":', '{"response":'), + (int, '{"other": 7}', '{"other": 7}'), + ], +) +async def test_non_basemodel_output_schema_reply( + model, + schema, + wire_text, + expected, + stream, + mock_acompletion, + mock_completion, + lite_llm_instance, +): + """Only wrapped complete replies are unwrapped; partials stay on the wire format.""" + request = LlmRequest(model=model) + request.set_output_schema(schema) + fragments = [wire_text[:5], wire_text[5:]] + if stream: + mock_completion.return_value = iter(_text_stream_chunks(fragments)) + else: + mock_acompletion.return_value = ModelResponse( + choices=[ + Choices( + message=ChatCompletionAssistantMessage( + role="assistant", content=wire_text + ) + ) + ] + ) + + responses = [ + r + async for r in lite_llm_instance.generate_content_async( + request, stream=stream + ) + ] + + if model.startswith("gemini/"): + expected = wire_text + assert responses[-1].content.parts[0].text == expected + if stream: + assert [ + r.content.parts[0].text for r in responses if r.partial + ] == fragments + + +@pytest.mark.parametrize("model", ["openai/gpt-4o", "gemini/gemini-2.5-flash"]) +def test_annotated_output_schema_response_format(model): + """Annotated model aliases convert without invoking model_dump.""" + schema = Annotated[_StructuredOutput, Field(description="d")] + + formatted = _to_litellm_response_format(schema, model) + + actual = ( + formatted["response_schema"] + if model.startswith("gemini/") + else formatted["json_schema"]["schema"] + ) + assert "value" in actual["$defs"]["_StructuredOutput"]["properties"] + if model.startswith("gemini/"): + assert actual["description"] == "d" + assert actual["$ref"] == "#/$defs/_StructuredOutput" + else: + assert actual["type"] == "object" + assert ( + actual["properties"]["response"]["$ref"] == "#/$defs/_StructuredOutput" + ) + assert actual["required"] == ["response"] + + +@pytest.mark.parametrize("model", ["openai/gpt-4o", "gemini/gemini-2.5-flash"]) +def test_map_json_schema_response_format(model, caplog): + """Strict targets skip map roots while Gemini retains the schema.""" + schema = {"type": "object", "additionalProperties": {"type": "integer"}} + + formatted = _to_litellm_response_format(schema, model) + + if model.startswith("gemini/"): + assert formatted == {"type": "json_object", "response_schema": schema} + else: + assert formatted is None + assert "Unsupported response_schema type" in caplog.text + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("model", ["openai/gpt-4o", "gemini/gemini-2.5-flash"]) +@pytest.mark.parametrize("prefer_response_schema", [False, True]) +async def test_response_json_schema_fallback( + stream, + model, + prefer_response_schema, + mock_acompletion, + mock_completion, + lite_llm_instance, +): + """JSON-schema-only requests use the provider format and unwrap when nested.""" + schema = {"type": "string", "enum": ["yes", "no"]} + config = types.GenerateContentConfig(response_json_schema=schema) + if prefer_response_schema: + config.response_schema = _StructuredOutput + request = LlmRequest(model=model, config=config) + wire_text = ( + '{"value": 7}' + if prefer_response_schema + else '"yes"' + if model.startswith("gemini/") + else '{"response": "yes"}' + ) + if stream: + mock_completion.return_value = iter(_text_stream_chunks([wire_text])) + else: + mock_acompletion.return_value = ModelResponse( + choices=[ + Choices( + message=ChatCompletionAssistantMessage( + role="assistant", content=wire_text + ) + ) + ] + ) + + responses = [ + r + async for r in lite_llm_instance.generate_content_async( + request, stream=stream + ) + ] + + call = mock_completion.call_args if stream else mock_acompletion.call_args + formatted = call.kwargs["response_format"] + expected_schema = _StructuredOutput if prefer_response_schema else schema + assert formatted == _to_litellm_response_format(expected_schema, model) + assert responses[-1].content.parts[0].text == ( + '{"value": 7}' if prefer_response_schema else '"yes"' + ) + + +@pytest.mark.parametrize("model", ["openai/gpt-4o", "gemini/gemini-2.5-flash"]) +@pytest.mark.parametrize( + "schema", + [ + object(), + Callable[[int], str], + ForwardRef("MissingSchema"), + Annotated[int, Field(discriminator="kind")], + ], +) +def test_unsupported_output_schema_warns_and_skips_format( + schema, model, caplog +): + """Unsupported schema conversion errors preserve the unconstrained fallback.""" + assert _to_litellm_response_format(schema, model) is None + assert _prepare_litellm_response_format(schema, model) == (None, False) + assert "Unsupported response_schema type" in caplog.text + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("guard", ["basemodel", "override"]) +async def test_structured_output_unwrap_guards( + stream, guard, mock_client, mock_acompletion, mock_completion +): + """Only replies produced under the adapter's own wrapper are unwrapped.""" + + class ResponseModel(BaseModel): + response: int + + wire_text = '{"response": 7}' + request = LlmRequest() + request.set_output_schema(ResponseModel if guard == "basemodel" else int) + caller_format = {"type": "json_object"} + model = LiteLlm( + model="openai/gpt-4o", + llm_client=mock_client, + **({"response_format": caller_format} if guard == "override" else {}), + ) + if stream: + mock_completion.return_value = iter(_text_stream_chunks([wire_text])) + else: + mock_acompletion.return_value = ModelResponse( + choices=[ + Choices( + message=ChatCompletionAssistantMessage( + role="assistant", content=wire_text + ) + ) + ] + ) + + responses = [ + r async for r in model.generate_content_async(request, stream=stream) + ] + + assert responses[-1].content.parts[0].text == wire_text + if guard == "override": + call = mock_completion.call_args if stream else mock_acompletion.call_args + assert call.kwargs["response_format"] == caller_format + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize( + "wire_text, expected", + [ + ('{"response": "héllo"}', '"héllo"'), + ('{"response": "héllo', '{"response": "héllo'), + ], +) +async def test_wrapped_reply_preserves_thoughts_and_unicode( + stream, + wire_text, + expected, + mock_acompletion, + mock_completion, + lite_llm_instance, +): + """Unwrapping preserves thought parts, Unicode, and malformed JSON text.""" + request = LlmRequest() + request.set_output_schema(str) + if stream: + mock_completion.return_value = iter( + _reasoning_stream_chunks([ + {"reasoning_content": "thinking..."}, + {"content": wire_text}, + ]) + ) + else: + mock_acompletion.return_value = ModelResponse( + choices=[ + Choices( + message=ChatCompletionAssistantMessage( + role="assistant", + content=wire_text, + reasoning_content="thinking...", + ) + ) + ] + ) + + responses = [ + r + async for r in lite_llm_instance.generate_content_async( + request, stream=stream + ) + ] + + parts = responses[-1].content.parts + assert len(parts) == 2 + assert parts[0].thought is True + assert parts[0].text == "thinking..." + assert not parts[1].thought + assert parts[1].text == expected