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

collab: Track active user counts independently for each model (#16624)

Browse source 
This PR fixes an issue where the active user count spanned individual
models.

We now track the active user counts on a per-model basis.

Release Notes:

- N/A
This commit is contained in:
Marshall Bowers 2024-08-21 17:19:47 -04:00 • committed by GitHub
parent f85ca387a7
commit 0229d3ccac
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 40 additions and 22 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -343,15 +343,27 @@ impl LlmDatabase {
.await
}
pub async fn get_active_user_count(&self, now: DateTimeUtc) -> Result<ActiveUserCount> {
/// Returns the active user count for the specified model.
pub async fn get_active_user_count(
&self,
provider: LanguageModelProvider,
model_name: &str,
now: DateTimeUtc,
) -> Result<ActiveUserCount> {
self.transaction(|tx| async move {
let minute_since = now - Duration::minutes(5);
let day_since = now - Duration::days(5);
let model = self
.models
.get(&(provider, model_name.to_string()))
.ok_or_else(|| anyhow!("unknown model {provider}:{model_name}"))?;
let users_in_recent_minutes = usage::Entity::find()
.filter(
usage::Column::Timestamp
.gte(minute_since.naive_utc())
usage::Column::ModelId
.eq(model.id)
.and(usage::Column::Timestamp.gte(minute_since.naive_utc()))
.and(usage::Column::IsStaff.eq(false)),
)
.select_only()
@ -362,8 +374,9 @@ impl LlmDatabase {
let users_in_recent_days = usage::Entity::find()
.filter(
usage::Column::Timestamp
.gte(day_since.naive_utc())
usage::Column::ModelId
.eq(model.id)
.and(usage::Column::Timestamp.gte(day_since.naive_utc()))
.and(usage::Column::IsStaff.eq(false)),
)
.select_only()