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

correctly clear cache when updating a user

Browse source
This commit is contained in:
Ishaan Jaff 2024-07-13 12:33:43 -07:00
parent bc91025307
commit 893ed4e5f1
1 changed files with 73 additions and 17 deletions
Download patch file Download diff file Expand all files Collapse all files

90
litellm/proxy/management_helpers/utils.py
View file

@ -8,10 +8,18 @@ from typing import Optional
from fastapi import Request
from litellm._logging import verbose_logger
from litellm.proxy._types import (
from litellm.proxy._types import ( # key request types; user request types; team request types; customer request types
DeleteCustomerRequest,
DeleteTeamRequest,
DeleteUserRequest,
KeyRequest,
LiteLLM_TeamTable,
ManagementEndpointLoggingPayload,
Member,
UpdateCustomerRequest,
UpdateKeyRequest,
UpdateTeamRequest,
UpdateUserRequest,
UserAPIKeyAuth,
)
from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
@ -77,6 +85,66 @@ async def add_new_member(
)
def _delete_user_id_from_cache(kwargs):
from litellm.proxy.proxy_server import user_api_key_cache
if kwargs.get("data") is not None:
update_user_request = kwargs.get("data")
if isinstance(update_user_request, UpdateUserRequest):
user_api_key_cache.delete_cache(key=update_user_request.user_id)
# delete user request
if isinstance(update_user_request, DeleteUserRequest):
for user_id in update_user_request.user_ids:
user_api_key_cache.delete_cache(key=user_id)
pass
def _delete_api_key_from_cache(kwargs):
from litellm.proxy.proxy_server import user_api_key_cache
if kwargs.get("data") is not None:
update_request = kwargs.get("data")
if isinstance(update_request, UpdateKeyRequest):
user_api_key_cache.delete_cache(key=update_request.key)
# delete key request
if isinstance(update_request, KeyRequest):
for key in update_request.keys:
user_api_key_cache.delete_cache(key=key)
pass
def _delete_team_id_from_cache(kwargs):
from litellm.proxy.proxy_server import user_api_key_cache
if kwargs.get("data") is not None:
update_request = kwargs.get("data")
if isinstance(update_request, UpdateTeamRequest):
user_api_key_cache.delete_cache(key=update_request.team_id)
# delete team request
if isinstance(update_request, DeleteTeamRequest):
for team_id in update_request.team_ids:
user_api_key_cache.delete_cache(key=team_id)
pass
def _delete_customer_id_from_cache(kwargs):
from litellm.proxy.proxy_server import user_api_key_cache
if kwargs.get("data") is not None:
update_request = kwargs.get("data")
if isinstance(update_request, UpdateCustomerRequest):
user_api_key_cache.delete_cache(key=update_request.user_id)
# delete customer request
if isinstance(update_request, DeleteCustomerRequest):
for user_id in update_request.user_ids:
user_api_key_cache.delete_cache(key=user_id)
pass
def management_endpoint_wrapper(func):
"""
This wrapper does the following:
@ -123,22 +191,10 @@ def management_endpoint_wrapper(func):
logging_payload=logging_payload,
parent_otel_span=parent_otel_span,
)
if _http_request:
_route = _http_request.url.path
# Flush user_api_key cache if this was an update/delete call to /key, /team, or /user
if _route in [
"/key/update",
"/key/delete",
"/team/update",
"/team/delete",
"/user/update",
"/user/delete",
"/customer/update",
"/customer/delete",
]:
from litellm.proxy.proxy_server import user_api_key_cache
user_api_key_cache.flush_cache()
_delete_api_key_from_cache(kwargs=kwargs)
_delete_user_id_from_cache(kwargs=kwargs)
_delete_team_id_from_cache(kwargs=kwargs)
_delete_customer_id_from_cache(kwargs=kwargs)
except Exception as e:
# Non-Blocking Exception
verbose_logger.debug("Error in management endpoint wrapper: %s", str(e))