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

This commit is contained in:
Krrish Dholakia 2024-07-19 17:06:49 -07:00
parent fa7037e48a
commit 35e640076b
8 changed files with 94 additions and 15 deletions

View file

@ -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(