8864 Add support for anyOf union type while handling null fields

This commit is contained in:
Nicholas Grabar 2025-03-25 22:37:28 -07:00
parent 122ee634f4
commit f68cc26f15
3 changed files with 95 additions and 20 deletions

2
.gitignore vendored
View file

@ -82,4 +82,4 @@ tests/llm_translation/test_vertex_key.json
litellm/proxy/migrations/0_init/migration.sql litellm/proxy/migrations/0_init/migration.sql
litellm/proxy/db/migrations/0_init/migration.sql litellm/proxy/db/migrations/0_init/migration.sql
litellm/proxy/db/migrations/* litellm/proxy/db/migrations/*
litellm/proxy/migrations/* litellm/proxy/migrations/*config.yaml

View file

@ -160,7 +160,7 @@ def _build_vertex_schema(parameters: dict):
# * https://github.com/pydantic/pydantic/issues/1270 # * https://github.com/pydantic/pydantic/issues/1270
# * https://stackoverflow.com/a/58841311 # * https://stackoverflow.com/a/58841311
# * https://github.com/pydantic/pydantic/discussions/4872 # * https://github.com/pydantic/pydantic/discussions/4872
convert_to_nullable(parameters) convert_anyof_null_to_nullable(parameters)
add_object_type(parameters) add_object_type(parameters)
# Postprocessing # Postprocessing
# 4. Suppress unnecessary title generation: # 4. Suppress unnecessary title generation:
@ -211,34 +211,39 @@ def unpack_defs(schema, defs):
continue continue
def convert_to_nullable(schema): def convert_anyof_null_to_nullable(schema):
anyof = schema.pop("anyOf", None) """ Converts null objects within anyOf by removing them and adding nullable to all remaining objects """
anyof = schema.get("anyOf", None)
if anyof is not None: if anyof is not None:
if len(anyof) != 2: contains_null = False
for atype in anyof:
if atype == {"type": "null"}:
# remove null type
anyof.remove(atype)
contains_null = True
if len(anyof) == 0:
# Edge case: response schema with only null type present is invalid in Vertex AI
raise ValueError( raise ValueError(
"Invalid input: Type Unions are not supported, except for `Optional` types. " "Invalid input: AnyOf schema with only null type is not supported. "
"Please provide an `Optional` type or a non-Union type." "Please provide a non-null type."
) )
a, b = anyof
if a == {"type": "null"}:
schema.update(b) if contains_null:
elif b == {"type": "null"}: # set all types to nullable following guidance found here: https://cloud.google.com/vertex-ai/generative-ai/docs/samples/generativeaionvertexai-gemini-controlled-generation-response-schema-3#generativeaionvertexai_gemini_controlled_generation_response_schema_3-python
schema.update(a) for atype in anyof:
else: atype["nullable"] = True
raise ValueError(
"Invalid input: Type Unions are not supported, except for `Optional` types. "
"Please provide an `Optional` type or a non-Union type."
)
schema["nullable"] = True
properties = schema.get("properties", None) properties = schema.get("properties", None)
if properties is not None: if properties is not None:
for name, value in properties.items(): for name, value in properties.items():
convert_to_nullable(value) convert_anyof_null_to_nullable(value)
items = schema.get("items", None) items = schema.get("items", None)
if items is not None: if items is not None:
convert_to_nullable(items) convert_anyof_null_to_nullable(items)
def add_object_type(schema): def add_object_type(schema):

View file

@ -12,6 +12,7 @@ import litellm
from litellm.llms.vertex_ai.common_utils import ( from litellm.llms.vertex_ai.common_utils import (
get_vertex_location_from_url, get_vertex_location_from_url,
get_vertex_project_id_from_url, get_vertex_project_id_from_url,
convert_anyof_null_to_nullable
) )
@ -41,3 +42,72 @@ async def test_get_vertex_location_from_url():
url = "https://invalid-url.com" url = "https://invalid-url.com"
location = get_vertex_location_from_url(url) location = get_vertex_location_from_url(url)
assert location is None assert location is None
def test_basic_anyof_conversion():
"""Test basic conversion of anyOf with 'null'."""
schema = {
"type": "object",
"properties": {
"example": {
"anyOf": [
{"type": "string"},
{"type": "null"}
]
}
}
}
convert_anyof_null_to_nullable(schema)
expected = {
"type": "object",
"properties": {
"example": {
"anyOf": [
{"type": "string", "nullable": True}
]
}
}
}
assert schema == expected
def test_nested_anyof_conversion():
"""Test nested conversion with 'anyOf' inside properties."""
schema = {
"type": "object",
"properties": {
"outer": {
"type": "object",
"properties": {
"inner": {
"anyOf": [
{"type": "array", "items": {"type": "string"}},
{"type": "string"},
{"type": "null"}
]
}
}
}
}
}
convert_anyof_null_to_nullable(schema)
expected = {
"type": "object",
"properties": {
"outer": {
"type": "object",
"properties": {
"inner": {
"anyOf": [
{"type": "array", "items": {"type": "string"}, "nullable": True},
{"type": "string", "nullable": True}
]
}
}
}
}
}
assert schema == expected