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

Track lifetime spending for each user and model (#16137)

Browse source 
Release Notes:

- N/A

Co-authored-by: Marshall <marshall@zed.dev>
This commit is contained in:
Max Brunsfeld 2024-08-12 17:15:26 -07:00 • committed by GitHub
parent 68ae347077
commit dbcd06642c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 121 additions and 11 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -1,3 +1,4 @@
use crate::db::UserId;
use chrono::Duration;
use rpc::LanguageModelProvider;
use sea_orm::QuerySelect;
@ -14,6 +15,7 @@ pub struct Usage {
pub input_tokens_this_month: usize,
pub output_tokens_this_month: usize,
pub spending_this_month: usize,
pub lifetime_spending: usize,
}
#[derive(Clone, Copy, Debug, Default)]
@ -63,7 +65,7 @@ impl LlmDatabase {
pub async fn get_usage(
&self,
user_id: i32,
user_id: UserId,
provider: LanguageModelProvider,
model_name: &str,
now: DateTimeUtc,
@ -83,6 +85,18 @@ impl LlmDatabase {
.all(&*tx)
.await?;
let (lifetime_input_tokens, lifetime_output_tokens) = 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)
});
let requests_this_minute =
self.get_usage_for_measure(&usages, now, UsageMeasure::RequestsPerMinute)?;
let tokens_this_minute =
@ -95,6 +109,8 @@ impl LlmDatabase {
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);
Ok(Usage {
requests_this_minute,
@ -103,6 +119,7 @@ impl LlmDatabase {
input_tokens_this_month,
output_tokens_this_month,
spending_this_month,
lifetime_spending,
})
})
.await
@ -111,7 +128,7 @@ impl LlmDatabase {
#[allow(clippy::too_many_arguments)]
pub async fn record_usage(
&self,
user_id: i32,
user_id: UserId,
is_staff: bool,
provider: LanguageModelProvider,
model_name: &str,
@ -194,6 +211,50 @@ impl LlmDatabase {
let spending_this_month =
calculate_spending(model, input_tokens_this_month, output_tokens_this_month);
// Update lifetime usage
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?;
let lifetime_usage = match lifetime_usage {
Some(usage) => {
lifetime_usage::Entity::update(lifetime_usage::ActiveModel {
id: ActiveValue::unchanged(usage.id),
input_tokens: ActiveValue::set(
usage.input_tokens + input_token_count as i64,
),
output_tokens: ActiveValue::set(
usage.output_tokens + output_token_count as i64,
),
..Default::default()
})
.exec(&*tx)
.await?
}
None => {
lifetime_usage::ActiveModel {
user_id: ActiveValue::set(user_id),
model_id: ActiveValue::set(model.id),
input_tokens: ActiveValue::set(input_token_count as i64),
output_tokens: ActiveValue::set(output_token_count as i64),
..Default::default()
}
.insert(&*tx)
.await?
}
};
let lifetime_spending = calculate_spending(
model,
lifetime_usage.input_tokens as usize,
lifetime_usage.output_tokens as usize,
);
Ok(Usage {
requests_this_minute,
tokens_this_minute,
@ -201,6 +262,7 @@ impl LlmDatabase {
input_tokens_this_month,
output_tokens_this_month,
spending_this_month,
lifetime_spending,
})
})
.await
@ -246,7 +308,7 @@ impl LlmDatabase {
#[allow(clippy::too_many_arguments)]
async fn update_usage_for_measure(
&self,
user_id: i32,
user_id: UserId,
is_staff: bool,
model_id: ModelId,
usages: &[usage::Model],