hyperhive/swarm-stt/src/engine.rs

//! The engine that turns audio bytes into text.

use std::path::PathBuf;
use std::process::Stdio;

use async_trait::async_trait;
use tokio::io::AsyncWriteExt as _;
use tokio::process::Command;

/// Turns the bytes of an audio file into its transcript.
#[async_trait]
pub trait Engine: Send + Sync {
    /// The text spoken in `audio`, or why it could not be told.
    async fn transcribe(&self, audio: &[u8]) -> Result<String, String>;
}

/// What a sender is told when the service could not run the tools at all.
/// The detail (paths, OS errors, tool output) goes to the log, never to the
/// sender.
const COULD_NOT_START: &str = "the service could not start the transcription";

/// `ffmpeg` decodes any container to 16 kHz mono WAV, then `whisper-cli` reads
/// it with the configured model.
#[derive(Debug, Clone)]
pub struct Whisper {
    /// The `ffmpeg` binary.
    pub ffmpeg: PathBuf,
    /// The `whisper-cli` binary (whisper.cpp).
    pub whisper: PathBuf,
    /// The ggml model file, such as `ggml-small.bin`.
    pub model: PathBuf,
    /// Where scratch files go; the caller creates it.
    pub scratch: PathBuf,
    /// Seconds of audio to decode; the rest is cut.
    pub max_secs: u64,
    /// How long one run (decode and transcribe together) may take. The service
    /// runs one job at a time, so a hung child would block every request.
    pub timeout: std::time::Duration,
}

#[async_trait]
impl Engine for Whisper {
    async fn transcribe(&self, audio: &[u8]) -> Result<String, String> {
        if !crate::sniff::is_audio(audio) {
            return Err(
                "the file does not look like audio (Ogg, WAV, FLAC, MP3, AMR, MP4/M4A, WebM)"
                    .to_owned(),
            );
        }
        let dir = tempfile_in(&self.scratch).await.map_err(|e| {
            tracing::warn!(error = %e, "making a scratch directory failed");
            COULD_NOT_START.to_owned()
        })?;
        // Dropping `run` on a timeout kills the child that was running.
        let result = tokio::time::timeout(self.timeout, self.run(&dir, audio))
            .await
            .unwrap_or_else(|_| {
                tracing::warn!(secs = self.timeout.as_secs(), "a transcription timed out");
                Err("the transcription timed out".to_owned())
            });
        let _ = tokio::fs::remove_dir_all(&dir).await;
        result
    }
}

impl Whisper {
    async fn run(&self, dir: &std::path::Path, audio: &[u8]) -> Result<String, String> {
        let wav = dir.join("audio.wav");
        let mut decode = Command::new(&self.ffmpeg)
            // Pipe-only: the audio is untrusted, and a demuxer that follows
            // references inside it (an HLS playlist, say) must not open files
            // or sockets.
            .args([
                "-nostdin",
                "-v",
                "error",
                "-protocol_whitelist",
                "pipe",
                "-i",
                "pipe:0",
                "-t",
            ])
            .arg(self.max_secs.to_string())
            .args(["-ar", "16000", "-ac", "1", "-y"])
            .arg(&wav)
            .stdin(Stdio::piped())
            .stdout(Stdio::null())
            .stderr(Stdio::piped())
            .kill_on_drop(true)
            .spawn()
            .map_err(|e| {
                tracing::warn!(error = %e, "starting ffmpeg failed");
                COULD_NOT_START.to_owned()
            })?;
        if let Some(mut stdin) = decode.stdin.take() {
            // ffmpeg may stop reading early (a cut at `-t`): a broken pipe is fine.
            let _ = stdin.write_all(audio).await;
        }
        let decoded = decode.wait_with_output().await.map_err(|e| {
            tracing::warn!(error = %e, "waiting for ffmpeg failed");
            COULD_NOT_START.to_owned()
        })?;
        if !decoded.status.success() {
            tracing::warn!(
                stderr = %String::from_utf8_lossy(&decoded.stderr).trim(),
                "ffmpeg could not decode the audio"
            );
            return Err("the audio could not be decoded".to_owned());
        }
        let heard = Command::new(&self.whisper)
            .arg("-m")
            .arg(&self.model)
            .arg("-f")
            .arg(&wav)
            .args(["-l", "auto", "-nt", "-np"])
            .stdin(Stdio::null())
            .kill_on_drop(true)
            .output()
            .await
            .map_err(|e| {
                tracing::warn!(error = %e, "starting whisper-cli failed");
                COULD_NOT_START.to_owned()
            })?;
        if !heard.status.success() {
            tracing::warn!(
                stderr = %String::from_utf8_lossy(&heard.stderr).trim(),
                "whisper-cli failed"
            );
            return Err("the transcription failed".to_owned());
        }
        Ok(String::from_utf8_lossy(&heard.stdout).trim().to_owned())
    }
}

