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