From 595f42d22a81321d46ffe335f2a70347e77d7bcf Mon Sep 17 00:00:00 2001 From: harish-berri Date: Thu, 23 Apr 2026 23:42:03 +0000 Subject: [PATCH] Refactor caching logic in auth_checks and user_api_key_auth to utilize CacheCodec for serialization and deserialization. Simplify cache retrieval by removing unnecessary type checks and streamline cache storage with consistent key formatting. --- litellm/proxy/auth/auth_checks.py | 25 +++++++++++-------------- litellm/proxy/auth/user_api_key_auth.py | 9 ++++++++- 2 files changed, 19 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 1577b5bda9c..760c0fb12d6 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1627,13 +1627,10 @@ async def _get_team_object_from_cache( if cached_team_obj is None: cached_team_obj = await user_api_key_cache.async_get_cache(key=key) - if cached_team_obj is not None: - if isinstance(cached_team_obj, dict): - return LiteLLM_TeamTableCachedObj(**cached_team_obj) - elif isinstance(cached_team_obj, LiteLLM_TeamTableCachedObj): - return cached_team_obj - - return None + if cached_team_obj is None: + return None + + return CacheCodec.deserialize(cached_team_obj, LiteLLM_TeamTableCachedObj) async def get_team_object( @@ -1709,9 +1706,12 @@ async def _cache_access_object( proxy_logging_obj: Optional[ProxyLogging] = None, ): key = "access_group_id:{}".format(access_group_id) + cache_payload = CacheCodec.serialize( + access_group_table, model_type=LiteLLM_AccessGroupTable + ) await user_api_key_cache.async_set_cache( key=key, - value=access_group_table, + value=cache_payload, ttl=DEFAULT_ACCESS_GROUP_CACHE_TTL, ) @@ -1758,13 +1758,10 @@ async def get_access_object( key = "access_group_id:{}".format(access_group_id) - # Always check cache first cached_access_obj = await user_api_key_cache.async_get_cache(key=key) - if cached_access_obj is not None: - if isinstance(cached_access_obj, dict): - return LiteLLM_AccessGroupTable(**cached_access_obj) - elif isinstance(cached_access_obj, LiteLLM_AccessGroupTable): - return cached_access_obj + deserialized = CacheCodec.deserialize(cached_access_obj, LiteLLM_AccessGroupTable) + if deserialized is not None: + return deserialized # Not in cache - fetch from DB try: diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index a6c8fd02ee1..855ce50ec0c 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -58,6 +58,7 @@ from litellm.proxy.auth.oauth2_check import Oauth2Handler from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.cache_coordinator import EventDrivenCacheCoordinator +from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, _safe_get_request_headers, @@ -1454,8 +1455,14 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 # Only cache when the key is a real team_id (non-team keys must not use key=None). if valid_token.team_id is not None and _team_obj is not None: + # Match get_team_object / spend counters: "team_id:{id}". Serialize for Redis + # (json.dumps) — same as _cache_team_object in auth_checks. + team_cache_key = f"team_id:{valid_token.team_id}" await user_api_key_cache.async_set_cache( - key=valid_token.team_id, value=_team_obj + key=team_cache_key, + value=CacheCodec.serialize( + _team_obj, model_type=LiteLLM_TeamTableCachedObj + ), ) # save team table in cache - used for tpm/rpm limiting - tpm_rpm_limiter.py # Fetch project object if key belongs to a project