swarm-stt: the speech-to-text core (hash cache, shared runs, ffmpeg and whisper-cli engine)
This commit is contained in:
parent
5443c1c9ba
commit
df4b2f1838
9 changed files with 663 additions and 0 deletions
|
|
@ -176,6 +176,10 @@ Swarm-level:
|
||||||
webhooks to todos on the agents they concern, through each agent's own
|
webhooks to todos on the agents they concern, through each agent's own
|
||||||
consumer on the `ROUTER_TODOS` stream. Design:
|
consumer on the `ROUTER_TODOS` stream. Design:
|
||||||
[`docs/4-implementation/swarm/notify-router.md`](docs/4-implementation/swarm/notify-router.md).
|
[`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
|
- **`swarm-review-classifier/`** — one `claude --print` per PR
|
||||||
classification, no tools, no forge credential. Design:
|
classification, no tools, no forge credential. Design:
|
||||||
[`docs/4-implementation/swarm/review-classifier.md`](docs/4-implementation/swarm/review-classifier.md).
|
[`docs/4-implementation/swarm/review-classifier.md`](docs/4-implementation/swarm/review-classifier.md).
|
||||||
|
|
|
||||||
12
Cargo.lock
generated
12
Cargo.lock
generated
|
|
@ -5364,6 +5364,18 @@ dependencies = [
|
||||||
"x509-cert",
|
"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]]
|
[[package]]
|
||||||
name = "swarm-wire"
|
name = "swarm-wire"
|
||||||
version = "0.1.0"
|
version = "0.1.0"
|
||||||
|
|
|
||||||
|
|
@ -46,6 +46,7 @@ members = [
|
||||||
"swarm-logs-auth",
|
"swarm-logs-auth",
|
||||||
"swarm-metrics-auth",
|
"swarm-metrics-auth",
|
||||||
"swarm-secret-client",
|
"swarm-secret-client",
|
||||||
|
"swarm-stt",
|
||||||
"swarm-wire",
|
"swarm-wire",
|
||||||
"swarmctl",
|
"swarmctl",
|
||||||
]
|
]
|
||||||
|
|
|
||||||
41
docs/4-implementation/swarm/stt.md
Normal file
41
docs/4-implementation/swarm/stt.md
Normal file
|
|
@ -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 `<sha256>.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.
|
||||||
22
swarm-stt/Cargo.toml
Normal file
22
swarm-stt/Cargo.toml
Normal file
|
|
@ -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
|
||||||
149
swarm-stt/src/cache.rs
Normal file
149
swarm-stt/src/cache.rs
Normal file
|
|
@ -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);
|
||||||
|
}
|
||||||
|
}
|
||||||
172
swarm-stt/src/engine.rs
Normal file
172
swarm-stt/src/engine.rs
Normal file
|
|
@ -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}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
20
swarm-stt/src/lib.rs
Normal file
20
swarm-stt/src/lib.rs
Normal file
|
|
@ -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};
|
||||||
242
swarm-stt/src/service.rs
Normal file
242
swarm-stt/src/service.rs
Normal file
|
|
@ -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);
|
||||||
|
}
|
||||||
|
}
|
||||||
Loading…
Reference in a new issue