Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Subtract FREE_TIER_MONTHLY_SPENDING_LIMIT from reported monthly spend (#19358)

Browse source 
Release Notes:

- N/A
This commit is contained in:
Antonio Scandurra 2024-10-17 13:09:50 +02:00 • committed by GitHub
parent 455f241c6a
commit 8c910540ed
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 47 additions and 52 deletions
Download patch file Download diff file Expand all files Collapse all files

5
crates/collab/src/api/billing.rs
View file

@ -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,

2
crates/collab/src/cents.rs
View file

@ -10,6 +10,8 @@
Copy,
derive_more::Add,
derive_more::AddAssign,
derive_more::Sub,
derive_more::SubAssign,
)]
pub struct Cents(pub u32);

4
crates/collab/src/llm.rs
View file

@ -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(),

2
crates/collab/src/llm/db/queries/usages.rs
View file

@ -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(),

86
crates/collab/src/llm/db/tests/billing_tests.rs
View file

@ -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);
}