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:
Marshall Bowers 2024-08-06 18:40:10 -04:00 committed by GitHub
parent 4f69336024
commit a54e16b7ea
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 165 additions and 1 deletions

View 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);
}