diff --git a/crates/transcription/Cargo.toml b/crates/transcription/Cargo.toml index 7530992..d42e695 100644 --- a/crates/transcription/Cargo.toml +++ b/crates/transcription/Cargo.toml @@ -6,3 +6,13 @@ description = "Speech-to-text engine wrappers, model management, and inference c [dependencies] kon-core = { path = "../core" } + +# Unified STT engine (Parakeet via ONNX, Whisper via whisper.cpp) +transcribe-rs = { version = "0.3", features = ["onnx", "whisper-cpp"] } + +# Async runtime for spawn_blocking +tokio = { version = "1", features = ["rt", "sync"] } + +# Model downloads +reqwest = { version = "0.12", features = ["stream"] } +futures-util = "0.3" diff --git a/crates/transcription/src/concurrency.rs b/crates/transcription/src/concurrency.rs new file mode 100644 index 0000000..d5e0199 --- /dev/null +++ b/crates/transcription/src/concurrency.rs @@ -0,0 +1,20 @@ +use std::sync::Arc; + +use kon_core::error::{KonError, Result}; +use kon_core::types::{AudioSamples, Transcript, TranscriptionOptions}; + +use crate::local_engine::LocalEngine; + +/// Runs inference on a blocking thread. Encapsulates all threading concerns. +/// Callers never see spawn_blocking — they call this async function. +pub async fn run_inference( + engine: Arc, + audio: AudioSamples, + options: TranscriptionOptions, +) -> Result { + tokio::task::spawn_blocking(move || { + engine.transcribe_sync(&audio, &options) + }) + .await + .map_err(|e| KonError::TranscriptionFailed(format!("Task join error: {e}")))? +} diff --git a/crates/transcription/src/lib.rs b/crates/transcription/src/lib.rs index 2d43764..6c82485 100644 --- a/crates/transcription/src/lib.rs +++ b/crates/transcription/src/lib.rs @@ -1,2 +1,7 @@ -// kon-transcription: transcribe-rs wrapper behind SpeechToText trait, -// unified model manager, and inference threading. +pub mod concurrency; +pub mod local_engine; +pub mod model_manager; + +pub use concurrency::run_inference; +pub use local_engine::{load_parakeet, load_whisper, LocalEngine}; +pub use model_manager::{download, is_downloaded, list_downloaded, model_dir, models_dir}; diff --git a/crates/transcription/src/local_engine.rs b/crates/transcription/src/local_engine.rs new file mode 100644 index 0000000..25355dc --- /dev/null +++ b/crates/transcription/src/local_engine.rs @@ -0,0 +1,118 @@ +use std::path::Path; +use std::sync::Mutex; + +use transcribe_rs::{SpeechModel, TranscribeOptions, TranscriptionResult}; + +use kon_core::error::{KonError, Result}; +use kon_core::types::{ + AudioSamples, EngineName, Segment, Transcript, TranscriptionOptions, +}; + +/// Wraps any transcribe-rs engine in Kon's SpeechToText trait. +/// Encapsulates threading: inference always runs on a blocking thread. +/// The rest of the app never imports transcribe-rs directly. +pub struct LocalEngine { + engine: Mutex>>, + engine_name: EngineName, +} + +impl LocalEngine { + pub fn new(engine_name: EngineName) -> Self { + Self { + engine: Mutex::new(None), + engine_name, + } + } + + pub fn load(&self, model: Box) { + let mut guard = self.engine.lock().unwrap_or_else(|e| e.into_inner()); + *guard = Some(model); + } + + pub fn name(&self) -> &EngineName { + &self.engine_name + } + + pub fn is_loaded(&self) -> bool { + let guard = self.engine.lock().unwrap_or_else(|e| e.into_inner()); + guard.is_some() + } + + /// Run transcription synchronously. Called from within spawn_blocking. + pub fn transcribe_sync( + &self, + audio: &AudioSamples, + options: &TranscriptionOptions, + ) -> Result { + let mut guard = + self.engine.lock().unwrap_or_else(|e| e.into_inner()); + let engine = guard + .as_mut() + .ok_or(KonError::EngineNotLoaded)?; + + let opts = TranscribeOptions { + language: options.language.clone(), + translate: false, + }; + + let result: TranscriptionResult = engine + .transcribe(audio.samples(), &opts) + .map_err(|e| KonError::TranscriptionFailed(e.to_string()))?; + + let segments = result + .segments + .unwrap_or_default() + .into_iter() + .map(|s| Segment { + start: s.start as f64, + end: s.end as f64, + text: s.text, + }) + .collect(); + + Ok(Transcript::new( + segments, + options + .language + .clone() + .unwrap_or_else(|| "en".to_string()), + audio.duration_secs(), + )) + } +} + +/// Load a Parakeet model from a directory path. +/// The directory should contain model_int8.onnx (+ .onnx_data) and tokenizer.json. +pub fn load_parakeet(model_dir: &Path) -> Result> { + use transcribe_rs::onnx::Quantization; + let model = transcribe_rs::onnx::parakeet::ParakeetModel::load( + model_dir, + &Quantization::Int8, + ) + .map_err(|e| { + KonError::TranscriptionFailed(format!("Failed to load Parakeet: {e}")) + })?; + Ok(Box::new(model)) +} + +/// Load a Whisper model from a GGML file path. +pub fn load_whisper(model_path: &Path) -> Result> { + let engine = transcribe_rs::whisper_cpp::WhisperEngine::load(model_path) + .map_err(|e| { + KonError::TranscriptionFailed(format!( + "Failed to load Whisper: {e}" + )) + })?; + Ok(Box::new(engine)) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn engine_reports_not_available_before_loading() { + let engine = LocalEngine::new(EngineName::new("test")); + assert!(!engine.is_loaded()); + } +} diff --git a/crates/transcription/src/model_manager.rs b/crates/transcription/src/model_manager.rs new file mode 100644 index 0000000..707e914 --- /dev/null +++ b/crates/transcription/src/model_manager.rs @@ -0,0 +1,165 @@ +use std::path::{Path, PathBuf}; + +use kon_core::error::{KonError, Result}; +use kon_core::model_registry::{find_model, ModelFile}; +use kon_core::types::{DownloadProgress, ModelId}; + +/// Resolve the models storage directory. +/// Windows: %LOCALAPPDATA%/kon/models +/// Unix: ~/.kon/models +pub fn models_dir() -> PathBuf { + if cfg!(target_os = "windows") { + let local_app_data = std::env::var("LOCALAPPDATA") + .unwrap_or_else(|_| ".".to_string()); + PathBuf::from(local_app_data).join("kon").join("models") + } else { + dirs_path().join("models") + } +} + +fn dirs_path() -> PathBuf { + if cfg!(target_os = "windows") { + let local_app_data = std::env::var("LOCALAPPDATA") + .unwrap_or_else(|_| ".".to_string()); + PathBuf::from(local_app_data).join("kon") + } else { + let home = + std::env::var("HOME").unwrap_or_else(|_| "/tmp".to_string()); + PathBuf::from(home).join(".kon") + } +} + +/// Get the directory path where a specific model's files are stored. +pub fn model_dir(id: &ModelId) -> PathBuf { + models_dir().join(id.as_str()) +} + +/// Check whether all files for a model have been downloaded. +pub fn is_downloaded(id: &ModelId) -> bool { + let entry = match find_model(id) { + Some(e) => e, + None => return false, + }; + let dir = model_dir(id); + entry.files.iter().all(|f| dir.join(f.filename).exists()) +} + +/// List all downloaded model IDs. +pub fn list_downloaded() -> Vec { + kon_core::model_registry::all_models() + .iter() + .filter(|m| is_downloaded(&m.id)) + .map(|m| m.id.clone()) + .collect() +} + +/// Download all files for a model, calling the progress callback per chunk. +/// Files are downloaded to a .part suffix and atomically renamed on completion. +pub async fn download( + id: &ModelId, + progress: impl Fn(DownloadProgress) + Send + 'static, +) -> Result<()> { + let entry = find_model(id) + .ok_or_else(|| KonError::ModelNotFound(id.clone()))?; + + let dir = model_dir(id); + std::fs::create_dir_all(&dir)?; + + for file in &entry.files { + let dest = dir.join(file.filename); + if dest.exists() { + continue; + } + download_file(file, &dest, id, &progress).await?; + } + + Ok(()) +} + +async fn download_file( + file: &ModelFile, + dest: &Path, + model_id: &ModelId, + progress: &(impl Fn(DownloadProgress) + Send), +) -> Result<()> { + use futures_util::StreamExt; + + let part_path = dest.with_extension( + dest.extension() + .map(|e| format!("{}.part", e.to_string_lossy())) + .unwrap_or_else(|| "part".to_string()), + ); + + let client = reqwest::Client::builder() + .timeout(std::time::Duration::from_secs(30)) + .build() + .map_err(|e| KonError::DownloadFailed(e.to_string()))?; + + let response = client + .get(file.url) + .send() + .await + .map_err(|e| KonError::DownloadFailed(e.to_string()))?; + + let total_bytes = response.content_length().unwrap_or(0); + let mut stream = response.bytes_stream(); + let mut downloaded: u64 = 0; + let mut last_percent: u8 = 0; + + let mut out = std::fs::File::create(&part_path)?; + + while let Some(chunk) = stream.next().await { + let chunk = chunk + .map_err(|e| KonError::DownloadFailed(e.to_string()))?; + std::io::Write::write_all(&mut out, &chunk)?; + downloaded += chunk.len() as u64; + + let percent = if total_bytes > 0 { + ((downloaded as f64 / total_bytes as f64) * 100.0) as u8 + } else { + 0 + }; + + if percent != last_percent { + last_percent = percent; + progress(DownloadProgress { + model_id: model_id.clone(), + file_name: file.filename.to_string(), + bytes_downloaded: downloaded, + total_bytes, + percent, + }); + } + } + + drop(out); + std::fs::rename(&part_path, dest)?; + + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn model_dir_returns_correct_path() { + let id = ModelId::new("whisper-tiny-en"); + let path = model_dir(&id); + assert!(path.to_string_lossy().contains("whisper-tiny-en")); + } + + #[test] + fn is_downloaded_returns_false_for_missing_model() { + let id = ModelId::new("nonexistent-model"); + assert!(!is_downloaded(&id)); + } + + #[test] + fn list_downloaded_returns_empty_when_no_models() { + let list = list_downloaded(); + // In test environment, no models are downloaded + // This just verifies the function doesn't panic + assert!(list.len() <= kon_core::model_registry::all_models().len()); + } +}