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

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

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

Release Notes:

- N/A
This commit is contained in:
Marshall Bowers 2025-04-25 23:11:44 -04:00 • committed by GitHub
parent 9a325a23e5
commit 539f4f1576
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 46 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

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

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

14
crates/collab/src/db/queries/billing_preferences.rs
View file

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