Apply rate limits in LLM service (#15997)

Release Notes:

- N/A

---------

Co-authored-by: Marshall <marshall@zed.dev>
Co-authored-by: Marshall Bowers <elliott.codes@gmail.com>

crates/collab/src/llm/db.rs

@@ -1,20 +1,26 @@
 mod ids;
 mod queries;
+mod seed;
 mod tables;
 
 #[cfg(test)]
 mod tests;
 
+use collections::HashMap;
 pub use ids::*;
+use rpc::LanguageModelProvider;
+pub use seed::*;
 pub use tables::*;
 
 #[cfg(test)]
 pub use tests::TestLlmDb;
+use usage_measure::UsageMeasure;
 
 use std::future::Future;
 use std::sync::Arc;
 
 use anyhow::anyhow;
+pub use queries::usages::ActiveUserCount;
 use sea_orm::prelude::*;
 pub use sea_orm::ConnectOptions;
 use sea_orm::{
@@ -31,6 +37,9 @@ pub struct LlmDatabase {
     pool: DatabaseConnection,
     #[allow(unused)]
     executor: Executor,
+    provider_ids: HashMap<LanguageModelProvider, ProviderId>,
+    models: HashMap<(LanguageModelProvider, String), model::Model>,
+    usage_measure_ids: HashMap<UsageMeasure, UsageMeasureId>,
     #[cfg(test)]
     runtime: Option<tokio::runtime::Runtime>,
 }
@@ -43,11 +52,28 @@ impl LlmDatabase {
             options: options.clone(),
             pool: sea_orm::Database::connect(options).await?,
             executor,
+            provider_ids: HashMap::default(),
+            models: HashMap::default(),
+            usage_measure_ids: HashMap::default(),
             #[cfg(test)]
             runtime: None,
         })
     }
 
+    pub async fn initialize(&mut self) -> Result<()> {
+        self.initialize_providers().await?;
+        self.initialize_models().await?;
+        self.initialize_usage_measures().await?;
+        Ok(())
+    }
+
+    pub fn model(&self, provider: LanguageModelProvider, name: &str) -> Result<&model::Model> {
+        Ok(self
+            .models
+            .get(&(provider, name.to_string()))
+            .ok_or_else(|| anyhow!("unknown model {provider:?}:{name}"))?)
+    }
+
     pub fn options(&self) -> &ConnectOptions {
         &self.options
     }

crates/collab/src/llm/db/ids.rs

@@ -6,3 +6,4 @@ use crate::id_type;
 id_type!(ModelId);
 id_type!(ProviderId);
 id_type!(UsageId);
+id_type!(UsageMeasureId);

crates/collab/src/llm/db/queries/providers.rs

@@ -1,66 +1,115 @@
-use sea_orm::sea_query::OnConflict;
-use sea_orm::QueryOrder;
-
 use super::*;
