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    pub fn model(&self, provider: LanguageModelProvider, name: &str) -> Result<&model::Model> {
 71        Ok(self
 72            .models
 73            .get(&(provider, name.to_string()))
 74            .ok_or_else(|| anyhow!("unknown model {provider:?}:{name}"))?)
 75    }
 76
 77    pub fn options(&self) -> &ConnectOptions {
 78        &self.options
 79    }
 80
 81    pub async fn transaction<F, Fut, T>(&self, f: F) -> Result<T>
 82    where
 83        F: Send + Fn(TransactionHandle) -> Fut,
 84        Fut: Send + Future<Output = Result<T>>,
 85    {
 86        let body = async {
 87            let (tx, result) = self.with_transaction(&f).await?;
 88            match result {
 89                Ok(result) => match tx.commit().await.map_err(Into::into) {
 90                    Ok(()) => return Ok(result),
 91                    Err(error) => {
 92                        return Err(error);
 93                    }
 94                },
 95                Err(error) => {
 96                    tx.rollback().await?;
 97                    return Err(error);
 98                }
 99            }
100        };
101
102        self.run(body).await
103    }
104
105    async fn with_transaction<F, Fut, T>(&self, f: &F) -> Result<(DatabaseTransaction, Result<T>)>
106    where
107        F: Send + Fn(TransactionHandle) -> Fut,
108        Fut: Send + Future<Output = Result<T>>,
109    {
110        let tx = self
111            .pool
112            .begin_with_config(Some(IsolationLevel::ReadCommitted), None)
113            .await?;
114
115        let mut tx = Arc::new(Some(tx));
116        let result = f(TransactionHandle(tx.clone())).await;
117        let Some(tx) = Arc::get_mut(&mut tx).and_then(|tx| tx.take()) else {
118            return Err(anyhow!(
119                "couldn't complete transaction because it's still in use"
120            ))?;
121        };
122
123        Ok((tx, result))
124    }
125
126    async fn run<F, T>(&self, future: F) -> Result<T>
127    where
128        F: Future<Output = Result<T>>,
129    {
130        #[cfg(test)]
131        {
132            if let Executor::Deterministic(executor) = &self.executor {
133                executor.simulate_random_delay().await;
134            }
135
136            self.runtime.as_ref().unwrap().block_on(future)
137        }
138
139        #[cfg(not(test))]
140        {
141            future.await
142        }
143    }
144}