diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index a42e56ef178..f882a54ae72 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -4,7 +4,7 @@ from collections.abc import Callable, Iterable, Mapping, Sequence from copy import deepcopy from datetime import datetime, timezone from enum import Enum -from functools import lru_cache +from functools import lru_cache, reduce from itertools import accumulate, chain, repeat, takewhile from types import MappingProxyType from typing import Any, Final, Literal, TypeVar, cast, get_type_hints @@ -913,22 +913,23 @@ def _root_cannot_be_vertex_parameters(root: Mapping[str, object]) -> bool: ) 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]: +def _root_ref_target(schema: Mapping[str, JsonValue], defs: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]: match schema.get("$ref"): - case str() as ref if ref.startswith(_ROOT_DEFS_REF_PREFIX) and ref not in seen_refs: + case str() as ref if ref.startswith(_ROOT_DEFS_REF_PREFIX): 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 + return target if isinstance(target, dict) else schema case _: return schema +def _resolve_root_ref(schema: Mapping[str, JsonValue], defs: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]: + resolved: Final = reduce(lambda current, _: _root_ref_target(current, defs), range(len(defs)), schema) + return resolved if _root_ref_target(resolved, defs) is resolved else schema + + def needs_parameters_json_schema(parameters: Mapping[str, JsonValue]) -> bool: defs: Final = parameters.get("$defs") - return _root_cannot_be_vertex_parameters( - _resolve_root_ref(parameters, defs if isinstance(defs, dict) else {}, frozenset()) - ) + return _root_cannot_be_vertex_parameters(_resolve_root_ref(parameters, defs if isinstance(defs, dict) else {})) def parameters_json_schema_for(parameters: Mapping[str, object]) -> Mapping[str, JsonValue] | None: 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 243dc9fdebd..e389fbb4eea 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 @@ -1624,6 +1624,12 @@ def test_map_function_keeps_integer_enum_when_tools_are_reused(): _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( @@ -1649,6 +1655,9 @@ _OBJECT_BRANCH: Final = {"type": "object", "properties": {"a": {"type": "string" ({"$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), + ({"$ref": "#/$defs/A", "$defs": {"A": {"$ref": "#/$defs/B"}, "B": {"$ref": "#/$defs/A", "anyOf": [_OBJECT_BRANCH]}}}, False), + (_deep_root_ref_chain({"anyOf": [_OBJECT_BRANCH]}), True), + (_deep_root_ref_chain(_OBJECT_BRANCH), False), ], ids=[ "object", @@ -1671,6 +1680,9 @@ _OBJECT_BRANCH: Final = {"type": "object", "properties": {"a": {"type": "string" "chained-root-ref-to-union", "chained-root-ref-to-object", "cyclic-root-ref", + "cyclic-root-ref-through-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: