scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

fix(vertex_ai): keep oneOf branches in Vertex and Gemini tool schemas as anyOf

Browse source 
Vertex and the Gemini API have no oneOf, so its branches were dropped and a field with required plus oneOf answered 400. Map oneOf to anyOf before the type-list step and give each branch its node's own keywords.
This commit is contained in:
mateo-berri 2026-10-30 11:02:34 -07:00
parent df47ec0171
commit a24ba9120d
3 changed files with 240 additions and 31 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

93
litellm/llms/vertex_ai/common_utils.py
View file

@ -833,21 +833,84 @@ def _fix_enum_types(schema, depth=0, send_integer_enums: bool = False):
_fix_enum_types(item, depth=depth + 1, send_integer_enums=send_integer_enums)
def _with_oneof_as_anyof(schema: Mapping[str, JsonValue], depth: int = 0) -> Mapping[str, object]:
if depth > DEFAULT_MAX_RECURSE_DEPTH:
raise ValueError(
f"Max depth of {DEFAULT_MAX_RECURSE_DEPTH} exceeded while processing schema. Please check the schema for excessive nesting."
)
node: Final = _oneof_node_as_anyof(schema)
return {key: _subschemas_with_oneof_as_anyof(key, value, depth + 1) for key, value in node.items()}
def _oneof_node_as_anyof(schema: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]:
match schema:
case {"oneOf": list() as branches, **keywords} if "anyOf" not in keywords:
return {"anyOf": [_branch_with_node_keywords(keywords, branch) for branch in branches]}
case _:
return schema
def _branch_with_node_keywords(keywords: Mapping[str, JsonValue], branch: JsonValue) -> JsonValue:
if not isinstance(branch, dict):
return branch
properties: Final = {**_object_keyword(keywords, "properties"), **_object_keyword(branch, "properties")}
named: Final = (
name
for name in (*_list_keyword(keywords, "required"), *_list_keyword(branch, "required"))
if isinstance(name, str)
)
required: Final[JsonValue] = [name for name in dict.fromkeys(named) if name in properties]
merged: Final = {
key: value for key, value in {**keywords, **branch}.items() if key not in ("properties", "required")
}
return {
**merged,
**({"properties": properties} if properties else {}),
**({"required": required} if required else {}),
}
def _object_keyword(schema: Mapping[str, JsonValue], key: str) -> Mapping[str, JsonValue]:
value: Final = schema.get(key)
return value if isinstance(value, dict) else {}
def _list_keyword(schema: Mapping[str, JsonValue], key: str) -> Sequence[JsonValue]:
value: Final = schema.get(key)
return value if isinstance(value, list) else ()
def _subschema_with_oneof_as_anyof(node: JsonValue, depth: int) -> object:
return _with_oneof_as_anyof(node, depth) if isinstance(node, dict) else node
def _subschemas_with_oneof_as_anyof(key: str, value: JsonValue, depth: int) -> object:
match key, value:
case "properties", dict():
return {name: _subschema_with_oneof_as_anyof(child, depth) for name, child in value.items()}
case "anyOf", list():
return [_subschema_with_oneof_as_anyof(branch, depth) for branch in value]
case "items", _:
return _subschema_with_oneof_as_anyof(value, depth)
case _:
return value
def build_vertex_schema(
parameters: dict, add_property_ordering: bool = False, send_integer_enums: bool = False
) -> dict[str, object]:
"""
This is a modified version of https://github.com/google-gemini/generative-ai-python/blob/8f77cc6ac99937cd3a81299ecf79608b91b06bbb/google/generativeai/types/content_types.py#L419
Updates the input parameters, removing extraneous fields, adjusting types, unwinding $defs, and adding propertyOrdering if specified, returning the updated parameters.
Builds a Vertex schema from the input parameters, removing extraneous fields, adjusting types, unwinding $defs, and adding propertyOrdering if specified.
Parameters:
parameters: dict - the json schema to build from
parameters: dict - the json schema to build from. Its $defs are popped and its $ref entries expanded in place
add_property_ordering: bool - whether to add propertyOrdering to the schema. This is only applicable to schemas for structured outputs. See
set_schema_property_ordering for more details.
send_integer_enums: bool - whether to send integer enums with string values instead of removing them
Returns:
parameters: dict - the input parameters, modified in place
dict - the Vertex schema
"""
# Get valid fields from Schema TypedDict
valid_schema_fields: Final = set(get_type_hints(Schema).keys())
@ -859,32 +922,34 @@ def build_vertex_schema(
# refs recursively and correctly detects/skips circular references.
unpack_defs(parameters, defs)
_convert_schema_types(parameters)
schema: Final = dict(_with_oneof_as_anyof(parameters))
_convert_schema_types(schema)
# 5. Nullable fields:
# * https://github.com/pydantic/pydantic/issues/1270
# * https://stackoverflow.com/a/58841311
# * https://github.com/pydantic/pydantic/discussions/4872
convert_anyof_null_to_nullable(parameters)
convert_anyof_null_to_nullable(schema)
# Handle empty strings in enum values - Gemini doesn't accept empty strings in enums
_fix_enum_empty_strings(parameters)
_fix_enum_empty_strings(schema)
# Remove enums for non-string typed fields (Gemini requires enum only on strings)
_fix_enum_types(parameters, send_integer_enums=send_integer_enums)
_fix_enum_types(schema, send_integer_enums=send_integer_enums)
# Handle empty items objects
process_items(parameters)
add_object_type(parameters)
process_items(schema)
add_object_type(schema)
# Postprocessing
# Filter out fields that don't exist in Schema
parameters = filter_schema_fields(parameters, valid_schema_fields)
filtered: Final = filter_schema_fields(schema, valid_schema_fields)
if add_property_ordering:
set_schema_property_ordering(parameters)
set_schema_property_ordering(filtered)
return parameters
return filtered
_build_vertex_schema = build_vertex_schema
@ -998,6 +1063,10 @@ def set_schema_property_ordering(schema: dict[str, object], depth: int = 0) -> d
items: Final = schema.get("items")
if isinstance(items, dict):
set_schema_property_ordering(cast("dict[str, object]", items), depth + 1) # cast-ok: JSON Schema child
any_of: Final = schema.get("anyOf")
for branch in any_of if isinstance(any_of, list) else ():
if isinstance(branch, dict):
set_schema_property_ordering(cast("dict[str, object]", branch), depth + 1) # cast-ok: JSON Schema child
return schema

16
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
View file

@ -1215,19 +1215,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
return self._map_function(value, optional_params, send_integer_enums)
def _map_response_schema(self, value: dict, send_integer_enums: bool = False) -> Mapping[str, object]:
old_schema = deepcopy(value)
if isinstance(old_schema, list):
for item in old_schema:
if isinstance(item, dict):
item = build_vertex_schema(
parameters=item, add_property_ordering=True, send_integer_enums=send_integer_enums
)
elif isinstance(old_schema, dict):
old_schema = build_vertex_schema(
parameters=old_schema, add_property_ordering=True, send_integer_enums=send_integer_enums
)
return old_schema
return build_vertex_schema(
parameters=deepcopy(value), add_property_ordering=True, send_integer_enums=send_integer_enums
)
def apply_response_schema_transformation(
self, value: dict, optional_params: dict, model: str, allow_integer_enums: bool = True

162
tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py
View file

@ -1,6 +1,7 @@
import copy
import functools
import time
from collections.abc import Iterator
from collections.abc import Callable, Iterator
from contextlib import contextmanager
from enum import Enum
from typing import Final
@ -794,6 +795,153 @@ def test_build_vertex_schema_array_branch_missing_items_in_anyof():
assert branch.get("items") == {"type": "object"}, f"array branch must have items synthesized; got {branch}"
def _branches() -> list[JsonValue]:
return [
{"properties": {"a": {"type": "string"}}, "required": ["a"]},
{"properties": {"b": {"type": "integer"}}, "required": ["b"]},
]
def _payment_union(union: str) -> dict[str, JsonValue]:
return {
"$defs": {
"Card": {
"type": "object",
"properties": {"kind": {"type": "string"}, "number": {"type": "string"}},
"required": ["kind", "number"],
},
"Bank": {
"type": "object",
"properties": {"kind": {"type": "string"}, "iban": {"type": "string"}},
"required": ["kind", "iban"],
},
},
"type": "object",
"properties": {
"payment": {
union: [{"$ref": "#/$defs/Card"}, {"$ref": "#/$defs/Bank"}],
"discriminator": {"propertyName": "kind"},
}
},
"required": ["payment"],
}
@pytest.mark.parametrize(
"shape",
[
lambda union: {"type": "object", "properties": {"f": {union: _branches()}}},
lambda union: {"type": "object", "properties": {"f": {"type": "array", "items": {union: _branches()}}}},
lambda union: {
"type": "object",
"properties": {"f": {"anyOf": [{"type": "object", "properties": {"g": {union: _branches()}}}]}},
},
lambda union: {
"type": "object",
"properties": {"f": {union: [{"type": "object", "properties": {"g": {union: _branches()}}}]}},
},
lambda union: {"type": "object", "properties": {"f": {union: ["x", {"type": "string"}]}}},
_payment_union,
],
ids=["untyped", "array-items", "inside-anyof", "inside-oneof", "non-object-branch", "pydantic-union"],
)
def test_vertex_schema_builds_oneof_branches_like_anyof_branches(shape: Callable[[str], dict[str, JsonValue]]):
assert build_vertex_schema(shape("oneOf")) == build_vertex_schema(shape("anyOf"))
@pytest.mark.parametrize(
("field", "expected"),
[
(
{
"type": "object",
"required": ["a"],
"oneOf": [{"properties": {"a": {"type": "string"}}}, {"properties": {"b": {"type": "string"}}}],
},
[
{"type": "object", "properties": {"a": {"type": "string"}}, "required": ["a"]},
{"type": "object", "properties": {"b": {"type": "string"}}},
],
),
(
{
"type": "object",
"properties": {"x": {"type": "string"}, "y": {"type": "integer"}},
"required": ["x"],
"oneOf": [{"required": ["x"]}, {"required": ["y"]}],
},
[
{"type": "object", "properties": {"x": {"type": "string"}, "y": {"type": "integer"}}, "required": ["x"]},
{
"type": "object",
"properties": {"x": {"type": "string"}, "y": {"type": "integer"}},
"required": ["x", "y"],
},
],
),
(
{"type": "object", "required": [["a"], "a"], "oneOf": [{"properties": {"a": {"type": "string"}}}]},
[{"type": "object", "properties": {"a": {"type": "string"}}, "required": ["a"]}],
),
(
{"type": "string", "description": "a letter", "oneOf": [{"enum": ["a"]}, {"enum": ["b"], "description": "b"}]},
[
{"type": "string", "description": "a letter", "enum": ["a"]},
{"type": "string", "description": "b", "enum": ["b"]},
],
),
(
{"type": ["object", "null"], "required": ["a"], "oneOf": [{"properties": {"a": {"type": "string"}}}]},
[{"anyOf": [{"type": "object", "properties": {"a": {"type": "string"}}, "required": ["a"], "nullable": True}]}],
),
],
ids=["required-on-the-node", "required-in-the-branches", "non-string-required", "scalar-node", "type-list-node"],
)
def test_vertex_schema_gives_every_oneof_branch_the_keywords_of_its_node(field: JsonValue, expected: JsonValue):
schema: Final = build_vertex_schema({"type": "object", "properties": {"f": field}})
assert schema["properties"]["f"] == {"anyOf": expected}
def test_vertex_schema_keeps_every_oneof_branch_with_its_properties():
schema: Final = build_vertex_schema(_payment_union("oneOf"))
assert [branch["properties"] for branch in schema["properties"]["payment"]["anyOf"]] == [
{"kind": {"type": "string"}, "number": {"type": "string"}},
{"kind": {"type": "string"}, "iban": {"type": "string"}},
]
@pytest.mark.parametrize(
("field", "expected"),
[
(
{"anyOf": [{"type": "string"}], "oneOf": [{"type": "integer"}, {"type": "boolean"}]},
{"anyOf": [{"type": "string"}]},
),
({"type": "string", "oneOf": {"type": "integer"}}, {"type": "string"}),
({"type": "object", "example": {"oneOf": ["x"]}}, {"type": "object", "example": {"oneOf": ["x"]}}),
],
ids=["anyof-wins-over-oneof", "non-list-oneof", "oneof-inside-example-data"],
)
def test_vertex_schema_leaves_oneof_alone_where_it_is_not_a_union_it_can_send(field: JsonValue, expected: JsonValue):
schema: Final = build_vertex_schema({"type": "object", "properties": {"f": field}})
assert schema["properties"]["f"] == expected
def test_vertex_schema_rejects_oneof_nested_past_the_max_depth():
leaf: Final[dict[str, JsonValue]] = {"type": "string"}
deepest: Final = functools.reduce(
lambda child, _: {"oneOf": [{"type": "object", "properties": {"nested": child}}]},
range(DEFAULT_MAX_RECURSE_DEPTH * 10),
leaf,
)
with pytest.raises(ValueError, match=f"Max depth of {DEFAULT_MAX_RECURSE_DEPTH} exceeded"):
build_vertex_schema({"type": "object", "properties": {"f": deepest}})
def test_vertex_ai_complex_response_schema():
from copy import deepcopy
@ -973,11 +1121,13 @@ def test_vertex_ai_complex_response_schema():
subfield2 = field3_items["properties"]["subfield2"]
assert subfield2["type"] == "array"
# Check subfield2 items structure
subfield2_items = subfield2["items"]
assert subfield2_items["type"] == "object"
assert "propertyOrdering" in subfield2_items
assert subfield2_items["propertyOrdering"] == ["Type1", "Type3", "Type2"]
subfield2_branches = subfield2["items"]["anyOf"]
assert [(branch["type"], branch["propertyOrdering"], branch["required"]) for branch in subfield2_branches] == [
("object", ["Type1", "Type3", "Type2"], ["Type1"]),
("object", ["Type1", "Type3", "Type2"], ["Type2"]),
("object", ["Type1", "Type3", "Type2"], ["Type3"]),
]
subfield2_items = subfield2_branches[0]
# Check Type1 structure
type1 = subfield2_items["properties"]["Type1"]