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

Marshall Bowers created 7 months ago

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

Release Notes:

- N/A

Change summary

crates/collab/src/api/billing.rs                            | 11 +++++
crates/collab/src/llm/db/tables/subscription_usage_meter.rs | 12 +++++++
2 files changed, 22 insertions(+), 1 deletion(-)

Detailed changes

crates/collab/src/api/billing.rs

@@ -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:?}")
                 }

crates/collab/src/llm/db/tables/subscription_usage_meter.rs

@@ -1,4 +1,5 @@
 use sea_orm::entity::prelude::*;
+use serde::Serialize;
 
 use crate::llm::db::ModelId;
 
@@ -9,6 +10,7 @@ pub struct Model {
     pub id: i32,
     pub subscription_usage_id: i32,
     pub model_id: ModelId,
+    pub mode: CompletionMode,
     pub requests: i32,
 }
 
@@ -41,3 +43,13 @@ impl Related<super::model::Entity> for Entity {
 }
 
 impl ActiveModelBehavior for ActiveModel {}
+
+#[derive(Eq, PartialEq, Copy, Clone, Debug, EnumIter, DeriveActiveEnum, Hash, Serialize)]
+#[sea_orm(rs_type = "String", db_type = "String(StringLen::None)")]
+#[serde(rename_all = "snake_case")]
+pub enum CompletionMode {
+    #[sea_orm(string_value = "normal")]
+    Normal,
+    #[sea_orm(string_value = "max")]
+    Max,
+}