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

collab: Update billing code for LLM usage billing (#18879)

Browse source 
This PR reworks our existing billing code in preparation for charging
based on LLM usage.

We aren't yet exercising the new billing-related code outside of
development.

There are some noteworthy changes for our existing LLM usage tracking:

- A new `monthly_usages` table has been added for tracking usage
per-user, per-model, per-month
- The per-month usage measures have been removed, in favor of the
`monthly_usages` table
- All of the per-month metrics in the Clickhouse rows have been changed
from a rolling 30-day window to a calendar month

Release Notes:

- N/A

---------

Co-authored-by: Antonio Scandurra <me@as-cii.com>
Co-authored-by: Richard <richard@zed.dev>
Co-authored-by: Max <max@zed.dev>
This commit is contained in:
Marshall Bowers 2024-10-08 18:29:38 -04:00 • committed by GitHub
parent a95fb8f1f9
commit f861479890
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
15 changed files with 390 additions and 132 deletions
Download patch file Download diff file Expand all files Collapse all files

8
crates/collab/src/llm/db.rs
View file

@ -97,6 +97,14 @@ impl LlmDatabase {
.ok_or_else(|| anyhow!("unknown model {provider:?}:{name}"))?)
}
pub fn model_by_id(&self, id: ModelId) -> Result<&model::Model> {
Ok(self
.models
.values()
.find(|model| model.id == id)
.ok_or_else(|| anyhow!("no model for ID {id:?}"))?)
}
pub fn options(&self) -> &ConnectOptions {
&self.options
}

230
crates/collab/src/llm/db/queries/usages.rs
View file

