scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-15 09:12:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 311

fix(vertex-ai): resolve long root $ref chains without recursion

Browse source
This commit is contained in:
mateo-berri 2026-11-04 19:49:41 -08:00
parent b67da07d7f
commit 47d9c8bab8
2 changed files with 22 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

19
litellm/llms/vertex_ai/common_utils.py
View file

@ -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:

12
tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py
View file

@ -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: