diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 5fb0a619744..c0116454ce6 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -2095,7 +2095,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 0013e678931..7a163d9b1d0 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -11,6 +11,7 @@ from typing import TYPE_CHECKING, Final, Literal, TypeVar, cast, get_type_hints import httpx from pydantic import JsonValue, TypeAdapter, ValidationError +from pydantic_core import to_jsonable_python from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm @@ -39,7 +40,12 @@ from litellm.types.llms.vertex_ai import ( ) from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import CountContextManagement, TokenCountMessagesFormat, TokenCountResponse -from litellm.utils import supports_response_schema, supports_system_messages +from litellm.utils import ( + remove_additional_properties, + remove_strict_from_schema, + supports_response_schema, + supports_system_messages, +) if TYPE_CHECKING: from litellm.llms.anthropic.count_tokens.transformation import AnthropicCountFields @@ -916,6 +922,132 @@ 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/" +_JSON_SCHEMA: Final = TypeAdapter(Mapping[str, JsonValue]) + + +def _is_absent_or_object_type(schema_type: object) -> bool: + match schema_type: + case None: + return True + case str(): + return schema_type.lower() == "object" + case [str() as only_type]: + return only_type.lower() == "object" + case _: + return False + + +def _root_cannot_be_vertex_parameters(root: Mapping[str, object]) -> bool: + return any( + keyword in root for keyword in _VERTEX_PARAMETERS_UNSUPPORTED_ROOT_KEYWORDS + ) or not _is_absent_or_object_type(root.get("type")) + + +def _root_ref_def_name(schema: Mapping[str, JsonValue], def_names: Mapping[str, str]) -> str | None: + match schema.get("$ref"): + case str() as ref if ref.startswith(_ROOT_DEFS_REF_PREFIX): + return def_names.get(ref.removeprefix(_ROOT_DEFS_REF_PREFIX)) + case _: + return None + + +def _root_ref_chain( + schema: Mapping[str, JsonValue], defs: Mapping[str, JsonValue] +) -> tuple[Mapping[str, JsonValue], ...] | None: + object_defs: Final = MappingProxyType({name: value for name, value in defs.items() if isinstance(value, dict)}) + def_names: Final = MappingProxyType({name: name for name in object_defs}) + next_names: Final = ((name, _root_ref_def_name(value, def_names)) for name, value in object_defs.items()) + next_def: Final = MappingProxyType({name: target for name, target in next_names if target is not None}) + first: Final = _root_ref_def_name(schema, def_names) + if first is None: + return (schema,) + names: Final = tuple(accumulate(range(len(next_def)), lambda name, _: next_def.get(name, name), initial=first)) + return None if names[-1] in next_def else (schema, *(object_defs[name] for name in frozenset(names))) + + +def needs_parameters_json_schema(parameters: Mapping[str, JsonValue]) -> bool: + defs: Final = parameters.get("$defs") + chain: Final = _root_ref_chain(parameters, defs if isinstance(defs, dict) else {}) + return chain is not None and any(_root_cannot_be_vertex_parameters(schema) for schema in chain) + + +def _is_union(schema: Mapping[str, JsonValue]) -> bool: + return "anyOf" in schema or "oneOf" in schema + + +def has_bare_root_union(json_schema: Mapping[str, JsonValue]) -> bool: + defs: Final = json_schema.get("$defs") + chain: Final = _root_ref_chain(json_schema, defs if isinstance(defs, dict) else {}) or (json_schema,) + return any(map(_is_union, chain)) and not any("properties" in schema for schema in chain) + + +_FORCING_MODES: Final = frozenset({"ANY", "VALIDATED"}) + + +def _forces(function_calling_config: JsonValue, name: JsonValue) -> bool: + match function_calling_config: + case {"mode": str() as mode, "allowed_function_names": [*names]} | { + "mode": str() as mode, + "allowedFunctionNames": [*names], + } if mode.upper() in _FORCING_MODES: + return name in names + case {"mode": str() as mode} if mode.upper() in _FORCING_MODES: + return True + case _: + return False + + +def _vertex_parameters(json_schema: Mapping[str, JsonValue], send_integer_enums: bool) -> Mapping[str, JsonValue]: + stripped: Final = remove_strict_from_schema(remove_additional_properties(deepcopy(dict(json_schema)))) + return _JSON_SCHEMA.validate_python(build_vertex_schema(stripped, send_integer_enums=send_integer_enums)) + + +def _declaration_for(declaration: JsonValue, function_calling_config: JsonValue, send_integer_enums: bool) -> JsonValue: + if not isinstance(declaration, dict): + return declaration + json_schema: Final = declaration.get("parametersJsonSchema") + if not isinstance(json_schema, dict) or not has_bare_root_union(json_schema): + return declaration + if not _forces(function_calling_config, declaration.get("name")): + return declaration + kept: Final = {key: value for key, value in declaration.items() if key != "parametersJsonSchema"} + return {**kept, "parameters": dict(_vertex_parameters(json_schema, send_integer_enums))} + + +def _tool_for(tool: JsonValue, function_calling_config: JsonValue, send_integer_enums: bool) -> JsonValue: + declarations: Final = tool.get("function_declarations") if isinstance(tool, dict) else None + if not isinstance(tool, dict) or not isinstance(declarations, list): + return tool + return { + **tool, + "function_declarations": [ + _declaration_for(item, function_calling_config, send_integer_enums) for item in declarations + ], + } + + +def with_forced_bare_root_unions_on_parameters( + tools: JsonValue, tool_config: JsonValue, send_integer_enums: bool +) -> JsonValue: + match tool_config: + case {"functionCallingConfig": function_calling_config} | {"function_calling_config": function_calling_config}: + return ( + [_tool_for(tool, function_calling_config, send_integer_enums) for tool in tools] + if isinstance(tools, list) + else tools + ) + case _: + return tools + + +def parameters_json_schema_for(parameters: Mapping[str, object]) -> Mapping[str, JsonValue] | None: + if "$ref" not in parameters and not _root_cannot_be_vertex_parameters(parameters): + return None + json_schema: Final = _JSON_SCHEMA.validate_python(to_jsonable_python(parameters)) + return json_schema if needs_parameters_json_schema(json_schema) else None + def build_json_schema(parameters: dict) -> dict: """ @@ -1656,7 +1788,7 @@ class VertexAITokenCounter(BaseTokenCounter): if isinstance(request_view, InferenceRequestViewError): 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 c29f2ecbcf8..b33e1d6fdb1 100644 --- a/litellm/llms/vertex_ai/count_tokens/context.py +++ b/litellm/llms/vertex_ai/count_tokens/context.py @@ -80,6 +80,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 ()) @@ -104,6 +105,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), ) @@ -119,11 +121,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 478291bd7db..be7628bf064 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -1338,6 +1338,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/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 765ccff23f5..97a36403c75 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -12,7 +12,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast from urllib.parse import quote import httpx -from pydantic import BaseModel, PositiveInt, TypeAdapter, ValidationError +from pydantic import BaseModel, JsonValue, PositiveInt, TypeAdapter, ValidationError from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm @@ -67,6 +67,8 @@ from ..common_utils import ( gemini_video_metadata_from_openai, get_supports_response_schema, get_supports_system_message, + with_forced_bare_root_unions_on_parameters, + with_integer_enums_that_fit, ) from .long_decode import with_default_output_budget @@ -1068,6 +1070,35 @@ def _pop_and_merge_extra_body(data: RequestBody, optional_params: dict) -> None: data_dict[k] = v +class _CallerToolConfig(TypedDict, total=False): + toolConfig: ReadOnly[JsonValue] + tool_config: ReadOnly[JsonValue] + + +_CALLER_TOOL_CONFIG: Final = TypeAdapter(_CallerToolConfig) + + +def _tools_fit_for_caller_tool_config(tools: Tools, extra_body: object, model: str) -> Tools: + extra_body_fields: Final = as_str_mapping(extra_body) + if extra_body_fields is None: + return tools + try: + caller: Final = _CALLER_TOOL_CONFIG.validate_python(extra_body_fields) + except ValidationError: + return tools + tool_config: Final = caller.get("toolConfig", caller.get("tool_config")) + if tool_config is None: + return tools + from .vertex_and_google_ai_studio_gemini import VertexGeminiConfig + + json_tools: Final = cast(JsonValue, tools) # cast-ok: LiteLLM built these tools as JSON request data + fitted: Final = with_integer_enums_that_fit( + partial(with_forced_bare_root_unions_on_parameters, json_tools, tool_config), + VertexGeminiConfig.is_gemini_3_or_newer(model), + ) + return cast(Tools, fitted) # cast-ok: same shape back + + def _has_google_maps_tool(tools: object) -> bool: """Return True if any tool object in the list has a 'googleMaps' key.""" if not isinstance(tools, list): @@ -1212,7 +1243,9 @@ def transform_request_body( if system_instructions is not None: data["system_instruction"] = system_instructions if tools is not None: - data["tools"] = tools + built_tools: Final = cast(Tools, tools) # cast-ok: popped from the untyped optional_params + caller_params: Final = cast(Mapping[str, object], optional_params) # cast-ok: untyped shared signature + data["tools"] = _tools_fit_for_caller_tool_config(built_tools, caller_params.get("extra_body"), model) if tool_choice is not None: data["toolConfig"] = tool_choice if include_server_side_tool_invocations: 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 e462e6a1488..572dcd5eecf 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 @@ -14,7 +14,8 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, 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 @@ -106,6 +107,8 @@ from ..common_utils import ( VertexAIError, build_json_schema, build_vertex_schema, + has_bare_root_union, + parameters_json_schema_for, supports_response_json_schema, with_integer_enums_that_fit, ) @@ -553,6 +556,48 @@ def mapped_gemini_speech_config( } +class _CallerFunction(TypedDict, total=False): + parameters: ReadOnly[Mapping[str, object]] + + +class _CallerTool(_CallerFunction, total=False): + function: ReadOnly[_CallerFunction] + + +class _ToolChoiceParam(TypedDict, total=False): + tool_choice: ReadOnly[object] + + +_CALLER_TOOL: Final = TypeAdapter(_CallerTool) +_CALLER_TOOLS: Final = TypeAdapter(tuple[object, ...]) +_TOOL_CHOICE_PARAM: Final = TypeAdapter(_ToolChoiceParam) + + +def _caller_parameters(tool: object) -> Mapping[str, object] | None: + try: + caller_tool: Final = _CALLER_TOOL.validate_python(tool) + except ValidationError: + return None + return caller_tool.get("function", caller_tool).get("parameters") + + +def _parameters_json_schema(tool: object, forced_function_names: frozenset[str]) -> Mapping[str, JsonValue] | None: + parameters: Final = _caller_parameters(tool) + json_schema: Final = None if parameters is None else parameters_json_schema_for(parameters) + if json_schema is None or declared_function_name(tool) not in forced_function_names: + return json_schema + return None if has_bare_root_union(json_schema) else json_schema + + +def _parameters_json_schemas( + tools: object, allow_parameters_json_schema: bool, forced_function_names: frozenset[str] +) -> tuple[Mapping[str, JsonValue] | None, ...]: + return tuple( + _parameters_json_schema(tool, forced_function_names) if allow_parameters_json_schema else None + for tool in _CALLER_TOOLS.validate_python(tools) + ) + + def _as_chat_tool(member: JsonValue) -> JsonValue: if not isinstance(member, dict) or "function" in member or "name" not in member: return member @@ -685,6 +730,16 @@ def _forced_function_names(name: str, tools: JsonValue) -> tuple[str, ...]: return tuple(qualified) or (name,) +def _function_names_forced_by(tool_choice: object, tools: JsonValue) -> frozenset[str]: + match tool_choice: + case "required": + return frozenset(chain.from_iterable(_names_each_tool_declares(_listed(tools)))) + case {"function": {"name": str() as name}}: + return frozenset(_forced_function_names(name, tools)) + case _: + return frozenset() + + class VertexAIBaseConfig: def get_mapped_special_auth_params(self) -> dict: """ @@ -1226,7 +1281,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): value: list[dict], optional_params: dict, send_integer_enums: bool = False, + allow_parameters_json_schema: bool = True, outside_declarations: JsonValue = None, + tool_choice: object = None, ) -> list[Tools]: """ Map OpenAI-style tools/functions to Vertex AI format. @@ -1235,6 +1292,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 @@ -1254,18 +1313,25 @@ 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( + _with_namespace_members_inlined(schema_value, outside_declarations), + allow_parameters_json_schema, + _function_names_forced_by(tool_choice, [*_listed(schema_value), *_listed(outside_declarations)]), + ) remove_additional_properties(schema_value) # remove 'strict' from tools remove_strict_from_schema(schema_value) inlined_tools: Final = _with_namespace_members_inlined(schema_value, outside_declarations) - for tool in cast(list[dict], inlined_tools): # 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], inlined_tools)[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. @@ -1328,8 +1394,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) @@ -1339,7 +1409,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: @@ -1415,8 +1485,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) @@ -1836,6 +1907,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): model: str, drop_params: bool, allow_integer_enums: bool = True, + allow_parameters_json_schema: bool = True, ) -> dict: request_params: Final = tool_params_without_hosted_mcp( non_default_params, model=model, route=self._route, drop_params=drop_params @@ -1898,7 +1970,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): self._map_function, value, optional_params, + allow_parameters_json_schema=allow_parameters_json_schema, outside_declarations=request_params.get("functions" if param == "tools" else "tools"), + tool_choice=_TOOL_CHOICE_PARAM.validate_python(request_params).get("tool_choice"), ), allow_integer_enums and self._is_gemini_3_or_newer(model), ) diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 968fb11e44a..2fbf9996f4d 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -2,6 +2,7 @@ from collections.abc import Mapping, Sequence from enum import Enum from typing import Final, Literal, Protocol +from pydantic import JsonValue from typing_extensions import ( NotRequired, ReadOnly, @@ -164,6 +165,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 d80661ea8fd..4696203c914 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -704,6 +704,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. On 2026-10-29 aiand's # deepseek-v4.1-flash billed 296 for each, and Fireworks' deepseek-v4p1-flash and OpenRouter's diff --git a/tests/unit/llms/gemini/test_gemini_common_utils.py b/tests/unit/llms/gemini/test_gemini_common_utils.py index 508bd264f19..51b2c36038d 100644 --- a/tests/unit/llms/gemini/test_gemini_common_utils.py +++ b/tests/unit/llms/gemini/test_gemini_common_utils.py @@ -82,6 +82,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 bc8da76656c..1c559ccbaa4 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 @@ -1698,6 +1698,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/gemini/test_transformation.py b/tests/unit/llms/vertex_ai/gemini/test_transformation.py index ac2d43f64a1..946f7c1cc6c 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_transformation.py +++ b/tests/unit/llms/vertex_ai/gemini/test_transformation.py @@ -485,3 +485,69 @@ def test_image_only_user_content_logs_the_blank_text_fixup_at_debug(caplog: pyte fixups: Final = tuple(r for r in caplog.records if "No text in user content" in r.getMessage()) assert tuple(r.levelno for r in fixups) == (logging.DEBUG,) assert contents[0]["parts"][-1] == {"text": " "} + + +_LOOKUP_TOOL: Final = { + "type": "function", + "function": { + "name": "lookup", + "parameters": { + "anyOf": [ + {"type": "object", "properties": {"a": {"type": "integer", "enum": [1, 2]}}, "required": ["a"]}, + {"type": "object", "properties": {"b": {"type": "string"}}, "required": ["b"]}, + ] + }, + }, +} +_WEATHER_TOOL: Final = { + "type": "function", + "function": { + "name": "weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + }, +} + + +def _declarations(model: str, tool_choice: str, extra_body: dict[str, object] | None) -> list[dict[str, object]]: + optional_params: Final = VertexGeminiConfig().map_openai_params( + non_default_params={"tools": [_LOOKUP_TOOL, _WEATHER_TOOL], "tool_choice": tool_choice}, + optional_params={}, + model=model, + drop_params=False, + ) + body: Final = transformation.transform_request_body( + messages=[{"role": "user", "content": "Find the record named Atlas"}], + model=model, + optional_params={**optional_params, **({} if extra_body is None else {"extra_body": extra_body})}, + custom_llm_provider="gemini", + litellm_params={}, + cached_content=None, + ) + return body["tools"][0]["function_declarations"] + + +@pytest.mark.parametrize( + ("extra_body", "forced"), + [ + ({"toolConfig": {"functionCallingConfig": {"mode": "ANY"}}}, True), + ({"tool_config": {"function_calling_config": {"mode": "VALIDATED"}}}, True), + ({"toolConfig": {"functionCallingConfig": {"mode": "any"}}}, True), + ({"toolConfig": {"functionCallingConfig": {"mode": "ANY", "allowedFunctionNames": ["lookup"]}}}, True), + ({"toolConfig": {"functionCallingConfig": {"mode": "ANY", "allowed_function_names": ["weather"]}}}, False), + ({"toolConfig": {"functionCallingConfig": {"mode": "AUTO"}}}, False), + ({"toolConfig": {"retrievalConfig": {"languageCode": "en"}}}, False), + (None, False), + ], + ids=["any", "validated-snake-case", "lowercase-any", "any-named-union-tool", "any-named-other-tool", "auto", "no-mode", "none"], +) +@pytest.mark.parametrize("model", ["gemini-2.5-flash", "gemini-4-argon"]) +def test_a_tool_call_forced_through_extra_body_keeps_a_bare_root_union_tool_as_tool_choice_does( + model: str, extra_body: dict[str, object] | None, forced: bool +) -> None: + lookup, weather = _declarations(model, "auto", extra_body) + forced_lookup, forced_weather = _declarations(model, "required", None) + free_lookup, _ = _declarations(model, "auto", None) + + assert "parameters" in forced_lookup and "parametersJsonSchema" in free_lookup + assert lookup == (forced_lookup if forced else free_lookup) + assert weather == forced_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 dbd71270978..31dbb9a4e9e 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 @@ -2897,6 +2897,199 @@ 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]}}}, + { + "$ref": "#/$defs/Lookup", + "$defs": {"Lookup": {"$ref": "#/$defs/Union"}, "Union": {"anyOf": [_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", + "chained-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: Mapping[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_tool_schema_built_with_tuples_as_parameters_json_schema() -> None: + request_params: Final = VertexGeminiConfig().map_openai_params( + non_default_params={ + "tools": [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"anyOf": ({**_LOOKUP_BY_A, "required": ("a",)}, _LOOKUP_BY_B)}, + }, + } + ] + }, + optional_params={}, + model="gemini-2.5-flash", + drop_params=False, + ) + + assert request_params["tools"] == [ + {"function_declarations": [{"name": "lookup", "parametersJsonSchema": {"anyOf": [_LOOKUP_BY_A, _LOOKUP_BY_B]}}]} + ] + + +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}, + ] + } + ] + + +_ROOT_ONE_OF: Final = {"oneOf": [_LOOKUP_BY_A, _LOOKUP_BY_B]} + + +@pytest.mark.parametrize( + ("lookup_parameters", "tool_choice", "expected_field"), + [ + (_ROOT_ONE_OF, "required", "parameters"), + (_ROOT_ONE_OF, {"type": "function", "function": {"name": "lookup"}}, "parameters"), + ( + {"$ref": "#/$defs/Lookup", "$defs": {"Lookup": {"anyOf": [_LOOKUP_BY_A, _LOOKUP_BY_B]}}}, + "required", + "parameters", + ), + ( + {"type": "object", "properties": _LOOKUP_PROPERTIES, "anyOf": [{"required": ["a"]}, {"required": ["b"]}]}, + "required", + "parametersJsonSchema", + ), + (_ROOT_ONE_OF, {"type": "function", "function": {"name": "weather"}}, "parametersJsonSchema"), + (_ROOT_ONE_OF, "auto", "parametersJsonSchema"), + ( + {"allOf": [{"type": "object", "properties": _LOOKUP_PROPERTIES}, {"required": ["b"]}]}, + "required", + "parametersJsonSchema", + ), + ], + ids=[ + "required", + "named-union-tool", + "required-root-ref-to-union", + "required-root-properties-beside-union", + "named-other-tool", + "auto", + "required-root-all-of", + ], +) +def test_vertex_keeps_a_union_root_tool_on_parameters_when_the_call_can_be_forced_to_it( + lookup_parameters: Mapping[str, object], tool_choice: str | dict[str, object], expected_field: str +) -> None: + request_params: Final = VertexGeminiConfig().map_openai_params( + non_default_params={ + "tool_choice": tool_choice, + "tools": [ + {"type": "function", "function": {"name": "lookup", "parameters": deepcopy(lookup_parameters)}}, + {"type": "function", "function": {"name": "weather", "parameters": deepcopy(_WEATHER_PARAMETERS)}}, + ], + }, + optional_params={}, + model="gemini-2.5-flash", + drop_params=False, + ) + lookup, weather = request_params["tools"][0]["function_declarations"] + + assert {field for field in ("parameters", "parametersJsonSchema") if field in lookup} == {expected_field} + assert weather == {"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 eb300d823aa..997bf6cf566 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 @@ -1,7 +1,7 @@ import copy import threading import time -from collections.abc import Iterator, Mapping +from collections.abc import Callable, Iterator, Mapping from contextlib import contextmanager from enum import Enum from typing import Final @@ -26,6 +26,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, @@ -298,6 +299,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 @@ -1628,6 +1660,109 @@ 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"]} +_DEEP_ROOT_REF_CHAIN_LENGTH: Final = 3000 + + +def _deep_root_ref_chain(end: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]: + links: Final = {f"d{i}": {"$ref": f"#/$defs/d{i + 1}"} for i in range(_DEEP_ROOT_REF_CHAIN_LENGTH)} + return {"$ref": "#/$defs/d0", "$defs": {**links, f"d{_DEEP_ROOT_REF_CHAIN_LENGTH}": end}} + + +@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), + ({"type": ["object"], "properties": {"a": {"type": "string"}}}, False), + ({"$ref": "#/$defs/Lookup", "$defs": {"Lookup": {"type": ["OBJECT"]}}}, False), + ({"$ref": "#/$defs/A", "$defs": {"A": {"$ref": "#/$defs/B"}, "B": {"anyOf": [_OBJECT_BRANCH]}}}, True), + ({"$ref": "#/$defs/A", "$defs": {"A": {"$ref": "#/$defs/B"}, "B": _OBJECT_BRANCH}}, False), + ({"$ref": "#/$defs/Lookup", "anyOf": [{"required": ["a"]}], "$defs": {"Lookup": _OBJECT_BRANCH}}, True), + ({"$ref": "#/$defs/A", "$defs": {"A": {"$ref": "#/$defs/B", "oneOf": [_OBJECT_BRANCH]}, "B": _OBJECT_BRANCH}}, True), + ({"$ref": "#/$defs/A", "$defs": {"A": {"$ref": "#/$defs/B"}, "B": {"$ref": "#/$defs/A"}}}, False), + ({"$ref": "#/$defs/A", "$defs": {"A": {"$ref": "#/$defs/B"}, "B": {"$ref": "#/$defs/A", "anyOf": [_OBJECT_BRANCH]}}}, False), + ({"$ref": "#/$defs/A", "anyOf": [{"required": ["a"]}], "$defs": {"A": {"$ref": "#/$defs/B"}, "B": {"$ref": "#/$defs/A"}}}, False), + (_deep_root_ref_chain({"anyOf": [_OBJECT_BRANCH]}), True), + (_deep_root_ref_chain(_OBJECT_BRANCH), False), + ], + 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", + "single-object-type-list", + "root-ref-to-single-object-type-list", + "chained-root-ref-to-union", + "chained-root-ref-to-object", + "root-ref-to-object-beside-a-union", + "chained-root-ref-through-a-union-to-object", + "cyclic-root-ref", + "cyclic-root-ref-through-a-union", + "cyclic-root-ref-beside-a-union", + "deep-root-ref-chain-to-union", + "deep-root-ref-chain-to-object", + ], +) +def test_needs_parameters_json_schema(parameters: Mapping[str, JsonValue], expected: bool) -> None: + assert needs_parameters_json_schema(parameters) is expected + + +def _missing_root_ref(size: int) -> Mapping[str, JsonValue]: + return {"$ref": "#/$defs/" + "m" * size, "$defs": {f"d{i}": {} for i in range(size // 16)}} + + +def _long_named_root_ref_cycle(size: int) -> Mapping[str, JsonValue]: + first, second = "a" * size, "b" * size + fillers: Final = {f"d{i}": {"type": "string"} for i in range(size // 16)} + cycle: Final = {first: {"$ref": f"#/$defs/{second}"}, second: {"$ref": f"#/$defs/{first}"}} + return {"$ref": f"#/$defs/{first}", "$defs": {**cycle, **fillers}} + + +def _root_ref_cost(parameters: Mapping[str, JsonValue]) -> float: + def once() -> float: + started: Final = time.perf_counter() + needs_parameters_json_schema(parameters) + return time.perf_counter() - started + + return min(once() for _ in range(3)) + + +@pytest.mark.parametrize("build", [_missing_root_ref, _long_named_root_ref_cycle], ids=["missing-ref", "ref-cycle"]) +def test_needs_parameters_json_schema_stays_linear_on_long_refs_and_many_defs( + build: Callable[[int], Mapping[str, JsonValue]], +) -> None: + """Tool parameters are caller-controlled, so 4x longer refs with 4x more $defs must not cost 16x. + Comparing two sizes keeps this honest on a slow box, and the quadratic walk this replaced grew 13x to 16x.""" + growth: Final = _root_ref_cost(build(256 * 1024)) / _root_ref_cost(build(64 * 1024)) + + assert growth < 10.0, f"cost grew {growth:.1f}x for 4x the input" + + class _Priority(int, Enum): LOW = 1 HIGH = 3