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

collab: Add a Cents type (#18935)

Browse source 
This PR adds a new `Cents` type that can be used to represent a monetary
value in cents.

This cuts down on the primitive obsession we were using when dealing
with money in the billing code.

Release Notes:

- N/A
This commit is contained in:
Marshall Bowers 2024-10-09 14:22:32 -04:00 • committed by GitHub
parent bc23d1e666
commit 817a41c4dc
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 120 additions and 39 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -4,7 +4,7 @@ mod telemetry;
mod token;
use crate::{
api::CloudflareIpCountryHeader, build_clickhouse_client, db::UserId, executor::Executor,
api::CloudflareIpCountryHeader, build_clickhouse_client, db::UserId, executor::Executor, Cents,
Config, Error, Result,
};
use anyhow::{anyhow, Context as _};
@ -439,12 +439,10 @@ fn normalize_model_name(known_models: Vec<String>, name: String) -> String {
}
/// The maximum monthly spending an individual user can reach before they have to pay.
pub const MONTHLY_SPENDING_LIMIT_IN_CENTS: usize = 5 * 100;
pub const MONTHLY_SPENDING_LIMIT: Cents = Cents::from_dollars(5);
/// The maximum lifetime spending an individual user can reach before being cut off.
///
/// Represented in cents.
const LIFETIME_SPENDING_LIMIT_IN_CENTS: usize = 1_000 * 100;
const LIFETIME_SPENDING_LIMIT: Cents = Cents::from_dollars(1_000);
async fn check_usage_limit(
state: &Arc<LlmState>,
@ -464,7 +462,7 @@ async fn check_usage_limit(
.await?;
if state.config.is_llm_billing_enabled() {
if usage.spending_this_month >= MONTHLY_SPENDING_LIMIT_IN_CENTS {
if usage.spending_this_month >= MONTHLY_SPENDING_LIMIT {
if !claims.has_llm_subscription.unwrap_or(false) {
return Err(Error::http(
StatusCode::PAYMENT_REQUIRED,
@ -475,7 +473,7 @@ async fn check_usage_limit(
}
// TODO: Remove this once we've rolled out monthly spending limits.
if usage.lifetime_spending >= LIFETIME_SPENDING_LIMIT_IN_CENTS {
if usage.lifetime_spending >= LIFETIME_SPENDING_LIMIT {
return Err(Error::http(
StatusCode::FORBIDDEN,
"Maximum spending limit reached.".to_string(),
@ -690,8 +688,8 @@ impl<S> Drop for TokenCountingStream<S> {
.cache_read_input_tokens_this_month
as u64,
output_tokens_this_month: usage.output_tokens_this_month as u64,
spending_this_month: usage.spending_this_month as u64,
lifetime_spending: usage.lifetime_spending as u64,
spending_this_month: usage.spending_this_month.0 as u64,
lifetime_spending: usage.lifetime_spending.0 as u64,
},
)
.await