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

collab: Take the mode into account when syncing usage to Stripe (#29606)

Browse source 
This PR makes it so we take the mode that was used into account when
syncing usage over to Stripe.

Release Notes:

- N/A
This commit is contained in:
Marshall Bowers 2025-04-29 12:05:55 -04:00 • committed by GitHub
parent 6386336eee
commit 5afb89ca93
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 22 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -26,6 +26,7 @@ use crate::api::events::SnowflakeRow;
use crate::db::billing_subscription::{
StripeCancellationReason, StripeSubscriptionStatus, SubscriptionKind,
};
use crate::llm::db::subscription_usage_meter::CompletionMode;
use crate::llm::{DEFAULT_MAX_MONTHLY_SPEND, FREE_TIER_MONTHLY_SPENDING_LIMIT};
use crate::rpc::{ResultExt as _, Server};
use crate::{AppState, Cents, Error, Result};
@ -1372,6 +1373,9 @@ async fn sync_model_request_usage_with_stripe(
let claude_3_7_sonnet = stripe_billing
.find_price_by_lookup_key("claude-3-7-sonnet-requests")
.await?;
let claude_3_7_sonnet_max = stripe_billing
.find_price_by_lookup_key("claude-3-7-sonnet-requests-max")
.await?;
for (usage_meter, usage) in usage_meters {
maybe!(async {
@ -1397,7 +1401,12 @@ async fn sync_model_request_usage_with_stripe(
let (price_id, meter_event_name) = match model.name.as_str() {
"claude-3-5-sonnet" => (&claude_3_5_sonnet.id, "claude_3_5_sonnet/requests"),
"claude-3-7-sonnet" => (&claude_3_7_sonnet.id, "claude_3_7_sonnet/requests"),
"claude-3-7-sonnet" => match usage_meter.mode {
CompletionMode::Normal => (&claude_3_7_sonnet.id, "claude_3_7_sonnet/requests"),
CompletionMode::Max => {
(&claude_3_7_sonnet_max.id, "claude_3_7_sonnet/requests/max")
}
},
model_name => {
bail!("Attempted to sync usage meter for unsupported model: {model_name:?}")
}