collab: Update billing preference endpoints with new preferences (#29458)

Marshall Bowers created 9 months ago

This PR updates the billing preference endpoints with the new
overage-related billing preferences.

Release Notes:

- N/A

Change summary

crates/collab/src/api/billing.rs                    | 33 ++++++++++++++
crates/collab/src/db/queries/billing_preferences.rs | 14 ++++++
2 files changed, 46 insertions(+), 1 deletion(-)

Detailed changes

crates/collab/src/api/billing.rs

@@ -65,6 +65,8 @@ struct GetBillingPreferencesParams {
 #[derive(Debug, Serialize)]
 struct BillingPreferencesResponse {
     max_monthly_llm_usage_spending_in_cents: i32,
+    model_request_overages_enabled: bool,
+    model_request_overages_spend_limit_in_cents: i32,
 }
 
 async fn get_billing_preferences(
@@ -81,16 +83,30 @@ async fn get_billing_preferences(
 
     Ok(Json(BillingPreferencesResponse {
         max_monthly_llm_usage_spending_in_cents: preferences
+            .as_ref()
             .map_or(DEFAULT_MAX_MONTHLY_SPEND.0 as i32, |preferences| {
                 preferences.max_monthly_llm_usage_spending_in_cents
             }),
+        model_request_overages_enabled: preferences.as_ref().map_or(false, |preferences| {
+            preferences.model_request_overages_enabled
+        }),
+        model_request_overages_spend_limit_in_cents: preferences
+            .as_ref()
+            .map_or(0, |preferences| {
+                preferences.model_request_overages_spend_limit_in_cents
+            }),
     }))
 }
 
 #[derive(Debug, Deserialize)]
 struct UpdateBillingPreferencesBody {
     github_user_id: i32,
+    #[serde(default)]
     max_monthly_llm_usage_spending_in_cents: i32,
+    #[serde(default)]
+    model_request_overages_enabled: bool,
+    #[serde(default)]
+    model_request_overages_spend_limit_in_cents: i32,
 }
 
 async fn update_billing_preferences(
@@ -106,6 +122,8 @@ async fn update_billing_preferences(
 
     let max_monthly_llm_usage_spending_in_cents =
         body.max_monthly_llm_usage_spending_in_cents.max(0);
+    let model_request_overages_spend_limit_in_cents =
+        body.model_request_overages_spend_limit_in_cents.max(0);
 
     let billing_preferences =
         if let Some(_billing_preferences) = app.db.get_billing_preferences(user.id).await? {
@@ -116,6 +134,12 @@ async fn update_billing_preferences(
                         max_monthly_llm_usage_spending_in_cents: ActiveValue::set(
                             max_monthly_llm_usage_spending_in_cents,
                         ),
+                        model_request_overages_enabled: ActiveValue::set(
+                            body.model_request_overages_enabled,
+                        ),
+                        model_request_overages_spend_limit_in_cents: ActiveValue::set(
+                            model_request_overages_spend_limit_in_cents,
+                        ),
                     },
                 )
                 .await?
@@ -125,18 +149,22 @@ async fn update_billing_preferences(
                     user.id,
                     &crate::db::CreateBillingPreferencesParams {
                         max_monthly_llm_usage_spending_in_cents,
+                        model_request_overages_enabled: body.model_request_overages_enabled,
+                        model_request_overages_spend_limit_in_cents,
                     },
                 )
                 .await?
         };
 
     SnowflakeRow::new(
-        "Spend Limit Updated",
+        "Billing Preferences Updated",
         Some(user.metrics_id),
         user.admin,
         None,
         json!({
             "user_id": user.id,
+            "model_request_overages_enabled": billing_preferences.model_request_overages_enabled,
+            "model_request_overages_spend_limit_in_cents": billing_preferences.model_request_overages_spend_limit_in_cents,
             "max_monthly_llm_usage_spending_in_cents": billing_preferences.max_monthly_llm_usage_spending_in_cents,
         }),
     )
@@ -149,6 +177,9 @@ async fn update_billing_preferences(
     Ok(Json(BillingPreferencesResponse {
         max_monthly_llm_usage_spending_in_cents: billing_preferences
             .max_monthly_llm_usage_spending_in_cents,
+        model_request_overages_enabled: billing_preferences.model_request_overages_enabled,
+        model_request_overages_spend_limit_in_cents: billing_preferences
+            .model_request_overages_spend_limit_in_cents,
     }))
 }
 

crates/collab/src/db/queries/billing_preferences.rs

@@ -3,11 +3,15 @@ use super::*;
 #[derive(Debug)]
 pub struct CreateBillingPreferencesParams {
     pub max_monthly_llm_usage_spending_in_cents: i32,
+    pub model_request_overages_enabled: bool,
+    pub model_request_overages_spend_limit_in_cents: i32,
 }
 
 #[derive(Debug, Default)]
 pub struct UpdateBillingPreferencesParams {
     pub max_monthly_llm_usage_spending_in_cents: ActiveValue<i32>,
+    pub model_request_overages_enabled: ActiveValue<bool>,
+    pub model_request_overages_spend_limit_in_cents: ActiveValue<i32>,
 }
 
 impl Database {
@@ -37,6 +41,12 @@ impl Database {
                 max_monthly_llm_usage_spending_in_cents: ActiveValue::set(
                     params.max_monthly_llm_usage_spending_in_cents,
                 ),
+                model_request_overages_enabled: ActiveValue::set(
+                    params.model_request_overages_enabled,
+                ),
+                model_request_overages_spend_limit_in_cents: ActiveValue::set(
+                    params.model_request_overages_spend_limit_in_cents,
+                ),
                 ..Default::default()
             })
             .exec_with_returning(&*tx)
@@ -59,6 +69,10 @@ impl Database {
                     max_monthly_llm_usage_spending_in_cents: params
                         .max_monthly_llm_usage_spending_in_cents
                         .clone(),
+                    model_request_overages_enabled: params.model_request_overages_enabled.clone(),
+                    model_request_overages_spend_limit_in_cents: params
+                        .model_request_overages_spend_limit_in_cents
+                        .clone(),
                     ..Default::default()
                 })
                 .filter(billing_preference::Column::UserId.eq(user_id))