From 5426b5a8f94a3878e584690a9500cd13e32a37cf Mon Sep 17 00:00:00 2001 From: Prathamesh Lohakare Date: Mon, 29 Dec 2025 22:26:04 +0530 Subject: [PATCH 1/3] fix: handle additionalProperties for openai models --- contributing/samples/gepa/experiment.py | 1 - contributing/samples/gepa/run_experiment.py | 1 - src/google/adk/models/lite_llm.py | 75 ++++++--- tests/unittests/models/test_litellm.py | 160 +++++++++++++++++--- 4 files changed, 192 insertions(+), 45 deletions(-) diff --git a/contributing/samples/gepa/experiment.py b/contributing/samples/gepa/experiment.py index 2f5d03a772..f68b349d9c 100644 --- a/contributing/samples/gepa/experiment.py +++ b/contributing/samples/gepa/experiment.py @@ -43,7 +43,6 @@ from tau_bench.types import EnvRunResult from tau_bench.types import RunConfig import tau_bench_agent as tau_bench_agent_lib - import utils diff --git a/contributing/samples/gepa/run_experiment.py b/contributing/samples/gepa/run_experiment.py index cfd850b3a3..1bc4ee58c8 100644 --- a/contributing/samples/gepa/run_experiment.py +++ b/contributing/samples/gepa/run_experiment.py @@ -25,7 +25,6 @@ from absl import flags import experiment from google.genai import types - import utils _OUTPUT_DIR = flags.DEFINE_string( diff --git a/src/google/adk/models/lite_llm.py b/src/google/adk/models/lite_llm.py index 140473982f..424d03fa52 100644 --- a/src/google/adk/models/lite_llm.py +++ b/src/google/adk/models/lite_llm.py @@ -1023,12 +1023,15 @@ def _model_response_to_chunk( if not func_name and not func_args: continue - yield FunctionChunk( - id=tool_call.id, - name=func_name, - args=func_args, - index=func_index, - ), finish_reason + yield ( + FunctionChunk( + id=tool_call.id, + name=func_name, + args=func_args, + index=func_index, + ), + finish_reason, + ) if finish_reason and not (message_content or tool_calls): yield None, finish_reason @@ -1040,12 +1043,17 @@ def _model_response_to_chunk( # finish_reason set. But this is not the case we are observing from litellm. # So we are sending it as a separate chunk to be set on the llm_response. if response.get("usage", None): - yield UsageMetadataChunk( - prompt_tokens=response["usage"].get("prompt_tokens", 0), - completion_tokens=response["usage"].get("completion_tokens", 0), - total_tokens=response["usage"].get("total_tokens", 0), - cached_prompt_tokens=_extract_cached_prompt_tokens(response["usage"]), - ), None + yield ( + UsageMetadataChunk( + prompt_tokens=response["usage"].get("prompt_tokens", 0), + completion_tokens=response["usage"].get("completion_tokens", 0), + total_tokens=response["usage"].get("total_tokens", 0), + cached_prompt_tokens=_extract_cached_prompt_tokens( + response["usage"] + ), + ), + None, + ) def _model_response_to_generate_content_response( @@ -1146,6 +1154,24 @@ def _message_to_generate_content_response( ) +def _enforce_closed_schema(schema: dict): + if not isinstance(schema, dict): + return + + if schema.get("type") == "object": + schema.setdefault("additionalProperties", False) + + for prop in schema.get("properties", {}).values(): + _enforce_closed_schema(prop) + + if "items" in schema: + _enforce_closed_schema(schema["items"]) + + if "$defs" in schema: + for def_schema in schema["$defs"].values(): + _enforce_closed_schema(def_schema) + + def _to_litellm_response_format( response_schema: types.SchemaUnion, model: str, @@ -1206,14 +1232,9 @@ def _to_litellm_response_format( # OpenAI-compatible format (default) per LiteLLM docs: # https://docs.litellm.ai/docs/completion/json_mode - if ( - isinstance(schema_dict, dict) - and schema_dict.get("type") == "object" - and "additionalProperties" not in schema_dict - ): + if isinstance(schema_dict, dict): # OpenAI structured outputs require explicit additionalProperties: false. - schema_dict = dict(schema_dict) - schema_dict["additionalProperties"] = False + _enforce_closed_schema(schema_dict) return { "type": "json_schema", @@ -1433,7 +1454,12 @@ def _warn_gemini_via_litellm(model_string: str) -> None: # Check if warning should be suppressed via environment variable if os.environ.get( "ADK_SUPPRESS_GEMINI_LITELLM_WARNINGS", "" - ).strip().lower() in ("1", "true", "yes", "on"): + ).strip().lower() in ( + "1", + "true", + "yes", + "on", + ): return warnings.warn( @@ -1541,9 +1567,12 @@ async def generate_content_async( logger.debug(_build_request_log(llm_request)) effective_model = llm_request.model or self.model - messages, tools, response_format, generation_params = ( - await _get_completion_inputs(llm_request, effective_model) - ) + ( + messages, + tools, + response_format, + generation_params, + ) = await _get_completion_inputs(llm_request, effective_model) normalized_messages = _normalize_ollama_chat_messages( messages, model=effective_model, diff --git a/tests/unittests/models/test_litellm.py b/tests/unittests/models/test_litellm.py index 4cf0329aa0..d0f6a768c3 100644 --- a/tests/unittests/models/test_litellm.py +++ b/tests/unittests/models/test_litellm.py @@ -13,11 +13,13 @@ # limitations under the Licens import contextlib +from enum import Enum import json import logging import os import sys import tempfile +from typing import List import unittest from unittest.mock import ANY from unittest.mock import AsyncMock @@ -218,8 +220,25 @@ ] +class _EnumType(str, Enum): + """Enum used to validate $ref handling.""" + + A = "A" + B = "B" + + +class _InnerObject(BaseModel): + enum_value: _EnumType + detail: str = Field(description="Some detail") + + +class _OuterObject(BaseModel): + items: List[_InnerObject] + name: str = Field(description="Name field") + + class _StructuredOutput(BaseModel): - value: int = Field(description="Value to emit") + objects: List[_OuterObject] class _ModelDumpOnly: @@ -321,8 +340,18 @@ def test_to_litellm_response_format_uses_json_schema_for_openai_model(): assert "json_schema" in formatted assert formatted["json_schema"]["name"] == "_StructuredOutput" assert formatted["json_schema"]["strict"] is True - assert formatted["json_schema"]["schema"]["additionalProperties"] is False - assert "additionalProperties" in formatted["json_schema"]["schema"] + + schema = formatted["json_schema"]["schema"] + assert schema["additionalProperties"] is False + assert schema["$defs"]["_OuterObject"]["additionalProperties"] is False + assert schema["$defs"]["_InnerObject"]["additionalProperties"] is False + + enum_ref = schema["$defs"]["_InnerObject"]["properties"]["enum_value"] + assert enum_ref == {"$ref": "#/$defs/_EnumType"} + + items_ref = schema["$defs"]["_OuterObject"]["properties"]["items"] + assert items_ref["type"] == "array" + assert items_ref["items"] == {"$ref": "#/$defs/_InnerObject"} def test_to_litellm_response_format_uses_response_schema_for_gemini_model(): @@ -333,7 +362,19 @@ def test_to_litellm_response_format_uses_response_schema_for_gemini_model(): assert formatted["type"] == "json_object" assert "response_schema" in formatted - assert formatted["response_schema"] == _StructuredOutput.model_json_schema() + + schema = formatted["response_schema"] + assert schema == _StructuredOutput.model_json_schema() + assert "additionalProperties" not in schema + assert "additionalProperties" not in schema["$defs"]["_OuterObject"] + assert "additionalProperties" not in schema["$defs"]["_InnerObject"] + + enum_ref = schema["$defs"]["_InnerObject"]["properties"]["enum_value"] + assert enum_ref == {"$ref": "#/$defs/_EnumType"} + + items_ref = schema["$defs"]["_OuterObject"]["properties"]["items"] + assert items_ref["type"] == "array" + assert items_ref["items"] == {"$ref": "#/$defs/_InnerObject"} def test_to_litellm_response_format_uses_response_schema_for_vertex_gemini(): @@ -344,7 +385,19 @@ def test_to_litellm_response_format_uses_response_schema_for_vertex_gemini(): assert formatted["type"] == "json_object" assert "response_schema" in formatted - assert formatted["response_schema"] == _StructuredOutput.model_json_schema() + + schema = formatted["response_schema"] + assert schema == _StructuredOutput.model_json_schema() + assert "additionalProperties" not in schema + assert "additionalProperties" not in schema["$defs"]["_OuterObject"] + assert "additionalProperties" not in schema["$defs"]["_InnerObject"] + + enum_ref = schema["$defs"]["_InnerObject"]["properties"]["enum_value"] + assert enum_ref == {"$ref": "#/$defs/_EnumType"} + + items_ref = schema["$defs"]["_OuterObject"]["properties"]["items"] + assert items_ref["type"] == "array" + assert items_ref["items"] == {"$ref": "#/$defs/_InnerObject"} def test_to_litellm_response_format_uses_json_schema_for_azure_openai(): @@ -353,12 +406,26 @@ def test_to_litellm_response_format_uses_json_schema_for_azure_openai(): _StructuredOutput, model="azure/gpt-4o" ) + print(f"{formatted=}") + assert formatted["type"] == "json_schema" assert "json_schema" in formatted + assert "schema" in formatted["json_schema"] + assert formatted["json_schema"]["name"] == "_StructuredOutput" assert formatted["json_schema"]["strict"] is True - assert formatted["json_schema"]["schema"]["additionalProperties"] is False - assert "additionalProperties" in formatted["json_schema"]["schema"] + + schema = formatted["json_schema"]["schema"] + assert schema["additionalProperties"] is False + assert schema["$defs"]["_OuterObject"]["additionalProperties"] is False + assert schema["$defs"]["_InnerObject"]["additionalProperties"] is False + + enum_ref = schema["$defs"]["_InnerObject"]["properties"]["enum_value"] + assert enum_ref == {"$ref": "#/$defs/_EnumType"} + + items_ref = schema["$defs"]["_OuterObject"]["properties"]["items"] + assert items_ref["type"] == "array" + assert items_ref["items"] == {"$ref": "#/$defs/_InnerObject"} def test_to_litellm_response_format_uses_json_schema_for_anthropic(): @@ -371,8 +438,18 @@ def test_to_litellm_response_format_uses_json_schema_for_anthropic(): assert "json_schema" in formatted assert formatted["json_schema"]["name"] == "_StructuredOutput" assert formatted["json_schema"]["strict"] is True - assert formatted["json_schema"]["schema"]["additionalProperties"] is False - assert "additionalProperties" in formatted["json_schema"]["schema"] + + schema = formatted["json_schema"]["schema"] + assert schema["additionalProperties"] is False + assert schema["$defs"]["_OuterObject"]["additionalProperties"] is False + assert schema["$defs"]["_InnerObject"]["additionalProperties"] is False + + enum_ref = schema["$defs"]["_InnerObject"]["properties"]["enum_value"] + assert enum_ref == {"$ref": "#/$defs/_EnumType"} + + items_ref = schema["$defs"]["_OuterObject"]["properties"]["items"] + assert items_ref["type"] == "array" + assert items_ref["items"] == {"$ref": "#/$defs/_InnerObject"} def test_to_litellm_response_format_with_dict_schema_for_openai(): @@ -383,11 +460,14 @@ def test_to_litellm_response_format_with_dict_schema_for_openai(): } formatted = _to_litellm_response_format(schema, model="gpt-4o") + print(f"{formatted=}") assert formatted["type"] == "json_schema" assert formatted["json_schema"]["name"] == "response" assert formatted["json_schema"]["strict"] is True - assert formatted["json_schema"]["schema"]["additionalProperties"] is False + + schema = formatted["json_schema"]["schema"] + assert schema["additionalProperties"] is False async def test_get_completion_inputs_uses_openai_format_for_openai_model(): @@ -401,13 +481,24 @@ async def test_get_completion_inputs_uses_openai_format_for_openai_model(): llm_request, model="gpt-4o-mini" ) + print(f"{response_format=}") + assert response_format["type"] == "json_schema" assert "json_schema" in response_format assert response_format["json_schema"]["name"] == "_StructuredOutput" assert response_format["json_schema"]["strict"] is True - assert ( - response_format["json_schema"]["schema"]["additionalProperties"] is False - ) + + schema = response_format["json_schema"]["schema"] + assert schema["additionalProperties"] is False + assert schema["$defs"]["_OuterObject"]["additionalProperties"] is False + assert schema["$defs"]["_InnerObject"]["additionalProperties"] is False + + enum_ref = schema["$defs"]["_InnerObject"]["properties"]["enum_value"] + assert enum_ref == {"$ref": "#/$defs/_EnumType"} + + items_ref = schema["$defs"]["_OuterObject"]["properties"]["items"] + assert items_ref["type"] == "array" + assert items_ref["items"] == {"$ref": "#/$defs/_InnerObject"} async def test_get_completion_inputs_uses_gemini_format_for_gemini_model(): @@ -424,6 +515,18 @@ async def test_get_completion_inputs_uses_gemini_format_for_gemini_model(): assert response_format["type"] == "json_object" assert "response_schema" in response_format + schema = response_format["response_schema"] + assert "additionalProperties" not in schema + assert "additionalProperties" not in schema["$defs"]["_OuterObject"] + assert "additionalProperties" not in schema["$defs"]["_InnerObject"] + + enum_ref = schema["$defs"]["_InnerObject"]["properties"]["enum_value"] + assert enum_ref == {"$ref": "#/$defs/_EnumType"} + + items_ref = schema["$defs"]["_OuterObject"]["properties"]["items"] + assert items_ref["type"] == "array" + assert items_ref["items"] == {"$ref": "#/$defs/_InnerObject"} + async def test_get_completion_inputs_uses_passed_model_for_response_format(): """Test that _get_completion_inputs uses the passed model parameter for response format. @@ -445,9 +548,18 @@ async def test_get_completion_inputs_uses_passed_model_for_response_format(): assert "json_schema" in response_format assert response_format["json_schema"]["name"] == "_StructuredOutput" assert response_format["json_schema"]["strict"] is True - assert ( - response_format["json_schema"]["schema"]["additionalProperties"] is False - ) + + schema = response_format["json_schema"]["schema"] + assert schema["additionalProperties"] is False + assert schema["$defs"]["_OuterObject"]["additionalProperties"] is False + assert schema["$defs"]["_InnerObject"]["additionalProperties"] is False + + enum_ref = schema["$defs"]["_InnerObject"]["properties"]["enum_value"] + assert enum_ref == {"$ref": "#/$defs/_EnumType"} + + items_ref = schema["$defs"]["_OuterObject"]["properties"]["items"] + assert items_ref["type"] == "array" + assert items_ref["items"] == {"$ref": "#/$defs/_InnerObject"} async def test_get_completion_inputs_uses_passed_model_for_gemini_format(): @@ -469,6 +581,18 @@ async def test_get_completion_inputs_uses_passed_model_for_gemini_format(): assert response_format["type"] == "json_object" assert "response_schema" in response_format + schema = response_format["response_schema"] + assert "additionalProperties" not in schema + assert "additionalProperties" not in schema["$defs"]["_OuterObject"] + assert "additionalProperties" not in schema["$defs"]["_InnerObject"] + + enum_ref = schema["$defs"]["_InnerObject"]["properties"]["enum_value"] + assert enum_ref == {"$ref": "#/$defs/_EnumType"} + + items_ref = schema["$defs"]["_OuterObject"]["properties"]["items"] + assert items_ref["type"] == "array" + assert items_ref["items"] == {"$ref": "#/$defs/_InnerObject"} + def test_schema_to_dict_filters_none_enum_values(): # Use model_construct to bypass strict enum validation. @@ -838,7 +962,6 @@ def completion(self, model, messages, tools, stream, **kwargs): @pytest.mark.asyncio async def test_generate_content_async(mock_acompletion, lite_llm_instance): - async for response in lite_llm_instance.generate_content_async( LLM_REQUEST_WITH_FUNCTION_DECLARATION ): @@ -1006,7 +1129,6 @@ async def test_generate_content_async_adds_fallback_user_message( def test_maybe_append_user_content( lite_llm_instance, llm_request, expected_output ): - lite_llm_instance._maybe_append_user_content(llm_request) assert len(llm_request.contents) == expected_output @@ -2572,7 +2694,6 @@ async def test_completion_with_drop_params(mock_completion, mock_client): async def test_generate_content_async_stream( mock_completion, lite_llm_instance ): - mock_completion.return_value = iter(STREAMING_MODEL_RESPONSE) responses = [ @@ -2621,7 +2742,6 @@ async def test_generate_content_async_stream( async def test_generate_content_async_stream_with_usage_metadata( mock_completion, lite_llm_instance ): - streaming_model_response_with_usage_metadata = [ *STREAMING_MODEL_RESPONSE, ModelResponse( From 319d4cb3a049244c1149ccca6f5c4b1b671ad469 Mon Sep 17 00:00:00 2001 From: Prathamesh Lohakare Date: Mon, 29 Dec 2025 22:42:32 +0530 Subject: [PATCH 2/3] remove unwanted statements --- tests/unittests/models/test_litellm.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/tests/unittests/models/test_litellm.py b/tests/unittests/models/test_litellm.py index d0f6a768c3..d607a52b97 100644 --- a/tests/unittests/models/test_litellm.py +++ b/tests/unittests/models/test_litellm.py @@ -406,8 +406,6 @@ def test_to_litellm_response_format_uses_json_schema_for_azure_openai(): _StructuredOutput, model="azure/gpt-4o" ) - print(f"{formatted=}") - assert formatted["type"] == "json_schema" assert "json_schema" in formatted assert "schema" in formatted["json_schema"] @@ -460,7 +458,6 @@ def test_to_litellm_response_format_with_dict_schema_for_openai(): } formatted = _to_litellm_response_format(schema, model="gpt-4o") - print(f"{formatted=}") assert formatted["type"] == "json_schema" assert formatted["json_schema"]["name"] == "response" @@ -481,8 +478,6 @@ async def test_get_completion_inputs_uses_openai_format_for_openai_model(): llm_request, model="gpt-4o-mini" ) - print(f"{response_format=}") - assert response_format["type"] == "json_schema" assert "json_schema" in response_format assert response_format["json_schema"]["name"] == "_StructuredOutput" From 527414758c54c50aecfddceccb0eda8b82e53d7a Mon Sep 17 00:00:00 2001 From: Prathamesh Lohakare Date: Mon, 29 Dec 2025 23:50:26 +0530 Subject: [PATCH 3/3] handle more schema options for openai response schema --- src/google/adk/models/lite_llm.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/google/adk/models/lite_llm.py b/src/google/adk/models/lite_llm.py index 424d03fa52..3f06e1edbe 100644 --- a/src/google/adk/models/lite_llm.py +++ b/src/google/adk/models/lite_llm.py @@ -1171,6 +1171,11 @@ def _enforce_closed_schema(schema: dict): for def_schema in schema["$defs"].values(): _enforce_closed_schema(def_schema) + for key in ("anyOf", "oneOf", "allOf"): + if key in schema: + for sub_schema in schema[key]: + _enforce_closed_schema(sub_schema) + def _to_litellm_response_format( response_schema: types.SchemaUnion,