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

refactor(proxy_server.py): fix merge error

Browse source
This commit is contained in:
Krrish Dholakia 2024-02-06 23:44:23 -08:00
parent 9e138b9e4e
commit 184e78772b
1 changed files with 53 additions and 44 deletions
Download patch file Download diff file Expand all files Collapse all files

95
litellm/proxy/proxy_server.py
View file

@ -474,7 +474,6 @@ async def user_api_key_auth(
key=valid_token.user_id, table_name="user" key=valid_token.user_id, table_name="user"
) )
verbose_proxy_logger.debug( verbose_proxy_logger.debug(
f"user_id_information: {user_id_information}" f"user_id_information: {user_id_information}"
) )
@ -884,57 +883,67 @@ async def update_database(
""" """
user_ids = [user_id, litellm_proxy_budget_name] user_ids = [user_id, litellm_proxy_budget_name]
data_list = [] data_list = []
for id in user_ids: try:
if id is None: for id in user_ids:
continue if id is None:
if prisma_client is not None: continue
existing_spend_obj = await prisma_client.get_data(user_id=id) if prisma_client is not None:
elif custom_db_client is not None and id != litellm_proxy_budget_name: existing_spend_obj = await prisma_client.get_data(user_id=id)
existing_spend_obj = await custom_db_client.get_data( elif (
key=id, table_name="user" custom_db_client is not None and id != litellm_proxy_budget_name
) ):
verbose_proxy_logger.debug( existing_spend_obj = await custom_db_client.get_data(
f"Updating existing_spend_obj: {existing_spend_obj}" key=id, table_name="user"
) )
if existing_spend_obj is None: verbose_proxy_logger.debug(
# if user does not exist in LiteLLM_UserTable, create a new user f"Updating existing_spend_obj: {existing_spend_obj}"
existing_spend = 0
max_user_budget = None
if litellm.max_user_budget is not None:
max_user_budget = litellm.max_user_budget
existing_spend_obj = LiteLLM_UserTable(
user_id=id, spend=0, max_budget=max_user_budget, user_email=None
) )
if existing_spend_obj is None: if existing_spend_obj is None:
# if user does not exist in LiteLLM_UserTable, create a new user
existing_spend = 0 existing_spend = 0
max_user_budget = None
if litellm.max_user_budget is not None:
max_user_budget = litellm.max_user_budget
existing_spend_obj = LiteLLM_UserTable( existing_spend_obj = LiteLLM_UserTable(
user_id=id, spend=0, max_budget=None, user_email=None user_id=id,
spend=0,
max_budget=max_user_budget,
user_email=None,
) )
else: if existing_spend_obj is None:
existing_spend = existing_spend_obj.spend 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 # Calculate the new cost by adding the existing cost and response_cost
existing_spend_obj.spend = existing_spend + response_cost existing_spend_obj.spend = existing_spend + response_cost
valid_token = user_api_key_cache.get_cache(key=id) valid_token = user_api_key_cache.get_cache(key=id)
if valid_token is not None and isinstance(valid_token, dict): if valid_token is not None and isinstance(valid_token, dict):
user_api_key_cache.set_cache( user_api_key_cache.set_cache(
key=id, value=existing_spend_obj.json() key=id, value=existing_spend_obj.json()
)
verbose_proxy_logger.debug(
f"new cost: {existing_spend_obj.spend}"
) )
data_list.append(existing_spend_obj)
verbose_proxy_logger.debug(f"new cost: {existing_spend_obj.spend}") # Update the cost column for the given user id
data_list.append(existing_spend_obj) if prisma_client is not None:
await prisma_client.update_data(
# Update the cost column for the given user id data_list=data_list,
if prisma_client is not None: query_type="update_many",
await prisma_client.update_data( table_name="user",
data_list=data_list, query_type="update_many", table_name="user" )
) elif custom_db_client is not None and user_id is not None:
elif custom_db_client is not None and user_id is not None: new_spend = data_list[0].spend
new_spend = data_list[0].spend await custom_db_client.update_data(
await custom_db_client.update_data( key=user_id, value={"spend": new_spend}, table_name="user"
key=user_id, value={"spend": new_spend}, table_name="user" )
)
except Exception as e: except Exception as e:
verbose_proxy_logger.info(f"Update User DB call failed to execute") verbose_proxy_logger.info(f"Update User DB call failed to execute")