From 958e939c531c782e58fcad491490ef594cbcebbb Mon Sep 17 00:00:00 2001 From: Debanshu Das Date: Thu, 1 Oct 2026 08:04:06 +0000 Subject: [PATCH] fix: merge top-level object unions in Vertex tool parameter schemas Signed-off-by: Debanshu Das --- litellm/llms/gemini/chat/transformation.py | 3 + litellm/llms/vertex_ai/common_utils.py | 87 ++++++++- .../vertex_and_google_ai_studio_gemini.py | 8 +- .../test_vertex_tool_parameters_union_root.py | 184 ++++++++++++++++++ 4 files changed, 277 insertions(+), 5 deletions(-) create mode 100644 tests/unit/llms/vertex_ai/test_vertex_tool_parameters_union_root.py diff --git a/litellm/llms/gemini/chat/transformation.py b/litellm/llms/gemini/chat/transformation.py index 42c9ef13730..f1130e14c48 100644 --- a/litellm/llms/gemini/chat/transformation.py +++ b/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 diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 14aebcaabaf..56f1f598e1c 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/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 @@ -647,7 +649,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 @@ -657,6 +734,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 """ @@ -690,12 +768,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: 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 36b5f2fb5e8..b6e3a64e351 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 @@ -242,6 +242,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 @@ -624,7 +627,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 diff --git a/tests/unit/llms/vertex_ai/test_vertex_tool_parameters_union_root.py b/tests/unit/llms/vertex_ai/test_vertex_tool_parameters_union_root.py new file mode 100644 index 00000000000..090adbc476e --- /dev/null +++ b/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