phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 11:14:04 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

(feat) track spend key-model, user-model, team-model

Browse source
This commit is contained in:
ishaan-jaff 2024-02-16 16:32:17 -08:00
parent 2c9d142e42
commit d65c6d3869
3 changed files with 31 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

24
litellm/proxy/proxy_server.py
View file

@ -937,6 +937,16 @@ async def update_database(
# Calculate the new cost by adding the existing cost and response_cost
existing_spend_obj.spend = existing_spend + response_cost
# track cost per model, for the given user
spend_per_model = existing_spend_obj.model_spend or {}
current_model = kwargs.get("model")
if current_model is not None and spend_per_model is not None:
if spend_per_model.get(current_model) is None:
spend_per_model[current_model] = response_cost
else:
spend_per_model[current_model] += response_cost
existing_spend_obj.model_spend = spend_per_model
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(
@ -1001,6 +1011,7 @@ async def update_database(
valid_token = user_api_key_cache.get_cache(key=token)
if valid_token is not None:
valid_token.spend = new_spend
valid_token.model_spend = spend_per_model
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
@ -1080,10 +1091,21 @@ async def update_database(
# Calculate the new cost by adding the existing cost and response_cost
new_spend = existing_spend + response_cost
# track cost per model, for the given team
spend_per_model = existing_spend_obj.model_spend or {}
current_model = kwargs.get("model")
if current_model is not None and spend_per_model is not None:
if spend_per_model.get(current_model) is None:
spend_per_model[current_model] = response_cost
else:
spend_per_model[current_model] += response_cost
verbose_proxy_logger.debug(f"new cost: {new_spend}")
# Update the cost column for the given token
await prisma_client.update_data(
team_id=team_id, data={"spend": new_spend}, table_name="team"
team_id=team_id,
data={"spend": new_spend, "model_spend": spend_per_model},
table_name="team",
)
elif custom_db_client is not None: