path: root/makima/src/server/state.rs

//! Application state holding shared ML models and database pool.

use std::sync::Arc;
use sqlx::PgPool;
use tokio::sync::Mutex;

use crate::listen::{DiarizationConfig, ParakeetEOU, ParakeetTDT, Sortformer};

/// Shared application state containing ML models and database pool.
///
/// Models are wrapped in `Mutex` for thread-safe mutable access during inference.
pub struct AppState {
    /// Speech-to-text model (Parakeet TDT)
    pub parakeet: Mutex<ParakeetTDT>,
    /// End-of-Utterance detection model for streaming
    pub parakeet_eou: Mutex<ParakeetEOU>,
    /// Speaker diarization model (Sortformer)
    pub sortformer: Mutex<Sortformer>,
    /// Optional database connection pool
    pub db_pool: Option<PgPool>,
}

impl AppState {
    /// Load all ML models from the specified directories.
    ///
    /// # Arguments
    /// * `parakeet_model_dir` - Path to the Parakeet TDT model directory
    /// * `parakeet_eou_dir` - Path to the Parakeet EOU model directory
    /// * `sortformer_model_path` - Path to the Sortformer diarization model file
    pub fn new(
        parakeet_model_dir: &str,
        parakeet_eou_dir: &str,
        sortformer_model_path: &str,
    ) -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
        let parakeet = ParakeetTDT::from_pretrained(parakeet_model_dir, None)?;
        let parakeet_eou = ParakeetEOU::from_pretrained(parakeet_eou_dir, None)?;
        let sortformer = Sortformer::with_config(
            sortformer_model_path,
            None,
            DiarizationConfig::callhome(),
        )?;

        Ok(Self {
            parakeet: Mutex::new(parakeet),
            parakeet_eou: Mutex::new(parakeet_eou),
            sortformer: Mutex::new(sortformer),
            db_pool: None,
        })
    }

    /// Set the database pool.
    pub fn with_db_pool(mut self, pool: PgPool) -> Self {
        self.db_pool = Some(pool);
        self
    }
}

/// Type alias for the shared application state.
pub type SharedState = Arc<AppState>;