diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b49890fe7..dda0d59c9 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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: diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index df840a9ee..101cf9b7f 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -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 diff --git a/schema.prisma b/schema.prisma index df840a9ee..101cf9b7f 100644 --- a/schema.prisma +++ b/schema.prisma @@ -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