fix(vertex-ai): resolve chained root refs and tuple schemas for parametersJsonSchema

A root $ref that points at another $ref now resolves to its end, so a chain
that ends in a union also goes out as parametersJsonSchema. A schema built
with tuples or sets is turned into its JSON form first, so it no longer falls
back to the parameters field. A type of ["object"] stays on the parameters
field, which already accepts it. Only tools whose root could need
parametersJsonSchema pay for the JSON conversion

litellm/llms/vertex_ai/common_utils.py

@@ -11,6 +11,7 @@ from typing import Any, Final, Literal, TypeVar, cast, get_type_hints
 
 import httpx
 from pydantic import JsonValue, TypeAdapter, ValidationError
+from pydantic_core import to_jsonable_python
 from typing_extensions import NotRequired, ReadOnly, TypedDict
 
 import litellm
@@ -891,25 +892,52 @@ _build_vertex_schema = build_vertex_schema
 
 _VERTEX_PARAMETERS_UNSUPPORTED_ROOT_KEYWORDS: Final = ("anyOf", "oneOf", "allOf")
 _ROOT_DEFS_REF_PREFIX: Final = "#/$defs/"
+_JSON_SCHEMA: Final = TypeAdapter(Mapping[str, JsonValue])
 
 
-def _root_ref_target(parameters: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]:
-    match parameters.get("$ref"), parameters.get("$defs"):
-        case str() as ref, dict() as defs if ref.startswith(_ROOT_DEFS_REF_PREFIX):
-            target: Final = defs.get(ref.removeprefix(_ROOT_DEFS_REF_PREFIX))
-            return target if isinstance(target, dict) else parameters
+def _is_absent_or_object_type(schema_type: object) -> bool:
+    match schema_type:
+        case None:
+            return True
+        case str():
+            return schema_type.lower() == "object"
+        case [str() as only_type]:
+            return only_type.lower() == "object"
         case _:
-            return parameters
+            return False
+
+
+def _root_cannot_be_vertex_parameters(root: Mapping[str, object]) -> bool:
+    return any(
+        keyword in root for keyword in _VERTEX_PARAMETERS_UNSUPPORTED_ROOT_KEYWORDS
+    ) or not _is_absent_or_object_type(root.get("type"))
+
+
+def _resolve_root_ref(
+    schema: Mapping[str, JsonValue], defs: Mapping[str, JsonValue], seen_refs: frozenset[str]
+) -> Mapping[str, JsonValue]:
+    match schema.get("$ref"):
+        case str() as ref if ref.startswith(_ROOT_DEFS_REF_PREFIX) and ref not in seen_refs:
+            target: Final = defs.get(ref.removeprefix(_ROOT_DEFS_REF_PREFIX))
+            return _resolve_root_ref(target, defs, seen_refs | {ref}) if isinstance(target, dict) else schema
+        case _:
+            return schema
 
 
 def needs_parameters_json_schema(parameters: Mapping[str, JsonValue]) -> bool:
