diff --git a/src/google/adk/models/lite_llm.py b/src/google/adk/models/lite_llm.py index e85772c5..dad5543f 100644 --- a/src/google/adk/models/lite_llm.py +++ b/src/google/adk/models/lite_llm.py @@ -1491,6 +1491,54 @@ def _message_to_generate_content_response( ) +def _enforce_strict_openai_schema(schema: dict[str, Any]) -> None: + """Recursively transforms a JSON schema for OpenAI strict structured outputs. + + OpenAI strict mode requires: + 1. additionalProperties: false on all object schemas (including nested/$defs). + 2. All properties listed in 'required' (no optional omissions). + 3. $ref nodes must have no sibling keywords (e.g., no 'description' next to + '$ref'). + + This function mutates the schema dict in place. + + Args: + schema: A JSON schema dictionary to transform. + """ + if not isinstance(schema, dict): + return + + # Strip sibling keywords from $ref nodes (OpenAI rejects them). + if "$ref" in schema: + for key in list(schema.keys()): + if key != "$ref": + del schema[key] + return + + # Ensure all object schemas have additionalProperties: false and list every + # property as required. + if schema.get("type") == "object" and "properties" in schema: + schema["additionalProperties"] = False + schema["required"] = sorted(schema["properties"].keys()) + + # Recurse into $defs (Pydantic's nested model definitions). + for defn in schema.get("$defs", {}).values(): + _enforce_strict_openai_schema(defn) + + # Recurse into property schemas. + for prop in schema.get("properties", {}).values(): + _enforce_strict_openai_schema(prop) + + # Recurse into combinators. + for key in ("anyOf", "oneOf", "allOf"): + for item in schema.get(key, []): + _enforce_strict_openai_schema(item) + + # Recurse into array item schemas. + if "items" in schema and isinstance(schema["items"], dict): + _enforce_strict_openai_schema(schema["items"]) + + def _to_litellm_response_format( response_schema: types.SchemaUnion, model: str, @@ -1515,7 +1563,7 @@ def _to_litellm_response_format( and schema_type.lower() in _LITELLM_STRUCTURED_TYPES ): return response_schema - schema_dict = dict(response_schema) + schema_dict = copy.deepcopy(response_schema) if "title" in schema_dict: schema_name = str(schema_dict["title"]) elif isinstance(response_schema, type) and issubclass( @@ -1526,14 +1574,18 @@ def _to_litellm_response_format( 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") + schema_dict = copy.deepcopy( + response_schema.model_dump(exclude_none=True, mode="json") + ) if "title" in schema_dict: schema_name = str(schema_dict["title"]) else: schema_dict = response_schema.__class__.model_json_schema() schema_name = response_schema.__class__.__name__ elif hasattr(response_schema, "model_dump"): - schema_dict = response_schema.model_dump(exclude_none=True, mode="json") + schema_dict = copy.deepcopy( + response_schema.model_dump(exclude_none=True, mode="json") + ) schema_name = response_schema.__class__.__name__ else: logger.warning( @@ -1551,14 +1603,8 @@ 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 - ): - # OpenAI structured outputs require explicit additionalProperties: false. - schema_dict = dict(schema_dict) - schema_dict["additionalProperties"] = False + if isinstance(schema_dict, dict): + _enforce_strict_openai_schema(schema_dict) return { "type": "json_schema", diff --git a/tests/unittests/models/test_litellm.py b/tests/unittests/models/test_litellm.py index 39f6b540..8e353efb 100644 --- a/tests/unittests/models/test_litellm.py +++ b/tests/unittests/models/test_litellm.py @@ -26,6 +26,7 @@ import warnings from google.adk.models.lite_llm import _append_fallback_user_content_if_missing from google.adk.models.lite_llm import _content_to_message_param +from google.adk.models.lite_llm import _enforce_strict_openai_schema from google.adk.models.lite_llm import _FILE_ID_REQUIRED_PROVIDERS from google.adk.models.lite_llm import _FINISH_REASON_MAPPING from google.adk.models.lite_llm import _function_declaration_to_tool_param @@ -394,6 +395,145 @@ def test_to_litellm_response_format_with_dict_schema_for_openai(): assert formatted["json_schema"]["schema"]["additionalProperties"] is False +class _InnerModel(BaseModel): + value: str = Field(description="A value") + optional_field: str | None = Field(default=None, description="Optional") + + +class _OuterModel(BaseModel): + inner: _InnerModel = Field(description="Nested model") + name: str + + +class _WithList(BaseModel): + items: list[_InnerModel] = Field(description="List of items") + label: str + + +def test_enforce_strict_openai_schema_adds_additional_properties_recursively(): + """additionalProperties: false must appear on all object schemas.""" + schema = _OuterModel.model_json_schema() + + _enforce_strict_openai_schema(schema) + + # Root level + assert schema["additionalProperties"] is False + # Nested model in $defs + inner_def = schema["$defs"]["_InnerModel"] + assert inner_def["additionalProperties"] is False + + +def test_enforce_strict_openai_schema_marks_all_properties_required(): + """All properties must appear in 'required', including optional fields.""" + schema = _InnerModel.model_json_schema() + + _enforce_strict_openai_schema(schema) + + assert sorted(schema["required"]) == ["optional_field", "value"] + + +def test_enforce_strict_openai_schema_strips_ref_sibling_keywords(): + """$ref nodes must have no sibling keywords like 'description'.""" + schema = _OuterModel.model_json_schema() + # Pydantic v2 generates {"$ref": "...", "description": "..."} for nested models + inner_prop = schema["properties"]["inner"] + assert "$ref" in inner_prop, "Expected Pydantic to generate a $ref property" + assert len(inner_prop) > 1, "Expected sibling keywords alongside $ref" + + _enforce_strict_openai_schema(schema) + + inner_prop = schema["properties"]["inner"] + assert list(inner_prop.keys()) == ["$ref"] + + +def test_enforce_strict_openai_schema_handles_array_items(): + """Array item schemas should also be recursively transformed.""" + schema = _WithList.model_json_schema() + + _enforce_strict_openai_schema(schema) + + assert schema["additionalProperties"] is False + inner_def = schema["$defs"]["_InnerModel"] + assert inner_def["additionalProperties"] is False + assert sorted(inner_def["required"]) == ["optional_field", "value"] + + +def test_enforce_strict_openai_schema_preserves_anyof_and_default(): + """anyOf structure and default value for Optional fields must be preserved.""" + schema = _InnerModel.model_json_schema() + + _enforce_strict_openai_schema(schema) + + opt_prop = schema["properties"]["optional_field"] + assert opt_prop["anyOf"] == [{"type": "string"}, {"type": "null"}] + assert opt_prop["default"] is None + + +def test_to_litellm_response_format_dict_input_not_mutated(): + """Passing a raw dict should not mutate the caller's original dict.""" + schema = { + "type": "object", + "properties": { + "nested": { + "type": "object", + "properties": {"x": {"type": "string"}}, + } + }, + } + import copy + + original = copy.deepcopy(schema) + + _to_litellm_response_format(schema, model="gpt-4o") + + assert schema == original, "Caller's input dict was mutated" + + +def test_to_litellm_response_format_instance_input_for_openai(): + """Passing a BaseModel instance should produce a valid strict schema.""" + instance = _OuterModel( + inner=_InnerModel(value="test", optional_field=None), name="foo" + ) + + formatted = _to_litellm_response_format(instance, model="gpt-4o") + + assert formatted["type"] == "json_schema" + schema = formatted["json_schema"]["schema"] + assert schema["additionalProperties"] is False + inner_def = schema["$defs"]["_InnerModel"] + assert inner_def["additionalProperties"] is False + assert sorted(inner_def["required"]) == ["optional_field", "value"] + + +def test_to_litellm_response_format_nested_pydantic_for_openai(): + """Nested Pydantic model should produce a valid OpenAI strict schema.""" + formatted = _to_litellm_response_format(_OuterModel, model="gpt-4o") + + assert formatted["type"] == "json_schema" + assert formatted["json_schema"]["strict"] is True + + schema = formatted["json_schema"]["schema"] + assert schema["additionalProperties"] is False + assert sorted(schema["required"]) == ["inner", "name"] + + # $defs inner model must also be strict + inner_def = schema["$defs"]["_InnerModel"] + assert inner_def["additionalProperties"] is False + assert sorted(inner_def["required"]) == ["optional_field", "value"] + + +def test_to_litellm_response_format_nested_pydantic_for_gemini_unchanged(): + """Gemini models should NOT get the strict OpenAI transformations.""" + formatted = _to_litellm_response_format( + _OuterModel, model="gemini/gemini-2.0-flash" + ) + + assert formatted["type"] == "json_object" + schema = formatted["response_schema"] + # Gemini path should pass through the raw Pydantic schema untouched + assert schema == _OuterModel.model_json_schema() + + async def test_get_completion_inputs_uses_openai_format_for_openai_model(): """Test that _get_completion_inputs produces OpenAI-compatible format.""" llm_request = LlmRequest(