Merge 958e939c53 into 169af2f883

litellm/llms/gemini/chat/transformation.py

@@ -77,6 +77,9 @@ class GoogleAIStudioGeminiConfig(VertexGeminiConfig):
     def get_config(cls):
         return super().get_config()
 
+    def _should_enforce_tool_parameters_object_root(self) -> bool:
+        return False
+
     def is_model_gemini_audio_model(self, model: str) -> bool:
         return "tts" in model
 

litellm/llms/vertex_ai/common_utils.py

@@ -1,3 +1,5 @@
+import itertools
+import json
 import re
 from collections.abc import Mapping
 from copy import deepcopy
@@ -683,7 +685,82 @@ def _fix_enum_types(schema, depth=0):
                 _fix_enum_types(item, depth=depth + 1)
 
 
-def _build_vertex_schema(parameters: dict, add_property_ordering: bool = False):
+_MAX_UNION_MERGE_BRANCHES: Final = 64
+_MAX_UNION_MERGE_PROPERTIES: Final = 256
+_MAX_UNION_MERGE_VARIANTS: Final = 512
+
+
+def _extract_union_variants(definition: object) -> tuple[object, ...]:
+    if isinstance(definition, dict) and tuple(definition.keys()) == ("anyOf",):
+        variants: Final = definition.get("anyOf")
+        if isinstance(variants, list) and variants:
+            return tuple(variants)
+    return (definition,)
+
+
+def _merge_property_definitions(definitions: tuple[object, ...]) -> object:
+    expanded: Final = tuple(
+        itertools.chain.from_iterable(_extract_union_variants(definition) for definition in definitions)
+    )
+    seen_keys: Final[set[str]] = set()  # mutable-ok: O(N) deduplication key set
+    unique: Final = tuple(
+        variant
+        for key, variant in ((json.dumps(v, sort_keys=True, default=str), v) for v in expanded)
+        if key not in seen_keys and not seen_keys.add(key)
+    )
+    if len(unique) == 1:
+        return unique[0]
+    return {"anyOf": list(unique)}
+
+
+def _merge_object_union_root(
+    schema: dict[str, object],  # mutable-ok: Vertex schema must stay a plain dict
+) -> dict[str, object]:  # mutable-ok: Vertex schema must stay a plain dict
+    branches: Final = schema.get("anyOf")
+    if not isinstance(branches, list) or not branches or len(branches) > _MAX_UNION_MERGE_BRANCHES:
+        return schema
+
+    if not all(isinstance(branch, dict) and branch.get("type") == "object" for branch in branches):
+        return schema
+
+    property_sources: Final = tuple(
+        branch["properties"] for branch in branches if isinstance(branch.get("properties"), dict)
+    )
+    grouped_defs: Final[dict[str, list[object]]] = {}  # mutable-ok: single-pass O(N) property grouping
+    for source in property_sources:
+        for name, defn in source.items():
+            grouped_defs.setdefault(name, []).append(defn)
+
+    all_defs: Final = tuple(itertools.chain.from_iterable(grouped_defs.values()))
+    total_variants: Final = sum(len(_extract_union_variants(defn)) for defn in all_defs)
+    if len(grouped_defs) > _MAX_UNION_MERGE_PROPERTIES or total_variants > _MAX_UNION_MERGE_VARIANTS:
+        return schema
+
+    properties: Final = {name: _merge_property_definitions(tuple(defs)) for name, defs in grouped_defs.items()}
+    required_sets: Final = tuple(
+        frozenset(branch["required"]) if isinstance(branch.get("required"), list) else frozenset()
+        for branch in branches
+    )
+    shared_required: Final = frozenset.intersection(*required_sets)
+    carried: Final = {key: value for key, value in schema.items() if key not in ("anyOf", "required")}
+    merged: Final = {
+        **carried,
+        "type": "object",
+        "properties": properties,
+    }
+    ordered_required: Final = [name for name in properties if name in shared_required]
+
+    if not ordered_required:
+        return merged
+
+    return {**merged, "required": ordered_required}
+
+
+def _build_vertex_schema(
+    parameters: dict,
+    add_property_ordering: bool = False,
+    enforce_object_root: 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
 
@@ -693,6 +770,7 @@ def _build_vertex_schema(parameters: dict, add_property_ordering: bool = False):
         parameters: dict - the json schema to build from
         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.
+        enforce_object_root: bool - whether to merge a root union of object branches into a single object.
     Returns:
         parameters: dict - the input parameters, modified in place
     """
@@ -726,12 +804,13 @@ def _build_vertex_schema(parameters: dict, add_property_ordering: bool = False):
     # Postprocessing
     # Filter out fields that don't exist in Schema
 
-    parameters = filter_schema_fields(parameters, valid_schema_fields)
+    filtered: Final = filter_schema_fields(parameters, valid_schema_fields)
+    rooted: Final = _merge_object_union_root(filtered) if enforce_object_root else filtered
 
     if add_property_ordering:
-        set_schema_property_ordering(parameters)
+        set_schema_property_ordering(rooted)
 
-    return parameters
+    return rooted
 
 
 def _build_json_schema(parameters: dict) -> dict:

litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py

@@ -244,6 +244,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
     def get_config(cls):
         return super().get_config()
 
+    def _should_enforce_tool_parameters_object_root(self) -> bool:
+        return True
+
     def get_json_schema_from_pydantic_object(self, response_format: type["BaseModel"] | dict | None) -> dict | None:
         """
         Override to use Pydantic's model_json_schema() instead of OpenAI's
@@ -625,7 +628,10 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
                     and _openai_function_object["parameters"] is not None
                     and isinstance(_openai_function_object["parameters"], dict)
                 ):  # OPENAI accepts JSON Schema, Google accepts OpenAPI schema.
