fix(vertex_ai): preserve oneOf unions and null branches that follow a removed null in schemas

litellm/llms/vertex_ai/common_utils.py

@@ -574,6 +574,42 @@ def _fix_enum_types(schema, depth=0):
                 _fix_enum_types(item, depth=depth + 1)
 
 
+def _convert_oneof_to_anyof(schema, depth=0):
+    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."
+        )
+
+    if not isinstance(schema, dict):
+        return
+
+    oneof = schema.get("oneOf", None)
+    if (
+        isinstance(oneof, list)
+        and "anyOf" not in schema
+        and "type" not in schema
+        and "properties" not in schema
+        and "items" not in schema
+    ):
+        schema.pop("oneOf")
+        schema["anyOf"] = oneof
+
+    properties: Final = schema.get("properties", None)
+    if properties is not None:
+        for value in properties.values():
+            _convert_oneof_to_anyof(value, depth=depth + 1)
+
+    items: Final = schema.get("items", None)
+    if items is not None:
+        _convert_oneof_to_anyof(items, depth=depth + 1)
+
+    anyof: Final = schema.get("anyOf", None)
+    if anyof is not None and isinstance(anyof, list):
+        for item in anyof:
+            if isinstance(item, dict):
+                _convert_oneof_to_anyof(item, depth=depth + 1)
+
+
 def _build_vertex_schema(parameters: dict, add_property_ordering: bool = False):
     """
     This is a modified version of https://github.com/google-gemini/generative-ai-python/blob/8f77cc6ac99937cd3a81299ecf79608b91b06bbb/google/generativeai/types/content_types.py#L419
@@ -597,6 +633,8 @@ def _build_vertex_schema(parameters: dict, add_property_ordering: bool = False):
     # refs recursively and correctly detects/skips circular references.
     unpack_defs(parameters, defs)
 
+    _convert_oneof_to_anyof(parameters)
+
     # 5. Nullable fields:
     #     * https://github.com/pydantic/pydantic/issues/1270
     #     * https://stackoverflow.com/a/58841311
@@ -782,7 +820,7 @@ def convert_anyof_null_to_nullable(schema, depth=0):
     anyof: Final = schema.get("anyOf", None)
     if anyof is not None:
         contains_null = False
-        for atype in anyof:
+        for atype in list(anyof):
             if isinstance(atype, dict) and atype.get("type") == "null":
                 # remove null type
                 anyof.remove(atype)

tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py

@@ -81,6 +81,86 @@ def test_basic_anyof_conversion():
     assert schema == expected
 
 
+def test_anyof_conversion_with_null_before_other_branches():
+    schema = {
+        "type": "object",
+        "properties": {
+            "example": {
+                "anyOf": [{"type": "null"}, {"type": "null"}, {"type": "string"}]
+            }
+        },
+    }
+
+    convert_anyof_null_to_nullable(schema)
+
+    expected = {
+        "type": "object",
+        "properties": {"example": {"anyOf": [{"type": "string", "nullable": True}]}},
+    }
+    assert schema == expected
+
+
+def test_build_vertex_schema_preserves_oneof_unions():
+    from litellm.llms.vertex_ai.common_utils import _build_vertex_schema
+
+    parameters = {
+        "$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": {
+                "oneOf": [{"$ref": "#/$defs/Card"}, {"$ref": "#/$defs/Bank"}],
+                "discriminator": {"propertyName": "kind"},
+            }
+        },
+        "required": ["payment"],
+    }
+
+    result = _build_vertex_schema(parameters)
+
+    assert result["properties"]["payment"] == {
+        "anyOf": [
+            {
+                "type": "object",
+                "properties": {"kind": {"type": "string"}, "number": {"type": "string"}},
+                "required": ["kind", "number"],
+            },
+            {
+                "type": "object",
+                "properties": {"kind": {"type": "string"}, "iban": {"type": "string"}},
+                "required": ["kind", "iban"],
+            },
+        ]
+    }
+
+
+def test_build_vertex_schema_converts_null_branch_inside_oneof():
+    from litellm.llms.vertex_ai.common_utils import _build_vertex_schema
+
+    parameters = {
+        "type": "object",
+        "properties": {
+            "value": {"oneOf": [{"type": "string"}, {"type": "null"}]}
+        },
+    }
+
+    result = _build_vertex_schema(parameters)
+
+    assert result["properties"]["value"] == {
+        "anyOf": [{"type": "string", "nullable": True}]
+    }
+
+
 def test_nested_anyof_conversion():
     """Test nested conversion with 'anyOf' inside properties."""
     schema = {