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 flush cache when updating user

Browse source
This commit is contained in:
Ishaan Jaff 2024-07-13 12:05:09 -07:00
parent 9d02d51a17
commit 670bf1b98d
1 changed files with 28 additions and 29 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -1,16 +1,21 @@
# What is this? # What is this?
## Helper utils for the management endpoints (keys/users/teams) ## Helper utils for the management endpoints (keys/users/teams)
import uuid
from datetime import datetime from datetime import datetime
from functools import wraps from functools import wraps
from litellm.proxy._types import UserAPIKeyAuth, ManagementEndpointLoggingPayload from typing import Optional
from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
from litellm._logging import verbose_logger
from fastapi import Request from fastapi import Request
from litellm.proxy._types import LiteLLM_TeamTable, Member, UserAPIKeyAuth from litellm._logging import verbose_logger
from litellm.proxy._types import (
LiteLLM_TeamTable,
ManagementEndpointLoggingPayload,
Member,
UserAPIKeyAuth,
)
from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
from litellm.proxy.utils import PrismaClient from litellm.proxy.utils import PrismaClient
import uuid
from typing import Optional
async def add_new_member( async def add_new_member(
@ -32,13 +37,8 @@ async def add_new_member(
await prisma_client.db.litellm_usertable.upsert( await prisma_client.db.litellm_usertable.upsert(
where={"user_id": new_member.user_id}, where={"user_id": new_member.user_id},
data={ data={
"update": { "update": {"teams": {"push": [team_id]}},
"teams": {"push": [team_id]} "create": {"user_id": new_member.user_id, "teams": [team_id]},
},
"create": {
"user_id": new_member.user_id,
"teams": [team_id]
}
}, },
) )
elif new_member.user_email is not None: elif new_member.user_email is not None:
@ -123,23 +123,22 @@ def management_endpoint_wrapper(func):
logging_payload=logging_payload, logging_payload=logging_payload,
parent_otel_span=parent_otel_span, 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
if _http_request: user_api_key_cache.flush_cache()
_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()
except Exception as e: except Exception as e:
# Non-Blocking Exception # Non-Blocking Exception
verbose_logger.debug("Error in management endpoint wrapper: %s", str(e)) verbose_logger.debug("Error in management endpoint wrapper: %s", str(e))