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

(fix) proxy - clear cache after /key/delete

Browse source
This commit is contained in:
ishaan-jaff 2024-02-02 18:40:18 -08:00
parent 4d6ffe4400
commit 9852462d1f
1 changed files with 27 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

25
litellm/proxy/proxy_server.py
View file

@ -527,7 +527,7 @@ async def user_api_key_auth(
) )
# Token passed all checks # Token passed all checks
api_key = valid_token.token api_key = hash_token(valid_token.token)
# Add hashed token to cache # Add hashed token to cache
user_api_key_cache.set_cache(key=api_key, value=valid_token, ttl=60) user_api_key_cache.set_cache(key=api_key, value=valid_token, ttl=60)
@ -1445,8 +1445,10 @@ async def generate_key_helper_fn(
saved_token["expires"], datetime saved_token["expires"], datetime
): ):
saved_token["expires"] = saved_token["expires"].isoformat() saved_token["expires"] = saved_token["expires"].isoformat()
if key_data["token"] is not None and isinstance(key_data["token"], str):
hashed_token = hash_token(key_data["token"])
user_api_key_cache.set_cache( user_api_key_cache.set_cache(
key=key_data["token"], key=hashed_token,
value=LiteLLM_VerificationToken(**saved_token), # type: ignore value=LiteLLM_VerificationToken(**saved_token), # type: ignore
ttl=60, ttl=60,
) )
@ -2665,13 +2667,32 @@ async def delete_key_fn(data: DeleteKeyRequest):
HTTPException: If an error occurs during key deletion. HTTPException: If an error occurs during key deletion.
""" """
try: try:
global user_api_key_cache
keys = data.keys keys = data.keys
if len(keys) == 0:
raise ProxyException(
message=f"No keys provided, passed in: keys={keys}",
type="auth_error",
param="keys",
code=status.HTTP_400_BAD_REQUEST,
)
result = await delete_verification_token(tokens=keys) result = await delete_verification_token(tokens=keys)
verbose_proxy_logger.debug("/key/delete - deleted_keys=", result) verbose_proxy_logger.debug("/key/delete - deleted_keys=", result)
number_deleted_keys = len(result["deleted_keys"]) number_deleted_keys = len(result["deleted_keys"])
assert len(keys) == number_deleted_keys assert len(keys) == number_deleted_keys
for key in keys:
user_api_key_cache.delete_cache(key)
# remove hash token from cache
hashed_token = hash_token(key)
user_api_key_cache.delete_cache(hashed_token)
verbose_proxy_logger.debug(
f"/keys/delete - cache after delete: {user_api_key_cache.in_memory_cache.cache_dict}"
)
return {"deleted_keys": keys} return {"deleted_keys": keys}
except Exception as e: except Exception as e:
if isinstance(e, HTTPException): if isinstance(e, HTTPException):