Skip to content

Commit 848fdbe

Browse files
GWealecopybara-github
authored andcommitted
fix: Filter out None values from enum lists in schema generation
Close #3552 Co-authored-by: George Weale <gweale@google.com> PiperOrigin-RevId: 834967220
1 parent 31cfa3b commit 848fdbe

2 files changed

Lines changed: 43 additions & 24 deletions

File tree

src/google/adk/models/lite_llm.py

Lines changed: 20 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -632,49 +632,45 @@ def _to_litellm_role(role: Optional[str]) -> Literal["user", "assistant"]:
632632
}
633633

634634

635-
def _schema_to_dict(schema: types.Schema) -> dict:
636-
"""Recursively converts a types.Schema to a pure-python dict
637-
638-
with all enum values written as lower-case strings.
635+
def _schema_to_dict(schema: types.Schema | dict[str, Any]) -> dict:
636+
"""Recursively converts a schema object or dict to a pure-python dict.
639637
640638
Args:
641639
schema: The schema to convert.
642640
643641
Returns:
644642
The dictionary representation of the schema.
645643
"""
646-
# Dump without json encoding so we still get Enum members
647-
schema_dict = schema.model_dump(exclude_none=True)
644+
schema_dict = (
645+
schema.model_dump(exclude_none=True)
646+
if isinstance(schema, types.Schema)
647+
else dict(schema)
648+
)
649+
enum_values = schema_dict.get("enum")
650+
if isinstance(enum_values, (list, tuple)):
651+
schema_dict["enum"] = [value for value in enum_values if value is not None]
648652

649-
# ---- normalise this level ------------------------------------------------
650-
if "type" in schema_dict:
651-
# schema_dict["type"] can be an Enum or a str
653+
if "type" in schema_dict and schema_dict["type"] is not None:
652654
t = schema_dict["type"]
653-
schema_dict["type"] = (t.value if isinstance(t, types.Type) else t).lower()
655+
schema_dict["type"] = (
656+
t.value if isinstance(t, types.Type) else str(t)
657+
).lower()
654658

655-
# ---- recurse into `items` -----------------------------------------------
656659
if "items" in schema_dict:
657-
schema_dict["items"] = _schema_to_dict(
658-
schema.items
659-
if isinstance(schema.items, types.Schema)
660-
else types.Schema.model_validate(schema_dict["items"])
660+
items = schema_dict["items"]
661+
schema_dict["items"] = (
662+
_schema_to_dict(items)
663+
if isinstance(items, (types.Schema, dict))
664+
else items
661665
)
662666

663-
# ---- recurse into `properties` ------------------------------------------
664667
if "properties" in schema_dict:
665668
new_props = {}
666669
for key, value in schema_dict["properties"].items():
667-
# value is a dict → rebuild a Schema object and recurse
668-
if isinstance(value, dict):
669-
new_props[key] = _schema_to_dict(types.Schema.model_validate(value))
670-
# value is already a Schema instance
671-
elif isinstance(value, types.Schema):
670+
if isinstance(value, (types.Schema, dict)):
672671
new_props[key] = _schema_to_dict(value)
673-
# plain dict without nested schemas
674672
else:
675673
new_props[key] = value
676-
if "type" in new_props[key]:
677-
new_props[key]["type"] = new_props[key]["type"].lower()
678674
schema_dict["properties"] = new_props
679675

680676
return schema_dict

tests/unittests/models/test_litellm.py

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
from google.adk.models.lite_llm import _model_response_to_chunk
2727
from google.adk.models.lite_llm import _model_response_to_generate_content_response
2828
from google.adk.models.lite_llm import _parse_tool_calls_from_text
29+
from google.adk.models.lite_llm import _schema_to_dict
2930
from google.adk.models.lite_llm import _split_message_content_and_tool_calls
3031
from google.adk.models.lite_llm import _to_litellm_response_format
3132
from google.adk.models.lite_llm import _to_litellm_role
@@ -286,6 +287,28 @@ def test_to_litellm_response_format_handles_genai_schema_instance():
286287
)
287288

288289

290+
def test_schema_to_dict_filters_none_enum_values():
291+
# Use model_construct to bypass strict enum validation.
292+
top_level_schema = types.Schema.model_construct(
293+
type=types.Type.STRING,
294+
enum=["ACTIVE", None, "INACTIVE"],
295+
)
296+
nested_schema = types.Schema.model_construct(
297+
type=types.Type.OBJECT,
298+
properties={
299+
"status": types.Schema.model_construct(
300+
type=types.Type.STRING, enum=["READY", None, "DONE"]
301+
),
302+
},
303+
)
304+
305+
assert _schema_to_dict(top_level_schema)["enum"] == ["ACTIVE", "INACTIVE"]
306+
assert _schema_to_dict(nested_schema)["properties"]["status"]["enum"] == [
307+
"READY",
308+
"DONE",
309+
]
310+
311+
289312
MULTIPLE_FUNCTION_CALLS_STREAM = [
290313
ModelResponse(
291314
choices=[

0 commit comments

Comments
 (0)