Skip to content

Commit

Permalink
collab: Don't record billing events if billing is not enabled (#19102)
Browse files Browse the repository at this point in the history
This PR adjusts the billing logic to not write any records to
`billing_events` if:

- The user is staff, as we don't want to bill staff members
- Billing is disabled (we currently enable billing based on the presence
of the Stripe API key)

Release Notes:

- N/A
  • Loading branch information
maxdeviant authored Oct 11, 2024
1 parent b739cfa commit c709b66
Show file tree
Hide file tree
Showing 2 changed files with 11 additions and 2 deletions.
10 changes: 9 additions & 1 deletion crates/collab/src/llm.rs
Original file line number Diff line number Diff line change
Expand Up @@ -639,6 +639,7 @@ 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 @@ -652,7 +653,14 @@ impl<S> Drop for TokenCountingStream<S> {
provider,
&model,
tokens,
claims.has_llm_subscription,
// 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
},
Cents(claims.max_monthly_spend_in_cents),
Utc::now(),
)
Expand Down
3 changes: 2 additions & 1 deletion crates/collab/src/llm/db/queries/usages.rs
Original file line number Diff line number Diff line change
Expand Up @@ -409,7 +409,8 @@ impl LlmDatabase {
monthly_usage.output_tokens as usize,
);

if spending_this_month > FREE_TIER_MONTHLY_SPENDING_LIMIT
if !is_staff
&& spending_this_month > FREE_TIER_MONTHLY_SPENDING_LIMIT
&& has_llm_subscription
&& spending_this_month <= max_monthly_spend
{
Expand Down

0 comments on commit c709b66

Please sign in to comment.