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(()) => return Ok(result),
114                    Err(error) => {
115                        return Err(error);
116                    }
117                },
118                Err(error) => {
119                    tx.rollback().await?;
120                    return Err(error);
121                }
122            }
123        };
124
125        self.run(body).await
126    }
127
128    async fn with_transaction<F, Fut, T>(&self, f: &F) -> Result<(DatabaseTransaction, Result<T>)>
129    where
130        F: Send + Fn(TransactionHandle) -> Fut,
131        Fut: Send + Future<Output = Result<T>>,
132    {
133        let tx = self
134            .pool
135            .begin_with_config(Some(IsolationLevel::ReadCommitted), None)
136            .await?;
137
138        let mut tx = Arc::new(Some(tx));
139        let result = f(TransactionHandle(tx.clone())).await;
140        let Some(tx) = Arc::get_mut(&mut tx).and_then(|tx| tx.take()) else {
141            return Err(anyhow!(
142                "couldn't complete transaction because it's still in use"
143            ))?;
144        };
145
146        Ok((tx, result))
147    }
148
149    async fn run<F, T>(&self, future: F) -> Result<T>
150    where
151        F: Future<Output = Result<T>>,
152    {
153        #[cfg(test)]
154        {
155            if let Executor::Deterministic(executor) = &self.executor {
156                executor.simulate_random_delay().await;
157            }
158
159            self.runtime.as_ref().unwrap().block_on(future)
160        }
161
162        #[cfg(not(test))]
163        {
164            future.await
165        }
166    }
167}