From 4fa1bd894da2f1ff2fa6efebd1432ac958dda31f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 11:05:51 -0700 Subject: [PATCH] fix(vertex-ai): send union-root tool schemas as parametersJsonSchema A function tool whose parameters root is a union (anyOf, oneOf or allOf at the root, a type list, or a root $ref to one of those) cannot be expressed in the Vertex and Gemini parameters field, so those tools answered 400. They now go out as parametersJsonSchema with the caller's JSON Schema unchanged, on both the tools and the legacy functions paths. Every other tool keeps the converted parameters it had before. The local token counter also reads parametersJsonSchema, so these tools still count --- litellm/litellm_core_utils/token_counter.py | 2 +- litellm/llms/vertex_ai/common_utils.py | 20 ++++ .../vertex_and_google_ai_studio_gemini.py | 46 +++++++- litellm/types/llms/vertex_ai.py | 4 +- .../litellm_core_utils/test_token_counter.py | 14 +++ ...test_vertex_and_google_ai_studio_gemini.py | 108 ++++++++++++++++++ .../vertex_ai/test_vertex_ai_common_utils.py | 45 ++++++++ 7 files changed, 231 insertions(+), 8 deletions(-) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 800a0939cec..38de7be20ee 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -1830,7 +1830,7 @@ def _format_single_function_definition(function: Mapping[str, object]) -> tuple[ if not function_name: return () function_description: Final = function.get("description") - parameters_value: Final = function.get("parameters") or {} + parameters_value: Final = function.get("parameters") or function.get("parametersJsonSchema") or {} parameters: Final = parameters_value if isinstance(parameters_value, Mapping) else {} properties: Final = parameters.get("properties") if isinstance(properties, Mapping) and properties: diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 0fa80cb2cad..21a80e39f26 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -889,6 +889,26 @@ def build_vertex_schema( _build_vertex_schema = build_vertex_schema +_VERTEX_PARAMETERS_UNSUPPORTED_ROOT_KEYWORDS: Final = ("anyOf", "oneOf", "allOf") +_ROOT_DEFS_REF_PREFIX: Final = "#/$defs/" + + +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 + case _: + return parameters + + +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") + ) + 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 1ef5267df14..ac5c84b6414 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 @@ -12,7 +12,8 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, Union, cast, get_args import httpx -from pydantic import JsonValue, TypeAdapter +from pydantic import JsonValue, TypeAdapter, ValidationError +from typing_extensions import ReadOnly, TypedDict import litellm from litellm import verbose_logger @@ -99,6 +100,7 @@ from ..common_utils import ( VertexAIError, build_json_schema, build_vertex_schema, + needs_parameters_json_schema, supports_response_json_schema, with_integer_enums_that_fit, ) @@ -510,6 +512,31 @@ def mapped_gemini_speech_config( } +class _CallerFunction(TypedDict, total=False): + parameters: ReadOnly[Mapping[str, JsonValue]] + + +class _CallerTool(_CallerFunction, total=False): + function: ReadOnly[_CallerFunction] + + +_CALLER_TOOL: Final = TypeAdapter(_CallerTool) +_CALLER_TOOLS: Final = TypeAdapter(tuple[object, ...]) + + +def _parameters_json_schema(tool: object) -> Mapping[str, JsonValue] | 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 + + +def _parameters_json_schemas(tools: object) -> tuple[Mapping[str, JsonValue] | None, ...]: + return tuple(_parameters_json_schema(tool) for tool in _CALLER_TOOLS.validate_python(tools)) + + class VertexAIBaseConfig: def get_mapped_special_auth_params(self) -> dict: """ @@ -1059,17 +1086,20 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): computerUse: dict | None = None # remove 'additionalProperties' from tools schema_value: Final[JsonValue] = deepcopy(cast(JsonValue, value)) # cast-ok: tools are JSON request data + parameters_json_schemas: Final = _parameters_json_schemas(schema_value) remove_additional_properties(schema_value) # remove 'strict' from tools remove_strict_from_schema(schema_value) - for tool in cast(list[dict], schema_value): # cast-ok: a deep copy keeps the caller's tool shape + for index, parameters_json_schema in enumerate(parameters_json_schemas): + tool = cast(list[dict], schema_value)[index] # cast-ok: a deep copy keeps the caller's tool shape openai_function_object: ChatCompletionToolParamFunctionChunk | None = None if "function" in tool: # tools list _openai_function_object = ChatCompletionToolParamFunctionChunk(**tool["function"]) if ( - "parameters" in _openai_function_object + parameters_json_schema is None + and "parameters" in _openai_function_object and _openai_function_object["parameters"] is not None and isinstance(_openai_function_object["parameters"], dict) ): # OPENAI accepts JSON Schema, Google accepts OpenAPI schema. @@ -1132,8 +1162,12 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): # Empty config - Gemini will use defaults computerUse = {} elif openai_function_object is not None: - gtool_func_declaration = FunctionDeclaration( - name=openai_function_object["name"], + gtool_func_declaration = ( + FunctionDeclaration(name=openai_function_object["name"]) + if parameters_json_schema is None + else FunctionDeclaration( + name=openai_function_object["name"], parametersJsonSchema=parameters_json_schema + ) ) _description = openai_function_object.get("description", None) _parameters = openai_function_object.get("parameters", None) @@ -1143,7 +1177,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): } if _description is not None: gtool_func_declaration["description"] = _description - if _parameters is not None: + if _parameters is not None and parameters_json_schema is None: gtool_func_declaration["parameters"] = _parameters gtool_func_declarations.append(gtool_func_declaration) else: diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index a938f6a7b1c..ec76e271964 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -1,7 +1,8 @@ -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from enum import Enum from typing import Final, Literal, Protocol +from pydantic import JsonValue from typing_extensions import ( ReadOnly, Required, @@ -163,6 +164,7 @@ class FunctionDeclaration(TypedDict, total=False): name: Required[str] description: str parameters: Schema | dict + parametersJsonSchema: ReadOnly[Mapping[str, JsonValue]] response: Schema diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index 697ac76b4e5..afd3a74aefb 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -693,6 +693,20 @@ def test_token_counter_counts_function_tools_like_chat_completions_tools_for_cha ) == token_counter_new(model=model, messages=list(_USER_MESSAGES), tools=[_openai_tool(function)]) +def test_token_counter_counts_a_gemini_parameters_json_schema_like_parameters(): + json_schema_function: Final = { + "name": _WEATHER_FUNCTION["name"], + "description": _WEATHER_FUNCTION["description"], + "parametersJsonSchema": _WEATHER_FUNCTION["parameters"], + } + + assert token_counter_new( + model="gemini-2.5-pro", messages=list(_USER_MESSAGES), tools=[_gemini_tool(json_schema_function)] + ) == token_counter_new( + model="gemini-2.5-pro", messages=list(_USER_MESSAGES), tools=[_gemini_tool(_WEATHER_FUNCTION)] + ) + + # Azure AI Foundry DeepSeek-V4-Flash-0731 billed 295 prompt tokens for the weather tool with each of # these tool_choice values on 2026-10-17, the same as without one @pytest.mark.parametrize("tool_choice", ["none", "required", {"type": "function", "function": {"name": "get_weather"}}]) diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 98d1b8f1300..944dd400e9b 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -2556,6 +2556,114 @@ def test_gemini_reasoning_effort_and_function_tool_share_request() -> None: ] +_LOOKUP_PROPERTIES: Final = {"a": {"type": "string"}, "b": {"type": "string"}} +_LOOKUP_BY_A: Final = {"type": "object", "properties": {"a": {"type": "string"}}, "required": ["a"]} +_LOOKUP_BY_B: Final = {"type": "object", "properties": {"b": {"type": "string"}}, "required": ["b"]} +_WEATHER_PARAMETERS: Final = { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + "additionalProperties": False, +} +_WEATHER_VERTEX_PARAMETERS: Final = {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]} + + +@pytest.mark.parametrize( + "parameters", + [ + { + "type": "object", + "properties": _LOOKUP_PROPERTIES, + "anyOf": [{"required": ["a"]}, {"required": ["b"]}], + "additionalProperties": False, + }, + { + "type": ["object", "null"], + "properties": _LOOKUP_PROPERTIES, + "anyOf": [{"required": ["a"]}, {"required": ["b"]}], + }, + {"anyOf": [{**_LOOKUP_BY_A, "additionalProperties": False}, _LOOKUP_BY_B]}, + { + "type": "object", + "properties": {"limit": {"type": "integer"}}, + "anyOf": [ + {"properties": {"query": {"type": "string"}}, "required": ["query"]}, + {"properties": {"url": {"type": "string"}}, "required": ["url"]}, + ], + }, + {"$ref": "#/$defs/Lookup", "$defs": {"Lookup": {"oneOf": [_LOOKUP_BY_A, _LOOKUP_BY_B]}}}, + {"allOf": [{"type": "object", "properties": _LOOKUP_PROPERTIES}, {"required": ["b"]}]}, + ], + ids=[ + "object-with-required-branches", + "nullable-object-with-required-branches", + "root-object-branches", + "object-with-property-branches", + "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] +) -> None: + route: Final = respx_mock.post( + url__regex=r"https://generativelanguage\.googleapis\.com/v1beta/models/gemini-2\.5-flash:generateContent.*" + ).mock( + return_value=httpx.Response( + 200, + json={ + "candidates": [{"content": {"parts": [{"text": "Ready"}]}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2}, + }, + ) + ) + completion( + model="gemini/gemini-2.5-flash", + messages=[{"role": "user", "content": "Look up the record named Atlas"}], + tools=[ + {"type": "function", "function": {"name": "lookup", "parameters": deepcopy(parameters)}}, + {"type": "function", "function": {"name": "weather", "parameters": deepcopy(_WEATHER_PARAMETERS)}}, + ], + api_key="test-api-key", + client=HTTPHandler(), + ) + request_body: Final = json.loads(route.calls.last.request.content) + + assert request_body["tools"] == [ + { + "function_declarations": [ + {"name": "lookup", "parametersJsonSchema": parameters}, + {"name": "weather", "parameters": _WEATHER_VERTEX_PARAMETERS}, + ] + } + ] + + +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( + non_default_params={ + "functions": [ + {"name": "lookup", "description": "Look up a record", "parameters": deepcopy(parameters)}, + {"name": "weather", "parameters": deepcopy(_WEATHER_PARAMETERS)}, + ] + }, + optional_params={}, + model="gemini-2.5-flash", + drop_params=False, + ) + + assert request_params["tools"] == [ + { + "function_declarations": [ + {"name": "lookup", "description": "Look up a record", "parametersJsonSchema": parameters}, + {"name": "weather", "parameters": _WEATHER_VERTEX_PARAMETERS}, + ] + } + ] + + @pytest.mark.parametrize( ("model", "expected_level"), [ 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..83d2c8d63fd 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 @@ -23,6 +23,7 @@ from litellm.llms.vertex_ai.common_utils import ( get_vertex_location_from_url, get_vertex_project_id_from_url, get_vertex_url, + needs_parameters_json_schema, pop_vertex_request_labels, set_schema_property_ordering, supports_response_json_schema, @@ -1622,6 +1623,50 @@ def test_map_function_keeps_integer_enum_when_tools_are_reused(): assert sent[0] == sent[1] +_OBJECT_BRANCH: Final = {"type": "object", "properties": {"a": {"type": "string"}}, "required": ["a"]} + + +@pytest.mark.parametrize( + ("parameters", "expected"), + [ + ({"type": "object", "properties": {"a": {"type": "string"}}}, False), + ({"type": "OBJECT", "properties": {"a": {"type": "string"}}}, False), + ({"properties": {"a": {"type": "string"}}}, False), + ({"type": "object", "properties": {"a": {"anyOf": [{"type": "string"}, {"type": "null"}]}}}, False), + ({"$ref": "#/$defs/Lookup", "$defs": {"Lookup": _OBJECT_BRANCH}}, False), + ({"$ref": "#/$defs/Missing", "$defs": {"Lookup": {"anyOf": [_OBJECT_BRANCH]}}}, False), + ({"$ref": "#/definitions/Lookup", "definitions": {"Lookup": {"anyOf": [_OBJECT_BRANCH]}}}, False), + ({"type": "object", "anyOf": [{"required": ["a"]}]}, True), + ({"anyOf": [_OBJECT_BRANCH]}, True), + ({"oneOf": [_OBJECT_BRANCH]}, True), + ({"allOf": [_OBJECT_BRANCH]}, True), + ({"type": ["object", "null"], "properties": {"a": {"type": "string"}}}, True), + ({"type": "string"}, True), + ({"$ref": "#/$defs/Lookup", "$defs": {"Lookup": {"anyOf": [_OBJECT_BRANCH]}}}, True), + ({"$ref": "#/$defs/Lookup", "$defs": {"Lookup": {"type": ["object", "null"]}}}, True), + ], + ids=[ + "object", + "uppercase-object", + "untyped-properties", + "union-inside-a-property", + "root-ref-to-object", + "root-ref-to-missing-def", + "root-ref-into-definitions", + "object-with-any-of", + "root-any-of", + "root-one-of", + "root-all-of", + "type-list", + "non-object-type", + "root-ref-to-union", + "root-ref-to-type-list", + ], +) +def test_needs_parameters_json_schema(parameters: dict[str, JsonValue], expected: bool) -> None: + assert needs_parameters_json_schema(parameters) is expected + + class _Priority(int, Enum): LOW = 1 HIGH = 3