-    root: Final = _root_ref_target(parameters)
-    schema_type: Final = root.get("type")
-    return any(keyword in root for keyword in _VERTEX_PARAMETERS_UNSUPPORTED_ROOT_KEYWORDS) or (
-        schema_type is not None and not (isinstance(schema_type, str) and schema_type.lower() == "object")
+    defs: Final = parameters.get("$defs")
+    return _root_cannot_be_vertex_parameters(
+        _resolve_root_ref(parameters, defs if isinstance(defs, dict) else {}, frozenset())
     )
 
 
+def parameters_json_schema_for(parameters: Mapping[str, object]) -> Mapping[str, JsonValue] | None:
+    if "$ref" not in parameters and not _root_cannot_be_vertex_parameters(parameters):
+        return None
+    json_schema: Final = _JSON_SCHEMA.validate_python(to_jsonable_python(parameters))
+    return json_schema if needs_parameters_json_schema(json_schema) else None
+
+
 def build_json_schema(parameters: dict) -> dict:
     """
     Build a JSON Schema for use with Gemini's responseJsonSchema parameter.

litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py

@@ -100,7 +100,7 @@ from ..common_utils import (
     VertexAIError,
     build_json_schema,
     build_vertex_schema,
-    needs_parameters_json_schema,
+    parameters_json_schema_for,
     supports_response_json_schema,
     with_integer_enums_that_fit,
 )
@@ -513,7 +513,7 @@ def mapped_gemini_speech_config(
 
 
 class _CallerFunction(TypedDict, total=False):
-    parameters: ReadOnly[Mapping[str, JsonValue]]
+    parameters: ReadOnly[Mapping[str, object]]
 
 
 class _CallerTool(_CallerFunction, total=False):
@@ -524,13 +524,17 @@ _CALLER_TOOL: Final = TypeAdapter(_CallerTool)
 _CALLER_TOOLS: Final = TypeAdapter(tuple[object, ...])
 
 
-def _parameters_json_schema(tool: object) -> Mapping[str, JsonValue] | None:
+def _caller_parameters(tool: object) -> Mapping[str, object] | None:
     try:
         caller_tool: Final = _CALLER_TOOL.validate_python(tool)
     except ValidationError:
         return None
-    parameters: Final = caller_tool.get("function", caller_tool).get("parameters")
-    return parameters if parameters is not None and needs_parameters_json_schema(parameters) else None
+    return caller_tool.get("function", caller_tool).get("parameters")
+
+
+def _parameters_json_schema(tool: object) -> Mapping[str, JsonValue] | None:
+    parameters: Final = _caller_parameters(tool)
+    return None if parameters is None else parameters_json_schema_for(parameters)
 
 
 def _parameters_json_schemas(tools: object) -> tuple[Mapping[str, JsonValue] | None, ...]:

tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py

@@ -2592,6 +2592,10 @@ _WEATHER_VERTEX_PARAMETERS: Final = {"type": "object", "properties": {"city": {"
             ],
         },
         {"$ref": "#/$defs/Lookup", "$defs": {"Lookup": {"oneOf": [_LOOKUP_BY_A, _LOOKUP_BY_B]}}},
+        {
+            "$ref": "#/$defs/Lookup",
+            "$defs": {"Lookup": {"$ref": "#/$defs/Union"}, "Union": {"anyOf": [_LOOKUP_BY_A, _LOOKUP_BY_B]}},
+        },
         {"allOf": [{"type": "object", "properties": _LOOKUP_PROPERTIES}, {"required": ["b"]}]},
     ],
     ids=[
@@ -2600,12 +2604,13 @@ _WEATHER_VERTEX_PARAMETERS: Final = {"type": "object", "properties": {"city": {"
         "root-object-branches",
         "object-with-property-branches",
         "root-ref-to-union",
+        "chained-root-ref-to-union",
         "root-all-of",
     ],
 )
 @pytest.mark.respx(assert_all_called=True)
 def test_gemini_sends_a_union_root_tool_schema_unchanged_as_parameters_json_schema(
-    respx_mock: respx.MockRouter, parameters: dict[str, object]
+    respx_mock: respx.MockRouter, parameters: Mapping[str, object]
 ) -> None:
     route: Final = respx_mock.post(
         url__regex=r"https://generativelanguage\.googleapis\.com/v1beta/models/gemini-2\.5-flash:generateContent.*"
@@ -2640,6 +2645,29 @@ def test_gemini_sends_a_union_root_tool_schema_unchanged_as_parameters_json_sche
     ]
 
 
+def test_vertex_sends_a_union_root_tool_schema_built_with_tuples_as_parameters_json_schema() -> None:
+    request_params: Final = VertexGeminiConfig().map_openai_params(
+        non_default_params={
+            "tools": [
+                {
+                    "type": "function",
+                    "function": {
+                        "name": "lookup",
+                        "parameters": {"anyOf": ({**_LOOKUP_BY_A, "required": ("a",)}, _LOOKUP_BY_B)},
+                    },
+                }
+            ]
+        },
+        optional_params={},
+        model="gemini-2.5-flash",
+        drop_params=False,
+    )
+
+    assert request_params["tools"] == [
+        {"function_declarations": [{"name": "lookup", "parametersJsonSchema": {"anyOf": [_LOOKUP_BY_A, _LOOKUP_BY_B]}}]}
+    ]
+
+
 def test_vertex_sends_a_union_root_legacy_function_schema_unchanged_as_parameters_json_schema() -> None:
     parameters: Final = {"anyOf": [{**_LOOKUP_BY_A, "additionalProperties": False}, _LOOKUP_BY_B]}
     request_params: Final = VertexGeminiConfig().map_openai_params(

tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py

@@ -1,6 +1,6 @@
 import copy
 import time
-from collections.abc import Iterator
+from collections.abc import Iterator, Mapping
 from contextlib import contextmanager
 from enum import Enum
 from typing import Final
@@ -1644,6 +1644,11 @@ _OBJECT_BRANCH: Final = {"type": "object", "properties": {"a": {"type": "string"
         ({"type": "string"}, True),
         ({"$ref": "#/$defs/Lookup", "$defs": {"Lookup": {"anyOf": [_OBJECT_BRANCH]}}}, True),
         ({"$ref": "#/$defs/Lookup", "$defs": {"Lookup": {"type": ["object", "null"]}}}, True),
+        ({"type": ["object"], "properties": {"a": {"type": "string"}}}, False),
+        ({"$ref": "#/$defs/Lookup", "$defs": {"Lookup": {"type": ["OBJECT"]}}}, False),
+        ({"$ref": "#/$defs/A", "$defs": {"A": {"$ref": "#/$defs/B"}, "B": {"anyOf": [_OBJECT_BRANCH]}}}, True),
+        ({"$ref": "#/$defs/A", "$defs": {"A": {"$ref": "#/$defs/B"}, "B": _OBJECT_BRANCH}}, False),
+        ({"$ref": "#/$defs/A", "$defs": {"A": {"$ref": "#/$defs/B"}, "B": {"$ref": "#/$defs/A"}}}, False),
     ],
     ids=[
         "object",
@@ -1661,9 +1666,14 @@ _OBJECT_BRANCH: Final = {"type": "object", "properties": {"a": {"type": "string"
         "non-object-type",
         "root-ref-to-union",
         "root-ref-to-type-list",
+        "single-object-type-list",
+        "root-ref-to-single-object-type-list",
+        "chained-root-ref-to-union",
+        "chained-root-ref-to-object",
+        "cyclic-root-ref",
     ],
 )
-def test_needs_parameters_json_schema(parameters: dict[str, JsonValue], expected: bool) -> None:
+def test_needs_parameters_json_schema(parameters: Mapping[str, JsonValue], expected: bool) -> None:
     assert needs_parameters_json_schema(parameters) is expected