first
This commit is contained in:
valenti
+116
@@ -0,0 +1,116 @@
|
||||
//! Startup asset bootstrap.
|
||||
//!
|
||||
//! The filler (end-of-utterance) and backchannel (during-utterance) clip banks
|
||||
//! are plain `*.raw` PCM files. If a clip is missing at startup it is
|
||||
//! synthesized with the configured TTS engine and written to disk, so a fresh
|
||||
//! checkout works without a manual generation step. Existing clips are never
|
||||
//! overwritten.
|
||||
|
||||
use crate::audio::SampleRingBuffer;
|
||||
use crate::tts::TtsEngine;
|
||||
use anyhow::{Context, Result};
|
||||
use std::path::Path;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::{info, warn};
|
||||
|
||||
/// Default end-of-utterance filler clips: `(name, text)`.
|
||||
pub const FILLER_CLIPS: &[(&str, &str)] = &[
|
||||
("uhm", "Uhm..."),
|
||||
("eh", "Eh..."),
|
||||
("mmm", "Mmm..."),
|
||||
("aver", "A ver..."),
|
||||
("bueno", "Bueno..."),
|
||||
("pues", "Pues..."),
|
||||
("vaya", "Vaya..."),
|
||||
];
|
||||
|
||||
/// Default during-utterance backchannel clips: `(name, text)`.
|
||||
pub const BACKCHANNEL_CLIPS: &[(&str, &str)] = &[
|
||||
("mmhmm", "Mm-hmm."),
|
||||
("aja", "Ajá."),
|
||||
("claro", "Claro."),
|
||||
("yaveo", "Ya veo."),
|
||||
];
|
||||
|
||||
/// Capacity (in samples) of the scratch buffer used to capture one synthesized
|
||||
/// clip. 30 s at 44.1 kHz is far beyond any filler / backchannel length.
|
||||
const SYNTH_BUFFER_SAMPLES: usize = 30 * 44_100;
|
||||
|
||||
/// Ensure every clip in `clips` exists as `<dir>/<name>.raw`, synthesizing the
|
||||
/// missing ones with `tts`. Creates `dir` if needed. Returns the number of
|
||||
/// clips generated. Never fails the caller: individual errors are logged.
|
||||
pub async fn ensure_clips(
|
||||
dir: &Path,
|
||||
clips: &[(&str, &str)],
|
||||
tts: &dyn TtsEngine,
|
||||
label: &str,
|
||||
) -> usize {
|
||||
if let Err(e) = std::fs::create_dir_all(dir) {
|
||||
warn!(
|
||||
label,
|
||||
dir = %dir.display(),
|
||||
error = %e,
|
||||
"cannot create clip directory — skipping auto-generation"
|
||||
);
|
||||
return 0;
|
||||
}
|
||||
|
||||
let missing: Vec<&(&str, &str)> = clips
|
||||
.iter()
|
||||
.filter(|(name, _)| !dir.join(format!("{name}.raw")).exists())
|
||||
.collect();
|
||||
if missing.is_empty() {
|
||||
return 0;
|
||||
}
|
||||
|
||||
info!(
|
||||
label,
|
||||
dir = %dir.display(),
|
||||
missing = missing.len(),
|
||||
engine = tts.name(),
|
||||
"generating missing clips at startup"
|
||||
);
|
||||
|
||||
let mut generated = 0usize;
|
||||
for (name, text) in missing {
|
||||
let path = dir.join(format!("{name}.raw"));
|
||||
match synthesize_to_file(tts, text, &path).await {
|
||||
Ok(samples) => {
|
||||
info!(
|
||||
label,
|
||||
name,
|
||||
samples,
|
||||
path = %path.display(),
|
||||
"generated clip"
|
||||
);
|
||||
generated += 1;
|
||||
}
|
||||
Err(e) => warn!(label, name, error = %e, "failed to generate clip"),
|
||||
}
|
||||
}
|
||||
generated
|
||||
}
|
||||
|
||||
/// Synthesize `text` and write the raw 16-bit LE mono PCM to `path`.
|
||||
async fn synthesize_to_file(tts: &dyn TtsEngine, text: &str, path: &Path) -> Result<usize> {
|
||||
let buf = SampleRingBuffer::new(SYNTH_BUFFER_SAMPLES);
|
||||
let cancel = CancellationToken::new();
|
||||
tts.synthesize(text, &buf, &cancel)
|
||||
.await
|
||||
.context("synthesis failed")?;
|
||||
|
||||
let n = buf.len();
|
||||
if n == 0 {
|
||||
anyhow::bail!("synthesis produced no samples");
|
||||
}
|
||||
let mut samples = vec![0i16; n];
|
||||
buf.pop_into(&mut samples);
|
||||
|
||||
let mut bytes = Vec::with_capacity(n * 2);
|
||||
for s in &samples {
|
||||
bytes.extend_from_slice(&s.to_le_bytes());
|
||||
}
|
||||
std::fs::write(path, &bytes)
|
||||
.with_context(|| format!("failed to write {}", path.display()))?;
|
||||
Ok(n)
|
||||
}
|
||||
@@ -0,0 +1,262 @@
|
||||
//! 20 ms audio frame mixer.
|
||||
//!
|
||||
//! Combines the TTS channel, the filler (end-of-utterance backchannel) channel
|
||||
//! and the backchannel (during-utterance acknowledgment) channel into a single
|
||||
//! mono PCM stream at 44.1 kHz / 16-bit. The mixer runs on a strict 20 ms ticker
|
||||
//! and must never block on network or model I/O.
|
||||
//!
|
||||
//! To absorb the uneven arrival of Fish Audio TTS chunks, the mixer applies a
|
||||
//! pre-buffer (jitter buffer) on the TTS channel: playback of TTS samples does
|
||||
//! not start until the TTS queue holds at least `tts_prebuffer_samples`
|
||||
//! (configured via `TTS_PREBUFFER_MS`). During this priming window only the
|
||||
//! filler / backchannel channels are emitted.
|
||||
//!
|
||||
//! The pre-buffer is only waited for at the start of a turn. If the queue runs
|
||||
//! dry mid-reply, playback pauses and then resumes as soon as any audio is
|
||||
//! queued again — without refilling the whole jitter buffer, which would add an
|
||||
//! audible pause on top of the actual underrun.
|
||||
|
||||
use crate::audio::ring_buffer::SampleRingBuffer;
|
||||
use crate::config::{OUTPUT_FRAME_BYTES, OUTPUT_FRAME_SAMPLES};
|
||||
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
|
||||
use std::sync::Arc;
|
||||
use tracing::{debug, info};
|
||||
|
||||
/// Number of consecutive empty TTS frames before the mixer pauses playback.
|
||||
/// This absorbs brief inter-chunk gaps without cutting playback short.
|
||||
/// 200 ms / 20 ms = 10 frames.
|
||||
///
|
||||
/// Once the turn has started, resuming after such a pause only requires *any*
|
||||
/// queued audio (see [`AudioMixer::pull_frame`]): the full pre-buffer is waited
|
||||
/// for just once, at the start of the turn, so a transient underrun never
|
||||
/// re-introduces the pre-buffer latency mid-reply.
|
||||
const TTS_DRAIN_GRACE_FRAMES: usize = 10;
|
||||
|
||||
/// Shared audio mixer holding the TTS, filler and backchannel sample queues.
|
||||
#[derive(Clone)]
|
||||
pub struct AudioMixer {
|
||||
/// Samples coming from the SGLang Drama 3 TTS stream.
|
||||
pub tts_queue: SampleRingBuffer,
|
||||
/// Samples coming from preloaded filler PCM (played at end-of-utterance).
|
||||
pub filler_queue: SampleRingBuffer,
|
||||
/// Samples coming from preloaded backchannel PCM (played while the user is
|
||||
/// speaking, so the agent signals it is listening).
|
||||
pub backchannel_queue: SampleRingBuffer,
|
||||
/// Filler gain when TTS is idle.
|
||||
filler_gain_idle: f32,
|
||||
/// Filler gain when TTS is active (ducking).
|
||||
filler_gain_duck: f32,
|
||||
/// Fixed, low gain applied to the backchannel channel so the cue stays in
|
||||
/// the background.
|
||||
backchannel_gain: f32,
|
||||
/// Whether the TTS pre-buffer has been reached and TTS is being played out.
|
||||
tts_primed: Arc<AtomicBool>,
|
||||
/// Count of consecutive frames where the TTS queue was empty while primed.
|
||||
tts_drain_frames: Arc<AtomicUsize>,
|
||||
/// Number of TTS samples that must be queued before playback starts
|
||||
/// (jitter buffer). Configured via `TTS_PREBUFFER_MS`.
|
||||
tts_prebuffer_samples: usize,
|
||||
/// Set once the TTS stream has finished producing audio. Allows the mixer
|
||||
/// to bypass the pre-buffer and drain short utterances whose total length
|
||||
/// is below `tts_prebuffer_samples` (otherwise they are never played).
|
||||
tts_stream_end: Arc<AtomicBool>,
|
||||
/// True from the moment a TTS request is dispatched until the stream has
|
||||
/// finished (or was cancelled). Independent of the queue contents, so it
|
||||
/// also covers the window before the first chunk arrives.
|
||||
tts_stream_active: Arc<AtomicBool>,
|
||||
/// Set the first time the pre-buffer is reached in a turn, i.e. when the
|
||||
/// response audio actually starts playing. Reset by `mark_stream_start`
|
||||
/// and `flush`. Used to decide whether a latency filler is still useful.
|
||||
tts_playback_started: Arc<AtomicBool>,
|
||||
}
|
||||
|
||||
impl AudioMixer {
|
||||
/// Create a new mixer with the given queue capacities (in samples), filler
|
||||
/// gain parameters and backchannel gain.
|
||||
pub fn new(
|
||||
queue_capacity: usize,
|
||||
tts_prebuffer_samples: usize,
|
||||
filler_gain_idle: f32,
|
||||
filler_gain_duck: f32,
|
||||
backchannel_gain: f32,
|
||||
) -> Self {
|
||||
Self {
|
||||
tts_queue: SampleRingBuffer::new(queue_capacity),
|
||||
filler_queue: SampleRingBuffer::new(queue_capacity),
|
||||
backchannel_queue: SampleRingBuffer::new(queue_capacity),
|
||||
filler_gain_idle,
|
||||
filler_gain_duck,
|
||||
backchannel_gain,
|
||||
tts_prebuffer_samples,
|
||||
tts_primed: Arc::new(AtomicBool::new(false)),
|
||||
tts_drain_frames: Arc::new(AtomicUsize::new(0)),
|
||||
tts_stream_end: Arc::new(AtomicBool::new(false)),
|
||||
tts_stream_active: Arc::new(AtomicBool::new(false)),
|
||||
tts_playback_started: Arc::new(AtomicBool::new(false)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Pull exactly one 20 ms output frame.
|
||||
///
|
||||
/// Returns `None` when nothing should be sent (all channels empty *and* TTS
|
||||
/// not yet primed). The returned `Vec<u8>` is always [`OUTPUT_FRAME_BYTES`]
|
||||
/// bytes long and contains little-endian 16-bit mono PCM.
|
||||
///
|
||||
/// Mixing formula:
|
||||
/// `mixed = tts + filler * filler_gain + backchannel * backchannel_gain`.
|
||||
/// The filler gain is dynamically attenuated to `filler_gain_duck` while
|
||||
/// TTS is being played out, otherwise it stays at `filler_gain_idle`; the
|
||||
/// backchannel gain is fixed and low so the cue stays in the background.
|
||||
pub fn pull_frame(&self) -> Option<Vec<u8>> {
|
||||
// TTS pre-buffer / priming logic.
|
||||
let primed = self.tts_primed.load(Ordering::Relaxed);
|
||||
let tts_len = self.tts_queue.len();
|
||||
if !primed {
|
||||
// Prime when the pre-buffer is reached, or when the TTS stream has
|
||||
// already ended and anything is queued. The latter handles short
|
||||
// utterances whose total length never reaches the pre-buffer.
|
||||
//
|
||||
// The full pre-buffer is only waited for on the *first* playback of
|
||||
// a turn. Afterwards (`tts_playback_started`) any queued audio is
|
||||
// enough, so resuming after a transient underrun is immediate
|
||||
// instead of waiting to refill the whole pre-buffer again.
|
||||
let stream_ended = self.tts_stream_end.load(Ordering::Relaxed);
|
||||
let need = if self.tts_playback_started.load(Ordering::Relaxed) {
|
||||
1
|
||||
} else {
|
||||
self.tts_prebuffer_samples
|
||||
};
|
||||
if tts_len >= need || (stream_ended && tts_len > 0) {
|
||||
info!(
|
||||
tts_queue_len = tts_len,
|
||||
prebuffer = self.tts_prebuffer_samples,
|
||||
stream_ended,
|
||||
"TTS pre-buffer reached — starting TTS playback"
|
||||
);
|
||||
self.tts_primed.store(true, Ordering::Relaxed);
|
||||
self.tts_drain_frames.store(0, Ordering::Relaxed);
|
||||
self.tts_playback_started.store(true, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
let tts_active = self.tts_primed.load(Ordering::Relaxed);
|
||||
|
||||
// Track consecutive empty TTS frames while primed so brief inter-chunk
|
||||
// gaps don't cut playback short. Once paused, playback resumes as soon
|
||||
// as any audio is queued (no full pre-buffer refill).
|
||||
if tts_active && tts_len == 0 {
|
||||
let drained = self.tts_drain_frames.fetch_add(1, Ordering::Relaxed) + 1;
|
||||
if drained == TTS_DRAIN_GRACE_FRAMES {
|
||||
info!("TTS queue underrun — pausing playback until audio resumes");
|
||||
self.tts_primed.store(false, Ordering::Relaxed);
|
||||
}
|
||||
} else if tts_len > 0 {
|
||||
self.tts_drain_frames.store(0, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
// Warn when the TTS buffer is getting low (risk of underrun / gap).
|
||||
if tts_active && tts_len < (self.tts_prebuffer_samples / 4) {
|
||||
debug!(
|
||||
tts_queue_len = tts_len,
|
||||
"TTS queue running low — possible audio gap"
|
||||
);
|
||||
}
|
||||
|
||||
// Determine filler gain based on whether TTS is actively playing.
|
||||
let filler_gain = if tts_active {
|
||||
self.filler_gain_duck
|
||||
} else {
|
||||
self.filler_gain_idle
|
||||
};
|
||||
|
||||
// Scratch buffers for one frame — allocation-free hot path.
|
||||
let mut tts_buf = [0i16; OUTPUT_FRAME_SAMPLES];
|
||||
let mut filler_buf = [0i16; OUTPUT_FRAME_SAMPLES];
|
||||
let mut backchannel_buf = [0i16; OUTPUT_FRAME_SAMPLES];
|
||||
|
||||
// Only consume TTS when primed (pre-buffer reached). This prevents
|
||||
// partial frames that cause audible drop-outs.
|
||||
let tts_n = if tts_active {
|
||||
self.tts_queue.pop_into(&mut tts_buf)
|
||||
} else {
|
||||
0
|
||||
};
|
||||
let filler_n = self.filler_queue.pop_into(&mut filler_buf);
|
||||
let backchannel_n = self.backchannel_queue.pop_into(&mut backchannel_buf);
|
||||
|
||||
// If no channel produced anything, skip sending a frame.
|
||||
if tts_n == 0 && filler_n == 0 && backchannel_n == 0 {
|
||||
return None;
|
||||
}
|
||||
|
||||
let mut out = vec![0u8; OUTPUT_FRAME_BYTES];
|
||||
// Mix into the output byte buffer directly to avoid an extra copy.
|
||||
let bytes = out.as_mut_slice();
|
||||
for i in 0..OUTPUT_FRAME_SAMPLES {
|
||||
let tts = tts_buf[i] as i32;
|
||||
// Apply filler / backchannel gains in floating point then round back.
|
||||
let filler_scaled = (filler_buf[i] as f32 * filler_gain).round() as i32;
|
||||
let backchannel_scaled =
|
||||
(backchannel_buf[i] as f32 * self.backchannel_gain).round() as i32;
|
||||
let mixed = tts + filler_scaled + backchannel_scaled;
|
||||
// Clamp to i16 range to prevent overflow / audible pops.
|
||||
let clamped = mixed.clamp(i16::MIN as i32, i16::MAX as i32) as i16;
|
||||
let lo = (clamped & 0xFF) as u8;
|
||||
let hi = ((clamped >> 8) & 0xFF) as u8;
|
||||
bytes[i * 2] = lo;
|
||||
bytes[i * 2 + 1] = hi;
|
||||
}
|
||||
Some(out)
|
||||
}
|
||||
|
||||
/// Signal that a new TTS stream is starting. Re-arms the pre-buffer so the
|
||||
/// first playback of the new turn waits for the full jitter buffer again
|
||||
/// (mid-turn underruns then resume immediately).
|
||||
pub fn mark_stream_start(&self) {
|
||||
self.tts_stream_end.store(false, Ordering::Relaxed);
|
||||
self.tts_stream_active.store(true, Ordering::Relaxed);
|
||||
self.tts_playback_started.store(false, Ordering::Relaxed);
|
||||
self.tts_primed.store(false, Ordering::Relaxed);
|
||||
self.tts_drain_frames.store(0, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Signal that the TTS stream has finished producing audio. The mixer will
|
||||
/// then drain whatever is queued even if it is below the pre-buffer
|
||||
/// threshold, so short utterances are never swallowed by the jitter buffer.
|
||||
pub fn mark_stream_end(&self) {
|
||||
self.tts_stream_end.store(true, Ordering::Relaxed);
|
||||
self.tts_stream_active.store(false, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Whether the agent is producing or playing audio: a TTS request is in
|
||||
/// flight, TTS samples are queued (or being played out), or a filler is
|
||||
/// still queued. Used to gate VAD-driven barge-in, so silence from the user
|
||||
/// does not flush an idle mixer.
|
||||
pub fn is_active(&self) -> bool {
|
||||
self.tts_stream_active.load(Ordering::Relaxed)
|
||||
|| self.tts_primed.load(Ordering::Relaxed)
|
||||
|| self.tts_queue.len() > 0
|
||||
|| self.filler_queue.len() > 0
|
||||
}
|
||||
|
||||
/// Whether the response audio has started playing since the last
|
||||
/// `mark_stream_start`. `false` means the agent is still silent (TTS
|
||||
/// request in flight, pre-buffer not reached yet), i.e. the latency filler
|
||||
/// would still be useful.
|
||||
pub fn tts_playback_started(&self) -> bool {
|
||||
self.tts_playback_started.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
/// Instantly clear all queues. Used for barge-in / interruption so that
|
||||
/// stale TTS, filler and backchannel audio is not played after the user
|
||||
/// speaks.
|
||||
pub fn flush(&self) {
|
||||
self.tts_queue.clear();
|
||||
self.filler_queue.clear();
|
||||
self.backchannel_queue.clear();
|
||||
self.tts_primed.store(false, Ordering::Relaxed);
|
||||
self.tts_drain_frames.store(0, Ordering::Relaxed);
|
||||
self.tts_stream_end.store(false, Ordering::Relaxed);
|
||||
self.tts_stream_active.store(false, Ordering::Relaxed);
|
||||
self.tts_playback_started.store(false, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
//! Audio subsystem: ring buffers and the 20 ms frame mixer.
|
||||
|
||||
pub mod mixer;
|
||||
pub mod ring_buffer;
|
||||
|
||||
pub use mixer::AudioMixer;
|
||||
pub use ring_buffer::SampleRingBuffer;
|
||||
@@ -0,0 +1,83 @@
|
||||
//! Thread-safe bounded sample ring buffer used by the audio mixer.
|
||||
//!
|
||||
//! Stores raw `i16` PCM samples (mono). Designed for a single producer
|
||||
//! (TTS stream / filler injector) and a single consumer (the 20 ms mixer
|
||||
//! ticker). A short fixed-capacity `VecDeque` guarded by a `parking_lot`
|
||||
//! mutex keeps the hot path allocation-free.
|
||||
|
||||
use parking_lot::Mutex;
|
||||
use std::collections::VecDeque;
|
||||
use std::sync::Arc;
|
||||
|
||||
/// A thread-safe, bounded ring buffer of mono `i16` PCM samples.
|
||||
#[derive(Clone)]
|
||||
pub struct SampleRingBuffer {
|
||||
inner: Arc<Mutex<VecDeque<i16>>>,
|
||||
capacity: usize,
|
||||
}
|
||||
|
||||
impl SampleRingBuffer {
|
||||
/// Create a new ring buffer that holds up to `capacity` samples.
|
||||
pub fn new(capacity: usize) -> Self {
|
||||
Self {
|
||||
inner: Arc::new(Mutex::new(VecDeque::with_capacity(capacity))),
|
||||
capacity,
|
||||
}
|
||||
}
|
||||
|
||||
/// Push a slice of samples, dropping the oldest samples when at capacity.
|
||||
/// Returns the number of samples actually accepted.
|
||||
pub fn push(&self, samples: &[i16]) -> usize {
|
||||
let mut buf = self.inner.lock();
|
||||
let available = self.capacity.saturating_sub(buf.len());
|
||||
let to_push = samples.len().min(available);
|
||||
if to_push == 0 {
|
||||
return 0;
|
||||
}
|
||||
// If we still cannot fit everything, drop from the front to make room.
|
||||
if samples.len() > available {
|
||||
let overflow = samples.len() - available;
|
||||
for _ in 0..overflow {
|
||||
buf.pop_front();
|
||||
}
|
||||
}
|
||||
buf.extend(&samples[..to_push]);
|
||||
to_push
|
||||
}
|
||||
|
||||
/// Pop up to `dst.len()` samples into `dst`. Returns the number of samples
|
||||
/// copied. Unfilled slots in `dst` are zeroed.
|
||||
pub fn pop_into(&self, dst: &mut [i16]) -> usize {
|
||||
let mut buf = self.inner.lock();
|
||||
let n = dst.len().min(buf.len());
|
||||
for (i, s) in buf.drain(..n).enumerate() {
|
||||
dst[i] = s;
|
||||
}
|
||||
// Zero-fill the remainder (silence) instead of leaving garbage.
|
||||
for slot in dst[n..].iter_mut() {
|
||||
*slot = 0;
|
||||
}
|
||||
n
|
||||
}
|
||||
|
||||
/// Current number of samples queued.
|
||||
pub fn len(&self) -> usize {
|
||||
self.inner.lock().len()
|
||||
}
|
||||
|
||||
/// Samples that can still be pushed before the buffer is full.
|
||||
pub fn remaining_capacity(&self) -> usize {
|
||||
let buf = self.inner.lock();
|
||||
self.capacity.saturating_sub(buf.len())
|
||||
}
|
||||
|
||||
/// Whether the buffer is empty.
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.inner.lock().is_empty()
|
||||
}
|
||||
|
||||
/// Drop all queued samples (used for barge-in flush).
|
||||
pub fn clear(&self) {
|
||||
self.inner.lock().clear();
|
||||
}
|
||||
}
|
||||
+424
@@ -0,0 +1,424 @@
|
||||
//! Server configuration: port, SGLang endpoint, model paths, and audio constants.
|
||||
|
||||
use crate::orchestrator::BackchannelSettings;
|
||||
use crate::vad::VadSettings;
|
||||
use std::path::PathBuf;
|
||||
|
||||
/// Sample rate of audio sent back to the browser (server -> client).
|
||||
/// Matches Fish Audio Drama 3's native output (44.1 kHz) so no TTS
|
||||
/// resampling is needed — the browser plays the PCM directly.
|
||||
pub const OUTPUT_SAMPLE_RATE: u32 = 44_100;
|
||||
/// Sample rate expected by the Parakeet STT engine.
|
||||
pub const STT_SAMPLE_RATE: u32 = 16_000;
|
||||
/// Frame duration in milliseconds used by the mixer ticker.
|
||||
pub const FRAME_DURATION_MS: u64 = 20;
|
||||
/// Number of samples per output frame (24 kHz * 20 ms).
|
||||
pub const OUTPUT_FRAME_SAMPLES: usize = (OUTPUT_SAMPLE_RATE as u64 * FRAME_DURATION_MS / 1000) as usize;
|
||||
/// Byte length of one output frame (16-bit mono).
|
||||
pub const OUTPUT_FRAME_BYTES: usize = OUTPUT_FRAME_SAMPLES * std::mem::size_of::<i16>();
|
||||
|
||||
/// EOU (end-of-utterance) silence window in milliseconds.
|
||||
pub const EOU_SILENCE_MS: u64 = 300;
|
||||
/// VAD RMS threshold above which audio is considered voice.
|
||||
pub const VAD_RMS_THRESHOLD: f32 = 0.01;
|
||||
|
||||
/// Default system prompt for the reasoning LLM. Written for *spoken* output:
|
||||
/// the reply is read aloud by the TTS engine, so formatting is discouraged.
|
||||
/// The model is asked to mirror the user's language (multilingual assistant).
|
||||
pub const DEFAULT_LLM_SYSTEM_PROMPT: &str =
|
||||
"You are a friendly, knowledgeable voice assistant speaking with a human \
|
||||
through a real-time audio channel. Always answer in the same language the \
|
||||
user used. Your reply is read aloud by a text-to-speech engine, so write \
|
||||
plain conversational prose: no markdown, no bullet points, no numbering, \
|
||||
no emojis, no code blocks and no stage directions. Keep replies short and \
|
||||
natural — one to three sentences — unless the user explicitly asks for a \
|
||||
longer explanation. If you do not know something, say so briefly.";
|
||||
|
||||
/// Appended to the system prompt when `LLM_DIRECT_ANSWERS=true`. Some
|
||||
/// reasoning models (Qwen3 without a vLLM reasoning parser) narrate their
|
||||
/// chain-of-thought directly in `content` instead of a separate field, and the
|
||||
/// TTS would read it aloud. This explicit instruction is the zero-cost,
|
||||
/// server-independent mitigation.
|
||||
pub const DIRECT_ANSWER_DIRECTIVE: &str =
|
||||
"Answer the user directly and never reveal your reasoning: do not output a \
|
||||
thinking process, an analysis, a plan, numbered steps, a deliberation or \
|
||||
any preamble before the answer. Reply with the answer itself only.";
|
||||
|
||||
/// Runtime configuration for the audio agent engine.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Config {
|
||||
/// TCP port the Axum server listens on.
|
||||
pub server_port: u16,
|
||||
/// Selected TTS engine plugin. Supported values: `fish` (Fish Audio /
|
||||
/// SGLang Drama 3), `kokoro`. See [`crate::tts::create_tts_engine`].
|
||||
pub tts_engine: String,
|
||||
/// Base URL of the SGLang-Omni / Fish Audio server (e.g. http://10.8.0.3:8000).
|
||||
pub sglang_url: String,
|
||||
/// Model identifier passed to the SGLang `/v1/audio/speech` endpoint.
|
||||
pub tts_model: String,
|
||||
/// Voice identifier for the TTS engine.
|
||||
pub tts_voice: String,
|
||||
/// Optional reference audio (WAV) path for voice cloning. The path must be
|
||||
/// accessible by the Fish Audio backend server.
|
||||
pub tts_reference_audio: Option<PathBuf>,
|
||||
/// Optional reference transcript. If this looks like a path to a `.txt`
|
||||
/// file, its contents are loaded; otherwise the string itself is used.
|
||||
pub tts_reference_text: Option<String>,
|
||||
/// Filesystem directory containing the Parakeet TDT ONNX pipeline
|
||||
/// (`nemo128.onnx`, `encoder-model*.onnx`, `decoder_joint-model.onnx`,
|
||||
/// `vocab.txt`, `config.json`).
|
||||
pub stt_model_path: PathBuf,
|
||||
/// Filesystem path to the Parakeet tokenizer (`tokenizer.model` SentencePiece
|
||||
/// or `tokenizer.json` HuggingFace). Optional: if missing, transcripts are
|
||||
/// emitted as a placeholder string.
|
||||
pub stt_tokenizer_path: Option<PathBuf>,
|
||||
/// Directory containing preloaded filler PCM files.
|
||||
pub fillers_dir: PathBuf,
|
||||
/// Whether the end-of-turn filler acknowledgment is used at all
|
||||
/// (`FILLER_ENABLED`). When disabled, no filler clip is ever queued.
|
||||
pub filler_enabled: bool,
|
||||
/// How long (ms) the agent waits for the response audio to start before
|
||||
/// playing the filler (`FILLER_DELAY_MS`). If the reply starts speaking
|
||||
/// within this window the filler is dropped — it is only a latency filler,
|
||||
/// so it must never overlap the beginning of the answer. Sent at most once
|
||||
/// per turn.
|
||||
pub filler_delay_ms: u64,
|
||||
/// TTS jitter buffer (ms) the mixer must fill before the response audio
|
||||
/// starts playing (`TTS_PREBUFFER_MS`, default `900`). Larger values absorb
|
||||
/// uneven TTS chunk arrival (fewer drop-outs) at the cost of added
|
||||
/// time-to-first-audio. The full buffer is only required at the start of a
|
||||
/// reply; after a mid-reply underrun the mixer resumes as soon as any audio
|
||||
/// is queued again.
|
||||
pub tts_prebuffer_ms: u64,
|
||||
/// Maximum number of TTS requests kept in flight at once
|
||||
/// (`TTS_PIPELINE_DEPTH`, default `3`).
|
||||
///
|
||||
/// The reply is synthesized chunk by chunk. Instead of waiting for a chunk
|
||||
/// to finish streaming before requesting the next one (which leaves an
|
||||
/// audible gap while the TTS engine "warms up" for every request), the
|
||||
/// engine keeps up to this many requests streaming concurrently and plays
|
||||
/// their audio back in the original order. This hides the per-request
|
||||
/// start-up latency behind the previous chunk's audio. A value of `1`
|
||||
/// restores the old strictly serial behaviour.
|
||||
pub tts_pipeline_depth: usize,
|
||||
/// Filler gain applied when TTS is idle.
|
||||
pub filler_gain_idle: f32,
|
||||
/// Filler gain applied while TTS is active (ducking).
|
||||
pub filler_gain_duck: f32,
|
||||
/// Optional TLS certificate (PEM) for serving HTTPS. Set `TLS_CERT` and
|
||||
/// `TLS_KEY` together to enable HTTPS; when unset the server speaks plain
|
||||
/// HTTP. HTTPS is required for browser microphone access from a non-
|
||||
/// localhost origin (WebRTC needs a secure context).
|
||||
pub tls_cert: Option<PathBuf>,
|
||||
/// TLS private key (PEM) matching [`Config::tls_cert`].
|
||||
pub tls_key: Option<PathBuf>,
|
||||
/// Voice Activity Detection settings (see [`crate::vad`]).
|
||||
pub vad: VadSettings,
|
||||
/// Backchannel settings (see [`crate::orchestrator::backchannel`]).
|
||||
pub backchannel: BackchannelSettings,
|
||||
/// Master switch for the reasoning LLM (`LLM_ENABLED`). When off the
|
||||
/// pipeline runs STT → TTS with no reasoning step.
|
||||
pub llm_enabled: bool,
|
||||
/// Base URL of the OpenAI-compatible endpoint, e.g.
|
||||
/// `http://10.8.0.4:8000/v1`. The client appends `/chat/completions`.
|
||||
pub llm_base_url: String,
|
||||
/// API key sent as `Authorization: Bearer <key>` (`LLM_API_KEY`). Servers
|
||||
/// that do not check it (vLLM, SGLang) accept any non-empty value.
|
||||
pub llm_api_key: String,
|
||||
/// Model identifier (`LLM_MODEL`), e.g. `Qwen/Qwen3.6-35B-A3B`.
|
||||
pub llm_model: String,
|
||||
/// Maximum tokens generated per reply (`LLM_MAX_TOKENS`).
|
||||
pub llm_max_tokens: u32,
|
||||
/// Sampling temperature (`LLM_TEMPERATURE`).
|
||||
pub llm_temperature: f32,
|
||||
/// System prompt prepended to every conversation (`LLM_SYSTEM_PROMPT`).
|
||||
pub llm_system_prompt: String,
|
||||
/// Stream the reply token-by-token (`LLM_STREAM`). Streaming starts the TTS
|
||||
/// on the first sentence (lower latency); disabling it waits for the full
|
||||
/// reply before speaking (simpler, more predictable).
|
||||
pub llm_stream: bool,
|
||||
/// Maximum number of past messages kept in the rolling history
|
||||
/// (`LLM_HISTORY_MESSAGES`).
|
||||
pub llm_history_messages: usize,
|
||||
/// Chain-of-thought handling on the request side (`LLM_ENABLE_THINKING`).
|
||||
/// `Some(false)` sends `chat_template_kwargs.enable_thinking = false`,
|
||||
/// which turns Qwen3-style thinking off on vLLM/SGLang (the reply is then
|
||||
/// the answer only). `None` (default) omits the field entirely so the
|
||||
/// request stays valid for strict OpenAI-compatible servers.
|
||||
pub llm_enable_thinking: Option<bool>,
|
||||
/// Strip the model's chain-of-thought before it reaches the chat box and
|
||||
/// the TTS (`LLM_STRIP_REASONING`, default `true`). Covers the separate
|
||||
/// `reasoning_content` delta and inline think-tag blocks. When `false` the
|
||||
/// raw output is passed through.
|
||||
pub llm_strip_reasoning: bool,
|
||||
/// Append [`DIRECT_ANSWER_DIRECTIVE`] to the system prompt
|
||||
/// (`LLM_DIRECT_ANSWERS`, default `true`).
|
||||
pub llm_direct_answers: bool,
|
||||
/// Minimum characters the reply must hold before the *first* chunk is
|
||||
/// handed to the TTS (`LLM_STREAM_FIRST_CHUNK_CHARS`, default `20`). Kept
|
||||
/// small so the agent starts speaking as soon as the model has produced a
|
||||
/// few words, without waiting for a whole sentence. `0` disables early
|
||||
/// chunking (chunks then end on sentence punctuation only).
|
||||
pub llm_stream_first_chunk_chars: usize,
|
||||
/// Same threshold for every following chunk (`LLM_STREAM_CHUNK_CHARS`,
|
||||
/// default `120`). Larger chunks give better prosody; smaller ones keep the
|
||||
/// audio flowing while the model is still generating. `0` disables.
|
||||
pub llm_stream_chunk_chars: usize,
|
||||
}
|
||||
|
||||
impl Default for Config {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
server_port: 8787,
|
||||
tts_engine: "fish".to_string(),
|
||||
sglang_url: "http://10.8.0.3:8000".to_string(),
|
||||
tts_model: "fishaudio/drama-3".to_string(),
|
||||
tts_voice: "default".to_string(),
|
||||
// Default voice-cloning reference. The transcript is auto-loaded
|
||||
// from the sibling `solis.txt` file when `tts_reference_text` is
|
||||
// left unset. Overridable at runtime from the playground.
|
||||
tts_reference_audio: Some(PathBuf::from("/root/pipecat-voice/voice_profiles/solis.wav")),
|
||||
tts_reference_text: None,
|
||||
stt_model_path: PathBuf::from("models/parakeet"),
|
||||
stt_tokenizer_path: None,
|
||||
fillers_dir: PathBuf::from("assets/fillers"),
|
||||
filler_enabled: true,
|
||||
filler_delay_ms: 2500,
|
||||
tts_prebuffer_ms: 900,
|
||||
tts_pipeline_depth: 3,
|
||||
filler_gain_idle: 1.0,
|
||||
filler_gain_duck: 0.15,
|
||||
tls_cert: None,
|
||||
tls_key: None,
|
||||
vad: VadSettings::default(),
|
||||
backchannel: BackchannelSettings::default(),
|
||||
llm_enabled: false,
|
||||
llm_base_url: String::new(),
|
||||
llm_api_key: "EMPTY".to_string(),
|
||||
llm_model: String::new(),
|
||||
llm_max_tokens: 1024,
|
||||
llm_temperature: 0.7,
|
||||
llm_system_prompt: DEFAULT_LLM_SYSTEM_PROMPT.to_string(),
|
||||
llm_stream: true,
|
||||
llm_history_messages: 16,
|
||||
llm_enable_thinking: None,
|
||||
llm_strip_reasoning: true,
|
||||
llm_direct_answers: true,
|
||||
llm_stream_first_chunk_chars: 20,
|
||||
llm_stream_chunk_chars: 120,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Config {
|
||||
/// Build a [`Config`] from environment variables, falling back to defaults.
|
||||
pub fn from_env() -> Self {
|
||||
let mut cfg = Self::default();
|
||||
if let Ok(port) = std::env::var("SERVER_PORT") {
|
||||
if let Ok(p) = port.parse::<u16>() {
|
||||
cfg.server_port = p;
|
||||
}
|
||||
}
|
||||
if let Ok(e) = std::env::var("TTS_ENGINE") {
|
||||
if !e.is_empty() {
|
||||
cfg.tts_engine = e;
|
||||
}
|
||||
}
|
||||
if let Ok(url) = std::env::var("SGLANG_URL") {
|
||||
cfg.sglang_url = url;
|
||||
}
|
||||
if let Ok(m) = std::env::var("TTS_MODEL") {
|
||||
cfg.tts_model = m;
|
||||
}
|
||||
if let Ok(v) = std::env::var("TTS_VOICE") {
|
||||
cfg.tts_voice = v;
|
||||
}
|
||||
if let Ok(p) = std::env::var("TTS_REFERENCE_AUDIO") {
|
||||
if !p.is_empty() {
|
||||
cfg.tts_reference_audio = Some(PathBuf::from(p));
|
||||
}
|
||||
}
|
||||
if let Ok(t) = std::env::var("TTS_REFERENCE_TEXT") {
|
||||
if !t.is_empty() {
|
||||
cfg.tts_reference_text = Some(t);
|
||||
}
|
||||
}
|
||||
if let Ok(p) = std::env::var("STT_MODEL_PATH") {
|
||||
cfg.stt_model_path = PathBuf::from(p);
|
||||
}
|
||||
if let Ok(p) = std::env::var("STT_TOKENIZER_PATH") {
|
||||
if !p.is_empty() {
|
||||
cfg.stt_tokenizer_path = Some(PathBuf::from(p));
|
||||
}
|
||||
}
|
||||
if let Ok(d) = std::env::var("FILLERS_DIR") {
|
||||
cfg.fillers_dir = PathBuf::from(d);
|
||||
}
|
||||
cfg.filler_enabled = env_bool("FILLER_ENABLED", cfg.filler_enabled);
|
||||
cfg.filler_delay_ms = env_u32("FILLER_DELAY_MS", cfg.filler_delay_ms as u32) as u64;
|
||||
cfg.tts_prebuffer_ms = env_u32("TTS_PREBUFFER_MS", cfg.tts_prebuffer_ms as u32) as u64;
|
||||
cfg.tts_pipeline_depth =
|
||||
env_u32("TTS_PIPELINE_DEPTH", cfg.tts_pipeline_depth as u32).max(1) as usize;
|
||||
if let Ok(p) = std::env::var("TLS_CERT") {
|
||||
if !p.is_empty() {
|
||||
cfg.tls_cert = Some(PathBuf::from(p));
|
||||
}
|
||||
}
|
||||
if let Ok(p) = std::env::var("TLS_KEY") {
|
||||
if !p.is_empty() {
|
||||
cfg.tls_key = Some(PathBuf::from(p));
|
||||
}
|
||||
}
|
||||
cfg.vad = VadSettings::from_env();
|
||||
cfg.backchannel = BackchannelSettings::from_env();
|
||||
cfg.read_llm_env();
|
||||
cfg
|
||||
}
|
||||
|
||||
/// Populate the LLM fields from the environment. Split out of `from_env`
|
||||
/// because the group is large and self-contained.
|
||||
fn read_llm_env(&mut self) {
|
||||
if let Ok(v) = std::env::var("LLM_BASE_URL") {
|
||||
if !v.trim().is_empty() {
|
||||
self.llm_base_url = v.trim().to_string();
|
||||
}
|
||||
}
|
||||
if let Ok(v) = std::env::var("LLM_API_KEY") {
|
||||
self.llm_api_key = v;
|
||||
}
|
||||
if let Ok(v) = std::env::var("LLM_MODEL") {
|
||||
if !v.trim().is_empty() {
|
||||
self.llm_model = v.trim().to_string();
|
||||
}
|
||||
}
|
||||
if let Ok(v) = std::env::var("LLM_SYSTEM_PROMPT") {
|
||||
if !v.trim().is_empty() {
|
||||
self.llm_system_prompt = v;
|
||||
}
|
||||
}
|
||||
self.llm_max_tokens = env_u32("LLM_MAX_TOKENS", self.llm_max_tokens);
|
||||
self.llm_temperature = env_f32("LLM_TEMPERATURE", self.llm_temperature);
|
||||
self.llm_stream = env_bool("LLM_STREAM", self.llm_stream);
|
||||
self.llm_history_messages =
|
||||
env_u32("LLM_HISTORY_MESSAGES", self.llm_history_messages as u32) as usize;
|
||||
self.llm_strip_reasoning = env_bool("LLM_STRIP_REASONING", self.llm_strip_reasoning);
|
||||
self.llm_enable_thinking = env_opt_bool("LLM_ENABLE_THINKING");
|
||||
self.llm_stream_first_chunk_chars = env_u32(
|
||||
"LLM_STREAM_FIRST_CHUNK_CHARS",
|
||||
self.llm_stream_first_chunk_chars as u32,
|
||||
) as usize;
|
||||
self.llm_stream_chunk_chars =
|
||||
env_u32("LLM_STREAM_CHUNK_CHARS", self.llm_stream_chunk_chars as u32) as usize;
|
||||
|
||||
// Direct-answer instruction: appended to whichever system prompt is in
|
||||
// use (built-in default or `LLM_SYSTEM_PROMPT` override).
|
||||
self.llm_direct_answers = env_bool("LLM_DIRECT_ANSWERS", self.llm_direct_answers);
|
||||
if self.llm_direct_answers && !self.llm_system_prompt.contains(DIRECT_ANSWER_DIRECTIVE) {
|
||||
self.llm_system_prompt.push_str("\n\n");
|
||||
self.llm_system_prompt.push_str(DIRECT_ANSWER_DIRECTIVE);
|
||||
}
|
||||
|
||||
// Enabled by default as soon as an endpoint + model are configured.
|
||||
let configured = !self.llm_base_url.is_empty() && !self.llm_model.is_empty();
|
||||
self.llm_enabled = env_bool("LLM_ENABLED", configured);
|
||||
}
|
||||
|
||||
/// TLS certificate path, but only if both cert and key are configured.
|
||||
pub fn tls_enabled(&self) -> Option<(&PathBuf, &PathBuf)> {
|
||||
match (&self.tls_cert, &self.tls_key) {
|
||||
(Some(c), Some(k)) => Some((c, k)),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Full URL of the SGLang audio speech endpoint.
|
||||
pub fn speech_endpoint(&self) -> String {
|
||||
format!("{}/v1/audio/speech", self.sglang_url.trim_end_matches('/'))
|
||||
}
|
||||
|
||||
/// Resolve the voice-cloning reference, if any.
|
||||
///
|
||||
/// Returns `(audio_path, transcript)`. The transcript resolution rules:
|
||||
/// 1. If `tts_reference_text` is set and points to an existing `.txt` file,
|
||||
/// its contents are used.
|
||||
/// 2. Otherwise the raw `tts_reference_text` string is used.
|
||||
/// 3. If no text is provided, we look for a `<audio>.txt` file next to the
|
||||
/// reference WAV.
|
||||
pub fn tts_reference(&self) -> Option<(String, String)> {
|
||||
let audio = self.tts_reference_audio.as_ref()?;
|
||||
let audio_path = audio.to_string_lossy().into_owned();
|
||||
|
||||
// Try the configured reference text first.
|
||||
if let Some(text) = &self.tts_reference_text {
|
||||
let resolved = resolve_text(text);
|
||||
return Some((audio_path, resolved));
|
||||
}
|
||||
|
||||
// Fallback: look for a .txt sibling of the WAV file.
|
||||
let mut txt_path = audio.clone();
|
||||
txt_path.set_extension("txt");
|
||||
if let Ok(content) = std::fs::read_to_string(&txt_path) {
|
||||
return Some((audio_path, content.trim().to_string()));
|
||||
}
|
||||
|
||||
// No text available — send empty string (Fish Audio may infer it).
|
||||
Some((audio_path, String::new()))
|
||||
}
|
||||
}
|
||||
|
||||
/// If `text` is a path to an existing file, return its contents; otherwise
|
||||
/// return the string unchanged.
|
||||
fn resolve_text(text: &str) -> String {
|
||||
if let Ok(content) = std::fs::read_to_string(text) {
|
||||
return content.trim().to_string();
|
||||
}
|
||||
text.to_string()
|
||||
}
|
||||
|
||||
// --- Environment parsing helpers (shared by the config sub-systems) ---------
|
||||
|
||||
/// Parse a boolean env var. Truthy: `1`, `true`, `yes`, `on` (case-insensitive).
|
||||
pub(crate) fn env_bool(name: &str, default: bool) -> bool {
|
||||
match std::env::var(name) {
|
||||
Ok(v) => matches!(
|
||||
v.trim().to_ascii_lowercase().as_str(),
|
||||
"1" | "true" | "yes" | "on"
|
||||
),
|
||||
Err(_) => default,
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse an optional boolean env var: `None` when unset or malformed.
|
||||
/// Truthy: `1`, `true`, `yes`, `on`; falsy: `0`, `false`, `no`, `off`
|
||||
/// (case-insensitive).
|
||||
pub(crate) fn env_opt_bool(name: &str) -> Option<bool> {
|
||||
match std::env::var(name).ok()?.trim().to_ascii_lowercase().as_str() {
|
||||
"1" | "true" | "yes" | "on" => Some(true),
|
||||
"0" | "false" | "no" | "off" => Some(false),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse an `f32` env var, falling back to `default` when unset or malformed.
|
||||
pub(crate) fn env_f32(name: &str, default: f32) -> f32 {
|
||||
std::env::var(name)
|
||||
.ok()
|
||||
.and_then(|v| v.trim().parse().ok())
|
||||
.unwrap_or(default)
|
||||
}
|
||||
|
||||
/// Parse a `u32` env var, falling back to `default` when unset or malformed.
|
||||
pub(crate) fn env_u32(name: &str, default: u32) -> u32 {
|
||||
std::env::var(name)
|
||||
.ok()
|
||||
.and_then(|v| v.trim().parse().ok())
|
||||
.unwrap_or(default)
|
||||
}
|
||||
|
||||
/// Parse a `u64` env var, falling back to `default` when unset or malformed.
|
||||
pub(crate) fn env_u64(name: &str, default: u64) -> u64 {
|
||||
std::env::var(name)
|
||||
.ok()
|
||||
.and_then(|v| v.trim().parse().ok())
|
||||
.unwrap_or(default)
|
||||
}
|
||||
+451
@@ -0,0 +1,451 @@
|
||||
//! Minimal OpenAI-compatible LLM client (chat completions).
|
||||
//!
|
||||
//! This is the "reasoning" stage of the speech-to-speech pipeline: the user
|
||||
//! transcript (from STT or typed in the chat box) is turned into a reply that
|
||||
//! the TTS engine then speaks. It works with any OpenAI-compatible server
|
||||
//! (vLLM, SGLang, llama.cpp, Ollama, OpenAI itself) thanks to the `/v1`
|
||||
//! convention: `POST {LLM_BASE_URL}/chat/completions`.
|
||||
//!
|
||||
//! Reasoning models (e.g. Qwen3) may leak their chain-of-thought either in a
|
||||
//! separate `reasoning_content` delta field (vLLM) or inline between think
|
||||
//! markers (llama.cpp / other servers), or — when the server runs no reasoning
|
||||
//! parser at all — as plain narration inside `content`. None of it should be
|
||||
//! displayed in the chat box nor read aloud, so three independent mechanisms
|
||||
//! are available, each switched from the environment:
|
||||
//! - `LLM_ENABLE_THINKING` — send `chat_template_kwargs.enable_thinking`
|
||||
//! (`false` turns the thinking mode off server-side; unset omits the field).
|
||||
//! - `LLM_STRIP_REASONING` — drop `reasoning_content` / inline think blocks.
|
||||
//! - `LLM_DIRECT_ANSWERS` — strengthen the system prompt (see
|
||||
//! [`crate::config::DIRECT_ANSWER_DIRECTIVE`]).
|
||||
//!
|
||||
//! Two entry points are provided, selectable at runtime via `LLM_STREAM`:
|
||||
//! - [`LlmClient::chat_stream`] — SSE token streaming (lowest time-to-first-audio).
|
||||
//! - [`LlmClient::chat`] — single blocking request (simplest, no partial output).
|
||||
|
||||
use crate::config::Config;
|
||||
use anyhow::{anyhow, Result};
|
||||
use futures_util::StreamExt;
|
||||
use reqwest::Client;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::json;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::warn;
|
||||
|
||||
/// One entry of the OpenAI `messages` array.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ChatMessage {
|
||||
pub role: String,
|
||||
pub content: String,
|
||||
}
|
||||
|
||||
impl ChatMessage {
|
||||
pub fn system(content: impl Into<String>) -> Self {
|
||||
Self { role: "system".into(), content: content.into() }
|
||||
}
|
||||
|
||||
pub fn user(content: impl Into<String>) -> Self {
|
||||
Self { role: "user".into(), content: content.into() }
|
||||
}
|
||||
|
||||
pub fn assistant(content: impl Into<String>) -> Self {
|
||||
Self { role: "assistant".into(), content: content.into() }
|
||||
}
|
||||
}
|
||||
|
||||
/// Client for an OpenAI-compatible `/chat/completions` endpoint.
|
||||
#[derive(Clone)]
|
||||
pub struct LlmClient {
|
||||
http: Client,
|
||||
/// Fully-qualified chat-completions URL (`{base}/chat/completions`).
|
||||
endpoint: String,
|
||||
model: String,
|
||||
api_key: String,
|
||||
max_tokens: u32,
|
||||
temperature: f32,
|
||||
/// `Some(v)` sends `chat_template_kwargs: {"enable_thinking": v}` — the
|
||||
/// vLLM/SGLang switch for Qwen3-style hybrid reasoning models. `None`
|
||||
/// omits the field so strict OpenAI-compatible servers keep working.
|
||||
enable_thinking: Option<bool>,
|
||||
/// Strip the chain-of-thought (`reasoning_content` / inline think blocks)
|
||||
/// from the reply before it is spoken and shown in the chat box.
|
||||
strip_reasoning: bool,
|
||||
}
|
||||
|
||||
impl LlmClient {
|
||||
pub fn new(
|
||||
base_url: &str,
|
||||
api_key: &str,
|
||||
model: &str,
|
||||
max_tokens: u32,
|
||||
temperature: f32,
|
||||
enable_thinking: Option<bool>,
|
||||
strip_reasoning: bool,
|
||||
) -> Self {
|
||||
let endpoint = format!("{}/chat/completions", base_url.trim_end_matches('/'));
|
||||
// No global request timeout: an SSE stream can legitimately stay open
|
||||
// for a while. Only the TCP connect is bounded.
|
||||
let http = Client::builder()
|
||||
.connect_timeout(Duration::from_secs(10))
|
||||
.build()
|
||||
.unwrap_or_default();
|
||||
Self {
|
||||
http,
|
||||
endpoint,
|
||||
model: model.to_string(),
|
||||
api_key: api_key.to_string(),
|
||||
max_tokens,
|
||||
temperature,
|
||||
enable_thinking,
|
||||
strip_reasoning,
|
||||
}
|
||||
}
|
||||
|
||||
/// Endpoint URL, for logging.
|
||||
pub fn endpoint(&self) -> &str {
|
||||
&self.endpoint
|
||||
}
|
||||
|
||||
/// Model identifier, for logging.
|
||||
pub fn model(&self) -> &str {
|
||||
&self.model
|
||||
}
|
||||
|
||||
/// Requested thinking mode, for logging (`None` = field not sent).
|
||||
pub fn enable_thinking(&self) -> Option<bool> {
|
||||
self.enable_thinking
|
||||
}
|
||||
|
||||
/// Whether chain-of-thought stripping is active, for logging.
|
||||
pub fn strip_reasoning(&self) -> bool {
|
||||
self.strip_reasoning
|
||||
}
|
||||
|
||||
fn request_body(&self, messages: &[ChatMessage], stream: bool) -> serde_json::Value {
|
||||
let mut body = json!({
|
||||
"model": self.model,
|
||||
"messages": messages,
|
||||
"max_tokens": self.max_tokens,
|
||||
"temperature": self.temperature,
|
||||
"stream": stream,
|
||||
});
|
||||
if let Some(enable) = self.enable_thinking {
|
||||
body["chat_template_kwargs"] = json!({ "enable_thinking": enable });
|
||||
}
|
||||
body
|
||||
}
|
||||
|
||||
/// The text a delta carries: the assistant `content` plus — only when
|
||||
/// stripping is off — the separate `reasoning_content` field. With
|
||||
/// stripping on, `reasoning_content` is dropped whole.
|
||||
fn delta_text(&self, delta: &serde_json::Value) -> String {
|
||||
let content = delta["content"].as_str().unwrap_or("");
|
||||
if self.strip_reasoning {
|
||||
return content.to_string();
|
||||
}
|
||||
let reasoning = delta["reasoning_content"].as_str().unwrap_or("");
|
||||
let mut out = String::with_capacity(content.len() + reasoning.len());
|
||||
out.push_str(content);
|
||||
out.push_str(reasoning);
|
||||
out
|
||||
}
|
||||
|
||||
/// Apply (or skip) chain-of-thought stripping to a complete reply.
|
||||
fn clean(&self, text: &str) -> String {
|
||||
if self.strip_reasoning {
|
||||
strip_think(text)
|
||||
} else {
|
||||
text.to_string()
|
||||
}
|
||||
}
|
||||
|
||||
/// Single blocking request: returns the full reply once complete.
|
||||
pub async fn chat(
|
||||
&self,
|
||||
messages: &[ChatMessage],
|
||||
cancel: &CancellationToken,
|
||||
) -> Result<String> {
|
||||
let body = self.request_body(messages, false);
|
||||
let req = self
|
||||
.http
|
||||
.post(&self.endpoint)
|
||||
.header("Authorization", format!("Bearer {}", self.api_key))
|
||||
.json(&body)
|
||||
.send();
|
||||
|
||||
let resp = tokio::select! {
|
||||
r = req => r?,
|
||||
_ = cancel.cancelled() => return Ok(String::new()),
|
||||
};
|
||||
|
||||
let status = resp.status();
|
||||
if !status.is_success() {
|
||||
let text = resp.text().await.unwrap_or_default();
|
||||
return Err(anyhow!("LLM HTTP {}: {}", status, text.trim()));
|
||||
}
|
||||
let value: serde_json::Value = resp.json().await?;
|
||||
let message = &value["choices"][0]["message"];
|
||||
Ok(self.clean(&self.delta_text(message)))
|
||||
}
|
||||
|
||||
/// Streaming request. Every clean (think-stripped) text delta is pushed to
|
||||
/// `delta_tx` as it arrives; the full clean reply is returned at the end.
|
||||
///
|
||||
/// When `cancel` fires, the stream is abandoned and the partial reply
|
||||
/// returned — used for barge-in.
|
||||
pub async fn chat_stream(
|
||||
&self,
|
||||
messages: &[ChatMessage],
|
||||
cancel: &CancellationToken,
|
||||
delta_tx: mpsc::UnboundedSender<String>,
|
||||
) -> Result<String> {
|
||||
let body = self.request_body(messages, true);
|
||||
let resp = self
|
||||
.http
|
||||
.post(&self.endpoint)
|
||||
.header("Authorization", format!("Bearer {}", self.api_key))
|
||||
.json(&body)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
let status = resp.status();
|
||||
if !status.is_success() {
|
||||
let text = resp.text().await.unwrap_or_default();
|
||||
return Err(anyhow!("LLM HTTP {}: {}", status, text.trim()));
|
||||
}
|
||||
|
||||
let mut stream = resp.bytes_stream();
|
||||
let mut buffer = String::new();
|
||||
// Only allocate the stripper when chain-of-thought filtering is on.
|
||||
let mut strip = self.strip_reasoning.then(ThinkStripper::new);
|
||||
let mut full = String::new();
|
||||
let mut done = false;
|
||||
|
||||
'outer: while !done {
|
||||
// Race the next chunk against barge-in: a cancelled turn drops the
|
||||
// HTTP response immediately instead of waiting for the next token.
|
||||
let chunk = tokio::select! {
|
||||
next = stream.next() => match next {
|
||||
Some(Ok(c)) => c,
|
||||
Some(Err(e)) => {
|
||||
warn!(error = %e, "LLM stream read error");
|
||||
break;
|
||||
}
|
||||
None => break,
|
||||
},
|
||||
_ = cancel.cancelled() => break,
|
||||
};
|
||||
buffer.push_str(&String::from_utf8_lossy(&chunk));
|
||||
|
||||
// SSE frames are newline-delimited; keep the trailing partial line.
|
||||
while let Some(pos) = buffer.find('\n') {
|
||||
let line: String = buffer.drain(..=pos).collect();
|
||||
let line = line.trim();
|
||||
let Some(data) = line.strip_prefix("data:") else {
|
||||
continue;
|
||||
};
|
||||
let data = data.trim();
|
||||
if data == "[DONE]" {
|
||||
done = true;
|
||||
break;
|
||||
}
|
||||
let Ok(value) = serde_json::from_str::<serde_json::Value>(data) else {
|
||||
continue;
|
||||
};
|
||||
// A separate `reasoning_content` delta (vLLM) is dropped unless
|
||||
// stripping is disabled — see `delta_text`.
|
||||
let content = self.delta_text(&value["choices"][0]["delta"]);
|
||||
if content.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let clean = match strip.as_mut() {
|
||||
Some(s) => s.push(&content),
|
||||
None => content,
|
||||
};
|
||||
if !clean.is_empty() {
|
||||
full.push_str(&clean);
|
||||
if delta_tx.send(clean).is_err() {
|
||||
break 'outer;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(s) = strip.as_mut() {
|
||||
let tail = s.finish();
|
||||
if !tail.is_empty() {
|
||||
full.push_str(&tail);
|
||||
let _ = delta_tx.send(tail);
|
||||
}
|
||||
}
|
||||
Ok(full)
|
||||
}
|
||||
}
|
||||
|
||||
/// Build the LLM client from the environment, or `None` when disabled /
|
||||
/// misconfigured (`LLM_ENABLED=false`, or no base URL / model).
|
||||
pub fn create_llm_client(config: &Config) -> Option<Arc<LlmClient>> {
|
||||
if !config.llm_enabled {
|
||||
return None;
|
||||
}
|
||||
let base = config.llm_base_url.trim();
|
||||
let model = config.llm_model.trim();
|
||||
if base.is_empty() || model.is_empty() {
|
||||
warn!("LLM enabled but LLM_BASE_URL / LLM_MODEL is empty — LLM disabled");
|
||||
return None;
|
||||
}
|
||||
Some(Arc::new(LlmClient::new(
|
||||
base,
|
||||
&config.llm_api_key,
|
||||
model,
|
||||
config.llm_max_tokens,
|
||||
config.llm_temperature,
|
||||
config.llm_enable_thinking,
|
||||
config.llm_strip_reasoning,
|
||||
)))
|
||||
}
|
||||
|
||||
// --- Chain-of-thought stripping --------------------------------------------
|
||||
|
||||
// Opening / closing chain-of-thought markers emitted by Qwen3-style reasoning
|
||||
// models. Written with `\u{..}` escapes so the literal tags never appear in
|
||||
// this source file as markup.
|
||||
const THINK_OPEN: &str = "\u{3C}think\u{3E}";
|
||||
const THINK_CLOSE: &str = "\u{3C}/think\u{3E}";
|
||||
|
||||
/// Remove every ` thinking…<｜end▁of▁thinking｜>` block from `text`.
|
||||
///
|
||||
/// An unterminated opening tag drops the remainder of the text (the model was
|
||||
/// still reasoning when the stream ended).
|
||||
fn strip_think(text: &str) -> String {
|
||||
let mut out = String::with_capacity(text.len());
|
||||
let mut rest = text;
|
||||
loop {
|
||||
match rest.find(THINK_OPEN) {
|
||||
Some(i) => {
|
||||
out.push_str(&rest[..i]);
|
||||
let after = &rest[i + THINK_OPEN.len()..];
|
||||
match after.find(THINK_CLOSE) {
|
||||
Some(j) => rest = &after[j + THINK_CLOSE.len()..],
|
||||
None => return out,
|
||||
}
|
||||
}
|
||||
None => {
|
||||
out.push_str(rest);
|
||||
return out;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Incremental think-stripper for a stream of deltas.
|
||||
///
|
||||
/// Recomputes the clean text from the accumulated raw text on every delta, so
|
||||
/// a ` thinking` block that spans several deltas is removed as a whole. A short
|
||||
/// tail is held back because the last few characters may belong to an opening
|
||||
/// tag split across deltas.
|
||||
struct ThinkStripper {
|
||||
raw: String,
|
||||
/// Bytes of clean text already returned to the caller.
|
||||
emitted: usize,
|
||||
}
|
||||
|
||||
impl ThinkStripper {
|
||||
fn new() -> Self {
|
||||
Self { raw: String::new(), emitted: 0 }
|
||||
}
|
||||
|
||||
fn push(&mut self, delta: &str) -> String {
|
||||
self.raw.push_str(delta);
|
||||
let clean = strip_think(&self.raw);
|
||||
let safe = floor_char_boundary(&clean, clean.len().saturating_sub(THINK_OPEN.len()));
|
||||
if safe <= self.emitted {
|
||||
return String::new();
|
||||
}
|
||||
let out = clean[self.emitted..safe].to_string();
|
||||
self.emitted = safe;
|
||||
out
|
||||
}
|
||||
|
||||
fn finish(&mut self) -> String {
|
||||
let clean = strip_think(&self.raw);
|
||||
if clean.len() <= self.emitted {
|
||||
return String::new();
|
||||
}
|
||||
let out = clean[self.emitted..].to_string();
|
||||
self.emitted = clean.len();
|
||||
out
|
||||
}
|
||||
}
|
||||
|
||||
/// Largest char boundary `<= i`.
|
||||
fn floor_char_boundary(s: &str, mut i: usize) -> usize {
|
||||
if i >= s.len() {
|
||||
return s.len();
|
||||
}
|
||||
while i > 0 && !s.is_char_boundary(i) {
|
||||
i -= 1;
|
||||
}
|
||||
i
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn strips_closed_think_block() {
|
||||
let text = format!("a{}b c{}d", THINK_OPEN, THINK_CLOSE);
|
||||
assert_eq!(strip_think(&text), "ad");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strips_unterminated_think_block() {
|
||||
// No opening tag yet: the partial text is kept verbatim.
|
||||
assert_eq!(strip_think("hello"), "hello");
|
||||
// Opening tag without a closing tag: everything after it is dropped.
|
||||
let text = format!("hello{}still going", THINK_OPEN);
|
||||
assert_eq!(strip_think(&text), "hello");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_body_thinking_is_opt_in() {
|
||||
let msgs = vec![ChatMessage::user("hi")];
|
||||
// Unset: the vendor-specific field must not be sent.
|
||||
let plain = LlmClient::new("http://x/v1", "k", "m", 16, 0.5, None, true);
|
||||
assert!(plain.request_body(&msgs, true).get("chat_template_kwargs").is_none());
|
||||
// Explicit: forwarded verbatim.
|
||||
let off = LlmClient::new("http://x/v1", "k", "m", 16, 0.5, Some(false), true);
|
||||
let body = off.request_body(&msgs, false);
|
||||
assert_eq!(body["chat_template_kwargs"]["enable_thinking"], false);
|
||||
assert_eq!(body["stream"], false);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stripper_is_stable_across_deltas() {
|
||||
// Split both markers in the middle to exercise the tail-holdback logic.
|
||||
let open = THINK_OPEN;
|
||||
let close = THINK_CLOSE;
|
||||
let parts = vec![
|
||||
"Ho".to_string(),
|
||||
"la ".to_string(),
|
||||
open[..4].to_string(),
|
||||
open[4..].to_string(),
|
||||
"plan".to_string(),
|
||||
"ning".to_string(),
|
||||
close[..4].to_string(),
|
||||
close[4..].to_string(),
|
||||
" answer".to_string(),
|
||||
" done.".to_string(),
|
||||
];
|
||||
let mut s = ThinkStripper::new();
|
||||
let mut out = String::new();
|
||||
for part in &parts {
|
||||
out.push_str(&s.push(part));
|
||||
}
|
||||
out.push_str(&s.finish());
|
||||
assert_eq!(out, "Hola answer done.");
|
||||
}
|
||||
}
|
||||
+268
@@ -0,0 +1,268 @@
|
||||
//! Real-Time Audio Agent Engine entry point.
|
||||
//!
|
||||
//! Axum server exposing a WebSocket endpoint for full-duplex audio streaming,
|
||||
//! a playground client at `/` and a speech-to-speech chat client at `/mivocero`.
|
||||
|
||||
mod assets;
|
||||
mod audio;
|
||||
mod config;
|
||||
mod llm;
|
||||
mod onnx;
|
||||
mod orchestrator;
|
||||
mod stt;
|
||||
mod tls;
|
||||
mod tts;
|
||||
mod vad;
|
||||
mod web;
|
||||
mod ws;
|
||||
|
||||
use crate::llm::create_llm_client;
|
||||
use crate::orchestrator::FillerBank;
|
||||
use crate::stt::{ParakeetStt, Tokenizer};
|
||||
use crate::tts::create_tts_engine;
|
||||
use crate::vad::SileroVad;
|
||||
use crate::web::{MIVOCERO_HTML, PLAYGROUND_HTML};
|
||||
use crate::ws::{ws_handler, AppState};
|
||||
use axum::http::header::{CONTENT_TYPE, CACHE_CONTROL};
|
||||
use axum::response::{Html, IntoResponse, Response};
|
||||
use axum::routing::get;
|
||||
use axum::Router;
|
||||
use std::net::SocketAddr;
|
||||
use std::sync::Arc;
|
||||
use tokio::signal;
|
||||
use tracing::{error, info};
|
||||
use tracing_subscriber::EnvFilter;
|
||||
|
||||
/// Fallback `RUST_LOG` used when the env var is not set: INFO everywhere,
|
||||
/// DEBUG for the TTS and VAD subsystems (per-chunk streaming / latency /
|
||||
/// speech-probability diagnostics).
|
||||
const DEFAULT_LOG_FILTER: &str = "info,audio_agent_engine::tts=debug,audio_agent_engine::vad=debug";
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> anyhow::Result<()> {
|
||||
// Load .env file if present (optional; existing env vars take precedence).
|
||||
if let Err(e) = dotenvy::dotenv() {
|
||||
tracing::debug!(error = %e, "no .env file loaded (using env vars only)");
|
||||
}
|
||||
|
||||
// Initialize tracing with env-filter (RUST_LOG). The default keeps INFO
|
||||
// globally but turns on DEBUG for the TTS subsystem, which is where the
|
||||
// per-chunk streaming / latency diagnostics live.
|
||||
tracing_subscriber::fmt()
|
||||
.with_env_filter(
|
||||
EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| EnvFilter::new(DEFAULT_LOG_FILTER)),
|
||||
)
|
||||
.init();
|
||||
|
||||
let config = Arc::new(config::Config::from_env());
|
||||
info!(port = config.server_port, sglang = %config.sglang_url, "starting audio-agent-engine");
|
||||
|
||||
// Load the Parakeet ONNX model. If it fails, the server still starts with
|
||||
// STT disabled (useful for testing the audio pipeline / playground).
|
||||
let stt = match ParakeetStt::load(&config.stt_model_path) {
|
||||
Ok(s) => Some(s),
|
||||
Err(e) => {
|
||||
tracing::warn!(error = %e, "STT model not loaded — running without speech recognition");
|
||||
None
|
||||
}
|
||||
};
|
||||
|
||||
// Optionally load the tokenizer (SentencePiece .model or HuggingFace .json).
|
||||
let tokenizer = match &config.stt_tokenizer_path {
|
||||
Some(path) => match Tokenizer::load(path) {
|
||||
Ok(t) => {
|
||||
info!(path = %path.display(), vocab = t.vocab_size(), "loaded tokenizer");
|
||||
Some(t)
|
||||
}
|
||||
Err(e) => {
|
||||
error!(error = %e, path = %path.display(), "failed to load tokenizer");
|
||||
None
|
||||
}
|
||||
},
|
||||
None => {
|
||||
info!("no STT tokenizer configured (transcripts will be placeholder)");
|
||||
None
|
||||
}
|
||||
};
|
||||
|
||||
// SGLang TTS client (built before the clip banks so missing filler /
|
||||
// backchannel clips can be synthesized with the configured voice).
|
||||
let tts = create_tts_engine(&config);
|
||||
|
||||
// Ensure the filler / backchannel clip banks exist, generating any missing
|
||||
// clip with the TTS engine (see `crate::assets`).
|
||||
let mut generated = assets::ensure_clips(
|
||||
&config.fillers_dir,
|
||||
assets::FILLER_CLIPS,
|
||||
tts.as_ref(),
|
||||
"filler",
|
||||
)
|
||||
.await;
|
||||
if config.backchannel.enabled {
|
||||
generated += assets::ensure_clips(
|
||||
&config.backchannel.dir,
|
||||
assets::BACKCHANNEL_CLIPS,
|
||||
tts.as_ref(),
|
||||
"backchannel",
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
// Preload filler PCM files (played at end-of-utterance).
|
||||
let fillers = Arc::new(FillerBank::load(&config.fillers_dir, "filler"));
|
||||
|
||||
// Preload backchannel PCM files (played while the user is speaking).
|
||||
let backchannel = Arc::new(FillerBank::load(&config.backchannel.dir, "backchannel"));
|
||||
if config.backchannel.enabled && backchannel.is_empty() {
|
||||
tracing::warn!(
|
||||
dir = %config.backchannel.dir.display(),
|
||||
"backchannel enabled but no clips found — the cue stays silent"
|
||||
);
|
||||
}
|
||||
if generated > 0 {
|
||||
info!(generated, "clip banks bootstrapped");
|
||||
}
|
||||
|
||||
// Silero VAD model. Shared across sessions (per-session state lives in the
|
||||
// detector); if it cannot be loaded the session keeps the RMS fallback.
|
||||
let vad = if config.vad.enabled {
|
||||
match SileroVad::load(&config.vad.model_path) {
|
||||
Ok(v) => Some(v),
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
error = %e,
|
||||
path = %config.vad.model_path.display(),
|
||||
"VAD model not loaded — falling back to RMS voice detection"
|
||||
);
|
||||
None
|
||||
}
|
||||
}
|
||||
} else {
|
||||
tracing::info!("VAD disabled (VAD_ENABLED=false)");
|
||||
None
|
||||
};
|
||||
|
||||
// Reasoning LLM (OpenAI-compatible). When disabled the session runs the
|
||||
// plain STT → TTS pipeline; when enabled the transcript is answered by the
|
||||
// model and the reply is spoken sentence by sentence.
|
||||
let llm = create_llm_client(&config);
|
||||
match &llm {
|
||||
Some(client) => info!(
|
||||
model = client.model(),
|
||||
endpoint = client.endpoint(),
|
||||
stream = config.llm_stream,
|
||||
thinking = ?client.enable_thinking(),
|
||||
strip_reasoning = client.strip_reasoning(),
|
||||
direct_answers = config.llm_direct_answers,
|
||||
"LLM reasoning enabled"
|
||||
),
|
||||
None => info!("LLM reasoning disabled (LLM_ENABLED=false or unconfigured)"),
|
||||
}
|
||||
|
||||
let app_state = AppState {
|
||||
config: config.clone(),
|
||||
tts,
|
||||
fillers,
|
||||
backchannel,
|
||||
stt,
|
||||
tokenizer,
|
||||
vad,
|
||||
llm,
|
||||
};
|
||||
|
||||
let app = Router::new()
|
||||
.route("/", get(playground))
|
||||
.route("/mivocero", get(mivocero))
|
||||
.route("/health", get(health))
|
||||
.route("/ws", get(ws_handler))
|
||||
.with_state(app_state);
|
||||
|
||||
// Bind dual-stack on the IPv6 wildcard: with net.ipv6.bindv6only=0 (default
|
||||
// on Linux) this accepts both IPv6 (::1 / `localhost`) and IPv4 clients.
|
||||
// Binding on 0.0.0.0 only would refuse connections that resolve `localhost`
|
||||
// to ::1, which breaks some preview/forwarding proxies.
|
||||
let addr = SocketAddr::from((std::net::Ipv6Addr::UNSPECIFIED, config.server_port));
|
||||
let listener = tokio::net::TcpListener::bind(addr).await?;
|
||||
|
||||
// HTTPS when both TLS_CERT and TLS_KEY are configured (needed for browser
|
||||
// mic access from a non-localhost origin); plain HTTP otherwise.
|
||||
match config.tls_enabled() {
|
||||
Some((cert, key)) => {
|
||||
let tls_config = tls::load_server_config(cert, key)?;
|
||||
info!(
|
||||
%addr,
|
||||
"listening (HTTPS) — open https://localhost:{}/ for the playground",
|
||||
config.server_port
|
||||
);
|
||||
tls::serve(listener, app, tls_config, shutdown_signal()).await?;
|
||||
}
|
||||
None => {
|
||||
info!(
|
||||
%addr,
|
||||
"listening (HTTP) — open http://localhost:{}/ for the playground",
|
||||
config.server_port
|
||||
);
|
||||
axum::serve(listener, app)
|
||||
.with_graceful_shutdown(shutdown_signal())
|
||||
.await?;
|
||||
}
|
||||
}
|
||||
|
||||
info!("server shut down");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// `GET /` — serve the playground HTML client.
|
||||
async fn playground() -> Response {
|
||||
html_response(PLAYGROUND_HTML)
|
||||
}
|
||||
|
||||
/// `GET /mivocero` — serve the speech-to-speech chat client.
|
||||
async fn mivocero() -> Response {
|
||||
html_response(MIVOCERO_HTML)
|
||||
}
|
||||
|
||||
/// Build an HTML response with caching disabled.
|
||||
fn html_response(body: &'static str) -> Response {
|
||||
let mut resp = Html(body).into_response();
|
||||
resp.headers_mut().insert(
|
||||
CACHE_CONTROL,
|
||||
"no-cache".parse().unwrap(),
|
||||
);
|
||||
resp.headers_mut().insert(
|
||||
CONTENT_TYPE,
|
||||
"text/html; charset=utf-8".parse().unwrap(),
|
||||
);
|
||||
resp
|
||||
}
|
||||
|
||||
/// `GET /health` — liveness probe.
|
||||
async fn health() -> &'static str {
|
||||
"ok\n"
|
||||
}
|
||||
|
||||
/// Graceful shutdown on Ctrl-C or SIGTERM.
|
||||
async fn shutdown_signal() {
|
||||
let ctrl_c = async {
|
||||
signal::ctrl_c()
|
||||
.await
|
||||
.expect("failed to install Ctrl-C handler");
|
||||
};
|
||||
|
||||
#[cfg(unix)]
|
||||
let terminate = async {
|
||||
signal::unix::signal(signal::unix::SignalKind::terminate())
|
||||
.expect("failed to install SIGTERM handler")
|
||||
.recv()
|
||||
.await;
|
||||
};
|
||||
|
||||
#[cfg(not(unix))]
|
||||
let terminate = std::future::pending::<()>();
|
||||
|
||||
tokio::select! {
|
||||
_ = ctrl_c => info!("received Ctrl-C, shutting down"),
|
||||
_ = terminate => info!("received SIGTERM, shutting down"),
|
||||
}
|
||||
}
|
||||
+25
@@ -0,0 +1,25 @@
|
||||
//! Process-wide ONNX Runtime environment.
|
||||
//!
|
||||
//! ONNX Runtime allows a single environment per process (`ort::init().build()`
|
||||
//! fails on the second call), yet every [`Session`](ort::session::Session) must
|
||||
//! be built from it. This module lazily creates that environment once and hands
|
||||
//! out cheap `Arc` clones, so the STT and VAD pipelines share it.
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use ort::environment::Environment;
|
||||
use parking_lot::Mutex;
|
||||
|
||||
static ENVIRONMENT: Mutex<Option<Environment>> = Mutex::new(None);
|
||||
|
||||
/// Return the shared ONNX Runtime environment, creating it on first use.
|
||||
pub fn environment() -> Result<Environment> {
|
||||
let mut guard = ENVIRONMENT.lock();
|
||||
if let Some(env) = guard.as_ref() {
|
||||
return Ok(env.clone());
|
||||
}
|
||||
let env = ort::init()
|
||||
.build()
|
||||
.context("failed to create ONNX Runtime environment")?;
|
||||
*guard = Some(env.clone());
|
||||
Ok(env)
|
||||
}
|
||||
@@ -0,0 +1,147 @@
|
||||
//! Backchannel: brief acknowledgments emitted *while the user is speaking*.
|
||||
//!
|
||||
//! This is distinct from a filler (see [`TurnManager`](crate::orchestrator::TurnManager)),
|
||||
//! which is played once the user stops, at end-of-utterance. A backchannel is a
|
||||
//! short "mm-hmm" / "ajá" mixed in at a reduced gain to signal that the agent is
|
||||
//! listening, without taking the turn. It is emitted only while the agent itself
|
||||
//! is silent (otherwise the user is being interrupted) and never more than
|
||||
//! `max_per_turn` times.
|
||||
//!
|
||||
//! Every behavioural knob is configured from the environment
|
||||
//! ([`BackchannelSettings`]), so the feature can be tuned or disabled without a
|
||||
//! rebuild.
|
||||
|
||||
use crate::config::{env_bool, env_f32, env_u32, env_u64};
|
||||
use crate::orchestrator::FillerBank;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
/// Runtime configuration for the backchannel subsystem, read from `.env`.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct BackchannelSettings {
|
||||
/// Master switch (`BACKCHANNEL_ENABLED`).
|
||||
pub enabled: bool,
|
||||
/// Directory of `*.raw` PCM clips (`BACKCHANNEL_DIR`).
|
||||
pub dir: PathBuf,
|
||||
/// Mix gain applied to backchannel clips (`BACKCHANNEL_GAIN`). Kept low so
|
||||
/// the cue stays in the background.
|
||||
pub gain: f32,
|
||||
/// Continuous speech required before the first backchannel
|
||||
/// (`BACKCHANNEL_MIN_SPEECH_MS`).
|
||||
pub min_speech_ms: u64,
|
||||
/// Minimum spacing between two backchannels (`BACKCHANNEL_INTERVAL_MS`).
|
||||
pub interval_ms: u64,
|
||||
/// Maximum backchannels per user turn (`BACKCHANNEL_MAX_PER_TURN`).
|
||||
pub max_per_turn: u32,
|
||||
}
|
||||
|
||||
impl Default for BackchannelSettings {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
enabled: true,
|
||||
dir: PathBuf::from("assets/backchannel"),
|
||||
gain: 0.35,
|
||||
min_speech_ms: 1_500,
|
||||
interval_ms: 3_000,
|
||||
max_per_turn: 2,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl BackchannelSettings {
|
||||
/// Build the settings from environment variables, falling back to defaults.
|
||||
pub fn from_env() -> Self {
|
||||
let mut s = Self::default();
|
||||
if let Ok(p) = std::env::var("BACKCHANNEL_DIR") {
|
||||
if !p.is_empty() {
|
||||
s.dir = PathBuf::from(p);
|
||||
}
|
||||
}
|
||||
s.enabled = env_bool("BACKCHANNEL_ENABLED", s.enabled);
|
||||
s.gain = env_f32("BACKCHANNEL_GAIN", s.gain);
|
||||
s.min_speech_ms = env_u64("BACKCHANNEL_MIN_SPEECH_MS", s.min_speech_ms);
|
||||
s.interval_ms = env_u64("BACKCHANNEL_INTERVAL_MS", s.interval_ms);
|
||||
s.max_per_turn = env_u32("BACKCHANNEL_MAX_PER_TURN", s.max_per_turn);
|
||||
|
||||
s.gain = s.gain.clamp(0.0, 1.0);
|
||||
// Never cue before the turn has meaningfully started.
|
||||
s.interval_ms = s.interval_ms.max(s.min_speech_ms);
|
||||
s
|
||||
}
|
||||
}
|
||||
|
||||
/// Per-session backchannel decision logic.
|
||||
///
|
||||
/// Fed by the VAD: [`on_speech_start`](Self::on_speech_start) when a user turn
|
||||
/// begins, [`on_speech_end`](Self::on_speech_end) when it closes, and
|
||||
/// [`maybe_emit`](Self::maybe_emit) on every frame while the user is speaking.
|
||||
pub struct BackchannelScheduler {
|
||||
bank: Arc<FillerBank>,
|
||||
settings: BackchannelSettings,
|
||||
/// Round-robin cursor over the clip bank.
|
||||
next_index: usize,
|
||||
/// When the current user turn started (only set while speaking).
|
||||
turn_start: Option<Instant>,
|
||||
/// When the last backchannel was emitted in this turn.
|
||||
last_emitted: Option<Instant>,
|
||||
/// Backchannels emitted in the current turn.
|
||||
emitted: u32,
|
||||
}
|
||||
|
||||
impl BackchannelScheduler {
|
||||
pub fn new(bank: Arc<FillerBank>, settings: BackchannelSettings) -> Self {
|
||||
Self {
|
||||
bank,
|
||||
settings,
|
||||
next_index: 0,
|
||||
turn_start: None,
|
||||
last_emitted: None,
|
||||
emitted: 0,
|
||||
}
|
||||
}
|
||||
|
||||
/// Begin a new user turn (VAD speech start).
|
||||
pub fn on_speech_start(&mut self) {
|
||||
self.turn_start = Some(Instant::now());
|
||||
self.last_emitted = None;
|
||||
self.emitted = 0;
|
||||
}
|
||||
|
||||
/// Close the current user turn (VAD speech end).
|
||||
pub fn on_speech_end(&mut self) {
|
||||
self.turn_start = None;
|
||||
self.last_emitted = None;
|
||||
self.emitted = 0;
|
||||
}
|
||||
|
||||
/// Emit a backchannel into `mixer` if the policy allows it right now.
|
||||
///
|
||||
/// Must be called while the user is speaking and the agent is silent.
|
||||
/// Returns the clip name when one was emitted.
|
||||
pub fn maybe_emit(&mut self, mixer: &crate::audio::AudioMixer) -> Option<String> {
|
||||
if !self.settings.enabled {
|
||||
return None;
|
||||
}
|
||||
let start = self.turn_start?;
|
||||
if self.emitted >= self.settings.max_per_turn {
|
||||
return None;
|
||||
}
|
||||
if start.elapsed() < Duration::from_millis(self.settings.min_speech_ms) {
|
||||
return None;
|
||||
}
|
||||
if let Some(last) = self.last_emitted {
|
||||
if last.elapsed() < Duration::from_millis(self.settings.interval_ms) {
|
||||
return None;
|
||||
}
|
||||
}
|
||||
|
||||
let (name, samples) = self.bank.nth(self.next_index)?;
|
||||
let name = name.to_string();
|
||||
mixer.backchannel_queue.push(samples);
|
||||
self.next_index = self.next_index.wrapping_add(1);
|
||||
self.last_emitted = Some(Instant::now());
|
||||
self.emitted += 1;
|
||||
Some(name)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
//! Turn orchestration: EOU handling, filler injection, TTS dispatch, barge-in.
|
||||
|
||||
pub mod backchannel;
|
||||
pub mod turn_manager;
|
||||
|
||||
pub use backchannel::{BackchannelScheduler, BackchannelSettings};
|
||||
pub use turn_manager::{spawn_orchestrator, FillerBank, TurnManager};
|
||||
@@ -0,0 +1,722 @@
|
||||
//! Turn orchestration: reacts to EOU events from the STT worker, schedules the
|
||||
//! latency filler, runs the reasoning LLM, fires the TTS request, and handles
|
||||
//! barge-in flush.
|
||||
//!
|
||||
//! Two pipeline shapes are supported:
|
||||
//! - **With LLM** (`LLM_ENABLED=true`): transcript → LLM → sentence chunks →
|
||||
//! TTS. The assistant text is also forwarded to the client so a chat box can
|
||||
//! display it.
|
||||
//! - **Without LLM**: transcript → TTS (the original pass-through behaviour).
|
||||
|
||||
use crate::audio::{AudioMixer, SampleRingBuffer};
|
||||
use crate::config::Config;
|
||||
use crate::llm::{ChatMessage, LlmClient};
|
||||
use crate::stt::{Tokenizer, Transcript};
|
||||
use crate::tts::TtsEngine;
|
||||
use serde_json::json;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::{debug, error, info, warn};
|
||||
|
||||
/// Capacity (samples) of the per-request staging buffer used by the TTS
|
||||
/// pipeline — 30 s at the mixer sample rate. Large enough to hold a whole
|
||||
/// chunk even when a later request finishes synthesizing while an earlier one
|
||||
/// is still playing out.
|
||||
const TTS_SLOT_CAPACITY_SAMPLES: usize = crate::config::OUTPUT_SAMPLE_RATE as usize * 30;
|
||||
|
||||
/// One in-flight synthesis request in the TTS pipeline.
|
||||
struct TtsSlot {
|
||||
/// Staging buffer the request streams its PCM into. The relay forwards it
|
||||
/// to the mixer queue in submission order.
|
||||
buffer: SampleRingBuffer,
|
||||
/// Task running the synthesis; finished means the chunk is fully staged.
|
||||
handle: tokio::task::JoinHandle<()>,
|
||||
}
|
||||
|
||||
/// Move every sample currently staged in `src` into `dst`, respecting `dst`'s
|
||||
/// free space so already-queued (i.e. currently playing) audio is never
|
||||
/// dropped. Returns the number of samples moved.
|
||||
fn pump(src: &SampleRingBuffer, dst: &SampleRingBuffer) -> usize {
|
||||
let mut tmp = [0i16; 8192];
|
||||
let mut moved = 0;
|
||||
loop {
|
||||
let want = dst
|
||||
.remaining_capacity()
|
||||
.min(src.len())
|
||||
.min(tmp.len());
|
||||
if want == 0 {
|
||||
break;
|
||||
}
|
||||
let n = src.pop_into(&mut tmp[..want]);
|
||||
if n == 0 {
|
||||
break;
|
||||
}
|
||||
dst.push(&tmp[..n]);
|
||||
moved += n;
|
||||
}
|
||||
moved
|
||||
}
|
||||
|
||||
/// Preloaded PCM samples keyed by name (e.g. "uhm"). Used for the filler bank
|
||||
/// injected at end-of-utterance (played round-robin) and, from a different
|
||||
/// directory, for the backchannel clips played while the user speaks.
|
||||
pub struct FillerBank {
|
||||
samples: HashMap<String, Vec<i16>>,
|
||||
/// Names in a stable order, for round-robin selection.
|
||||
names: Vec<String>,
|
||||
}
|
||||
|
||||
impl FillerBank {
|
||||
/// Load every `*.raw` file in `dir` as raw 16-bit LE mono PCM. `label` is
|
||||
/// only used in log messages (e.g. "filler", "backchannel").
|
||||
pub fn load(dir: &std::path::Path, label: &str) -> Self {
|
||||
let mut samples = HashMap::new();
|
||||
if let Ok(entries) = std::fs::read_dir(dir) {
|
||||
for entry in entries.flatten() {
|
||||
let path = entry.path();
|
||||
if path.extension().and_then(|e| e.to_str()) != Some("raw") {
|
||||
continue;
|
||||
}
|
||||
let name = path
|
||||
.file_stem()
|
||||
.and_then(|n| n.to_str())
|
||||
.unwrap_or("unknown")
|
||||
.to_string();
|
||||
match std::fs::read(&path) {
|
||||
Ok(bytes) => {
|
||||
let pcm = bytes
|
||||
.chunks_exact(2)
|
||||
.map(|c| i16::from_le_bytes([c[0], c[1]]))
|
||||
.collect::<Vec<i16>>();
|
||||
info!(label, name, samples = pcm.len(), "loaded {label} clip");
|
||||
samples.insert(name, pcm);
|
||||
}
|
||||
Err(e) => warn!(label, path = %path.display(), error = %e, "failed to read {label} clip"),
|
||||
}
|
||||
}
|
||||
if samples.is_empty() {
|
||||
warn!(label, dir = %dir.display(), "no {label} clips found");
|
||||
}
|
||||
} else {
|
||||
warn!(label, dir = %dir.display(), "{label} directory not found");
|
||||
}
|
||||
|
||||
let mut names: Vec<String> = samples.keys().cloned().collect();
|
||||
names.sort();
|
||||
Self { samples, names }
|
||||
}
|
||||
|
||||
/// Get the samples for a clip by name. Kept for reference/debugging.
|
||||
#[allow(dead_code)]
|
||||
pub fn get(&self, name: &str) -> Option<&Vec<i16>> {
|
||||
self.samples.get(name)
|
||||
}
|
||||
|
||||
/// Clip at position `index` (wrapping), with its name. Used for the
|
||||
/// round-robin filler and backchannel selection.
|
||||
pub fn nth(&self, index: usize) -> Option<(&str, &Vec<i16>)> {
|
||||
if self.names.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let name = &self.names[index % self.names.len()];
|
||||
self.samples.get(name).map(|s| (name.as_str(), s))
|
||||
}
|
||||
|
||||
/// Whether the bank holds no clips.
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.names.is_empty()
|
||||
}
|
||||
}
|
||||
|
||||
/// Orchestrates a single conversation turn.
|
||||
pub struct TurnManager {
|
||||
mixer: AudioMixer,
|
||||
tts: Arc<dyn TtsEngine>,
|
||||
fillers: Arc<FillerBank>,
|
||||
/// Optional tokenizer to decode STT token IDs into text. If `None`, a
|
||||
/// placeholder string is sent to the TTS (useful for testing the pipeline
|
||||
/// without a real tokenizer).
|
||||
tokenizer: Option<Tokenizer>,
|
||||
/// Cancellation token for the currently running TTS request.
|
||||
cancel: Arc<parking_lot::Mutex<Option<CancellationToken>>>,
|
||||
/// Round-robin cursor over the filler bank (end-of-utterance fillers).
|
||||
filler_cursor: AtomicUsize,
|
||||
/// Reasoning LLM client. `None` => direct STT → TTS pipeline.
|
||||
llm: Option<Arc<LlmClient>>,
|
||||
/// Server configuration (LLM prompt / streaming / history window).
|
||||
config: Arc<Config>,
|
||||
/// Rolling conversation history (user + assistant turns).
|
||||
history: parking_lot::Mutex<Vec<ChatMessage>>,
|
||||
/// Outbound text-event channel (JSON strings) consumed by the session
|
||||
/// sender task and forwarded to the browser (chat box).
|
||||
out_tx: Option<mpsc::UnboundedSender<String>>,
|
||||
}
|
||||
|
||||
impl TurnManager {
|
||||
pub fn new(
|
||||
mixer: AudioMixer,
|
||||
tts: Arc<dyn TtsEngine>,
|
||||
fillers: Arc<FillerBank>,
|
||||
tokenizer: Option<Tokenizer>,
|
||||
llm: Option<Arc<LlmClient>>,
|
||||
config: Arc<Config>,
|
||||
out_tx: Option<mpsc::UnboundedSender<String>>,
|
||||
) -> Self {
|
||||
Self {
|
||||
mixer,
|
||||
tts,
|
||||
fillers,
|
||||
tokenizer,
|
||||
cancel: Arc::new(parking_lot::Mutex::new(None)),
|
||||
filler_cursor: AtomicUsize::new(0),
|
||||
llm,
|
||||
config,
|
||||
history: parking_lot::Mutex::new(Vec::new()),
|
||||
out_tx,
|
||||
}
|
||||
}
|
||||
|
||||
/// Pick the next filler clip round-robin from the bank.
|
||||
fn next_filler(&self) -> Option<(String, Vec<i16>)> {
|
||||
let i = self.filler_cursor.fetch_add(1, Ordering::Relaxed);
|
||||
self.fillers.nth(i).map(|(name, s)| (name.to_string(), s.clone()))
|
||||
}
|
||||
|
||||
/// Schedule the single end-of-turn filler acknowledgment.
|
||||
///
|
||||
/// The filler is a latency mask: it buys time while the reply is being
|
||||
/// generated, so it is only useful when the response audio is *late*. It is
|
||||
/// therefore not played unconditionally — if the reply starts playing
|
||||
/// within `FILLER_DELAY_MS`, the filler is dropped instead of speaking over
|
||||
/// the beginning of the answer. At most one clip is emitted per turn and a
|
||||
/// barge-in (`cancel`) drops the pending one. Disabled entirely by
|
||||
/// `FILLER_ENABLED=false`.
|
||||
fn schedule_filler(&self, cancel: &CancellationToken) {
|
||||
if !self.config.filler_enabled {
|
||||
return;
|
||||
}
|
||||
let Some((name, samples)) = self.next_filler() else {
|
||||
debug!("filler enabled but the bank is empty — skipping");
|
||||
return;
|
||||
};
|
||||
|
||||
let delay_ms = self.config.filler_delay_ms;
|
||||
let mixer = self.mixer.clone();
|
||||
let cancel = cancel.clone();
|
||||
let samples_len = samples.len();
|
||||
tokio::spawn(async move {
|
||||
tokio::select! {
|
||||
_ = tokio::time::sleep(Duration::from_millis(delay_ms)) => {}
|
||||
// Barge-in / new turn: never play a stale acknowledgment.
|
||||
_ = cancel.cancelled() => return,
|
||||
}
|
||||
if mixer.tts_playback_started() {
|
||||
debug!(
|
||||
filler = %name,
|
||||
delay_ms, "response audio already playing — filler skipped"
|
||||
);
|
||||
return;
|
||||
}
|
||||
mixer.filler_queue.push(&samples);
|
||||
debug!(filler = %name, samples = samples_len, delay_ms, "pushed filler (delayed)");
|
||||
});
|
||||
}
|
||||
|
||||
/// Send a JSON event to the client, if a channel is attached.
|
||||
fn emit(&self, value: serde_json::Value) {
|
||||
if let Some(tx) = &self.out_tx {
|
||||
let _ = tx.send(value.to_string());
|
||||
}
|
||||
}
|
||||
|
||||
/// React to an EOU / transcript event from the STT worker.
|
||||
pub async fn on_eou(&self, transcript: &Transcript) {
|
||||
info!(
|
||||
tokens = transcript.token_ids.len(),
|
||||
duration_ms = transcript.duration_ms,
|
||||
text = %transcript.text,
|
||||
"EOU detected"
|
||||
);
|
||||
|
||||
// The Parakeet pipeline decodes text internally; fall back to the
|
||||
// external tokenizer only if no text was produced.
|
||||
let text = if !transcript.text.trim().is_empty() {
|
||||
transcript.text.trim().to_string()
|
||||
} else {
|
||||
match &self.tokenizer {
|
||||
Some(tok) => tok.decode(&transcript.token_ids),
|
||||
None => format!("<transcript {} tokens>", transcript.token_ids.len()),
|
||||
}
|
||||
};
|
||||
if text.is_empty() {
|
||||
debug!("empty transcript after decode, skipping turn");
|
||||
return;
|
||||
}
|
||||
self.handle_user_turn(text).await;
|
||||
}
|
||||
|
||||
/// Handle a user turn typed in the chat box (no STT involved).
|
||||
pub async fn on_user_text(&self, text: &str) {
|
||||
let text = text.trim();
|
||||
if text.is_empty() {
|
||||
return;
|
||||
}
|
||||
self.handle_user_turn(text.to_string()).await;
|
||||
}
|
||||
|
||||
/// Shared entry point for every user turn, whether spoken or typed.
|
||||
///
|
||||
/// 1. Forward the recognized turn to the client chat box.
|
||||
/// 2. Cancel any in-flight reply (barge-in within the turn).
|
||||
/// 3. Run the LLM (if configured) or send the text straight to the TTS.
|
||||
///
|
||||
/// The filler is not queued here: it is scheduled once the response is
|
||||
/// dispatched, so it only plays if the audio is late (see
|
||||
/// [`Self::schedule_filler`]).
|
||||
async fn handle_user_turn(&self, text: String) {
|
||||
self.emit(json!({ "type": "user", "text": text }));
|
||||
|
||||
self.cancel_current_tts();
|
||||
|
||||
match &self.llm {
|
||||
Some(llm) => self.run_llm_turn(llm.clone(), text).await,
|
||||
None => self.dispatch_tts(&text),
|
||||
}
|
||||
}
|
||||
|
||||
/// Run one full LLM turn: stream the reply, chunk it into sentences and
|
||||
/// synthesize each sentence as soon as it is complete.
|
||||
async fn run_llm_turn(&self, llm: Arc<LlmClient>, user_text: String) {
|
||||
// Build the message list from the rolling history.
|
||||
let messages = {
|
||||
let mut hist = self.history.lock();
|
||||
hist.push(ChatMessage::user(&user_text));
|
||||
let max = self.config.llm_history_messages.max(2);
|
||||
if hist.len() > max {
|
||||
let drop = hist.len() - max;
|
||||
hist.drain(..drop);
|
||||
}
|
||||
let mut msgs = Vec::with_capacity(hist.len() + 1);
|
||||
let system = self.config.llm_system_prompt.trim();
|
||||
if !system.is_empty() {
|
||||
msgs.push(ChatMessage::system(system));
|
||||
}
|
||||
msgs.extend(hist.iter().cloned());
|
||||
msgs
|
||||
};
|
||||
|
||||
let cancel = CancellationToken::new();
|
||||
*self.cancel.lock() = Some(cancel.clone());
|
||||
self.mixer.mark_stream_start();
|
||||
self.schedule_filler(&cancel);
|
||||
debug!(
|
||||
endpoint = llm.endpoint(),
|
||||
model = llm.model(),
|
||||
messages = messages.len(),
|
||||
stream = self.config.llm_stream,
|
||||
"dispatching LLM turn"
|
||||
);
|
||||
|
||||
// TTS sink: a *pipelined* synthesizer. Each chunk is requested as soon
|
||||
// as the LLM produces it instead of waiting for the previous chunk's
|
||||
// stream to end, so the per-request start-up latency overlaps the
|
||||
// previous chunk's playback. Requests write into their own staging
|
||||
// buffer; a relay forwards those buffers to the mixer queue strictly in
|
||||
// submission order, which keeps playback in order and gapless.
|
||||
let (sentence_tx, mut sentence_rx) = mpsc::unbounded_channel::<String>();
|
||||
let tts = self.tts.clone();
|
||||
let tts_queue = self.mixer.tts_queue.clone();
|
||||
let cancel_tts = cancel.clone();
|
||||
let mixer_for_tts = self.mixer.clone();
|
||||
let pipeline_depth = self.config.tts_pipeline_depth;
|
||||
let tts_task = tokio::spawn(async move {
|
||||
let (slot_tx, mut slot_rx) = mpsc::unbounded_channel::<TtsSlot>();
|
||||
let permits = Arc::new(tokio::sync::Semaphore::new(pipeline_depth));
|
||||
|
||||
// Collector: start one synthesis task per chunk immediately, so its
|
||||
// start-up latency runs concurrently with the chunks already
|
||||
// playing. `permits` caps how many requests are in flight at once.
|
||||
let collector = {
|
||||
let tts = tts.clone();
|
||||
let cancel = cancel_tts.clone();
|
||||
let permits = permits.clone();
|
||||
tokio::spawn(async move {
|
||||
while let Some(sentence) = sentence_rx.recv().await {
|
||||
if cancel.is_cancelled() {
|
||||
break;
|
||||
}
|
||||
let buffer = SampleRingBuffer::new(TTS_SLOT_CAPACITY_SAMPLES);
|
||||
let buf_task = buffer.clone();
|
||||
let permits = permits.clone();
|
||||
let cancel_task = cancel.clone();
|
||||
let tts_task_engine = tts.clone();
|
||||
let handle = tokio::spawn(async move {
|
||||
let _permit = tokio::select! {
|
||||
p = permits.acquire_owned() => match p {
|
||||
Ok(p) => p,
|
||||
Err(_) => return,
|
||||
},
|
||||
_ = cancel_task.cancelled() => return,
|
||||
};
|
||||
if cancel_task.is_cancelled() {
|
||||
return;
|
||||
}
|
||||
if let Err(e) =
|
||||
tts_task_engine.synthesize(&sentence, &buf_task, &cancel_task).await
|
||||
{
|
||||
error!(error = %e, "TTS synthesis failed (pipelined)");
|
||||
}
|
||||
});
|
||||
if slot_tx.send(TtsSlot { buffer, handle }).is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
})
|
||||
};
|
||||
|
||||
// Relay: forward each chunk's staged audio to the mixer queue in
|
||||
// order, waiting for the current chunk to finish before advancing.
|
||||
while let Some(slot) = slot_rx.recv().await {
|
||||
let started = std::time::Instant::now();
|
||||
debug!(
|
||||
staged = slot.buffer.len(),
|
||||
finished = slot.handle.is_finished(),
|
||||
"TTS pipeline: relaying chunk to mixer"
|
||||
);
|
||||
let mut relayed = 0usize;
|
||||
loop {
|
||||
relayed += pump(&slot.buffer, &tts_queue);
|
||||
if cancel_tts.is_cancelled() {
|
||||
break;
|
||||
}
|
||||
if slot.handle.is_finished() && slot.buffer.is_empty() {
|
||||
debug!(
|
||||
relayed,
|
||||
elapsed_ms = started.elapsed().as_millis(),
|
||||
queue_len = tts_queue.len(),
|
||||
"TTS pipeline chunk fully relayed to the mixer"
|
||||
);
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(2)).await;
|
||||
}
|
||||
if cancel_tts.is_cancelled() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
let _ = collector.await;
|
||||
if !cancel_tts.is_cancelled() {
|
||||
mixer_for_tts.mark_stream_end();
|
||||
}
|
||||
});
|
||||
|
||||
// LLM producer: streams clean text deltas into a channel.
|
||||
let (delta_tx, mut delta_rx) = mpsc::unbounded_channel::<String>();
|
||||
let stream_mode = self.config.llm_stream;
|
||||
let llm_cancel = cancel.clone();
|
||||
let llm_task = tokio::spawn(async move {
|
||||
if stream_mode {
|
||||
llm.chat_stream(&messages, &llm_cancel, delta_tx).await
|
||||
} else {
|
||||
let res = llm.chat(&messages, &llm_cancel).await;
|
||||
if let Ok(text) = &res {
|
||||
let _ = delta_tx.send(text.clone());
|
||||
}
|
||||
res
|
||||
}
|
||||
});
|
||||
|
||||
// Consume deltas: emit `assistant_delta` events and feed the splitter,
|
||||
// so the TTS starts as soon as the first few words are available.
|
||||
let mut splitter = SentenceSplitter::new(
|
||||
self.config.llm_stream_first_chunk_chars,
|
||||
self.config.llm_stream_chunk_chars,
|
||||
);
|
||||
let mut full = String::new();
|
||||
while let Some(delta) = delta_rx.recv().await {
|
||||
if cancel.is_cancelled() {
|
||||
break;
|
||||
}
|
||||
self.emit(json!({ "type": "assistant_delta", "text": delta }));
|
||||
full.push_str(&delta);
|
||||
for sentence in splitter.push(&delta) {
|
||||
if sentence_tx.send(sentence).is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
if !cancel.is_cancelled() {
|
||||
if let Some(rest) = splitter.finish() {
|
||||
let _ = sentence_tx.send(rest);
|
||||
}
|
||||
}
|
||||
drop(sentence_tx);
|
||||
let _ = tts_task.await;
|
||||
|
||||
match llm_task.await {
|
||||
Ok(Ok(_)) => {}
|
||||
Ok(Err(e)) => {
|
||||
error!(error = %e, "LLM request failed");
|
||||
self.emit(json!({ "type": "error", "text": e.to_string() }));
|
||||
}
|
||||
Err(e) => {
|
||||
error!(error = %e, "LLM task join failed");
|
||||
self.emit(json!({ "type": "error", "text": format!("LLM task failed: {e}") }));
|
||||
}
|
||||
}
|
||||
|
||||
if !cancel.is_cancelled() && !full.trim().is_empty() {
|
||||
self.history.lock().push(ChatMessage::assistant(full.clone()));
|
||||
self.emit(json!({ "type": "assistant_end" }));
|
||||
}
|
||||
|
||||
let mut guard = self.cancel.lock();
|
||||
if guard.as_ref().map(|c| c.is_cancelled()).unwrap_or(false) {
|
||||
*guard = None;
|
||||
}
|
||||
}
|
||||
|
||||
/// Handle an explicit client-side interrupt (barge-in).
|
||||
/// Flushes all queued audio and cancels the in-flight LLM/TTS request.
|
||||
pub fn interrupt(&self) {
|
||||
info!("interrupt received, flushing audio");
|
||||
self.cancel_current_tts();
|
||||
self.mixer.flush();
|
||||
}
|
||||
|
||||
/// Whether the agent is currently speaking (TTS in flight or queued audio).
|
||||
/// Used by the VAD-driven barge-in to avoid flushing an idle mixer.
|
||||
pub fn is_speaking(&self) -> bool {
|
||||
self.mixer.is_active()
|
||||
}
|
||||
|
||||
/// Update the voice-cloning reference at runtime. Pass `None` for
|
||||
/// `audio_path` to clear the reference.
|
||||
pub fn set_reference(&self, audio_path: Option<String>, text: Option<String>) {
|
||||
self.tts.set_reference(audio_path, text);
|
||||
}
|
||||
|
||||
/// Synthesize `text` directly with the TTS engine, bypassing the LLM.
|
||||
/// Used by the playground's "Speak" box to test the TTS pipeline.
|
||||
pub async fn speak(&self, text: &str) {
|
||||
let text = text.trim();
|
||||
if text.is_empty() {
|
||||
return;
|
||||
}
|
||||
info!(text_len = text.len(), "speak request (direct TTS)");
|
||||
|
||||
// Cancel any in-flight TTS and clear the mixer queues so the new
|
||||
// utterance starts cleanly (no overlap with the previous one).
|
||||
self.cancel_current_tts();
|
||||
self.mixer.flush();
|
||||
|
||||
self.dispatch_tts(text);
|
||||
}
|
||||
|
||||
/// Fire one TTS request for `text` into the mixer's TTS queue.
|
||||
fn dispatch_tts(&self, text: &str) {
|
||||
let cancel = CancellationToken::new();
|
||||
*self.cancel.lock() = Some(cancel.clone());
|
||||
|
||||
let tts = self.tts.clone();
|
||||
let queue = self.mixer.tts_queue.clone();
|
||||
let cancel_holder = self.cancel.clone();
|
||||
let mixer = self.mixer.clone();
|
||||
let text_owned = text.to_string();
|
||||
self.mixer.mark_stream_start();
|
||||
self.schedule_filler(&cancel);
|
||||
debug!(
|
||||
engine = tts.name(),
|
||||
sample_rate = tts.sample_rate(),
|
||||
text_len = text.len(),
|
||||
"dispatching TTS synthesis"
|
||||
);
|
||||
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = tts.synthesize(&text_owned, &queue, &cancel).await {
|
||||
error!(error = %e, "TTS synthesis failed");
|
||||
}
|
||||
// Signal completion (unless cancelled by a barge-in) so the mixer
|
||||
// drains short utterances instead of waiting for the pre-buffer.
|
||||
if !cancel.is_cancelled() {
|
||||
mixer.mark_stream_end();
|
||||
debug!(queue_len = queue.len(), "TTS stream end signalled to mixer");
|
||||
}
|
||||
let mut guard = cancel_holder.lock();
|
||||
if guard.as_ref().map(|c| c.is_cancelled()).unwrap_or(false) {
|
||||
*guard = None;
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
fn cancel_current_tts(&self) {
|
||||
if let Some(token) = self.cancel.lock().take() {
|
||||
token.cancel();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Spawn an orchestrator loop that consumes transcripts from `rx` and drives
|
||||
/// the turn manager.
|
||||
pub fn spawn_orchestrator(
|
||||
manager: Arc<TurnManager>,
|
||||
mut rx: mpsc::UnboundedReceiver<Transcript>,
|
||||
) {
|
||||
tokio::spawn(async move {
|
||||
while let Some(transcript) = rx.recv().await {
|
||||
manager.on_eou(&transcript).await;
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/// Splits a streaming reply into speech-sized chunks so the TTS can start
|
||||
/// speaking while the model is still generating the rest.
|
||||
///
|
||||
/// A chunk always ends on `.`, `!`, `?`, `;` or a newline. On top of that, a
|
||||
/// word-boundary flush keeps latency bounded when the model runs on without
|
||||
/// punctuation: the *first* chunk is cut once [`first_chunk_chars`] characters
|
||||
/// are buffered (kept small, so the agent starts on the first few words), and
|
||||
/// every following chunk at [`chunk_chars`] (larger, for better prosody). A
|
||||
/// threshold of `0` disables that flush.
|
||||
///
|
||||
/// [`first_chunk_chars`]: SentenceSplitter::first_chunk_chars
|
||||
/// [`chunk_chars`]: SentenceSplitter::chunk_chars
|
||||
struct SentenceSplitter {
|
||||
buf: String,
|
||||
/// Bytes buffered before the first chunk is flushed on a word boundary.
|
||||
first_chunk_chars: usize,
|
||||
/// Same, for every chunk after the first.
|
||||
chunk_chars: usize,
|
||||
/// How many chunks were handed out so far (drives the threshold above).
|
||||
emitted: usize,
|
||||
}
|
||||
|
||||
impl SentenceSplitter {
|
||||
fn new(first_chunk_chars: usize, chunk_chars: usize) -> Self {
|
||||
Self {
|
||||
buf: String::new(),
|
||||
first_chunk_chars,
|
||||
chunk_chars,
|
||||
emitted: 0,
|
||||
}
|
||||
}
|
||||
|
||||
/// Feed a delta and return every complete chunk it produced.
|
||||
fn push(&mut self, delta: &str) -> Vec<String> {
|
||||
self.buf.push_str(delta);
|
||||
let mut out = Vec::new();
|
||||
|
||||
let mut start = 0usize;
|
||||
for (i, b) in self.buf.bytes().enumerate() {
|
||||
// ASCII terminators only: safe to slice at `i + 1`.
|
||||
if matches!(b, b'.' | b'!' | b'?' | b';' | b'\n') {
|
||||
let segment = self.buf[start..=i].trim();
|
||||
if !segment.is_empty() {
|
||||
out.push(segment.to_string());
|
||||
self.emitted += 1;
|
||||
}
|
||||
start = i + 1;
|
||||
}
|
||||
}
|
||||
if start > 0 {
|
||||
self.buf.drain(..start);
|
||||
}
|
||||
|
||||
// Word-boundary flush: cut at the first whitespace at or past the
|
||||
// threshold, so the chunk holds at least that many bytes.
|
||||
let threshold = if self.emitted == 0 {
|
||||
self.first_chunk_chars
|
||||
} else {
|
||||
self.chunk_chars
|
||||
};
|
||||
if threshold > 0 && self.buf.len() >= threshold {
|
||||
let safe = floor_char_boundary(&self.buf, threshold);
|
||||
if let Some(rel) = self.buf[safe..].find(' ') {
|
||||
let end = safe + rel;
|
||||
let segment = self.buf[..end].trim().to_string();
|
||||
if !segment.is_empty() {
|
||||
out.push(segment);
|
||||
self.emitted += 1;
|
||||
}
|
||||
self.buf.drain(..end);
|
||||
}
|
||||
}
|
||||
|
||||
out
|
||||
}
|
||||
|
||||
/// Flush whatever is left once the stream has ended.
|
||||
fn finish(&mut self) -> Option<String> {
|
||||
let rest = self.buf.trim().to_string();
|
||||
self.buf.clear();
|
||||
if rest.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(rest)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Largest char boundary `<= i`.
|
||||
fn floor_char_boundary(s: &str, mut i: usize) -> usize {
|
||||
if i >= s.len() {
|
||||
return s.len();
|
||||
}
|
||||
while i > 0 && !s.is_char_boundary(i) {
|
||||
i -= 1;
|
||||
}
|
||||
i
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn splits_on_sentence_terminators() {
|
||||
let mut s = SentenceSplitter::new(20, 120);
|
||||
assert_eq!(
|
||||
s.push("Hola. ¿Cómo estás? Bien"),
|
||||
vec!["Hola.".to_string(), "¿Cómo estás?".to_string()]
|
||||
);
|
||||
assert_eq!(s.finish(), Some("Bien".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn splits_across_deltas() {
|
||||
let mut s = SentenceSplitter::new(20, 120);
|
||||
assert!(s.push("Hola").is_empty());
|
||||
assert_eq!(s.push(" mundo!"), vec!["Hola mundo!".to_string()]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn first_chunk_flushes_early_on_word_boundary() {
|
||||
let mut s = SentenceSplitter::new(20, 120);
|
||||
let out = s.push("Ciao come stai oggi tutto bene");
|
||||
assert_eq!(out.len(), 1);
|
||||
assert!(out[0].len() >= 20, "chunk too short: {:?}", out[0]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn later_chunks_use_the_larger_threshold() {
|
||||
let mut s = SentenceSplitter::new(20, 120);
|
||||
// First chunk is small and punctuation-free.
|
||||
let first = s.push("Ciao come stai oggi tutto bene");
|
||||
assert_eq!(first.len(), 1);
|
||||
// The next run must reach the larger threshold before flushing.
|
||||
let second = s.push(&"palabra ".repeat(20)); // 160 bytes
|
||||
assert_eq!(second.len(), 1);
|
||||
assert!(second[0].len() >= 120, "chunk too short: {}", second[0].len());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn zero_threshold_disables_word_boundary_flush() {
|
||||
let mut s = SentenceSplitter::new(0, 0);
|
||||
assert!(s.push(&"palabra ".repeat(40)).is_empty());
|
||||
assert!(!s.finish().unwrap_or_default().is_empty());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
//! Speech-to-Text subsystem (NVIDIA Parakeet TDT via ONNX Runtime).
|
||||
|
||||
pub mod parakeet;
|
||||
pub mod tokenizer;
|
||||
|
||||
pub use parakeet::{ParakeetStt, SttWorker, Transcript};
|
||||
pub use tokenizer::Tokenizer;
|
||||
@@ -0,0 +1,705 @@
|
||||
//! NVIDIA Parakeet TDT 0.6B v3 STT engine via ONNX Runtime (pure Rust).
|
||||
//!
|
||||
//! Implements the full NeMo TDT speech-to-text pipeline in Rust, mirroring the
|
||||
//! reference implementation in `istupakov/onnx-asr`:
|
||||
//!
|
||||
//! ```text
|
||||
//! 16 kHz mono waveform
|
||||
//! └─ nemo128.onnx → mel features [1, 128, T]
|
||||
//! └─ encoder-model → encoder outputs [1, 1024, T']
|
||||
//! └─ decoder_joint (greedy TDT) → token IDs
|
||||
//! └─ vocab.txt → text
|
||||
//! ```
|
||||
//!
|
||||
//! The TDT decoder is transitive: at each encoder frame it predicts a token and
|
||||
//! a "duration" (how many frames to advance). A special blank token (id 8192,
|
||||
//! `<blk>`) means "no token emitted for this frame".
|
||||
//!
|
||||
//! A worker task consumes microphone PCM frames, runs RMS-based Voice Activity
|
||||
//! Detection (VAD), and triggers inference once an end-of-utterance (EOU)
|
||||
//! silence window is detected.
|
||||
|
||||
use crate::config::{EOU_SILENCE_MS, STT_SAMPLE_RATE, VAD_RMS_THRESHOLD};
|
||||
use anyhow::{Context, Result};
|
||||
use ndarray::{Array1, ArrayD, IxDyn};
|
||||
use ort::environment::Environment;
|
||||
use ort::session::builder::GraphOptimizationLevel;
|
||||
use ort::session::{Session, SessionInputValue};
|
||||
use ort::value::{Value, ValueType};
|
||||
use parking_lot::Mutex;
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::mpsc;
|
||||
use tracing::{debug, error, info};
|
||||
|
||||
// --- Model files (relative to the model directory) -------------------------
|
||||
|
||||
const CONFIG_FILE: &str = "config.json";
|
||||
const VOCAB_FILE: &str = "vocab.txt";
|
||||
const PREPROCESSOR_FILE: &str = "nemo128.onnx";
|
||||
/// Quantized encoder (preferred when present: ~2-3x faster on CPU, no external
|
||||
/// data file). Drop this file to fall back to the fp32 encoder.
|
||||
const ENCODER_INT8_FILE: &str = "encoder-model.int8.onnx";
|
||||
const ENCODER_FILE: &str = "encoder-model.onnx";
|
||||
const DECODER_FILE: &str = "decoder_joint-model.onnx";
|
||||
|
||||
// --- ONNX tensor names (fixed by the exported graphs) ----------------------
|
||||
|
||||
const PP_IN_WAVEFORMS: &str = "waveforms";
|
||||
const PP_IN_LENS: &str = "waveforms_lens";
|
||||
const PP_OUT_FEATURES: &str = "features";
|
||||
const PP_OUT_LENS: &str = "features_lens";
|
||||
|
||||
const ENC_IN_SIGNAL: &str = "audio_signal";
|
||||
const ENC_IN_LENGTH: &str = "length";
|
||||
const ENC_OUT_OUTPUTS: &str = "outputs";
|
||||
const ENC_OUT_LENGTHS: &str = "encoded_lengths";
|
||||
|
||||
const DEC_IN_ENCODER: &str = "encoder_outputs";
|
||||
const DEC_IN_TARGETS: &str = "targets";
|
||||
const DEC_IN_TARGET_LENGTH: &str = "target_length";
|
||||
const DEC_IN_STATE1: &str = "input_states_1";
|
||||
const DEC_IN_STATE2: &str = "input_states_2";
|
||||
const DEC_OUT_OUTPUTS: &str = "outputs";
|
||||
const DEC_OUT_STATE1: &str = "output_states_1";
|
||||
const DEC_OUT_STATE2: &str = "output_states_2";
|
||||
|
||||
/// A single decoded turn from the STT engine.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Transcript {
|
||||
/// Token IDs emitted by the TDT decoder.
|
||||
pub token_ids: Vec<i64>,
|
||||
/// Decoded, human-readable transcript.
|
||||
pub text: String,
|
||||
/// Approximate duration of the utterance in milliseconds.
|
||||
pub duration_ms: u64,
|
||||
}
|
||||
|
||||
/// The three ONNX sessions making up the Parakeet TDT pipeline.
|
||||
struct Sessions {
|
||||
/// Keeps the ONNX Runtime environment alive for the lifetime of the sessions.
|
||||
_environment: Environment,
|
||||
preprocessor: Session,
|
||||
encoder: Session,
|
||||
decoder_joint: Session,
|
||||
}
|
||||
|
||||
/// Parakeet TDT ONNX pipeline. Cloning is cheap (Arc + Mutex).
|
||||
#[derive(Clone)]
|
||||
pub struct ParakeetStt {
|
||||
sessions: Arc<Mutex<Sessions>>,
|
||||
/// Token id -> piece, with the SentencePiece word marker (`▁`) replaced by a
|
||||
/// regular space.
|
||||
vocab: Arc<Vec<String>>,
|
||||
/// Index of the blank token (`<blk>`).
|
||||
blank_idx: i64,
|
||||
/// Number of token logits (excludes the trailing duration logits).
|
||||
vocab_size: usize,
|
||||
/// Maximum number of tokens emitted per encoder frame before forcing an advance.
|
||||
max_tokens_per_step: usize,
|
||||
/// Encoder hidden size (1024 for Parakeet TDT 0.6B).
|
||||
d_model: usize,
|
||||
/// Number of decoder prediction-network state layers.
|
||||
state_layers: usize,
|
||||
/// Hidden size of each decoder state layer.
|
||||
state_hidden: usize,
|
||||
}
|
||||
|
||||
impl ParakeetStt {
|
||||
/// Load the Parakeet TDT ONNX pipeline from `model_dir`, which must contain
|
||||
/// `nemo128.onnx`, an `encoder-model*.onnx`, `decoder_joint-model.onnx`,
|
||||
/// `vocab.txt` and `config.json`.
|
||||
pub fn load(model_dir: &Path) -> Result<Self> {
|
||||
info!(dir = %model_dir.display(), "loading Parakeet TDT ONNX models");
|
||||
|
||||
// `config.json` carries `max_tokens_per_step` (default 10).
|
||||
let config_json: serde_json::Value = serde_json::from_slice(
|
||||
&std::fs::read(model_dir.join(CONFIG_FILE))
|
||||
.with_context(|| format!("failed to read {}", model_dir.join(CONFIG_FILE).display()))?,
|
||||
)
|
||||
.context("invalid Parakeet config.json")?;
|
||||
let max_tokens_per_step = config_json
|
||||
.get("max_tokens_per_step")
|
||||
.and_then(|v| v.as_u64())
|
||||
.unwrap_or(10) as usize;
|
||||
|
||||
// Vocabulary: id -> piece, with `▁` (U+2581) mapped to a real space.
|
||||
let vocab = Arc::new(load_vocab(&model_dir.join(VOCAB_FILE))?);
|
||||
let vocab_size = vocab.len();
|
||||
let blank_idx = vocab
|
||||
.iter()
|
||||
.position(|p| p == "<blk>")
|
||||
.map(|i| i as i64)
|
||||
.context("vocab.txt does not contain the <blk> token")?;
|
||||
|
||||
// Prefer the int8 encoder when available; fall back to fp32.
|
||||
let int8_path = model_dir.join(ENCODER_INT8_FILE);
|
||||
let encoder_path = if int8_path.exists() {
|
||||
int8_path
|
||||
} else {
|
||||
model_dir.join(ENCODER_FILE)
|
||||
};
|
||||
|
||||
let environment = crate::onnx::environment()?;
|
||||
|
||||
let preprocessor = build_session(&environment, &model_dir.join(PREPROCESSOR_FILE))?;
|
||||
let encoder = build_session(&environment, &encoder_path)?;
|
||||
let decoder_joint = build_session(&environment, &model_dir.join(DECODER_FILE))?;
|
||||
|
||||
// Introspect the graphs for the dimensions we need at runtime.
|
||||
let d_model = output_dim(&encoder, ENC_OUT_OUTPUTS, 1)?;
|
||||
let decoder_out_dim = output_dim(&decoder_joint, DEC_OUT_OUTPUTS, 3)?;
|
||||
let num_durations = decoder_out_dim
|
||||
.checked_sub(vocab_size)
|
||||
.context("decoder output is smaller than the vocabulary")?;
|
||||
let (state_layers, state_hidden) = input_state_dims(&decoder_joint, DEC_IN_STATE1)?;
|
||||
|
||||
info!(
|
||||
encoder = %encoder_path.display(),
|
||||
vocab_size,
|
||||
blank_idx,
|
||||
num_durations,
|
||||
d_model,
|
||||
state_layers,
|
||||
state_hidden,
|
||||
"Parakeet TDT pipeline loaded"
|
||||
);
|
||||
|
||||
Ok(Self {
|
||||
sessions: Arc::new(Mutex::new(Sessions {
|
||||
_environment: environment,
|
||||
preprocessor,
|
||||
encoder,
|
||||
decoder_joint,
|
||||
})),
|
||||
vocab,
|
||||
blank_idx,
|
||||
vocab_size,
|
||||
max_tokens_per_step,
|
||||
d_model,
|
||||
state_layers,
|
||||
state_hidden,
|
||||
})
|
||||
}
|
||||
|
||||
/// Run the full pipeline on a complete utterance (16 kHz mono float PCM
|
||||
/// normalized to [-1, 1]). Returns the emitted token IDs and decoded text.
|
||||
pub fn infer(&self, audio: &[f32]) -> Result<(Vec<i64>, String)> {
|
||||
let n = audio.len();
|
||||
if n == 0 {
|
||||
return Ok((Vec::new(), String::new()));
|
||||
}
|
||||
|
||||
let mut guard = self.sessions.lock();
|
||||
let sessions = &mut *guard;
|
||||
|
||||
// --- 1. Mel-feature preprocessor (nemo128.onnx) ---------------------
|
||||
let (feature_shape, feature_data, feature_len) = {
|
||||
let waveforms = Value::from_array(ArrayD::from_shape_vec(IxDyn(&[1, n]), audio.to_vec())?)
|
||||
.context("failed to build waveform tensor")?;
|
||||
let waveforms_lens = Value::from_array(Array1::from_vec(vec![n as i64]))
|
||||
.context("failed to build waveform length tensor")?;
|
||||
|
||||
let outputs = sessions
|
||||
.preprocessor
|
||||
.run(vec![
|
||||
(PP_IN_WAVEFORMS.to_string(), SessionInputValue::from(waveforms)),
|
||||
(PP_IN_LENS.to_string(), SessionInputValue::from(waveforms_lens)),
|
||||
])
|
||||
.context("nemo128 preprocessor inference failed")?;
|
||||
|
||||
let (shape, data) = outputs
|
||||
.get(PP_OUT_FEATURES)
|
||||
.context("preprocessor produced no 'features' output")?
|
||||
.try_extract_tensor::<f32>()
|
||||
.context("failed to extract preprocessor features")?;
|
||||
let (_, lens) = outputs
|
||||
.get(PP_OUT_LENS)
|
||||
.context("preprocessor produced no 'features_lens' output")?
|
||||
.try_extract_tensor::<i64>()
|
||||
.context("failed to extract preprocessor feature lengths")?;
|
||||
|
||||
let shape: Vec<usize> = shape.iter().map(|&d| d as usize).collect();
|
||||
(shape, data.to_vec(), lens[0])
|
||||
};
|
||||
|
||||
// --- 2. Encoder -----------------------------------------------------
|
||||
// Encoder output layout is [1, d_model, T'] (already transposed by the
|
||||
// exported graph). We re-pack it to [T', d_model] so each frame is
|
||||
// contiguous for the decoder loop.
|
||||
let (frames, encoded_len) = {
|
||||
let signal = Value::from_array(
|
||||
ArrayD::from_shape_vec(IxDyn(&feature_shape), feature_data)
|
||||
.context("failed to reshape mel features")?,
|
||||
)
|
||||
.context("failed to build encoder input tensor")?;
|
||||
let length = Value::from_array(Array1::from_vec(vec![feature_len]))
|
||||
.context("failed to build encoder length tensor")?;
|
||||
|
||||
let outputs = sessions
|
||||
.encoder
|
||||
.run(vec![
|
||||
(ENC_IN_SIGNAL.to_string(), SessionInputValue::from(signal)),
|
||||
(ENC_IN_LENGTH.to_string(), SessionInputValue::from(length)),
|
||||
])
|
||||
.context("encoder inference failed")?;
|
||||
|
||||
let (shape, data) = outputs
|
||||
.get(ENC_OUT_OUTPUTS)
|
||||
.context("encoder produced no 'outputs'")?
|
||||
.try_extract_tensor::<f32>()
|
||||
.context("failed to extract encoder outputs")?;
|
||||
let (_, lens) = outputs
|
||||
.get(ENC_OUT_LENGTHS)
|
||||
.context("encoder produced no 'encoded_lengths'")?
|
||||
.try_extract_tensor::<i64>()
|
||||
.context("failed to extract encoded lengths")?;
|
||||
|
||||
let t = shape[2] as usize;
|
||||
let mut frames = vec![0f32; t * self.d_model];
|
||||
for tt in 0..t {
|
||||
for d in 0..self.d_model {
|
||||
frames[tt * self.d_model + d] = data[d * t + tt];
|
||||
}
|
||||
}
|
||||
// Clamp: low-precision encoders may report more frames than present.
|
||||
let encoded_len = (lens[0] as usize).min(t);
|
||||
(frames, encoded_len)
|
||||
};
|
||||
|
||||
// --- 3. Greedy TDT decoding ----------------------------------------
|
||||
let mut tokens: Vec<i64> = Vec::new();
|
||||
let mut state1 = vec![0f32; self.state_layers * self.state_hidden];
|
||||
let mut state2 = state1.clone();
|
||||
let mut prev_token = self.blank_idx;
|
||||
let mut emitted = 0usize;
|
||||
let mut t = 0usize;
|
||||
|
||||
while t < encoded_len {
|
||||
let frame = &frames[t * self.d_model..(t + 1) * self.d_model];
|
||||
|
||||
let enc_frame = Value::from_array(
|
||||
ArrayD::from_shape_vec(IxDyn(&[1, self.d_model, 1]), frame.to_vec())?,
|
||||
)
|
||||
.context("failed to build decoder encoder-frame tensor")?;
|
||||
let targets = Value::from_array(ArrayD::from_shape_vec(
|
||||
IxDyn(&[1, 1]),
|
||||
vec![prev_token as i32],
|
||||
)?)
|
||||
.context("failed to build decoder targets tensor")?;
|
||||
let target_length = Value::from_array(Array1::from_vec(vec![1i32]))
|
||||
.context("failed to build decoder target-length tensor")?;
|
||||
let s1 = Value::from_array(ArrayD::from_shape_vec(
|
||||
IxDyn(&[self.state_layers, 1, self.state_hidden]),
|
||||
state1.clone(),
|
||||
)?)
|
||||
.context("failed to build decoder state1 tensor")?;
|
||||
let s2 = Value::from_array(ArrayD::from_shape_vec(
|
||||
IxDyn(&[self.state_layers, 1, self.state_hidden]),
|
||||
state2.clone(),
|
||||
)?)
|
||||
.context("failed to build decoder state2 tensor")?;
|
||||
|
||||
let (token, step, new_state1, new_state2) = {
|
||||
let outputs = sessions
|
||||
.decoder_joint
|
||||
.run(vec![
|
||||
(DEC_IN_ENCODER.to_string(), SessionInputValue::from(enc_frame)),
|
||||
(DEC_IN_TARGETS.to_string(), SessionInputValue::from(targets)),
|
||||
(DEC_IN_TARGET_LENGTH.to_string(), SessionInputValue::from(target_length)),
|
||||
(DEC_IN_STATE1.to_string(), SessionInputValue::from(s1)),
|
||||
(DEC_IN_STATE2.to_string(), SessionInputValue::from(s2)),
|
||||
])
|
||||
.context("decoder_joint inference failed")?;
|
||||
|
||||
let (_, logits) = outputs
|
||||
.get(DEC_OUT_OUTPUTS)
|
||||
.context("decoder produced no 'outputs'")?
|
||||
.try_extract_tensor::<f32>()
|
||||
.context("failed to extract decoder logits")?;
|
||||
if logits.len() < self.vocab_size {
|
||||
anyhow::bail!("decoder logits shorter than vocabulary");
|
||||
}
|
||||
// Token logits come first, then the duration logits (TDT).
|
||||
let token = argmax(&logits[..self.vocab_size]) as i64;
|
||||
let step = argmax(&logits[self.vocab_size..]);
|
||||
|
||||
let (_, next_s1) = outputs
|
||||
.get(DEC_OUT_STATE1)
|
||||
.context("decoder produced no 'output_states_1'")?
|
||||
.try_extract_tensor::<f32>()
|
||||
.context("failed to extract decoder state1")?;
|
||||
let (_, next_s2) = outputs
|
||||
.get(DEC_OUT_STATE2)
|
||||
.context("decoder produced no 'output_states_2'")?
|
||||
.try_extract_tensor::<f32>()
|
||||
.context("failed to extract decoder state2")?;
|
||||
|
||||
(token, step, next_s1.to_vec(), next_s2.to_vec())
|
||||
};
|
||||
|
||||
// Emit the token unless it is the blank symbol.
|
||||
if token != self.blank_idx {
|
||||
state1 = new_state1;
|
||||
state2 = new_state2;
|
||||
tokens.push(token);
|
||||
prev_token = token;
|
||||
emitted += 1;
|
||||
}
|
||||
|
||||
// Advance: a positive duration jumps several frames; otherwise we
|
||||
// only move on when a blank was emitted or we hit the per-step cap.
|
||||
if step > 0 {
|
||||
t += step;
|
||||
emitted = 0;
|
||||
} else if token == self.blank_idx || emitted == self.max_tokens_per_step {
|
||||
t += 1;
|
||||
emitted = 0;
|
||||
}
|
||||
}
|
||||
|
||||
let text = decode_text(&self.vocab, &tokens);
|
||||
Ok((tokens, text))
|
||||
}
|
||||
}
|
||||
|
||||
/// Build a session with full graph optimization.
|
||||
fn build_session(environment: &Environment, path: &Path) -> Result<Session> {
|
||||
let builder = Session::builder(environment)
|
||||
.map_err(|e| anyhow::anyhow!("failed to create ONNX session builder: {e}"))?;
|
||||
builder
|
||||
.with_optimization_level(GraphOptimizationLevel::Level3)
|
||||
.map_err(|e| anyhow::anyhow!("failed to set ONNX optimization level: {e}"))?
|
||||
.commit_from_file(path)
|
||||
.with_context(|| format!("failed to load ONNX model {}", path.display()))
|
||||
}
|
||||
|
||||
/// Read the size of `axis` from the named output of `session`.
|
||||
fn output_dim(session: &Session, name: &str, axis: usize) -> Result<usize> {
|
||||
let outlet = session
|
||||
.outputs()
|
||||
.iter()
|
||||
.find(|o| o.name() == name)
|
||||
.with_context(|| format!("model has no output named '{name}'"))?;
|
||||
match outlet.dtype() {
|
||||
ValueType::Tensor { shape, .. } => {
|
||||
let dim = *shape
|
||||
.get(axis)
|
||||
.with_context(|| format!("output '{name}' has no axis {axis}"))?;
|
||||
usize::try_from(dim).with_context(|| format!("output '{name}' axis {axis} is dynamic"))
|
||||
}
|
||||
_ => anyhow::bail!("output '{name}' is not a tensor"),
|
||||
}
|
||||
}
|
||||
|
||||
/// Read `(layers, hidden)` from a decoder state input shaped `[layers, batch, hidden]`.
|
||||
fn input_state_dims(session: &Session, name: &str) -> Result<(usize, usize)> {
|
||||
let outlet = session
|
||||
.inputs()
|
||||
.iter()
|
||||
.find(|i| i.name() == name)
|
||||
.with_context(|| format!("decoder has no input named '{name}'"))?;
|
||||
match outlet.dtype() {
|
||||
ValueType::Tensor { shape, .. } => {
|
||||
let layers = usize::try_from(*shape.first().context("state input has no axes")?)
|
||||
.context("state input layer count is dynamic")?;
|
||||
let hidden = usize::try_from(*shape.get(2).context("state input has < 3 axes")?)
|
||||
.context("state input hidden size is dynamic")?;
|
||||
Ok((layers, hidden))
|
||||
}
|
||||
_ => anyhow::bail!("input '{name}' is not a tensor"),
|
||||
}
|
||||
}
|
||||
|
||||
/// Index of the maximum value (ties resolved to the lowest index).
|
||||
fn argmax(values: &[f32]) -> usize {
|
||||
let mut best = 0usize;
|
||||
let mut best_val = f32::NEG_INFINITY;
|
||||
for (i, &v) in values.iter().enumerate() {
|
||||
if v > best_val {
|
||||
best_val = v;
|
||||
best = i;
|
||||
}
|
||||
}
|
||||
best
|
||||
}
|
||||
|
||||
/// Load `vocab.txt` (lines of `<piece> <id>`) into an id-indexed table, mapping
|
||||
/// the SentencePiece word marker `▁` (U+2581) to a real space.
|
||||
fn load_vocab(path: &Path) -> Result<Vec<String>> {
|
||||
let text = std::fs::read_to_string(path)
|
||||
.with_context(|| format!("failed to read {}", path.display()))?;
|
||||
|
||||
let mut entries: Vec<(usize, String)> = Vec::new();
|
||||
let mut max_id = 0usize;
|
||||
for line in text.lines() {
|
||||
if line.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let (piece, id) = line
|
||||
.rsplit_once(' ')
|
||||
.with_context(|| format!("malformed vocab line: {line:?}"))?;
|
||||
let id: usize = id
|
||||
.trim()
|
||||
.parse()
|
||||
.with_context(|| format!("invalid token id in vocab line: {line:?}"))?;
|
||||
max_id = max_id.max(id);
|
||||
entries.push((id, piece.replace('\u{2581}', " ")));
|
||||
}
|
||||
|
||||
let mut vocab = vec![String::new(); max_id + 1];
|
||||
for (id, piece) in entries {
|
||||
vocab[id] = piece;
|
||||
}
|
||||
Ok(vocab)
|
||||
}
|
||||
|
||||
/// Decode token IDs to text and normalize SentencePiece spacing.
|
||||
fn decode_text(vocab: &[String], ids: &[i64]) -> String {
|
||||
let mut joined = String::new();
|
||||
for &id in ids {
|
||||
if id >= 0 {
|
||||
if let Some(piece) = vocab.get(id as usize) {
|
||||
joined.push_str(piece);
|
||||
}
|
||||
}
|
||||
}
|
||||
normalize_spaces(&joined)
|
||||
}
|
||||
|
||||
/// Collapse SentencePiece spacing: drop a leading space, keep a space only when
|
||||
/// it precedes a word character, and drop spacing before punctuation. This
|
||||
/// mirrors the `\A\s|\s\B|(\s)\b` substitution used by the reference decoder.
|
||||
fn normalize_spaces(text: &str) -> String {
|
||||
let chars: Vec<char> = text.chars().collect();
|
||||
let mut out = String::with_capacity(text.len());
|
||||
for (i, &c) in chars.iter().enumerate() {
|
||||
if !c.is_whitespace() {
|
||||
out.push(c);
|
||||
continue;
|
||||
}
|
||||
if i == 0 {
|
||||
continue; // leading whitespace
|
||||
}
|
||||
let next_is_word = chars
|
||||
.get(i + 1)
|
||||
.map(|&n| n.is_alphanumeric() || n == '_')
|
||||
.unwrap_or(false);
|
||||
if next_is_word {
|
||||
out.push(' ');
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Input of the STT worker channel.
|
||||
enum SttInput {
|
||||
/// A streaming mic frame: the worker runs its own RMS voice detection.
|
||||
Frame(Vec<i16>),
|
||||
/// A complete utterance already segmented by an external VAD (see
|
||||
/// [`crate::vad`]), which bypasses the internal RMS detection.
|
||||
Utterance(Vec<i16>),
|
||||
}
|
||||
|
||||
/// Handle to the running STT worker.
|
||||
#[derive(Clone)]
|
||||
pub struct SttWorker {
|
||||
tx: mpsc::UnboundedSender<SttInput>,
|
||||
}
|
||||
|
||||
impl SttWorker {
|
||||
/// Spawn the STT worker task. `stt` is the loaded Parakeet pipeline.
|
||||
/// `transcript_tx` receives one [`Transcript`] per detected EOU.
|
||||
pub fn spawn(
|
||||
stt: ParakeetStt,
|
||||
transcript_tx: mpsc::UnboundedSender<Transcript>,
|
||||
) -> Self {
|
||||
let (tx, mut rx) = mpsc::unbounded_channel::<SttInput>();
|
||||
|
||||
tokio::spawn(async move {
|
||||
let mut utterance: Vec<f32> = Vec::new();
|
||||
let mut in_utterance = false;
|
||||
let mut silence_samples: u64 = 0;
|
||||
let eou_samples = (STT_SAMPLE_RATE as u64 * EOU_SILENCE_MS) / 1000;
|
||||
|
||||
while let Some(input) = rx.recv().await {
|
||||
// Frames arrive already resampled to 16 kHz by the session layer.
|
||||
// Externally segmented utterances need no further detection.
|
||||
let frame = match input {
|
||||
SttInput::Utterance(frame) => {
|
||||
let audio: Vec<f32> =
|
||||
frame.iter().map(|&s| s as f32 / 32768.0).collect();
|
||||
debug!(
|
||||
samples = audio.len(),
|
||||
"VAD-segmented utterance received"
|
||||
);
|
||||
let stt = stt.clone();
|
||||
let tx = transcript_tx.clone();
|
||||
tokio::task::spawn_blocking(move || {
|
||||
run_inference_blocking(&stt, &audio, tx);
|
||||
});
|
||||
continue;
|
||||
}
|
||||
SttInput::Frame(frame) => frame,
|
||||
};
|
||||
|
||||
// Convert i16 -> f32 normalized.
|
||||
let normalized: Vec<f32> =
|
||||
frame.iter().map(|&s| s as f32 / 32768.0).collect();
|
||||
let rms = compute_rms(&normalized);
|
||||
|
||||
if rms > VAD_RMS_THRESHOLD {
|
||||
in_utterance = true;
|
||||
silence_samples = 0;
|
||||
utterance.extend_from_slice(&normalized);
|
||||
} else if in_utterance {
|
||||
silence_samples += normalized.len() as u64;
|
||||
utterance.extend_from_slice(&normalized);
|
||||
if silence_samples >= eou_samples {
|
||||
// EOU detected: run inference on a blocking thread.
|
||||
let utterance_clone = std::mem::take(&mut utterance);
|
||||
let stt = stt.clone();
|
||||
let tx = transcript_tx.clone();
|
||||
|
||||
tokio::task::spawn_blocking(move || {
|
||||
run_inference_blocking(&stt, &utterance_clone, tx);
|
||||
});
|
||||
|
||||
in_utterance = false;
|
||||
silence_samples = 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
info!("STT worker channel closed, stopping");
|
||||
});
|
||||
|
||||
Self { tx }
|
||||
}
|
||||
|
||||
/// Push a frame of 16 kHz mono i16 PCM samples into the worker, which will
|
||||
/// segment it with the built-in RMS voice detection.
|
||||
pub fn push_frame(&self, samples: Vec<i16>) {
|
||||
let _ = self.tx.send(SttInput::Frame(samples));
|
||||
}
|
||||
|
||||
/// Push a complete 16 kHz mono i16 utterance segmented by the external VAD.
|
||||
/// Transcribed as-is, without any further voice-activity detection.
|
||||
pub fn push_utterance(&self, samples: Vec<i16>) {
|
||||
if samples.is_empty() {
|
||||
return;
|
||||
}
|
||||
let _ = self.tx.send(SttInput::Utterance(samples));
|
||||
}
|
||||
}
|
||||
|
||||
/// Blocking inference helper used inside `spawn_blocking`.
|
||||
fn run_inference_blocking(
|
||||
stt: &ParakeetStt,
|
||||
audio: &[f32],
|
||||
tx: mpsc::UnboundedSender<Transcript>,
|
||||
) {
|
||||
let (token_ids, text) = match stt.infer(audio) {
|
||||
Ok(result) => result,
|
||||
Err(e) => {
|
||||
error!(error = %e, "Parakeet inference failed");
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let duration_ms = (audio.len() as u64 * 1000) / STT_SAMPLE_RATE as u64;
|
||||
debug!(tokens = token_ids.len(), duration_ms, text = %text, "EOU transcript ready");
|
||||
|
||||
let _ = tx.send(Transcript {
|
||||
token_ids,
|
||||
text,
|
||||
duration_ms,
|
||||
});
|
||||
}
|
||||
|
||||
/// Compute the Root-Mean-Square energy of a normalized f32 frame.
|
||||
fn compute_rms(samples: &[f32]) -> f32 {
|
||||
if samples.is_empty() {
|
||||
return 0.0;
|
||||
}
|
||||
let sum_sq: f32 = samples.iter().map(|&s| s * s).sum();
|
||||
(sum_sq / samples.len() as f32).sqrt()
|
||||
}
|
||||
|
||||
/// Convert a duration in milliseconds to a sample count at the STT rate.
|
||||
#[allow(dead_code)]
|
||||
pub fn ms_to_samples(ms: u64) -> usize {
|
||||
((STT_SAMPLE_RATE as u64 * ms) / 1000) as usize
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_normalize_spaces() {
|
||||
assert_eq!(normalize_spaces(" Hello world"), "Hello world");
|
||||
assert_eq!(normalize_spaces(" Hello world!"), "Hello world!");
|
||||
assert_eq!(normalize_spaces("a b c"), "a b c");
|
||||
assert_eq!(normalize_spaces(" trailing "), "trailing");
|
||||
assert_eq!(normalize_spaces("a , b"), "a, b");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_load_real_vocab() {
|
||||
let path = Path::new("models/parakeet/vocab.txt");
|
||||
if !path.exists() {
|
||||
eprintln!("{} not found, skipping", path.display());
|
||||
return;
|
||||
}
|
||||
let vocab = load_vocab(path).expect("vocab load failed");
|
||||
assert_eq!(vocab.len(), 8193);
|
||||
assert_eq!(vocab[8192], "<blk>");
|
||||
assert_eq!(vocab[0], "<unk>");
|
||||
}
|
||||
|
||||
/// Read a canonical 16-bit PCM mono WAV file into normalized f32 samples.
|
||||
fn read_wav_mono16(path: &Path) -> Vec<f32> {
|
||||
let bytes = std::fs::read(path).expect("failed to read wav");
|
||||
assert_eq!(&bytes[0..4], b"RIFF", "not a RIFF file");
|
||||
assert_eq!(&bytes[8..12], b"WAVE", "not a WAVE file");
|
||||
// Walk the chunk list to find the `data` chunk.
|
||||
let mut pos = 12usize;
|
||||
while pos + 8 <= bytes.len() {
|
||||
let id = &bytes[pos..pos + 4];
|
||||
let size = u32::from_le_bytes(bytes[pos + 4..pos + 8].try_into().unwrap()) as usize;
|
||||
let body = pos + 8;
|
||||
if id == b"data" {
|
||||
return bytes[body..body + size]
|
||||
.chunks_exact(2)
|
||||
.map(|c| i16::from_le_bytes([c[0], c[1]]) as f32 / 32768.0)
|
||||
.collect();
|
||||
}
|
||||
pos = body + size + (size & 1); // chunks are word-aligned
|
||||
}
|
||||
panic!("no data chunk found");
|
||||
}
|
||||
|
||||
/// End-to-end decode test. Requires the full model directory. Run with:
|
||||
/// `PARAKEET_TEST_WAV=/path/to/16k_mono.wav cargo test -- --ignored`
|
||||
#[test]
|
||||
#[ignore = "requires the full Parakeet ONNX model directory + a 16 kHz WAV"]
|
||||
fn test_infer_real_audio() {
|
||||
let model_dir = std::env::var("PARAKEET_MODEL_DIR")
|
||||
.unwrap_or_else(|_| "models/parakeet".to_string());
|
||||
let wav = std::env::var("PARAKEET_TEST_WAV")
|
||||
.unwrap_or_else(|_| "/tmp/stt_test/solis16k.wav".to_string());
|
||||
|
||||
if !Path::new(&model_dir).join(PREPROCESSOR_FILE).exists() {
|
||||
eprintln!("model dir {model_dir} not found, skipping");
|
||||
return;
|
||||
}
|
||||
let stt = ParakeetStt::load(Path::new(&model_dir)).expect("model load failed");
|
||||
let audio = read_wav_mono16(Path::new(&wav));
|
||||
eprintln!("running inference on {} samples", audio.len());
|
||||
|
||||
let (tokens, text) = stt.infer(&audio).expect("inference failed");
|
||||
eprintln!("tokens: {} | text: {text:?}", tokens.len());
|
||||
assert!(!tokens.is_empty(), "no tokens emitted");
|
||||
assert!(!text.trim().is_empty(), "empty transcript");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,333 @@
|
||||
//! Tokenizer decoder per Parakeet TDT (SentencePiece BPE).
|
||||
//!
|
||||
//! Supporta due formati di input:
|
||||
//! - `tokenizer.model` — protobuf SentencePiece (formato nativo di Parakeet)
|
||||
//! - `tokenizer.json` — HuggingFace BPE (`model.vocab` e `added_tokens`)
|
||||
//!
|
||||
//! La decodifica mappa ogni ID nella sua stringa "piece", sostituisce il
|
||||
//! marcatore di spazio SentencePiece `▁` (U+2581) con uno spazio vero,
|
||||
//! rimuove i token speciali (`<s>`, `</s>`, `<unk>`, `<pad>`) e pulisce gli
|
||||
//! spazi iniziali.
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
|
||||
/// Decoder tokenizer id -> testo.
|
||||
#[derive(Clone)]
|
||||
pub struct Tokenizer {
|
||||
/// Mappa id -> stringa piece.
|
||||
id_to_piece: HashMap<u32, String>,
|
||||
}
|
||||
|
||||
/// Token speciali da scartare in fase di decodifica.
|
||||
const SPECIAL_TOKENS: &[&str] = &["<s>", "</s>", "<unk>", "<pad>", "<blk>", "<mask>"];
|
||||
|
||||
impl Tokenizer {
|
||||
/// Carica il tokenizer da un file, scegliendo il parser in base
|
||||
/// all'estensione (`.model` -> SentencePiece, `.json` -> HuggingFace).
|
||||
pub fn load(path: &Path) -> Result<Self> {
|
||||
let ext = path
|
||||
.extension()
|
||||
.and_then(|e| e.to_str())
|
||||
.unwrap_or("")
|
||||
.to_lowercase();
|
||||
match ext.as_str() {
|
||||
"model" => Self::load_sentencepiece(path),
|
||||
"json" => Self::load_huggingface_json(path),
|
||||
other => anyhow::bail!(
|
||||
"formato tokenizer non supportato: '{}' (usa .model o .json)",
|
||||
other
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
/// Decodifica una sequenza di token ID in testo.
|
||||
pub fn decode(&self, ids: &[i64]) -> String {
|
||||
let mut pieces: Vec<String> = Vec::with_capacity(ids.len());
|
||||
for &id in ids {
|
||||
if id < 0 {
|
||||
continue;
|
||||
}
|
||||
if let Some(piece) = self.id_to_piece.get(&(id as u32)) {
|
||||
pieces.push(piece.clone());
|
||||
}
|
||||
}
|
||||
pieces_to_text(&pieces)
|
||||
}
|
||||
|
||||
/// Numero di pezzi nel vocabolario.
|
||||
pub fn vocab_size(&self) -> usize {
|
||||
self.id_to_piece.len()
|
||||
}
|
||||
}
|
||||
|
||||
/// Converte una lista di piece BPE in testo leggibile.
|
||||
fn pieces_to_text(pieces: &[String]) -> String {
|
||||
// SentencePiece usa '▁' (U+2581) come marcatore di inizio parola/spazio.
|
||||
// Uniamo i pezzi e sostituiamo '▁' con spazio, poi scartiamo i token
|
||||
// speciali e gli spazi ridondanti.
|
||||
let joined: String = pieces.join("");
|
||||
|
||||
// Rimuovi token speciali: sia quelli noti (<s>, </s>, ...) sia qualsiasi
|
||||
// token nel formato <...> o <|...|> usato da Parakeet/NVIDIA.
|
||||
let mut cleaned = joined;
|
||||
for tok in SPECIAL_TOKENS {
|
||||
cleaned = cleaned.replace(tok, "");
|
||||
}
|
||||
// Rimuovi token speciali generici nel formato <token> o <|token|>.
|
||||
cleaned = remove_angle_bracket_tokens(&cleaned);
|
||||
|
||||
// Sostituisci il marcatore SentencePiece con spazio.
|
||||
let mut text = cleaned.replace('▁', " ");
|
||||
|
||||
// Collassa spazi multipli e togli spazi iniziali.
|
||||
text = text.split_whitespace().collect::<Vec<_>>().join(" ");
|
||||
text
|
||||
}
|
||||
|
||||
/// Rimuove token speciali nel formato `<...>` o `<|...|>` dalla stringa.
|
||||
fn remove_angle_bracket_tokens(s: &str) -> String {
|
||||
let mut result = String::with_capacity(s.len());
|
||||
let mut chars = s.chars().peekable();
|
||||
while let Some(c) = chars.next() {
|
||||
if c == '<' {
|
||||
let mut depth = 1i32;
|
||||
let mut inner = String::new();
|
||||
inner.push(c);
|
||||
let mut found_close = false;
|
||||
while let Some(&nc) = chars.peek() {
|
||||
inner.push(nc);
|
||||
chars.next();
|
||||
match nc {
|
||||
'>' => {
|
||||
depth -= 1;
|
||||
if depth == 0 {
|
||||
found_close = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
'<' => depth += 1,
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
if found_close {
|
||||
// Token speciale: scartalo.
|
||||
continue;
|
||||
}
|
||||
// Nessuna chiusura: ripristina il testo originale.
|
||||
result.push_str(&inner);
|
||||
} else {
|
||||
result.push(c);
|
||||
}
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Parser SentencePiece (raw protobuf wire format)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
impl Tokenizer {
|
||||
/// Parsa un file `tokenizer.model` di SentencePiece.
|
||||
///
|
||||
/// Il proto `SentencePieceModel` ha come campo 1 un repeated `SentencePiece`
|
||||
/// che a sua volta ha come campo 1 la stringa `piece`. Leggiamo solo questi
|
||||
/// due campi, ignorando tutto il resto.
|
||||
fn load_sentencepiece(path: &Path) -> Result<Self> {
|
||||
let bytes = std::fs::read(path)
|
||||
.with_context(|| format!("impossibile leggere {}", path.display()))?;
|
||||
let pieces = parse_sentencepiece_pieces(&bytes)
|
||||
.with_context(|| format!("parse fallito di {}", path.display()))?;
|
||||
|
||||
let mut id_to_piece = HashMap::with_capacity(pieces.len());
|
||||
for (id, piece) in pieces.into_iter().enumerate() {
|
||||
id_to_piece.insert(id as u32, piece);
|
||||
}
|
||||
Ok(Self { id_to_piece })
|
||||
}
|
||||
|
||||
/// Parsa un file `tokenizer.json` di HuggingFace (formato BPE).
|
||||
fn load_huggingface_json(path: &Path) -> Result<Self> {
|
||||
let text = std::fs::read_to_string(path)
|
||||
.with_context(|| format!("impossibile leggere {}", path.display()))?;
|
||||
let json: serde_json::Value = serde_json::from_str(&text)
|
||||
.with_context(|| format!("JSON non valido in {}", path.display()))?;
|
||||
|
||||
let vocab = json
|
||||
.get("model")
|
||||
.and_then(|m| m.get("vocab"))
|
||||
.and_then(|v| v.as_object())
|
||||
.context("tokenizer.json manca di model.vocab")?;
|
||||
|
||||
let mut id_to_piece = HashMap::with_capacity(vocab.len());
|
||||
for (piece, id_val) in vocab {
|
||||
let id = id_val
|
||||
.as_u64()
|
||||
.context("id vocabolario non numerico")? as u32;
|
||||
id_to_piece.insert(id, piece.clone());
|
||||
}
|
||||
Ok(Self { id_to_piece })
|
||||
}
|
||||
}
|
||||
|
||||
/// Estrae la lista dei `piece` (in ordine di ID) dal protobuf SentencePiece.
|
||||
fn parse_sentencepiece_pieces(data: &[u8]) -> Result<Vec<String>> {
|
||||
let mut pieces = Vec::new();
|
||||
let mut pos = 0;
|
||||
while pos < data.len() {
|
||||
let (field, wire, next) = read_key(data, pos)?;
|
||||
pos = next;
|
||||
if field == 1 && wire == 2 {
|
||||
// Field 1 = repeated SentencePiece (length-delimited message).
|
||||
let (len, next) = read_varint(data, pos)?;
|
||||
pos = next;
|
||||
let end = pos + len as usize;
|
||||
if end > data.len() {
|
||||
anyhow::bail!("truncated SentencePiece message");
|
||||
}
|
||||
let msg = &data[pos..end];
|
||||
if let Some(piece) = parse_piece_message(msg) {
|
||||
pieces.push(piece);
|
||||
}
|
||||
pos = end;
|
||||
} else {
|
||||
// Salta il campo in base al wire type.
|
||||
pos = skip_field(data, pos, wire)?;
|
||||
}
|
||||
}
|
||||
Ok(pieces)
|
||||
}
|
||||
|
||||
/// Parsa un messaggio `SentencePiece` e restituisce il campo `piece` (field 1).
|
||||
fn parse_piece_message(data: &[u8]) -> Option<String> {
|
||||
let mut pos = 0;
|
||||
while pos < data.len() {
|
||||
let (field, wire, next) = read_key(data, pos).ok()?;
|
||||
pos = next;
|
||||
if field == 1 && wire == 2 {
|
||||
let (len, next) = read_varint(data, pos).ok()?;
|
||||
pos = next;
|
||||
let end = pos + len as usize;
|
||||
if end > data.len() {
|
||||
return None;
|
||||
}
|
||||
return Some(String::from_utf8_lossy(&data[pos..end]).into_owned());
|
||||
}
|
||||
pos = skip_field(data, pos, wire).ok()?;
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
// --- Funzioni helper protobuf wire-format ---
|
||||
|
||||
/// Legge una chiave protobuf: restituisce (field_number, wire_type, nuova_pos).
|
||||
fn read_key(data: &[u8], pos: usize) -> Result<(u32, u8, usize)> {
|
||||
let (key, next) = read_varint(data, pos)?;
|
||||
let wire = (key & 0x7) as u8;
|
||||
let field = (key >> 3) as u32;
|
||||
Ok((field, wire, next))
|
||||
}
|
||||
|
||||
/// Legge un varint protobuf (base-128).
|
||||
fn read_varint(data: &[u8], pos: usize) -> Result<(u64, usize)> {
|
||||
let mut result: u64 = 0;
|
||||
let mut shift: u32 = 0;
|
||||
let mut p = pos;
|
||||
loop {
|
||||
if p >= data.len() {
|
||||
anyhow::bail!("truncated varint");
|
||||
}
|
||||
let byte = data[p];
|
||||
p += 1;
|
||||
result |= ((byte & 0x7F) as u64) << shift;
|
||||
if byte & 0x80 == 0 {
|
||||
break;
|
||||
}
|
||||
shift += 7;
|
||||
if shift >= 64 {
|
||||
anyhow::bail!("varint troppo lungo");
|
||||
}
|
||||
}
|
||||
Ok((result, p))
|
||||
}
|
||||
|
||||
/// Salta un campo protobuf in base al wire type.
|
||||
fn skip_field(data: &[u8], pos: usize, wire: u8) -> Result<usize> {
|
||||
match wire {
|
||||
0 => {
|
||||
// varint
|
||||
let (_, next) = read_varint(data, pos)?;
|
||||
Ok(next)
|
||||
}
|
||||
1 => {
|
||||
// 64-bit
|
||||
if pos + 8 > data.len() {
|
||||
anyhow::bail!("truncated 64-bit field");
|
||||
}
|
||||
Ok(pos + 8)
|
||||
}
|
||||
2 => {
|
||||
// length-delimited
|
||||
let (len, next) = read_varint(data, pos)?;
|
||||
Ok(next + len as usize)
|
||||
}
|
||||
5 => {
|
||||
// 32-bit
|
||||
if pos + 4 > data.len() {
|
||||
anyhow::bail!("truncated 32-bit field");
|
||||
}
|
||||
Ok(pos + 4)
|
||||
}
|
||||
_ => anyhow::bail!("wire type non supportato: {}", wire),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_pieces_to_text() {
|
||||
let pieces = vec![
|
||||
"▁Hello".to_string(),
|
||||
"▁world".to_string(),
|
||||
"!".to_string(),
|
||||
];
|
||||
// "!" non ha il marcatore ▁ quindi si attacca a "world".
|
||||
assert_eq!(pieces_to_text(&pieces), "Hello world!");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_pieces_to_text_special_tokens() {
|
||||
let pieces = vec![
|
||||
"<s>".to_string(),
|
||||
"▁hi".to_string(),
|
||||
"</s>".to_string(),
|
||||
];
|
||||
assert_eq!(pieces_to_text(&pieces), "hi");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_remove_angle_bracket_tokens() {
|
||||
assert_eq!(remove_angle_bracket_tokens("hello<|endoftext|>world"), "helloworld");
|
||||
assert_eq!(remove_angle_bracket_tokens("<|nospeech|>abc<unk>"), "abc");
|
||||
assert_eq!(remove_angle_bracket_tokens("no tokens here"), "no tokens here");
|
||||
}
|
||||
|
||||
/// Verifica che il parser SentencePiece legga il tokenizer reale di Parakeet.
|
||||
#[test]
|
||||
fn test_load_real_parakeet_tokenizer() {
|
||||
let path = std::path::Path::new("models/tokenizer.model");
|
||||
if !path.exists() {
|
||||
eprintln!("models/tokenizer.model non trovato, skip test");
|
||||
return;
|
||||
}
|
||||
let tok = Tokenizer::load(path).expect("caricamento tokenizer fallito");
|
||||
// Il vocabolario di Parakeet ha ~1024 pezzi.
|
||||
assert!(tok.vocab_size() > 500, "vocabolario troppo piccolo: {}", tok.vocab_size());
|
||||
// I primi pezzi sono tipicamente token speciali.
|
||||
let decoded = tok.decode(&[1, 2, 3]);
|
||||
eprintln!("decoded ids 1,2,3 = {:?}", decoded);
|
||||
}
|
||||
}
|
||||
+97
@@ -0,0 +1,97 @@
|
||||
//! Optional TLS termination (HTTPS / WSS) for the Axum server.
|
||||
//!
|
||||
//! Enabled by setting both `TLS_CERT` and `TLS_KEY` (PEM files) in the
|
||||
//! environment / `.env`. It serves the exact same [`Router`] as the plain-HTTP
|
||||
//! path, but over `tokio-rustls`, so the browser sees a *secure context* —
|
||||
//! which is what `getUserMedia()` (the microphone) requires when the page is
|
||||
//! not on `localhost`.
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use axum::body::Body;
|
||||
use axum::Router;
|
||||
use hyper::server::conn::http1;
|
||||
use hyper::service::service_fn;
|
||||
use hyper_util::rt::TokioIo;
|
||||
use rustls::ServerConfig;
|
||||
use rustls_pki_types::pem::PemObject;
|
||||
use rustls_pki_types::{CertificateDer, PrivateKeyDer};
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
use tokio::net::TcpListener;
|
||||
use tokio_rustls::TlsAcceptor;
|
||||
use tower::Service;
|
||||
use tracing::{debug, info};
|
||||
|
||||
/// Load a PEM certificate chain + private key and build a rustls server config.
|
||||
pub fn load_server_config(cert_path: &Path, key_path: &Path) -> Result<Arc<ServerConfig>> {
|
||||
let certs = CertificateDer::pem_file_iter(cert_path)
|
||||
.with_context(|| format!("opening certificate {}", cert_path.display()))?
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.with_context(|| format!("parsing certificate {}", cert_path.display()))?;
|
||||
if certs.is_empty() {
|
||||
anyhow::bail!("no certificates found in {}", cert_path.display());
|
||||
}
|
||||
let key = PrivateKeyDer::from_pem_file(key_path)
|
||||
.with_context(|| format!("parsing private key {}", key_path.display()))?;
|
||||
|
||||
let config = ServerConfig::builder()
|
||||
.with_no_client_auth()
|
||||
.with_single_cert(certs, key)
|
||||
.context("building rustls server config")?;
|
||||
Ok(Arc::new(config))
|
||||
}
|
||||
|
||||
/// Accept connections on `listener` and serve `app` over TLS until `shutdown`
|
||||
/// resolves. Each connection is handled on its own task; HTTP/1 upgrades are
|
||||
/// enabled so WebSocket handshakes work.
|
||||
pub async fn serve(
|
||||
listener: TcpListener,
|
||||
app: Router,
|
||||
tls_config: Arc<ServerConfig>,
|
||||
shutdown: impl std::future::Future<Output = ()>,
|
||||
) -> Result<()> {
|
||||
let acceptor = TlsAcceptor::from(tls_config);
|
||||
tokio::pin!(shutdown);
|
||||
|
||||
loop {
|
||||
let (stream, peer) = tokio::select! {
|
||||
_ = &mut shutdown => break,
|
||||
accepted = listener.accept() => accepted?,
|
||||
};
|
||||
|
||||
let acceptor = acceptor.clone();
|
||||
let app = app.clone();
|
||||
tokio::spawn(async move {
|
||||
let stream = match acceptor.accept(stream).await {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
debug!(%peer, error = %e, "TLS handshake failed");
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let io = TokioIo::new(stream);
|
||||
// Bridge hyper's request type (with `Incoming` body) to the Axum
|
||||
// router, which expects `axum::body::Body`.
|
||||
let service = service_fn(move |req: hyper::Request<hyper::body::Incoming>| {
|
||||
let mut router = app.clone();
|
||||
async move {
|
||||
let (parts, body) = req.into_parts();
|
||||
let req = hyper::Request::from_parts(parts, Body::new(body));
|
||||
router.call(req).await
|
||||
}
|
||||
});
|
||||
|
||||
if let Err(e) = http1::Builder::new()
|
||||
.serve_connection(io, service)
|
||||
.with_upgrades()
|
||||
.await
|
||||
{
|
||||
debug!(%peer, error = %e, "connection error");
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
info!("TLS listener stopped");
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,273 @@
|
||||
//! Fish Audio / SGLang Drama 3 TTS plugin.
|
||||
//!
|
||||
//! POSTs text to `/v1/audio/speech` (OpenAI-compatible) requesting PCM
|
||||
//! streaming. Fish Audio Drama 3 outputs at 44.1 kHz, which matches the
|
||||
//! mixer's output rate, so samples are pushed directly into the TTS queue
|
||||
//! without resampling. A [`CancellationToken`] allows immediate interruption
|
||||
//! (barge-in).
|
||||
|
||||
use crate::audio::SampleRingBuffer;
|
||||
use crate::config::Config;
|
||||
use crate::tts::TtsEngine;
|
||||
use anyhow::{Context, Result};
|
||||
use async_trait::async_trait;
|
||||
use futures_util::StreamExt;
|
||||
use reqwest::Client;
|
||||
use serde::Serialize;
|
||||
use std::sync::Arc;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::{debug, error, info};
|
||||
|
||||
/// A voice-cloning reference: path to a WAV file + its transcript.
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
struct Reference {
|
||||
audio_path: String,
|
||||
text: String,
|
||||
}
|
||||
|
||||
/// Request body for the SGLang `/v1/audio/speech` endpoint.
|
||||
#[derive(Debug, Serialize)]
|
||||
struct SpeechRequest<'a> {
|
||||
model: &'a str,
|
||||
input: &'a str,
|
||||
voice: &'a str,
|
||||
response_format: &'a str,
|
||||
stream: bool,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
references: Option<Vec<Reference>>,
|
||||
}
|
||||
|
||||
/// Streaming TTS client backed by Fish Audio / SGLang-Omni Drama 3.
|
||||
#[derive(Clone)]
|
||||
pub struct FishAudioTts {
|
||||
http: Client,
|
||||
endpoint: String,
|
||||
model: String,
|
||||
voice: String,
|
||||
/// Current voice-cloning reference (config default, overridable at runtime).
|
||||
reference: Arc<parking_lot::RwLock<Option<Reference>>>,
|
||||
}
|
||||
|
||||
impl FishAudioTts {
|
||||
pub fn new(config: &Config) -> Self {
|
||||
let reference = config.tts_reference().map(|(audio_path, text)| Reference {
|
||||
audio_path,
|
||||
text,
|
||||
});
|
||||
if let Some(r) = &reference {
|
||||
info!(
|
||||
audio_path = %r.audio_path,
|
||||
text_len = r.text.len(),
|
||||
"TTS voice-cloning reference configured"
|
||||
);
|
||||
}
|
||||
Self {
|
||||
http: Client::builder()
|
||||
.build()
|
||||
.expect("failed to build reqwest client"),
|
||||
endpoint: config.speech_endpoint(),
|
||||
model: config.tts_model.clone(),
|
||||
voice: config.tts_voice.clone(),
|
||||
reference: Arc::new(parking_lot::RwLock::new(reference)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Override the voice-cloning reference at runtime. Pass `None` to clear.
|
||||
pub fn set_reference(&self, audio_path: Option<String>, text: Option<String>) {
|
||||
let mut guard = self.reference.write();
|
||||
match (audio_path, text) {
|
||||
(Some(audio_path), text) => {
|
||||
let text = text.unwrap_or_default();
|
||||
info!(audio_path = %audio_path, text_len = text.len(), "TTS reference updated");
|
||||
*guard = Some(Reference { audio_path, text });
|
||||
}
|
||||
_ => {
|
||||
info!("TTS reference cleared");
|
||||
*guard = None;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl TtsEngine for FishAudioTts {
|
||||
fn name(&self) -> &str {
|
||||
"fish-audio"
|
||||
}
|
||||
|
||||
fn sample_rate(&self) -> u32 {
|
||||
// Fish Audio Drama 3 outputs natively at 44.1 kHz, matching the mixer.
|
||||
44_100
|
||||
}
|
||||
|
||||
async fn synthesize(
|
||||
&self,
|
||||
text: &str,
|
||||
tts_queue: &SampleRingBuffer,
|
||||
cancel: &CancellationToken,
|
||||
) -> Result<()> {
|
||||
let references = self.reference.read().clone().map(|r| vec![r]);
|
||||
|
||||
let body = SpeechRequest {
|
||||
model: &self.model,
|
||||
input: text,
|
||||
voice: &self.voice,
|
||||
response_format: "pcm",
|
||||
stream: true,
|
||||
references,
|
||||
};
|
||||
|
||||
info!(model = %self.model, text_len = text.len(), "SGLang TTS request");
|
||||
debug!(
|
||||
endpoint = %self.endpoint,
|
||||
model = %self.model,
|
||||
voice = %self.voice,
|
||||
response_format = "pcm",
|
||||
stream = true,
|
||||
has_reference = body.references.is_some(),
|
||||
text = %text,
|
||||
"TTS request dispatched"
|
||||
);
|
||||
|
||||
let req = self
|
||||
.http
|
||||
.post(&self.endpoint)
|
||||
.json(&body)
|
||||
.build()
|
||||
.context("failed to build SGLang request")?;
|
||||
|
||||
let response = self
|
||||
.http
|
||||
.execute(req)
|
||||
.await
|
||||
.context("SGLang request failed")?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
let status = response.status();
|
||||
let err_text = response.text().await.unwrap_or_default();
|
||||
error!(%status, body = %err_text, "SGLang returned an error");
|
||||
anyhow::bail!("SGLang error: {} {}", status, err_text);
|
||||
}
|
||||
|
||||
// Read the actual sample rate from the response header for logging.
|
||||
let tts_sr = response
|
||||
.headers()
|
||||
.get("x-sample-rate")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.and_then(|s| s.parse::<u32>().ok())
|
||||
.unwrap_or(44100);
|
||||
info!(tts_sample_rate = tts_sr, "TTS response sample rate");
|
||||
|
||||
let mut byte_stream = response.bytes_stream();
|
||||
let mut carry: Option<u8> = None;
|
||||
let mut chunk_count = 0usize;
|
||||
let mut total_samples = 0usize;
|
||||
let t0 = std::time::Instant::now();
|
||||
let mut last_log = t0;
|
||||
|
||||
while let Some(chunk) = tokio::select! {
|
||||
c = byte_stream.next() => c,
|
||||
_ = cancel.cancelled() => {
|
||||
info!("TTS stream cancelled (barge-in)");
|
||||
return Ok(());
|
||||
}
|
||||
} {
|
||||
let chunk = match chunk {
|
||||
Ok(b) => b,
|
||||
Err(e) => {
|
||||
error!(error = %e, "error reading TTS stream chunk");
|
||||
return Err(e.into());
|
||||
}
|
||||
};
|
||||
let n = ingest_pcm_chunk(&chunk, &mut carry, tts_queue);
|
||||
chunk_count += 1;
|
||||
total_samples += n;
|
||||
if chunk_count == 1 {
|
||||
debug!(
|
||||
ttfb_ms = t0.elapsed().as_millis(),
|
||||
bytes = chunk.len(),
|
||||
"TTS time-to-first-byte"
|
||||
);
|
||||
}
|
||||
debug!(
|
||||
chunk = chunk_count,
|
||||
bytes = chunk.len(),
|
||||
samples = n,
|
||||
queue_len = tts_queue.len(),
|
||||
elapsed_ms = t0.elapsed().as_millis(),
|
||||
"TTS chunk received"
|
||||
);
|
||||
if last_log.elapsed() > std::time::Duration::from_millis(500) {
|
||||
last_log = std::time::Instant::now();
|
||||
info!(
|
||||
chunk = chunk_count,
|
||||
total_samples,
|
||||
queue_len = tts_queue.len(),
|
||||
elapsed_ms = t0.elapsed().as_millis(),
|
||||
"TTS stream progress"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
info!(
|
||||
chunk_count,
|
||||
total_samples,
|
||||
duration_s = format!("{:.2}", t0.elapsed().as_secs_f32()),
|
||||
"TTS stream complete"
|
||||
);
|
||||
debug!(
|
||||
chunk_count,
|
||||
total_samples,
|
||||
avg_chunk_bytes = if chunk_count > 0 { (total_samples * 2) / chunk_count } else { 0 },
|
||||
audio_s = format!("{:.2}", total_samples as f32 / tts_sr as f32),
|
||||
elapsed_ms = t0.elapsed().as_millis(),
|
||||
queue_len = tts_queue.len(),
|
||||
"TTS stream summary"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn set_reference(&self, audio_path: Option<String>, text: Option<String>) {
|
||||
FishAudioTts::set_reference(self, audio_path, text);
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse 16-bit LE PCM samples from a chunk and push them directly to the
|
||||
/// queue (Fish Audio's 44.1 kHz output matches the mixer rate).
|
||||
/// Returns the number of samples pushed.
|
||||
fn ingest_pcm_chunk(
|
||||
chunk: &[u8],
|
||||
carry: &mut Option<u8>,
|
||||
tts_queue: &SampleRingBuffer,
|
||||
) -> usize {
|
||||
let mut samples: Vec<i16> = Vec::with_capacity(chunk.len() / 2 + 1);
|
||||
let mut iter = chunk.iter().copied();
|
||||
|
||||
if let Some(lo) = carry.take() {
|
||||
if let Some(hi) = iter.next() {
|
||||
samples.push(i16::from_le_bytes([lo, hi]));
|
||||
} else {
|
||||
*carry = Some(lo);
|
||||
}
|
||||
}
|
||||
|
||||
loop {
|
||||
let lo = match iter.next() {
|
||||
Some(b) => b,
|
||||
None => break,
|
||||
};
|
||||
let hi = match iter.next() {
|
||||
Some(b) => b,
|
||||
None => {
|
||||
*carry = Some(lo);
|
||||
break;
|
||||
}
|
||||
};
|
||||
samples.push(i16::from_le_bytes([lo, hi]));
|
||||
}
|
||||
|
||||
if !samples.is_empty() {
|
||||
tts_queue.push(&samples);
|
||||
}
|
||||
samples.len()
|
||||
}
|
||||
@@ -0,0 +1,128 @@
|
||||
//! Kokoro TTS plugin (stub).
|
||||
//!
|
||||
//! Kokoro is a lightweight, open-source TTS model. This stub shows the pattern for
|
||||
//! adding a new engine plugin: implement [`TtsEngine`] and wire the HTTP or
|
||||
//! local inference calls.
|
||||
//!
|
||||
//! # Configuration
|
||||
//! Set `TTS_ENGINE=kokoro` and provide:
|
||||
//! - `KOKORO_URL` — base URL of the Kokoro HTTP server (default: `http://localhost:8880`)
|
||||
//! - `KOKORO_VOICE` — voice name (default: `af_heart`)
|
||||
//!
|
||||
//! # Resampling note
|
||||
//! Kokoro typically outputs at 24 kHz. Since the mixer runs at
|
||||
//! [`crate::config::OUTPUT_SAMPLE_RATE`] (44.1 kHz), this plugin must
|
||||
//! resample internally before pushing samples. A basic linear resampler is
|
||||
//! provided below; replace with a polyphase/FIR resampler for higher quality.
|
||||
|
||||
use crate::audio::SampleRingBuffer;
|
||||
use crate::config::{Config, OUTPUT_SAMPLE_RATE};
|
||||
use crate::tts::TtsEngine;
|
||||
use anyhow::Result;
|
||||
use async_trait::async_trait;
|
||||
use std::sync::Arc;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::info;
|
||||
|
||||
/// Kokoro native output sample rate (Hz).
|
||||
const KOKORO_SAMPLE_RATE: u32 = 24_000;
|
||||
|
||||
/// Kokoro TTS plugin.
|
||||
#[derive(Clone)]
|
||||
pub struct KokoroTts {
|
||||
endpoint: String,
|
||||
voice: String,
|
||||
http: reqwest::Client,
|
||||
_config: Arc<Config>,
|
||||
}
|
||||
|
||||
impl KokoroTts {
|
||||
pub fn new(config: &Config) -> Self {
|
||||
let endpoint = std::env::var("KOKORO_URL")
|
||||
.unwrap_or_else(|_| "http://localhost:8880".to_string());
|
||||
let voice = std::env::var("KOKORO_VOICE")
|
||||
.unwrap_or_else(|_| "af_heart".to_string());
|
||||
info!(endpoint = %endpoint, voice = %voice, "Kokoro TTS plugin initialized");
|
||||
Self {
|
||||
endpoint,
|
||||
voice,
|
||||
http: reqwest::Client::builder()
|
||||
.build()
|
||||
.expect("failed to build reqwest client"),
|
||||
_config: Arc::new(config.clone()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl TtsEngine for KokoroTts {
|
||||
fn name(&self) -> &str {
|
||||
"kokoro"
|
||||
}
|
||||
|
||||
fn sample_rate(&self) -> u32 {
|
||||
// We resample to the mixer rate before pushing.
|
||||
OUTPUT_SAMPLE_RATE
|
||||
}
|
||||
|
||||
async fn synthesize(
|
||||
&self,
|
||||
text: &str,
|
||||
tts_queue: &SampleRingBuffer,
|
||||
cancel: &CancellationToken,
|
||||
) -> Result<()> {
|
||||
// TODO: implement the actual Kokoro HTTP streaming call.
|
||||
//
|
||||
// Example request (adjust to be completed by the user):
|
||||
//
|
||||
// ```ignore
|
||||
// let resp = self.http
|
||||
// .post(format!("{}/v1/audio/speech", self.endpoint))
|
||||
// .json(&serde_json::json!({
|
||||
// "model": "kokoro",
|
||||
// "input": text,
|
||||
// "voice": self.voice,
|
||||
// "response_format": "pcm",
|
||||
// "stream": true,
|
||||
// }))
|
||||
// .send()
|
||||
// .await?;
|
||||
//
|
||||
// let mut stream = resp.bytes_stream();
|
||||
// while let Some(chunk) = stream.next().await {
|
||||
// if cancel.is_cancelled() { break; }
|
||||
// let bytes = chunk?;
|
||||
// // Kokoro outputs 16-bit mono PCM at 24 kHz.
|
||||
// let samples = pcm_to_i16(&bytes);
|
||||
// let resampled = resample_24k_to_output(&samples);
|
||||
// tts_queue.push(&resampled);
|
||||
// }
|
||||
// ```
|
||||
let _ = (text, tts_queue, cancel);
|
||||
anyhow::bail!(
|
||||
"Kokoro TTS plugin is not yet implemented — \
|
||||
complete the HTTP streaming call in src/tts/kokoro.rs"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// Resample a 24 kHz mono `i16` buffer to [`OUTPUT_SAMPLE_RATE` using linear
|
||||
/// interpolation. Replace with a windowed-sinc FIR resampler for higher fidelity.
|
||||
#[allow(dead_code)]
|
||||
fn resample_24k_to_output(samples: &[i16]) -> Vec<i16> {
|
||||
if samples.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
let ratio = OUTPUT_SAMPLE_RATE as f64 / KOKORO_SAMPLE_RATE as f64;
|
||||
let out_len = (samples.len() as f64 * ratio).round() as usize;
|
||||
let mut out = Vec::with_capacity(out_len);
|
||||
for i in 0..out_len {
|
||||
let src_pos = i as f64 / ratio;
|
||||
let idx = src_pos.floor() as usize;
|
||||
let frac = src_pos - idx as f64;
|
||||
let a = samples[idx.min(samples.len() - 1)] as f32;
|
||||
let b = samples[(idx + 1).min(samples.len() - 1)] as f32;
|
||||
out.push((a + (b - a) * frac as f32).round() as i16);
|
||||
}
|
||||
out
|
||||
}
|
||||
@@ -0,0 +1,82 @@
|
||||
//! Pluggable Text-to-Speech engine subsystem.
|
||||
//!
|
||||
//! The engine is selected at runtime via the `TTS_ENGINE` environment variable
|
||||
//! (see [`Config::tts_engine`]). Every plugin implements the [`TtsEngine`]
|
||||
//! trait so the orchestration layer is engine-agnostic.
|
||||
//!
|
||||
//! Currently bundled plugins:
|
||||
//! - [`fish_audio::FishAudioTts`] — Fish Audio / SGLang Drama 3 (OpenAI-compatible HTTP).
|
||||
//! - [`kokoro::KokoroTts`] — Kokoro TTS plugin (stub; implement the HTTP/API calls).
|
||||
//!
|
||||
//! To add a new engine, create a module under `src/tts/`, implement
|
||||
//! [`TtsEngine`], and register it in [`create_tts_engine`].
|
||||
|
||||
use crate::audio::SampleRingBuffer;
|
||||
use crate::config::Config;
|
||||
use anyhow::Result;
|
||||
use async_trait::async_trait;
|
||||
use std::sync::Arc;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::info;
|
||||
|
||||
pub mod fish_audio;
|
||||
pub mod kokoro;
|
||||
|
||||
use fish_audio::FishAudioTts;
|
||||
use kokoro::KokoroTts;
|
||||
|
||||
/// Trait implemented by every TTS engine plugin.
|
||||
///
|
||||
/// Implementations must stream **16-bit mono PCM samples at
|
||||
/// [`crate::config::OUTPUT_SAMPLE_RATE`]** into the provided `tts_queue`.
|
||||
/// If the engine's native sample rate differs, it must resample internally
|
||||
/// before pushing samples (the mixer runs at a fixed rate).
|
||||
#[async_trait]
|
||||
pub trait TtsEngine: Send + Sync {
|
||||
/// Human-readable engine name (used for logging).
|
||||
fn name(&self) -> &str;
|
||||
|
||||
/// The sample rate (Hz) at which this engine pushes samples into the
|
||||
/// mixer queue. Must equal [`crate::config::OUTPUT_SAMPLE_RATE`].
|
||||
fn sample_rate(&self) -> u32;
|
||||
|
||||
/// Synthesize `text` and stream PCM samples into `tts_queue`.
|
||||
///
|
||||
/// The implementation **must** honor `cancel` for barge-in: when the token
|
||||
/// is cancelled, stop pushing samples and return `Ok(())` promptly.
|
||||
async fn synthesize(
|
||||
&self,
|
||||
text: &str,
|
||||
tts_queue: &SampleRingBuffer,
|
||||
cancel: &CancellationToken,
|
||||
) -> Result<()>;
|
||||
|
||||
/// Update the voice-cloning reference at runtime. Default: no-op (engines
|
||||
/// that support voice cloning override this).
|
||||
fn set_reference(&self, _audio_path: Option<String>, _text: Option<String>) {}
|
||||
}
|
||||
|
||||
/// Build the TTS engine selected by the `TTS_ENGINE` env var.
|
||||
///
|
||||
/// Supported values (case-insensitive): `fish`, `fishaudio`, `fish_audio`,
|
||||
/// `sglang` → Fish Audio Drama 3; `kokoro` → Kokoro TTS.
|
||||
///
|
||||
/// # Panics
|
||||
/// Panics if `TTS_ENGINE` is set to an unknown value.
|
||||
pub fn create_tts_engine(config: &Config) -> Arc<dyn TtsEngine> {
|
||||
match config.tts_engine.to_lowercase().as_str() {
|
||||
"fish" | "fishaudio" | "fish_audio" | "sglang" => {
|
||||
info!(engine = "fish", "selected TTS engine");
|
||||
Arc::new(FishAudioTts::new(config))
|
||||
}
|
||||
"kokoro" => {
|
||||
info!(engine = "kokoro", "selected TTS engine");
|
||||
Arc::new(KokoroTts::new(config))
|
||||
}
|
||||
other => panic!(
|
||||
"unknown TTS_ENGINE='{}' — supported values: fish, kokoro",
|
||||
other
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
+110
@@ -0,0 +1,110 @@
|
||||
//! Voice Activity Detection (Silero VAD v5 via ONNX Runtime).
|
||||
//!
|
||||
//! The detector runs on the 16 kHz mono stream produced by the session layer
|
||||
//! (the same stream fed to the STT engine) and emits three kinds of events:
|
||||
//!
|
||||
//! - [`VadEvent::SpeechOnset`] — fast onset, used to interrupt the agent
|
||||
//! (barge-in) as soon as the user starts talking.
|
||||
//! - [`VadEvent::SpeechStart`] — confirmed utterance start.
|
||||
//! - [`VadEvent::SpeechEnd`] — utterance finished, carrying the captured audio
|
||||
//! (with a pre-roll pad) that is handed to the STT engine.
|
||||
//!
|
||||
//! Every behavioural knob is configured from the environment (see
|
||||
//! [`VadSettings`]), so the feature can be tuned or disabled without a rebuild.
|
||||
|
||||
pub mod silero;
|
||||
|
||||
pub use silero::{SileroVad, VadDetector, VadEvent};
|
||||
|
||||
use crate::config::{env_bool, env_f32, env_u32};
|
||||
use std::path::PathBuf;
|
||||
|
||||
/// Runtime configuration for the VAD subsystem, read from `.env`.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct VadSettings {
|
||||
/// Master switch (`VAD_ENABLED`). When false the legacy RMS segmentation
|
||||
/// inside the STT worker is used instead.
|
||||
pub enabled: bool,
|
||||
/// Path to `silero_vad.onnx` (`VAD_MODEL_PATH`).
|
||||
pub model_path: PathBuf,
|
||||
/// Speech probability above which a window counts as voice
|
||||
/// (`VAD_THRESHOLD`, default 0.5).
|
||||
pub threshold: f32,
|
||||
/// Probability below which a window counts as silence
|
||||
/// (`VAD_NEG_THRESHOLD`). Must be lower than `threshold`; the gap is the
|
||||
/// hysteresis that prevents rapid onset/offset flapping.
|
||||
pub neg_threshold: f32,
|
||||
/// Speech required to confirm an utterance (`VAD_MIN_SPEECH_MS`).
|
||||
pub min_speech_ms: u32,
|
||||
/// Silence required to close an utterance — the EOU window
|
||||
/// (`VAD_MIN_SILENCE_MS`).
|
||||
pub min_silence_ms: u32,
|
||||
/// Audio kept before the detected onset so the first phoneme is not
|
||||
/// clipped (`VAD_SPEECH_PAD_MS`).
|
||||
pub speech_pad_ms: u32,
|
||||
/// Speech required to fire the barge-in onset (`VAD_BARGE_IN_MIN_SPEECH_MS`).
|
||||
pub barge_in_min_speech_ms: u32,
|
||||
/// Feed the STT engine with VAD-segmented utterances instead of letting it
|
||||
/// run its own RMS segmentation (`VAD_SEGMENT_STT`).
|
||||
pub segment_stt: bool,
|
||||
/// Interrupt the agent when the user starts speaking over it
|
||||
/// (`VAD_BARGE_IN`).
|
||||
pub barge_in: bool,
|
||||
/// Log the speech probability of every analysis window (`VAD_DEBUG`).
|
||||
pub debug: bool,
|
||||
}
|
||||
|
||||
impl Default for VadSettings {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
enabled: true,
|
||||
model_path: PathBuf::from("models/silero_vad.onnx"),
|
||||
threshold: 0.5,
|
||||
neg_threshold: 0.35,
|
||||
min_speech_ms: 250,
|
||||
min_silence_ms: 300,
|
||||
speech_pad_ms: 30,
|
||||
barge_in_min_speech_ms: 150,
|
||||
segment_stt: true,
|
||||
barge_in: true,
|
||||
debug: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl VadSettings {
|
||||
/// Build the settings from environment variables, falling back to defaults.
|
||||
pub fn from_env() -> Self {
|
||||
let mut s = Self::default();
|
||||
if let Ok(p) = std::env::var("VAD_MODEL_PATH") {
|
||||
if !p.is_empty() {
|
||||
s.model_path = PathBuf::from(p);
|
||||
}
|
||||
}
|
||||
s.enabled = env_bool("VAD_ENABLED", s.enabled);
|
||||
s.segment_stt = env_bool("VAD_SEGMENT_STT", s.segment_stt);
|
||||
s.barge_in = env_bool("VAD_BARGE_IN", s.barge_in);
|
||||
s.debug = env_bool("VAD_DEBUG", s.debug);
|
||||
s.threshold = env_f32("VAD_THRESHOLD", s.threshold);
|
||||
s.neg_threshold = env_f32("VAD_NEG_THRESHOLD", s.neg_threshold);
|
||||
s.min_speech_ms = env_u32("VAD_MIN_SPEECH_MS", s.min_speech_ms);
|
||||
s.min_silence_ms = env_u32("VAD_MIN_SILENCE_MS", s.min_silence_ms);
|
||||
s.speech_pad_ms = env_u32("VAD_SPEECH_PAD_MS", s.speech_pad_ms);
|
||||
s.barge_in_min_speech_ms = env_u32("VAD_BARGE_IN_MIN_SPEECH_MS", s.barge_in_min_speech_ms);
|
||||
|
||||
// Hysteresis must stay below the speech threshold, and the fast
|
||||
// barge-in onset can never be slower than the utterance onset.
|
||||
s.neg_threshold = s.neg_threshold.min(s.threshold);
|
||||
s.barge_in_min_speech_ms = s.barge_in_min_speech_ms.min(s.min_speech_ms);
|
||||
s
|
||||
}
|
||||
}
|
||||
|
||||
/// Samples retained ahead of a detected onset (pre-roll + triggering speech).
|
||||
pub(crate) fn preroll_samples(settings: &VadSettings) -> usize {
|
||||
let ms = settings
|
||||
.speech_pad_ms
|
||||
.saturating_add(settings.min_speech_ms)
|
||||
.saturating_add(silero::WINDOW_MS);
|
||||
(crate::config::STT_SAMPLE_RATE as u64 * ms as u64 / 1000) as usize
|
||||
}
|
||||
@@ -0,0 +1,399 @@
|
||||
//! Silero VAD v5 ONNX wrapper and the onset/offset state machine.
|
||||
//!
|
||||
//! The model consumes fixed 32 ms windows (512 samples at 16 kHz) preceded by
|
||||
//! a 64-sample context carried over from the previous window, plus a recurrent
|
||||
//! state tensor. It returns the speech probability of the window.
|
||||
|
||||
use crate::config::STT_SAMPLE_RATE;
|
||||
use crate::vad::{preroll_samples, VadSettings};
|
||||
use anyhow::{bail, Context, Result};
|
||||
use ndarray::{Array1, ArrayD, IxDyn};
|
||||
use ort::environment::Environment;
|
||||
use ort::session::builder::GraphOptimizationLevel;
|
||||
use ort::session::{Session, SessionInputValue};
|
||||
use ort::value::{Value, ValueType};
|
||||
use parking_lot::Mutex;
|
||||
use std::collections::VecDeque;
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
use tracing::{debug, error, info};
|
||||
|
||||
/// Analysis window of the Silero v5 model at 16 kHz (512 samples = 32 ms).
|
||||
pub const WINDOW_SAMPLES: usize = 512;
|
||||
/// Duration of one analysis window in milliseconds.
|
||||
pub(crate) const WINDOW_MS: u32 = 32;
|
||||
/// Samples of the previous window prepended to the current one (v5 context).
|
||||
const CONTEXT_SAMPLES: usize = 64;
|
||||
/// Fallback recurrent state shape `[layers, batch, hidden]`.
|
||||
const STATE_SHAPE: [usize; 3] = [2, 1, 128];
|
||||
|
||||
const IN_INPUT: &str = "input";
|
||||
const IN_STATE: &str = "state";
|
||||
const IN_SR: &str = "sr";
|
||||
const OUT_PROB: &str = "output";
|
||||
const OUT_STATE: &str = "stateN";
|
||||
|
||||
struct Inner {
|
||||
/// Keeps the ONNX Runtime environment alive for the session.
|
||||
_environment: Environment,
|
||||
session: Session,
|
||||
/// `state` shape as declared by the graph (dynamic axes resolved to 1).
|
||||
state_shape: Vec<usize>,
|
||||
/// Rank of the `sr` input (0 = scalar, as exported by Silero v5).
|
||||
sr_rank: usize,
|
||||
}
|
||||
|
||||
/// Shared handle to the Silero VAD ONNX model. Cloning is cheap (`Arc`), the
|
||||
/// per-session recurrent state lives in [`VadDetector`].
|
||||
#[derive(Clone)]
|
||||
pub struct SileroVad {
|
||||
inner: Arc<Mutex<Inner>>,
|
||||
}
|
||||
|
||||
impl SileroVad {
|
||||
/// Load `silero_vad.onnx` from `path`.
|
||||
pub fn load(path: &Path) -> Result<Self> {
|
||||
info!(model = %path.display(), "loading Silero VAD model");
|
||||
|
||||
let environment = crate::onnx::environment()?;
|
||||
let session = Session::builder(&environment)
|
||||
.map_err(|e| anyhow::anyhow!("failed to create ONNX session builder: {e}"))?
|
||||
.with_optimization_level(GraphOptimizationLevel::Level3)
|
||||
.map_err(|e| anyhow::anyhow!("failed to set ONNX optimization level: {e}"))?
|
||||
.commit_from_file(path)
|
||||
.with_context(|| format!("failed to load ONNX model {}", path.display()))?;
|
||||
|
||||
// Introspect the recurrent-state and sample-rate inputs so we build
|
||||
// tensors matching the exported graph instead of hardcoding shapes.
|
||||
let state_shape = declared_shape(&session, IN_STATE).unwrap_or_else(|_| STATE_SHAPE.to_vec());
|
||||
let sr_rank = declared_shape(&session, IN_SR)
|
||||
.map(|s| s.len())
|
||||
.unwrap_or(0);
|
||||
|
||||
info!(
|
||||
model = %path.display(),
|
||||
window_samples = WINDOW_SAMPLES,
|
||||
window_ms = WINDOW_MS,
|
||||
state_shape = ?state_shape,
|
||||
sr_rank,
|
||||
"Silero VAD model loaded"
|
||||
);
|
||||
|
||||
Ok(Self {
|
||||
inner: Arc::new(Mutex::new(Inner {
|
||||
_environment: environment,
|
||||
session,
|
||||
state_shape,
|
||||
sr_rank,
|
||||
})),
|
||||
})
|
||||
}
|
||||
|
||||
/// Number of `f32` elements of the recurrent state, as declared by the
|
||||
/// graph. Detectors must allocate their state buffer with this size.
|
||||
pub fn state_len(&self) -> usize {
|
||||
self.inner.lock().state_shape.iter().product()
|
||||
}
|
||||
|
||||
/// Speech probability of one window. `state` is updated in place and
|
||||
/// `context` must hold the last [`CONTEXT_SAMPLES`] samples of the previous
|
||||
/// window (zeros on the first call).
|
||||
fn probability(&self, state: &mut [f32], context: &[f32], window: &[f32]) -> Result<f32> {
|
||||
if window.len() != WINDOW_SAMPLES {
|
||||
bail!("VAD window must be {WINDOW_SAMPLES} samples, got {}", window.len());
|
||||
}
|
||||
|
||||
let mut input = Vec::with_capacity(CONTEXT_SAMPLES + WINDOW_SAMPLES);
|
||||
input.extend_from_slice(context);
|
||||
input.extend_from_slice(window);
|
||||
let input_len = input.len();
|
||||
|
||||
let input_t = Value::from_array(ArrayD::from_shape_vec(IxDyn(&[1, input_len]), input)?)
|
||||
.context("failed to build VAD input tensor")?;
|
||||
|
||||
let mut guard = self.inner.lock();
|
||||
let inner = &mut *guard;
|
||||
|
||||
let state_t = Value::from_array(ArrayD::from_shape_vec(
|
||||
IxDyn(&inner.state_shape),
|
||||
state.to_vec(),
|
||||
)?)
|
||||
.context("failed to build VAD state tensor")?;
|
||||
|
||||
let sr = STT_SAMPLE_RATE as i64;
|
||||
let sr_t = if inner.sr_rank == 0 {
|
||||
Value::from_array(ArrayD::from_shape_vec(IxDyn(&[]), vec![sr])?)
|
||||
} else {
|
||||
Value::from_array(Array1::from_vec(vec![sr]))
|
||||
}
|
||||
.context("failed to build VAD sample-rate tensor")?;
|
||||
|
||||
let outputs = inner
|
||||
.session
|
||||
.run(vec![
|
||||
(IN_INPUT.to_string(), SessionInputValue::from(input_t)),
|
||||
(IN_STATE.to_string(), SessionInputValue::from(state_t)),
|
||||
(IN_SR.to_string(), SessionInputValue::from(sr_t)),
|
||||
])
|
||||
.context("Silero VAD inference failed")?;
|
||||
|
||||
let prob = outputs
|
||||
.get(OUT_PROB)
|
||||
.context("VAD model produced no 'output'")?
|
||||
.try_extract_tensor::<f32>()
|
||||
.context("failed to extract VAD probability")?
|
||||
.1
|
||||
.first()
|
||||
.copied()
|
||||
.unwrap_or(0.0);
|
||||
|
||||
let (_, new_state) = outputs
|
||||
.get(OUT_STATE)
|
||||
.context("VAD model produced no 'stateN'")?
|
||||
.try_extract_tensor::<f32>()
|
||||
.context("failed to extract VAD state")?;
|
||||
if new_state.len() == state.len() {
|
||||
state.copy_from_slice(new_state);
|
||||
}
|
||||
|
||||
Ok(prob.clamp(0.0, 1.0))
|
||||
}
|
||||
}
|
||||
|
||||
/// Read the declared shape of a tensor input, resolving dynamic axes to 1.
|
||||
fn declared_shape(session: &Session, name: &str) -> Result<Vec<usize>> {
|
||||
let input = session
|
||||
.inputs()
|
||||
.iter()
|
||||
.find(|i| i.name() == name)
|
||||
.with_context(|| format!("model has no input named '{name}'"))?;
|
||||
match input.dtype() {
|
||||
ValueType::Tensor { shape, .. } => Ok(shape
|
||||
.iter()
|
||||
.map(|&d| if d < 0 { 1 } else { d as usize })
|
||||
.collect()),
|
||||
other => bail!("input '{name}' is not a tensor: {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
/// Events emitted by [`VadDetector::push`].
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum VadEvent {
|
||||
/// The user just started speaking (after `barge_in_min_speech_ms`). Used to
|
||||
/// interrupt the agent.
|
||||
SpeechOnset { prob: f32 },
|
||||
/// Utterance start confirmed (after `min_speech_ms`).
|
||||
SpeechStart { prob: f32 },
|
||||
/// Utterance closed (after `min_silence_ms` of silence). `audio` is the
|
||||
/// captured 16 kHz mono PCM, including the pre-roll pad.
|
||||
SpeechEnd {
|
||||
audio: Vec<i16>,
|
||||
duration_ms: u64,
|
||||
avg_prob: f32,
|
||||
},
|
||||
}
|
||||
|
||||
/// Per-session VAD: feeds the model and turns probabilities into speech
|
||||
/// onset/offset events.
|
||||
pub struct VadDetector {
|
||||
model: SileroVad,
|
||||
settings: VadSettings,
|
||||
/// Recurrent model state (`state_shape` elements).
|
||||
state: Vec<f32>,
|
||||
/// Last [`CONTEXT_SAMPLES`] samples of the previous window.
|
||||
context: Vec<f32>,
|
||||
/// Samples received but not yet consumed by a full window.
|
||||
pending: Vec<i16>,
|
||||
/// Raw pre-roll ring, so the first phonemes are not clipped at onset.
|
||||
preroll: VecDeque<i16>,
|
||||
preroll_capacity: usize,
|
||||
/// Audio of the utterance being captured.
|
||||
utterance: Vec<i16>,
|
||||
/// Whether the detector is inside a confirmed utterance.
|
||||
in_speech: bool,
|
||||
/// Fast onset already reported for the current speech segment.
|
||||
onset_reported: bool,
|
||||
/// Utterance start already reported for the current speech segment.
|
||||
start_reported: bool,
|
||||
/// Consecutive speech milliseconds while idle.
|
||||
speech_run_ms: u32,
|
||||
/// Consecutive silence milliseconds while in speech.
|
||||
silence_run_ms: u32,
|
||||
/// Running probability average over the current utterance.
|
||||
prob_sum: f32,
|
||||
speech_windows: u64,
|
||||
/// Windows analysed since the detector was created (for logging).
|
||||
window_count: u64,
|
||||
}
|
||||
|
||||
impl VadDetector {
|
||||
/// Create a detector for one session.
|
||||
pub fn new(model: SileroVad, settings: VadSettings) -> Self {
|
||||
let preroll_capacity = preroll_samples(&settings).max(WINDOW_SAMPLES);
|
||||
let state_len = model.state_len().max(STATE_SHAPE.iter().product());
|
||||
Self {
|
||||
model,
|
||||
settings,
|
||||
state: vec![0.0; state_len],
|
||||
context: vec![0.0; CONTEXT_SAMPLES],
|
||||
pending: Vec::new(),
|
||||
preroll: VecDeque::with_capacity(preroll_capacity),
|
||||
preroll_capacity,
|
||||
utterance: Vec::new(),
|
||||
in_speech: false,
|
||||
onset_reported: false,
|
||||
start_reported: false,
|
||||
speech_run_ms: 0,
|
||||
silence_run_ms: 0,
|
||||
prob_sum: 0.0,
|
||||
speech_windows: 0,
|
||||
window_count: 0,
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether the detector is currently inside a confirmed utterance. Used by
|
||||
/// the backchannel scheduler to know when the user is still speaking.
|
||||
pub fn in_speech(&self) -> bool {
|
||||
self.in_speech
|
||||
}
|
||||
|
||||
/// Feed 16 kHz mono PCM and collect the events it triggers.
|
||||
///
|
||||
/// Any number of samples may be pushed; internally they are buffered into
|
||||
/// fixed 512-sample windows, so callers can pass arbitrary frame sizes.
|
||||
pub fn push(&mut self, samples: &[i16]) -> Vec<VadEvent> {
|
||||
let mut events = Vec::new();
|
||||
if samples.is_empty() {
|
||||
return events;
|
||||
}
|
||||
|
||||
// Keep the pre-roll ring up to date over the raw stream.
|
||||
for &s in samples {
|
||||
if self.preroll.len() >= self.preroll_capacity {
|
||||
self.preroll.pop_front();
|
||||
}
|
||||
self.preroll.push_back(s);
|
||||
}
|
||||
|
||||
self.pending.extend_from_slice(samples);
|
||||
while self.pending.len() >= WINDOW_SAMPLES {
|
||||
let window: Vec<i16> = self.pending.drain(..WINDOW_SAMPLES).collect();
|
||||
let prob = match self.analyse(&window) {
|
||||
Ok(p) => p,
|
||||
Err(e) => {
|
||||
error!(error = %e, "VAD inference failed — treating window as silence");
|
||||
0.0
|
||||
}
|
||||
};
|
||||
self.window_count += 1;
|
||||
self.step(prob, &window, &mut events);
|
||||
}
|
||||
|
||||
events
|
||||
}
|
||||
|
||||
/// Run the model on one window and refresh the recurrent context.
|
||||
fn analyse(&mut self, window: &[i16]) -> Result<f32> {
|
||||
let f32_window: Vec<f32> = window.iter().map(|&s| s as f32 / 32768.0).collect();
|
||||
let model = self.model.clone();
|
||||
let prob = model.probability(&mut self.state, &self.context, &f32_window)?;
|
||||
self.context.clear();
|
||||
self.context
|
||||
.extend_from_slice(&f32_window[f32_window.len() - CONTEXT_SAMPLES..]);
|
||||
Ok(prob)
|
||||
}
|
||||
|
||||
/// Advance the onset/offset state machine by one window.
|
||||
fn step(&mut self, prob: f32, window: &[i16], events: &mut Vec<VadEvent>) {
|
||||
if self.settings.debug {
|
||||
debug!(
|
||||
window = self.window_count,
|
||||
prob = format!("{prob:.3}"),
|
||||
in_speech = self.in_speech,
|
||||
"VAD window"
|
||||
);
|
||||
}
|
||||
|
||||
if !self.in_speech {
|
||||
if prob >= self.settings.threshold {
|
||||
self.speech_run_ms = self.speech_run_ms.saturating_add(WINDOW_MS);
|
||||
} else {
|
||||
self.speech_run_ms = 0;
|
||||
}
|
||||
|
||||
if !self.onset_reported && self.speech_run_ms >= self.settings.barge_in_min_speech_ms {
|
||||
self.onset_reported = true;
|
||||
debug!(
|
||||
speech_ms = self.speech_run_ms,
|
||||
prob = format!("{prob:.3}"),
|
||||
"VAD speech onset"
|
||||
);
|
||||
events.push(VadEvent::SpeechOnset { prob });
|
||||
}
|
||||
|
||||
if !self.start_reported && self.speech_run_ms >= self.settings.min_speech_ms {
|
||||
self.start_reported = true;
|
||||
self.in_speech = true;
|
||||
// The pre-roll ring already holds this window (it is refreshed
|
||||
// before the pending buffer is drained in `push`), so seed the
|
||||
// utterance from it instead of appending again.
|
||||
self.utterance.clear();
|
||||
self.utterance.extend(self.preroll.iter().copied());
|
||||
self.prob_sum = prob;
|
||||
self.speech_windows = 1;
|
||||
self.silence_run_ms = 0;
|
||||
debug!(
|
||||
speech_ms = self.speech_run_ms,
|
||||
prob = format!("{prob:.3}"),
|
||||
preroll_samples = self.utterance.len(),
|
||||
"VAD speech start"
|
||||
);
|
||||
events.push(VadEvent::SpeechStart { prob });
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
self.utterance.extend_from_slice(window);
|
||||
self.prob_sum += prob;
|
||||
self.speech_windows += 1;
|
||||
|
||||
if prob < self.settings.neg_threshold {
|
||||
self.silence_run_ms = self.silence_run_ms.saturating_add(WINDOW_MS);
|
||||
} else {
|
||||
self.silence_run_ms = 0;
|
||||
}
|
||||
|
||||
if self.silence_run_ms >= self.settings.min_silence_ms {
|
||||
let audio = std::mem::take(&mut self.utterance);
|
||||
let avg_prob = if self.speech_windows > 0 {
|
||||
self.prob_sum / self.speech_windows as f32
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
let duration_ms = (audio.len() as u64 * 1000) / STT_SAMPLE_RATE as u64;
|
||||
debug!(
|
||||
duration_ms,
|
||||
samples = audio.len(),
|
||||
avg_prob = format!("{avg_prob:.3}"),
|
||||
"VAD speech end"
|
||||
);
|
||||
events.push(VadEvent::SpeechEnd {
|
||||
audio,
|
||||
duration_ms,
|
||||
avg_prob,
|
||||
});
|
||||
self.reset_segment();
|
||||
}
|
||||
}
|
||||
|
||||
/// Clear the per-segment state after an utterance is closed.
|
||||
fn reset_segment(&mut self) {
|
||||
self.in_speech = false;
|
||||
self.onset_reported = false;
|
||||
self.start_reported = false;
|
||||
self.speech_run_ms = 0;
|
||||
self.silence_run_ms = 0;
|
||||
self.prob_sum = 0.0;
|
||||
self.speech_windows = 0;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,525 @@
|
||||
<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8" />
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0" />
|
||||
<title>Mivocero — Speech-to-Speech Agent</title>
|
||||
<style>
|
||||
:root {
|
||||
--bg: #0b0f1a;
|
||||
--panel: #141b2d;
|
||||
--accent: #4f8cff;
|
||||
--accent-dim: #2a4a8a;
|
||||
--ok: #3ddc97;
|
||||
--warn: #ffb454;
|
||||
--err: #ff5470;
|
||||
--text: #e6edf7;
|
||||
--muted: #8a94a6;
|
||||
--border: #243049;
|
||||
}
|
||||
* { box-sizing: border-box; }
|
||||
html, body { height: 100%; }
|
||||
body {
|
||||
margin: 0;
|
||||
font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif;
|
||||
background: var(--bg);
|
||||
color: var(--text);
|
||||
display: flex;
|
||||
justify-content: center;
|
||||
padding: 18px;
|
||||
}
|
||||
.app {
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
width: 100%;
|
||||
max-width: 820px;
|
||||
height: calc(100vh - 36px);
|
||||
background: var(--panel);
|
||||
border: 1px solid var(--border);
|
||||
border-radius: 14px;
|
||||
overflow: hidden;
|
||||
}
|
||||
header {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
gap: 12px;
|
||||
padding: 14px 18px;
|
||||
border-bottom: 1px solid var(--border);
|
||||
}
|
||||
header h1 { font-size: 16px; margin: 0; flex: 1; }
|
||||
.status { display: inline-flex; align-items: center; gap: 8px; font-size: 13px; color: var(--muted); }
|
||||
.dot { width: 10px; height: 10px; border-radius: 50%; background: var(--muted); }
|
||||
.dot.on { background: var(--ok); box-shadow: 0 0 8px var(--ok); }
|
||||
.dot.err { background: var(--err); box-shadow: 0 0 8px var(--err); }
|
||||
button {
|
||||
background: var(--accent-dim);
|
||||
color: var(--text);
|
||||
border: 1px solid var(--accent);
|
||||
border-radius: 8px;
|
||||
padding: 8px 14px;
|
||||
font-size: 13px;
|
||||
cursor: pointer;
|
||||
transition: all .15s;
|
||||
}
|
||||
button:hover:not(:disabled) { background: var(--accent); }
|
||||
button:disabled { opacity: .4; cursor: not-allowed; }
|
||||
button.danger { border-color: var(--err); background: rgba(255,84,112,.15); }
|
||||
button.danger:hover:not(:disabled) { background: var(--err); }
|
||||
|
||||
#chat {
|
||||
flex: 1;
|
||||
overflow-y: auto;
|
||||
padding: 18px;
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
gap: 12px;
|
||||
}
|
||||
.msg { display: flex; flex-direction: column; gap: 4px; max-width: 80%; }
|
||||
.msg .who { font-size: 11px; text-transform: uppercase; letter-spacing: .5px; color: var(--muted); }
|
||||
.msg .bubble {
|
||||
padding: 10px 13px;
|
||||
border-radius: 12px;
|
||||
font-size: 14px;
|
||||
line-height: 1.45;
|
||||
white-space: pre-wrap;
|
||||
word-break: break-word;
|
||||
}
|
||||
.msg.user { align-self: flex-end; align-items: flex-end; }
|
||||
.msg.user .bubble { background: var(--accent-dim); border: 1px solid var(--accent); }
|
||||
.msg.agent { align-self: flex-start; }
|
||||
.msg.agent .bubble { background: var(--bg); border: 1px solid var(--border); }
|
||||
.msg.error .bubble { border-color: var(--err); color: var(--err); }
|
||||
.msg .bubble.pending::after {
|
||||
content: "▍";
|
||||
animation: blink 1s step-start infinite;
|
||||
color: var(--accent);
|
||||
}
|
||||
@keyframes blink { 50% { opacity: 0; } }
|
||||
.hint { color: var(--muted); font-size: 13px; text-align: center; margin: auto; }
|
||||
|
||||
footer { border-top: 1px solid var(--border); padding: 12px 18px; }
|
||||
.meters { display: flex; align-items: center; gap: 14px; margin-bottom: 10px; }
|
||||
.meter { height: 8px; background: var(--bg); border-radius: 4px; overflow: hidden; flex: 1; }
|
||||
.meter > div { height: 100%; width: 0%; background: linear-gradient(90deg, var(--ok), var(--warn), var(--err)); transition: width .05s; }
|
||||
.meters span { font-size: 11px; color: var(--muted); white-space: nowrap; }
|
||||
.composer { display: flex; gap: 8px; }
|
||||
.composer input {
|
||||
flex: 1;
|
||||
background: var(--bg);
|
||||
border: 1px solid var(--border);
|
||||
border-radius: 8px;
|
||||
padding: 10px 12px;
|
||||
color: var(--text);
|
||||
font-size: 14px;
|
||||
}
|
||||
#log {
|
||||
display: none;
|
||||
max-height: 110px;
|
||||
overflow-y: auto;
|
||||
margin-top: 10px;
|
||||
padding: 8px;
|
||||
background: var(--bg);
|
||||
border: 1px solid var(--border);
|
||||
border-radius: 8px;
|
||||
font-family: ui-monospace, Menlo, monospace;
|
||||
font-size: 11px;
|
||||
color: var(--muted);
|
||||
white-space: pre-wrap;
|
||||
}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div class="app">
|
||||
<header>
|
||||
<h1>🎙️ Mivocero</h1>
|
||||
<span class="status"><span class="dot" id="connDot"></span><span id="connText">disconnected</span></span>
|
||||
<button id="btnConnect" class="ok">Connect</button>
|
||||
<button id="btnInterrupt" class="danger" disabled>Stop</button>
|
||||
</header>
|
||||
|
||||
<div id="chat">
|
||||
<div class="hint" id="emptyHint">
|
||||
Press <b>Connect</b>, allow microphone access and start talking.<br />
|
||||
The agent replies with synthesized speech and the transcript appears here.
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<footer>
|
||||
<div class="meters">
|
||||
<span>🎤</span><div class="meter"><div id="micMeter"></div></div>
|
||||
<span>🔊</span><div class="meter"><div id="outMeter"></div></div>
|
||||
</div>
|
||||
<form class="composer" id="composer">
|
||||
<input type="text" id="textInput" placeholder="Or type a message…" disabled autocomplete="off" />
|
||||
<button type="submit" id="btnSend" disabled>Send</button>
|
||||
</form>
|
||||
<div id="log"></div>
|
||||
</footer>
|
||||
</div>
|
||||
|
||||
<script>
|
||||
// ---------------------------------------------------------------------------
|
||||
// Mivocero — speech-to-speech agent client
|
||||
//
|
||||
// mic -> AudioWorklet (44.1 kHz, 16-bit LE mono, 20 ms frames) -> WebSocket
|
||||
// WebSocket binary frames -> playback queue -> AudioWorklet -> speakers
|
||||
//
|
||||
// Text frames received from the server are JSON chat events:
|
||||
// {"type":"user","text":"…"} recognized user speech
|
||||
// {"type":"assistant_delta","text":"…"} incremental assistant reply
|
||||
// {"type":"assistant_end"} assistant finished
|
||||
// {"type":"error","text":"…"} pipeline error
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
const SAMPLE_RATE_OUT = 44100; // server <-> client (matches the TTS rate)
|
||||
const SAMPLE_RATE_IN = 44100;
|
||||
const FRAME_MS = 20;
|
||||
const FRAME_SAMPLES = SAMPLE_RATE_IN * FRAME_MS / 1000; // 882
|
||||
|
||||
const $ = (id) => document.getElementById(id);
|
||||
const wsDot = $('connDot');
|
||||
const wsText = $('connText');
|
||||
const btnConnect = $('btnConnect');
|
||||
const btnInterrupt = $('btnInterrupt');
|
||||
const micMeter = $('micMeter');
|
||||
const outMeter = $('outMeter');
|
||||
const chatEl = $('chat');
|
||||
const emptyHint = $('emptyHint');
|
||||
const composer = $('composer');
|
||||
const textInput = $('textInput');
|
||||
const btnSend = $('btnSend');
|
||||
const logEl = $('log');
|
||||
|
||||
let ws = null;
|
||||
let audioCtx = null;
|
||||
let micStream = null;
|
||||
let captureNode = null;
|
||||
let playbackNode = null;
|
||||
let outputGain = null;
|
||||
let agentBubble = null; // bubble being streamed by the assistant
|
||||
let agentText = ''; // accumulated assistant text for that bubble
|
||||
|
||||
function log(msg) {
|
||||
const ts = new Date().toLocaleTimeString();
|
||||
logEl.textContent = `[${ts}] ${msg}\n` + logEl.textContent;
|
||||
if (logEl.textContent.length > 4000) logEl.textContent = logEl.textContent.slice(0, 4000);
|
||||
}
|
||||
|
||||
function setConn(state, text) {
|
||||
wsDot.className = 'dot ' + (state === 'on' ? 'on' : state === 'err' ? 'err' : '');
|
||||
wsText.textContent = text;
|
||||
}
|
||||
|
||||
function atBottom() {
|
||||
return chatEl.scrollHeight - chatEl.scrollTop - chatEl.clientHeight < 60;
|
||||
}
|
||||
function scrollDown(force) {
|
||||
if (force || atBottom()) chatEl.scrollTop = chatEl.scrollHeight;
|
||||
}
|
||||
|
||||
// --- Chat rendering --------------------------------------------------------
|
||||
|
||||
function addMessage(role, text, cls = '') {
|
||||
emptyHint?.remove();
|
||||
const wrap = document.createElement('div');
|
||||
wrap.className = 'msg ' + role + (cls ? ' ' + cls : '');
|
||||
const who = document.createElement('div');
|
||||
who.className = 'who';
|
||||
who.textContent = role === 'user' ? 'You' : role === 'error' ? 'Error' : 'Agent';
|
||||
const bubble = document.createElement('div');
|
||||
bubble.className = 'bubble';
|
||||
bubble.textContent = text;
|
||||
wrap.append(who, bubble);
|
||||
chatEl.appendChild(wrap);
|
||||
scrollDown(true);
|
||||
return bubble;
|
||||
}
|
||||
|
||||
function startAgentMessage() {
|
||||
agentText = '';
|
||||
agentBubble = addMessage('agent', '');
|
||||
agentBubble.classList.add('pending');
|
||||
}
|
||||
|
||||
function appendAgentDelta(text) {
|
||||
if (!agentBubble) startAgentMessage();
|
||||
agentText += text;
|
||||
agentBubble.textContent = agentText;
|
||||
scrollDown(false);
|
||||
}
|
||||
|
||||
function endAgentMessage() {
|
||||
if (agentBubble) agentBubble.classList.remove('pending');
|
||||
agentBubble = null;
|
||||
agentText = '';
|
||||
}
|
||||
|
||||
// --- Audio worklets --------------------------------------------------------
|
||||
|
||||
const CAPTURE_PROCESSOR = `
|
||||
class CaptureProcessor extends AudioWorkletProcessor {
|
||||
constructor() {
|
||||
super();
|
||||
this.buf = [];
|
||||
this.frameSize = ${FRAME_SAMPLES};
|
||||
this.targetRate = ${SAMPLE_RATE_IN};
|
||||
this.ratio = this.targetRate / sampleRate;
|
||||
this.phase = 0;
|
||||
this.prev = 0;
|
||||
}
|
||||
process(inputs) {
|
||||
const inp = inputs[0][0];
|
||||
if (!inp) return true;
|
||||
let acc = 0;
|
||||
for (let i = 0; i < inp.length; i++) {
|
||||
const cur = inp[i];
|
||||
acc += cur * cur;
|
||||
while (this.phase < i + 1) {
|
||||
const frac = this.phase - Math.floor(this.phase);
|
||||
const s = this.prev + (cur - this.prev) * frac;
|
||||
this.buf.push(s);
|
||||
if (this.buf.length >= this.frameSize) this.flush();
|
||||
this.phase += this.ratio;
|
||||
}
|
||||
this.prev = cur;
|
||||
}
|
||||
this.phase -= inp.length;
|
||||
const rms = Math.sqrt(acc / inp.length);
|
||||
this.port.postMessage({ rms });
|
||||
return true;
|
||||
}
|
||||
flush() {
|
||||
const n = this.frameSize;
|
||||
const buf = new ArrayBuffer(n * 2);
|
||||
const dv = new DataView(buf);
|
||||
for (let i = 0; i < n; i++) {
|
||||
let v = Math.max(-1, Math.min(1, this.buf[i]));
|
||||
dv.setInt16(i * 2, v < 0 ? v * 0x8000 : v * 0x7FFF, true);
|
||||
}
|
||||
this.buf.splice(0, n);
|
||||
this.port.postMessage({ pcm: buf }, [buf]);
|
||||
}
|
||||
}
|
||||
registerProcessor('capture-processor', CaptureProcessor);
|
||||
`;
|
||||
|
||||
const PLAYBACK_PROCESSOR = `
|
||||
class PlaybackProcessor extends AudioWorkletProcessor {
|
||||
constructor() {
|
||||
super();
|
||||
this.buf = new Float32Array(48000 * 2); // ~2 s ring @ 44.1 kHz
|
||||
this.head = 0;
|
||||
this.tail = 0;
|
||||
this.preBuffer = Math.floor(sampleRate * 0.30); // 300 ms jitter buffer
|
||||
this.buffering = true;
|
||||
this.underrunSamples = 0;
|
||||
this.underrunThreshold = Math.floor(sampleRate * 0.10);
|
||||
this.port.onmessage = (e) => {
|
||||
const pcm = e.data;
|
||||
for (let i = 0; i < pcm.length; i++) {
|
||||
this.buf[this.head] = pcm[i];
|
||||
this.head = (this.head + 1) % this.buf.length;
|
||||
if (this.head === this.tail) this.tail = (this.tail + 1) % this.buf.length;
|
||||
}
|
||||
};
|
||||
}
|
||||
available() { return (this.head - this.tail + this.buf.length) % this.buf.length; }
|
||||
process(_inputs, outputs) {
|
||||
const out = outputs[0][0];
|
||||
let level = 0;
|
||||
for (let i = 0; i < out.length; i++) {
|
||||
let s = 0;
|
||||
if (this.buffering) {
|
||||
if (this.available() >= this.preBuffer) { this.buffering = false; this.underrunSamples = 0; }
|
||||
}
|
||||
if (!this.buffering) {
|
||||
if (this.tail !== this.head) {
|
||||
s = this.buf[this.tail];
|
||||
this.tail = (this.tail + 1) % this.buf.length;
|
||||
this.underrunSamples = 0;
|
||||
} else {
|
||||
this.underrunSamples++;
|
||||
if (this.underrunSamples >= this.underrunThreshold) { this.buffering = true; this.underrunSamples = 0; }
|
||||
}
|
||||
}
|
||||
out[i] = s;
|
||||
level += s * s;
|
||||
}
|
||||
this.port.postMessage({ rms: Math.sqrt(level / out.length) });
|
||||
return true;
|
||||
}
|
||||
}
|
||||
registerProcessor('playback-processor', PlaybackProcessor);
|
||||
`;
|
||||
|
||||
async function startAudio() {
|
||||
audioCtx = new (window.AudioContext || window.webkitAudioContext)({ sampleRate: SAMPLE_RATE_OUT });
|
||||
if (audioCtx.state !== 'running') {
|
||||
try { await audioCtx.resume(); } catch (e) {}
|
||||
}
|
||||
await audioCtx.audioWorklet.addModule(URL.createObjectURL(new Blob([CAPTURE_PROCESSOR], { type: 'application/javascript' })));
|
||||
await audioCtx.audioWorklet.addModule(URL.createObjectURL(new Blob([PLAYBACK_PROCESSOR], { type: 'application/javascript' })));
|
||||
|
||||
micStream = await navigator.mediaDevices.getUserMedia({
|
||||
audio: { channelCount: 1, echoCancellation: true, noiseSuppression: true }
|
||||
});
|
||||
const source = audioCtx.createMediaStreamSource(micStream);
|
||||
captureNode = new AudioWorkletNode(audioCtx, 'capture-processor');
|
||||
captureNode.port.onmessage = (e) => {
|
||||
if (e.data.rms !== undefined) {
|
||||
micMeter.style.width = Math.min(100, e.data.rms * 200) + '%';
|
||||
} else if (e.data.pcm && ws && ws.readyState === WebSocket.OPEN) {
|
||||
ws.send(e.data.pcm);
|
||||
}
|
||||
};
|
||||
source.connect(captureNode);
|
||||
const silent = audioCtx.createGain();
|
||||
silent.gain.value = 0;
|
||||
captureNode.connect(silent).connect(audioCtx.destination);
|
||||
|
||||
playbackNode = new AudioWorkletNode(audioCtx, 'playback-processor');
|
||||
playbackNode.port.onmessage = (e) => {
|
||||
if (e.data.rms !== undefined) outMeter.style.width = Math.min(100, e.data.rms * 200) + '%';
|
||||
};
|
||||
outputGain = audioCtx.createGain();
|
||||
outputGain.gain.value = 3.0; // Fish Audio output level is fairly low
|
||||
playbackNode.connect(outputGain).connect(audioCtx.destination);
|
||||
}
|
||||
|
||||
function stopAudio() {
|
||||
try { captureNode?.disconnect(); } catch (e) {}
|
||||
try { playbackNode?.disconnect(); } catch (e) {}
|
||||
try { outputGain?.disconnect(); } catch (e) {}
|
||||
micStream?.getTracks().forEach((t) => t.stop());
|
||||
audioCtx?.close();
|
||||
audioCtx = null;
|
||||
micStream = null;
|
||||
captureNode = playbackNode = outputGain = null;
|
||||
micMeter.style.width = '0%';
|
||||
outMeter.style.width = '0%';
|
||||
}
|
||||
|
||||
function feedPlayback(arrayBuffer) {
|
||||
const view = new DataView(arrayBuffer);
|
||||
const n = Math.floor(view.byteLength / 2);
|
||||
const f32 = new Float32Array(n);
|
||||
for (let i = 0; i < n; i++) f32[i] = view.getInt16(i * 2, true) / 32768;
|
||||
playbackNode?.port.postMessage(f32);
|
||||
}
|
||||
|
||||
// --- WebSocket -------------------------------------------------------------
|
||||
|
||||
function handleServerText(text) {
|
||||
let event = null;
|
||||
try { event = JSON.parse(text); } catch (e) { log('server: ' + text); return; }
|
||||
switch (event.type) {
|
||||
case 'user':
|
||||
endAgentMessage(); // a new turn starts
|
||||
addMessage('user', event.text);
|
||||
break;
|
||||
case 'assistant_delta':
|
||||
appendAgentDelta(event.text);
|
||||
break;
|
||||
case 'assistant_end':
|
||||
endAgentMessage();
|
||||
break;
|
||||
case 'error':
|
||||
endAgentMessage();
|
||||
addMessage('error', event.text, 'error');
|
||||
log('error: ' + event.text);
|
||||
break;
|
||||
default:
|
||||
log('server: ' + text);
|
||||
}
|
||||
}
|
||||
|
||||
async function connect() {
|
||||
if (!navigator.mediaDevices?.getUserMedia) {
|
||||
const why = window.isSecureContext
|
||||
? 'navigator.mediaDevices is unavailable in this browser'
|
||||
: `insecure origin (${location.origin}) — the microphone requires HTTPS or localhost`;
|
||||
addMessage('error', 'Microphone unavailable: ' + why, 'error');
|
||||
setConn('err', 'mic error');
|
||||
return;
|
||||
}
|
||||
try {
|
||||
await startAudio();
|
||||
} catch (e) {
|
||||
addMessage('error', 'Microphone access failed: ' + (e.name ? e.name + ': ' : '') + (e.message || e), 'error');
|
||||
setConn('err', 'mic error');
|
||||
return;
|
||||
}
|
||||
|
||||
const proto = location.protocol === 'https:' ? 'wss:' : 'ws:';
|
||||
const url = `${proto}//${location.host}/ws`;
|
||||
log('connecting to ' + url);
|
||||
ws = new WebSocket(url);
|
||||
ws.binaryType = 'arraybuffer';
|
||||
|
||||
ws.onopen = () => {
|
||||
setConn('on', 'connected');
|
||||
log('connected');
|
||||
btnConnect.disabled = true;
|
||||
btnInterrupt.disabled = false;
|
||||
textInput.disabled = false;
|
||||
btnSend.disabled = false;
|
||||
};
|
||||
ws.onmessage = (e) => {
|
||||
if (e.data instanceof ArrayBuffer) {
|
||||
feedPlayback(e.data);
|
||||
} else {
|
||||
handleServerText(e.data);
|
||||
}
|
||||
};
|
||||
ws.onerror = () => { log('websocket error'); setConn('err', 'error'); };
|
||||
ws.onclose = () => {
|
||||
log('disconnected');
|
||||
setConn('', 'disconnected');
|
||||
btnConnect.disabled = false;
|
||||
btnInterrupt.disabled = true;
|
||||
textInput.disabled = true;
|
||||
btnSend.disabled = true;
|
||||
endAgentMessage();
|
||||
stopAudio();
|
||||
};
|
||||
}
|
||||
|
||||
function disconnect() {
|
||||
try { ws?.close(); } catch (e) {}
|
||||
stopAudio();
|
||||
}
|
||||
|
||||
function sendJson(obj) {
|
||||
if (ws && ws.readyState === WebSocket.OPEN) {
|
||||
ws.send(JSON.stringify(obj));
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
btnConnect.onclick = () => (ws ? disconnect() : connect());
|
||||
btnInterrupt.onclick = () => {
|
||||
if (sendJson({ type: 'interrupt' })) {
|
||||
log('→ interrupt');
|
||||
endAgentMessage();
|
||||
}
|
||||
};
|
||||
|
||||
composer.onsubmit = (e) => {
|
||||
e.preventDefault();
|
||||
const text = textInput.value.trim();
|
||||
if (!text) return;
|
||||
if (sendJson({ type: 'user_text', text })) {
|
||||
endAgentMessage();
|
||||
addMessage('user', text);
|
||||
textInput.value = '';
|
||||
log('→ user_text');
|
||||
}
|
||||
};
|
||||
|
||||
// Optional verbose log: click the header status to toggle.
|
||||
wsText.style.cursor = 'pointer';
|
||||
wsText.onclick = () => { logEl.style.display = logEl.style.display === 'block' ? 'none' : 'block'; };
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
@@ -0,0 +1,9 @@
|
||||
//! Static web assets (the playground and chat HTML clients).
|
||||
|
||||
/// The playground single-page app (audio pipeline test bench), embedded into
|
||||
/// the binary at compile time.
|
||||
pub const PLAYGROUND_HTML: &str = include_str!("playground.html");
|
||||
|
||||
/// The `/mivocero` speech-to-speech chat client, embedded into the binary at
|
||||
/// compile time.
|
||||
pub const MIVOCERO_HTML: &str = include_str!("mivocero.html");
|
||||
@@ -0,0 +1,586 @@
|
||||
<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8" />
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0" />
|
||||
<title>Audio Agent Engine — Playground</title>
|
||||
<style>
|
||||
:root {
|
||||
--bg: #0b0f1a;
|
||||
--panel: #141b2d;
|
||||
--accent: #4f8cff;
|
||||
--accent-dim: #2a4a8a;
|
||||
--ok: #3ddc97;
|
||||
--warn: #ffb454;
|
||||
--err: #ff5470;
|
||||
--text: #e6edf7;
|
||||
--muted: #8a94a6;
|
||||
--border: #243049;
|
||||
}
|
||||
* { box-sizing: border-box; }
|
||||
body {
|
||||
margin: 0;
|
||||
font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif;
|
||||
background: var(--bg);
|
||||
color: var(--text);
|
||||
min-height: 100vh;
|
||||
padding: 24px;
|
||||
}
|
||||
h1 { margin: 0 0 4px; font-size: 22px; }
|
||||
.sub { color: var(--muted); font-size: 13px; margin-bottom: 20px; }
|
||||
.grid { display: grid; grid-template-columns: 1fr 1fr; gap: 16px; max-width: 1100px; }
|
||||
@media (max-width: 820px) { .grid { grid-template-columns: 1fr; } }
|
||||
.card {
|
||||
background: var(--panel);
|
||||
border: 1px solid var(--border);
|
||||
border-radius: 12px;
|
||||
padding: 18px;
|
||||
}
|
||||
.card h2 { margin: 0 0 14px; font-size: 15px; color: var(--muted); text-transform: uppercase; letter-spacing: .5px; }
|
||||
.row { display: flex; gap: 10px; align-items: center; flex-wrap: wrap; }
|
||||
button {
|
||||
background: var(--accent-dim);
|
||||
color: var(--text);
|
||||
border: 1px solid var(--accent);
|
||||
border-radius: 8px;
|
||||
padding: 9px 16px;
|
||||
font-size: 14px;
|
||||
cursor: pointer;
|
||||
transition: all .15s;
|
||||
}
|
||||
button:hover:not(:disabled) { background: var(--accent); }
|
||||
button:disabled { opacity: .4; cursor: not-allowed; }
|
||||
button.danger { border-color: var(--err); background: rgba(255,84,112,.15); }
|
||||
button.danger:hover:not(:disabled) { background: var(--err); }
|
||||
button.ok { border-color: var(--ok); background: rgba(61,220,151,.15); }
|
||||
button.ok:hover:not(:disabled) { background: var(--ok); color: #062; }
|
||||
input[type=text] {
|
||||
background: var(--bg);
|
||||
border: 1px solid var(--border);
|
||||
border-radius: 8px;
|
||||
padding: 9px 12px;
|
||||
color: var(--text);
|
||||
font-size: 14px;
|
||||
flex: 1;
|
||||
min-width: 180px;
|
||||
}
|
||||
.status { display: inline-flex; align-items: center; gap: 8px; font-size: 13px; }
|
||||
.dot { width: 10px; height: 10px; border-radius: 50%; background: var(--muted); }
|
||||
.dot.on { background: var(--ok); box-shadow: 0 0 8px var(--ok); }
|
||||
.dot.err { background: var(--err); box-shadow: 0 0 8px var(--err); }
|
||||
.meter { height: 10px; background: var(--bg); border-radius: 5px; overflow: hidden; margin: 6px 0 14px; }
|
||||
.meter > div { height: 100%; width: 0%; background: linear-gradient(90deg, var(--ok), var(--warn), var(--err)); transition: width .04s; }
|
||||
label.check { display: inline-flex; align-items: center; gap: 6px; font-size: 13px; color: var(--muted); cursor: pointer; }
|
||||
#log {
|
||||
background: var(--bg);
|
||||
border: 1px solid var(--border);
|
||||
border-radius: 8px;
|
||||
padding: 10px;
|
||||
font-family: ui-monospace, "SF Mono", Menlo, monospace;
|
||||
font-size: 12px;
|
||||
height: 220px;
|
||||
overflow-y: auto;
|
||||
white-space: pre-wrap;
|
||||
word-break: break-word;
|
||||
}
|
||||
.log-line { opacity: .85; }
|
||||
.log-line.in { color: var(--accent); }
|
||||
.log-line.out { color: var(--ok); }
|
||||
.log-line.err { color: var(--err); }
|
||||
.log-line.sys { color: var(--muted); }
|
||||
.stats { display: flex; gap: 18px; font-size: 12px; color: var(--muted); flex-wrap: wrap; }
|
||||
.stats b { color: var(--text); font-weight: 600; }
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
|
||||
<h1>🎙️ Audio Agent Engine — Playground</h1>
|
||||
<div class="sub">Full-duplex WebSocket PCM streaming: mic in → STT/EOU → filler + Drama 3 TTS → mixed audio out (24 kHz / 16-bit mono).</div>
|
||||
|
||||
<div class="grid">
|
||||
<!-- Connection card -->
|
||||
<div class="card">
|
||||
<h2>Connection</h2>
|
||||
<div class="row">
|
||||
<input type="text" id="wsUrl" value="" placeholder="ws://host:8787/ws" />
|
||||
</div>
|
||||
<div class="row" style="margin-top:10px;">
|
||||
<button id="btnConnect" class="ok">Connect & Start Mic</button>
|
||||
<button id="btnDisconnect" disabled>Disconnect</button>
|
||||
<button id="btnInterrupt" class="danger" disabled>Interrupt (barge-in)</button>
|
||||
</div>
|
||||
<div style="margin-top:12px;">
|
||||
<span class="status"><span class="dot" id="connDot"></span><span id="connText">disconnected</span></span>
|
||||
</div>
|
||||
<div style="margin-top:16px;">
|
||||
<label class="check"><input type="checkbox" id="chkSendSilence" checked /> Send silence when idle (keepalive)</label>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- Text-to-speech card -->
|
||||
<div class="card">
|
||||
<h2>Text-to-Speech (test TTS pipeline)</h2>
|
||||
<div class="row">
|
||||
<textarea id="ttsText" rows="12" style="width:100%;padding:10px;border:1px solid var(--border);border-radius:6px;background:var(--bg);color:var(--text);font-size:13px;resize:vertical;">Cristóbal Colón fue un navegante y cartógrafo genovés que, al servicio de la Corona de Castilla, realizó cuatro famosos viajes a través del océano Atlántico, propiciando el encuentro entre Europa y América en 1492. Aunque su objetivo original era abrir una ruta comercial hacia las Indias Orientales navegando hacia el oeste, su llegada a las islas del Caribe cambió drásticamente el curso de la historia universal.
|
||||
|
||||
Resumen Cronológico de su Vida
|
||||
|
||||
• Primeros años (1451–1476): Nació en Génova (actual Italia) en el seno de una familia de tejedores. Desde joven se dedicó a la navegación comercial en el mar Mediterráneo y más tarde se trasladó a Portugal.
|
||||
• El Proyecto (1480s): Basándose en la idea de que la Tierra era una esfera, calculó (con un error matemático considerable a la baja) que la distancia entre Europa y Asia por el oeste era mucho más corta de lo que se creía. Tras ser rechazado por el rey de Portugal, buscó el patrocinio de los Reyes Católicos de España, Isabel I de Castilla y Fernando II de Aragón.
|
||||
• Las Capitulaciones de Santa Fe (1492): Los reyes financiaron su expedición y le otorgaron los títulos de "Almirante del Mar Océano", "Visorrey" y "Gobernador" de las tierras que descubriese, además del diezmo de las riquezas obtenidas.</textarea>
|
||||
</div>
|
||||
<div class="row" style="margin-top:10px;">
|
||||
<button id="btnSpeak" disabled>Speak</button>
|
||||
<span class="stats"><span>💡 Sends text to the server for Drama 3 TTS</span></span>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- Voice cloning (reference) card -->
|
||||
<div class="card">
|
||||
<h2>Voice Cloning (reference audio)</h2>
|
||||
<div class="row">
|
||||
<input type="text" id="refAudio" value="/root/pipecat-voice/voice_profiles/solis.wav" placeholder="Reference WAV path (on Fish server)" />
|
||||
</div>
|
||||
<div class="row" style="margin-top:6px;">
|
||||
<input type="text" id="refText" value="" placeholder="Reference transcript (or leave empty to load .txt sibling)" />
|
||||
</div>
|
||||
<div class="row" style="margin-top:10px;">
|
||||
<button id="btnSetRef" disabled>Set Reference</button>
|
||||
<button id="btnClearRef" disabled>Clear Reference</button>
|
||||
<span class="stats" id="refStatus"><span>not set</span></span>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- Audio levels + stats (full width) -->
|
||||
<div class="card" style="margin-top:16px;max-width:1100px;">
|
||||
<h2>Audio Levels</h2>
|
||||
<div style="font-size:12px;color:var(--muted);">Mic input</div>
|
||||
<div class="meter"><div id="micMeter"></div></div>
|
||||
<div style="font-size:12px;color:var(--muted);">Server output (TTS + filler)</div>
|
||||
<div class="meter"><div id="outMeter"></div></div>
|
||||
<div class="stats" id="stats" style="margin-top:10px;">
|
||||
<span>🎤 frames sent: <b id="statSent">0</b></span>
|
||||
<span>🔊 frames recv: <b id="statRecv">0</b></span>
|
||||
<span>⚡ input rate: <b id="statInRate">—</b></span>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- Debug log -->
|
||||
<div class="card" style="margin-top:16px;max-width:1100px;">
|
||||
<div class="row" style="justify-content:space-between;">
|
||||
<h2 style="margin:0;">Event Log</h2>
|
||||
<button id="btnClearLog" style="padding:4px 12px;font-size:12px;">Clear</button>
|
||||
</div>
|
||||
<div id="log" style="margin-top:10px;"></div>
|
||||
</div>
|
||||
|
||||
<script>
|
||||
// ---------------------------------------------------------------------------
|
||||
// Audio Agent Engine — Playground client
|
||||
//
|
||||
// Pipeline:
|
||||
// mic (AudioContext) -> AudioWorklet (resample to 44.1kHz, 16-bit LE mono)
|
||||
// -> WebSocket binary frames
|
||||
// WebSocket binary frames -> playback queue -> AudioWorklet -> speakers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
const SAMPLE_RATE_OUT = 44100; // server -> client (matches Fish Audio native rate)
|
||||
const SAMPLE_RATE_IN = 44100; // client -> server
|
||||
const FRAME_MS = 20;
|
||||
const FRAME_SAMPLES = SAMPLE_RATE_IN * FRAME_MS / 1000; // 882
|
||||
|
||||
const $ = (id) => document.getElementById(id);
|
||||
|
||||
// ---- DOM refs ----
|
||||
const wsUrlInput = $('wsUrl');
|
||||
const btnConnect = $('btnConnect');
|
||||
const btnDisconnect = $('btnDisconnect');
|
||||
const btnInterrupt = $('btnInterrupt');
|
||||
const connDot = $('connDot');
|
||||
const connText = $('connText');
|
||||
const micMeter = $('micMeter');
|
||||
const outMeter = $('outMeter');
|
||||
const logEl = $('log');
|
||||
const statSent = $('statSent');
|
||||
const statRecv = $('statRecv');
|
||||
const statInRate = $('statInRate');
|
||||
const chkSendSilence = $('chkSendSilence');
|
||||
const btnClearLog = $('btnClearLog');
|
||||
const ttsText = $('ttsText');
|
||||
const btnSpeak = $('btnSpeak');
|
||||
const refAudio = $('refAudio');
|
||||
const refText = $('refText');
|
||||
const btnSetRef = $('btnSetRef');
|
||||
const btnClearRef = $('btnClearRef');
|
||||
const refStatus = $('refStatus');
|
||||
|
||||
// Default WS URL based on current page host.
|
||||
if (!wsUrlInput.value) {
|
||||
const proto = location.protocol === 'https:' ? 'wss:' : 'ws:';
|
||||
wsUrlInput.value = `${proto}//${location.host}/ws`;
|
||||
}
|
||||
|
||||
let ws = null;
|
||||
let audioCtx = null;
|
||||
let micStream = null;
|
||||
let captureNode = null; // AudioWorkletNode sending mic frames
|
||||
let playbackNode = null; // AudioWorkletNode playing received frames
|
||||
let outputGain = null; // master output gain
|
||||
let statSentN = 0, statRecvN = 0;
|
||||
|
||||
function log(msg, cls = 'sys') {
|
||||
const line = document.createElement('div');
|
||||
line.className = 'log-line ' + cls;
|
||||
const ts = new Date().toLocaleTimeString();
|
||||
line.textContent = `[${ts}] ${msg}`;
|
||||
logEl.appendChild(line);
|
||||
logEl.scrollTop = logEl.scrollHeight;
|
||||
}
|
||||
btnClearLog.onclick = () => { logEl.innerHTML = ''; };
|
||||
|
||||
function setConn(state, text) {
|
||||
connDot.className = 'dot ' + (state === 'on' ? 'on' : state === 'err' ? 'err' : '');
|
||||
connText.textContent = text;
|
||||
}
|
||||
|
||||
// AudioWorklet processor source (capture). Resamples from the context rate
|
||||
// to SAMPLE_RATE_IN and emits 16-bit LE PCM frames of FRAME_SAMPLES.
|
||||
const CAPTURE_PROCESSOR = `
|
||||
class CaptureProcessor extends AudioWorkletProcessor {
|
||||
constructor() {
|
||||
super();
|
||||
this.buf = [];
|
||||
this.frameSize = ${FRAME_SAMPLES};
|
||||
this.targetRate = ${SAMPLE_RATE_IN};
|
||||
this.ratio = this.targetRate / sampleRate;
|
||||
this.phase = 0;
|
||||
this.prev = 0;
|
||||
}
|
||||
process(inputs) {
|
||||
const inp = inputs[0][0];
|
||||
if (!inp) return true;
|
||||
let acc = 0;
|
||||
// Downsample via linear interpolation accumulating into frame buffer.
|
||||
for (let i = 0; i < inp.length; i++) {
|
||||
const cur = inp[i];
|
||||
acc += cur * cur;
|
||||
// Walk the output timeline at targetRate.
|
||||
while (this.phase < i + 1) {
|
||||
const frac = this.phase - Math.floor(this.phase);
|
||||
const s = this.prev + (cur - this.prev) * frac;
|
||||
this.buf.push(s);
|
||||
if (this.buf.length >= this.frameSize) {
|
||||
this.flush();
|
||||
}
|
||||
this.phase += this.ratio;
|
||||
}
|
||||
this.prev = cur;
|
||||
}
|
||||
this.phase -= inp.length;
|
||||
// RMS -> meter
|
||||
const rms = Math.sqrt(acc / inp.length);
|
||||
this.port.postMessage({ rms });
|
||||
return true;
|
||||
}
|
||||
flush() {
|
||||
const n = this.frameSize;
|
||||
const buf = new ArrayBuffer(n * 2);
|
||||
const dv = new DataView(buf);
|
||||
for (let i = 0; i < n; i++) {
|
||||
let v = Math.max(-1, Math.min(1, this.buf[i]));
|
||||
dv.setInt16(i * 2, v < 0 ? v * 0x8000 : v * 0x7FFF, true);
|
||||
}
|
||||
this.buf.splice(0, n);
|
||||
this.port.postMessage({ pcm: buf }, [buf]);
|
||||
}
|
||||
}
|
||||
registerProcessor('capture-processor', CaptureProcessor);
|
||||
`;
|
||||
|
||||
// AudioWorklet processor (playback). Pulls samples from a shared ring buffer
|
||||
// filled by the main thread when WebSocket frames arrive.
|
||||
//
|
||||
// A jitter buffer (pre-buffer of ~200 ms) absorbs the uneven arrival of
|
||||
// Fish Audio TTS chunks so the user does not hear brief drop-outs.
|
||||
const PLAYBACK_PROCESSOR = `
|
||||
class PlaybackProcessor extends AudioWorkletProcessor {
|
||||
constructor() {
|
||||
super();
|
||||
this.buf = new Float32Array(48000 * 2); // ~2s ring @ 44.1kHz
|
||||
this.head = 0; // write index
|
||||
this.tail = 0; // read index
|
||||
// Pre-buffer threshold: start playing once we have this many samples.
|
||||
// 500 ms @ 44.1 kHz ≈ 22050 samples. Absorbs WebSocket network jitter.
|
||||
this.preBuffer = Math.floor(sampleRate * 0.50);
|
||||
this.buffering = true;
|
||||
this.underrunSamples = 0; // consecutive silent samples before re-buffering
|
||||
this.underrunThreshold = Math.floor(sampleRate * 0.10); // 100 ms
|
||||
this.port.onmessage = (e) => {
|
||||
const pcm = e.data;
|
||||
const len = pcm.length;
|
||||
for (let i = 0; i < len; i++) {
|
||||
this.buf[this.head] = pcm[i];
|
||||
this.head = (this.head + 1) % this.buf.length;
|
||||
if (this.head === this.tail) this.tail = (this.tail + 1) % this.buf.length;
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
// Number of samples currently available in the ring buffer.
|
||||
available() {
|
||||
return (this.head - this.tail + this.buf.length) % this.buf.length;
|
||||
}
|
||||
|
||||
process(_inputs, outputs) {
|
||||
const out = outputs[0][0];
|
||||
let level = 0;
|
||||
for (let i = 0; i < out.length; i++) {
|
||||
let s = 0;
|
||||
if (this.buffering) {
|
||||
// Wait until we have enough buffered audio before starting playback.
|
||||
if (this.available() >= this.preBuffer) {
|
||||
this.buffering = false;
|
||||
this.underrunSamples = 0;
|
||||
}
|
||||
}
|
||||
if (!this.buffering) {
|
||||
if (this.tail !== this.head) {
|
||||
s = this.buf[this.tail];
|
||||
this.tail = (this.tail + 1) % this.buf.length;
|
||||
this.underrunSamples = 0;
|
||||
} else {
|
||||
// Brief underrun: output silence but keep playing so playback
|
||||
// resumes immediately when the next chunk arrives. Only go back
|
||||
// to buffering if the underrun lasts > 100 ms.
|
||||
this.underrunSamples++;
|
||||
if (this.underrunSamples >= this.underrunThreshold) {
|
||||
this.buffering = true;
|
||||
this.underrunSamples = 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
out[i] = s;
|
||||
level += s * s;
|
||||
}
|
||||
this.port.postMessage({ rms: Math.sqrt(level / out.length) });
|
||||
return true;
|
||||
}
|
||||
}
|
||||
registerProcessor('playback-processor', PlaybackProcessor);
|
||||
`;
|
||||
|
||||
async function startAudio() {
|
||||
audioCtx = new (window.AudioContext || window.webkitAudioContext)({ sampleRate: SAMPLE_RATE_OUT });
|
||||
// Ensure the context is running (browsers suspend it until a user gesture).
|
||||
if (audioCtx.state !== 'running') {
|
||||
try { await audioCtx.resume(); } catch(e) {}
|
||||
}
|
||||
await audioCtx.audioWorklet.addModule(URL.createObjectURL(new Blob([CAPTURE_PROCESSOR], { type: 'application/javascript' })));
|
||||
await audioCtx.audioWorklet.addModule(URL.createObjectURL(new Blob([PLAYBACK_PROCESSOR], { type: 'application/javascript' })));
|
||||
|
||||
// Mic capture
|
||||
micStream = await navigator.mediaDevices.getUserMedia({ audio: { channelCount: 1, echoCancellation: true, noiseSuppression: true } });
|
||||
const source = audioCtx.createMediaStreamSource(micStream);
|
||||
captureNode = new AudioWorkletNode(audioCtx, 'capture-processor');
|
||||
captureNode.port.onmessage = (e) => {
|
||||
if (e.data.rms !== undefined) {
|
||||
micMeter.style.width = Math.min(100, e.data.rms * 200) + '%';
|
||||
} else if (e.data.pcm) {
|
||||
if (ws && ws.readyState === WebSocket.OPEN) {
|
||||
ws.send(e.data.pcm);
|
||||
statSentN++;
|
||||
statSent.textContent = statSentN;
|
||||
}
|
||||
}
|
||||
};
|
||||
source.connect(captureNode);
|
||||
// captureNode needs to be connected to output to be processed? No — connect to a silent gain.
|
||||
const silent = audioCtx.createGain();
|
||||
silent.gain.value = 0;
|
||||
captureNode.connect(silent).connect(audioCtx.destination);
|
||||
|
||||
// Playback with a master gain stage.
|
||||
playbackNode = new AudioWorkletNode(audioCtx, 'playback-processor');
|
||||
playbackNode.port.onmessage = (e) => {
|
||||
if (e.data.rms !== undefined) {
|
||||
outMeter.style.width = Math.min(100, e.data.rms * 200) + '%';
|
||||
}
|
||||
};
|
||||
outputGain = audioCtx.createGain();
|
||||
outputGain.gain.value = 3.0; // boost TTS level (Fish Audio output is somewhat low)
|
||||
playbackNode.connect(outputGain).connect(audioCtx.destination);
|
||||
|
||||
statInRate.textContent = audioCtx.sampleRate + ' Hz';
|
||||
log(`AudioContext started at ${audioCtx.sampleRate} Hz, capture target ${SAMPLE_RATE_IN} Hz`, 'sys');
|
||||
}
|
||||
|
||||
function stopAudio() {
|
||||
try { captureNode?.disconnect(); } catch(e){}
|
||||
try { playbackNode?.disconnect(); } catch(e){}
|
||||
try { outputGain?.disconnect(); } catch(e){}
|
||||
micStream?.getTracks().forEach(t => t.stop());
|
||||
audioCtx?.close();
|
||||
audioCtx = null;
|
||||
micStream = null;
|
||||
captureNode = playbackNode = outputGain = null;
|
||||
}
|
||||
|
||||
// Convert a received 16-bit LE mono PCM ArrayBuffer to Float32 and feed playback.
|
||||
function feedPlayback(arrayBuffer) {
|
||||
const view = new DataView(arrayBuffer);
|
||||
const n = Math.floor(view.byteLength / 2);
|
||||
const f32 = new Float32Array(n);
|
||||
for (let i = 0; i < n; i++) {
|
||||
f32[i] = view.getInt16(i * 2, true) / 32768;
|
||||
}
|
||||
playbackNode?.port.postMessage(f32);
|
||||
statRecvN++;
|
||||
statRecv.textContent = statRecvN;
|
||||
notifyAudioActivity();
|
||||
}
|
||||
|
||||
async function connect() {
|
||||
// Pre-flight: microphone APIs require a *secure context* (HTTPS or localhost).
|
||||
// Over plain HTTP on a non-localhost origin `navigator.mediaDevices` is
|
||||
// undefined and getUserMedia is impossible — fail with a clear hint instead
|
||||
// of a generic error.
|
||||
if (!navigator.mediaDevices?.getUserMedia) {
|
||||
const why = window.isSecureContext
|
||||
? 'navigator.mediaDevices is unavailable in this browser'
|
||||
: `insecure origin (${location.origin}) — the microphone requires HTTPS or localhost`;
|
||||
log('Mic unavailable: ' + why, 'err');
|
||||
log('Hint: open the page over HTTPS, or via a localhost tunnel (ssh -L 8787:localhost:8787).', 'err');
|
||||
setConn('err', 'mic error');
|
||||
return;
|
||||
}
|
||||
try {
|
||||
await startAudio();
|
||||
} catch (e) {
|
||||
log('Mic access failed: ' + (e.name ? e.name + ': ' : '') + (e.message || e), 'err');
|
||||
if (e.name === 'NotAllowedError')
|
||||
log('Hint: allow microphone access for this site (address-bar permission icon).', 'err');
|
||||
else if (e.name === 'NotFoundError')
|
||||
log('Hint: no microphone device was found.', 'err');
|
||||
setConn('err', 'mic error');
|
||||
return;
|
||||
}
|
||||
|
||||
const url = wsUrlInput.value.trim();
|
||||
log('Connecting to ' + url + ' …', 'sys');
|
||||
ws = new WebSocket(url);
|
||||
ws.binaryType = 'arraybuffer';
|
||||
|
||||
ws.onopen = () => {
|
||||
setConn('on', 'connected');
|
||||
log('WebSocket connected', 'out');
|
||||
btnConnect.disabled = true;
|
||||
btnDisconnect.disabled = false;
|
||||
btnInterrupt.disabled = false;
|
||||
btnSpeak.disabled = false;
|
||||
btnSetRef.disabled = false;
|
||||
btnClearRef.disabled = false;
|
||||
};
|
||||
ws.onmessage = (e) => {
|
||||
if (e.data instanceof ArrayBuffer) {
|
||||
feedPlayback(e.data);
|
||||
} else {
|
||||
log('server: ' + e.data, 'in');
|
||||
}
|
||||
};
|
||||
ws.onerror = () => {
|
||||
log('WebSocket error', 'err');
|
||||
setConn('err', 'error');
|
||||
};
|
||||
ws.onclose = () => {
|
||||
log('WebSocket closed', 'sys');
|
||||
setConn('', 'disconnected');
|
||||
btnConnect.disabled = false;
|
||||
btnDisconnect.disabled = true;
|
||||
btnInterrupt.disabled = true;
|
||||
btnSpeak.disabled = true;
|
||||
btnSetRef.disabled = true;
|
||||
btnClearRef.disabled = true;
|
||||
refStatus.querySelector('span').textContent = 'not set';
|
||||
stopAudio();
|
||||
};
|
||||
}
|
||||
|
||||
function disconnect() {
|
||||
try { ws?.close(); } catch(e){}
|
||||
stopAudio();
|
||||
}
|
||||
|
||||
btnConnect.onclick = connect;
|
||||
btnDisconnect.onclick = disconnect;
|
||||
btnInterrupt.onclick = () => {
|
||||
if (ws && ws.readyState === WebSocket.OPEN) {
|
||||
ws.send('INTERRUPT');
|
||||
log('→ sent INTERRUPT (barge-in)', 'out');
|
||||
}
|
||||
};
|
||||
let speakActive = false;
|
||||
let speakIdleTimer = null;
|
||||
|
||||
btnSpeak.onclick = () => {
|
||||
if (speakActive) return;
|
||||
const text = ttsText.value.trim();
|
||||
if (!text) return;
|
||||
if (ws && ws.readyState === WebSocket.OPEN) {
|
||||
speakActive = true;
|
||||
btnSpeak.disabled = true;
|
||||
btnSpeak.textContent = 'Speaking...';
|
||||
ws.send(text);
|
||||
log('→ TTS: "' + text + '"', 'out');
|
||||
}
|
||||
};
|
||||
|
||||
// Re-enable Speak button when the server stops sending audio for >1.5s.
|
||||
function notifyAudioActivity() {
|
||||
if (!speakActive) return;
|
||||
clearTimeout(speakIdleTimer);
|
||||
speakIdleTimer = setTimeout(() => {
|
||||
speakActive = false;
|
||||
btnSpeak.disabled = false;
|
||||
btnSpeak.textContent = 'Speak';
|
||||
}, 1500);
|
||||
}
|
||||
|
||||
btnSetRef.onclick = () => {
|
||||
const path = refAudio.value.trim();
|
||||
const text = refText.value.trim();
|
||||
if (!path) { log('Reference audio path is empty', 'err'); return; }
|
||||
if (ws && ws.readyState === WebSocket.OPEN) {
|
||||
ws.send('SET_REFERENCE:' + path + '|' + text);
|
||||
refStatus.querySelector('span').textContent = text ? 'set (with transcript)' : 'set (no transcript)';
|
||||
log('→ SET_REFERENCE: ' + path + (text ? ' | ' + text.slice(0, 60) : ''), 'out');
|
||||
}
|
||||
};
|
||||
|
||||
btnClearRef.onclick = () => {
|
||||
if (ws && ws.readyState === WebSocket.OPEN) {
|
||||
ws.send('CLEAR_REFERENCE');
|
||||
refStatus.querySelector('span').textContent = 'not set';
|
||||
log('→ CLEAR_REFERENCE', 'out');
|
||||
}
|
||||
};
|
||||
|
||||
// Keep-alive: send silence frames when the mic produces nothing (optional).
|
||||
let keepAliveTimer = null;
|
||||
chkSendSilence.addEventListener('change', () => {
|
||||
if (chkSendSilence.checked) {
|
||||
keepAliveTimer = setInterval(() => {
|
||||
if (ws && ws.readyState === WebSocket.OPEN && audioCtx) {
|
||||
// one silent 20ms frame @ 44.1kHz
|
||||
const buf = new ArrayBuffer(FRAME_SAMPLES * 2);
|
||||
ws.send(buf);
|
||||
}
|
||||
}, 500);
|
||||
} else {
|
||||
clearInterval(keepAliveTimer);
|
||||
}
|
||||
});
|
||||
chkSendSilence.dispatchEvent(new Event('change'));
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
@@ -0,0 +1,42 @@
|
||||
//! WebSocket upgrade handler. Wires an incoming connection to a per-session
|
||||
//! audio loop.
|
||||
|
||||
use crate::config::Config;
|
||||
use crate::llm::LlmClient;
|
||||
use crate::orchestrator::FillerBank;
|
||||
use crate::stt::{ParakeetStt, Tokenizer};
|
||||
use crate::tts::TtsEngine;
|
||||
use crate::vad::SileroVad;
|
||||
use crate::ws::session::run_session;
|
||||
use axum::extract::{State, WebSocketUpgrade};
|
||||
use axum::response::IntoResponse;
|
||||
use std::sync::Arc;
|
||||
|
||||
/// Shared application state passed to every WebSocket handler.
|
||||
#[derive(Clone)]
|
||||
pub struct AppState {
|
||||
pub config: Arc<Config>,
|
||||
pub tts: Arc<dyn TtsEngine>,
|
||||
pub fillers: Arc<FillerBank>,
|
||||
/// Preloaded backchannel clips, played while the user is speaking (see
|
||||
/// [`crate::orchestrator::backchannel`]).
|
||||
pub backchannel: Arc<FillerBank>,
|
||||
/// Optional STT engine. If `None`, the session runs without speech
|
||||
/// recognition (useful for testing the audio pipeline without a model).
|
||||
pub stt: Option<ParakeetStt>,
|
||||
pub tokenizer: Option<Tokenizer>,
|
||||
/// Optional Silero VAD model. `None` (disabled or failed to load) makes the
|
||||
/// session fall back to the RMS voice detection inside the STT worker.
|
||||
pub vad: Option<SileroVad>,
|
||||
/// Optional reasoning LLM. `None` (disabled or misconfigured) makes the
|
||||
/// session run STT → TTS with no reasoning step (see [`crate::llm`]).
|
||||
pub llm: Option<Arc<LlmClient>>,
|
||||
}
|
||||
|
||||
/// `GET /ws` — upgrades the connection to a WebSocket and runs the session.
|
||||
pub async fn ws_handler(
|
||||
ws: WebSocketUpgrade,
|
||||
State(state): State<AppState>,
|
||||
) -> impl IntoResponse {
|
||||
ws.on_upgrade(move |socket| run_session(socket, state))
|
||||
}
|
||||
@@ -0,0 +1,6 @@
|
||||
//! WebSocket subsystem: connection upgrade and per-session audio loop.
|
||||
|
||||
pub mod handler;
|
||||
pub mod session;
|
||||
|
||||
pub use handler::{ws_handler, AppState};
|
||||
@@ -0,0 +1,361 @@
|
||||
//! Per-session audio loop: receives microphone PCM, feeds the STT worker, and
|
||||
//! streams mixed 20 ms frames back to the browser on a strict ticker.
|
||||
|
||||
use crate::audio::AudioMixer;
|
||||
use crate::config::{FRAME_DURATION_MS, OUTPUT_FRAME_BYTES, OUTPUT_SAMPLE_RATE, STT_SAMPLE_RATE};
|
||||
use crate::orchestrator::{spawn_orchestrator, BackchannelScheduler, TurnManager};
|
||||
use crate::stt::{SttWorker, Transcript};
|
||||
use crate::vad::{VadDetector, VadEvent};
|
||||
use crate::ws::handler::AppState;
|
||||
use axum::extract::ws::{Message, WebSocket};
|
||||
use futures_util::{SinkExt, StreamExt};
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::time::{interval, MissedTickBehavior};
|
||||
use tracing::{debug, error, info, warn};
|
||||
|
||||
/// Capacity (in samples) of the mixer queues. ~4 seconds of audio at 44.1 kHz.
|
||||
/// Must hold at least 2 Fish Audio chunks (~2s each) to bridge inter-chunk gaps.
|
||||
const MIXER_QUEUE_CAPACITY: usize = OUTPUT_SAMPLE_RATE as usize * 4;
|
||||
|
||||
/// Run a single WebSocket session to completion.
|
||||
pub async fn run_session(socket: WebSocket, state: AppState) {
|
||||
info!("new WebSocket session started");
|
||||
|
||||
let (mut ws_sender, mut ws_receiver) = socket.split();
|
||||
|
||||
// Per-session mixer and turn manager.
|
||||
let mixer = AudioMixer::new(
|
||||
MIXER_QUEUE_CAPACITY,
|
||||
(OUTPUT_SAMPLE_RATE as u64 * state.config.tts_prebuffer_ms / 1000) as usize,
|
||||
state.config.filler_gain_idle,
|
||||
state.config.filler_gain_duck,
|
||||
state.config.backchannel.gain,
|
||||
);
|
||||
// Outbound text events (chat box): the turn manager pushes JSON strings,
|
||||
// the sender task forwards them on the same WebSocket as the audio frames.
|
||||
let (out_tx, out_rx) = mpsc::unbounded_channel::<String>();
|
||||
let turn_manager = Arc::new(TurnManager::new(
|
||||
mixer.clone(),
|
||||
state.tts.clone(),
|
||||
state.fillers.clone(),
|
||||
state.tokenizer.clone(),
|
||||
state.llm.clone(),
|
||||
state.config.clone(),
|
||||
Some(out_tx),
|
||||
));
|
||||
|
||||
let vad_settings = state.config.vad.clone();
|
||||
let mut vad_detector = state.vad.as_ref().map(|model| {
|
||||
info!(
|
||||
threshold = vad_settings.threshold,
|
||||
neg_threshold = vad_settings.neg_threshold,
|
||||
min_speech_ms = vad_settings.min_speech_ms,
|
||||
min_silence_ms = vad_settings.min_silence_ms,
|
||||
speech_pad_ms = vad_settings.speech_pad_ms,
|
||||
segment_stt = vad_settings.segment_stt,
|
||||
barge_in = vad_settings.barge_in,
|
||||
debug = vad_settings.debug,
|
||||
"VAD active for this session"
|
||||
);
|
||||
VadDetector::new(model.clone(), vad_settings.clone())
|
||||
});
|
||||
|
||||
// STT transcript channel: worker -> orchestrator.
|
||||
let stt_worker = state.stt.as_ref().map(|stt| {
|
||||
let (transcript_tx, transcript_rx) = mpsc::unbounded_channel::<Transcript>();
|
||||
let worker = SttWorker::spawn(stt.clone(), transcript_tx);
|
||||
spawn_orchestrator(turn_manager.clone(), transcript_rx);
|
||||
worker
|
||||
});
|
||||
|
||||
// Channel for passing received mic frames to the STT feed task.
|
||||
let (mic_tx, mut mic_rx) = mpsc::unbounded_channel::<Vec<i16>>();
|
||||
|
||||
// --- Task 1: receiver loop (mic frames + control text) ---
|
||||
let tm_for_interrupt = turn_manager.clone();
|
||||
let receiver_handle = tokio::spawn(async move {
|
||||
while let Some(msg) = ws_receiver.next().await {
|
||||
match msg {
|
||||
Ok(Message::Binary(bytes)) => {
|
||||
// Parse raw 16-bit LE mono PCM.
|
||||
if bytes.len() % 2 != 0 {
|
||||
warn!("odd-length audio frame, dropping trailing byte");
|
||||
}
|
||||
let samples: Vec<i16> = bytes
|
||||
.chunks_exact(2)
|
||||
.map(|c| i16::from_le_bytes([c[0], c[1]]))
|
||||
.collect();
|
||||
// Resample to 16 kHz for STT (handles 16 or 24 kHz input).
|
||||
let resampled = resample_to_16k(&samples, OUTPUT_SAMPLE_RATE);
|
||||
let _ = mic_tx.send(resampled);
|
||||
}
|
||||
Ok(Message::Text(text)) => {
|
||||
handle_control_message(&text, &tm_for_interrupt).await;
|
||||
}
|
||||
Ok(Message::Close(_)) => break,
|
||||
Ok(_) => {}
|
||||
Err(e) => {
|
||||
error!(error = %e, "websocket recv error");
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
info!("websocket receiver ended");
|
||||
});
|
||||
|
||||
// --- Task 2: mic processing (VAD-driven segmentation + STT feed) ---
|
||||
let mic_feed_handle = {
|
||||
let worker_for_mic = stt_worker.clone();
|
||||
let tm_for_vad = turn_manager.clone();
|
||||
let settings = vad_settings.clone();
|
||||
// Backchannel scheduler: emits short acknowledgments while the user
|
||||
// speaks (see `crate::orchestrator::backchannel`).
|
||||
let mut backchannel = BackchannelScheduler::new(
|
||||
state.backchannel.clone(),
|
||||
state.config.backchannel.clone(),
|
||||
);
|
||||
let mixer_for_backchannel = mixer.clone();
|
||||
tokio::spawn(async move {
|
||||
let mut frames: u64 = 0;
|
||||
let mut samples: u64 = 0;
|
||||
|
||||
while let Some(frame) = mic_rx.recv().await {
|
||||
frames += 1;
|
||||
samples += frame.len() as u64;
|
||||
|
||||
let Some(detector) = vad_detector.as_mut() else {
|
||||
// No VAD: the STT worker segments with its RMS detector; if
|
||||
// STT is disabled we only report mic activity, one line per
|
||||
// ~1 s (50 frames of 20 ms).
|
||||
match &worker_for_mic {
|
||||
Some(worker) => worker.push_frame(frame),
|
||||
None if frames % 50 == 0 => info!(
|
||||
mic_frames = frames,
|
||||
mic_samples = samples,
|
||||
"mic audio received (STT disabled — not transcribed)"
|
||||
),
|
||||
None => {}
|
||||
}
|
||||
continue;
|
||||
};
|
||||
|
||||
for event in detector.push(&frame) {
|
||||
match event {
|
||||
VadEvent::SpeechOnset { prob } => {
|
||||
// The user started talking: cut the agent off unless
|
||||
// it is already silent.
|
||||
if settings.barge_in && tm_for_vad.is_speaking() {
|
||||
info!(
|
||||
prob = format!("{prob:.3}"),
|
||||
"VAD barge-in — interrupting agent"
|
||||
);
|
||||
tm_for_vad.interrupt();
|
||||
}
|
||||
}
|
||||
VadEvent::SpeechStart { prob } => {
|
||||
debug!(
|
||||
prob = format!("{prob:.3}"),
|
||||
"VAD speech confirmed (min speech reached)"
|
||||
);
|
||||
backchannel.on_speech_start();
|
||||
}
|
||||
VadEvent::SpeechEnd {
|
||||
audio,
|
||||
duration_ms,
|
||||
avg_prob,
|
||||
} => {
|
||||
backchannel.on_speech_end();
|
||||
if settings.segment_stt {
|
||||
match &worker_for_mic {
|
||||
Some(worker) => {
|
||||
info!(
|
||||
duration_ms,
|
||||
samples = audio.len(),
|
||||
avg_prob = format!("{avg_prob:.3}"),
|
||||
"VAD utterance → STT"
|
||||
);
|
||||
worker.push_utterance(audio);
|
||||
}
|
||||
None => info!(
|
||||
duration_ms,
|
||||
"VAD utterance detected (STT disabled — not transcribed)"
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// While the user keeps speaking and the agent is silent, emit a
|
||||
// brief acknowledgment at most at the configured intervals. The
|
||||
// `!is_speaking()` gate avoids cuing on top of the agent's own
|
||||
// TTS (that would be an interruption, not a backchannel).
|
||||
if detector.in_speech() && !tm_for_vad.is_speaking() {
|
||||
if let Some(name) = backchannel.maybe_emit(&mixer_for_backchannel) {
|
||||
info!(clip = %name, "backchannel emitted during user speech");
|
||||
}
|
||||
}
|
||||
|
||||
// With VAD segmentation the worker must not also receive the raw
|
||||
// frames, or the audio would be segmented twice.
|
||||
if !settings.segment_stt {
|
||||
if let Some(worker) = &worker_for_mic {
|
||||
worker.push_frame(frame);
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
};
|
||||
|
||||
// --- Task 3: 20 ms sender ticker (mixed audio -> browser) ---
|
||||
let sender_handle = tokio::spawn(async move {
|
||||
let mut ticker = interval(std::time::Duration::from_millis(FRAME_DURATION_MS));
|
||||
// Burst: if ticks are missed (event loop busy), catch up by sending
|
||||
// multiple frames back-to-back. The browser's jitter buffer absorbs
|
||||
// the burst. Skip would drop frames, causing audible micro-gaps.
|
||||
ticker.set_missed_tick_behavior(MissedTickBehavior::Burst);
|
||||
|
||||
// Text events (chat box) are interleaved with the audio frames on the
|
||||
// same socket: they are rare and small, so a `select!` is enough. The
|
||||
// branch is disabled once the channel closes to avoid a busy loop.
|
||||
let mut out_rx = out_rx;
|
||||
let mut events_open = true;
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = ticker.tick() => {
|
||||
let frame = mixer
|
||||
.pull_frame()
|
||||
.unwrap_or_else(|| vec![0u8; OUTPUT_FRAME_BYTES]);
|
||||
if ws_sender.send(Message::Binary(frame)).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
event = out_rx.recv(), if events_open => {
|
||||
match event {
|
||||
Some(text) => {
|
||||
if ws_sender.send(Message::Text(text)).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
None => events_open = false,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
info!("websocket sender ended");
|
||||
});
|
||||
|
||||
// Wait for the receiver to finish (client disconnect), then clean up.
|
||||
let _ = receiver_handle.await;
|
||||
// Cancel any in-flight TTS and flush audio so the synthesis task stops
|
||||
// consuming resources after the client disconnects.
|
||||
turn_manager.interrupt();
|
||||
// Dropping the sender will close the socket; abort the other tasks.
|
||||
sender_handle.abort();
|
||||
mic_feed_handle.abort();
|
||||
|
||||
info!("session ended");
|
||||
}
|
||||
|
||||
/// Handle inbound text messages: control signals or text-to-speak.
|
||||
///
|
||||
/// JSON control messages (used by the `/mivocero` chat client):
|
||||
/// - `{"type":"user_text","text":"…"}`: run a user turn through the LLM.
|
||||
/// - `{"type":"interrupt"}`: flush audio and cancel the in-flight reply.
|
||||
///
|
||||
/// Legacy plain-text controls (used by the playground):
|
||||
/// - `INTERRUPT` / `BAARGEIN`: flush audio and cancel TTS (barge-in).
|
||||
/// - `SET_REFERENCE:<path>|<text>`: set the voice-cloning reference.
|
||||
/// - `CLEAR_REFERENCE`: remove the voice-cloning reference.
|
||||
/// - Any other text: synthesize it as speech (TTS pipeline test).
|
||||
async fn handle_control_message(text: &str, turn_manager: &Arc<TurnManager>) {
|
||||
let trimmed = text.trim();
|
||||
|
||||
if trimmed.starts_with('{') {
|
||||
if let Ok(value) = serde_json::from_str::<serde_json::Value>(trimmed) {
|
||||
match value["type"].as_str() {
|
||||
Some("user_text") => {
|
||||
if let Some(user_text) = value["text"].as_str() {
|
||||
info!(text = user_text, "client text turn");
|
||||
// Run the turn in its own task: awaiting it here would
|
||||
// block the receive loop for the whole LLM reply and
|
||||
// make the session deaf to `interrupt` messages.
|
||||
let tm = turn_manager.clone();
|
||||
let user_text = user_text.to_string();
|
||||
tokio::spawn(async move {
|
||||
tm.on_user_text(&user_text).await;
|
||||
});
|
||||
}
|
||||
return;
|
||||
}
|
||||
Some("interrupt") => {
|
||||
info!("client requested interrupt (json)");
|
||||
turn_manager.interrupt();
|
||||
return;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if trimmed == "INTERRUPT" || trimmed == "interrupt" || trimmed == "BAARGEIN" || trimmed == "bargein" {
|
||||
info!("client requested interrupt");
|
||||
turn_manager.interrupt();
|
||||
return;
|
||||
}
|
||||
if trimmed == "CLEAR_REFERENCE" || trimmed == "clear_reference" {
|
||||
info!("client cleared voice reference");
|
||||
turn_manager.set_reference(None, None);
|
||||
return;
|
||||
}
|
||||
if let Some(rest) = trimmed.strip_prefix("SET_REFERENCE:") {
|
||||
let (path, text) = match rest.split_once('|') {
|
||||
Some((p, t)) => (p.trim().to_string(), Some(t.to_string())),
|
||||
None => (rest.trim().to_string(), None),
|
||||
};
|
||||
if path.is_empty() {
|
||||
turn_manager.set_reference(None, None);
|
||||
} else {
|
||||
info!(audio_path = %path, "client set voice reference");
|
||||
turn_manager.set_reference(Some(path), text);
|
||||
}
|
||||
return;
|
||||
}
|
||||
if !trimmed.is_empty() {
|
||||
info!(text = trimmed, "client text-to-speech request");
|
||||
turn_manager.speak(trimmed).await;
|
||||
}
|
||||
}
|
||||
|
||||
/// Resample a mono `i16` PCM buffer from `input_rate` to 16 kHz.
|
||||
///
|
||||
/// Uses linear interpolation. Handles the common 16 kHz and 24 kHz cases
|
||||
/// efficiently. For other rates, falls back to a generic ratio interpolator.
|
||||
fn resample_to_16k(samples: &[i16], input_rate: u32) -> Vec<i16> {
|
||||
if input_rate == STT_SAMPLE_RATE || samples.is_empty() {
|
||||
return samples.to_vec();
|
||||
}
|
||||
if input_rate == 0 {
|
||||
return samples.to_vec();
|
||||
}
|
||||
|
||||
let ratio = STT_SAMPLE_RATE as f64 / input_rate as f64;
|
||||
let out_len = ((samples.len() as f64) * ratio).round() as usize;
|
||||
if out_len == 0 {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
let mut out = Vec::with_capacity(out_len);
|
||||
for i in 0..out_len {
|
||||
let src_pos = i as f64 / ratio;
|
||||
let idx = src_pos.floor() as usize;
|
||||
let frac = src_pos - idx as f64;
|
||||
let a = samples[idx.min(samples.len() - 1)] as f32;
|
||||
let b = samples[(idx + 1).min(samples.len() - 1)] as f32;
|
||||
let v = a + (b - a) * frac as f32;
|
||||
out.push(v.round() as i16);
|
||||
}
|
||||
out
|
||||
}
|
||||
Reference in New Issue
Block a user