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