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

collab: Track cache writes/reads in LLM usage (#18834)

Browse source 
This PR extends the LLM usage tracking to support tracking usage for
cache writes and reads for Anthropic models.

Release Notes:

- N/A

---------

Co-authored-by: Antonio Scandurra <me@as-cii.com>
Co-authored-by: Antonio <antonio@zed.dev>
This commit is contained in:
Marshall Bowers 2024-10-07 17:32:49 -04:00 • committed by GitHub
parent c5d252b837
commit d55f025906
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 241 additions and 39 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -14,6 +14,8 @@ pub struct Usage {
pub tokens_this_minute: usize,
pub tokens_this_day: usize,
pub input_tokens_this_month: usize,
pub cache_creation_input_tokens_this_month: usize,
pub cache_read_input_tokens_this_month: usize,
pub output_tokens_this_month: usize,
pub spending_this_month: usize,
pub lifetime_spending: usize,
@ -160,17 +162,14 @@ impl LlmDatabase {
.all(&*tx)
.await?;
let (lifetime_input_tokens, lifetime_output_tokens) = lifetime_usage::Entity::find()
let lifetime_usage = lifetime_usage::Entity::find()
.filter(
lifetime_usage::Column::UserId
.eq(user_id)
.and(lifetime_usage::Column::ModelId.eq(model.id)),
)
.one(&*tx)
.await?
.map_or((0, 0), |usage| {
(usage.input_tokens as usize, usage.output_tokens as usize)
});
.await?;
let requests_this_minute =
self.get_usage_for_measure(&usages, now, UsageMeasure::RequestsPerMinute)?;
@ -180,18 +179,44 @@ impl LlmDatabase {
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, output_tokens_this_month);
let lifetime_spending =
calculate_spending(model, lifetime_input_tokens, lifetime_output_tokens);
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 {
calculate_spending(
model,
lifetime_usage.input_tokens as usize,
lifetime_usage.cache_creation_input_tokens as usize,
lifetime_usage.cache_read_input_tokens as usize,
lifetime_usage.output_tokens as usize,
)
} else {
0
};
Ok(Usage {
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,
spending_this_month,
lifetime_spending,
@ -208,6 +233,8 @@ impl LlmDatabase {
provider: LanguageModelProvider,
model_name: &str,
input_token_count: usize,
cache_creation_input_tokens: usize,
cache_read_input_tokens: usize,
output_token_count: usize,
now: DateTimeUtc,
) -> Result<Usage> {
@ -235,6 +262,10 @@ impl LlmDatabase {
&tx,
)
.await?;
let total_token_count = input_token_count
+ cache_read_input_tokens
+ cache_creation_input_tokens
+ output_token_count;
let tokens_this_minute = self
.update_usage_for_measure(
user_id,
@ -243,7 +274,7 @@ impl LlmDatabase {
&usages,
UsageMeasure::TokensPerMinute,
now,
input_token_count + output_token_count,
total_token_count,
&tx,
)
.await?;
@ -255,7 +286,7 @@ impl LlmDatabase {
&usages,
UsageMeasure::TokensPerDay,
now,
input_token_count + output_token_count,
total_token_count,
&tx,
)
.await?;
@ -271,6 +302,30 @@ impl LlmDatabase {
&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,
@ -283,8 +338,13 @@ impl LlmDatabase {
&tx,
)
.await?;
let spending_this_month =
calculate_spending(model, input_tokens_this_month, output_tokens_this_month);
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,
);
// Update lifetime usage
let lifetime_usage = lifetime_usage::Entity::find()
@ -303,6 +363,12 @@ impl LlmDatabase {
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,
),
@ -327,6 +393,8 @@ impl LlmDatabase {
let lifetime_spending = calculate_spending(
model,
lifetime_usage.input_tokens as usize,
lifetime_usage.cache_creation_input_tokens as usize,
lifetime_usage.cache_read_input_tokens as usize,
lifetime_usage.output_tokens as usize,
);
@ -335,6 +403,8 @@ impl LlmDatabase {
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,
spending_this_month,
lifetime_spending,
@ -501,13 +571,24 @@ impl LlmDatabase {
fn calculate_spending(
model: &model::Model,
input_tokens_this_month: usize,
cache_creation_input_tokens_this_month: usize,
cache_read_input_tokens_this_month: usize,
output_tokens_this_month: usize,
) -> usize {
let input_token_cost =
input_tokens_this_month * model.price_per_million_input_tokens as usize / 1_000_000;
let cache_creation_input_token_cost = cache_creation_input_tokens_this_month
* model.price_per_million_cache_creation_input_tokens as usize
/ 1_000_000;
let cache_read_input_token_cost = cache_read_input_tokens_this_month
* model.price_per_million_cache_read_input_tokens as usize
/ 1_000_000;
let output_token_cost =
output_tokens_this_month * model.price_per_million_output_tokens as usize / 1_000_000;
input_token_cost + output_token_cost
input_token_cost
+ cache_creation_input_token_cost
+ cache_read_input_token_cost
+ output_token_cost
}
const MINUTE_BUCKET_COUNT: usize = 12;
@ -521,6 +602,8 @@ impl UsageMeasure {
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,
}
}
@ -531,6 +614,8 @@ impl UsageMeasure {
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),
}
}

2
crates/collab/src/llm/db/tables/lifetime_usage.rs
View file

@ -9,6 +9,8 @@ pub struct Model {
pub user_id: UserId,
pub model_id: ModelId,
pub input_tokens: i64,
pub cache_creation_input_tokens: i64,
pub cache_read_input_tokens: i64,
pub output_tokens: i64,
}

2
crates/collab/src/llm/db/tables/model.rs
View file

@ -14,6 +14,8 @@ pub struct Model {
pub max_tokens_per_minute: i64,
pub max_tokens_per_day: i64,
pub price_per_million_input_tokens: i32,
pub price_per_million_cache_creation_input_tokens: i32,
pub price_per_million_cache_read_input_tokens: i32,
pub price_per_million_output_tokens: i32,
}

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

@ -10,6 +10,8 @@ pub enum UsageMeasure {
TokensPerMinute,
TokensPerDay,
InputTokensPerMonth,
CacheCreationInputTokensPerMonth,
CacheReadInputTokensPerMonth,
OutputTokensPerMonth,
}

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

@ -33,12 +33,12 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
let user_id = UserId::from_proto(123);
let now = t0;
db.record_usage(user_id, false, provider, model, 1000, 0, now)
db.record_usage(user_id, false, provider, model, 1000, 0, 0, 0, now)
.await
.unwrap();
let now = t0 + Duration::seconds(10);
db.record_usage(user_id, false, provider, model, 2000, 0, now)
db.record_usage(user_id, false, provider, model, 2000, 0, 0, 0, now)
.await
.unwrap();
@ -50,6 +50,8 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
tokens_this_minute: 3000,
tokens_this_day: 3000,
input_tokens_this_month: 3000,
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,
@ -65,6 +67,8 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
tokens_this_minute: 2000,
tokens_this_day: 3000,
input_tokens_this_month: 3000,
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,
@ -72,7 +76,7 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
);
let now = t0 + Duration::seconds(60);
db.record_usage(user_id, false, provider, model, 3000, 0, now)
db.record_usage(user_id, false, provider, model, 3000, 0, 0, 0, now)
.await
.unwrap();
@ -84,6 +88,8 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
tokens_this_minute: 5000,
tokens_this_day: 6000,
input_tokens_this_month: 6000,
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,
@ -100,13 +106,15 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
tokens_this_minute: 0,
tokens_this_day: 5000,
input_tokens_this_month: 6000,
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,
}
);
db.record_usage(user_id, false, provider, model, 4000, 0, now)
db.record_usage(user_id, false, provider, model, 4000, 0, 0, 0, now)
.await
.unwrap();
@ -118,6 +126,8 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
tokens_this_minute: 4000,
tokens_this_day: 9000,
input_tokens_this_month: 10000,
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,
@ -134,6 +144,50 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
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,
}
);
// Test cache creation input tokens
db.record_usage(user_id, false, provider, model, 1000, 500, 0, 0, now)
.await
.unwrap();
let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
assert_eq!(
usage,
Usage {
requests_this_minute: 1,
tokens_this_minute: 1500,
tokens_this_day: 1500,
input_tokens_this_month: 10000,
cache_creation_input_tokens_this_month: 500,
cache_read_input_tokens_this_month: 0,
output_tokens_this_month: 0,
spending_this_month: 0,
lifetime_spending: 0,
}
);
// Test cache read input tokens
db.record_usage(user_id, false, provider, model, 1000, 0, 300, 0, now)
.await
.unwrap();
let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
assert_eq!(
usage,
Usage {
requests_this_minute: 2,
tokens_this_minute: 2800,
tokens_this_day: 2800,
input_tokens_this_month: 11000,
cache_creation_input_tokens_this_month: 500,
cache_read_input_tokens_this_month: 300,
output_tokens_this_month: 0,
spending_this_month: 0,
lifetime_spending: 0,

4
crates/collab/src/llm/telemetry.rs
View file

@ -12,11 +12,15 @@ pub struct LlmUsageEventRow {
pub model: String,
pub provider: String,
pub input_token_count: u64,
pub cache_creation_input_token_count: u64,
pub cache_read_input_token_count: u64,
pub output_token_count: u64,
pub requests_this_minute: u64,
pub tokens_this_minute: u64,
pub tokens_this_day: u64,
pub input_tokens_this_month: u64,
pub cache_creation_input_tokens_this_month: u64,
pub cache_read_input_tokens_this_month: u64,
pub output_tokens_this_month: u64,
pub spending_this_month: u64,
pub lifetime_spending: u64,