1use super::*;
 2
 3#[derive(Debug)]
 4pub struct CreateBillingPreferencesParams {
 5    pub max_monthly_llm_usage_spending_in_cents: i32,
 6    pub model_request_overages_enabled: bool,
 7    pub model_request_overages_spend_limit_in_cents: i32,
 8}
 9
10#[derive(Debug, Default)]
11pub struct UpdateBillingPreferencesParams {
12    pub max_monthly_llm_usage_spending_in_cents: ActiveValue<i32>,
13    pub model_request_overages_enabled: ActiveValue<bool>,
14    pub model_request_overages_spend_limit_in_cents: ActiveValue<i32>,
15}
16
17impl Database {
18    /// Returns the billing preferences for the given user, if they exist.
19    pub async fn get_billing_preferences(
20        &self,
21        user_id: UserId,
22    ) -> Result<Option<billing_preference::Model>> {
23        self.transaction(|tx| async move {
24            Ok(billing_preference::Entity::find()
25                .filter(billing_preference::Column::UserId.eq(user_id))
26                .one(&*tx)
27                .await?)
28        })
29        .await
30    }
31
32    /// Creates new billing preferences for the given user.
33    pub async fn create_billing_preferences(
34        &self,
35        user_id: UserId,
36        params: &CreateBillingPreferencesParams,
37    ) -> Result<billing_preference::Model> {
38        self.transaction(|tx| async move {
39            let preferences = billing_preference::Entity::insert(billing_preference::ActiveModel {
40                user_id: ActiveValue::set(user_id),
41                max_monthly_llm_usage_spending_in_cents: ActiveValue::set(
42                    params.max_monthly_llm_usage_spending_in_cents,
43                ),
44                model_request_overages_enabled: ActiveValue::set(
45                    params.model_request_overages_enabled,
46                ),
47                model_request_overages_spend_limit_in_cents: ActiveValue::set(
48                    params.model_request_overages_spend_limit_in_cents,
49                ),
50                ..Default::default()
51            })
52            .exec_with_returning(&*tx)
53            .await?;
54
55            Ok(preferences)
56        })
57        .await
58    }
59
60    /// Updates the billing preferences for the given user.
61    pub async fn update_billing_preferences(
62        &self,
63        user_id: UserId,
64        params: &UpdateBillingPreferencesParams,
65    ) -> Result<billing_preference::Model> {
66        self.transaction(|tx| async move {
67            let preferences = billing_preference::Entity::update_many()
68                .set(billing_preference::ActiveModel {
69                    max_monthly_llm_usage_spending_in_cents: params
70                        .max_monthly_llm_usage_spending_in_cents
71                        .clone(),
72                    model_request_overages_enabled: params.model_request_overages_enabled.clone(),
73                    model_request_overages_spend_limit_in_cents: params
74                        .model_request_overages_spend_limit_in_cents
75                        .clone(),
76                    ..Default::default()
77                })
78                .filter(billing_preference::Column::UserId.eq(user_id))
79                .exec_with_returning(&*tx)
80                .await?;
81
82            Ok(preferences
83                .into_iter()
84                .next()
85                .ok_or_else(|| anyhow!("billing preferences not found"))?)
86        })
87        .await
88    }
89}