mirror of
https://github.com/zed-industries/zed.git
synced 2024-09-20 02:47:34 +03:00
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
This commit is contained in:
parent
4f69336024
commit
a54e16b7ea
@ -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;
|
||||
|
57
crates/collab/src/llm/db/queries/usages.rs
Normal file
57
crates/collab/src/llm/db/queries/usages.rs
Normal file
@ -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 {}
|
||||
|
40
crates/collab/src/llm/db/tables/usage.rs
Normal file
40
crates/collab/src/llm/db/tables/usage.rs
Normal file
@ -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;
|
||||
|
24
crates/collab/src/llm/db/tests/usage_tests.rs
Normal file
24
crates/collab/src/llm/db/tests/usage_tests.rs
Normal file
@ -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);
|
||||
}
|
Loading…
Reference in New Issue
Block a user