forked from phoenix/litellm-mirror
fix(proxy_server.py): prisma client fixes for high traffic
This commit is contained in:
Krrish Dholakia
parent
d1549cb2f3
commit
b6adeec347
6 changed files with 224 additions and 114 deletions
|
|
@ -432,15 +432,25 @@ 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:
|
||||
if prisma_client is not None:
|
||||
user_id_information = await prisma_client.get_data(
|
||||
user_id_list=[valid_token.user_id, litellm_proxy_budget_name],
|
||||
table_name="user",
|
||||
query_type="find_all",
|
||||
)
|
||||
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_id_information = user_api_key_cache.get_cache(
|
||||
key=valid_token.user_id
|
||||
)
|
||||
if user_id_information is None:
|
||||
if prisma_client is not None:
|
||||
user_id_information = await prisma_client.get_data(
|
||||
user_id_list=[
|
||||
valid_token.user_id,
|
||||
litellm_proxy_budget_name,
|
||||
],
|
||||
table_name="user",
|
||||
query_type="find_all",
|
||||
)
|
||||
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(
|
||||
|
|
@ -544,7 +554,7 @@ async def user_api_key_auth(
|
|||
api_key = valid_token.token
|
||||
|
||||
# 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=600)
|
||||
valid_token_dict = _get_pydantic_json_dict(valid_token)
|
||||
valid_token_dict.pop("token", None)
|
||||
"""
|
||||
|
|
@ -837,114 +847,127 @@ async def update_database(
|
|||
- Update that user's row
|
||||
- Update litellm-proxy-budget row (global proxy spend)
|
||||
"""
|
||||
user_ids = [user_id, litellm_proxy_budget_name]
|
||||
data_list = []
|
||||
for id in user_ids:
|
||||
if id is None:
|
||||
continue
|
||||
try:
|
||||
user_ids = [user_id, litellm_proxy_budget_name]
|
||||
data_list = []
|
||||
for id in user_ids:
|
||||
if id is None:
|
||||
continue
|
||||
if prisma_client is not None:
|
||||
existing_spend_obj = await prisma_client.get_data(user_id=id)
|
||||
elif (
|
||||
custom_db_client is not None and id != litellm_proxy_budget_name
|
||||
):
|
||||
existing_spend_obj = await custom_db_client.get_data(
|
||||
key=id, table_name="user"
|
||||
)
|
||||
verbose_proxy_logger.debug(
|
||||
f"Updating existing_spend_obj: {existing_spend_obj}"
|
||||
)
|
||||
if existing_spend_obj is None:
|
||||
existing_spend = 0
|
||||
existing_spend_obj = LiteLLM_UserTable(
|
||||
user_id=id, spend=0, max_budget=None, user_email=None
|
||||
)
|
||||
else:
|
||||
existing_spend = existing_spend_obj.spend
|
||||
|
||||
# Calculate the new cost by adding the existing cost and response_cost
|
||||
existing_spend_obj.spend = existing_spend + response_cost
|
||||
|
||||
verbose_proxy_logger.debug(f"new cost: {existing_spend_obj.spend}")
|
||||
data_list.append(existing_spend_obj)
|
||||
|
||||
# Update the cost column for the given user id
|
||||
if prisma_client is not None:
|
||||
existing_spend_obj = await prisma_client.get_data(user_id=id)
|
||||
elif custom_db_client is not None and id != litellm_proxy_budget_name:
|
||||
existing_spend_obj = await custom_db_client.get_data(
|
||||
key=id, table_name="user"
|
||||
await prisma_client.update_data(
|
||||
data_list=data_list, query_type="update_many", table_name="user"
|
||||
)
|
||||
verbose_proxy_logger.debug(
|
||||
f"Updating existing_spend_obj: {existing_spend_obj}"
|
||||
)
|
||||
if existing_spend_obj is None:
|
||||
existing_spend = 0
|
||||
existing_spend_obj = LiteLLM_UserTable(
|
||||
user_id=id, spend=0, max_budget=None, user_email=None
|
||||
elif custom_db_client is not None and user_id is not None:
|
||||
new_spend = data_list[0].spend
|
||||
await custom_db_client.update_data(
|
||||
key=user_id, value={"spend": new_spend}, table_name="user"
|
||||
)
|
||||
else:
|
||||
existing_spend = existing_spend_obj.spend
|
||||
|
||||
# Calculate the new cost by adding the existing cost and response_cost
|
||||
existing_spend_obj.spend = existing_spend + response_cost
|
||||
|
||||
verbose_proxy_logger.debug(f"new cost: {existing_spend_obj.spend}")
|
||||
data_list.append(existing_spend_obj)
|
||||
|
||||
# Update the cost column for the given user id
|
||||
if prisma_client is not None:
|
||||
await prisma_client.update_data(
|
||||
data_list=data_list, query_type="update_many", table_name="user"
|
||||
)
|
||||
elif custom_db_client is not None and user_id is not None:
|
||||
new_spend = data_list[0].spend
|
||||
await custom_db_client.update_data(
|
||||
key=user_id, value={"spend": new_spend}, table_name="user"
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.info(f"Update User DB call failed to execute")
|
||||
|
||||
### UPDATE KEY SPEND ###
|
||||
async def _update_key_db():
|
||||
verbose_proxy_logger.debug(
|
||||
f"adding spend to key db. Response cost: {response_cost}. Token: {token}."
|
||||
)
|
||||
if prisma_client is not None:
|
||||
# Fetch the existing cost for the given token
|
||||
existing_spend_obj = await prisma_client.get_data(token=token)
|
||||
try:
|
||||
verbose_proxy_logger.debug(
|
||||
f"_update_key_db: existing spend: {existing_spend_obj}"
|
||||
f"adding spend to key db. Response cost: {response_cost}. Token: {token}."
|
||||
)
|
||||
if existing_spend_obj is None:
|
||||
existing_spend = 0
|
||||
else:
|
||||
existing_spend = existing_spend_obj.spend
|
||||
# Calculate the new cost by adding the existing cost and response_cost
|
||||
new_spend = existing_spend + response_cost
|
||||
if prisma_client is not None:
|
||||
# Fetch the existing cost for the given token
|
||||
existing_spend_obj = await prisma_client.get_data(token=token)
|
||||
verbose_proxy_logger.debug(
|
||||
f"_update_key_db: existing spend: {existing_spend_obj}"
|
||||
)
|
||||
if existing_spend_obj is None:
|
||||
existing_spend = 0
|
||||
else:
|
||||
existing_spend = existing_spend_obj.spend
|
||||
# Calculate the new cost by adding the existing cost and response_cost
|
||||
new_spend = existing_spend + response_cost
|
||||
|
||||
verbose_proxy_logger.debug(f"new cost: {new_spend}")
|
||||
# Update the cost column for the given token
|
||||
await prisma_client.update_data(token=token, data={"spend": new_spend})
|
||||
verbose_proxy_logger.debug(f"new cost: {new_spend}")
|
||||
# Update the cost column for the given token
|
||||
await prisma_client.update_data(
|
||||
token=token, data={"spend": new_spend}
|
||||
)
|
||||
|
||||
valid_token = user_api_key_cache.get_cache(key=token)
|
||||
if valid_token is not None:
|
||||
valid_token.spend = new_spend
|
||||
user_api_key_cache.set_cache(key=token, value=valid_token)
|
||||
elif custom_db_client is not None:
|
||||
# Fetch the existing cost for the given token
|
||||
existing_spend_obj = await custom_db_client.get_data(
|
||||
key=token, table_name="key"
|
||||
)
|
||||
verbose_proxy_logger.debug(
|
||||
f"_update_key_db existing spend: {existing_spend_obj}"
|
||||
)
|
||||
if existing_spend_obj is None:
|
||||
existing_spend = 0
|
||||
else:
|
||||
existing_spend = existing_spend_obj.spend
|
||||
# Calculate the new cost by adding the existing cost and response_cost
|
||||
new_spend = existing_spend + response_cost
|
||||
valid_token = user_api_key_cache.get_cache(key=token)
|
||||
if valid_token is not None:
|
||||
valid_token.spend = new_spend
|
||||
user_api_key_cache.set_cache(key=token, value=valid_token)
|
||||
elif custom_db_client is not None:
|
||||
# Fetch the existing cost for the given token
|
||||
existing_spend_obj = await custom_db_client.get_data(
|
||||
key=token, table_name="key"
|
||||
)
|
||||
verbose_proxy_logger.debug(
|
||||
f"_update_key_db existing spend: {existing_spend_obj}"
|
||||
)
|
||||
if existing_spend_obj is None:
|
||||
existing_spend = 0
|
||||
else:
|
||||
existing_spend = existing_spend_obj.spend
|
||||
# Calculate the new cost by adding the existing cost and response_cost
|
||||
new_spend = existing_spend + response_cost
|
||||
|
||||
verbose_proxy_logger.debug(f"new cost: {new_spend}")
|
||||
# Update the cost column for the given token
|
||||
await custom_db_client.update_data(
|
||||
key=token, value={"spend": new_spend}, table_name="key"
|
||||
)
|
||||
verbose_proxy_logger.debug(f"new cost: {new_spend}")
|
||||
# Update the cost column for the given token
|
||||
await custom_db_client.update_data(
|
||||
key=token, value={"spend": new_spend}, table_name="key"
|
||||
)
|
||||
|
||||
valid_token = user_api_key_cache.get_cache(key=token)
|
||||
if valid_token is not None:
|
||||
valid_token.spend = new_spend
|
||||
user_api_key_cache.set_cache(key=token, value=valid_token)
|
||||
valid_token = user_api_key_cache.get_cache(key=token)
|
||||
if valid_token is not None:
|
||||
valid_token.spend = new_spend
|
||||
user_api_key_cache.set_cache(key=token, value=valid_token)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.info(f"Update Key DB Call failed to execute")
|
||||
|
||||
### UPDATE SPEND LOGS ###
|
||||
async def _insert_spend_log_to_db():
|
||||
# Helper to generate payload to log
|
||||
verbose_proxy_logger.debug("inserting spend log to db")
|
||||
payload = get_logging_payload(
|
||||
kwargs=kwargs,
|
||||
response_obj=completion_response,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
try:
|
||||
# Helper to generate payload to log
|
||||
verbose_proxy_logger.debug("inserting spend log to db")
|
||||
payload = get_logging_payload(
|
||||
kwargs=kwargs,
|
||||
response_obj=completion_response,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
|
||||
payload["spend"] = response_cost
|
||||
if prisma_client is not None:
|
||||
await prisma_client.insert_data(data=payload, table_name="spend")
|
||||
payload["spend"] = response_cost
|
||||
if prisma_client is not None:
|
||||
await prisma_client.insert_data(data=payload, table_name="spend")
|
||||
|
||||
elif custom_db_client is not None:
|
||||
await custom_db_client.insert_data(payload, table_name="spend")
|
||||
elif custom_db_client is not None:
|
||||
await custom_db_client.insert_data(payload, table_name="spend")
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.info(f"Update Spend Logs DB failed to execute")
|
||||
|
||||
asyncio.create_task(_update_user_db())
|
||||
asyncio.create_task(_update_key_db())
|
||||
|
|
@ -1534,7 +1557,7 @@ async def generate_key_helper_fn(
|
|||
user_api_key_cache.set_cache(
|
||||
key=hashed_token,
|
||||
value=LiteLLM_VerificationToken(**saved_token), # type: ignore
|
||||
ttl=60,
|
||||
ttl=600,
|
||||
)
|
||||
if prisma_client is not None:
|
||||
## CREATE USER (If necessary)
|
||||
|
|
@ -2979,6 +3002,9 @@ async def spend_user_fn(
|
|||
"/spend/logs",
|
||||
tags=["budget & spend Tracking"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
responses={
|
||||
200: {"model": List[LiteLLM_SpendLogs]},
|
||||
},
|
||||
)
|
||||
async def view_spend_logs(
|
||||
api_key: Optional[str] = fastapi.Query(
|
||||
|
|
@ -3049,7 +3075,8 @@ async def view_spend_logs(
|
|||
query_type="find_unique",
|
||||
key_val={"key": "request_id", "value": request_id},
|
||||
)
|
||||
return [spend_log]
|
||||
response = LiteLLM_SpendLogs_ResponseObject(response=[spend_log])
|
||||
return response
|
||||
elif user_id is not None:
|
||||
spend_log = await prisma_client.get_data(
|
||||
table_name="spend",
|
||||
|
|
@ -3065,7 +3092,7 @@ async def view_spend_logs(
|
|||
table_name="spend", query_type="find_all"
|
||||
)
|
||||
|
||||
return spend_logs
|
||||
return spend_log
|
||||
|
||||
return None
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue