Merge pull request #4799 from BerriAI/litellm_proxy_team_cache_update

fix(user_api_key_auth.py): update valid token cache with updated team object cache

litellm/proxy/_new_secret_config.yaml

@@ -1,13 +1,6 @@
 model_list:
-  - model_name: bad-azure-model
+  - model_name: azure-chatgpt
     litellm_params:
       model: azure/chatgpt-v-2
-      azure_ad_token: ""
-      api_base: os.environ/AZURE_API_BASE
-
-  - model_name: good-openai-model
-    litellm_params:
-      model: gpt-3.5-turbo
-
-litellm_settings:
-  fallbacks: [{"bad-azure-model": ["good-openai-model"]}]
+      api_key: os.environ/AZURE_API_KEY
+      api_base: os.environ/AZURE_API_BASE

litellm/proxy/_types.py

@@ -886,6 +886,7 @@ 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=())
 
@@ -1238,6 +1239,9 @@ class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken):
     end_user_rpm_limit: Optional[int] = None
     end_user_max_budget: Optional[float] = None
 
+    # Time stamps
+    last_refreshed_at: Optional[float] = None  # last time joint view was pulled from db
+
 
 class UserAPIKeyAuth(
     LiteLLM_VerificationTokenView

litellm/proxy/auth/auth_checks.py

@@ -59,7 +59,7 @@ def common_checks(
     6. [OPTIONAL] If 'litellm.max_budget' is set (>0), is proxy under budget
     """
     _model = request_body.get("model", None)
-    if team_object is not None and team_object.blocked == True:
+    if team_object is not None and team_object.blocked is True:
         raise Exception(
             f"Team={team_object.team_id} is blocked. Update via `/team/unblock` if your admin."
         )
@@ -349,6 +349,15 @@ async def get_user_object(
         )
 
 
+async def _cache_team_object(
+    team_id: str,
+    team_table: LiteLLM_TeamTable,
+    user_api_key_cache: DualCache,
+):
+    key = "team_id:{}".format(team_id)
+    await user_api_key_cache.async_set_cache(key=key, value=team_table)
+
+
 @log_to_opentelemetry
 async def get_team_object(
     team_id: str,
@@ -386,7 +395,9 @@ async def get_team_object(
 
         _response = LiteLLM_TeamTable(**response.dict())
         # save the team object to cache
-        await user_api_key_cache.async_set_cache(key=key, value=_response)
+        await _cache_team_object(
+            team_id=team_id, team_table=_response, user_api_key_cache=user_api_key_cache
+        )
 
         return _response
     except Exception as e:

litellm/proxy/auth/user_api_key_auth.py

@@ -453,6 +453,32 @@ async def user_api_key_auth(
 
             return valid_token
 
+        if (
+            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)
+            )
+
+            if (
+                team_obj.last_refreshed_at is not None
+                and valid_token.last_refreshed_at is not None
+                and team_obj.last_refreshed_at > valid_token.last_refreshed_at
+            ):
+                team_obj_dict = team_obj.__dict__
+
+                for k, v in team_obj_dict.items():
+                    field_name = f"team_{k}"
+                    if field_name in valid_token.__fields__:
+                        setattr(valid_token, field_name, v)
+
         try:
             is_master_key_valid = secrets.compare_digest(api_key, master_key)  # type: ignore
         except Exception as e:
@@ -504,7 +530,6 @@ async def user_api_key_auth(
             raise Exception("No connected db.")
 
         ## check for cache hit (In-Memory Cache)
-        original_api_key = api_key  # (Patch: For DynamoDB Backwards Compatibility)
         _user_role = None
         if api_key.startswith("sk-"):
             api_key = hash_token(token=api_key)
@@ -521,10 +546,13 @@ async def user_api_key_auth(
                     parent_otel_span=parent_otel_span,
                     proxy_logging_obj=proxy_logging_obj,
                 )
+
                 if _valid_token is not None:
+                    ## update cached token
                     valid_token = UserAPIKeyAuth(
                         **_valid_token.model_dump(exclude_none=True)
                     )
+
             verbose_proxy_logger.debug("Token from db: %s", valid_token)
         elif valid_token is not None and isinstance(valid_token, UserAPIKeyAuth):
             verbose_proxy_logger.debug("API Key Cache Hit!")

litellm/proxy/management_endpoints/team_endpoints.py

@@ -328,11 +328,13 @@ async def update_team(
     }'
     ```
     """
+    from litellm.proxy.auth.auth_checks import _cache_team_object
     from litellm.proxy.proxy_server import (
         _duration_in_seconds,
         create_audit_log_for_update,
         litellm_proxy_admin_name,
         prisma_client,
+        user_api_key_cache,
     )
 
     if prisma_client is None:
@@ -361,11 +363,22 @@ async def update_team(
         # set the budget_reset_at in DB
         updated_kv["budget_reset_at"] = reset_at
 
-    team_row = await prisma_client.update_data(
-        update_key_values=updated_kv,
-        data=updated_kv,
-        table_name="team",
-        team_id=data.team_id,
+    team_row: Optional[
+        LiteLLM_TeamTable
+    ] = await prisma_client.db.litellm_teamtable.update(
+        where={"team_id": data.team_id}, data=updated_kv  # type: ignore
+    )
+
+    if team_row is None or team_row.team_id is None:
+        raise HTTPException(
+            status_code=400,
+            detail={"error": "Team doesn't exist. Got={}".format(team_row)},
+        )
+
+    await _cache_team_object(
+        team_id=team_row.team_id,
+        team_table=team_row,
+        user_api_key_cache=user_api_key_cache,
     )
 
     # Enterprise Feature - Audit Logging. Enable with litellm.store_audit_logs = True
@@ -392,7 +405,7 @@ async def update_team(
             )
         )
 
-    return team_row
+    return {"team_id": team_row.team_id, "data": team_row}
 
 
 @router.post(

litellm/proxy/utils.py

@@ -1331,7 +1331,9 @@ class PrismaClient:
                             response["team_models"] = []
                         if response["team_blocked"] is None:
                             response["team_blocked"] = False
-                        response = LiteLLM_VerificationTokenView(**response)
+                        response = LiteLLM_VerificationTokenView(
+                            **response, last_refreshed_at=time.time()
+                        )
                         # for prisma we need to cast the expires time to str
                         if response.expires is not None and isinstance(
                             response.expires, datetime

litellm/tests/test_user_api_key_auth.py

@@ -44,3 +44,48 @@ def test_check_valid_ip(
     request = Request(client_ip)
 
     assert _check_valid_ip(allowed_ips, request) == expected_result  # type: ignore
+
+
+@pytest.mark.asyncio
+async def test_check_blocked_team():
+    """
+    cached valid_token obj has team_blocked = true
+
+    cached team obj has team_blocked = false
+
+    assert team is not blocked
+    """
+    import asyncio
+    import time
+
+    from fastapi import Request
+    from starlette.datastructures import URL
+
+    from litellm.proxy._types import LiteLLM_TeamTable, UserAPIKeyAuth
+    from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+    from litellm.proxy.proxy_server import hash_token, user_api_key_cache
+
+    _team_id = "1234"
+    user_key = "sk-12345678"
+
+    valid_token = UserAPIKeyAuth(
+        team_id=_team_id,
+        team_blocked=True,
+        token=hash_token(user_key),
+        last_refreshed_at=time.time(),
+    )
+    await asyncio.sleep(1)
+    team_obj = LiteLLM_TeamTable(
+        team_id=_team_id, blocked=False, last_refreshed_at=time.time()
+    )
+    user_api_key_cache.set_cache(key=hash_token(user_key), value=valid_token)
+    user_api_key_cache.set_cache(key="team_id:{}".format(_team_id), value=team_obj)
+
+    setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache)
+    setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
+    setattr(litellm.proxy.proxy_server, "prisma_client", "hello-world")
+
+    request = Request(scope={"type": "http"})
+    request._url = URL(url="/chat/completions")
+
+    await user_api_key_auth(request=request, api_key="Bearer " + user_key)