From 8c910540edae12bd8e12ba89a3bcaf6abc87aa42 Mon Sep 17 00:00:00 2001 From: Antonio Scandurra Date: Thu, 17 Oct 2024 13:09:50 +0200 Subject: [PATCH] Subtract FREE_TIER_MONTHLY_SPENDING_LIMIT from reported monthly spend (#19358) Release Notes: - N/A --- crates/collab/src/api/billing.rs | 5 +- crates/collab/src/cents.rs | 2 + crates/collab/src/llm.rs | 4 +- crates/collab/src/llm/db/queries/usages.rs | 2 +- .../collab/src/llm/db/tests/billing_tests.rs | 86 ++++++++----------- 5 files changed, 47 insertions(+), 52 deletions(-) diff --git a/crates/collab/src/api/billing.rs b/crates/collab/src/api/billing.rs index 82a6786ff0d7c..5e167a668c4f5 100644 --- a/crates/collab/src/api/billing.rs +++ b/crates/collab/src/api/billing.rs @@ -19,7 +19,7 @@ use stripe::{ }; use util::ResultExt; -use crate::llm::DEFAULT_MAX_MONTHLY_SPEND; +use crate::llm::{DEFAULT_MAX_MONTHLY_SPEND, FREE_TIER_MONTHLY_SPENDING_LIMIT}; use crate::rpc::{ResultExt as _, Server}; use crate::{ db::{ @@ -702,7 +702,8 @@ async fn get_monthly_spend( let monthly_spend = llm_db .get_user_spending_for_month(user.id, Utc::now()) - .await?; + .await? + .saturating_sub(FREE_TIER_MONTHLY_SPENDING_LIMIT); Ok(Json(GetMonthlySpendResponse { monthly_spend_in_cents: monthly_spend.0 as i32, diff --git a/crates/collab/src/cents.rs b/crates/collab/src/cents.rs index 917177bc51219..defbcea4e26d3 100644 --- a/crates/collab/src/cents.rs +++ b/crates/collab/src/cents.rs @@ -10,6 +10,8 @@ Copy, derive_more::Add, derive_more::AddAssign, + derive_more::Sub, + derive_more::SubAssign, )] pub struct Cents(pub u32); diff --git a/crates/collab/src/llm.rs b/crates/collab/src/llm.rs index e54c9e133e3fc..b782db58706df 100644 --- a/crates/collab/src/llm.rs +++ b/crates/collab/src/llm.rs @@ -469,7 +469,9 @@ async fn check_usage_limit( )); } - if usage.spending_this_month >= Cents(claims.max_monthly_spend_in_cents) { + if (usage.spending_this_month - FREE_TIER_MONTHLY_SPENDING_LIMIT) + >= Cents(claims.max_monthly_spend_in_cents) + { return Err(Error::Http( StatusCode::FORBIDDEN, "Maximum spending limit reached for this month.".to_string(), diff --git a/crates/collab/src/llm/db/queries/usages.rs b/crates/collab/src/llm/db/queries/usages.rs index 1e9204b02b21c..5883bcef57132 100644 --- a/crates/collab/src/llm/db/queries/usages.rs +++ b/crates/collab/src/llm/db/queries/usages.rs @@ -412,7 +412,7 @@ impl LlmDatabase { if !is_staff && spending_this_month > FREE_TIER_MONTHLY_SPENDING_LIMIT && has_llm_subscription - && spending_this_month <= max_monthly_spend + && (spending_this_month - FREE_TIER_MONTHLY_SPENDING_LIMIT) <= max_monthly_spend { billing_event::ActiveModel { id: ActiveValue::not_set(), diff --git a/crates/collab/src/llm/db/tests/billing_tests.rs b/crates/collab/src/llm/db/tests/billing_tests.rs index 2e71df5129764..88551dd5f85e9 100644 --- a/crates/collab/src/llm/db/tests/billing_tests.rs +++ b/crates/collab/src/llm/db/tests/billing_tests.rs @@ -1,10 +1,7 @@ use crate::{ db::UserId, llm::{ - db::{ - queries::{providers::ModelParams, usages::Usage}, - LlmDatabase, TokenUsage, - }, + db::{queries::providers::ModelParams, LlmDatabase, TokenUsage}, FREE_TIER_MONTHLY_SPENDING_LIMIT, }, test_llm_db, Cents, @@ -76,29 +73,9 @@ async fn test_billing_limit_exceeded(db: &mut LlmDatabase) { // Verify the recorded usage and spending let recorded_usage = db.get_usage(user_id, provider, model, now).await.unwrap(); - // Verify that we exceeded the free tier usage - assert!( - recorded_usage.spending_this_month > FREE_TIER_MONTHLY_SPENDING_LIMIT, - "Expected spending to exceed free tier limit" - ); - - assert_eq!( - recorded_usage, - Usage { - requests_this_minute: 1, - tokens_this_minute: tokens_to_use, - tokens_this_day: tokens_to_use, - tokens_this_month: TokenUsage { - input: tokens_to_use, - input_cache_creation: 0, - input_cache_read: 0, - output: 0, - }, - spending_this_month: Cents::new(1050), - lifetime_spending: Cents::new(1050), - } - ); + assert_eq!(recorded_usage.spending_this_month, Cents::new(1050)); + assert!(recorded_usage.spending_this_month > FREE_TIER_MONTHLY_SPENDING_LIMIT); // Verify that there is one `billing_event` record let billing_events = db.get_billing_events().await.unwrap(); @@ -111,7 +88,35 @@ async fn test_billing_limit_exceeded(db: &mut LlmDatabase) { assert_eq!(billing_event.input_cache_read_tokens, 0); assert_eq!(billing_event.output_tokens, 0); - let tokens_to_exceed = 20_000_000; // This will cost $1.00 more, pushing us from $10.50 to $11.50, which is over the $11 monthly maximum limit + // Record usage that puts us at $20.50 + let usage_2 = TokenUsage { + input: 200_000_000, // This will cost $10 more, pushing us from $10.50 to $20.50, + input_cache_creation: 0, + input_cache_read: 0, + output: 0, + }; + db.record_usage( + user_id, + false, + provider, + model, + usage_2, + true, + max_monthly_spend, + now, + ) + .await + .unwrap(); + + // Verify the updated usage and spending + let updated_usage = db.get_usage(user_id, provider, model, now).await.unwrap(); + assert_eq!(updated_usage.spending_this_month, Cents::new(2050)); + + // Verify that there are now two billing events + let billing_events = db.get_billing_events().await.unwrap(); + assert_eq!(billing_events.len(), 2); + + let tokens_to_exceed = 20_000_000; // This will cost $1.00 more, pushing us from $20.50 to $21.50, which is over the $11 monthly maximum limit let usage_exceeding = TokenUsage { input: tokens_to_exceed, input_cache_creation: 0, @@ -132,27 +137,12 @@ async fn test_billing_limit_exceeded(db: &mut LlmDatabase) { ) .await .unwrap(); - - // Verify that there is still one billing record - let billing_events = db.get_billing_events().await.unwrap(); - assert_eq!(billing_events.len(), 1); - // Verify the updated usage and spending let updated_usage = db.get_usage(user_id, provider, model, now).await.unwrap(); - assert_eq!( - updated_usage, - Usage { - requests_this_minute: 2, - tokens_this_minute: tokens_to_use + tokens_to_exceed, - tokens_this_day: tokens_to_use + tokens_to_exceed, - tokens_this_month: TokenUsage { - input: tokens_to_use + tokens_to_exceed, - input_cache_creation: 0, - input_cache_read: 0, - output: 0, - }, - spending_this_month: Cents::new(1150), - lifetime_spending: Cents::new(1150), - } - ); + assert_eq!(updated_usage.spending_this_month, Cents::new(2150)); + + // Verify that we never exceed the user max spending for the user + // and avoid charging them. + let billing_events = db.get_billing_events().await.unwrap(); + assert_eq!(billing_events.len(), 2); }