+use sea_orm::QueryOrder;
+use std::str::FromStr;
+use strum::IntoEnumIterator as _;
+
+pub struct ModelRateLimits {
+    pub max_requests_per_minute: i32,
+    pub max_tokens_per_minute: i32,
+    pub max_tokens_per_day: i32,
+}
 
 impl LlmDatabase {
-    pub async fn initialize_providers(&self) -> Result<()> {
+    pub async fn initialize_providers(&mut self) -> Result<()> {
+        self.provider_ids = self
+            .transaction(|tx| async move {
+                let existing_providers = provider::Entity::find().all(&*tx).await?;
+
+                let mut new_providers = LanguageModelProvider::iter()
+                    .filter(|provider| {
+                        !existing_providers
+                            .iter()
+                            .any(|p| p.name == provider.to_string())
+                    })
+                    .map(|provider| provider::ActiveModel {
+                        name: ActiveValue::set(provider.to_string()),
+                        ..Default::default()
+                    })
+                    .peekable();
+
+                if new_providers.peek().is_some() {
+                    provider::Entity::insert_many(new_providers)
+                        .exec(&*tx)
+                        .await?;
+                }
+
+                let all_providers: HashMap<_, _> = provider::Entity::find()
+                    .all(&*tx)
+                    .await?
+                    .iter()
+                    .filter_map(|provider| {
+                        LanguageModelProvider::from_str(&provider.name)
+                            .ok()
+                            .map(|p| (p, provider.id))
+                    })
+                    .collect();
+
+                Ok(all_providers)
+            })
+            .await?;
+        Ok(())
+    }
+
+    pub async fn initialize_models(&mut self) -> Result<()> {
+        let all_provider_ids = &self.provider_ids;
+        self.models = self
+            .transaction(|tx| async move {
+                let all_models: HashMap<_, _> = model::Entity::find()
+                    .all(&*tx)
+                    .await?
+                    .into_iter()
+                    .filter_map(|model| {
+                        let provider = all_provider_ids.iter().find_map(|(provider, id)| {
+                            if *id == model.provider_id {
+                                Some(provider)
+                            } else {
+                                None
+                            }
+                        })?;
+                        Some(((*provider, model.name.clone()), model))
+                    })
+                    .collect();
+                Ok(all_models)
+            })
+            .await?;
+        Ok(())
+    }
+
+    pub async fn insert_models(
+        &mut self,
+        models: &[(LanguageModelProvider, String, ModelRateLimits)],
+    ) -> Result<()> {
+        let all_provider_ids = &self.provider_ids;
         self.transaction(|tx| async move {
-            let providers_and_models = vec![
-                ("anthropic", "claude-3-5-sonnet"),
-                ("anthropic", "claude-3-opus"),
-                ("anthropic", "claude-3-sonnet"),
-                ("anthropic", "claude-3-haiku"),
-            ];
-
-            for (provider_name, model_name) in providers_and_models {
-                let insert_provider = provider::Entity::insert(provider::ActiveModel {
-                    name: ActiveValue::set(provider_name.to_owned()),
+            model::Entity::insert_many(models.into_iter().map(|(provider, name, rate_limits)| {
+                let provider_id = all_provider_ids[&provider];
+                model::ActiveModel {
+                    provider_id: ActiveValue::set(provider_id),
+                    name: ActiveValue::set(name.clone()),
+                    max_requests_per_minute: ActiveValue::set(rate_limits.max_requests_per_minute),
+                    max_tokens_per_minute: ActiveValue::set(rate_limits.max_tokens_per_minute),
+                    max_tokens_per_day: ActiveValue::set(rate_limits.max_tokens_per_day),
                     ..Default::default()
-                })
-                .on_conflict(
-                    OnConflict::columns([provider::Column::Name])
-                        .update_column(provider::Column::Name)
-                        .to_owned(),
-                );
-
-                let provider = if tx.support_returning() {
-                    insert_provider.exec_with_returning(&*tx).await?
-                } else {
-                    insert_provider.exec_without_returning(&*tx).await?;
-                    provider::Entity::find()
-                        .filter(provider::Column::Name.eq(provider_name))
-                        .one(&*tx)
-                        .await?
-                        .ok_or_else(|| anyhow!("failed to insert provider"))?
-                };
-
-                model::Entity::insert(model::ActiveModel {
-                    provider_id: ActiveValue::set(provider.id),
-                    name: ActiveValue::set(model_name.to_owned()),
-                    ..Default::default()
-                })
-                .on_conflict(
-                    OnConflict::columns([model::Column::ProviderId, model::Column::Name])
-                        .update_column(model::Column::Name)
-                        .to_owned(),
-                )
-                .exec_without_returning(&*tx)
-                .await?;
-            }
-
+                }
+            }))
+            .exec_without_returning(&*tx)
+            .await?;
             Ok(())
         })
-        .await
+        .await?;
+        self.initialize_models().await
     }
 
     /// Returns the list of LLM providers.
-    pub async fn list_providers(&self) -> Result<Vec<provider::Model>> {
+    pub async fn list_providers(&self) -> Result<Vec<LanguageModelProvider>> {
         self.transaction(|tx| async move {
             Ok(provider::Entity::find()
                 .order_by_asc(provider::Column::Name)
                 .all(&*tx)
-                .await?)
+                .await?
+                .into_iter()
+                .filter_map(|p| LanguageModelProvider::from_str(&p.name).ok())
+                .collect())
         })
         .await
     }

crates/collab/src/llm/db/queries/usages.rs

@@ -1,57 +1,318 @@
+use chrono::Duration;
 use rpc::LanguageModelProvider;
+use sea_orm::QuerySelect;
+use std::{iter, str::FromStr};
+use strum::IntoEnumIterator as _;
 
 use super::*;
 
+#[derive(Debug, PartialEq, Clone, Copy)]
+pub struct Usage {
+    pub requests_this_minute: usize,
+    pub tokens_this_minute: usize,
+    pub tokens_this_day: usize,
+    pub tokens_this_month: usize,
+}
+
+#[derive(Clone, Copy, Debug, Default)]
+pub struct ActiveUserCount {
+    pub users_in_recent_minutes: usize,
+    pub users_in_recent_days: usize,
+}
+
 impl LlmDatabase {
-    pub async fn find_or_create_usage(
+    pub async fn initialize_usage_measures(&mut self) -> Result<()> {
+        let all_measures = self
+            .transaction(|tx| async move {
+                let existing_measures = usage_measure::Entity::find().all(&*tx).await?;
+
+                let new_measures = UsageMeasure::iter()
+                    .filter(|measure| {
+                        !existing_measures
+                            .iter()
+                            .any(|m| m.name == measure.to_string())
+                    })
+                    .map(|measure| usage_measure::ActiveModel {
+                        name: ActiveValue::set(measure.to_string()),
+                        ..Default::default()
+                    })
+                    .collect::<Vec<_>>();
+
+                if !new_measures.is_empty() {
+                    usage_measure::Entity::insert_many(new_measures)
+                        .exec(&*tx)
+                        .await?;
+                }
+
+                Ok(usage_measure::Entity::find().all(&*tx).await?)
+            })
+            .await?;
+
+        self.usage_measure_ids = all_measures
+            .into_iter()
+            .filter_map(|measure| {
+                UsageMeasure::from_str(&measure.name)
+                    .ok()
+                    .map(|um| (um, measure.id))
+            })
+            .collect();
+        Ok(())
+    }
+
+    pub async fn get_usage(
         &self,
         user_id: i32,
         provider: LanguageModelProvider,
         model_name: &str,
-    ) -> Result<usage::Model> {
+        now: DateTimeUtc,
+    ) -> Result<Usage> {
         self.transaction(|tx| async move {
-            let provider_name = match provider {
-                LanguageModelProvider::Anthropic => "anthropic",
-                LanguageModelProvider::OpenAi => "open_ai",
-                LanguageModelProvider::Google => "google",
-                LanguageModelProvider::Zed => "zed",
-            };
+            let model = self
+                .models
+                .get(&(provider, model_name.to_string()))
+                .ok_or_else(|| anyhow!("unknown model {provider}:{model_name}"))?;
 
-            let model = model::Entity::find()
-                .inner_join(provider::Entity)
-                .filter(
-                    provider::Column::Name
-                        .eq(provider_name)
-                        .and(model::Column::Name.eq(model_name)),
-                )
-                .one(&*tx)
-                .await?
-                // TODO: Create the model, if one doesn't exist.
-                .ok_or_else(|| anyhow!("no model found for {provider_name}:{model_name}"))?;
-            let model_id = model.id;
-
-            let existing_usage = usage::Entity::find()
+            let usages = usage::Entity::find()
                 .filter(
                     usage::Column::UserId
                         .eq(user_id)
-                        .and(usage::Column::ModelId.eq(model_id)),
+                        .and(usage::Column::ModelId.eq(model.id)),
                 )
-                .one(&*tx)
+                .all(&*tx)
                 .await?;
-            if let Some(usage) = existing_usage {
-                return Ok(usage);
-            }
 
-            let usage = usage::Entity::insert(usage::ActiveModel {
-                user_id: ActiveValue::set(user_id),
-                model_id: ActiveValue::set(model_id),
-                ..Default::default()
+            let requests_this_minute =
+                self.get_usage_for_measure(&usages, now, UsageMeasure::RequestsPerMinute)?;
+            let tokens_this_minute =
+                self.get_usage_for_measure(&usages, now, UsageMeasure::TokensPerMinute)?;
+            let tokens_this_day =
+                self.get_usage_for_measure(&usages, now, UsageMeasure::TokensPerDay)?;
+            let tokens_this_month =
+                self.get_usage_for_measure(&usages, now, UsageMeasure::TokensPerMonth)?;
+
+            Ok(Usage {
+                requests_this_minute,
+                tokens_this_minute,
+                tokens_this_day,
+                tokens_this_month,
             })
-            .exec_with_returning(&*tx)
-            .await?;
-
-            Ok(usage)
         })
         .await
     }
+
+    pub async fn record_usage(
+        &self,
+        user_id: i32,
+        provider: LanguageModelProvider,
+        model_name: &str,
+        token_count: usize,
+        now: DateTimeUtc,
+    ) -> Result<()> {
+        self.transaction(|tx| async move {
+            let model = self.model(provider, model_name)?;
+
+            let usages = usage::Entity::find()
+                .filter(
+                    usage::Column::UserId
+                        .eq(user_id)
+                        .and(usage::Column::ModelId.eq(model.id)),
+                )
+                .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?;
+
+            Ok(())
+        })
+        .await
+    }
+
+    pub async fn get_active_user_count(&self, now: DateTimeUtc) -> Result<ActiveUserCount> {
+        self.transaction(|tx| async move {
+            let minute_since = now - Duration::minutes(5);
+            let day_since = now - Duration::days(5);
+
+            let users_in_recent_minutes = usage::Entity::find()
+                .filter(usage::Column::Timestamp.gte(minute_since.naive_utc()))
+                .group_by(usage::Column::UserId)
+                .count(&*tx)
+                .await? as usize;
+
+            let users_in_recent_days = usage::Entity::find()
+                .filter(usage::Column::Timestamp.gte(day_since.naive_utc()))
+                .group_by(usage::Column::UserId)
+                .count(&*tx)
+                .await? as usize;
+
+            Ok(ActiveUserCount {
+                users_in_recent_minutes,
+                users_in_recent_days,
+            })
+        })
+        .await
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    async fn update_usage_for_measure(
+        &self,
+        user_id: i32,
+        model_id: ModelId,
+        usages: &[usage::Model],
+        usage_measure: UsageMeasure,
+        now: DateTimeUtc,
+        usage_to_add: usize,
+        tx: &DatabaseTransaction,
+    ) -> Result<()> {
+        let now = now.naive_utc();
+        let measure_id = *self
+            .usage_measure_ids
+            .get(&usage_measure)
+            .ok_or_else(|| anyhow!("usage measure {usage_measure} not found"))?;
+
+        let mut id = None;
+        let mut timestamp = now;
+        let mut buckets = vec![0_i64];
+
+        if let Some(old_usage) = usages.iter().find(|usage| usage.measure_id == measure_id) {
+            id = Some(old_usage.id);
+            let (live_buckets, buckets_since) =
+                Self::get_live_buckets(old_usage, now, usage_measure);
+            if !live_buckets.is_empty() {
+                buckets.clear();
+                buckets.extend_from_slice(live_buckets);
+                buckets.extend(iter::repeat(0).take(buckets_since));
+                timestamp =
+                    old_usage.timestamp + (usage_measure.bucket_duration() * buckets_since as i32);
+            }
+        }
+
+        *buckets.last_mut().unwrap() += usage_to_add as i64;
+
+        let mut model = usage::ActiveModel {
+            user_id: ActiveValue::set(user_id),
+            model_id: ActiveValue::set(model_id),
+            measure_id: ActiveValue::set(measure_id),
+            timestamp: ActiveValue::set(timestamp),
+            buckets: ActiveValue::set(buckets),
+            ..Default::default()
+        };
+
+        if let Some(id) = id {
+            model.id = ActiveValue::unchanged(id);
+            model.update(tx).await?;
+        } else {
+            usage::Entity::insert(model)
+                .exec_without_returning(tx)
+                .await?;
+        }
+
+        Ok(())
+    }
+
+    fn get_usage_for_measure(
+        &self,
+        usages: &[usage::Model],
+        now: DateTimeUtc,
+        usage_measure: UsageMeasure,
+    ) -> Result<usize> {
+        let now = now.naive_utc();
+        let measure_id = *self
+            .usage_measure_ids
+            .get(&usage_measure)
+            .ok_or_else(|| anyhow!("usage measure {usage_measure} not found"))?;
+        let Some(usage) = usages.iter().find(|usage| usage.measure_id == measure_id) else {
+            return Ok(0);
+        };
+
+        let (live_buckets, _) = Self::get_live_buckets(usage, now, usage_measure);
+        Ok(live_buckets.iter().sum::<i64>() as _)
+    }
+
+    fn get_live_buckets(
+        usage: &usage::Model,
+        now: chrono::NaiveDateTime,
+        measure: UsageMeasure,
+    ) -> (&[i64], usize) {
+        let seconds_since_usage = (now - usage.timestamp).num_seconds().max(0);
+        let buckets_since_usage =
+            seconds_since_usage as f32 / measure.bucket_duration().num_seconds() as f32;
+        let buckets_since_usage = buckets_since_usage.ceil() as usize;
+        let mut live_buckets = &[] as &[i64];
+        if buckets_since_usage < measure.bucket_count() {
+            let expired_bucket_count =
+                (usage.buckets.len() + buckets_since_usage).saturating_sub(measure.bucket_count());
+            live_buckets = &usage.buckets[expired_bucket_count..];
+            while live_buckets.first() == Some(&0) {
+                live_buckets = &live_buckets[1..];
+            }
+        }
+        (live_buckets, buckets_since_usage)
+    }
+}
+
+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 {
+        match self {
+            UsageMeasure::RequestsPerMinute => MINUTE_BUCKET_COUNT,
+            UsageMeasure::TokensPerMinute => MINUTE_BUCKET_COUNT,
+            UsageMeasure::TokensPerDay => DAY_BUCKET_COUNT,
+            UsageMeasure::TokensPerMonth => MONTH_BUCKET_COUNT,
+        }
+    }
+
+    fn total_duration(&self) -> Duration {
+        match self {
+            UsageMeasure::RequestsPerMinute => Duration::minutes(1),
+            UsageMeasure::TokensPerMinute => Duration::minutes(1),
+            UsageMeasure::TokensPerDay => Duration::hours(24),
+            UsageMeasure::TokensPerMonth => Duration::days(30),
+        }
+    }
+
+    fn bucket_duration(&self) -> Duration {
+        self.total_duration() / self.bucket_count() as i32
+    }
 }

crates/collab/src/llm/db/seed.rs

@@ -0,0 +1,45 @@
+use super::*;
+use crate::{Config, Result};
+use queries::providers::ModelRateLimits;
+
+pub async fn seed_database(_config: &Config, db: &mut LlmDatabase, _force: bool) -> Result<()> {
+    db.insert_models(&[
+        (
+            LanguageModelProvider::Anthropic,
+            "claude-3-5-sonnet".into(),
+            ModelRateLimits {
+                max_requests_per_minute: 5,
+                max_tokens_per_minute: 20_000,
+                max_tokens_per_day: 300_000,
+            },
+        ),
+        (
+            LanguageModelProvider::Anthropic,
+            "claude-3-opus".into(),
+            ModelRateLimits {
+                max_requests_per_minute: 5,
+                max_tokens_per_minute: 10_000,
+                max_tokens_per_day: 300_000,
+            },
+        ),
+        (
+            LanguageModelProvider::Anthropic,
+            "claude-3-sonnet".into(),
+            ModelRateLimits {
+                max_requests_per_minute: 5,
+                max_tokens_per_minute: 20_000,
+                max_tokens_per_day: 300_000,
+            },
+        ),
+        (
+            LanguageModelProvider::Anthropic,
+            "claude-3-haiku".into(),
+            ModelRateLimits {
+                max_requests_per_minute: 5,
+                max_tokens_per_minute: 25_000,
+                max_tokens_per_day: 300_000,
+            },
+        ),
+    ])
+    .await
+}

crates/collab/src/llm/db/tables.rs

@@ -1,3 +1,4 @@
 pub mod model;
 pub mod provider;
 pub mod usage;
+pub mod usage_measure;

crates/collab/src/llm/db/tables/model.rs

@@ -10,6 +10,9 @@ pub struct Model {
     pub id: ModelId,
     pub provider_id: ProviderId,
     pub name: String,
+    pub max_requests_per_minute: i32,
+    pub max_tokens_per_minute: i32,
+    pub max_tokens_per_day: i32,
 }
 
 #[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]

crates/collab/src/llm/db/tables/provider.rs

@@ -1,6 +1,5 @@
-use sea_orm::entity::prelude::*;
-
 use crate::llm::db::ProviderId;
+use sea_orm::entity::prelude::*;
 
 /// An LLM provider.
 #[derive(Clone, Debug, PartialEq, DeriveEntityModel)]

crates/collab/src/llm/db/tables/usage.rs

@@ -1,24 +1,20 @@
+use crate::llm::db::{ModelId, UsageId, UsageMeasureId};
 use sea_orm::entity::prelude::*;
 
-use crate::llm::db::ModelId;
-
 /// An LLM usage record.
 #[derive(Clone, Debug, PartialEq, DeriveEntityModel)]
 #[sea_orm(table_name = "usages")]
 pub struct Model {
     #[sea_orm(primary_key)]
-    pub id: i32,
+    pub id: UsageId,
     /// The ID of the Zed user.
     ///
     /// Corresponds to the `users` table in the primary collab database.
     pub user_id: i32,
     pub model_id: ModelId,
-    pub requests_this_minute: i32,
-    pub tokens_this_minute: i64,
-    pub requests_this_day: i32,
-    pub tokens_this_day: i64,
-    pub requests_this_month: i32,
-    pub tokens_this_month: i64,
+    pub measure_id: UsageMeasureId,
+    pub timestamp: DateTime,
+    pub buckets: Vec<i64>,
 }
 
 #[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
@@ -29,6 +25,12 @@ pub enum Relation {
         to = "super::model::Column::Id"
     )]
     Model,
+    #[sea_orm(
+        belongs_to = "super::usage_measure::Entity",
+        from = "Column::MeasureId",
+        to = "super::usage_measure::Column::Id"
+    )]
+    UsageMeasure,
 }
 
 impl Related<super::model::Entity> for Entity {
@@ -37,4 +39,10 @@ impl Related<super::model::Entity> for Entity {
     }
 }
 
+impl Related<super::usage_measure::Entity> for Entity {
+    fn to() -> RelationDef {
+        Relation::UsageMeasure.def()
+    }
+}
+
 impl ActiveModelBehavior for ActiveModel {}

crates/collab/src/llm/db/tables/usage_measure.rs

@@ -0,0 +1,35 @@
+use crate::llm::db::UsageMeasureId;
+use sea_orm::entity::prelude::*;
+
+#[derive(
+    Copy, Clone, Debug, PartialEq, Eq, Hash, strum::EnumString, strum::Display, strum::EnumIter,
+)]
+#[strum(serialize_all = "snake_case")]
+pub enum UsageMeasure {
+    RequestsPerMinute,
+    TokensPerMinute,
+    TokensPerDay,
+    TokensPerMonth,
+}
+
+#[derive(Clone, Debug, PartialEq, DeriveEntityModel)]
+#[sea_orm(table_name = "usage_measures")]
+pub struct Model {
+    #[sea_orm(primary_key)]
+    pub id: UsageMeasureId,
+    pub name: String,
+}
+
+#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
+pub enum Relation {
+    #[sea_orm(has_many = "super::usage::Entity")]
+    Usages,
+}
+
+impl Related<super::usage::Entity> for Entity {
+    fn to() -> RelationDef {
+        Relation::Usages.def()
+    }
+}
+
+impl ActiveModelBehavior for ActiveModel {}

