diff --git a/litellm/llms/base_llm/managed_resources/base_managed_resource.py b/litellm/llms/base_llm/managed_resources/base_managed_resource.py index 2663a3c42e7..fc988a8ec3a 100644 --- a/litellm/llms/base_llm/managed_resources/base_managed_resource.py +++ b/litellm/llms/base_llm/managed_resources/base_managed_resource.py @@ -4,18 +4,29 @@ import base64 import json +import math from abc import ABC, abstractmethod from collections.abc import Mapping +from datetime import timedelta from typing import TYPE_CHECKING, Any, Final, Generic, Protocol, TypeVar, cast, runtime_checkable from litellm import verbose_logger from litellm._internal_context import with_service_target +from litellm.caching.dual_cache import DualCache +from litellm.caching.redis_batch import active_request_redis_batch +from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.llms.base_llm.managed_resources.isolation import ( build_list_page, build_owner_filter, can_access_resource, resolve_resource_owner_id, ) +from litellm.redis.stores.managed_vector_resources import ( + ManagedVectorResourceStore, + TransportManagedVectorResourceStore, + vector_resource_key, +) +from litellm.redis.transport.contract import RedisFailure from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth from litellm.types.utils import SpecialEnums @@ -66,6 +77,12 @@ class _ManagedResourceTable(Protocol[ResourceObjectType]): async def delete(self, *, where: Mapping[str, object]) -> object: ... +class _LocalVectorResourceCache(Protocol): + async def async_get_cache(self, key: str, *, local_only: bool) -> object: ... + + async def async_set_cache(self, key: str, value: object, *, local_only: bool) -> None: ... + + class BaseManagedResource(ABC, Generic[ResourceObjectType]): """ Base class for managing resources with target_model_names support. @@ -95,6 +112,50 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): def _resource_table(self) -> _ManagedResourceTable[ResourceObjectType]: return getattr(self.prisma_client.db, self.table_name) + def _metadata_cache(self) -> DualCache: + if isinstance(self.internal_usage_cache, DualCache): + return self.internal_usage_cache + return self.internal_usage_cache.dual_cache + + async def _cache_vector_resource(self, unified_resource_id: str, record: Mapping[str, object] | None) -> None: + cache: Final = self._metadata_cache() + local_cache: Final = require_attributes(cache, _LocalVectorResourceCache, "async_get_cache", "async_set_cache") + await local_cache.async_set_cache(vector_resource_key(unified_resource_id), record, local_only=True) + if cache.redis_cache is None: + return + store: Final[ManagedVectorResourceStore] = TransportManagedVectorResourceStore(cache.redis_cache.transport()) + try: + seconds: Final = ( + int(cache.default_in_memory_ttl) + if cache.default_in_memory_ttl is not None and not math.isnan(cache.default_in_memory_ttl) + else cache.redis_cache.default_ttl + ) + ttl: Final = timedelta(seconds=seconds) + except (OverflowError, ValueError): + return + if record is None: + await store.invalidate_vector_resource(unified_resource_id, ttl=ttl) + else: + await store.write_vector_resource(unified_resource_id, record, ttl=ttl) + + async def _cached_vector_resource(self, unified_resource_id: str) -> object: + cache: Final = self._metadata_cache() + local_cache: Final = require_attributes(cache, _LocalVectorResourceCache, "async_get_cache", "async_set_cache") + local: Final = await local_cache.async_get_cache(vector_resource_key(unified_resource_id), local_only=True) + if local is not None or cache.redis_cache is None: + return local + request_batch: Final = active_request_redis_batch(cache.redis_cache) + if request_batch is not None and request_batch.read_as_missing(vector_resource_key(unified_resource_id)): + return None + store: Final[ManagedVectorResourceStore] = TransportManagedVectorResourceStore(cache.redis_cache.transport()) + result: Final = await store.read_vector_resource(unified_resource_id) + if result is None or isinstance(result, RedisFailure): + return None + value: Final = result.value + if value is not None: + await local_cache.async_set_cache(vector_resource_key(unified_resource_id), value, local_only=True) + return value + # ============================================================================ # ABSTRACT METHODS # ============================================================================ @@ -201,11 +262,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): # Store in cache if resource_object is not None: - await self.internal_usage_cache.async_set_cache( - key=unified_resource_id, - value=cache_data, - litellm_parent_otel_span=litellm_parent_otel_span, - ) + await self._cache_vector_resource(unified_resource_id, cache_data) # Prepare database data db_data: Final[dict[str, object]] = { @@ -263,10 +320,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): # Check cache first result: Final = cast( dict | None, - await self.internal_usage_cache.async_get_cache( - key=unified_resource_id, - litellm_parent_otel_span=litellm_parent_otel_span, - ), + await self._cached_vector_resource(unified_resource_id), ) if result: @@ -305,11 +359,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): raise Exception(f"LiteLLM Managed {self.resource_type} with id={unified_resource_id} not found") # Delete from cache - await self.internal_usage_cache.async_set_cache( - key=unified_resource_id, - value=None, - litellm_parent_otel_span=litellm_parent_otel_span, - ) + await self._cache_vector_resource(unified_resource_id, None) # Delete from database await table.delete(where={"unified_resource_id": unified_resource_id}) @@ -355,7 +405,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): async def get_model_resource_id_mapping( self, resource_ids: list[str], - litellm_parent_otel_span: Span, + litellm_parent_otel_span: Span | None, ) -> dict[str, dict[str, str]]: """ Get model-specific resource IDs for a list of unified resource IDs. diff --git a/litellm/redis/stores/managed_vector_resources.py b/litellm/redis/stores/managed_vector_resources.py new file mode 100644 index 00000000000..732174dfb2e --- /dev/null +++ b/litellm/redis/stores/managed_vector_resources.py @@ -0,0 +1,71 @@ +import ast +import json +from collections.abc import Mapping +from dataclasses import dataclass +from datetime import timedelta +from typing import Final, Protocol + +from pydantic import TypeAdapter + +from litellm.redis.transport.contract import RedisFailure, RedisTransport + + +@dataclass(frozen=True, slots=True) +class ManagedVectorResourceValue: + value: object + + +_VALUE: Final[TypeAdapter[ManagedVectorResourceValue]] = TypeAdapter(ManagedVectorResourceValue) + + +def vector_resource_key(unified_resource_id: str) -> str: + return unified_resource_id + + +class ManagedVectorResourceStore(Protocol): + async def read_vector_resource( + self, unified_resource_id: str + ) -> ManagedVectorResourceValue | None | RedisFailure: ... + + async def write_vector_resource( + self, unified_resource_id: str, record: Mapping[str, object], *, ttl: timedelta | None + ) -> None | RedisFailure: ... + + async def invalidate_vector_resource( + self, unified_resource_id: str, *, ttl: timedelta | None + ) -> None | RedisFailure: ... + + +@dataclass(frozen=True, slots=True) +class TransportManagedVectorResourceStore: + transport: RedisTransport + + async def read_vector_resource(self, unified_resource_id: str) -> ManagedVectorResourceValue | None | RedisFailure: + raw: Final = await self.transport.get(vector_resource_key(unified_resource_id)) + if raw is None or isinstance(raw, RedisFailure): + return raw + try: + return _VALUE.validate_python({"value": json.loads(raw.decode("utf-8"))}) + except (ValueError, UnicodeError, RecursionError): + try: + return _VALUE.validate_python({"value": ast.literal_eval(raw.decode("utf-8"))}) + except (ValueError, SyntaxError, UnicodeError, RecursionError) as exc: + return RedisFailure(kind="decode", message=str(exc)) + + async def write_vector_resource( + self, unified_resource_id: str, record: Mapping[str, object], *, ttl: timedelta | None + ) -> None | RedisFailure: + try: + encoded: Final = json.dumps(record) + except (TypeError, ValueError) as exc: + return RedisFailure(kind="decode", message=str(exc)) + return await self._write(unified_resource_id, encoded, ttl) + + async def invalidate_vector_resource( + self, unified_resource_id: str, *, ttl: timedelta | None + ) -> None | RedisFailure: + return await self._write(unified_resource_id, "null", ttl) + + async def _write(self, unified_resource_id: str, encoded: str, ttl: timedelta | None) -> None | RedisFailure: + result: Final = await self.transport.set_seconds(vector_resource_key(unified_resource_id), encoded, ttl=ttl) + return result if isinstance(result, RedisFailure) else None diff --git a/tests/unit/llms/base_llm/test_base_managed_resource.py b/tests/unit/llms/base_llm/test_base_managed_resource.py index a6c518616f4..ebe93afe322 100644 --- a/tests/unit/llms/base_llm/test_base_managed_resource.py +++ b/tests/unit/llms/base_llm/test_base_managed_resource.py @@ -2,15 +2,16 @@ Integration tests for BaseManagedResource listing and access control. """ -from typing import List from unittest.mock import AsyncMock, MagicMock import pytest +from litellm.caching.dual_cache import DualCache from litellm.llms.base_llm.managed_resources.base_managed_resource import ( BaseManagedResource, ) from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.utils import InternalUsageCache class _StubResource(BaseManagedResource): @@ -24,18 +25,14 @@ class _StubResource(BaseManagedResource): def table_name(self) -> str: return "litellm_test_resource_table" - def get_unified_resource_id_format( - self, resource_object, target_model_names_list - ) -> str: + def get_unified_resource_id_format(self, resource_object, target_model_names_list) -> str: return "test" - async def create_resource_for_model( - self, llm_router, model, request_data, litellm_parent_otel_span - ): + async def create_resource_for_model(self, llm_router, model, request_data, litellm_parent_otel_span): return {"id": "test"} -def _make_resource(records: List = None) -> _StubResource: +def _make_resource(records: list[object] | None = None) -> _StubResource: cache = MagicMock() cache.async_get_cache = AsyncMock(return_value=None) @@ -68,9 +65,7 @@ async def test_list_user_filters_by_user_id(): await resource.list_user_resources(user_api_key_dict=user) - where = resource.prisma_client.db.litellm_test_resource_table.find_many.await_args.kwargs[ - "where" - ] + where = resource.prisma_client.db.litellm_test_resource_table.find_many.await_args.kwargs["where"] assert where["created_by"] == "alice" assert "team_id" not in where @@ -82,9 +77,7 @@ async def test_list_service_account_filters_by_team_id(): await resource.list_user_resources(user_api_key_dict=service_account) - where = resource.prisma_client.db.litellm_test_resource_table.find_many.await_args.kwargs[ - "where" - ] + where = resource.prisma_client.db.litellm_test_resource_table.find_many.await_args.kwargs["where"] assert where["team_id"] == "team-eng" assert "created_by" not in where @@ -114,18 +107,17 @@ async def test_list_identity_less_caller_returns_empty_without_query(): [("team-eng", True), ("team-sales", False), (None, False)], ) async def test_can_access_uses_team_id_for_service_account(caller_team_id, expected): - cache = MagicMock() - cache.async_get_cache = AsyncMock( - return_value={ + cache = DualCache() + await cache.async_set_cache( + key="rid", + value={ "created_by": None, "team_id": "team-eng", - } + }, ) prisma = MagicMock() - resource = _StubResource(internal_usage_cache=cache, prisma_client=prisma) + resource = _StubResource(internal_usage_cache=InternalUsageCache(cache), prisma_client=prisma) - caller = ( - UserAPIKeyAuth(team_id=caller_team_id) if caller_team_id else UserAPIKeyAuth() - ) + caller = UserAPIKeyAuth(team_id=caller_team_id) if caller_team_id else UserAPIKeyAuth() assert await resource.can_user_access_unified_resource_id("rid", caller) is expected diff --git a/tests/unit/redis/stores/test_managed_vector_resources.py b/tests/unit/redis/stores/test_managed_vector_resources.py new file mode 100644 index 00000000000..624ab3bf21c --- /dev/null +++ b/tests/unit/redis/stores/test_managed_vector_resources.py @@ -0,0 +1,286 @@ +import base64 +import json +import sys +import time +from collections.abc import Mapping +from dataclasses import dataclass +from datetime import timedelta +from types import SimpleNamespace +from typing import Final, cast +from unittest.mock import patch + +import fakeredis +import pytest +from litellm_enterprise.proxy.hooks.managed_vector_stores import PROXY_LiteLLMManagedVectorStores +from redis.exceptions import TimeoutError as RedisTimeoutError + +from litellm.caching.dual_cache import DualCache +from litellm.proxy.utils import InternalUsageCache, PrismaClient +from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth +from litellm.types.vector_stores import VectorStoreCreateResponse +from tests._support.fake_redis import fake_redis_cache + + +@dataclass(frozen=True) +class _Row: + data: Mapping[str, object] + + @property + def resource_object(self) -> object: + return self.data.get("resource_object") + + def model_dump(self) -> dict[str, object]: + return dict(self.data) + + +class _Table: + def __init__(self, row: _Row | None = None, raw: fakeredis.FakeAsyncRedis | None = None) -> None: + self.row = row + self.reads = 0 + self.raw = raw + + async def create(self, data: Mapping[str, object]) -> _Row: + self.row = _Row(data) + return self.row + + async def find_first(self, where: Mapping[str, object]) -> _Row | None: + self.reads += 1 + return self.row if self.row and self.row.data["unified_resource_id"] == where["unified_resource_id"] else None + + async def delete(self, where: Mapping[str, object]) -> None: + if self.raw is not None: + assert await self.raw.get(f"tenant:{where['unified_resource_id']}") == b"null" + self.row = None + + +def _resource() -> VectorStoreCreateResponse: + return cast(VectorStoreCreateResponse, {"id": "vs-é", "object": "vector_store", "name": "文", "usage_bytes": 2**70}) + + +def _hook(cache: DualCache, table: _Table) -> PROXY_LiteLLMManagedVectorStores: + return PROXY_LiteLLMManagedVectorStores( + InternalUsageCache(cache), + cast(PrismaClient, SimpleNamespace(db=SimpleNamespace(litellm_managedvectorstoretable=table))), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ttl", [1, 12.9, None, 0.5, 0, -1]) +async def test_vector_wire_matches_legacy_with_namespace_and_effective_ttl( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis, ttl: float | None +) -> None: + cache: Final = fake_redis_cache(server, namespace="tenant") + dual: Final = DualCache(redis_cache=cache, default_in_memory_ttl=ttl, default_redis_ttl=999) + dual.update_cache_ttl(ttl, None) + table: Final = _Table() + hook: Final = _hook(dual, table) + resource: Final = _resource() + unified_id: Final = hook.generate_unified_resource_id([resource], ["model-a"]) + assert ( + base64.urlsafe_b64decode(unified_id + "=" * (-len(unified_id) % 4)) + .decode() + .startswith("litellm_proxy:vector_store;unified_id,") + ) + record: Final = { + "unified_resource_id": unified_id, + "resource_object": resource, + "model_mappings": {"model-a": "vs-é"}, + "flat_model_resource_ids": ["vs-é"], + "created_by": "owner", + "team_id": "team-a", + "updated_by": "owner", + "extra": "文", + } + with patch("time.time", return_value=time.time()): + await hook.store_unified_resource_id( + unified_id, + resource, + None, + {"model-a": "vs-é"}, + UserAPIKeyAuth(user_id="owner", team_id="team-a"), + {"extra": "文"}, + ) + await cache.async_set_cache("legacy", record, **({"ttl": ttl} if ttl is not None else {})) + assert await raw.get(f"tenant:{unified_id}") == await raw.get("tenant:legacy") + assert await raw.pttl(f"tenant:{unified_id}") == await raw.pttl("tenant:legacy") + assert await raw.get(f"tenant:{unified_id}") == ( + None if ttl is not None and ttl < 1 else json.dumps(record).encode() + ) + assert await cache.async_get_cache(unified_id) == (None if ttl is not None and ttl < 1 else record) + assert await raw.get(f"tenant:tenant:{unified_id}") is None + assert table.row is not None + assert table.row.data == { + **record, + "resource_object": json.dumps(resource), + "model_mappings": json.dumps(record["model_mappings"]), + } + + +@pytest.mark.asyncio +async def test_old_vector_writer_new_reader_preserves_access_mapping_and_expiry( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + resource_id: Final = ( + base64.urlsafe_b64encode( + b"litellm_proxy:vector_store;unified_id,uuid;target_model_names,model;resource_id,vs-1;model_id,model" + ) + .decode() + .rstrip("=") + ) + record: Final = { + "unified_resource_id": resource_id, + "resource_object": _resource(), + "model_mappings": {"model": "vs-é"}, + "created_by": "owner", + "team_id": "team-a", + } + with patch("time.time", return_value=1_800_000_000): + cache: Final = fake_redis_cache(server, namespace="tenant") + dual: Final = DualCache(redis_cache=cache, default_in_memory_ttl=1) + table: Final = _Table(_Row({**record, "model_mappings": json.dumps(record["model_mappings"])})) + hook: Final = _hook(dual, table) + await cache.async_set_cache(resource_id, record, ttl=1) + assert await hook.get_unified_resource_id(resource_id) == record + assert await hook.can_user_access_unified_resource_id(resource_id, UserAPIKeyAuth(user_id="owner")) is True + assert await hook.can_user_access_unified_resource_id(resource_id, UserAPIKeyAuth(team_id="team-a")) is True + assert await hook.can_user_access_unified_resource_id(resource_id, UserAPIKeyAuth(team_id="team-b")) is False + assert await hook.get_model_resource_id_mapping([resource_id], None) == {resource_id: {"model": "vs-é"}} + assert table.reads == 0 + with patch("time.time", return_value=1_800_000_002): + assert await hook.get_model_resource_id_mapping([resource_id], None) == {resource_id: {"model": "vs-é"}} + assert table.reads == 1 + assert await raw.get(f"tenant:{resource_id}") is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("encoded", [b"null", b"{}", b"broken", b"\xff"]) +async def test_vector_null_or_malformed_value_falls_back_without_backfill( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis, encoded: bytes +) -> None: + row: Final = _Row({"unified_resource_id": "rid", "resource_object": _resource()}) + table: Final = _Table(row) + hook: Final = _hook(DualCache(redis_cache=fake_redis_cache(server)), table) + await raw.set("rid", encoded) + assert await hook.get_unified_resource_id("rid") == row.model_dump() + assert await raw.get("rid") == encoded + assert table.reads == 1 + + +@pytest.mark.asyncio +async def test_recursive_vector_metadata_retains_legacy_database_fallback( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + row: Final = _Row({"unified_resource_id": "rid", "resource_object": _resource()}) + table: Final = _Table(row) + cache: Final = fake_redis_cache(server) + hook: Final = _hook(DualCache(redis_cache=cache), table) + await raw.set("rid", b"[" * 10000 + b"0" + b"]" * 10000) + limit: Final = sys.getrecursionlimit() + try: + sys.setrecursionlimit(1000) + assert await cache.async_get_cache("rid") is None + assert await hook.get_unified_resource_id("rid") == row.model_dump() + assert table.reads == 1 + finally: + sys.setrecursionlimit(limit) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("encoded", [b"(1, 2)", b"{'item'}", b"{1: 'value'}"]) +async def test_python_literal_vector_value_retains_legacy_read_outcome( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis, encoded: bytes +) -> None: + cache: Final = fake_redis_cache(server) + table: Final = _Table(_Row({"unified_resource_id": "rid", "resource_object": _resource()})) + hook: Final = _hook(DualCache(redis_cache=cache), table) + await raw.set("rid", encoded) + assert await hook.get_unified_resource_id("rid") == await cache.async_get_cache("rid") + assert table.reads == 0 + + +@pytest.mark.asyncio +async def test_vector_invalidation_precedes_database_delete_and_keeps_null_ttl( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + row: Final = _Row({"unified_resource_id": "rid", "resource_object": _resource()}) + table: Final = _Table(row, raw) + hook: Final = _hook( + DualCache(redis_cache=fake_redis_cache(server, namespace="tenant"), default_in_memory_ttl=1), table + ) + with patch("time.time", return_value=time.time()): + assert await hook.delete_unified_resource_id("rid") == _resource() + assert await raw.get("tenant:rid") == b"null" + assert await raw.pttl("tenant:rid") == 1000 + assert await hook.get_unified_resource_id("rid") is None + assert table.reads == 2 + + +@pytest.mark.asyncio +async def test_vector_outage_keeps_local_and_database_writes(server: fakeredis.FakeServer) -> None: + cache: Final = fake_redis_cache(server) + dual: Final = DualCache(redis_cache=cache) + table: Final = _Table() + hook: Final = _hook(dual, table) + server.connected = False + await hook.store_unified_resource_id("rid", _resource(), None, {"model": "vs-é"}, UserAPIKeyAuth(user_id="owner")) + cached: Final = await hook.get_unified_resource_id("rid") + assert cached is not None + assert cached["resource_object"] == _resource() + assert table.reads == 0 + dual.in_memory_cache.flush_cache() + assert table.row is not None + for _ in range(8): + assert await hook.get_unified_resource_id("rid") == table.row.model_dump() + assert cache._circuit_breaker.state == "open" # pyright: ignore[reportPrivateUsage] # asserts the breaker opened after the outage + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ttl", [float("nan"), float("inf"), float("-inf")]) +async def test_invalid_vector_cache_ttl_does_not_block_database_write( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis, ttl: float +) -> None: + table: Final = _Table() + dual: Final = DualCache(redis_cache=fake_redis_cache(server), default_in_memory_ttl=ttl) + hook: Final = _hook(dual, table) + with patch("time.time", return_value=1_800_000_000): + await dual.async_set_cache("legacy", _resource()) + await hook.store_unified_resource_id( + "rid", _resource(), None, {"model": "vs-é"}, UserAPIKeyAuth(user_id="owner") + ) + assert table.row is not None + assert table.row.data["unified_resource_id"] == "rid" + assert (await raw.get("rid") is None) == (await raw.get("legacy") is None) + assert await raw.pttl("rid") == await raw.pttl("legacy") + + +@pytest.mark.asyncio +async def test_vector_unknown_applied_write_does_not_undo_cache_or_database( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + cache: Final = fake_redis_cache(server) + table: Final = _Table() + hook: Final = _hook(DualCache(redis_cache=cache, default_in_memory_ttl=1), table) + + async def applied_then_timeout(name: str, value: str, *, ex: timedelta | None, nx: bool) -> bool: + await raw.set(name, value, ex=ex, nx=nx) + raise RedisTimeoutError("reply lost after apply") + + with patch.object(cache.init_async_client(), "set", side_effect=applied_then_timeout): + await hook.store_unified_resource_id( + "rid", _resource(), None, {"model": "vs-é"}, UserAPIKeyAuth(user_id="owner") + ) + assert await raw.get("rid") == json.dumps(await hook.get_unified_resource_id("rid")).encode() + assert table.row is not None + + +@pytest.mark.asyncio +async def test_absent_vector_object_does_not_overwrite_cache( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + table: Final = _Table() + hook: Final = _hook(DualCache(redis_cache=fake_redis_cache(server)), table) + await raw.set("rid", b"previous") + await hook.store_unified_resource_id("rid", None, None, {"model": "vs-é"}, UserAPIKeyAuth()) + assert await raw.get("rid") == b"previous" + assert table.row is not None + assert "resource_object" not in table.row.data