//! Application state holding shared ML models and database pool. use std::sync::Arc; use sqlx::PgPool; use tokio::sync::{broadcast, Mutex}; use uuid::Uuid; use crate::listen::{DiarizationConfig, ParakeetEOU, ParakeetTDT, Sortformer}; /// Notification payload for file updates (broadcast to WebSocket subscribers). #[derive(Debug, Clone)] pub struct FileUpdateNotification { /// ID of the updated file pub file_id: Uuid, /// New version number after update pub version: i32, /// List of fields that were updated pub updated_fields: Vec, /// Source of the update: "user", "llm", or "system" pub updated_by: String, } /// 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, /// End-of-Utterance detection model for streaming pub parakeet_eou: Mutex, /// Speaker diarization model (Sortformer) pub sortformer: Mutex, /// Optional database connection pool pub db_pool: Option, /// Broadcast channel for file update notifications pub file_updates: broadcast::Sender, } 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> { 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(), )?; // Create broadcast channel with buffer for 256 messages let (file_updates, _) = broadcast::channel(256); Ok(Self { parakeet: Mutex::new(parakeet), parakeet_eou: Mutex::new(parakeet_eou), sortformer: Mutex::new(sortformer), db_pool: None, file_updates, }) } /// Set the database pool. pub fn with_db_pool(mut self, pool: PgPool) -> Self { self.db_pool = Some(pool); self } /// Broadcast a file update notification to all subscribers. /// /// This is a no-op if there are no subscribers (ignores send errors). pub fn broadcast_file_update(&self, notification: FileUpdateNotification) { // Ignore send errors - they just mean no one is listening let _ = self.file_updates.send(notification); } } /// Type alias for the shared application state. pub type SharedState = Arc;