mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-15 09:12:18 +00:00
fix(vertex-ai): send union-root tool schemas as parametersJsonSchema (#51835)
* 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 * fix(vertex-ai): resolve chained root refs and tuple schemas for parametersJsonSchema A root $ref that points at another $ref now resolves to its end, so a chain that ends in a union also goes out as parametersJsonSchema. A schema built with tuples or sets is turned into its JSON form first, so it no longer falls back to the parameters field. A type of ["object"] stays on the parameters field, which already accepts it. Only tools whose root could need parametersJsonSchema pay for the JSON conversion * fix(vertex-ai): resolve long root $ref chains without recursion * 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. * fix(vertex-ai): resolve a root $ref chain in linear time Walk $defs by name once with cycle detection, so a long missing ref or a long-named ref cycle beside many defs no longer costs quadratic CPU * fix(vertex-ai): send a root $ref tool as parametersJsonSchema when a union sits beside the ref Check the root and every ref on its chain for anyOf, oneOf, allOf, or a non-object type, since a rule next to a $ref was dropped when only the chain's end was checked * fix(vertex-ai): keep a bare root union tool on parameters when the call can be forced to it * fix(vertex-ai): keep a bare root union tool on parameters when extra_body forces the call * fix(vertex-ai): keep integer enums on a tool that extra_body forces back to parameters * fix(vertex-ai): keep a union beside a looping root $ref on parameters and read the forcing mode in any case A root union next to a $ref chain that loops back on itself went out as parametersJsonSchema, and Google answers 400 "circular ref loop" for that, where base sent parameters and got 200. A looping chain now stays on parameters whatever sits next to the $ref An extra_body toolConfig with mode "any" forces the call on the Gemini API, but only "ANY" and "VALIDATED" were read as forcing, so a bare root union tool moved to parametersJsonSchema and the forced call answered 400. The mode is now compared in upper case --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
quietjules
GitHub
parent
d8e8621b7b
commit
bc3113fff2
13 changed files with 744 additions and 14 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"),
|
||||
[
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue