Detailed changes
@@ -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);
@@ -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);
@@ -3,5 +3,6 @@ use serde::{Deserialize, Serialize};
use crate::id_type;
-id_type!(ProviderId);
id_type!(ModelId);
+id_type!(ProviderId);
+id_type!(UsageId);
@@ -1,3 +1,4 @@
use super::*;
pub mod providers;
+pub mod usages;
@@ -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
+ }
+}
@@ -1,2 +1,3 @@
pub mod model;
pub mod provider;
+pub mod usage;
@@ -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 {}
@@ -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 {}
@@ -1,4 +1,5 @@
mod provider_tests;
+mod usage_tests;
use gpui::BackgroundExecutor;
use parking_lot::Mutex;
@@ -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);
+}