From 2add6a0629cab769b54247dd081ea1de784618dc Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 4 Nov 2026 21:04:55 -0800 Subject: [PATCH] fix(vertex-ai): keep parameters for Vertex batch rows and Vertex countTokens Vertex batch merges every row's keys into one shape, so a parametersJsonSchema tool in one row breaks every parameters row in the file. Vertex countTokens accepts parametersJsonSchema but leaves it out of the count. Both keep sending parameters, and the Gemini API count route keeps parametersJsonSchema, which it counts in full. --- litellm/llms/vertex_ai/common_utils.py | 2 +- .../llms/vertex_ai/count_tokens/context.py | 5 +- .../llms/vertex_ai/files/transformation.py | 1 + .../vertex_and_google_ai_studio_gemini.py | 32 ++++++++++--- .../llms/gemini/test_gemini_common_utils.py | 31 +++++++++++++ .../test_vertex_ai_files_transformation.py | 46 +++++++++++++++++++ .../vertex_ai/test_vertex_ai_common_utils.py | 31 +++++++++++++ 7 files changed, 140 insertions(+), 8 deletions(-) diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index f882a54ae72..faa21c812e6 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -1636,7 +1636,7 @@ class VertexAITokenCounter(BaseTokenCounter): if isinstance(request_view, VertexGeminiCountContextError): return count_context_error_response(request_view, request_model, model_to_use) count_context: Final = vertex_gemini_count_context( - system=request_view.system, tools=tools, model=model_to_use + system=request_view.system, tools=tools, model=model_to_use, allow_parameters_json_schema=False ) if isinstance(count_context, VertexGeminiCountContextError): return count_context_error_response(count_context, request_model, model_to_use) diff --git a/litellm/llms/vertex_ai/count_tokens/context.py b/litellm/llms/vertex_ai/count_tokens/context.py index fc7d6545624..eb29add17db 100644 --- a/litellm/llms/vertex_ai/count_tokens/context.py +++ b/litellm/llms/vertex_ai/count_tokens/context.py @@ -79,6 +79,7 @@ def _is_function_tool(tool: Mapping[str, JsonValue]) -> bool: def _map_function_tools( tools: object | None, model: str, + allow_parameters_json_schema: bool, ) -> tuple[Mapping[str, object], ...] | None | VertexGeminiCountContextError: try: tool_values: Final = _TOOLS.validate_python(tools or ()) @@ -103,6 +104,7 @@ def _map_function_tools( value=[dict(tool) for tool in chat_tools], optional_params={}, send_integer_enums=send_integer_enums, + allow_parameters_json_schema=allow_parameters_json_schema, ), VertexGeminiConfig.is_gemini_3_or_newer(model), ) @@ -118,11 +120,12 @@ def vertex_gemini_count_context( system: object | None, tools: object | None, model: str, + allow_parameters_json_schema: bool = True, ) -> VertexGeminiCountContext | VertexGeminiCountContextError: system_parts: Final = _system_parts(system) if isinstance(system_parts, VertexGeminiCountContextError): return system_parts - function_tools: Final = _map_function_tools(tools, model) + function_tools: Final = _map_function_tools(tools, model, allow_parameters_json_schema) if isinstance(function_tools, VertexGeminiCountContextError): return function_tools return VertexGeminiCountContext( diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 9159c9a4070..79aa7da3bec 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -1321,6 +1321,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): optional_params={}, drop_params=drops_unsupported_params(openai_request_body.get("drop_params", deployment_drop_params)), allow_integer_enums=False, + allow_parameters_json_schema=False, ) return vertex_params 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 639df5f222e..af97bf75d5a 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 @@ -537,8 +537,13 @@ def _parameters_json_schema(tool: object) -> Mapping[str, JsonValue] | None: return None if parameters is None else parameters_json_schema_for(parameters) -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)) +def _parameters_json_schemas( + tools: object, allow_parameters_json_schema: bool +) -> tuple[Mapping[str, JsonValue] | None, ...]: + return tuple( + _parameters_json_schema(tool) if allow_parameters_json_schema else None + for tool in _CALLER_TOOLS.validate_python(tools) + ) class VertexAIBaseConfig: @@ -1063,7 +1068,13 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): return googleSearch, googleSearchRetrieval, enterpriseWebSearch, urlContext - def _map_function(self, value: list[dict], optional_params: dict, send_integer_enums: bool = False) -> list[Tools]: + def _map_function( + self, + value: list[dict], + optional_params: dict, + send_integer_enums: bool = False, + allow_parameters_json_schema: bool = True, + ) -> list[Tools]: """ Map OpenAI-style tools/functions to Vertex AI format. @@ -1071,6 +1082,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): value: List of tool definitions optional_params: Request-scoped parameters to store retrieval config send_integer_enums: Whether to send integer enums with string values + allow_parameters_json_schema: Whether a tool whose parameters root Vertex cannot take as `parameters` + goes as `parametersJsonSchema` Returns: List of mapped tools in Vertex AI format @@ -1090,7 +1103,7 @@ 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) + parameters_json_schemas: Final = _parameters_json_schemas(schema_value, allow_parameters_json_schema) remove_additional_properties(schema_value) # remove 'strict' from tools remove_strict_from_schema(schema_value) @@ -1257,8 +1270,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): value: list[dict[str, object]], # mutable-ok: mirrors override contract optional_params: dict[str, object], # mutable-ok: mirrors override contract send_integer_enums: bool = False, + allow_parameters_json_schema: bool = True, ) -> list[Tools]: # mutable-ok: mirrors override contract - return self._map_function(value, optional_params, send_integer_enums) + return self._map_function(value, optional_params, send_integer_enums, allow_parameters_json_schema) def _map_response_schema(self, value: dict, send_integer_enums: bool = False) -> Mapping[str, object]: old_schema = deepcopy(value) @@ -1677,6 +1691,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): model: str, drop_params: bool, allow_integer_enums: bool = True, + allow_parameters_json_schema: bool = True, ) -> dict: self._apply_include_server_side_tool_invocations(non_default_params, optional_params) supports_sampling_params: Final = self.supports_sampling_params(model) @@ -1732,7 +1747,12 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): elif (param == "tools" or param == "functions") and isinstance(value, list) and value: # Pass optional_params so _map_function can add toolConfig if needed mapped_tools = with_integer_enums_that_fit( - partial(self._map_function, value, optional_params), + partial( + self._map_function, + value, + optional_params, + allow_parameters_json_schema=allow_parameters_json_schema, + ), allow_integer_enums and self._is_gemini_3_or_newer(model), ) optional_params = self._add_tools_to_optional_params(optional_params, mapped_tools) diff --git a/tests/unit/llms/gemini/test_gemini_common_utils.py b/tests/unit/llms/gemini/test_gemini_common_utils.py index 5f2af34e513..42ada8ea246 100644 --- a/tests/unit/llms/gemini/test_gemini_common_utils.py +++ b/tests/unit/llms/gemini/test_gemini_common_utils.py @@ -80,6 +80,37 @@ async def test_google_ai_studio_count_tokens_forwards_system_and_tools() -> None } +@pytest.mark.asyncio +async def test_google_ai_studio_count_tokens_sends_a_union_root_tool_as_parameters_json_schema() -> None: + union_root: Final = { + "type": "object", + "properties": {"a": {"type": "string"}, "b": {"type": "string"}}, + "anyOf": [{"required": ["a"]}, {"required": ["b"]}], + } + peer: Final = _GeminiCountTokensPeer() + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(peer)) + + with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client): + result: Final = await GoogleAIStudioTokenCounter().count_tokens( + model_to_use="gemini-2.5-flash", + messages=None, + contents=[{"role": "user", "parts": [{"text": "Look up Atlas."}]}], + deployment={"litellm_params": {"api_key": "test-key", "api_base": "https://count-tokens.test"}}, + request_model="gemini/gemini-2.5-flash", + tools=[{"name": "lookup", "input_schema": union_root}], + ) + + assert result is not None + assert result.total_tokens == 11 + assert peer.request_body == { + "generateContentRequest": { + "model": "models/gemini-2.5-flash", + "contents": [{"role": "user", "parts": [{"text": "Look up Atlas."}]}], + "tools": [{"functionDeclarations": [{"name": "lookup", "parametersJsonSchema": union_root}]}], + } + } + + @pytest.mark.asyncio async def test_google_ai_studio_count_tokens_rejects_image_system_without_request() -> None: peer: Final = _GeminiCountTokensPeer() diff --git a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py index 98e204dd44e..8c5ddf81bbd 100644 --- a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py +++ b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py @@ -1692,6 +1692,52 @@ class TestVertexBatchIntegerEnums: assert properties["mode"]["enum"] == ["fast", "slow"] +_FLAT_TOOL_SCHEMA: Final = {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]} + + +def _chat_entry_with_tool(custom_id: str, name: str, parameters: Mapping[str, object]) -> dict[str, object]: + return { + "custom_id": custom_id, + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "gemini-2.5-flash", + "messages": [{"role": "user", "content": "Look it up"}], + "tools": [{"type": "function", "function": {"name": name, "parameters": parameters}}], + }, + } + + +class TestVertexBatchToolParameters: + @pytest.mark.parametrize( + "root_schema", + [ + { + "type": "object", + "properties": {"a": {"type": "string"}, "b": {"type": "string"}}, + "anyOf": [{"required": ["a"]}, {"required": ["b"]}], + }, + {"type": ["object", "null"], "properties": {"a": {"type": "string"}}}, + ], + ids=["anyof-root", "type-list-root"], + ) + def test_rows_never_send_parameters_json_schema_since_vertex_batch_merges_keys_across_rows( + self, root_schema: Mapping[str, object] + ): + rows: Final = _wrap_entries( + [ + _chat_entry_with_tool("root-1", "lookup", root_schema), + _chat_entry_with_tool("flat-1", "get_weather", _FLAT_TOOL_SCHEMA), + ] + ) + + declarations: Final = [row["request"]["tools"][0]["function_declarations"][0] for row in rows] + assert [ + (declaration["name"], "parameters" in declaration, "parametersJsonSchema" in declaration) + for declaration in declarations + ] == [("lookup", True, False), ("get_weather", True, False)] + + class TestVertexEmbeddingsBatchOutputTranslation: """Vertex Gemini Embedding batch output rows must come back as OpenAI batch rows.""" 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 e389fbb4eea..1dc6abce68c 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 @@ -296,6 +296,37 @@ async def test_vertex_gemini_count_applies_the_context_edits_inference_applies() assert peer.requests[0].content.decode().count(_TOOL_OUTPUT) == 1 +@pytest.mark.asyncio +async def test_vertex_gemini_count_sends_a_union_root_tool_as_parameters_since_vertex_counts_only_that_field() -> None: + with _vertex_count_tokens_peer(200) as peer: + result: Final = await VertexAITokenCounter().count_tokens( + model_to_use="gemini-2.5-flash", + messages=[{"role": "user", "content": "Look up Atlas."}], + contents=None, + deployment=_VERTEX_DEPLOYMENT, + request_model="vertex_ai/gemini-2.5-flash", + tools=[ + { + "name": "lookup", + "input_schema": { + "type": "object", + "properties": {"a": {"type": "string"}, "b": {"type": "string"}}, + "anyOf": [{"required": ["a"]}, {"required": ["b"]}], + }, + } + ], + messages_format="anthropic", + ) + + assert result is not None + assert result.total_tokens == 37 + sent_tools: Final = _JSON_BODY.validate_json(peer.requests[0].content)["tools"] + assert isinstance(sent_tools, list) and isinstance(sent_tools[0], dict) + declarations: Final = sent_tools[0]["functionDeclarations"] + assert isinstance(declarations, list) and isinstance(declarations[0], dict) + assert ("parameters" in declarations[0], "parametersJsonSchema" in declarations[0]) == (True, False) + + @pytest.mark.parametrize("location", ["us", "eu", "us-central1", "europe-west1", "global"]) def test_validate_vertex_location_accepts_valid(location): assert validate_vertex_location(location) == location