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

Add telemetry for LLM usage (#16049)

Browse source 
Release Notes:

- N/A

Co-authored-by: Marshall <marshall@zed.dev>
This commit is contained in:
Max Brunsfeld 2024-08-09 15:15:57 -07:00 • committed by GitHub
parent 423c7b999a
commit 8688b2ad19
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 137 additions and 60 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -107,7 +107,7 @@ impl LlmDatabase {
model_name: &str,
token_count: usize,
now: DateTimeUtc,
) -> Result<()> {
) -> Result<Usage> {
self.transaction(|tx| async move {
let model = self.model(provider, model_name)?;
@ -120,48 +120,57 @@ impl LlmDatabase {
.all(&*tx)
.await?;
self.update_usage_for_measure(
user_id,
model.id,
&usages,
UsageMeasure::RequestsPerMinute,
now,
1,
&tx,
)
.await?;
self.update_usage_for_measure(
user_id,
model.id,
&usages,
UsageMeasure::TokensPerMinute,
now,
token_count,
&tx,
)
.await?;
self.update_usage_for_measure(
user_id,
model.id,
&usages,
UsageMeasure::TokensPerDay,
now,
token_count,
&tx,
)
.await?;
self.update_usage_for_measure(
user_id,
model.id,
&usages,
UsageMeasure::TokensPerMonth,
now,
token_count,
&tx,
)
.await?;
let requests_this_minute = self
.update_usage_for_measure(
user_id,
model.id,
&usages,
UsageMeasure::RequestsPerMinute,
now,
1,
&tx,
)
.await?;
let tokens_this_minute = self
.update_usage_for_measure(
user_id,
model.id,
&usages,
UsageMeasure::TokensPerMinute,
now,
token_count,
&tx,
)
.await?;
let tokens_this_day = self
.update_usage_for_measure(
user_id,
model.id,
&usages,
UsageMeasure::TokensPerDay,
now,
token_count,
&tx,
)
.await?;
let tokens_this_month = self
.update_usage_for_measure(
user_id,
model.id,
&usages,
UsageMeasure::TokensPerMonth,
now,
token_count,
&tx,
)
.await?;
Ok(())
Ok(Usage {
requests_this_minute,
tokens_this_minute,
tokens_this_day,
tokens_this_month,
})
})
.await
}
@ -205,7 +214,7 @@ impl LlmDatabase {
now: DateTimeUtc,
usage_to_add: usize,
tx: &DatabaseTransaction,
) -> Result<()> {
) -> Result<usize> {
let now = now.naive_utc();
let measure_id = *self
.usage_measure_ids
@ -230,6 +239,7 @@ impl LlmDatabase {
}
*buckets.last_mut().unwrap() += usage_to_add as i64;
let total_usage = buckets.iter().sum::<i64>() as usize;
let mut model = usage::ActiveModel {
user_id: ActiveValue::set(user_id),
@ -249,7 +259,7 @@ impl LlmDatabase {
.await?;
}
Ok(())
Ok(total_usage)
}
fn get_usage_for_measure(

25
crates/collab/src/llm/telemetry.rs Normal file
View file

@ -0,0 +1,25 @@
use anyhow::Result;
use serde::Serialize;
#[derive(Serialize, Debug, clickhouse::Row)]
pub struct LlmUsageEventRow {
pub time: i64,
pub user_id: i32,
pub is_staff: bool,
pub plan: String,
pub model: String,
pub provider: String,
pub 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 tokens_this_month: u64,
}
pub async fn report_llm_usage(client: &clickhouse::Client, row: LlmUsageEventRow) -> Result<()> {
let mut insert = client.insert("llm_usage_events")?;
insert.write(&row).await?;
insert.end().await?;
Ok(())
}