phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-25 10:44:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

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

Browse source
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
Download patch file Download diff file Expand all files Collapse all files

22
litellm/proxy/auth/user_api_key_auth.py
View file

@ -453,6 +453,27 @@ 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)
)
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 +525,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)