From 7ea4aed35ba70ec5a38dc1b3b0a9808183c2bab1 Mon Sep 17 00:00:00 2001 From: George Weale Date: Thu, 6 Nov 2025 11:28:30 -0800 Subject: [PATCH] fix: Add support for structured output schemas in LiteLLM models Add `_to_litellm_response_format` to convert ADK's `response_schema` types (Pydantic models, JSON schema dicts) into the format needed by LiteLLM for JSON object/schema constraints Close #1967 Co-authored-by: George Weale PiperOrigin-RevId: 829037987 --- .../litellm_structured_output/__init__.py | 15 ++++ .../litellm_structured_output/agent.py | 47 ++++++++++ src/google/adk/models/lite_llm.py | 47 +++++++++- tests/unittests/models/test_litellm.py | 85 +++++++++++++++++++ 4 files changed, 191 insertions(+), 3 deletions(-) create mode 100644 contributing/samples/litellm_structured_output/__init__.py create mode 100644 contributing/samples/litellm_structured_output/agent.py diff --git a/contributing/samples/litellm_structured_output/__init__.py b/contributing/samples/litellm_structured_output/__init__.py new file mode 100644 index 00000000..c48963cd --- /dev/null +++ b/contributing/samples/litellm_structured_output/__init__.py @@ -0,0 +1,15 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from . import agent diff --git a/contributing/samples/litellm_structured_output/agent.py b/contributing/samples/litellm_structured_output/agent.py new file mode 100644 index 00000000..8fdd5f66 --- /dev/null +++ b/contributing/samples/litellm_structured_output/agent.py @@ -0,0 +1,47 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Sample agent showing LiteLLM structured output support.""" + +from __future__ import annotations + +from google.adk import Agent +from google.adk.models.lite_llm import LiteLlm +from pydantic import BaseModel +from pydantic import Field + + +class CitySummary(BaseModel): + """Simple structure used to verify LiteLLM JSON schema handling.""" + + city: str = Field(description="Name of the city being described.") + highlights: list[str] = Field( + description="Bullet points summarising the city's key highlights.", + ) + recommended_visit_length_days: int = Field( + description="Recommended number of days for a typical visit.", + ) + + +root_agent = Agent( + name="litellm_structured_output_agent", + model=LiteLlm(model="gemini-2.5-flash"), + description="Generates structured travel recommendations for a given city.", + instruction=""" +Produce a JSON object that follows the CitySummary schema. +Only include fields that appear in the schema and ensure highlights +contains short bullet points. +""".strip(), + output_schema=CitySummary, +) diff --git a/src/google/adk/models/lite_llm.py b/src/google/adk/models/lite_llm.py index 7ffcb317..7b41fff4 100644 --- a/src/google/adk/models/lite_llm.py +++ b/src/google/adk/models/lite_llm.py @@ -63,6 +63,7 @@ logger = logging.getLogger("google_adk." + __name__) _NEW_LINE = "\n" _EXCLUDED_PART_FIELD = {"inline_data": {"data"}} +_LITELLM_STRUCTURED_TYPES = {"json_object", "json_schema"} # Mapping of LiteLLM finish_reason strings to FinishReason enum values # Note: tool_calls/function_call map to STOP because: @@ -673,12 +674,50 @@ def _message_to_generate_content_response( ) +def _to_litellm_response_format( + response_schema: types.SchemaUnion, +) -> Optional[Dict[str, Any]]: + """Converts ADK response schema objects into LiteLLM-compatible payloads.""" + + if isinstance(response_schema, dict): + schema_type = response_schema.get("type") + if ( + isinstance(schema_type, str) + and schema_type.lower() in _LITELLM_STRUCTURED_TYPES + ): + return response_schema + schema_dict = dict(response_schema) + elif isinstance(response_schema, type) and issubclass( + response_schema, BaseModel + ): + schema_dict = response_schema.model_json_schema() + elif isinstance(response_schema, BaseModel): + if isinstance(response_schema, types.Schema): + # GenAI Schema instances already represent JSON schema definitions. + schema_dict = response_schema.model_dump(exclude_none=True, mode="json") + else: + schema_dict = response_schema.__class__.model_json_schema() + elif hasattr(response_schema, "model_dump"): + schema_dict = response_schema.model_dump(exclude_none=True, mode="json") + else: + logger.warning( + "Unsupported response_schema type %s for LiteLLM structured outputs.", + type(response_schema), + ) + return None + + return { + "type": "json_object", + "response_schema": schema_dict, + } + + def _get_completion_inputs( llm_request: LlmRequest, ) -> Tuple[ List[Message], Optional[List[Dict]], - Optional[types.SchemaUnion], + Optional[Dict[str, Any]], Optional[Dict], ]: """Converts an LlmRequest to litellm inputs and extracts generation params. @@ -721,9 +760,11 @@ def _get_completion_inputs( ] # 3. Handle response format - response_format: Optional[types.SchemaUnion] = None + response_format: Optional[Dict[str, Any]] = None if llm_request.config and llm_request.config.response_schema: - response_format = llm_request.config.response_schema + response_format = _to_litellm_response_format( + llm_request.config.response_schema + ) # 4. Extract generation parameters generation_params: Optional[Dict] = None diff --git a/tests/unittests/models/test_litellm.py b/tests/unittests/models/test_litellm.py index aefa2c67..ef15b078 100644 --- a/tests/unittests/models/test_litellm.py +++ b/tests/unittests/models/test_litellm.py @@ -21,9 +21,11 @@ import warnings from google.adk.models.lite_llm import _content_to_message_param from google.adk.models.lite_llm import _FINISH_REASON_MAPPING from google.adk.models.lite_llm import _function_declaration_to_tool_param +from google.adk.models.lite_llm import _get_completion_inputs from google.adk.models.lite_llm import _get_content from google.adk.models.lite_llm import _message_to_generate_content_response from google.adk.models.lite_llm import _model_response_to_chunk +from google.adk.models.lite_llm import _to_litellm_response_format from google.adk.models.lite_llm import _to_litellm_role from google.adk.models.lite_llm import FunctionChunk from google.adk.models.lite_llm import LiteLlm @@ -40,6 +42,8 @@ from litellm.types.utils import Choices from litellm.types.utils import Delta from litellm.types.utils import ModelResponse from litellm.types.utils import StreamingChoices +from pydantic import BaseModel +from pydantic import Field import pytest LLM_REQUEST_WITH_FUNCTION_DECLARATION = LlmRequest( @@ -179,6 +183,87 @@ STREAMING_MODEL_RESPONSE = [ ), ] + +class _StructuredOutput(BaseModel): + value: int = Field(description="Value to emit") + + +class _ModelDumpOnly: + """Test helper that mimics objects exposing only model_dump.""" + + def __init__(self): + self._schema = { + "type": "object", + "properties": {"foo": {"type": "string"}}, + } + + def model_dump(self, *, exclude_none=True, mode="json"): + # The method signature matches pydantic BaseModel.model_dump to simulate + # google.genai schema-like objects. + del exclude_none + del mode + return self._schema + + +def test_get_completion_inputs_formats_pydantic_schema_for_litellm(): + llm_request = LlmRequest( + config=types.GenerateContentConfig(response_schema=_StructuredOutput) + ) + + _, _, response_format, _ = _get_completion_inputs(llm_request) + + assert response_format == { + "type": "json_object", + "response_schema": _StructuredOutput.model_json_schema(), + } + + +def test_to_litellm_response_format_passes_preformatted_dict(): + response_format = { + "type": "json_object", + "response_schema": { + "type": "object", + "properties": {"foo": {"type": "string"}}, + }, + } + + assert _to_litellm_response_format(response_format) == response_format + + +def test_to_litellm_response_format_wraps_json_schema_dict(): + schema = { + "type": "object", + "properties": {"foo": {"type": "string"}}, + } + + formatted = _to_litellm_response_format(schema) + assert formatted["type"] == "json_object" + assert formatted["response_schema"] == schema + + +def test_to_litellm_response_format_handles_model_dump_object(): + schema_obj = _ModelDumpOnly() + + formatted = _to_litellm_response_format(schema_obj) + + assert formatted["type"] == "json_object" + assert formatted["response_schema"] == schema_obj.model_dump() + + +def test_to_litellm_response_format_handles_genai_schema_instance(): + schema_instance = types.Schema( + type=types.Type.OBJECT, + properties={"foo": types.Schema(type=types.Type.STRING)}, + required=["foo"], + ) + + formatted = _to_litellm_response_format(schema_instance) + assert formatted["type"] == "json_object" + assert formatted["response_schema"] == schema_instance.model_dump( + exclude_none=True, mode="json" + ) + + MULTIPLE_FUNCTION_CALLS_STREAM = [ ModelResponse( choices=[