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 <gweale@google.com>
PiperOrigin-RevId: 829037987

Author: GWeale

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

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,
+)

src/google/adk/models/lite_llm.py

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

tests/unittests/models/test_litellm.py

@@ -21,9 +21,11 @@
 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 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 @@
     ),
 ]
 
+
+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=[