-                    _openai_function_object["parameters"] = _build_vertex_schema(_openai_function_object["parameters"])
+                    _openai_function_object["parameters"] = _build_vertex_schema(
+                        _openai_function_object["parameters"],
+                        enforce_object_root=self._should_enforce_tool_parameters_object_root(),
+                    )
 
                 openai_function_object = _openai_function_object
 

tests/unit/llms/vertex_ai/test_vertex_tool_parameters_union_root.py

@@ -0,0 +1,184 @@
+from typing import Final
+
+from litellm.llms.gemini.chat.transformation import GoogleAIStudioGeminiConfig
+from litellm.llms.vertex_ai.common_utils import _build_vertex_schema
+from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
+    VertexGeminiConfig,
+)
+
+
+def test_build_vertex_schema_merges_object_union_root_for_tool_parameters():
+    schema: Final = {
+        "anyOf": [
+            {"type": "object", "properties": {"query": {"type": "string"}}},
+            {"type": "object", "properties": {"url": {"type": "string"}}},
+        ]
+    }
+
+    result: Final = _build_vertex_schema(schema, enforce_object_root=True)
+
+    assert result["type"] == "object"
+    assert "anyOf" not in result
+    assert set(result["properties"]) == {"query", "url"}
+
+
+def test_build_vertex_schema_combines_conflicting_branch_property_schemas():
+    schema: Final = {
+        "anyOf": [
+            {
+                "type": "object",
+                "properties": {
+                    "target": {"type": "string"},
+                    "lang": {"type": "string"},
+                },
+            },
+            {
+                "type": "object",
+                "properties": {
+                    "target": {"type": "integer"},
+                    "lang": {"type": "string"},
+                },
+            },
+        ]
+    }
+
+    result: Final = _build_vertex_schema(schema, enforce_object_root=True)
+
+    assert result["properties"]["lang"] == {"type": "string"}
+    assert result["properties"]["target"] == {"anyOf": [{"type": "string"}, {"type": "integer"}]}
+
+
+def test_build_vertex_schema_keeps_only_required_fields_shared_by_every_branch():
+    schema: Final = {
+        "anyOf": [
+            {
+                "type": "object",
+                "properties": {
+                    "query": {"type": "string"},
+                    "lang": {"type": "string"},
+                },
+                "required": ["query", "lang"],
+            },
+            {
+                "type": "object",
+                "properties": {
+                    "url": {"type": "string"},
+                    "lang": {"type": "string"},
+                },
+                "required": ["url", "lang"],
+            },
+        ]
+    }
+
+    result: Final = _build_vertex_schema(schema, enforce_object_root=True)
+
+    assert result["required"] == ["lang"]
+
+
+def test_build_vertex_schema_leaves_non_object_union_root_alone():
+    schema: Final = {
+        "anyOf": [
+            {"type": "string"},
+            {"type": "object", "properties": {"url": {"type": "string"}}},
+        ]
+    }
+
+    result: Final = _build_vertex_schema(schema, enforce_object_root=True)
+
+    assert "anyOf" in result
+
+
+def test_build_vertex_schema_leaves_response_schema_root_untouched():
+    schema: Final = {"anyOf": [{"type": "object", "properties": {"query": {"type": "string"}}}]}
+
+    result: Final = _build_vertex_schema(schema)
+
+    assert "type" not in result
+
+
+def test_build_vertex_schema_still_strips_fields_beside_nested_anyof():
+    schema: Final = {
+        "type": "object",
+        "properties": {
+            "value": {
+                "anyOf": [{"type": "string"}, {"type": "null"}],
+                "default": "x",
+            }
+        },
+    }
+
+    result: Final = _build_vertex_schema(schema, enforce_object_root=True)
+
+    assert "default" not in result["properties"]["value"]
+
+
+def test_vertex_ai_map_tool_with_top_level_anyof():
+    v: Final = VertexGeminiConfig()
+    value: Final = [
+        {
+            "type": "function",
+            "function": {
+                "name": "search",
+                "description": "Search by query or by url",
+                "parameters": {
+                    "anyOf": [
+                        {
+                            "type": "object",
+                            "properties": {"query": {"type": "string"}},
+                        },
+                        {
+                            "type": "object",
+                            "properties": {"url": {"type": "string"}},
+                        },
+                    ]
+                },
+            },
+        }
+    ]
+
+    tools: Final = v._map_function(value=value, optional_params={})
+
+    parameters: Final = tools[0]["function_declarations"][0]["parameters"]
+    assert parameters["type"] == "object"
+    assert "anyOf" not in parameters
+    assert set(parameters["properties"]) == {"query", "url"}
+
+
+def test_google_ai_studio_map_tool_preserves_top_level_anyof():
+    g: Final = GoogleAIStudioGeminiConfig()
+    value: Final = [
+        {
+            "type": "function",
+            "function": {
+                "name": "search",
+                "description": "Search by query or by url",
+                "parameters": {
+                    "anyOf": [
+                        {
+                            "type": "object",
+                            "properties": {"query": {"type": "string"}},
+                        },
+                        {
+                            "type": "object",
+                            "properties": {"url": {"type": "string"}},
+                        },
+                    ]
+                },
+            },
+        }
+    ]
+
+    tools: Final = g._map_function(value=value, optional_params={})
+
+    parameters: Final = tools[0]["function_declarations"][0]["parameters"]
+    assert "anyOf" in parameters
+    assert "type" not in parameters
+
+
+def test_build_vertex_schema_skips_merge_when_branch_limit_exceeded():
+    schema: Final = {"anyOf": [{"type": "object", "properties": {f"field_{i}": {"type": "string"}}} for i in range(65)]}
+
+    result: Final = _build_vertex_schema(schema, enforce_object_root=True)
+
+    assert "anyOf" in result
+    assert "type" not in result