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

@@ -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)

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(

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
 

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)

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()

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."""
 

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