From a24ba9120d68bc864a09a3c2bcb1451e3eb2c3d8 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Fri, 30 Oct 2026 11:02:34 -0700
Subject: [PATCH] fix(vertex_ai): keep oneOf branches in Vertex and Gemini tool
 schemas as anyOf

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.
---
 litellm/llms/vertex_ai/common_utils.py        |  93 ++++++++--
 .../vertex_and_google_ai_studio_gemini.py     |  16 +-
 .../vertex_ai/test_vertex_ai_common_utils.py  | 162 +++++++++++++++++-
 3 files changed, 240 insertions(+), 31 deletions(-)

diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py
index 0fa80cb2cad..5f48ca1c550 100644
--- a/litellm/llms/vertex_ai/common_utils.py
+++ b/litellm/llms/vertex_ai/common_utils.py
@@ -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
 
 
diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
index c506192f253..ac238b5ee57 100644
--- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
+++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
@@ -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
diff --git a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py
index 152ab2ffc13..fb8382bba50 100644
--- a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py
+++ b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py
@@ -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"]