diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 21a80e39f26..a42e56ef178 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -11,6 +11,7 @@ from typing import Any, 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 @@ -891,25 +892,52 @@ _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 _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 +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 parameters + 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 _resolve_root_ref( + schema: Mapping[str, JsonValue], defs: Mapping[str, JsonValue], seen_refs: frozenset[str] +) -> Mapping[str, JsonValue]: + match schema.get("$ref"): + case str() as ref if ref.startswith(_ROOT_DEFS_REF_PREFIX) and ref not in seen_refs: + target: Final = defs.get(ref.removeprefix(_ROOT_DEFS_REF_PREFIX)) + return _resolve_root_ref(target, defs, seen_refs | {ref}) if isinstance(target, dict) else schema + case _: + return schema 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") + defs: Final = parameters.get("$defs") + return _root_cannot_be_vertex_parameters( + _resolve_root_ref(parameters, defs if isinstance(defs, dict) else {}, frozenset()) ) +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: """ Build a JSON Schema for use with Gemini's responseJsonSchema parameter. 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 ac5c84b6414..639df5f222e 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 @@ -100,7 +100,7 @@ from ..common_utils import ( VertexAIError, build_json_schema, build_vertex_schema, - needs_parameters_json_schema, + parameters_json_schema_for, supports_response_json_schema, with_integer_enums_that_fit, ) @@ -513,7 +513,7 @@ def mapped_gemini_speech_config( class _CallerFunction(TypedDict, total=False): - parameters: ReadOnly[Mapping[str, JsonValue]] + parameters: ReadOnly[Mapping[str, object]] class _CallerTool(_CallerFunction, total=False): @@ -524,13 +524,17 @@ _CALLER_TOOL: Final = TypeAdapter(_CallerTool) _CALLER_TOOLS: Final = TypeAdapter(tuple[object, ...]) -def _parameters_json_schema(tool: object) -> Mapping[str, JsonValue] | None: +def _caller_parameters(tool: object) -> Mapping[str, object] | 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 + return caller_tool.get("function", caller_tool).get("parameters") + + +def _parameters_json_schema(tool: object) -> Mapping[str, JsonValue] | None: + parameters: Final = _caller_parameters(tool) + return None if parameters is None else parameters_json_schema_for(parameters) def _parameters_json_schemas(tools: object) -> tuple[Mapping[str, JsonValue] | None, ...]: 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 944dd400e9b..17ea76d1272 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 @@ -2592,6 +2592,10 @@ _WEATHER_VERTEX_PARAMETERS: Final = {"type": "object", "properties": {"city": {" ], }, {"$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=[ @@ -2600,12 +2604,13 @@ _WEATHER_VERTEX_PARAMETERS: Final = {"type": "object", "properties": {"city": {" "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: dict[str, object] + 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.*" @@ -2640,6 +2645,29 @@ def test_gemini_sends_a_union_root_tool_schema_unchanged_as_parameters_json_sche ] +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( 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 83d2c8d63fd..243dc9fdebd 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,6 +1,6 @@ import copy import time -from collections.abc import Iterator +from collections.abc import Iterator, Mapping from contextlib import contextmanager from enum import Enum from typing import Final @@ -1644,6 +1644,11 @@ _OBJECT_BRANCH: Final = {"type": "object", "properties": {"a": {"type": "string" ({"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/A", "$defs": {"A": {"$ref": "#/$defs/B"}, "B": {"$ref": "#/$defs/A"}}}, False), ], ids=[ "object", @@ -1661,9 +1666,14 @@ _OBJECT_BRANCH: Final = {"type": "object", "properties": {"a": {"type": "string" "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", + "cyclic-root-ref", ], ) -def test_needs_parameters_json_schema(parameters: dict[str, JsonValue], expected: bool) -> None: +def test_needs_parameters_json_schema(parameters: Mapping[str, JsonValue], expected: bool) -> None: assert needs_parameters_json_schema(parameters) is expected