diff --git a/.gitignore b/.gitignore index dab6d4ec81..a9d14852eb 100644 --- a/.gitignore +++ b/.gitignore @@ -82,4 +82,4 @@ tests/llm_translation/test_vertex_key.json litellm/proxy/migrations/0_init/migration.sql litellm/proxy/db/migrations/0_init/migration.sql litellm/proxy/db/migrations/* -litellm/proxy/migrations/* \ No newline at end of file +litellm/proxy/migrations/*config.yaml diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index a3f91fbacc..9a01c15233 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -160,7 +160,7 @@ def _build_vertex_schema(parameters: dict): # * https://github.com/pydantic/pydantic/issues/1270 # * https://stackoverflow.com/a/58841311 # * https://github.com/pydantic/pydantic/discussions/4872 - convert_to_nullable(parameters) + convert_anyof_null_to_nullable(parameters) add_object_type(parameters) # Postprocessing # 4. Suppress unnecessary title generation: @@ -211,34 +211,39 @@ def unpack_defs(schema, defs): continue -def convert_to_nullable(schema): - anyof = schema.pop("anyOf", None) +def convert_anyof_null_to_nullable(schema): + """ 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 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( - "Invalid input: Type Unions are not supported, except for `Optional` types. " - "Please provide an `Optional` type or a non-Union type." + "Invalid input: AnyOf schema with only null type is not supported. " + "Please provide a non-null type." ) - a, b = anyof - if a == {"type": "null"}: - schema.update(b) - elif b == {"type": "null"}: - schema.update(a) - else: - 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 + + + if contains_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 + for atype in anyof: + atype["nullable"] = True + properties = schema.get("properties", None) if properties is not None: for name, value in properties.items(): - convert_to_nullable(value) + convert_anyof_null_to_nullable(value) items = schema.get("items", None) if items is not None: - convert_to_nullable(items) + convert_anyof_null_to_nullable(items) def add_object_type(schema): diff --git a/tests/litellm/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/litellm/llms/vertex_ai/test_vertex_ai_common_utils.py index e89355443f..1c6ff65ce6 100644 --- a/tests/litellm/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/litellm/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -12,6 +12,7 @@ import litellm from litellm.llms.vertex_ai.common_utils import ( get_vertex_location_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" location = get_vertex_location_from_url(url) 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 \ No newline at end of file