swarm-stt: the speech-to-text core (hash cache, shared runs, ffmpeg and whisper-cli engine)

swarm-stt/Cargo.toml

@@ -0,0 +1,22 @@
+[package]
+name = "swarm-stt"
+version.workspace = true
+edition.workspace = true
+
+[dependencies]
+# The `Engine` trait: a fake in the tests, whisper in the daemon.
+async-trait.workspace = true
+# The SHA-256 of the audio, which keys the cache and the in-flight map.
+sha2.workspace = true
+thiserror.workspace = true
+tokio.workspace = true
+tracing.workspace = true
+
+[dev-dependencies]
+# A scratch cache directory per test.
+tempfile = "3"
+# `start_paused` for the test that overlaps two requests.
+tokio = { workspace = true, features = ["test-util"] }
+
+[lints]
+workspace = true

swarm-stt/src/cache.rs

@@ -0,0 +1,149 @@
+//! The transcript cache: one `<sha256>.txt` file per memo in a directory, kept
+//! for a TTL counted from the file's modification time.
+
+use std::fmt::Write as _;
+use std::io::ErrorKind;
+use std::path::{Path, PathBuf};
+use std::time::{Duration, SystemTime};
+
+use sha2::{Digest as _, Sha256};
+
+/// The lowercase hex SHA-256 of `bytes`: the cache key.
+#[must_use]
+pub fn key(bytes: &[u8]) -> String {
+    Sha256::digest(bytes)
+        .iter()
+        .fold(String::with_capacity(64), |mut hex, byte| {
+            let _ = write!(hex, "{byte:02x}");
+            hex
+        })
+}
+
+/// A directory of transcripts that expire `ttl` after they were written.
+#[derive(Debug, Clone)]
+pub struct Cache {
+    dir: PathBuf,
+    ttl: Duration,
+}
+
+impl Cache {
+    /// A cache in `dir` (created on first write).
+    #[must_use]
+    pub fn new(dir: impl Into<PathBuf>, ttl: Duration) -> Self {
+        Self {
+            dir: dir.into(),
+            ttl,
+        }
+    }
+
+    fn path(&self, key: &str) -> PathBuf {
+        self.dir.join(format!("{key}.txt"))
+    }
+
+    /// The transcript stored under `key`, unless it is missing or expired. An
+    /// expired file is removed.
+    pub async fn get(&self, key: &str) -> Option<String> {
+        let path = self.path(key);
+        let age = age_of(&path).await?;
+        if age >= self.ttl {
+            let _ = tokio::fs::remove_file(&path).await;
+            return None;
+        }
+        tokio::fs::read_to_string(&path).await.ok()
+    }
+
+    /// Stores `text` under `key`. The file is written beside its name and
+    /// renamed into place, so a reader never sees half a transcript.
+    pub async fn put(&self, key: &str, text: &str) -> std::io::Result<()> {
+        tokio::fs::create_dir_all(&self.dir).await?;
+        let path = self.path(key);
+        let tmp = self.dir.join(format!("{key}.tmp"));
+        tokio::fs::write(&tmp, text).await?;
+        tokio::fs::rename(&tmp, &path).await
+    }
+
+    /// Removes every expired transcript. Returns how many it removed.
+    pub async fn purge(&self) -> usize {
+        let Ok(mut entries) = tokio::fs::read_dir(&self.dir).await else {
+            return 0;
+        };
+        let mut removed = 0;
+        while let Ok(Some(entry)) = entries.next_entry().await {
+            let path = entry.path();
+            if path.extension().is_some_and(|ext| ext == "txt")
+                && age_of(&path).await.is_some_and(|age| age >= self.ttl)
+                && tokio::fs::remove_file(&path).await.is_ok()
+            {
+                removed += 1;
+            }
+        }
+        removed
+    }
+}
+
+/// How long ago `path` was last written, or `None` when it cannot be told.
+async fn age_of(path: &Path) -> Option<Duration> {
+    let modified = match tokio::fs::metadata(path).await {
+        Ok(meta) => meta.modified().ok()?,
+        Err(e) if e.kind() == ErrorKind::NotFound => return None,
+        Err(e) => {
+            tracing::warn!(path = %path.display(), error = %e, "reading a cached transcript's age failed");
+            return None;
+        }
+    };
+    // A file dated in the future counts as brand new.
+    Some(
+        SystemTime::now()
+            .duration_since(modified)
+            .unwrap_or_default(),
+    )
+}
+
+#[cfg(test)]
+mod tests {
+    use super::{Cache, key};
+    use std::time::Duration;
+
+    #[test]
+    fn the_key_is_the_sha256_of_the_bytes() {
+        assert_eq!(
+            key(b"abc"),
+            "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"
+        );
+    }
+
+    #[tokio::test]
+    async fn a_stored_transcript_is_found_until_it_expires() {
+        let dir = tempfile::tempdir().unwrap();
+        let fresh = Cache::new(dir.path(), Duration::from_mins(1));
+        fresh.put("k", "hello").await.unwrap();
+        assert_eq!(fresh.get("k").await.as_deref(), Some("hello"));
+
+        let expired = Cache::new(dir.path(), Duration::ZERO);
+        assert_eq!(expired.get("k").await, None);
+        assert!(
+            !dir.path().join("k.txt").exists(),
+            "an expired file is removed"
+        );
+    }
+
+    #[tokio::test]
+    async fn purge_removes_only_expired_transcripts() {
+        let dir = tempfile::tempdir().unwrap();
+        let cache = Cache::new(dir.path(), Duration::from_mins(1));
+        cache.put("keep", "a").await.unwrap();
+        assert_eq!(cache.purge().await, 0);
+        assert_eq!(cache.get("keep").await.as_deref(), Some("a"));
+
+        let expiring = Cache::new(dir.path(), Duration::ZERO);
+        assert_eq!(expiring.purge().await, 1);
+        assert_eq!(cache.get("keep").await, None);
+    }
+
+    #[tokio::test]
+    async fn a_missing_directory_is_a_miss_and_an_empty_purge() {
+        let cache = Cache::new("/nonexistent/swarm-stt-test", Duration::from_mins(1));
+        assert_eq!(cache.get("k").await, None);
+        assert_eq!(cache.purge().await, 0);
+    }
+}

