diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 800a0939cec..38de7be20ee 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -1830,7 +1830,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 0fa80cb2cad..21a80e39f26 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -889,6 +889,26 @@ 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/" + + +def _root_ref_target(parameters: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]: + match parameters.get("$ref"), parameters.get("$defs"): + case str() as ref, dict() as defs if ref.startswith(_ROOT_DEFS_REF_PREFIX): + target: Final = defs.get(ref.removeprefix(_ROOT_DEFS_REF_PREFIX)) + return target if isinstance(target, dict) else parameters + case _: + return parameters + + +def needs_parameters_json_schema(parameters: Mapping[str, JsonValue]) -> bool: + root: Final = _root_ref_target(parameters) + schema_type: Final = root.get("type") + return any(keyword in root for keyword in _VERTEX_PARAMETERS_UNSUPPORTED_ROOT_KEYWORDS) or ( + schema_type is not None and not (isinstance(schema_type, str) and schema_type.lower() == "object") + ) + def build_json_schema(parameters: dict) -> dict: """ 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 1ef5267df14..ac5c84b6414 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 @@ -12,7 +12,8 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, 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 @@ -99,6 +100,7 @@ from ..common_utils import ( VertexAIError, build_json_schema, build_vertex_schema, + needs_parameters_json_schema, supports_response_json_schema, with_integer_enums_that_fit, ) @@ -510,6 +512,31 @@ def mapped_gemini_speech_config( } +class _CallerFunction(TypedDict, total=False): + parameters: ReadOnly[Mapping[str, JsonValue]] + + +class _CallerTool(_CallerFunction, total=False): + function: ReadOnly[_CallerFunction] + + +_CALLER_TOOL: Final = TypeAdapter(_CallerTool) +_CALLER_TOOLS: Final = TypeAdapter(tuple[object, ...]) + + +def _parameters_json_schema(tool: object) -> Mapping[str, JsonValue] | None: + try: + caller_tool: Final = _CALLER_TOOL.validate_python(tool) + except ValidationError: + return None + parameters: Final = caller_tool.get("function", caller_tool).get("parameters") + return parameters if parameters is not None and needs_parameters_json_schema(parameters) else None + + +def _parameters_json_schemas(tools: object) -> tuple[Mapping[str, JsonValue] | None, ...]: + return tuple(_parameters_json_schema(tool) for tool in _CALLER_TOOLS.validate_python(tools)) + + class VertexAIBaseConfig: def get_mapped_special_auth_params(self) -> dict: """ @@ -1059,17 +1086,20 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): computerUse: dict | None = None # remove 'additionalProperties' from tools schema_value: Final[JsonValue] = deepcopy(cast(JsonValue, value)) # cast-ok: tools are JSON request data + parameters_json_schemas: Final = _parameters_json_schemas(schema_value) remove_additional_properties(schema_value) # remove 'strict' from tools remove_strict_from_schema(schema_value) - for tool in cast(list[dict], schema_value): # 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], schema_value)[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. @@ -1132,8 +1162,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) @@ -1143,7 +1177,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: diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index a938f6a7b1c..ec76e271964 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -1,7 +1,8 @@ -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from enum import Enum from typing import Final, Literal, Protocol +from pydantic import JsonValue from typing_extensions import ( ReadOnly, Required, @@ -163,6 +164,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 697ac76b4e5..afd3a74aefb 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -693,6 +693,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 @pytest.mark.parametrize("tool_choice", ["none", "required", {"type": "function", "function": {"name": "get_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 98d1b8f1300..944dd400e9b 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 @@ -2556,6 +2556,114 @@ 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]}}}, + {"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", + "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: dict[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_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}, + ] + } + ] + + @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 152ab2ffc13..83d2c8d63fd 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 @@ -23,6 +23,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, @@ -1622,6 +1623,50 @@ 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"]} + + +@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), + ], + 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", + ], +) +def test_needs_parameters_json_schema(parameters: dict[str, JsonValue], expected: bool) -> None: + assert needs_parameters_json_schema(parameters) is expected + + class _Priority(int, Enum): LOW = 1 HIGH = 3