Skip to content

Commit

Permalink
collab: Unconditionally execute billing checks (#19432)
Browse files Browse the repository at this point in the history
This PR removes the conditional checks around the billing-related
enforcement for LLM completions.

These were just in place to prevent executing any billing code before we
had rolled it out. Now that it is rolled out, we don't need this
conditional execution anymore.

Release Notes:

- N/A
  • Loading branch information
maxdeviant authored and mikayla-maki committed Oct 18, 2024
1 parent 4d6271b commit 897593e
Show file tree
Hide file tree
Showing 2 changed files with 21 additions and 35 deletions.
4 changes: 0 additions & 4 deletions crates/collab/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
52 changes: 21 additions & 31 deletions crates/collab/src/llm.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
));
}
}

Expand Down Expand Up @@ -627,7 +625,6 @@ where
impl<S> Drop for TokenCountingStream<S> {
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);
Expand All @@ -641,14 +638,7 @@ impl<S> Drop for TokenCountingStream<S> {
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(),
)
Expand Down

0 comments on commit 897593e

Please sign in to comment.