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

(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:

4
litellm/proxy/schema.prisma
View file

@ -24,6 +24,8 @@ model LiteLLM_TeamTable {
budget_reset_at DateTime?
created_at DateTime @default(now()) @map("created_at")
updated_at DateTime @default(now()) @updatedAt @map("updated_at")
model_spend Json @default("{}")
model_max_budget Json @default("{}")
}
// Track spend, rate limit, budget Users
@ -41,6 +43,8 @@ model LiteLLM_UserTable {
budget_duration String?
budget_reset_at DateTime?
allowed_cache_controls String[] @default([])
model_spend Json @default("{}")
model_max_budget Json @default("{}")
}
// Generate Tokens for Proxy

4
schema.prisma
View file

@ -24,6 +24,8 @@ model LiteLLM_TeamTable {
budget_reset_at DateTime?
created_at DateTime @default(now()) @map("created_at")
updated_at DateTime @default(now()) @updatedAt @map("updated_at")
model_spend Json @default("{}")
model_max_budget Json @default("{}")
}
// Track spend, rate limit, budget Users
@ -41,6 +43,8 @@ model LiteLLM_UserTable {
budget_duration String?
budget_reset_at DateTime?
allowed_cache_controls String[] @default([])
model_spend Json @default("{}")
model_max_budget Json @default("{}")
}
// Generate Tokens for Proxy