diff --git a/AGENTS.md b/AGENTS.md index 3508021fb..cb4f22906 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -176,6 +176,10 @@ Swarm-level: webhooks to todos on the agents they concern, through each agent's own consumer on the `ROUTER_TODOS` stream. Design: [`docs/4-implementation/swarm/notify-router.md`](docs/4-implementation/swarm/notify-router.md). +- **`swarm-stt/`** — speech to text for voice memos: a cache keyed by the + SHA-256 of the audio, one shared run per memo, `ffmpeg` and `whisper-cli` + behind an `Engine` trait. Design: + [`docs/4-implementation/swarm/stt.md`](docs/4-implementation/swarm/stt.md). - **`swarm-review-classifier/`** — one `claude --print` per PR classification, no tools, no forge credential. Design: [`docs/4-implementation/swarm/review-classifier.md`](docs/4-implementation/swarm/review-classifier.md). diff --git a/Cargo.lock b/Cargo.lock index 2a5d4f33d..0325825db 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5364,6 +5364,18 @@ dependencies = [ "x509-cert", ] +[[package]] +name = "swarm-stt" +version = "0.1.0" +dependencies = [ + "async-trait", + "sha2 0.11.1", + "tempfile", + "thiserror 2.0.21", + "tokio", + "tracing", +] + [[package]] name = "swarm-wire" version = "0.1.0" diff --git a/Cargo.toml b/Cargo.toml index a243805e5..eaaae612a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -46,6 +46,7 @@ members = [ "swarm-logs-auth", "swarm-metrics-auth", "swarm-secret-client", + "swarm-stt", "swarm-wire", "swarmctl", ] diff --git a/docs/4-implementation/swarm/stt.md b/docs/4-implementation/swarm/stt.md new file mode 100644 index 000000000..394095ad2 --- /dev/null +++ b/docs/4-implementation/swarm/stt.md @@ -0,0 +1,41 @@ +# Speech to text + +`swarm-stt` turns the voice memos people send agents into text, once per memo +for the whole swarm. This page covers the core in the `swarm-stt` crate. The +daemon around it, its deploy options and the matrix side follow in later +changes. + +## What a request does + +`Service::transcribe` takes the bytes of an audio file and answers its text: + +1. It rejects empty audio, and audio over `max_bytes`, before any work. +2. It hashes the bytes with SHA-256. The hash names the cache entry. +3. A cache hit answers at once with `cached: true`. +4. A miss joins the run already in flight for that hash, or starts one. Two + requests for one memo that arrive together share one engine run, and only + the run's starter sees `cached: false`. +5. The engine runs one at a time, because whisper uses every core it gets. +6. The text goes into the cache, then back to every waiting request. A failed + run caches nothing, so the next request tries again. + +## The cache + +One `.txt` file per memo, in a directory. An entry lives for the TTL +counted from its write, and a read of an expired entry removes it. `purge` +removes every expired entry. A write goes to a temporary name first and then moves +into place with a rename, so a reader never sees half a transcript. + +## The engine + +The `Engine` trait takes bytes and returns text. `Whisper` is the real one: +`ffmpeg` decodes any container to 16 kHz mono WAV, cutting it at `max_secs`, +and `whisper-cli` from whisper.cpp reads the WAV with the configured model and +language detection on. A scratch directory per run holds the WAV and goes away +with the run. + +## Settings the daemon will expose + +The model (default `small`), the cache TTL (default one day), the size and +length caps. The design and the choices behind them are in the discussion on +issue 6383. diff --git a/swarm-stt/Cargo.toml b/swarm-stt/Cargo.toml new file mode 100644 index 000000000..a65bc3abc --- /dev/null +++ b/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 diff --git a/swarm-stt/src/cache.rs b/swarm-stt/src/cache.rs new file mode 100644 index 000000000..06329ed71 --- /dev/null +++ b/swarm-stt/src/cache.rs @@ -0,0 +1,149 @@ +//! The transcript cache: one `.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, 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 { + 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 { + 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); + } +} diff --git a/swarm-stt/src/engine.rs b/swarm-stt/src/engine.rs new file mode 100644 index 000000000..b79c0cc35 --- /dev/null +++ b/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; +} + +/// `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 { + 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 { + 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 { + 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}" + ); + } +} diff --git a/swarm-stt/src/lib.rs b/swarm-stt/src/lib.rs new file mode 100644 index 000000000..a751f9d33 --- /dev/null +++ b/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}; diff --git a/swarm-stt/src/service.rs b/swarm-stt/src/service.rs new file mode 100644 index 000000000..633bbb733 --- /dev/null +++ b/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>>; + +/// Transcribes voice memos once each. +pub struct Service { + engine: E, + cache: Cache, + config: Config, + flights: Mutex>, + /// One engine run at a time: whisper uses every core it is given. + engine_slot: Semaphore, +} + +impl Service { + /// 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 { + 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 { + 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: ` after `delay`, counting its runs. + struct Fake { + runs: Arc, + delay: Duration, + fail: bool, + } + + #[async_trait] + impl Engine for Fake { + async fn transcribe(&self, audio: &[u8]) -> Result { + 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, Arc) { + 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); + } +}