mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-09 08:18:12 +00:00
fix(auth-cache): type UserApiKeyCache values as plain data
Some checks failed
Cost map pins recheck / rerun-cost-map-guard (push) Has been cancelled
LiteLLM Rust DB / db-queries (push) Has been cancelled
LiteLLM Rust / rust-lint (push) Has been cancelled
LiteLLM Rust / rust-test (push) Has been cancelled
LiteLLM Rust / rust-wheel (push) Has been cancelled
Terraform Modules / fmt, validate, test (aws) (push) Has been cancelled
Terraform Modules / fmt, validate, test (gcp) (push) Has been cancelled
Some checks failed
Cost map pins recheck / rerun-cost-map-guard (push) Has been cancelled
LiteLLM Rust DB / db-queries (push) Has been cancelled
LiteLLM Rust / rust-lint (push) Has been cancelled
LiteLLM Rust / rust-test (push) Has been cancelled
LiteLLM Rust / rust-wheel (push) Has been cancelled
Terraform Modules / fmt, validate, test (aws) (push) Has been cancelled
Terraform Modules / fmt, validate, test (gcp) (push) Has been cancelled
Every cache setter now takes a pydantic model or plain data (str, number, bool, None, sequences, str-keyed mappings), which is what Redis can hold and what the Rust auth cache accepts. No production writer passes anything else, and pyright rejects a new caller that does
This commit is contained in:
parent
6438d6a797
commit
e3df2dac15
3 changed files with 24 additions and 5 deletions
|
|
@ -16,7 +16,8 @@ are encoded; pass a Pydantic model or convert with e.g. ``dataclasses.asdict`` f
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, TypeVar
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Any, TypeAlias, TypeVar
|
||||
|
||||
from pydantic import BaseModel, ValidationError
|
||||
|
||||
|
|
@ -24,6 +25,9 @@ from litellm._logging import verbose_proxy_logger
|
|||
|
||||
T = TypeVar("T", bound=BaseModel)
|
||||
|
||||
CacheData: TypeAlias = "str | int | float | bool | None | Sequence[CacheData] | Mapping[str, CacheData]"
|
||||
CacheValue: TypeAlias = "BaseModel | CacheData"
|
||||
|
||||
|
||||
class CacheCodec:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@ from litellm.caching.dual_cache import DualCache
|
|||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
from litellm.constants import DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL
|
||||
from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec
|
||||
from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec, CacheValue
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from opentelemetry.trace import Span
|
||||
|
|
@ -223,14 +223,14 @@ class UserApiKeyCache(DualCache):
|
|||
return None
|
||||
return decoded
|
||||
|
||||
def set_cache(self, key: str | None, value: object, local_only: bool = False, **kwargs: object):
|
||||
def set_cache(self, key: str | None, value: CacheValue, local_only: bool = False, **kwargs: object):
|
||||
model_type: Final = cast(type[BaseModel] | None, kwargs.pop("model_type", None))
|
||||
payload: Final[object] = CacheCodec.serialize(value, model_type=model_type)
|
||||
if key is not None and is_user_key_cache_key(key):
|
||||
return self.key_object_cache.set_cache(key=key, value=payload, local_only=local_only, **kwargs)
|
||||
return super().set_cache(key=key, value=payload, local_only=local_only, **kwargs)
|
||||
|
||||
async def async_set_cache(self, key: str | None, value: object, local_only: bool = False, **kwargs: object):
|
||||
async def async_set_cache(self, key: str | None, value: CacheValue, local_only: bool = False, **kwargs: object):
|
||||
"""Inside a request the Redis SET rides the request's pipeline (memory is written at once); anywhere
|
||||
else, or with options the pipeline does not carry, it goes to Redis directly as before."""
|
||||
model_type: Final = cast(type[BaseModel] | None, kwargs.pop("model_type", None))
|
||||
|
|
@ -291,7 +291,7 @@ class UserApiKeyCache(DualCache):
|
|||
self.budget_windows_absent_cache.flush_cache()
|
||||
|
||||
async def async_set_cache_pipeline(
|
||||
self, cache_list: Sequence[tuple[str, object]], local_only: bool = False, **kwargs: object
|
||||
self, cache_list: Sequence[tuple[str, CacheValue]], local_only: bool = False, **kwargs: object
|
||||
) -> None:
|
||||
"""
|
||||
Batch writes with the same Codec boundary as ``async_set_cache`` without
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import datetime
|
||||
import hashlib
|
||||
import json
|
||||
from typing import Any, Final
|
||||
|
|
@ -251,6 +252,20 @@ class TestUserApiKeyCache:
|
|||
with pytest.raises(TypeError):
|
||||
fake.set_cache("k2", {"ok": NotSerializable()})
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_model_holding_a_datetime_is_cached_as_its_json_dump(self):
|
||||
cache = UserApiKeyCache()
|
||||
key_obj = UserAPIKeyAuth(token="abc", expires=datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc))
|
||||
|
||||
await cache.async_set_cache("k", key_obj, model_type=UserAPIKeyAuth)
|
||||
|
||||
stored = cache.get_cache("k", local_only=True)
|
||||
assert stored == key_obj.model_dump(mode="json")
|
||||
assert json.loads(json.dumps(stored)) == stored
|
||||
cached = await cache.async_get_cache("k", model_type=UserAPIKeyAuth)
|
||||
assert cached is not None
|
||||
assert cached.model_dump(mode="json") == key_obj.model_dump(mode="json")
|
||||
|
||||
|
||||
class TestUserKeyObjectPartition:
|
||||
"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue