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

collab: Unconditionally execute billing checks (#19432)

Browse source 
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
This commit is contained in:
Marshall Bowers 2024-10-18 15:55:28 -04:00 • committed by GitHub
parent 11a82e3347
commit b44bed0115
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 21 additions and 35 deletions
Download patch file Download diff file Expand all files Collapse all files

4
crates/collab/src/lib.rs
View file

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

52
crates/collab/src/llm.rs
View file

@ -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<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);
@ -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(),
)