crates/collab/src/llm/db/tests.rs

@@ -6,7 +6,6 @@ use parking_lot::Mutex;
 use rand::prelude::*;
 use sea_orm::ConnectionTrait;
 use sqlx::migrate::MigrateDatabase;
-use std::sync::Arc;
 use std::time::Duration;
 
 use crate::migrations::run_database_migrations;
@@ -14,47 +13,11 @@ use crate::migrations::run_database_migrations;
 use super::*;
 
 pub struct TestLlmDb {
-    pub db: Option<Arc<LlmDatabase>>,
+    pub db: Option<LlmDatabase>,
     pub connection: Option<sqlx::AnyConnection>,
 }
 
 impl TestLlmDb {
-    pub fn sqlite(background: BackgroundExecutor) -> Self {
-        let url = "sqlite::memory:";
-        let runtime = tokio::runtime::Builder::new_current_thread()
-            .enable_io()
-            .enable_time()
-            .build()
-            .unwrap();
-
-        let mut db = runtime.block_on(async {
-            let mut options = ConnectOptions::new(url);
-            options.max_connections(5);
-            let db = LlmDatabase::new(options, Executor::Deterministic(background))
-                .await
-                .unwrap();
-            let sql = include_str!(concat!(
-                env!("CARGO_MANIFEST_DIR"),
-                "/migrations_llm.sqlite/20240806182921_test_schema.sql"
-            ));
-            db.pool
-                .execute(sea_orm::Statement::from_string(
-                    db.pool.get_database_backend(),
-                    sql,
-                ))
-                .await
-                .unwrap();
-            db
-        });
-
-        db.runtime = Some(runtime);
-
-        Self {
-            db: Some(Arc::new(db)),
-            connection: None,
-        }
-    }
-
     pub fn postgres(background: BackgroundExecutor) -> Self {
         static LOCK: Mutex<()> = Mutex::new(());
 
@@ -91,29 +54,26 @@ impl TestLlmDb {
         db.runtime = Some(runtime);
 
         Self {
-            db: Some(Arc::new(db)),
+            db: Some(db),
             connection: None,
         }
     }
 
-    pub fn db(&self) -> &Arc<LlmDatabase> {
-        self.db.as_ref().unwrap()
+    pub fn db(&mut self) -> &mut LlmDatabase {
+        self.db.as_mut().unwrap()
     }
 }
 
 #[macro_export]
-macro_rules! test_both_llm_dbs {
-    ($test_name:ident, $postgres_test_name:ident, $sqlite_test_name:ident) => {
-        #[cfg(target_os = "macos")]
+macro_rules! test_llm_db {
+    ($test_name:ident, $postgres_test_name:ident) => {
         #[gpui::test]
         async fn $postgres_test_name(cx: &mut gpui::TestAppContext) {
-            let test_db = $crate::llm::db::TestLlmDb::postgres(cx.executor().clone());
-            $test_name(test_db.db()).await;
-        }
+            if !cfg!(target_os = "macos") {
+                return;
+            }
 
-        #[gpui::test]
-        async fn $sqlite_test_name(cx: &mut gpui::TestAppContext) {
-            let test_db = $crate::llm::db::TestLlmDb::sqlite(cx.executor().clone());
+            let mut test_db = $crate::llm::db::TestLlmDb::postgres(cx.executor().clone());
             $test_name(test_db.db()).await;
         }
     };

crates/collab/src/llm/db/tests/provider_tests.rs

@@ -1,17 +1,15 @@
-use std::sync::Arc;
-
 use pretty_assertions::assert_eq;
+use rpc::LanguageModelProvider;
 
 use crate::llm::db::LlmDatabase;
-use crate::test_both_llm_dbs;
+use crate::test_llm_db;
 
-test_both_llm_dbs!(
+test_llm_db!(
     test_initialize_providers,
-    test_initialize_providers_postgres,
-    test_initialize_providers_sqlite
+    test_initialize_providers_postgres
 );
 
-async fn test_initialize_providers(db: &Arc<LlmDatabase>) {
+async fn test_initialize_providers(db: &mut LlmDatabase) {
     let initial_providers = db.list_providers().await.unwrap();
     assert_eq!(initial_providers, vec![]);
 
@@ -22,9 +20,13 @@ async fn test_initialize_providers(db: &Arc<LlmDatabase>) {
 
     let providers = db.list_providers().await.unwrap();
 
-    let provider_names = providers
-        .into_iter()
-        .map(|provider| provider.name)
-        .collect::<Vec<_>>();
-    assert_eq!(provider_names, vec!["anthropic".to_string()]);
+    assert_eq!(
+        providers,
+        &[
+            LanguageModelProvider::Anthropic,
+            LanguageModelProvider::Google,
+            LanguageModelProvider::OpenAi,
+            LanguageModelProvider::Zed
+        ]
+    )
 }

crates/collab/src/llm/db/tests/usage_tests.rs

@@ -1,24 +1,120 @@
-use std::sync::Arc;
-
+use crate::{
+    llm::db::{queries::providers::ModelRateLimits, queries::usages::Usage, LlmDatabase},
+    test_llm_db,
+};
+use chrono::{Duration, Utc};
 use pretty_assertions::assert_eq;
 use rpc::LanguageModelProvider;
 
-use crate::llm::db::LlmDatabase;
-use crate::test_both_llm_dbs;
+test_llm_db!(test_tracking_usage, test_tracking_usage_postgres);
 
-test_both_llm_dbs!(
-    test_find_or_create_usage,
-    test_find_or_create_usage_postgres,
-    test_find_or_create_usage_sqlite
-);
+async fn test_tracking_usage(db: &mut LlmDatabase) {
+    let provider = LanguageModelProvider::Anthropic;
+    let model = "claude-3-5-sonnet";
 
-async fn test_find_or_create_usage(db: &Arc<LlmDatabase>) {
-    db.initialize_providers().await.unwrap();
+    db.initialize().await.unwrap();
+    db.insert_models(&[(
+        provider,
+        model.to_string(),
+        ModelRateLimits {
+            max_requests_per_minute: 5,
+            max_tokens_per_minute: 10_000,
+            max_tokens_per_day: 50_000,
+        },
+    )])
+    .await
+    .unwrap();
 
-    let usage = db
-        .find_or_create_usage(123, LanguageModelProvider::Anthropic, "claude-3-5-sonnet")
+    let t0 = Utc::now();
+    let user_id = 123;
+
+    let now = t0;
+    db.record_usage(user_id, provider, model, 1000, now)
         .await
         .unwrap();
 
-    assert_eq!(usage.user_id, 123);
+    let now = t0 + Duration::seconds(10);
+    db.record_usage(user_id, provider, model, 2000, 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: 3000,
+            tokens_this_day: 3000,
+            tokens_this_month: 3000,
+        }
+    );
+
+    let now = t0 + Duration::seconds(60);
+    let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
+    assert_eq!(
+        usage,
+        Usage {
+            requests_this_minute: 1,
+            tokens_this_minute: 2000,
+            tokens_this_day: 3000,
+            tokens_this_month: 3000,
+        }
+    );
+
+    let now = t0 + Duration::seconds(60);
+    db.record_usage(user_id, provider, model, 3000, 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: 5000,
+            tokens_this_day: 6000,
+            tokens_this_month: 6000,
+        }
+    );
+
+    let t1 = t0 + Duration::hours(24);
+    let now = t1;
+    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: 5000,
+            tokens_this_month: 6000,
+        }
+    );
+
+    db.record_usage(user_id, provider, model, 4000, 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: 4000,
+            tokens_this_day: 9000,
+            tokens_this_month: 10000,
+        }
+    );
+
+    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,
+            tokens_this_month: 9000,
+        }
+    );
 }