1mod ids;
2mod queries;
3mod seed;
4mod tables;
5
6#[cfg(test)]
7mod tests;
8
9use collections::HashMap;
10pub use ids::*;
11use rpc::LanguageModelProvider;
12pub use seed::*;
13pub use tables::*;
14
15#[cfg(test)]
16pub use tests::TestLlmDb;
17use usage_measure::UsageMeasure;
18
19use std::future::Future;
20use std::sync::Arc;
21
22use anyhow::anyhow;
23pub use queries::usages::ActiveUserCount;
24use sea_orm::prelude::*;
25pub use sea_orm::ConnectOptions;
26use sea_orm::{
27 ActiveValue, DatabaseConnection, DatabaseTransaction, IsolationLevel, TransactionTrait,
28};
29
30use crate::db::TransactionHandle;
31use crate::executor::Executor;
32use crate::Result;
33
34/// The database for the LLM service.
35pub struct LlmDatabase {
36 options: ConnectOptions,
37 pool: DatabaseConnection,
38 #[allow(unused)]
39 executor: Executor,
40 provider_ids: HashMap<LanguageModelProvider, ProviderId>,
41 models: HashMap<(LanguageModelProvider, String), model::Model>,
42 usage_measure_ids: HashMap<UsageMeasure, UsageMeasureId>,
43 #[cfg(test)]
44 runtime: Option<tokio::runtime::Runtime>,
45}
46
47impl LlmDatabase {
48 /// Connects to the database with the given options
49 pub async fn new(options: ConnectOptions, executor: Executor) -> Result<Self> {
50 sqlx::any::install_default_drivers();
51 Ok(Self {
52 options: options.clone(),
53 pool: sea_orm::Database::connect(options).await?,
54 executor,
55 provider_ids: HashMap::default(),
56 models: HashMap::default(),
57 usage_measure_ids: HashMap::default(),
58 #[cfg(test)]
59 runtime: None,
60 })
61 }
62
63 pub async fn initialize(&mut self) -> Result<()> {
64 self.initialize_providers().await?;
65 self.initialize_models().await?;
66 self.initialize_usage_measures().await?;
67 Ok(())
68 }
69
70 /// Returns the list of all known models, with their [`LanguageModelProvider`].
71 pub fn all_models(&self) -> Vec<(LanguageModelProvider, model::Model)> {
72 self.models
73 .iter()
74 .map(|((model_provider, _model_name), model)| (*model_provider, model.clone()))
75 .collect::<Vec<_>>()
76 }
77
78 /// Returns the names of the known models for the given [`LanguageModelProvider`].
79 pub fn model_names_for_provider(&self, provider: LanguageModelProvider) -> Vec<String> {
80 self.models
81 .keys()
82 .filter_map(|(model_provider, model_name)| {
83 if model_provider == &provider {
84 Some(model_name)
85 } else {
86 None
87 }
88 })
89 .cloned()
90 .collect::<Vec<_>>()
91 }
92
93 pub fn model(&self, provider: LanguageModelProvider, name: &str) -> Result<&model::Model> {
94 Ok(self
95 .models
96 .get(&(provider, name.to_string()))
97 .ok_or_else(|| anyhow!("unknown model {provider:?}:{name}"))?)
98 }
99
100 pub fn options(&self) -> &ConnectOptions {
101 &self.options
102 }
103
104 pub async fn transaction<F, Fut, T>(&self, f: F) -> Result<T>
105 where
106 F: Send + Fn(TransactionHandle) -> Fut,
107 Fut: Send + Future<Output = Result<T>>,
108 {
109 let body = async {
110 let (tx, result) = self.with_transaction(&f).await?;
111 match result {
112 Ok(result) => match tx.commit().await.map_err(Into::into) {
113 Ok(()) => Ok(result),
114 Err(error) => Err(error),
115 },
116 Err(error) => {
117 tx.rollback().await?;
118 Err(error)
119 }
120 }
121 };
122
123 self.run(body).await
124 }
125
126 async fn with_transaction<F, Fut, T>(&self, f: &F) -> Result<(DatabaseTransaction, Result<T>)>
127 where
128 F: Send + Fn(TransactionHandle) -> Fut,
129 Fut: Send + Future<Output = Result<T>>,
130 {
131 let tx = self
132 .pool
133 .begin_with_config(Some(IsolationLevel::ReadCommitted), None)
134 .await?;
135
136 let mut tx = Arc::new(Some(tx));
137 let result = f(TransactionHandle(tx.clone())).await;
138 let Some(tx) = Arc::get_mut(&mut tx).and_then(|tx| tx.take()) else {
139 return Err(anyhow!(
140 "couldn't complete transaction because it's still in use"
141 ))?;
142 };
143
144 Ok((tx, result))
145 }
146
147 async fn run<F, T>(&self, future: F) -> Result<T>
148 where
149 F: Future<Output = Result<T>>,
150 {
151 #[cfg(test)]
152 {
153 if let Executor::Deterministic(executor) = &self.executor {
154 executor.simulate_random_delay().await;
155 }
156
157 self.runtime.as_ref().unwrap().block_on(future)
158 }
159
160 #[cfg(not(test))]
161 {
162 future.await
163 }
164 }
165}