diff --git a/crates/collab/src/lib.rs b/crates/collab/src/lib.rs index c82968efec..78f514adb7 100644 --- a/crates/collab/src/lib.rs +++ b/crates/collab/src/lib.rs @@ -198,10 +198,6 @@ impl Config { } } - pub fn is_llm_billing_enabled(&self) -> bool { - self.stripe_api_key.is_some() - } - #[cfg(test)] pub fn test() -> Self { Self { diff --git a/crates/collab/src/llm.rs b/crates/collab/src/llm.rs index b782db5870..9ee31ab3d1 100644 --- a/crates/collab/src/llm.rs +++ b/crates/collab/src/llm.rs @@ -460,29 +460,27 @@ async fn check_usage_limit( ) .await?; - if state.config.is_llm_billing_enabled() { - if usage.spending_this_month >= FREE_TIER_MONTHLY_SPENDING_LIMIT { - if !claims.has_llm_subscription { - return Err(Error::http( - StatusCode::PAYMENT_REQUIRED, - "Maximum spending limit reached for this month.".to_string(), - )); - } + if usage.spending_this_month >= FREE_TIER_MONTHLY_SPENDING_LIMIT { + if !claims.has_llm_subscription { + return Err(Error::http( + StatusCode::PAYMENT_REQUIRED, + "Maximum spending limit reached for this month.".to_string(), + )); + } - if (usage.spending_this_month - FREE_TIER_MONTHLY_SPENDING_LIMIT) - >= Cents(claims.max_monthly_spend_in_cents) - { - return Err(Error::Http( - StatusCode::FORBIDDEN, - "Maximum spending limit reached for this month.".to_string(), - [( - HeaderName::from_static(MAX_LLM_MONTHLY_SPEND_REACHED_HEADER_NAME), - HeaderValue::from_static("true"), - )] - .into_iter() - .collect(), - )); - } + if (usage.spending_this_month - FREE_TIER_MONTHLY_SPENDING_LIMIT) + >= Cents(claims.max_monthly_spend_in_cents) + { + return Err(Error::Http( + StatusCode::FORBIDDEN, + "Maximum spending limit reached for this month.".to_string(), + [( + HeaderName::from_static(MAX_LLM_MONTHLY_SPEND_REACHED_HEADER_NAME), + HeaderValue::from_static("true"), + )] + .into_iter() + .collect(), + )); } } @@ -627,7 +625,6 @@ where impl Drop for TokenCountingStream { fn drop(&mut self) { let state = self.state.clone(); - let is_llm_billing_enabled = state.config.is_llm_billing_enabled(); let claims = self.claims.clone(); let provider = self.provider; let model = std::mem::take(&mut self.model); @@ -641,14 +638,7 @@ impl Drop for TokenCountingStream { provider, &model, tokens, - // We're passing `false` here if LLM billing is not enabled - // so that we don't write any records to the - // `billing_events` table until we're ready to bill users. - if is_llm_billing_enabled { - claims.has_llm_subscription - } else { - false - }, + claims.has_llm_subscription, Cents(claims.max_monthly_spend_in_cents), Utc::now(), )