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
10 changed files with 165 additions and 1 deletions
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…
Add table
Add a link
Reference in a new issue