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 magnotia_core::constants::inference_thread_count;
use magnotia_core::error::{MagnotiaError, Result};
use magnotia_core::types::{Segment, TranscriptionOptions};

use crate::transcriber::{Transcriber, TranscriberCapabilities};

#[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) -> std::result::Result<Self, WhisperBackendError> {
        let ctx = WhisperContext::new_with_params(model_path, WhisperContextParameters::default())
            .map_err(|e| WhisperBackendError::Load(e.to_string()))?;
        Ok(Self { ctx })
    }
}

impl Transcriber for WhisperRsBackend {
    fn capabilities(&self) -> TranscriberCapabilities {
        TranscriberCapabilities {
            sample_rate: magnotia_core::constants::WHISPER_SAMPLE_RATE,
            channels: 1,
            supports_initial_prompt: true,
        }
    }

    /// Synchronously transcribe 16 kHz mono f32 PCM.
    ///
    /// `options.initial_prompt` is piped directly to whisper-rs — this
    /// is the only backend path that honours it; `SpeechModelAdapter`
    /// discards it (Parakeet has no equivalent).
    fn transcribe_sync(
        &mut self,
        samples: &[f32],
        options: &TranscriptionOptions,
    ) -> Result<Vec<Segment>> {
        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| {
            MagnotiaError::TranscriptionFailed(WhisperBackendError::State(e.to_string()).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(inference_thread_count() as i32);
        params.set_print_special(false);
        params.set_print_progress(false);
        params.set_print_realtime(false);

        state.full(params, samples).map_err(|e| {
            MagnotiaError::TranscriptionFailed(
                WhisperBackendError::Transcribe(e.to_string()).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| {
                    MagnotiaError::TranscriptionFailed(
                        WhisperBackendError::Transcribe(e.to_string()).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"));
    }
}