swarm-stt/src/engine.rs

@@ -0,0 +1,172 @@
+//! 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>;
+}
+
+/// `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,
+}
+
+#[async_trait]
+impl Engine for Whisper {
+    async fn transcribe(&self, audio: &[u8]) -> Result<String, String> {
+        let dir = tempfile_in(&self.scratch)
+            .await
+            .map_err(|e| format!("making a scratch directory: {e}"))?;
+        let result = self.run(&dir, audio).await;
+        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)
+            .args(["-nostdin", "-v", "error", "-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| format!("starting ffmpeg: {e}"))?;
+        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| format!("waiting for ffmpeg: {e}"))?;
+        if !decoded.status.success() {
+            return Err(format!(
+                "ffmpeg could not decode the audio: {}",
+                String::from_utf8_lossy(&decoded.stderr).trim()
+            ));
+        }
+        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| format!("starting whisper-cli: {e}"))?;
+        if !heard.status.success() {
+            return Err(format!(
+                "whisper-cli failed: {}",
+                String::from_utf8_lossy(&heard.stderr).trim()
+            ));
+        }
+        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,
+        }
+    }
+
+    // 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"  hello there \n").await.unwrap();
+        assert_eq!(text, "hello there");
+        assert_eq!(
+            std::fs::read_dir(dir.path().join("scratch"))
+                .unwrap()
+                .count(),
+            0,
+            "the scratch directory is removed"
+        );
+    }
+
+    #[tokio::test]
+    async fn an_ffmpeg_failure_is_reported_with_its_stderr() {
+        let dir = tempfile::tempdir().unwrap();
+        let engine = whisper(dir.path(), "echo 'bad data' >&2; exit 1", WHISPER_ECHO);
+        let err = engine.transcribe(b"x").await.unwrap_err();
+        assert!(
+            err.contains("could not decode") && err.contains("bad data"),
+            "{err}"
+        );
+    }
+
+    #[tokio::test]
+    async fn a_whisper_failure_is_reported_with_its_stderr() {
+        let dir = tempfile::tempdir().unwrap();
+        let engine = whisper(dir.path(), FFMPEG_OK, "echo 'no model' >&2; exit 2");
+        let err = engine.transcribe(b"x").await.unwrap_err();
+        assert!(
+            err.contains("whisper-cli failed") && err.contains("no model"),
+            "{err}"
+        );
+    }
+}

swarm-stt/src/lib.rs

