phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

fix(proxy_server.py): update user cache to with new spend

Browse source
This commit is contained in:
Krrish Dholakia 2024-02-06 23:06:05 -08:00
parent 4174471dac
commit fd9c7a90af
2 changed files with 23 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

2
litellm/_logging.py
View file

@ -3,7 +3,7 @@ import logging
set_verbose = False
# Create a handler for the logger (you may need to adapt this based on your needs)
handler = logging.FileHandler("log_file.txt")
handler = logging.StreamHandler()
handler.setLevel(logging.DEBUG)
# Create a formatter and set it for the handler

29
litellm/proxy/proxy_server.py
View file

@ -432,10 +432,18 @@ async def user_api_key_auth(
# Check 2. If user_id for this token is in budget
## Check 2.5 If global proxy is in budget
if valid_token.user_id is not None:
user_id_information = user_api_key_cache.get_cache(
key=valid_token.user_id
)
if user_id_information is None:
user_id_list = [valid_token.user_id, litellm_proxy_budget_name]
user_id_information = None
for id in user_id_list:
value = user_api_key_cache.get_cache(key=id)
if value is not None:
if user_id_information is None:
user_id_information = []
user_id_information.append(value)
if user_id_information is None or (
isinstance(user_id_information, list)
and len(user_id_information) < 2
):
if prisma_client is not None:
user_id_information = await prisma_client.get_data(
user_id_list=[
@ -445,13 +453,14 @@ async def user_api_key_auth(
table_name="user",
query_type="find_all",
)
for _id in user_id_information:
user_api_key_cache.set_cache(
key=_id["user_id"], value=_id, ttl=600
)
if custom_db_client is not None:
user_id_information = await custom_db_client.get_data(
key=valid_token.user_id, table_name="user"
)
user_api_key_cache.set_cache(
key=valid_token.user_id, value=user_id_information, ttl=600
)
verbose_proxy_logger.debug(
f"user_id_information: {user_id_information}"
@ -879,6 +888,12 @@ async def update_database(
# Calculate the new cost by adding the existing cost and response_cost
existing_spend_obj.spend = existing_spend + response_cost
valid_token = user_api_key_cache.get_cache(key=id)
if valid_token is not None and isinstance(valid_token, dict):
user_api_key_cache.set_cache(
key=id, value=existing_spend_obj.json()
)
verbose_proxy_logger.debug(f"new cost: {existing_spend_obj.spend}")
data_list.append(existing_spend_obj)