Lumotia/crates/transcription/src/whisper_rs_backend.rs

//! Direct whisper-rs backend. Owns a WhisperContext; each call builds a
//! fresh WhisperState (state can be reused, but fresh-per-call is simpler
//! and matches the transcribe-rs call style we are replacing).
//!
//! Exists because transcribe-rs does not expose set_initial_prompt; this
//! wrapper is the only path that can pipe per-capture vocabulary context
//! into Whisper.

use std::path::Path;

use whisper_rs::{FullParams, SamplingStrategy, WhisperContext, WhisperContextParameters};

use kon_core::types::{Segment, TranscriptionOptions};

#[derive(Debug, thiserror::Error)]
pub enum WhisperBackendError {
    #[error("whisper-rs load failed: {0}")]
    Load(String),
    #[error("whisper-rs state creation failed: {0}")]
    State(String),
    #[error("whisper-rs transcribe failed: {0}")]
    Transcribe(String),
}

pub struct WhisperRsBackend {
    ctx: WhisperContext,
}

impl WhisperRsBackend {
    pub fn load(model_path: &Path) -> Result<Self, WhisperBackendError> {
        let ctx = WhisperContext::new_with_params(model_path, WhisperContextParameters::default())
            .map_err(|e| WhisperBackendError::Load(e.to_string()))?;
        Ok(Self { ctx })
    }

    /// Synchronously transcribe 16 kHz mono f32 PCM.
    ///
    /// `options.initial_prompt` is piped directly to whisper-rs.
    pub fn transcribe_sync(
        &self,
        samples: &[f32],
        options: &TranscriptionOptions,
    ) -> Result<Vec<Segment>, WhisperBackendError> {
        tracing::info!(
            language = ?options.language,
            has_initial_prompt = options.initial_prompt.as_deref().map(|p| !p.is_empty()).unwrap_or(false),
            "WhisperRsBackend::transcribe_sync entering"
        );

        let mut state = self
            .ctx
            .create_state()
            .map_err(|e| WhisperBackendError::State(e.to_string()))?;

        let mut params = FullParams::new(SamplingStrategy::Greedy { best_of: 1 });
        if let Some(lang) = options.language.as_deref() {
            if !lang.is_empty() {
                params.set_language(Some(lang));
            }
        }
        if let Some(prompt) = options.initial_prompt.as_deref() {
            if !prompt.is_empty() {
                params.set_initial_prompt(prompt);
            }
        }
        params.set_n_threads(num_cpus::get() as i32);
        params.set_print_special(false);
        params.set_print_progress(false);
        params.set_print_realtime(false);

        state
            .full(params, samples)
            .map_err(|e| WhisperBackendError::Transcribe(e.to_string()))?;

        let n = state.full_n_segments();

        let mut out = Vec::with_capacity(n.max(0) as usize);
        for i in 0..n {
            let Some(seg) = state.get_segment(i) else {
                continue;
            };
            let text = seg
                .to_str()
                .map_err(|e| WhisperBackendError::Transcribe(e.to_string()))?
                .to_string();
            // whisper-rs timestamps are centiseconds (10ms units). Convert to seconds (f64).
            let start = seg.start_timestamp() as f64 * 0.01;
            let end = seg.end_timestamp() as f64 * 0.01;
            out.push(Segment { start, end, text });
        }
        Ok(out)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn backend_error_displays() {
        let e = WhisperBackendError::Load("oops".into());
        assert!(e.to_string().contains("oops"));
    }
}