fix(auth_checks.py): fix redis usage for team cached objects

litellm/proxy/_new_secret_config.yaml

@@ -3,7 +3,7 @@ model_list:
     litellm_params:
       model: "*"
 
-# litellm_settings:
-#   cache: true
-#   cache_params:
-#     type: redis
+litellm_settings:
+  cache: true
+  cache_params:
+    type: redis

litellm/proxy/_types.py

@@ -910,7 +910,6 @@ class LiteLLM_TeamTable(TeamBase):
     budget_duration: Optional[str] = None
     budget_reset_at: Optional[datetime] = None
     model_id: Optional[int] = None
-    last_refreshed_at: Optional[float] = None
 
     model_config = ConfigDict(protected_namespaces=())
 
@@ -936,6 +935,10 @@ class LiteLLM_TeamTable(TeamBase):
         return values
 
 
+class LiteLLM_TeamTableCachedObj(LiteLLM_TeamTable):
+    last_refreshed_at: Optional[float] = None
+
+
 class TeamRequest(LiteLLMBase):
     teams: List[str]
 

litellm/proxy/auth/auth_checks.py

@@ -8,6 +8,7 @@ Run checks for:
 2. If user is in budget 
 3. If end_user ('user' passed to /chat/completions, /embeddings endpoint) is in budget 
 """
+import time
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, Literal, Optional
 
@@ -19,6 +20,7 @@ from litellm.proxy._types import (
     LiteLLM_JWTAuth,
     LiteLLM_OrganizationTable,
     LiteLLM_TeamTable,
+    LiteLLM_TeamTableCachedObj,
     LiteLLM_UserTable,
     LiteLLMRoutes,
     LitellmUserRoles,
@@ -368,11 +370,15 @@ async def get_user_object(
 
 async def _cache_team_object(
     team_id: str,
-    team_table: LiteLLM_TeamTable,
+    team_table: LiteLLM_TeamTableCachedObj,
     user_api_key_cache: DualCache,
     proxy_logging_obj: Optional[ProxyLogging],
 ):
     key = "team_id:{}".format(team_id)
+
+    ## CACHE REFRESH TIME!
+    team_table.last_refreshed_at = time.time()
+
     value = team_table.model_dump_json(exclude_unset=True)
     await user_api_key_cache.async_set_cache(key=key, value=value)
 
@@ -390,7 +396,7 @@ async def get_team_object(
     user_api_key_cache: DualCache,
     parent_otel_span: Optional[Span] = None,
     proxy_logging_obj: Optional[ProxyLogging] = None,
-) -> LiteLLM_TeamTable:
+) -> LiteLLM_TeamTableCachedObj:
     """
     - Check if team id in proxy Team Table
     - if valid, return LiteLLM_TeamTable object with defined limits
@@ -404,20 +410,26 @@ async def get_team_object(
     # check if in cache
     key = "team_id:{}".format(team_id)
 
-    cached_team_obj: Optional[LiteLLM_TeamTable] = None
+    cached_team_obj: Optional[LiteLLM_TeamTableCachedObj] = None
+
     ## CHECK REDIS CACHE ##
-    if proxy_logging_obj is not None:
-        cached_team_obj = await proxy_logging_obj.internal_usage_cache.async_get_cache(
+    if (
+        proxy_logging_obj is not None
+        and proxy_logging_obj.internal_usage_cache.redis_cache is not None
+    ):
+        cached_team_obj = (
+            await proxy_logging_obj.internal_usage_cache.redis_cache.async_get_cache(
                 key=key
             )
+        )
 
     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_TeamTable(**cached_team_obj)
-        elif isinstance(cached_team_obj, LiteLLM_TeamTable):
+            return LiteLLM_TeamTableCachedObj(**cached_team_obj)
+        elif isinstance(cached_team_obj, LiteLLM_TeamTableCachedObj):
             return cached_team_obj
     # else, check db
     try:
@@ -428,7 +440,7 @@ async def get_team_object(
         if response is None:
             raise Exception
 
-        _response = LiteLLM_TeamTable(**response.dict())
+        _response = LiteLLM_TeamTableCachedObj(**response.dict())
         # save the team object to cache
         await _cache_team_object(
             team_id=team_id,

litellm/proxy/auth/user_api_key_auth.py

@@ -461,14 +461,14 @@ async def user_api_key_auth(
             valid_token is not None
             and isinstance(valid_token, UserAPIKeyAuth)
             and valid_token.team_id is not None
-            and user_api_key_cache.get_cache(
-                key="team_id:{}".format(valid_token.team_id)
-            )
-            is not None
         ):
             ## UPDATE TEAM VALUES BASED ON CACHED TEAM OBJECT - allows `/team/update` values to work for cached token
-            team_obj: LiteLLM_TeamTable = user_api_key_cache.get_cache(
-                key="team_id:{}".format(valid_token.team_id)
+            team_obj: LiteLLM_TeamTableCachedObj = await get_team_object(
+                team_id=valid_token.team_id,
+                prisma_client=prisma_client,
+                user_api_key_cache=user_api_key_cache,
+                parent_otel_span=parent_otel_span,
+                proxy_logging_obj=proxy_logging_obj,
             )
 
             if (

litellm/proxy/management_endpoints/team_endpoints.py

@@ -18,6 +18,7 @@ from litellm.proxy._types import (
     LiteLLM_AuditLogs,
     LiteLLM_ModelTable,
     LiteLLM_TeamTable,
+    LiteLLM_TeamTableCachedObj,
     LitellmTableNames,
     LitellmUserRoles,
     Member,
@@ -379,7 +380,7 @@ async def update_team(
 
     await _cache_team_object(
         team_id=team_row.team_id,
-        team_table=team_row,
+        team_table=LiteLLM_TeamTableCachedObj(**team_row.model_dump()),
         user_api_key_cache=user_api_key_cache,
         proxy_logging_obj=proxy_logging_obj,
     )