collab: Add usages table to LLM database (#15884)

This PR adds a `usages` table to the LLM database.

We'll use this to track usage for rate-limiting purposes.

Release Notes:

- N/A

crates/collab/migrations_llm.sqlite/20240806182921_test_schema.sql

@@ -14,3 +14,19 @@ create table models (
 create unique index uix_models_on_provider_id_name on models (provider_id, name);
 create index ix_models_on_provider_id on models (provider_id);
 create index ix_models_on_name on models (name);
+
+create table if not exists usages (
+    id integer primary key autoincrement,
+    user_id integer not null,
+    model_id integer not null references models (id) on delete cascade,
+    requests_this_minute integer not null default 0,
+    tokens_this_minute integer not null default 0,
+    requests_this_day integer not null default 0,
+    tokens_this_day integer not null default 0,
+    requests_this_month integer not null default 0,
+    tokens_this_month integer not null default 0
+);
+
+create index ix_usages_on_user_id on usages (user_id);
+create index ix_usages_on_model_id on usages (model_id);
+create unique index uix_usages_on_user_id_model_id on usages (user_id, model_id);

crates/collab/migrations_llm/20240806213401_create_usages.sql

@@ -0,0 +1,15 @@
+create table if not exists usages (
+    id serial primary key,
+    user_id integer not null,
+    model_id integer not null references models (id) on delete cascade,
+    requests_this_minute integer not null default 0,
+    tokens_this_minute bigint not null default 0,
+    requests_this_day integer not null default 0,
+    tokens_this_day bigint not null default 0,
+    requests_this_month integer not null default 0,
+    tokens_this_month bigint not null default 0
+);
+
+create index ix_usages_on_user_id on usages (user_id);
+create index ix_usages_on_model_id on usages (model_id);
+create unique index uix_usages_on_user_id_model_id on usages (user_id, model_id);

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

@@ -3,5 +3,6 @@ use serde::{Deserialize, Serialize};
 
 use crate::id_type;
 
-id_type!(ProviderId);
 id_type!(ModelId);
+id_type!(ProviderId);
+id_type!(UsageId);

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

@@ -1,3 +1,4 @@
 use super::*;
 
 pub mod providers;
+pub mod usages;

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

@@ -0,0 +1,57 @@
+use rpc::LanguageModelProvider;
+
+use super::*;
+
+impl LlmDatabase {
+    pub async fn find_or_create_usage(
+        &self,
+        user_id: i32,
+        provider: LanguageModelProvider,
+        model_name: &str,
+    ) -> Result<usage::Model> {
+        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 = 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()
+                .filter(
+                    usage::Column::UserId
+                        .eq(user_id)
+                        .and(usage::Column::ModelId.eq(model_id)),
+                )
+                .one(&*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()
+            })
+            .exec_with_returning(&*tx)
+            .await?;
+
+            Ok(usage)
+        })
+        .await
+    }
+}

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

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

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

@@ -20,6 +20,8 @@ pub enum Relation {
         to = "super::provider::Column::Id"
     )]
     Provider,
+    #[sea_orm(has_many = "super::usage::Entity")]
+    Usages,
 }
 
 impl Related<super::provider::Entity> for Entity {
@@ -28,4 +30,10 @@ impl Related<super::provider::Entity> for Entity {
     }
 }
 
+impl Related<super::usage::Entity> for Entity {
+    fn to() -> RelationDef {
+        Relation::Usages.def()
+    }
+}
+
 impl ActiveModelBehavior for ActiveModel {}

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

@@ -0,0 +1,40 @@
+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,
+    /// 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,
+}
+
+#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
+pub enum Relation {
+    #[sea_orm(
+        belongs_to = "super::model::Entity",
+        from = "Column::ModelId",
+        to = "super::model::Column::Id"
+    )]
+    Model,
+}
+
+impl Related<super::model::Entity> for Entity {
+    fn to() -> RelationDef {
+        Relation::Model.def()
+    }
+}
+
+impl ActiveModelBehavior for ActiveModel {}

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

@@ -1,4 +1,5 @@
 mod provider_tests;
+mod usage_tests;
 
 use gpui::BackgroundExecutor;
 use parking_lot::Mutex;

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

@@ -0,0 +1,24 @@
+use std::sync::Arc;
+
+use pretty_assertions::assert_eq;
+use rpc::LanguageModelProvider;
+
+use crate::llm::db::LlmDatabase;
+use crate::test_both_llm_dbs;
+
+test_both_llm_dbs!(
+    test_find_or_create_usage,
+    test_find_or_create_usage_postgres,
+    test_find_or_create_usage_sqlite
+);
+
+async fn test_find_or_create_usage(db: &Arc<LlmDatabase>) {
+    db.initialize_providers().await.unwrap();
+
+    let usage = db
+        .find_or_create_usage(123, LanguageModelProvider::Anthropic, "claude-3-5-sonnet")
+        .await
+        .unwrap();
+
+    assert_eq!(usage.user_id, 123);
+}