From a54e16b7ea8f4b949882f8e9d9edb5abd0edae66 Mon Sep 17 00:00:00 2001 From: Marshall Bowers Date: Tue, 6 Aug 2024 18:40:10 -0400 Subject: [PATCH] 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 --- .../20240806182921_test_schema.sql | 16 ++++++ .../20240806213401_create_usages.sql | 15 +++++ crates/collab/src/llm/db/ids.rs | 3 +- crates/collab/src/llm/db/queries.rs | 1 + crates/collab/src/llm/db/queries/usages.rs | 57 +++++++++++++++++++ crates/collab/src/llm/db/tables.rs | 1 + crates/collab/src/llm/db/tables/model.rs | 8 +++ crates/collab/src/llm/db/tables/usage.rs | 40 +++++++++++++ crates/collab/src/llm/db/tests.rs | 1 + crates/collab/src/llm/db/tests/usage_tests.rs | 24 ++++++++ 10 files changed, 165 insertions(+), 1 deletion(-) create mode 100644 crates/collab/migrations_llm/20240806213401_create_usages.sql create mode 100644 crates/collab/src/llm/db/queries/usages.rs create mode 100644 crates/collab/src/llm/db/tables/usage.rs create mode 100644 crates/collab/src/llm/db/tests/usage_tests.rs diff --git a/crates/collab/migrations_llm.sqlite/20240806182921_test_schema.sql b/crates/collab/migrations_llm.sqlite/20240806182921_test_schema.sql index 7b6feb0302..ab854d3e43 100644 --- a/crates/collab/migrations_llm.sqlite/20240806182921_test_schema.sql +++ b/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); diff --git a/crates/collab/migrations_llm/20240806213401_create_usages.sql b/crates/collab/migrations_llm/20240806213401_create_usages.sql new file mode 100644 index 0000000000..913f0a1add --- /dev/null +++ b/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); diff --git a/crates/collab/src/llm/db/ids.rs b/crates/collab/src/llm/db/ids.rs index d4613e9c7f..2b256651f8 100644 --- a/crates/collab/src/llm/db/ids.rs +++ b/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); diff --git a/crates/collab/src/llm/db/queries.rs b/crates/collab/src/llm/db/queries.rs index 3e02c17a6a..ded7a54cfb 100644 --- a/crates/collab/src/llm/db/queries.rs +++ b/crates/collab/src/llm/db/queries.rs @@ -1,3 +1,4 @@ use super::*; pub mod providers; +pub mod usages; diff --git a/crates/collab/src/llm/db/queries/usages.rs b/crates/collab/src/llm/db/queries/usages.rs new file mode 100644 index 0000000000..4b672fa6ac --- /dev/null +++ b/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 { + 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 + } +} diff --git a/crates/collab/src/llm/db/tables.rs b/crates/collab/src/llm/db/tables.rs index 89b42283de..87307eacfa 100644 --- a/crates/collab/src/llm/db/tables.rs +++ b/crates/collab/src/llm/db/tables.rs @@ -1,2 +1,3 @@ pub mod model; pub mod provider; +pub mod usage; diff --git a/crates/collab/src/llm/db/tables/model.rs b/crates/collab/src/llm/db/tables/model.rs index 7242365acf..c8ff1ce47e 100644 --- a/crates/collab/src/llm/db/tables/model.rs +++ b/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 for Entity { @@ -28,4 +30,10 @@ impl Related for Entity { } } +impl Related for Entity { + fn to() -> RelationDef { + Relation::Usages.def() + } +} + impl ActiveModelBehavior for ActiveModel {} diff --git a/crates/collab/src/llm/db/tables/usage.rs b/crates/collab/src/llm/db/tables/usage.rs new file mode 100644 index 0000000000..afb4f7e03a --- /dev/null +++ b/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 for Entity { + fn to() -> RelationDef { + Relation::Model.def() + } +} + +impl ActiveModelBehavior for ActiveModel {} diff --git a/crates/collab/src/llm/db/tests.rs b/crates/collab/src/llm/db/tests.rs index d5f2c8e19c..1e76d85522 100644 --- a/crates/collab/src/llm/db/tests.rs +++ b/crates/collab/src/llm/db/tests.rs @@ -1,4 +1,5 @@ mod provider_tests; +mod usage_tests; use gpui::BackgroundExecutor; use parking_lot::Mutex; diff --git a/crates/collab/src/llm/db/tests/usage_tests.rs b/crates/collab/src/llm/db/tests/usage_tests.rs new file mode 100644 index 0000000000..ee2bcdbe01 --- /dev/null +++ b/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) { + 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); +}