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 std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;

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

use lumotia_core::error::{Error, Result};
use lumotia_core::hardware::vulkan_loader_available;
use lumotia_core::tuning::{inference_thread_count, Workload};
use lumotia_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: lumotia_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>> {
        self.transcribe_sync_inner(samples, options, None)
    }

    /// Cancellable variant. Installs the abort flag as whisper-rs's
    /// abort callback so a `drain_inference` timeout in the live session
    /// can break the worker out of the decoding loop instead of letting
    /// it run to completion against a stuck backend.
    fn transcribe_sync_with_abort(
        &mut self,
        samples: &[f32],
        options: &TranscriptionOptions,
        abort_flag: Arc<AtomicBool>,
    ) -> Result<Vec<Segment>> {
        self.transcribe_sync_inner(samples, options, Some(abort_flag))
    }
}

impl WhisperRsBackend {
    fn transcribe_sync_inner(
        &mut self,
        samples: &[f32],
        options: &TranscriptionOptions,
        abort_flag: Option<Arc<AtomicBool>>,
    ) -> Result<Vec<Segment>> {
        tracing::info!(
            language = ?options.language,
            has_initial_prompt = options.initial_prompt.as_deref().map(|p| !p.is_empty()).unwrap_or(false),
            has_abort_flag = abort_flag.is_some(),
            "WhisperRsBackend::transcribe_sync entering"
        );

        let mut state = self.ctx.create_state().map_err(|e| {
            Error::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);
            }
        }
        let gpu_offloaded = cfg!(feature = "whisper-vulkan") && vulkan_loader_available();
        params.set_n_threads(inference_thread_count(Workload::Whisper, gpu_offloaded) as i32);
        params.set_print_special(false);
        params.set_print_progress(false);
        params.set_print_realtime(false);

        // Wire the per-task abort flag into whisper-rs. The closure is
        // polled by whisper.cpp between decode steps; returning true
        // tells the backend to abort the current inference. This is the
        // cancellation route that closes the orphaned-thread loophole
        // when a live session is stopped or wedges.
        if let Some(flag) = abort_flag {
            params.set_abort_callback_safe(move || flag.load(Ordering::Relaxed));
        }

        state.full(params, samples).map_err(|e| {
            Error::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| {
                    Error::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"));
    }
}