@@ -0,0 +1,20 @@
+//! The speech-to-text core of the swarm's `swarm-stt` service.
+//!
+//! [`Service::transcribe`] takes the bytes of a voice memo and answers its
+//! text. Three things stand between the bytes and the engine:
+//!
+//! - a size cap, so one oversized file cannot occupy the engine;
+//! - a cache keyed by the SHA-256 of the bytes, so the same memo sent to
+//!   several agents is transcribed once ([`cache`]); and
+//! - an in-flight map, so two requests for one memo that arrive together share
+//!   one engine run instead of starting two.
+//!
+//! The engine is behind the [`Engine`] trait. [`engine::Whisper`] decodes with
+//! `ffmpeg` and transcribes with `whisper-cli`; the tests use a fake.
+
+pub mod cache;
+pub mod engine;
+mod service;
+
+pub use engine::Engine;
+pub use service::{Config, Error, Reply, Service};

swarm-stt/src/service.rs

@@ -0,0 +1,242 @@
+use std::collections::HashMap;
+use std::sync::{Arc, Mutex};
+
+use tokio::sync::{OnceCell, Semaphore};
+
+use crate::cache::{self, Cache};
+use crate::engine::Engine;
+
+/// What [`Service`] enforces around the engine.
+#[derive(Debug, Clone)]
+pub struct Config {
+    /// The largest audio file taken, in bytes.
+    pub max_bytes: usize,
+}
+
+/// Why a request got no transcript.
+#[derive(Debug, thiserror::Error)]
+pub enum Error {
+    /// The audio is larger than [`Config::max_bytes`].
+    #[error("the audio is {size} bytes, over the {max} byte limit")]
+    TooLarge {
+        /// The size received.
+        size: usize,
+        /// The limit.
+        max: usize,
+    },
+    /// The audio has no bytes.
+    #[error("the audio is empty")]
+    Empty,
+    /// The engine failed; the message says how.
+    #[error("{0}")]
+    Engine(String),
+}
+
+/// A transcript, and whether it came from the cache or from a run another
+/// request had already started.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct Reply {
+    /// The text spoken in the audio (empty for silence).
+    pub text: String,
+    /// `true` when no engine run was started for this request.
+    pub cached: bool,
+}
+
+type Flight = Arc<OnceCell<Result<String, String>>>;
+
+/// Transcribes voice memos once each.
+pub struct Service<E> {
+    engine: E,
+    cache: Cache,
+    config: Config,
+    flights: Mutex<HashMap<String, Flight>>,
+    /// One engine run at a time: whisper uses every core it is given.
+    engine_slot: Semaphore,
+}
+
+impl<E: Engine> Service<E> {
+    /// A service running `engine`, caching in `cache`.
+    #[must_use]
+    pub fn new(engine: E, cache: Cache, config: Config) -> Self {
+        Self {
+            engine,
+            cache,
+            config,
+            flights: Mutex::new(HashMap::new()),
+            engine_slot: Semaphore::new(1),
+        }
+    }
+
+    /// The transcript of `audio`: from the cache, from a run already in
+    /// flight for the same bytes, or from a new engine run.
+    pub async fn transcribe(&self, audio: &[u8]) -> Result<Reply, Error> {
+        if audio.is_empty() {
+            return Err(Error::Empty);
+        }
+        if audio.len() > self.config.max_bytes {
+            return Err(Error::TooLarge {
+                size: audio.len(),
+                max: self.config.max_bytes,
+            });
+        }
+        let key = cache::key(audio);
+        if let Some(text) = self.cache.get(&key).await {
+            return Ok(Reply { text, cached: true });
+        }
+        let flight = Arc::clone(
+            self.flights
+                .lock()
+                .expect("the flights lock is never held across a panic")
+                .entry(key.clone())
+                .or_default(),
+        );
+        let mut ran = false;
+        let outcome = flight
+            .get_or_init(|| async {
+                ran = true;
+                self.run(&key, audio).await
+            })
+            .await
+            .clone();
+        // The first caller to finish clears the entry; a later request finds
+        // the cache or starts a new run.
+        self.flights
+            .lock()
+            .expect("the flights lock is never held across a panic")
+            .remove(&key);
+        outcome
+            .map(|text| Reply { text, cached: !ran })
+            .map_err(Error::Engine)
+    }
+
+    async fn run(&self, key: &str, audio: &[u8]) -> Result<String, String> {
+        let _slot = self
+            .engine_slot
+            .acquire()
+            .await
+            .map_err(|_| "the engine is shut down".to_owned())?;
+        let text = self.engine.transcribe(audio).await?;
+        if let Err(e) = self.cache.put(key, &text).await {
+            tracing::warn!(error = %e, "caching a transcript failed");
+        }
+        Ok(text)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+    use std::sync::atomic::{AtomicUsize, Ordering};
+    use std::time::Duration;
+
+    use async_trait::async_trait;
+
+    use super::{Config, Error, Reply, Service};
+    use crate::cache::Cache;
+    use crate::engine::Engine;
+
+    /// Answers `heard: <bytes>` after `delay`, counting its runs.
+    struct Fake {
+        runs: Arc<AtomicUsize>,
+        delay: Duration,
+        fail: bool,
+    }
+
+    #[async_trait]
+    impl Engine for Fake {
+        async fn transcribe(&self, audio: &[u8]) -> Result<String, String> {
+            self.runs.fetch_add(1, Ordering::SeqCst);
+            tokio::time::sleep(self.delay).await;
+            if self.fail {
+                Err("engine broke".to_owned())
+            } else {
+                Ok(format!("heard: {}", String::from_utf8_lossy(audio)))
+            }
+        }
+    }
+
+    fn service(
+        dir: &std::path::Path,
+        delay: Duration,
+        fail: bool,
+    ) -> (Service<Fake>, Arc<AtomicUsize>) {
+        let runs = Arc::new(AtomicUsize::new(0));
+        let fake = Fake {
+            runs: Arc::clone(&runs),
+            delay,
+            fail,
+        };
+        let cache = Cache::new(dir, Duration::from_hours(24));
+        (Service::new(fake, cache, Config { max_bytes: 16 }), runs)
+    }
+
+    #[tokio::test]
+    async fn the_same_memo_is_transcribed_once_and_then_served_from_the_cache() {
+        let dir = tempfile::tempdir().unwrap();
+        let (service, runs) = service(dir.path(), Duration::ZERO, false);
+        let first = service.transcribe(b"memo").await.unwrap();
+        let second = service.transcribe(b"memo").await.unwrap();
+        assert_eq!(
+            first,
+            Reply {
+                text: "heard: memo".into(),
+                cached: false
+            }
+        );
+        assert_eq!(
+            second,
+            Reply {
+                text: "heard: memo".into(),
+                cached: true
+            }
+        );
+        assert_eq!(runs.load(Ordering::SeqCst), 1);
+    }
+
+    #[tokio::test(start_paused = true)]
+    async fn requests_arriving_together_for_one_memo_share_one_run() {
+        let dir = tempfile::tempdir().unwrap();
+        let (service, runs) = service(dir.path(), Duration::from_secs(5), false);
+        let (a, b, c) = tokio::join!(
+            service.transcribe(b"memo"),
+            service.transcribe(b"memo"),
+            service.transcribe(b"memo")
+        );
+        let replies = [a.unwrap(), b.unwrap(), c.unwrap()];
+        assert!(replies.iter().all(|r| r.text == "heard: memo"));
+        assert_eq!(replies.iter().filter(|r| !r.cached).count(), 1);
+        assert_eq!(runs.load(Ordering::SeqCst), 1);
+    }
+
+    #[tokio::test]
+    async fn different_memos_get_their_own_runs() {
+        let dir = tempfile::tempdir().unwrap();
+        let (service, runs) = service(dir.path(), Duration::ZERO, false);
+        service.transcribe(b"one").await.unwrap();
+        service.transcribe(b"two").await.unwrap();
+        assert_eq!(runs.load(Ordering::SeqCst), 2);
+    }
+
+    #[tokio::test]
+    async fn empty_and_oversized_audio_never_reach_the_engine() {
+        let dir = tempfile::tempdir().unwrap();
+        let (service, runs) = service(dir.path(), Duration::ZERO, false);
+        assert!(matches!(service.transcribe(b"").await, Err(Error::Empty)));
+        let big = [0u8; 17];
+        assert!(matches!(
+            service.transcribe(&big).await,
+            Err(Error::TooLarge { size: 17, max: 16 })
+        ));
+        assert_eq!(runs.load(Ordering::SeqCst), 0);
+    }
+
+    #[tokio::test]
+    async fn a_failed_run_is_not_cached_and_the_next_request_runs_again() {
+        let dir = tempfile::tempdir().unwrap();
+        let (service, runs) = service(dir.path(), Duration::ZERO, true);
+        let err = service.transcribe(b"memo").await.unwrap_err();
+        assert!(matches!(err, Error::Engine(ref m) if m == "engine broke"));
+        service.transcribe(b"memo").await.unwrap_err();
+        assert_eq!(runs.load(Ordering::SeqCst), 2);
+    }
+}