@ -1,5 +1,5 @@
use crate::db::UserId;
use chrono::Duration;
use chrono::{Datelike, Duration};
use futures::StreamExt as _;
use rpc::LanguageModelProvider;
use sea_orm::QuerySelect;
@ -140,6 +140,46 @@ impl LlmDatabase {
.await
}
pub async fn get_user_spending_for_month(
&self,
user_id: UserId,
now: DateTimeUtc,
) -> Result<usize> {
self.transaction(|tx| async move {
let month = now.date_naive().month() as i32;
let year = now.date_naive().year();
let mut monthly_usages = monthly_usage::Entity::find()
.filter(
monthly_usage::Column::UserId
.eq(user_id)
.and(monthly_usage::Column::Month.eq(month))
.and(monthly_usage::Column::Year.eq(year)),
)
.stream(&*tx)
.await?;
let mut monthly_spending_in_cents = 0;
while let Some(usage) = monthly_usages.next().await {
let usage = usage?;
let Ok(model) = self.model_by_id(usage.model_id) else {
continue;
};
monthly_spending_in_cents += calculate_spending(
model,
usage.input_tokens as usize,
usage.cache_creation_input_tokens as usize,
usage.cache_read_input_tokens as usize,
usage.output_tokens as usize,
);
}
Ok(monthly_spending_in_cents)
})
.await
}
pub async fn get_usage(
&self,
user_id: UserId,
@ -162,6 +202,18 @@ impl LlmDatabase {
.all(&*tx)
.await?;
let month = now.date_naive().month() as i32;
let year = now.date_naive().year();
let monthly_usage = monthly_usage::Entity::find()
.filter(
monthly_usage::Column::UserId
.eq(user_id)
.and(monthly_usage::Column::ModelId.eq(model.id))
.and(monthly_usage::Column::Month.eq(month))
.and(monthly_usage::Column::Year.eq(year)),
)
.one(&*tx)
.await?;
let lifetime_usage = lifetime_usage::Entity::find()
.filter(
lifetime_usage::Column::UserId
@ -177,28 +229,18 @@ impl LlmDatabase {
self.get_usage_for_measure(&usages, now, UsageMeasure::TokensPerMinute)?;
let tokens_this_day =
self.get_usage_for_measure(&usages, now, UsageMeasure::TokensPerDay)?;
let input_tokens_this_month =
self.get_usage_for_measure(&usages, now, UsageMeasure::InputTokensPerMonth)?;
let cache_creation_input_tokens_this_month = self.get_usage_for_measure(
&usages,
now,
UsageMeasure::CacheCreationInputTokensPerMonth,
)?;
let cache_read_input_tokens_this_month = self.get_usage_for_measure(
&usages,
now,
UsageMeasure::CacheReadInputTokensPerMonth,
)?;
let output_tokens_this_month =
self.get_usage_for_measure(&usages, now, UsageMeasure::OutputTokensPerMonth)?;
let spending_this_month = calculate_spending(
model,
input_tokens_this_month,
cache_creation_input_tokens_this_month,
cache_read_input_tokens_this_month,
output_tokens_this_month,
);
let lifetime_spending = if let Some(lifetime_usage) = lifetime_usage {
let spending_this_month = if let Some(monthly_usage) = &monthly_usage {
calculate_spending(
model,
monthly_usage.input_tokens as usize,
monthly_usage.cache_creation_input_tokens as usize,
monthly_usage.cache_read_input_tokens as usize,
monthly_usage.output_tokens as usize,
)
} else {
0
};
let lifetime_spending = if let Some(lifetime_usage) = &lifetime_usage {
calculate_spending(
model,
lifetime_usage.input_tokens as usize,
@ -214,10 +256,18 @@ impl LlmDatabase {
requests_this_minute,
tokens_this_minute,
tokens_this_day,
input_tokens_this_month,
cache_creation_input_tokens_this_month,
cache_read_input_tokens_this_month,
output_tokens_this_month,
input_tokens_this_month: monthly_usage
.as_ref()
.map_or(0, |usage| usage.input_tokens as usize),
cache_creation_input_tokens_this_month: monthly_usage
.as_ref()
.map_or(0, |usage| usage.cache_creation_input_tokens as usize),
cache_read_input_tokens_this_month: monthly_usage
.as_ref()
.map_or(0, |usage| usage.cache_read_input_tokens as usize),
output_tokens_this_month: monthly_usage
.as_ref()
.map_or(0, |usage| usage.output_tokens as usize),
spending_this_month,
lifetime_spending,
})
@ -290,60 +340,68 @@ impl LlmDatabase {
&tx,
)
.await?;
let input_tokens_this_month = self
.update_usage_for_measure(
user_id,
is_staff,
model.id,
&usages,
UsageMeasure::InputTokensPerMonth,
now,
input_token_count,
&tx,
)
.await?;
let cache_creation_input_tokens_this_month = self
.update_usage_for_measure(
user_id,
is_staff,
model.id,
&usages,
UsageMeasure::CacheCreationInputTokensPerMonth,
now,
cache_creation_input_tokens,
&tx,
)
.await?;
let cache_read_input_tokens_this_month = self
.update_usage_for_measure(
user_id,
is_staff,
model.id,
&usages,
UsageMeasure::CacheReadInputTokensPerMonth,
now,
cache_read_input_tokens,
&tx,
)
.await?;
let output_tokens_this_month = self
.update_usage_for_measure(
user_id,
is_staff,
model.id,
&usages,
UsageMeasure::OutputTokensPerMonth,
now,
output_token_count,
&tx,
let month = now.date_naive().month() as i32;
let year = now.date_naive().year();
// Update monthly usage
let monthly_usage = monthly_usage::Entity::find()
.filter(
monthly_usage::Column::UserId
.eq(user_id)
.and(monthly_usage::Column::ModelId.eq(model.id))
.and(monthly_usage::Column::Month.eq(month))
.and(monthly_usage::Column::Year.eq(year)),
)
.one(&*tx)
.await?;
let monthly_usage = match monthly_usage {
Some(usage) => {
monthly_usage::Entity::update(monthly_usage::ActiveModel {
id: ActiveValue::unchanged(usage.id),
input_tokens: ActiveValue::set(
usage.input_tokens + input_token_count as i64,
),
cache_creation_input_tokens: ActiveValue::set(
usage.cache_creation_input_tokens + cache_creation_input_tokens as i64,
),
cache_read_input_tokens: ActiveValue::set(
usage.cache_read_input_tokens + cache_read_input_tokens as i64,
),
output_tokens: ActiveValue::set(
usage.output_tokens + output_token_count as i64,
),
..Default::default()
})
.exec(&*tx)
.await?
}
None => {
monthly_usage::ActiveModel {
user_id: ActiveValue::set(user_id),
model_id: ActiveValue::set(model.id),
month: ActiveValue::set(month),
year: ActiveValue::set(year),
input_tokens: ActiveValue::set(input_token_count as i64),
cache_creation_input_tokens: ActiveValue::set(
cache_creation_input_tokens as i64,
),
cache_read_input_tokens: ActiveValue::set(cache_read_input_tokens as i64),
output_tokens: ActiveValue::set(output_token_count as i64),
..Default::default()
}
.insert(&*tx)
.await?
}
};
let spending_this_month = calculate_spending(
model,
input_tokens_this_month,
cache_creation_input_tokens_this_month,
cache_read_input_tokens_this_month,
output_tokens_this_month,
monthly_usage.input_tokens as usize,
monthly_usage.cache_creation_input_tokens as usize,
monthly_usage.cache_read_input_tokens as usize,
monthly_usage.output_tokens as usize,
);
// Update lifetime usage
@ -406,10 +464,11 @@ impl LlmDatabase {
requests_this_minute,
tokens_this_minute,
tokens_this_day,
input_tokens_this_month,
cache_creation_input_tokens_this_month,
cache_read_input_tokens_this_month,
output_tokens_this_month,
input_tokens_this_month: monthly_usage.input_tokens as usize,
cache_creation_input_tokens_this_month: monthly_usage.cache_creation_input_tokens
as usize,
cache_read_input_tokens_this_month: monthly_usage.cache_read_input_tokens as usize,
output_tokens_this_month: monthly_usage.output_tokens as usize,
spending_this_month,
lifetime_spending,
})
@ -597,7 +656,6 @@ fn calculate_spending(
const MINUTE_BUCKET_COUNT: usize = 12;
const DAY_BUCKET_COUNT: usize = 48;
const MONTH_BUCKET_COUNT: usize = 30;
impl UsageMeasure {
fn bucket_count(&self) -> usize {
@ -605,10 +663,6 @@ impl UsageMeasure {
UsageMeasure::RequestsPerMinute => MINUTE_BUCKET_COUNT,
UsageMeasure::TokensPerMinute => MINUTE_BUCKET_COUNT,
UsageMeasure::TokensPerDay => DAY_BUCKET_COUNT,
UsageMeasure::InputTokensPerMonth => MONTH_BUCKET_COUNT,
UsageMeasure::CacheCreationInputTokensPerMonth => MONTH_BUCKET_COUNT,
UsageMeasure::CacheReadInputTokensPerMonth => MONTH_BUCKET_COUNT,
UsageMeasure::OutputTokensPerMonth => MONTH_BUCKET_COUNT,
}
}
@ -617,10 +671,6 @@ impl UsageMeasure {
UsageMeasure::RequestsPerMinute => Duration::minutes(1),
UsageMeasure::TokensPerMinute => Duration::minutes(1),
UsageMeasure::TokensPerDay => Duration::hours(24),
UsageMeasure::InputTokensPerMonth => Duration::days(30),
UsageMeasure::CacheCreationInputTokensPerMonth => Duration::days(30),
UsageMeasure::CacheReadInputTokensPerMonth => Duration::days(30),
UsageMeasure::OutputTokensPerMonth => Duration::days(30),
}
}

1
crates/collab/src/llm/db/tables.rs
View file

@ -1,5 +1,6 @@
pub mod lifetime_usage;
pub mod model;
pub mod monthly_usage;
pub mod provider;
pub mod revoked_access_token;
pub mod usage;

22
crates/collab/src/llm/db/tables/monthly_usage.rs Normal file
View file

@ -0,0 +1,22 @@
use crate::{db::UserId, llm::db::ModelId};
use sea_orm::entity::prelude::*;
#[derive(Clone, Debug, PartialEq, DeriveEntityModel)]
#[sea_orm(table_name = "monthly_usages")]
pub struct Model {
#[sea_orm(primary_key)]
pub id: i32,
pub user_id: UserId,
pub model_id: ModelId,
pub month: i32,
pub year: i32,
pub input_tokens: i64,
pub cache_creation_input_tokens: i64,
pub cache_read_input_tokens: i64,
pub output_tokens: i64,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {}
impl ActiveModelBehavior for ActiveModel {}

4
crates/collab/src/llm/db/tables/usage_measure.rs
View file

@ -9,10 +9,6 @@ pub enum UsageMeasure {
RequestsPerMinute,
TokensPerMinute,
TokensPerDay,
InputTokensPerMonth,
CacheCreationInputTokensPerMonth,
CacheReadInputTokensPerMonth,
OutputTokensPerMonth,
}
#[derive(Clone, Debug, PartialEq, DeriveEntityModel)]

32
crates/collab/src/llm/db/tests/usage_tests.rs
View file

@ -6,7 +6,7 @@ use crate::{
},
test_llm_db,
};
use chrono::{Duration, Utc};
use chrono::{DateTime, Duration, Utc};
use pretty_assertions::assert_eq;
use rpc::LanguageModelProvider;
@ -29,7 +29,10 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
.await
.unwrap();
let t0 = Utc::now();
// We're using a fixed datetime to prevent flakiness based on the clock.
let t0 = DateTime::parse_from_rfc3339("2024-08-08T22:46:33Z")
.unwrap()
.with_timezone(&Utc);
let user_id = UserId::from_proto(123);
let now = t0;
@ -134,23 +137,10 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
}
);
let t2 = t0 + Duration::days(30);
let now = t2;
let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
assert_eq!(
usage,
Usage {
requests_this_minute: 0,
tokens_this_minute: 0,
tokens_this_day: 0,
input_tokens_this_month: 9000,
cache_creation_input_tokens_this_month: 0,
cache_read_input_tokens_this_month: 0,
output_tokens_this_month: 0,
spending_this_month: 0,
lifetime_spending: 0,
}
);
// We're using a fixed datetime to prevent flakiness based on the clock.
let now = DateTime::parse_from_rfc3339("2024-10-08T22:15:58Z")
.unwrap()
.with_timezone(&Utc);
// Test cache creation input tokens
db.record_usage(user_id, false, provider, model, 1000, 500, 0, 0, now)
@ -164,7 +154,7 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
requests_this_minute: 1,
tokens_this_minute: 1500,
tokens_this_day: 1500,
input_tokens_this_month: 10000,
input_tokens_this_month: 1000,
cache_creation_input_tokens_this_month: 500,
cache_read_input_tokens_this_month: 0,
output_tokens_this_month: 0,
@ -185,7 +175,7 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
requests_this_minute: 2,
tokens_this_minute: 2800,
tokens_this_day: 2800,
input_tokens_this_month: 11000,
input_tokens_this_month: 2000,
cache_creation_input_tokens_this_month: 500,
cache_read_input_tokens_this_month: 300,
output_tokens_this_month: 0,

8
crates/collab/src/llm/token.rs
View file

@ -22,6 +22,12 @@ pub struct LlmTokenClaims {
pub is_staff: bool,
#[serde(default)]
pub has_llm_closed_beta_feature_flag: bool,
// This field is temporarily optional so it can be added
// in a backwards-compatible way. We can make it required
// once all of the LLM tokens have cycled (~1 hour after
// this change has been deployed).
#[serde(default)]
pub has_llm_subscription: Option<bool>,
pub plan: rpc::proto::Plan,
}
@ -33,6 +39,7 @@ impl LlmTokenClaims {
github_user_login: String,
is_staff: bool,
has_llm_closed_beta_feature_flag: bool,
has_llm_subscription: bool,
plan: rpc::proto::Plan,
config: &Config,
) -> Result<String> {
@ -50,6 +57,7 @@ impl LlmTokenClaims {
github_user_login: Some(github_user_login),
is_staff,
has_llm_closed_beta_feature_flag,
has_llm_subscription: Some(has_llm_subscription),
plan,
};