/// A fresh empty directory under `parent`, named for this process and a counter.
async fn tempfile_in(parent: &std::path::Path) -> std::io::Result<PathBuf> {
    use std::sync::atomic::{AtomicU64, Ordering};
    static NEXT: AtomicU64 = AtomicU64::new(0);
    let dir = parent.join(format!(
        "run-{}-{}",
        std::process::id(),
        NEXT.fetch_add(1, Ordering::Relaxed)
    ));
    tokio::fs::create_dir_all(&dir).await?;
    Ok(dir)
}

#[cfg(test)]
mod tests {
    use super::{Engine as _, Whisper};
    use std::os::unix::fs::PermissionsExt as _;
    use std::path::Path;

    fn script(path: &Path, body: &str) {
        std::fs::write(path, format!("#!/bin/sh\n{body}\n")).unwrap();
        std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o755)).unwrap();
    }

    fn whisper(dir: &Path, ffmpeg: &str, whisper: &str) -> Whisper {
        script(&dir.join("ffmpeg"), ffmpeg);
        script(&dir.join("whisper-cli"), whisper);
        Whisper {
            ffmpeg: dir.join("ffmpeg"),
            whisper: dir.join("whisper-cli"),
            model: dir.join("model.bin"),
            scratch: dir.join("scratch"),
            max_secs: 300,
            timeout: std::time::Duration::from_secs(5),
        }
    }

    // The fake ffmpeg writes the last argument (the wav path), the fake
    // whisper-cli prints what it finds in the file named by `-f`.
    const FFMPEG_OK: &str = r#"for last; do :; done; cat > "$last""#;
    const WHISPER_ECHO: &str =
        r#"while [ $# -gt 0 ]; do [ "$1" = -f ] && f=$2; shift; done; cat "$f""#;

    #[tokio::test]
    async fn audio_goes_through_ffmpeg_then_whisper_and_the_text_is_trimmed() {
        let dir = tempfile::tempdir().unwrap();
        let engine = whisper(dir.path(), FFMPEG_OK, WHISPER_ECHO);
        let text = engine.transcribe(b"OggS  hello there \n").await.unwrap();
        assert_eq!(text, "OggS  hello there");
        assert_eq!(
            std::fs::read_dir(dir.path().join("scratch"))
                .unwrap()
                .count(),
            0,
            "the scratch directory is removed"
        );
    }

    #[tokio::test]
    async fn a_decode_that_never_ends_times_out_and_leaves_no_scratch_directory() {
        let dir = tempfile::tempdir().unwrap();
        let mut engine = whisper(dir.path(), "sleep 60", WHISPER_ECHO);
        engine.timeout = std::time::Duration::from_millis(200);
        let err = engine.transcribe(b"OggSx").await.unwrap_err();
        assert!(err.contains("timed out"), "{err}");
        assert_eq!(
            std::fs::read_dir(dir.path().join("scratch"))
                .unwrap()
                .count(),
            0
        );
    }

    #[tokio::test]
    async fn ffmpeg_is_started_pipe_only() {
        let dir = tempfile::tempdir().unwrap();
        let record = dir.path().join("args");
        let ffmpeg = format!(
            r#"echo "$@" > {}; for last; do :; done; cat > "$last""#,
            record.display()
        );
        let engine = whisper(dir.path(), &ffmpeg, WHISPER_ECHO);
        engine.transcribe(b"OggSx").await.unwrap();
        let args = std::fs::read_to_string(record).unwrap();
        assert!(
            args.contains("-protocol_whitelist pipe -i pipe:0"),
            "{args}"
        );
    }

    #[tokio::test]
    async fn a_file_that_is_not_audio_never_reaches_ffmpeg() {
        let dir = tempfile::tempdir().unwrap();
        let record = dir.path().join("ran");
        let ffmpeg = format!("touch {}", record.display());
        let engine = whisper(dir.path(), &ffmpeg, WHISPER_ECHO);
        let err = engine
            .transcribe(b"#EXTM3U\nfile:///etc/passwd")
            .await
            .unwrap_err();
        assert!(err.contains("does not look like audio"), "{err}");
        assert!(!record.exists(), "ffmpeg was started");
    }

    #[tokio::test]
    async fn an_ffmpeg_failure_tells_the_sender_briefly_and_keeps_the_detail_out() {
        let dir = tempfile::tempdir().unwrap();
        let engine = whisper(dir.path(), "echo 'bad data' >&2; exit 1", WHISPER_ECHO);
        let err = engine.transcribe(b"OggSx").await.unwrap_err();
        assert_eq!(err, "the audio could not be decoded");
    }

    #[tokio::test]
    async fn a_whisper_failure_tells_the_sender_briefly_and_keeps_the_detail_out() {
        let dir = tempfile::tempdir().unwrap();
        let engine = whisper(dir.path(), FFMPEG_OK, "echo 'no model' >&2; exit 2");
        let err = engine.transcribe(b"OggSx").await.unwrap_err();
        assert_eq!(err, "the transcription failed");
    }
}