osobhandClaude Opus 4.7
fff1b7acd5
The canonical voice loop now lives in zeroclaw-channel-voice (`~/projects/zeroclaw/crates/zeroclaw-channel-voice`, binary `voice_server`). It routes the LLM path through zeroclaw's agent runtime — multi-turn history, tools, memory, provider routing — instead of the OpenAI-compatible direct path here. Same WS wire protocol so `examples/converse_client.rs` drives both; no client-side migration needed. This binary is intentionally kept buildable for: 1. Reproducing perf_history.md Phase 8.10 benches. 2. Standalone (no-agent) use when zeroclaw isn't desired. Module doc + main() startup banner updated to point at the new home. Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
1854 lines
78 KiB
Rust
1854 lines
78 KiB
Rust
//! Rust Unmute MVP — WebSocket conversational server.
|
||
//!
|
||
//! ## ⚠ Superseded by zeroclaw-channel-voice
|
||
//!
|
||
//! The canonical voice-loop server lives at
|
||
//! `~/projects/zeroclaw/crates/zeroclaw-channel-voice` (binary
|
||
//! `voice_server`). It routes the LLM path through the zeroclaw agent
|
||
//! runtime — multi-turn history, tools, memory, provider routing —
|
||
//! instead of the OpenAI-compatible direct path used here. Same WS
|
||
//! wire protocol so `examples/converse_client.rs` drives both.
|
||
//!
|
||
//! This binary is kept for two reasons:
|
||
//! 1. Reproducibility of the perf benches in `docs/perf_history.md`
|
||
//! Phase 8.10 (Moonshine TTFA -80 %, mock-LLM 3-turn loop, etc.).
|
||
//! 2. Standalone use when zeroclaw isn't desired (one-shot, no agent).
|
||
//!
|
||
//! For new conversational deployments, prefer:
|
||
//!
|
||
//! ```bash
|
||
//! cd ~/projects/zeroclaw
|
||
//! cargo run -p zeroclaw-channel-voice --release --bin voice_server -- \
|
||
//! --bind 127.0.0.1:18090
|
||
//! ```
|
||
//!
|
||
//! ## Original behavior
|
||
//!
|
||
//! Loads CSM-1B + Kyutai STT + an OpenAI-compatible LLM at startup.
|
||
//! Each WebSocket connection is a turn-based voice conversation:
|
||
//!
|
||
//! ```text
|
||
//! Client -> Server binary frames: 16-bit LE PCM @ 24 kHz mono
|
||
//! Client -> Server text "EOT" : signal end-of-turn
|
||
//! Server -> Client text {"event":"transcript","text":"..."}
|
||
//! Server -> Client binary frames: 16-bit LE PCM @ 24 kHz mono
|
||
//! Server -> Client text {"event":"done","assistant":"..."}
|
||
//! ```
|
||
//!
|
||
//! Conversation history is kept per-connection. Multiple turns supported
|
||
//! over a single socket; each turn ends when the client sends "EOT".
|
||
//!
|
||
//! ## Status
|
||
//!
|
||
//! Half-duplex turn-based. Auto end-of-turn (semantic VAD via the
|
||
//! Kyutai 1B en/fr `extra_heads` outputs) is the obvious 6c.3 follow-up.
|
||
//! Barge-in (user interrupts assistant) is also deferred.
|
||
//!
|
||
//! ## Usage
|
||
//!
|
||
//! ```bash
|
||
//! export OPENAI_API_KEY=...
|
||
//! cargo run -p rtx-csm --release --features metal --example converse_server -- \
|
||
//! --bind 127.0.0.1:18090 \
|
||
//! --llm-base "https://api.openai.com/v1" --llm-model gpt-4o-mini \
|
||
//! --system "You are a concise voice assistant. Keep answers under 2 sentences."
|
||
//! ```
|
||
//!
|
||
//! Drive it with `examples/converse_client.rs`.
|
||
|
||
// Without the `asr` feature, AsrEngine has only the Kyutai variant so
|
||
// `if let AsrEngine::Kyutai(_) = ...` patterns are technically
|
||
// irrefutable. With `asr`, they're refutable. Silence the warning so the
|
||
// code compiles cleanly under both cfgs.
|
||
#![allow(irrefutable_let_patterns)]
|
||
|
||
use anyhow::{Context, Result};
|
||
use async_trait::async_trait;
|
||
use axum::{
|
||
Router,
|
||
extract::{
|
||
State, WebSocketUpgrade,
|
||
ws::{Message, WebSocket},
|
||
},
|
||
http::{HeaderMap, StatusCode},
|
||
response::IntoResponse,
|
||
routing::get,
|
||
};
|
||
use clap::Parser;
|
||
use futures_util::{StreamExt, stream};
|
||
use rtx_csm::{
|
||
GenerateOptions, Generator,
|
||
converse::{Converse, ConverseOptions, FlushPolicy},
|
||
error::Result as CsmResult,
|
||
llm_client::{ChatMessage, GenConfig, LlmClient, OpenAiCompatibleClient, TokenStream},
|
||
stt::{ASR_DELAY_FRAMES, AsrEvent, SAMPLE_RATE as STT_SR, Stt},
|
||
};
|
||
use std::net::SocketAddr;
|
||
use std::path::PathBuf;
|
||
use std::sync::Arc;
|
||
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
|
||
use tokio::sync::Mutex;
|
||
|
||
const PCM_RATE: u32 = 24_000;
|
||
|
||
#[derive(Debug, Parser)]
|
||
struct Cli {
|
||
#[arg(long, default_value = "127.0.0.1:18090")]
|
||
bind: SocketAddr,
|
||
#[arg(long, default_value = "https://api.openai.com/v1")]
|
||
llm_base: String,
|
||
#[arg(long, default_value = "gpt-4o-mini")]
|
||
llm_model: String,
|
||
#[arg(long)]
|
||
llm_api_key: Option<String>,
|
||
/// Skip the real LLM; canned echo responses for end-to-end testing
|
||
/// without an API key.
|
||
#[arg(long)]
|
||
mock_llm: bool,
|
||
#[arg(
|
||
long,
|
||
default_value = "You are a concise voice assistant. Reply in one or two short sentences."
|
||
)]
|
||
system: String,
|
||
#[arg(long, default_value_t = 0)]
|
||
speaker: u32,
|
||
#[arg(long)]
|
||
cpu: bool,
|
||
|
||
/// Optional emotion control tag prepended to every TTS turn's text, e.g.
|
||
/// `--emotion-hint "[whisper]"`. Only meaningful when the loaded LoRA
|
||
/// adapter was fine-tuned with matching tags (see Phase 12.2 +
|
||
/// `docs/personal_voice_training_guide.md`); on the un-adapted base it
|
||
/// just adds extra cosmetic prefix tokens. Acts as a fallback when
|
||
/// `--reactive-emotion` is also set and detection produces no label.
|
||
#[arg(long)]
|
||
emotion_hint: Option<String>,
|
||
|
||
/// Detect the user's emotion from each incoming audio buffer and use
|
||
/// the resulting tag as the per-turn emotion_hint. Default classifier
|
||
/// is the Phase 13.3 prosody-rule placeholder; pair with
|
||
/// `--use-emotion2vec` for the real emotion2vec_plus_base classifier
|
||
/// (Phase 13.8). Static `--emotion-hint` is the fallback when
|
||
/// detection returns Neutral.
|
||
#[arg(long, default_value_t = false)]
|
||
reactive_emotion: bool,
|
||
|
||
/// With `--reactive-emotion`, swap the prosody-rule classifier for
|
||
/// the emotion2vec_plus_base candle port. Adds a one-time ~1.1 GB
|
||
/// download on first run and ~150 ms per turn vs <1 ms for prosody
|
||
/// rules, in exchange for SOTA-comparable accuracy.
|
||
#[arg(long, default_value_t = false)]
|
||
use_emotion2vec: bool,
|
||
|
||
/// When `--reactive-emotion` produces a non-Neutral label, ALSO append
|
||
/// a per-turn signal to the LLM-facing user message so the LLM's
|
||
/// response text adapts (not just TTS prosody). Per turn only — the
|
||
/// annotation is NOT pushed to persistent history, so subsequent
|
||
/// turns aren't biased by stale signals. Has no effect without
|
||
/// `--reactive-emotion`.
|
||
#[arg(long, default_value_t = false)]
|
||
emotion_aware_llm: bool,
|
||
|
||
/// Bearer token required on the WebSocket Authorization header. If
|
||
/// unset the server is open (suitable for local dev only). Reads
|
||
/// RTX_AUTH_TOKEN env var if not provided.
|
||
#[arg(long)]
|
||
auth_token: Option<String>,
|
||
|
||
/// Per-connection rate limit: max user audio seconds per 60-second
|
||
/// rolling window. Excess audio frames are rejected with an error
|
||
/// event and the connection closes. 0 disables.
|
||
#[arg(long, default_value_t = 600)]
|
||
rate_audio_secs_per_min: u32,
|
||
|
||
/// Per-connection rate limit: max turns per 60-second rolling window.
|
||
/// 0 disables.
|
||
#[arg(long, default_value_t = 60)]
|
||
rate_turns_per_min: u32,
|
||
|
||
/// Optional path to a quantized Q8/Q4_K_M GGUF for CSM-1B (output of
|
||
/// `examples/quantize`). When set, loads the quantized TTS model
|
||
/// instead of FP — typically ~3× faster on Metal at minimal quality
|
||
/// loss. Mimi codec stays FP regardless.
|
||
#[arg(long)]
|
||
quantized_gguf: Option<PathBuf>,
|
||
|
||
/// LoRA voice adapter (safetensors, output of `examples/lora_train` or
|
||
/// `examples/lora_train_emotional`). Loaded into the backbone (FP or
|
||
/// quantized) before the conversation loop starts so every assistant
|
||
/// utterance uses the trained voice / prosody. Combines cleanly with
|
||
/// `--quantized-gguf` (Phase 12.6): the adapter rides on top of the Q8
|
||
/// base in F32 — Phase 12.1's `csm_quantized` LoRA hooks make this work
|
||
/// with the same shared `apply_lora_adapter` helper used for the FP path.
|
||
#[arg(long)]
|
||
lora: Option<PathBuf>,
|
||
/// LoRA rank override. When omitted, auto-detected from adapter metadata
|
||
/// (Phase 12.5); falls back to 8 for older adapters.
|
||
#[arg(long)]
|
||
lora_rank: Option<usize>,
|
||
/// LoRA alpha override. Auto-detected from metadata when omitted.
|
||
#[arg(long)]
|
||
lora_alpha: Option<f32>,
|
||
/// Force extended LoRA coverage (q+k+v+output_proj + MLP). Auto-detected
|
||
/// from metadata when omitted; setting it on a classic q+v adapter just
|
||
/// allocates extra unused B=0 slots.
|
||
#[arg(long, default_value_t = false)]
|
||
extended_lora: bool,
|
||
|
||
/// Path to converted AudioSeal generator safetensors (run
|
||
/// `examples/audioseal_convert` first). When set together with
|
||
/// `--watermark-detector`, every assistant utterance is watermarked
|
||
/// for provenance before being sent over the WebSocket.
|
||
#[arg(long)]
|
||
watermark_generator: Option<PathBuf>,
|
||
/// Path to converted AudioSeal detector safetensors. Required for the
|
||
/// watermarker even when only embedding (the detector is part of
|
||
/// AudioSealWatermarker construction).
|
||
#[arg(long)]
|
||
watermark_detector: Option<PathBuf>,
|
||
/// 16-bit watermark message (decimal or 0xHEX). Defaults to 0.
|
||
/// Used by both AudioSeal (`--watermark-{generator,detector}`)
|
||
/// and SilentCipher (`--watermark-silentcipher`).
|
||
#[arg(long, default_value = "0")]
|
||
watermark_message: String,
|
||
|
||
/// Use Sesame's actual production watermarker — SilentCipher —
|
||
/// instead of AudioSeal. Pass a directory containing the three
|
||
/// released checkpoint files (`enc_c.ckpt`, `dec_c.ckpt`,
|
||
/// `dec_m_0.ckpt`). The 16 kHz model is the one we wire here;
|
||
/// the 44.1 kHz model is unsupported.
|
||
///
|
||
/// If you don't have a local copy, leave this empty and pass
|
||
/// the magic value `hf` to download from `sony/silentcipher`
|
||
/// at boot.
|
||
///
|
||
/// Mutex with `--watermark-generator` / `--watermark-detector`
|
||
/// (AudioSeal). Pick one watermarker.
|
||
#[arg(long)]
|
||
watermark_silentcipher: Option<String>,
|
||
|
||
/// Use the VAD-enabled STT variant (`kyutai/stt-1b-en_fr-candle`).
|
||
/// Adds 4 extra prediction heads emitting per-frame probabilities;
|
||
/// /v1/converse uses head-2 probability > `--vad-threshold` for K
|
||
/// consecutive frames as auto-end-of-turn (no client "EOT" needed).
|
||
#[arg(long)]
|
||
vad: bool,
|
||
/// VAD probability threshold for auto end-of-turn (default 0.5).
|
||
#[arg(long, default_value_t = 0.5)]
|
||
vad_threshold: f32,
|
||
/// Number of consecutive frames above threshold required to fire
|
||
/// end-of-turn. Higher = less sensitive to noise.
|
||
#[arg(long, default_value_t = 4)]
|
||
vad_consecutive: u32,
|
||
|
||
/// Provider-specific JSON merged into every LLM request body.
|
||
/// Example: `--llm-extra-body '{"thinking":{"type":"disabled"}}'`
|
||
/// disables Z.AI glm-4.5/4.6/4.7 reasoning_content traversal so
|
||
/// `content` emits immediately (the cost is one bypassed thinking
|
||
/// pass; latency drops 30-40s on glm-4.5 voice loops). Validated
|
||
/// against Z.AI's coding endpoint; see
|
||
/// `reference_zai_coding_setup.md` for the full provider matrix.
|
||
#[arg(long)]
|
||
llm_extra_body: Option<String>,
|
||
|
||
/// Stream TTS chunks to the client as they're generated rather than
|
||
/// buffering each full sentence. Cuts time-to-first-audio by ~3-5s
|
||
/// (the first ~320ms chunk leaves the server within ~500-1000ms of
|
||
/// the LLM's first sentence boundary). Tradeoff: post-process
|
||
/// (HPF/LUFS) and watermark are SKIPPED — they require
|
||
/// full-sentence context. Combining `--stream-tts` with
|
||
/// `--watermark-*` is rejected at boot.
|
||
#[arg(long)]
|
||
stream_tts: bool,
|
||
|
||
/// Use Moonshine-tiny (pure candle, English-only) instead of Kyutai
|
||
/// STT 1B for transcription. ~34× faster than realtime on M-series
|
||
/// Metal (10 s of audio → ~300 ms transcribe). Batch-only — fires
|
||
/// the transcript at EOT, not incrementally. Loses semantic VAD.
|
||
///
|
||
/// **Pure-Rust win**: unlike `--whisper`, Moonshine integrates
|
||
/// cleanly without ggml/protobuf runtime conflicts. No feature
|
||
/// flag needed (ships in the default build).
|
||
///
|
||
/// English-only — for multilingual deploys, stick with Kyutai.
|
||
#[arg(long)]
|
||
moonshine: bool,
|
||
|
||
/// Use Whisper-tiny (whisper-rs / whisper.cpp) instead of Kyutai STT
|
||
/// 1B for transcription. ~50× faster on Metal in isolation (10 s of
|
||
/// audio → 0.2 s transcribe) but batch-only — transcript fires at
|
||
/// EOT, not incrementally. Loses semantic VAD (`--vad`).
|
||
///
|
||
/// **WARNING — perf regression**: linking whisper-rs's C++ runtime
|
||
/// into the same binary as candle/CSM costs ~2-3× across all CSM
|
||
/// inference (recv_phase, tts_per_utterance, total_turn) even when
|
||
/// this flag isn't set. We don't yet understand the conflict
|
||
/// (BLAS / threadpool / Metal init contention is suspected). For
|
||
/// production deploys, build WITHOUT `--features asr` and accept
|
||
/// Kyutai STT's 1× realtime cost; or run Whisper as a sidecar
|
||
/// process via IPC (not yet implemented).
|
||
///
|
||
/// Requires the crate built with `--features asr-metal` (or `asr`).
|
||
#[arg(long)]
|
||
whisper: bool,
|
||
/// In streaming mode, the number of 80 ms Mimi frames per chunk.
|
||
/// Default 2 = 160 ms.
|
||
///
|
||
/// Phase 9.2 measurement (Q8 + stream + Moonshine + mock LLM):
|
||
/// chunk_frames=4 llm_to_first_audio 533 ms e2e 939 ms
|
||
/// chunk_frames=2 llm_to_first_audio 374 ms e2e 784 ms ← default
|
||
/// chunk_frames=1 llm_to_first_audio 283 ms e2e 664 ms
|
||
/// Per-utterance TTS gen is comparable across all three (~5.7-5.9 s
|
||
/// for the 4-sentence mock LLM), so smaller chunks don't add
|
||
/// meaningful decode overhead. Drop to 1 for tightest TTFA at the
|
||
/// cost of slightly more per-chunk send overhead.
|
||
#[arg(long, default_value_t = 2)]
|
||
stream_chunk_frames: usize,
|
||
|
||
/// Enable an RMS-amplitude VAD pre-filter on STT work. Each
|
||
/// incoming 200 ms binary chunk gets its RMS computed; chunks
|
||
/// below `--vad-gate-threshold` skip the (expensive) Kyutai
|
||
/// step_pcm call entirely. Pure-Rust energy detector — no external
|
||
/// runtime, no protobuf conflict (the original plan was Silero V5
|
||
/// via `ort`, but `ort` 3.21 protobuf clashes with sentencepiece
|
||
/// 3.14; energy VAD sidesteps that).
|
||
///
|
||
/// Real-world voice-agent audio is 30-50% silence (pauses,
|
||
/// breathing, room tone), giving proportional `recv_phase`
|
||
/// reduction. Quality tradeoff: energy VAD misses quiet speech
|
||
/// (whispering, distant). Acceptable for mic-distance loops.
|
||
///
|
||
/// Orthogonal to `--vad` (Kyutai's semantic-VAD-for-EOT).
|
||
#[arg(long)]
|
||
vad_gate: bool,
|
||
/// RMS threshold for `--vad-gate` in [0, 1] (post-normalization
|
||
/// f32 PCM range). 0.01 default; lower = more chunks pass through
|
||
/// STT (less aggressive). Tune per microphone gain.
|
||
#[arg(long, default_value_t = 0.01)]
|
||
vad_gate_threshold: f32,
|
||
}
|
||
|
||
/// Multi-sentence mock LLM. Reads the most recent user message and
|
||
/// returns a fixed multi-sentence acknowledgment with a small delay
|
||
/// between sentence boundaries so the streaming pipeline runs long
|
||
/// enough to be barge-in-testable.
|
||
struct MockLlm;
|
||
|
||
#[async_trait]
|
||
impl LlmClient for MockLlm {
|
||
async fn generate_stream(
|
||
&self,
|
||
messages: Vec<ChatMessage>,
|
||
_config: GenConfig,
|
||
) -> CsmResult<TokenStream> {
|
||
let user_text = messages
|
||
.iter()
|
||
.rev()
|
||
.find_map(|m| {
|
||
if matches!(m.role, rtx_csm::llm_client::Role::User) {
|
||
Some(m.content.clone())
|
||
} else {
|
||
None
|
||
}
|
||
})
|
||
.unwrap_or_default();
|
||
let response = format!(
|
||
"I heard you. You said: {}. That is interesting. Tell me more about it.",
|
||
user_text
|
||
);
|
||
// Stream tokens one whitespace-split chunk at a time, with a small
|
||
// sleep between chunks so the pipeline takes a few seconds end to
|
||
// end (lets barge-in tests fire mid-stream).
|
||
let chunks: Vec<String> = response
|
||
.split_inclusive(' ')
|
||
.map(|s| s.to_string())
|
||
.collect();
|
||
let s = stream::iter(chunks).then(|s| async move {
|
||
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
|
||
Ok::<_, rtx_csm::CsmError>(s)
|
||
});
|
||
Ok(Box::pin(s))
|
||
}
|
||
}
|
||
|
||
enum AnyLlm {
|
||
Real(OpenAiCompatibleClient),
|
||
Mock(MockLlm),
|
||
}
|
||
|
||
#[async_trait]
|
||
impl LlmClient for AnyLlm {
|
||
async fn generate_stream(
|
||
&self,
|
||
messages: Vec<ChatMessage>,
|
||
config: GenConfig,
|
||
) -> CsmResult<TokenStream> {
|
||
match self {
|
||
AnyLlm::Real(c) => c.generate_stream(messages, config).await,
|
||
AnyLlm::Mock(c) => c.generate_stream(messages, config).await,
|
||
}
|
||
}
|
||
}
|
||
|
||
#[derive(Default)]
|
||
struct Metrics {
|
||
/// Successful turns completed.
|
||
turns_total: AtomicU64,
|
||
/// Errors raised during a turn (auth, STT, LLM, TTS).
|
||
errors_total: AtomicU64,
|
||
/// Connections opened.
|
||
connections_total: AtomicU64,
|
||
/// Connections currently active.
|
||
connections_active: AtomicU64,
|
||
/// Sum of STT latencies in ms (for avg = sum / count).
|
||
stt_latency_ms_sum: AtomicU64,
|
||
stt_latency_ms_count: AtomicU64,
|
||
/// Sum of TTS latencies in ms (per-utterance, summed across utterances).
|
||
tts_latency_ms_sum: AtomicU64,
|
||
tts_latency_ms_count: AtomicU64,
|
||
/// End-to-end turn latency (audio-in → first-audio-out): sum + count.
|
||
e2e_first_ms_sum: AtomicU64,
|
||
e2e_first_ms_count: AtomicU64,
|
||
/// Receive phase: turn start → EOT (or VAD auto-EOT) received.
|
||
recv_phase_ms_sum: AtomicU64,
|
||
recv_phase_ms_count: AtomicU64,
|
||
/// LLM-to-first-audio: from conv_fut start (after STT done) to first
|
||
/// PCM chunk reaching the pump. Bundles LLM stream + first sentence
|
||
/// TTS gen.
|
||
llm_to_first_audio_ms_sum: AtomicU64,
|
||
llm_to_first_audio_ms_count: AtomicU64,
|
||
/// Full conversation phase: conv_fut start → all sentences TTS'd.
|
||
conv_total_ms_sum: AtomicU64,
|
||
conv_total_ms_count: AtomicU64,
|
||
/// Total turn wall time: turn start → done event sent.
|
||
total_turn_ms_sum: AtomicU64,
|
||
total_turn_ms_count: AtomicU64,
|
||
/// Reactive emotion observability (Phase 13.6). Each turn that ran
|
||
/// the detector increments `_calls_total`; the per-label counters
|
||
/// add to one of the five labels. `emotion_aware_llm_applied_total`
|
||
/// counts turns where the LLM-history augmentation actually fired
|
||
/// (only when the label is non-Neutral AND --emotion-aware-llm is on).
|
||
reactive_emotion_calls_total: AtomicU64,
|
||
reactive_emotion_neutral_total: AtomicU64,
|
||
reactive_emotion_calm_total: AtomicU64,
|
||
reactive_emotion_sad_total: AtomicU64,
|
||
reactive_emotion_angry_total: AtomicU64,
|
||
reactive_emotion_excited_total: AtomicU64,
|
||
/// Phase 13.10 raw emotion2vec classes — added when the 9→5 fold
|
||
/// was removed. The original 5 buckets above are still emitted
|
||
/// (prosody-rule SER produces those) so these are emotion2vec-only
|
||
/// in practice.
|
||
reactive_emotion_disgusted_total: AtomicU64,
|
||
reactive_emotion_fearful_total: AtomicU64,
|
||
reactive_emotion_happy_total: AtomicU64,
|
||
reactive_emotion_surprised_total: AtomicU64,
|
||
reactive_emotion_unk_total: AtomicU64,
|
||
emotion_aware_llm_applied_total: AtomicU64,
|
||
}
|
||
|
||
/// ASR backend wrapped behind a single Mutex so the receive loop can
|
||
/// branch on its kind. Whisper variant is feature-gated; without `asr`
|
||
/// at compile time only Kyutai + Moonshine are available.
|
||
///
|
||
/// Moonshine is pure candle (no external runtime), so it ships in the
|
||
/// default build alongside Kyutai. English-only.
|
||
enum AsrEngine {
|
||
Kyutai(Stt),
|
||
Moonshine(MoonshineAsr),
|
||
#[cfg(feature = "asr")]
|
||
Whisper(rtx_csm::asr::WhisperAsr),
|
||
}
|
||
|
||
/// Bundle of Moonshine encoder + decoder + tokenizer + config so the
|
||
/// receive-loop's batch transcribe call has everything it needs.
|
||
struct MoonshineAsr {
|
||
encoder: rtx_csm::moonshine::Encoder,
|
||
decoder: rtx_csm::moonshine::Decoder,
|
||
tokenizer: tokenizers::Tokenizer,
|
||
cfg: rtx_csm::moonshine::MoonshineConfig,
|
||
device: candle_core::Device,
|
||
}
|
||
|
||
impl MoonshineAsr {
|
||
/// Resample 24 kHz user audio to 16 kHz, encode, greedy decode,
|
||
/// detokenize. Returns the transcript text. Batch-only (no
|
||
/// streaming, no per-frame VAD).
|
||
fn transcribe_24k(&self, samples_24k: &[f32]) -> anyhow::Result<String> {
|
||
let samples_16k = rtx_csm::audio_io::resample(samples_24k, 24_000, 16_000)
|
||
.map_err(|e| anyhow::anyhow!("resample 24k -> 16k: {e}"))?;
|
||
let pcm = candle_core::Tensor::from_vec(
|
||
samples_16k.clone(),
|
||
(1, 1, samples_16k.len()),
|
||
&self.device,
|
||
)?;
|
||
let enc = self.encoder.forward(&pcm)?;
|
||
let token_ids = self.decoder.generate_cached(
|
||
&enc,
|
||
&self.cfg,
|
||
self.cfg.max_position_embeddings.min(180),
|
||
)?;
|
||
let text = self
|
||
.tokenizer
|
||
.decode(&token_ids, true)
|
||
.map_err(|e| anyhow::anyhow!("detok: {e}"))?;
|
||
Ok(text)
|
||
}
|
||
}
|
||
|
||
struct Shared {
|
||
generator: Mutex<Generator>,
|
||
stt: Mutex<AsrEngine>,
|
||
llm: AnyLlm,
|
||
system_prompt: String,
|
||
speaker: u32,
|
||
auth_token: Option<String>,
|
||
rate_audio_secs_per_min: u32,
|
||
rate_turns_per_min: u32,
|
||
/// VAD: when Some, the WS receive loop monitors per-frame end-of-turn
|
||
/// probability and auto-fires EOT when prs[2][0] > threshold for
|
||
/// `consecutive` frames in a row.
|
||
vad_eot: Option<(f32, u32)>,
|
||
/// Provider-specific JSON merged into every LLM request body. None
|
||
/// means the OpenAI-schema fields are sent as-is.
|
||
llm_extra_body: serde_json::Map<String, serde_json::Value>,
|
||
/// When true, conv_fut uses Converse::run_streaming for low-TTFA
|
||
/// chunked PCM delivery. Skips post-process + watermark.
|
||
stream_tts: bool,
|
||
/// Mimi frames per stream chunk (4 = 320ms default).
|
||
stream_chunk_frames: usize,
|
||
/// True when the ASR backend is batch-only (Whisper or Moonshine).
|
||
/// The receive loop skips parallel STT in this case and runs one
|
||
/// transcribe call at EOT. (Kyutai uses incremental step_pcm.)
|
||
batch_asr: bool,
|
||
/// Optional VAD gate. When `Some`, each incoming binary chunk is
|
||
/// classified speech/silence before being sent to STT; silence
|
||
/// chunks skip step_pcm. Behind a Mutex because is_speech() takes
|
||
/// `&mut self` to update internal counters.
|
||
vad_gate: Option<Mutex<rtx_csm::stt::VadGate>>,
|
||
/// Optional control-token tag prepended to every TTS turn's text. Only
|
||
/// useful when the loaded LoRA was fine-tuned with the same tag (Phase 12.2
|
||
/// plumbing); on the un-adapted base it is a no-op cosmetic prefix.
|
||
emotion_hint: Option<String>,
|
||
/// Pre-built reactive-emotion classifier shared across turns. Boxed
|
||
/// `dyn` so the same call site handles both ProsodyDetector (cheap,
|
||
/// zero-init) and Emotion2Vec (loaded once on boot, ~150 ms/turn).
|
||
/// `Some` ⇔ `--reactive-emotion` was set; `None` disables the path.
|
||
emotion_detector: Option<Box<dyn rtx_csm::ser::EmotionDetector + Send + Sync>>,
|
||
/// When true (and `reactive_emotion` is also true), append a per-turn
|
||
/// emotion-signal annotation to the LLM-facing user message so the
|
||
/// response text adapts in addition to the TTS prosody. Per-turn,
|
||
/// not pushed to persistent history (Phase 13.5).
|
||
emotion_aware_llm: bool,
|
||
metrics: Metrics,
|
||
}
|
||
|
||
/// Sliding-window rate limiter (60-second window). Stores timestamps
|
||
/// (Instant) of recent events; trims to the window on each `try_consume`.
|
||
struct RateBucket {
|
||
/// (timestamp, weight) for each event in the window.
|
||
events: std::collections::VecDeque<(std::time::Instant, u32)>,
|
||
weight_in_window: u32,
|
||
limit: u32,
|
||
window: std::time::Duration,
|
||
}
|
||
|
||
impl RateBucket {
|
||
fn new(limit: u32) -> Self {
|
||
Self {
|
||
events: std::collections::VecDeque::new(),
|
||
weight_in_window: 0,
|
||
limit,
|
||
window: std::time::Duration::from_secs(60),
|
||
}
|
||
}
|
||
/// Returns true if `weight` units fit within the budget; records the
|
||
/// event if so.
|
||
fn try_consume(&mut self, weight: u32) -> bool {
|
||
if self.limit == 0 {
|
||
return true; // disabled
|
||
}
|
||
let now = std::time::Instant::now();
|
||
while let Some((t, w)) = self.events.front() {
|
||
if now.duration_since(*t) > self.window {
|
||
self.weight_in_window = self.weight_in_window.saturating_sub(*w);
|
||
self.events.pop_front();
|
||
} else {
|
||
break;
|
||
}
|
||
}
|
||
if self.weight_in_window + weight > self.limit {
|
||
return false;
|
||
}
|
||
self.weight_in_window += weight;
|
||
self.events.push_back((now, weight));
|
||
true
|
||
}
|
||
}
|
||
|
||
#[tokio::main]
|
||
async fn main() -> Result<()> {
|
||
tracing_subscriber::fmt().init();
|
||
tracing::warn!(
|
||
"converse_server is superseded by zeroclaw-channel-voice (see \
|
||
~/projects/zeroclaw/crates/zeroclaw-channel-voice). This binary \
|
||
remains for reproducibility of docs/perf_history.md Phase 8.10 \
|
||
benches and standalone (no-agent) use."
|
||
);
|
||
let cli = Cli::parse();
|
||
let llm = if cli.mock_llm {
|
||
tracing::info!("LLM: mock (echoes user transcript)");
|
||
AnyLlm::Mock(MockLlm)
|
||
} else {
|
||
let api_key = cli
|
||
.llm_api_key
|
||
.clone()
|
||
.or_else(|| std::env::var("OPENAI_API_KEY").ok())
|
||
.ok_or_else(|| {
|
||
anyhow::anyhow!("set --llm-api-key or OPENAI_API_KEY (or pass --mock-llm)")
|
||
})?;
|
||
AnyLlm::Real(OpenAiCompatibleClient::new(
|
||
&cli.llm_base,
|
||
api_key,
|
||
&cli.llm_model,
|
||
))
|
||
};
|
||
let device = if cli.cpu {
|
||
candle_core::Device::Cpu
|
||
} else {
|
||
Generator::default_device()?
|
||
};
|
||
tracing::info!("device: {device:?}");
|
||
|
||
let mut generator = if let Some(gguf) = cli.quantized_gguf.as_ref() {
|
||
tracing::info!("loading CSM-1B (quantized, {})...", gguf.display());
|
||
Generator::load_csm_1b_quantized(gguf, &device, /* enable_cfg */ false)?
|
||
} else {
|
||
tracing::info!("loading CSM-1B (FP)...");
|
||
Generator::load_csm_1b(&device)?
|
||
};
|
||
|
||
if let Some(lora_path) = cli.lora.as_ref() {
|
||
let extended_override = if cli.extended_lora { Some(true) } else { None };
|
||
rtx_csm::training::apply_lora_adapter(
|
||
&mut generator,
|
||
lora_path,
|
||
cli.lora_rank,
|
||
cli.lora_alpha,
|
||
extended_override,
|
||
&device,
|
||
)?;
|
||
}
|
||
|
||
if let (Some(gen_path), Some(det_path)) = (
|
||
cli.watermark_generator.as_ref(),
|
||
cli.watermark_detector.as_ref(),
|
||
) {
|
||
let msg_str = cli.watermark_message.trim();
|
||
let message: u16 = if let Some(rest) = msg_str
|
||
.strip_prefix("0x")
|
||
.or_else(|| msg_str.strip_prefix("0X"))
|
||
{
|
||
u16::from_str_radix(rest, 16)?
|
||
} else {
|
||
msg_str.parse::<u16>()?
|
||
};
|
||
let gen_vb = unsafe {
|
||
candle_nn::VarBuilder::from_mmaped_safetensors(
|
||
&[gen_path],
|
||
candle_core::DType::F32,
|
||
&device,
|
||
)
|
||
}?;
|
||
let det_vb = unsafe {
|
||
candle_nn::VarBuilder::from_mmaped_safetensors(
|
||
&[det_path],
|
||
candle_core::DType::F32,
|
||
&device,
|
||
)
|
||
}?;
|
||
let inner = rtx_csm::AudioSealWatermarker::from_var_builders(
|
||
gen_vb,
|
||
det_vb,
|
||
device.clone(),
|
||
message,
|
||
)?;
|
||
// CSM is 24 kHz; AudioSeal is 16 kHz. Wrap so the watermarker
|
||
// resamples around its forward pass.
|
||
let wm = rtx_csm::ResampledWatermarker::new(inner, generator.config.sample_rate, 16_000);
|
||
generator.set_watermarker(Box::new(wm));
|
||
tracing::info!(
|
||
"watermarker installed (message=0x{:04X}, model 16 kHz, output {} Hz)",
|
||
message,
|
||
generator.config.sample_rate,
|
||
);
|
||
} else if cli.watermark_generator.is_some() || cli.watermark_detector.is_some() {
|
||
anyhow::bail!(
|
||
"--watermark-generator and --watermark-detector must be set together \
|
||
(AudioSealWatermarker requires both)."
|
||
);
|
||
}
|
||
|
||
// SilentCipher (Sesame's actual production watermarker). Mutex
|
||
// with --watermark-generator/-detector (AudioSeal): the Generator
|
||
// can only carry one watermarker.
|
||
if cli.watermark_silentcipher.is_some()
|
||
&& (cli.watermark_generator.is_some() || cli.watermark_detector.is_some())
|
||
{
|
||
anyhow::bail!(
|
||
"--watermark-silentcipher and --watermark-generator/-detector are mutually \
|
||
exclusive — pick one watermarker (Generator only carries one)."
|
||
);
|
||
}
|
||
if let Some(silent_path) = cli.watermark_silentcipher.as_ref() {
|
||
let msg_str = cli.watermark_message.trim();
|
||
let message: u16 = if let Some(rest) = msg_str
|
||
.strip_prefix("0x")
|
||
.or_else(|| msg_str.strip_prefix("0X"))
|
||
{
|
||
u16::from_str_radix(rest, 16)?
|
||
} else {
|
||
msg_str.parse::<u16>()?
|
||
};
|
||
// Resolve checkpoint paths: either an `hf` magic value (download
|
||
// from sony/silentcipher), or a local directory containing the
|
||
// three .ckpt files.
|
||
let (enc_c, dec_c, dec_m_0) = if silent_path == "hf" {
|
||
tracing::info!("downloading SilentCipher 16 kHz from sony/silentcipher...");
|
||
let api = hf_hub::api::sync::Api::new()?;
|
||
let repo = api.model("sony/silentcipher".to_string());
|
||
let dir = "16_khz/97561_iteration";
|
||
let enc = repo.get(&format!("{dir}/enc_c.ckpt"))?;
|
||
let dec = repo.get(&format!("{dir}/dec_c.ckpt"))?;
|
||
let dem = repo.get(&format!("{dir}/dec_m_0.ckpt"))?;
|
||
(enc, dec, dem)
|
||
} else {
|
||
let dir = std::path::PathBuf::from(silent_path);
|
||
(
|
||
dir.join("enc_c.ckpt"),
|
||
dir.join("dec_c.ckpt"),
|
||
dir.join("dec_m_0.ckpt"),
|
||
)
|
||
};
|
||
let cfg_sc = rtx_csm::silentcipher::SilentCipherConfig::sixteen_khz();
|
||
let inner = rtx_csm::silentcipher::SilentCipherWatermarker::from_ckpts(
|
||
cfg_sc, &enc_c, &dec_c, &dec_m_0, &device,
|
||
)?;
|
||
let wm = rtx_csm::silentcipher::SilentCipherWatermark::new(inner, message as u32);
|
||
// CSM is 24 kHz; SilentCipher 16 kHz. Wrap with the existing
|
||
// ResampledWatermarker so callers feed 24 kHz directly.
|
||
let wrapped = rtx_csm::ResampledWatermarker::new(wm, generator.config.sample_rate, 16_000);
|
||
generator.set_watermarker(Box::new(wrapped));
|
||
tracing::info!(
|
||
"SilentCipher watermarker installed (message=0x{:04X}, model 16 kHz, output {} Hz)",
|
||
message,
|
||
generator.config.sample_rate,
|
||
);
|
||
}
|
||
|
||
if cli.stream_tts && generator.watermarker.is_some() {
|
||
anyhow::bail!(
|
||
"--stream-tts cannot combine with --watermark-* (watermark needs \
|
||
full-sentence context). Pick one: low-latency streaming OR \
|
||
provenance-marked audio."
|
||
);
|
||
}
|
||
|
||
if cli.moonshine && cli.whisper {
|
||
anyhow::bail!("--moonshine and --whisper are mutually exclusive (pick one ASR backend).");
|
||
}
|
||
if cli.moonshine && cli.vad {
|
||
anyhow::bail!(
|
||
"--moonshine and --vad cannot combine: Moonshine has no per-frame probability \
|
||
output."
|
||
);
|
||
}
|
||
let stt: AsrEngine = if cli.moonshine {
|
||
tracing::info!("loading Moonshine-tiny (pure candle, English-only)...");
|
||
let api = hf_hub::api::sync::Api::new()?;
|
||
let repo = api.model("UsefulSensors/moonshine-tiny".to_string());
|
||
let weights = repo.get("model.safetensors")?;
|
||
let tok_path = repo.get("tokenizer.json")?;
|
||
let cfg = rtx_csm::moonshine::MoonshineConfig::tiny();
|
||
let (encoder, decoder) = rtx_csm::moonshine::load_full(&weights, &device, &cfg)?;
|
||
let tokenizer = rtx_csm::moonshine::load_tokenizer(&tok_path)
|
||
.map_err(|e| anyhow::anyhow!("moonshine tokenizer: {e}"))?;
|
||
AsrEngine::Moonshine(MoonshineAsr {
|
||
encoder,
|
||
decoder,
|
||
tokenizer,
|
||
cfg,
|
||
device: device.clone(),
|
||
})
|
||
} else if cli.whisper {
|
||
if cli.vad {
|
||
anyhow::bail!(
|
||
"--whisper and --vad cannot combine: Whisper has no per-frame probability \
|
||
output. Pick one."
|
||
);
|
||
}
|
||
#[cfg(feature = "asr")]
|
||
{
|
||
tracing::info!("loading Whisper-tiny (whisper-rs)...");
|
||
AsrEngine::Whisper(rtx_csm::asr::WhisperAsr::load_default()?)
|
||
}
|
||
#[cfg(not(feature = "asr"))]
|
||
{
|
||
anyhow::bail!(
|
||
"--whisper requires the `asr-metal` (or `asr`) feature. Rebuild with \
|
||
`cargo build --features metal,asr-metal --example converse_server`."
|
||
);
|
||
}
|
||
} else if cli.vad {
|
||
tracing::info!("loading Kyutai STT 1B en/fr-candle with VAD (~3 GB)...");
|
||
AsrEngine::Kyutai(Stt::load_default_with_vad(&device)?)
|
||
} else {
|
||
tracing::info!("loading Kyutai STT 1B en/fr (~3 GB)...");
|
||
AsrEngine::Kyutai(Stt::load_default(&device)?)
|
||
};
|
||
tracing::info!("models loaded");
|
||
|
||
// Warm up the Generator: a longer throwaway generate() so the full
|
||
// Metal kernel set (backbone + depth decoder + Mimi decode + every
|
||
// sampler path) compiles up-front. Empirically eliminates ~4-6 s
|
||
// from the first-sentence TTS gen latency. Cost is paid once at
|
||
// boot. Phase 8.2 raised this from 200 ms to 2000 ms (~25 frames)
|
||
// because at 200 ms only the early code paths fired — variable
|
||
// first-sentence cost still showed up on first user turn.
|
||
let warmup_t = std::time::Instant::now();
|
||
{
|
||
let opts = GenerateOptions {
|
||
max_audio_ms: 2000,
|
||
seed: 0,
|
||
..GenerateOptions::default()
|
||
};
|
||
match generator.generate(
|
||
"Warming up the speech model with a slightly longer prompt.",
|
||
cli.speaker,
|
||
&[],
|
||
opts,
|
||
) {
|
||
Ok(_) => {
|
||
tracing::info!(
|
||
"generator warm-up complete ({}ms)",
|
||
warmup_t.elapsed().as_millis()
|
||
);
|
||
}
|
||
Err(e) => {
|
||
tracing::warn!(
|
||
"generator warm-up failed (non-fatal, first turn will be slower): {e}"
|
||
);
|
||
}
|
||
}
|
||
}
|
||
|
||
let auth_token = cli
|
||
.auth_token
|
||
.or_else(|| std::env::var("RTX_AUTH_TOKEN").ok());
|
||
if auth_token.is_none() {
|
||
tracing::warn!(
|
||
"auth disabled — anyone who can reach this address can use the service. \
|
||
Set --auth-token or RTX_AUTH_TOKEN for production."
|
||
);
|
||
}
|
||
let llm_extra_body = match cli.llm_extra_body.as_deref() {
|
||
Some(s) => {
|
||
let parsed: serde_json::Value = serde_json::from_str(s)
|
||
.with_context(|| format!("--llm-extra-body must be a JSON object: {s}"))?;
|
||
match parsed {
|
||
serde_json::Value::Object(m) => m,
|
||
other => anyhow::bail!("--llm-extra-body must be a JSON object, got {}", other),
|
||
}
|
||
}
|
||
None => serde_json::Map::new(),
|
||
};
|
||
if !llm_extra_body.is_empty() {
|
||
tracing::info!(
|
||
"LLM extra_body: {} field(s) — merged into every request",
|
||
llm_extra_body.len()
|
||
);
|
||
}
|
||
|
||
// Build the reactive-emotion classifier once at boot. Cheap for
|
||
// ProsodyDetector; pays the ~1.1 GB download + Metal allocation
|
||
// upfront for Emotion2Vec so per-turn cost is just ~150 ms forward.
|
||
let emotion_detector: Option<Box<dyn rtx_csm::ser::EmotionDetector + Send + Sync>> =
|
||
if cli.reactive_emotion {
|
||
if cli.use_emotion2vec {
|
||
let api = hf_hub::api::sync::Api::new()?;
|
||
let path = api
|
||
.model("emotion2vec/emotion2vec_plus_base".into())
|
||
.get("model.pt")
|
||
.context("download emotion2vec_plus_base/model.pt")?;
|
||
let model = rtx_csm::emotion2vec::Emotion2Vec::load_from_pickle(&path, &device)?;
|
||
tracing::info!("loaded emotion2vec_plus_base from {}", path.display());
|
||
Some(Box::new(model))
|
||
} else {
|
||
Some(Box::new(rtx_csm::ser::ProsodyDetector::default()))
|
||
}
|
||
} else {
|
||
None
|
||
};
|
||
|
||
let shared = Arc::new(Shared {
|
||
generator: Mutex::new(generator),
|
||
stt: Mutex::new(stt),
|
||
llm,
|
||
system_prompt: cli.system,
|
||
speaker: cli.speaker,
|
||
auth_token,
|
||
rate_audio_secs_per_min: cli.rate_audio_secs_per_min,
|
||
rate_turns_per_min: cli.rate_turns_per_min,
|
||
vad_eot: if cli.vad {
|
||
Some((cli.vad_threshold, cli.vad_consecutive))
|
||
} else {
|
||
None
|
||
},
|
||
llm_extra_body,
|
||
stream_tts: cli.stream_tts,
|
||
stream_chunk_frames: cli.stream_chunk_frames.max(1),
|
||
batch_asr: cli.whisper || cli.moonshine,
|
||
vad_gate: if cli.vad_gate {
|
||
tracing::info!(
|
||
"energy VAD gate enabled (rms_threshold={:.4})",
|
||
cli.vad_gate_threshold
|
||
);
|
||
Some(Mutex::new(rtx_csm::stt::VadGate::new(
|
||
cli.vad_gate_threshold,
|
||
)))
|
||
} else {
|
||
None
|
||
},
|
||
emotion_hint: cli.emotion_hint.clone(),
|
||
emotion_detector,
|
||
emotion_aware_llm: cli.emotion_aware_llm,
|
||
metrics: Metrics::default(),
|
||
});
|
||
|
||
let app = Router::new()
|
||
.route("/health", get(|| async { "ok" }))
|
||
.route("/metrics", get(metrics_handler))
|
||
.route("/v1/converse", get(ws_handler))
|
||
.with_state(shared);
|
||
|
||
let listener = tokio::net::TcpListener::bind(&cli.bind).await?;
|
||
tracing::info!("listening on http://{}/v1/converse (WebSocket)", cli.bind);
|
||
let shutdown = shutdown_signal();
|
||
axum::serve(listener, app)
|
||
.with_graceful_shutdown(shutdown)
|
||
.await?;
|
||
tracing::info!("server stopped");
|
||
Ok(())
|
||
}
|
||
|
||
/// SIGINT (Ctrl+C) or SIGTERM stops accepting new connections; in-flight
|
||
/// turns finish naturally before the server exits.
|
||
async fn shutdown_signal() {
|
||
let ctrl_c = async {
|
||
let _ = tokio::signal::ctrl_c().await;
|
||
};
|
||
#[cfg(unix)]
|
||
let term = async {
|
||
let mut sig = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
|
||
.expect("install SIGTERM handler");
|
||
sig.recv().await;
|
||
};
|
||
#[cfg(not(unix))]
|
||
let term = std::future::pending::<()>();
|
||
tokio::select! {
|
||
_ = ctrl_c => tracing::info!("received SIGINT, shutting down gracefully"),
|
||
_ = term => tracing::info!("received SIGTERM, shutting down gracefully"),
|
||
}
|
||
}
|
||
|
||
/// Prometheus-style /metrics endpoint. Counter + summary lines.
|
||
async fn metrics_handler(State(shared): State<Arc<Shared>>) -> impl IntoResponse {
|
||
let m = &shared.metrics;
|
||
let stt_count = m.stt_latency_ms_count.load(Ordering::Relaxed).max(1);
|
||
let tts_count = m.tts_latency_ms_count.load(Ordering::Relaxed).max(1);
|
||
let e2e_count = m.e2e_first_ms_count.load(Ordering::Relaxed).max(1);
|
||
let recv_count = m.recv_phase_ms_count.load(Ordering::Relaxed).max(1);
|
||
let llm_to_audio_count = m.llm_to_first_audio_ms_count.load(Ordering::Relaxed).max(1);
|
||
let conv_count = m.conv_total_ms_count.load(Ordering::Relaxed).max(1);
|
||
let total_count = m.total_turn_ms_count.load(Ordering::Relaxed).max(1);
|
||
let body = format!(
|
||
"# TYPE rtx_csm_turns_total counter\n\
|
||
rtx_csm_turns_total {}\n\
|
||
# TYPE rtx_csm_errors_total counter\n\
|
||
rtx_csm_errors_total {}\n\
|
||
# TYPE rtx_csm_connections_total counter\n\
|
||
rtx_csm_connections_total {}\n\
|
||
# TYPE rtx_csm_connections_active gauge\n\
|
||
rtx_csm_connections_active {}\n\
|
||
# TYPE rtx_csm_stt_latency_ms_avg gauge\n\
|
||
rtx_csm_stt_latency_ms_avg {}\n\
|
||
# TYPE rtx_csm_tts_latency_ms_avg gauge\n\
|
||
rtx_csm_tts_latency_ms_avg {}\n\
|
||
# TYPE rtx_csm_e2e_first_audio_ms_avg gauge\n\
|
||
rtx_csm_e2e_first_audio_ms_avg {}\n\
|
||
# TYPE rtx_csm_recv_phase_ms_avg gauge\n\
|
||
rtx_csm_recv_phase_ms_avg {}\n\
|
||
# TYPE rtx_csm_llm_to_first_audio_ms_avg gauge\n\
|
||
rtx_csm_llm_to_first_audio_ms_avg {}\n\
|
||
# TYPE rtx_csm_conv_total_ms_avg gauge\n\
|
||
rtx_csm_conv_total_ms_avg {}\n\
|
||
# TYPE rtx_csm_total_turn_ms_avg gauge\n\
|
||
rtx_csm_total_turn_ms_avg {}\n\
|
||
# TYPE rtx_csm_reactive_emotion_calls_total counter\n\
|
||
rtx_csm_reactive_emotion_calls_total {}\n\
|
||
# TYPE rtx_csm_reactive_emotion_total counter\n\
|
||
rtx_csm_reactive_emotion_total{{label=\"neutral\"}} {}\n\
|
||
rtx_csm_reactive_emotion_total{{label=\"calm\"}} {}\n\
|
||
rtx_csm_reactive_emotion_total{{label=\"sad\"}} {}\n\
|
||
rtx_csm_reactive_emotion_total{{label=\"angry\"}} {}\n\
|
||
rtx_csm_reactive_emotion_total{{label=\"excited\"}} {}\n\
|
||
rtx_csm_reactive_emotion_total{{label=\"disgusted\"}} {}\n\
|
||
rtx_csm_reactive_emotion_total{{label=\"fearful\"}} {}\n\
|
||
rtx_csm_reactive_emotion_total{{label=\"happy\"}} {}\n\
|
||
rtx_csm_reactive_emotion_total{{label=\"surprised\"}} {}\n\
|
||
rtx_csm_reactive_emotion_total{{label=\"unk\"}} {}\n\
|
||
# TYPE rtx_csm_emotion_aware_llm_applied_total counter\n\
|
||
rtx_csm_emotion_aware_llm_applied_total {}\n",
|
||
m.turns_total.load(Ordering::Relaxed),
|
||
m.errors_total.load(Ordering::Relaxed),
|
||
m.connections_total.load(Ordering::Relaxed),
|
||
m.connections_active.load(Ordering::Relaxed),
|
||
m.stt_latency_ms_sum.load(Ordering::Relaxed) / stt_count,
|
||
m.tts_latency_ms_sum.load(Ordering::Relaxed) / tts_count,
|
||
m.e2e_first_ms_sum.load(Ordering::Relaxed) / e2e_count,
|
||
m.recv_phase_ms_sum.load(Ordering::Relaxed) / recv_count,
|
||
m.llm_to_first_audio_ms_sum.load(Ordering::Relaxed) / llm_to_audio_count,
|
||
m.conv_total_ms_sum.load(Ordering::Relaxed) / conv_count,
|
||
m.total_turn_ms_sum.load(Ordering::Relaxed) / total_count,
|
||
m.reactive_emotion_calls_total.load(Ordering::Relaxed),
|
||
m.reactive_emotion_neutral_total.load(Ordering::Relaxed),
|
||
m.reactive_emotion_calm_total.load(Ordering::Relaxed),
|
||
m.reactive_emotion_sad_total.load(Ordering::Relaxed),
|
||
m.reactive_emotion_angry_total.load(Ordering::Relaxed),
|
||
m.reactive_emotion_excited_total.load(Ordering::Relaxed),
|
||
m.reactive_emotion_disgusted_total.load(Ordering::Relaxed),
|
||
m.reactive_emotion_fearful_total.load(Ordering::Relaxed),
|
||
m.reactive_emotion_happy_total.load(Ordering::Relaxed),
|
||
m.reactive_emotion_surprised_total.load(Ordering::Relaxed),
|
||
m.reactive_emotion_unk_total.load(Ordering::Relaxed),
|
||
m.emotion_aware_llm_applied_total.load(Ordering::Relaxed),
|
||
);
|
||
(
|
||
StatusCode::OK,
|
||
[(
|
||
axum::http::header::CONTENT_TYPE,
|
||
"text/plain; version=0.0.4",
|
||
)],
|
||
body,
|
||
)
|
||
}
|
||
|
||
async fn ws_handler(
|
||
ws: WebSocketUpgrade,
|
||
State(shared): State<Arc<Shared>>,
|
||
headers: HeaderMap,
|
||
) -> impl IntoResponse {
|
||
if let Some(expected) = shared.auth_token.as_ref() {
|
||
let supplied = headers
|
||
.get("authorization")
|
||
.and_then(|v| v.to_str().ok())
|
||
.and_then(|h| h.strip_prefix("Bearer "))
|
||
.unwrap_or("");
|
||
if supplied != expected {
|
||
shared.metrics.errors_total.fetch_add(1, Ordering::Relaxed);
|
||
return (StatusCode::UNAUTHORIZED, "unauthorized").into_response();
|
||
}
|
||
}
|
||
ws.on_upgrade(move |socket| handle_connection(socket, shared))
|
||
.into_response()
|
||
}
|
||
|
||
async fn handle_connection(mut socket: WebSocket, shared: Arc<Shared>) {
|
||
shared
|
||
.metrics
|
||
.connections_total
|
||
.fetch_add(1, Ordering::Relaxed);
|
||
shared
|
||
.metrics
|
||
.connections_active
|
||
.fetch_add(1, Ordering::Relaxed);
|
||
tracing::info!("WS connection opened");
|
||
let mut history: Vec<ChatMessage> = Vec::new();
|
||
history.push(ChatMessage::system(&shared.system_prompt));
|
||
|
||
// Per-connection rate limits.
|
||
let mut audio_bucket = RateBucket::new(shared.rate_audio_secs_per_min);
|
||
let mut turns_bucket = RateBucket::new(shared.rate_turns_per_min);
|
||
// If the previous turn ended in a barge-in, the carried-over user audio
|
||
// bytes seed the next turn so we don't drop them.
|
||
let mut carry_over: Option<Vec<f32>> = None;
|
||
|
||
'session: loop {
|
||
let turn_start = std::time::Instant::now();
|
||
// Charge the turn bucket up front; reject if over budget.
|
||
if !turns_bucket.try_consume(1) {
|
||
shared.metrics.errors_total.fetch_add(1, Ordering::Relaxed);
|
||
send_text(
|
||
&mut socket,
|
||
"{\"event\":\"error\",\"msg\":\"rate limit: turns per minute exceeded\"}",
|
||
)
|
||
.await;
|
||
break 'session;
|
||
}
|
||
// Reset Kyutai STT state for each new turn (Whisper is stateless
|
||
// so this is a no-op for Whisper).
|
||
{
|
||
let mut stt = shared.stt.lock().await;
|
||
if let AsrEngine::Kyutai(s) = &mut *stt {
|
||
if let Err(e) = s.reset() {
|
||
send_text(
|
||
&mut socket,
|
||
&format!("{{\"event\":\"error\",\"msg\":\"stt reset: {e}\"}}"),
|
||
)
|
||
.await;
|
||
break 'session;
|
||
}
|
||
}
|
||
}
|
||
|
||
let mut user_audio_24k: Vec<f32> = Vec::new();
|
||
// Barge-in carryover from the previous turn.
|
||
if let Some(co) = carry_over.take() {
|
||
user_audio_24k.extend(co);
|
||
}
|
||
// Same reset (some legacy duplication; harmless and Kyutai-only).
|
||
{
|
||
let mut stt = shared.stt.lock().await;
|
||
if let AsrEngine::Kyutai(s) = &mut *stt {
|
||
if let Err(e) = s.reset() {
|
||
send_text(
|
||
&mut socket,
|
||
&format!("{{\"event\":\"error\",\"msg\":\"stt reset: {e}\"}}"),
|
||
)
|
||
.await;
|
||
break 'session;
|
||
}
|
||
}
|
||
}
|
||
// VAD state (only used when --vad is enabled).
|
||
let mut consecutive_eot: u32 = 0;
|
||
// Tracks how much of user_audio_24k has been pushed to STT.
|
||
let mut stt_streaming_offset: usize = 0;
|
||
// Accumulate Word/EndWord events as they arrive.
|
||
let mut pending_word: Option<Vec<u32>> = None;
|
||
let mut transcript_words: Vec<String> = Vec::new();
|
||
// If we have carry-over audio, feed it first (it's already in
|
||
// user_audio_24k from the barge-in stash). Kyutai-only — Whisper
|
||
// skips parallel ingest entirely. STT runs on the blocking pool
|
||
// so the async worker can poll other tasks during the ~80ms
|
||
// per-frame work (multi-tenant concurrency win).
|
||
if !user_audio_24k.is_empty() && !shared.batch_asr {
|
||
let carry_slice = user_audio_24k.clone();
|
||
stt_streaming_offset = carry_slice.len();
|
||
let shared_carry = shared.clone();
|
||
let events = tokio::task::spawn_blocking(move || {
|
||
let mut guard = shared_carry.stt.blocking_lock();
|
||
match &mut *guard {
|
||
AsrEngine::Kyutai(s) => match s.step_pcm(&carry_slice) {
|
||
Ok(es) => es,
|
||
Err(e) => {
|
||
tracing::warn!("carry-over step_pcm: {e}");
|
||
Vec::new()
|
||
}
|
||
},
|
||
AsrEngine::Moonshine(_) => Vec::new(),
|
||
#[cfg(feature = "asr")]
|
||
AsrEngine::Whisper(_) => Vec::new(),
|
||
}
|
||
})
|
||
.await
|
||
.unwrap_or_else(|e| {
|
||
tracing::warn!("carry-over spawn_blocking: {e}");
|
||
Vec::new()
|
||
});
|
||
collect_transcript_events(
|
||
&events,
|
||
&shared.stt,
|
||
&mut pending_word,
|
||
&mut transcript_words,
|
||
)
|
||
.await;
|
||
} else if !user_audio_24k.is_empty() {
|
||
// Whisper mode: just track that carry-over is in the buffer.
|
||
stt_streaming_offset = user_audio_24k.len();
|
||
}
|
||
// Receive frames until EOT.
|
||
loop {
|
||
match socket.recv().await {
|
||
Some(Ok(Message::Binary(bytes))) => {
|
||
if bytes.len() % 2 != 0 {
|
||
continue;
|
||
}
|
||
let n = bytes.len() / 2;
|
||
user_audio_24k.reserve(n);
|
||
for c in bytes.chunks_exact(2) {
|
||
let s = i16::from_le_bytes([c[0], c[1]]);
|
||
user_audio_24k.push(s as f32 / i16::MAX as f32);
|
||
}
|
||
// Kyutai path: feed new samples into STT incrementally
|
||
// so the transcript builds as audio arrives. STT
|
||
// runs on the blocking pool — frees the async worker
|
||
// for other connections during the ~80 ms/frame
|
||
// compute (multi-tenant concurrency win; no impact
|
||
// on single-connection latency since STT is
|
||
// hardware-bound). Whisper path skips this — audio
|
||
// accumulates for the post-EOT batch transcribe.
|
||
//
|
||
// VAD gate (when --vad-gate): classify the new slice
|
||
// before STT; if silence, skip step_pcm entirely.
|
||
// Saves ~80-300 ms per silent chunk depending on
|
||
// chunk size.
|
||
let events = if !shared.batch_asr {
|
||
let new_slice = user_audio_24k[stt_streaming_offset..].to_vec();
|
||
stt_streaming_offset = user_audio_24k.len();
|
||
|
||
let skip_stt = if let Some(gate) = shared.vad_gate.as_ref() {
|
||
let mut g = gate.lock().await;
|
||
!g.is_speech(&new_slice)
|
||
} else {
|
||
false
|
||
};
|
||
|
||
if skip_stt {
|
||
Vec::new()
|
||
} else {
|
||
let shared_step = shared.clone();
|
||
tokio::task::spawn_blocking(move || {
|
||
let mut guard = shared_step.stt.blocking_lock();
|
||
match &mut *guard {
|
||
AsrEngine::Kyutai(s) => match s.step_pcm(&new_slice) {
|
||
Ok(es) => es,
|
||
Err(e) => {
|
||
tracing::warn!("step_pcm: {e}");
|
||
Vec::new()
|
||
}
|
||
},
|
||
AsrEngine::Moonshine(_) => Vec::new(),
|
||
#[cfg(feature = "asr")]
|
||
AsrEngine::Whisper(_) => Vec::new(),
|
||
}
|
||
})
|
||
.await
|
||
.unwrap_or_else(|e| {
|
||
tracing::warn!("step_pcm spawn_blocking: {e}");
|
||
Vec::new()
|
||
})
|
||
} // close `if !skip_stt` else-arm
|
||
} else {
|
||
stt_streaming_offset = user_audio_24k.len();
|
||
Vec::new()
|
||
};
|
||
collect_transcript_events(
|
||
&events,
|
||
&shared.stt,
|
||
&mut pending_word,
|
||
&mut transcript_words,
|
||
)
|
||
.await;
|
||
// VAD: watch Step events for end-of-turn probability.
|
||
if let Some((threshold, k)) = shared.vad_eot {
|
||
for ev in &events {
|
||
if let Some(pr) = Stt::end_of_turn_probability(ev) {
|
||
if pr > threshold {
|
||
consecutive_eot += 1;
|
||
if consecutive_eot >= k {
|
||
tracing::info!(
|
||
"VAD auto-EOT (pr={pr:.3} > {threshold} for {k} frames)"
|
||
);
|
||
send_text(&mut socket, "{\"event\":\"vad_eot\"}").await;
|
||
break;
|
||
}
|
||
} else {
|
||
consecutive_eot = 0;
|
||
}
|
||
}
|
||
}
|
||
if consecutive_eot >= k {
|
||
break;
|
||
}
|
||
}
|
||
}
|
||
Some(Ok(Message::Text(t))) if t.trim() == "EOT" => break,
|
||
Some(Ok(Message::Text(_))) => continue,
|
||
Some(Ok(Message::Close(_))) | None => break 'session,
|
||
Some(Ok(_)) => continue,
|
||
Some(Err(e)) => {
|
||
tracing::warn!("WS recv: {e}");
|
||
break 'session;
|
||
}
|
||
}
|
||
}
|
||
// Mark the end of the audio-receive phase. Everything after this
|
||
// is post-EOT processing (STT flush + LLM + TTS).
|
||
let recv_done_t = std::time::Instant::now();
|
||
let recv_phase_ms = turn_start.elapsed().as_millis() as u64;
|
||
shared
|
||
.metrics
|
||
.recv_phase_ms_sum
|
||
.fetch_add(recv_phase_ms, Ordering::Relaxed);
|
||
shared
|
||
.metrics
|
||
.recv_phase_ms_count
|
||
.fetch_add(1, Ordering::Relaxed);
|
||
// Drain the asr_delay buffer so any trailing Kyutai words
|
||
// flush. Whisper has no asr_delay buffer to drain.
|
||
let final_events = {
|
||
let mut stt = shared.stt.lock().await;
|
||
match &mut *stt {
|
||
AsrEngine::Kyutai(s) => s.finish().unwrap_or_default(),
|
||
AsrEngine::Moonshine(_) => Vec::new(),
|
||
#[cfg(feature = "asr")]
|
||
AsrEngine::Whisper(_) => Vec::new(),
|
||
}
|
||
};
|
||
collect_transcript_events(
|
||
&final_events,
|
||
&shared.stt,
|
||
&mut pending_word,
|
||
&mut transcript_words,
|
||
)
|
||
.await;
|
||
|
||
if user_audio_24k.is_empty() {
|
||
send_text(&mut socket, "{\"event\":\"error\",\"msg\":\"empty audio\"}").await;
|
||
continue 'session;
|
||
}
|
||
|
||
// Charge the audio-seconds bucket for this turn.
|
||
let audio_secs = (user_audio_24k.len() as f32 / PCM_RATE as f32).ceil() as u32;
|
||
if !audio_bucket.try_consume(audio_secs) {
|
||
shared.metrics.errors_total.fetch_add(1, Ordering::Relaxed);
|
||
send_text(
|
||
&mut socket,
|
||
"{\"event\":\"error\",\"msg\":\"rate limit: audio seconds per minute exceeded\"}",
|
||
)
|
||
.await;
|
||
break 'session;
|
||
}
|
||
|
||
// -- STT: transcribe the turn -------------------------------------
|
||
// STT runs at 24 kHz natively (Mimi sample rate). Pad with 2s
|
||
// silence suffix so the asr_delay buffer flushes the last words.
|
||
// (We use the same SAMPLE_RATE constant as the STT module.)
|
||
debug_assert_eq!(STT_SR, PCM_RATE);
|
||
user_audio_24k.extend(std::iter::repeat(0.0f32).take((PCM_RATE as f32 * 2.0) as usize));
|
||
|
||
// Build user_text. Kyutai has been streaming words incrementally
|
||
// — just join them. Whisper is batch-only, so we fire one
|
||
// transcribe call here over the full accumulated audio. (With
|
||
// Whisper-tiny on Metal, 10s of audio takes ~200ms — much
|
||
// faster than Kyutai's 5s parallel-STT-during-receive.)
|
||
// Build user_text. Kyutai has been streaming words incrementally
|
||
// — just join them. Batch backends (Whisper / Moonshine) fire
|
||
// one transcribe call here over the full accumulated audio.
|
||
let stt_t = std::time::Instant::now();
|
||
let user_text = if shared.batch_asr {
|
||
let stt = shared.stt.lock().await;
|
||
match &*stt {
|
||
AsrEngine::Moonshine(m) => match m.transcribe_24k(&user_audio_24k) {
|
||
Ok(t) => t.trim().to_string(),
|
||
Err(e) => {
|
||
tracing::warn!("moonshine transcribe: {e}");
|
||
String::new()
|
||
}
|
||
},
|
||
#[cfg(feature = "asr")]
|
||
AsrEngine::Whisper(w) => match w.transcribe_24k(&user_audio_24k) {
|
||
Ok(t) => t.trim().to_string(),
|
||
Err(e) => {
|
||
tracing::warn!("whisper transcribe: {e}");
|
||
String::new()
|
||
}
|
||
},
|
||
AsrEngine::Kyutai(_) => transcript_words.join(" ").trim().to_string(),
|
||
}
|
||
} else {
|
||
transcript_words.join(" ").trim().to_string()
|
||
};
|
||
let stt_ms = stt_t.elapsed().as_millis() as u64;
|
||
shared
|
||
.metrics
|
||
.stt_latency_ms_sum
|
||
.fetch_add(stt_ms, Ordering::Relaxed);
|
||
shared
|
||
.metrics
|
||
.stt_latency_ms_count
|
||
.fetch_add(1, Ordering::Relaxed);
|
||
let _ = ASR_DELAY_FRAMES;
|
||
|
||
let transcript_msg = serde_json::json!({"event":"transcript","text":&user_text});
|
||
send_text(&mut socket, &transcript_msg.to_string()).await;
|
||
if user_text.is_empty() {
|
||
shared.metrics.errors_total.fetch_add(1, Ordering::Relaxed);
|
||
send_text(
|
||
&mut socket,
|
||
"{\"event\":\"error\",\"msg\":\"empty transcript\"}",
|
||
)
|
||
.await;
|
||
continue 'session;
|
||
}
|
||
|
||
// -- Reactive emotion detection (Phase 13.4) ----------------------
|
||
// Reuses the audio buffer STT just consumed. ProsodyDetector
|
||
// operates at 16 kHz so we resample once on the fly; cheap for a
|
||
// ~10 s buffer and only paid when --reactive-emotion is on.
|
||
let reactive_tag: Option<String> = if let Some(det) = shared.emotion_detector.as_ref() {
|
||
// Trim the 2 s silence pad we appended for STT before
|
||
// classifying — silence drags the voiced ratio down.
|
||
let trim_to = user_audio_24k
|
||
.len()
|
||
.saturating_sub((PCM_RATE as f32 * 2.0) as usize);
|
||
let speech = &user_audio_24k[..trim_to];
|
||
shared
|
||
.metrics
|
||
.reactive_emotion_calls_total
|
||
.fetch_add(1, Ordering::Relaxed);
|
||
match rtx_csm::audio_io::resample(speech, PCM_RATE, 16_000) {
|
||
Ok(speech_16k) => {
|
||
let label = det
|
||
.classify(&speech_16k)
|
||
.unwrap_or(rtx_csm::ser::EmotionLabel::Neutral);
|
||
// Per-label counters.
|
||
let bucket = match label {
|
||
rtx_csm::ser::EmotionLabel::Neutral => {
|
||
&shared.metrics.reactive_emotion_neutral_total
|
||
}
|
||
rtx_csm::ser::EmotionLabel::Calm => {
|
||
&shared.metrics.reactive_emotion_calm_total
|
||
}
|
||
rtx_csm::ser::EmotionLabel::Sad => {
|
||
&shared.metrics.reactive_emotion_sad_total
|
||
}
|
||
rtx_csm::ser::EmotionLabel::Angry => {
|
||
&shared.metrics.reactive_emotion_angry_total
|
||
}
|
||
rtx_csm::ser::EmotionLabel::Excited => {
|
||
&shared.metrics.reactive_emotion_excited_total
|
||
}
|
||
rtx_csm::ser::EmotionLabel::Disgusted => {
|
||
&shared.metrics.reactive_emotion_disgusted_total
|
||
}
|
||
rtx_csm::ser::EmotionLabel::Fearful => {
|
||
&shared.metrics.reactive_emotion_fearful_total
|
||
}
|
||
rtx_csm::ser::EmotionLabel::Happy => {
|
||
&shared.metrics.reactive_emotion_happy_total
|
||
}
|
||
rtx_csm::ser::EmotionLabel::Surprised => {
|
||
&shared.metrics.reactive_emotion_surprised_total
|
||
}
|
||
rtx_csm::ser::EmotionLabel::Unk => {
|
||
&shared.metrics.reactive_emotion_unk_total
|
||
}
|
||
};
|
||
bucket.fetch_add(1, Ordering::Relaxed);
|
||
if label == rtx_csm::ser::EmotionLabel::Neutral {
|
||
None // fall back to static --emotion-hint
|
||
} else {
|
||
let tag = label.as_tag().to_string();
|
||
tracing::info!("reactive-emotion: detected {tag}");
|
||
Some(tag)
|
||
}
|
||
}
|
||
Err(e) => {
|
||
tracing::warn!("reactive-emotion resample failed: {e}");
|
||
None
|
||
}
|
||
}
|
||
} else {
|
||
None
|
||
};
|
||
let resolved_hint = reactive_tag.clone().or_else(|| shared.emotion_hint.clone());
|
||
|
||
// -- LLM + TTS: stream response back as audio chunks --------------
|
||
history.push(ChatMessage::user(&user_text));
|
||
let opts = ConverseOptions {
|
||
speaker: shared.speaker,
|
||
flush: FlushPolicy::Punctuation,
|
||
generate: GenerateOptions {
|
||
max_audio_ms: 8_000,
|
||
emotion_hint: resolved_hint,
|
||
..GenerateOptions::default()
|
||
},
|
||
..ConverseOptions::default()
|
||
};
|
||
let gen_cfg = GenConfig {
|
||
max_tokens: Some(160),
|
||
temperature: 0.7,
|
||
extra_body: shared.llm_extra_body.clone(),
|
||
..GenConfig::default()
|
||
};
|
||
// Cancel signal: set by the pump when barge-in is detected;
|
||
// checked inside the TTS callback to abort early.
|
||
let cancel = Arc::new(AtomicBool::new(false));
|
||
let cancel_for_cb = cancel.clone();
|
||
// Bridging channel: TTS callback pushes encoded PCM, pump_and_watch
|
||
// forwards to the WebSocket while concurrently watching for
|
||
// user-audio frames (barge-in).
|
||
let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<Vec<u8>>();
|
||
|
||
// Build the LLM-facing history. When --emotion-aware-llm is on AND
|
||
// we have a non-Neutral reactive tag (Phase 13.4), augment THIS
|
||
// turn's user message with a tone signal so the LLM adapts its
|
||
// text. The augmented message is per-turn-only — it never lands
|
||
// in persistent `history`, so subsequent turns aren't biased by
|
||
// stale signals. Phase 13.5.
|
||
let history_clone = if shared.emotion_aware_llm && reactive_tag.is_some() {
|
||
let tag = reactive_tag.as_deref().unwrap();
|
||
// Strip surrounding [] from the tag for the prose annotation.
|
||
let label = tag.trim_matches(|c| c == '[' || c == ']');
|
||
let mut h = history.clone();
|
||
if let Some(last) = h.last_mut() {
|
||
if matches!(last.role, rtx_csm::llm_client::Role::User) {
|
||
last.content = format!(
|
||
"{}\n\n[user audio tone: {label} — adjust your response in tone and content to match]",
|
||
last.content
|
||
);
|
||
shared
|
||
.metrics
|
||
.emotion_aware_llm_applied_total
|
||
.fetch_add(1, Ordering::Relaxed);
|
||
}
|
||
}
|
||
h
|
||
} else {
|
||
history.clone()
|
||
};
|
||
let metrics_for_tts = shared.clone();
|
||
// Capture the moment conv_fut starts so we can derive
|
||
// llm_to_first_audio (LLM stream + first-sentence TTS) and
|
||
// conv_total (whole conversation) latencies.
|
||
let conv_start_t = std::time::Instant::now();
|
||
let mut llm_to_first_audio_ms: Option<u64> = None;
|
||
let stream_tts = shared.stream_tts;
|
||
let stream_chunk_frames = shared.stream_chunk_frames;
|
||
let cancel_for_chunk_cb = cancel.clone();
|
||
// Move-captured handles for the spawned conv task. Everything
|
||
// here is Send + 'static (Arc clones, owned vectors, simple
|
||
// values, channel senders) so we can lift conv_fut out of the
|
||
// current task. Why spawn: tokio::join! polls cooperatively in
|
||
// ONE task, so a sync TTS gen blocks the whole runtime worker
|
||
// and prevents pump_fut from draining chunks. Spawn lets the
|
||
// runtime reschedule pump_fut onto a different worker while
|
||
// conv_fut is mid-synthesize.
|
||
let shared_conv = shared.clone();
|
||
let conv_fut = async move {
|
||
let mut g = shared_conv.generator.lock().await;
|
||
let mut conv = Converse::new(&shared_conv.llm, &mut *g);
|
||
let tx_cb = tx.clone();
|
||
let r = if stream_tts {
|
||
// Streaming: PCM chunks delivered as they're produced
|
||
// (low TTFA). Per-sentence callback only updates metrics
|
||
// — the audio has already shipped via on_chunk.
|
||
let tx_chunk = tx.clone();
|
||
let metrics_for_sentence = metrics_for_tts.clone();
|
||
conv.run_streaming(
|
||
history_clone,
|
||
gen_cfg,
|
||
opts,
|
||
stream_chunk_frames,
|
||
move |chunk: &[f32]| {
|
||
if cancel_for_chunk_cb.load(Ordering::Relaxed) {
|
||
return Err(rtx_csm::CsmError::Config(
|
||
"barge-in: TTS cancelled".into(),
|
||
));
|
||
}
|
||
let mut buf = Vec::with_capacity(chunk.len() * 2);
|
||
for &s in chunk {
|
||
let v = (s.clamp(-1.0, 1.0) * i16::MAX as f32) as i16;
|
||
buf.extend_from_slice(&v.to_le_bytes());
|
||
}
|
||
let _ = tx_chunk.send(buf);
|
||
Ok(())
|
||
},
|
||
move |u| {
|
||
metrics_for_sentence
|
||
.metrics
|
||
.tts_latency_ms_sum
|
||
.fetch_add(u.tts_latency_ms as u64, Ordering::Relaxed);
|
||
metrics_for_sentence
|
||
.metrics
|
||
.tts_latency_ms_count
|
||
.fetch_add(1, Ordering::Relaxed);
|
||
Ok(())
|
||
},
|
||
)
|
||
.await
|
||
} else {
|
||
conv.run(history_clone, gen_cfg, opts, move |u| {
|
||
if cancel_for_cb.load(Ordering::Relaxed) {
|
||
return Err(rtx_csm::CsmError::Config("barge-in: TTS cancelled".into()));
|
||
}
|
||
metrics_for_tts
|
||
.metrics
|
||
.tts_latency_ms_sum
|
||
.fetch_add(u.tts_latency_ms as u64, Ordering::Relaxed);
|
||
metrics_for_tts
|
||
.metrics
|
||
.tts_latency_ms_count
|
||
.fetch_add(1, Ordering::Relaxed);
|
||
let mut buf = Vec::with_capacity(u.audio.len() * 2);
|
||
for &s in &u.audio {
|
||
let v = (s.clamp(-1.0, 1.0) * i16::MAX as f32) as i16;
|
||
buf.extend_from_slice(&v.to_le_bytes());
|
||
}
|
||
let _ = tx_cb.send(buf);
|
||
Ok(())
|
||
})
|
||
.await
|
||
};
|
||
drop(tx); // close so pump exits naturally
|
||
r
|
||
};
|
||
// pump_and_watch: send TTS chunks to the WS while concurrently
|
||
// watching for incoming binary frames (barge-in) or "EOT" text
|
||
// (user wants to stop). Returns the carried-over user audio when
|
||
// barge-in fires so the next turn can start with it pre-buffered.
|
||
enum PumpResult {
|
||
Done,
|
||
BargeIn(Vec<f32>),
|
||
Disconnected,
|
||
}
|
||
let mut first_audio_recorded = false;
|
||
let pump_fut = async {
|
||
loop {
|
||
tokio::select! {
|
||
biased; // poll audio_rx first so a fast TTS doesn't get starved
|
||
chunk = rx.recv() => {
|
||
match chunk {
|
||
Some(buf) => {
|
||
if !first_audio_recorded {
|
||
first_audio_recorded = true;
|
||
let e2e_ms = turn_start.elapsed().as_millis() as u64;
|
||
shared
|
||
.metrics
|
||
.e2e_first_ms_sum
|
||
.fetch_add(e2e_ms, Ordering::Relaxed);
|
||
shared
|
||
.metrics
|
||
.e2e_first_ms_count
|
||
.fetch_add(1, Ordering::Relaxed);
|
||
// llm_to_first_audio: from conv_fut
|
||
// start to the first PCM chunk arriving
|
||
// here. Bundles LLM stream + first
|
||
// sentence's TTS gen.
|
||
let lta = conv_start_t.elapsed().as_millis() as u64;
|
||
llm_to_first_audio_ms = Some(lta);
|
||
shared
|
||
.metrics
|
||
.llm_to_first_audio_ms_sum
|
||
.fetch_add(lta, Ordering::Relaxed);
|
||
shared
|
||
.metrics
|
||
.llm_to_first_audio_ms_count
|
||
.fetch_add(1, Ordering::Relaxed);
|
||
}
|
||
if socket.send(Message::Binary(buf)).await.is_err() {
|
||
return PumpResult::Disconnected;
|
||
}
|
||
}
|
||
None => return PumpResult::Done,
|
||
}
|
||
}
|
||
incoming = socket.recv() => {
|
||
match incoming {
|
||
Some(Ok(Message::Binary(bytes))) if !bytes.is_empty() => {
|
||
cancel.store(true, Ordering::Relaxed);
|
||
// Decode the barge-in PCM so we can carry it
|
||
// forward as the first audio of the next turn.
|
||
let mut samples = Vec::with_capacity(bytes.len() / 2);
|
||
for c in bytes.chunks_exact(2) {
|
||
let s = i16::from_le_bytes([c[0], c[1]]);
|
||
samples.push(s as f32 / i16::MAX as f32);
|
||
}
|
||
// Drain pending audio chunks so we don't keep
|
||
// sending the cancelled TTS to the client.
|
||
while rx.try_recv().is_ok() {}
|
||
// Send a 50ms exponential-decay fade-out
|
||
// chunk so the audio doesn't cut to a click.
|
||
// (The client's last received sample is
|
||
// arbitrary; we synthesize a smooth ramp
|
||
// toward silence from low-amplitude noise.)
|
||
let fade_len = (PCM_RATE as f32 * 0.05) as usize;
|
||
let mut fade_buf = Vec::with_capacity(fade_len * 2);
|
||
for i in 0..fade_len {
|
||
let t = i as f32 / fade_len as f32;
|
||
let env = (-4.0 * t).exp(); // 1.0 → ~0.018
|
||
// Soft pseudo-noise so the fade isn't dead silence.
|
||
let n = ((i as u32).wrapping_mul(2654435761) as f32
|
||
/ u32::MAX as f32
|
||
- 0.5)
|
||
* 0.02;
|
||
let s = n * env;
|
||
let v = (s.clamp(-1.0, 1.0) * i16::MAX as f32) as i16;
|
||
fade_buf.extend_from_slice(&v.to_le_bytes());
|
||
}
|
||
let _ = socket
|
||
.send(Message::Binary(fade_buf.into()))
|
||
.await;
|
||
let _ = socket
|
||
.send(Message::Text("{\"event\":\"barge_in\"}".into()))
|
||
.await;
|
||
tracing::info!(
|
||
"barge-in detected ({} samples carried over, fade-out sent)",
|
||
samples.len()
|
||
);
|
||
return PumpResult::BargeIn(samples);
|
||
}
|
||
Some(Ok(Message::Text(t))) if t.trim() == "EOT" => {
|
||
cancel.store(true, Ordering::Relaxed);
|
||
while rx.try_recv().is_ok() {}
|
||
return PumpResult::Done;
|
||
}
|
||
Some(Ok(Message::Close(_))) | None => {
|
||
return PumpResult::Disconnected;
|
||
}
|
||
Some(Ok(_)) => continue,
|
||
Some(Err(_)) => return PumpResult::Disconnected,
|
||
}
|
||
}
|
||
}
|
||
}
|
||
};
|
||
// Spawn conv_fut as a separate task so the runtime can schedule
|
||
// pump_fut on another worker while conv_fut's sync TTS gen
|
||
// blocks. Both finish before we move on; pump_fut blocks until
|
||
// the channel closes (which happens when conv_fut drops `tx`).
|
||
let conv_handle = tokio::spawn(conv_fut);
|
||
let pump_res = pump_fut.await;
|
||
let conv_res: rtx_csm::error::Result<String> = match conv_handle.await {
|
||
Ok(r) => r,
|
||
Err(join_err) => Err(rtx_csm::CsmError::Config(format!(
|
||
"conv task join: {join_err}"
|
||
))),
|
||
};
|
||
let conv_total_ms = conv_start_t.elapsed().as_millis() as u64;
|
||
shared
|
||
.metrics
|
||
.conv_total_ms_sum
|
||
.fetch_add(conv_total_ms, Ordering::Relaxed);
|
||
shared
|
||
.metrics
|
||
.conv_total_ms_count
|
||
.fetch_add(1, Ordering::Relaxed);
|
||
let mut barge_in_audio: Option<Vec<f32>> = None;
|
||
match pump_res {
|
||
PumpResult::Done => {}
|
||
PumpResult::BargeIn(samples) => {
|
||
barge_in_audio = Some(samples);
|
||
}
|
||
PumpResult::Disconnected => break 'session,
|
||
}
|
||
// If barge-in cancelled TTS the conv result will be an Err. That's
|
||
// expected — don't surface it as an error to the client.
|
||
let assistant_text = if cancel.load(Ordering::Relaxed) {
|
||
// Cancelled — assistant turn is incomplete. Don't store in history;
|
||
// the user's next turn will replace what they were responding to.
|
||
String::new()
|
||
} else {
|
||
match conv_res {
|
||
Ok(t) => t,
|
||
Err(e) => {
|
||
shared.metrics.errors_total.fetch_add(1, Ordering::Relaxed);
|
||
send_text(
|
||
&mut socket,
|
||
&format!("{{\"event\":\"error\",\"msg\":\"converse: {e}\"}}"),
|
||
)
|
||
.await;
|
||
continue 'session;
|
||
}
|
||
}
|
||
};
|
||
|
||
if !assistant_text.is_empty() {
|
||
history.push(ChatMessage::assistant(&assistant_text));
|
||
shared.metrics.turns_total.fetch_add(1, Ordering::Relaxed);
|
||
let done_msg = serde_json::json!({"event":"done","assistant":&assistant_text});
|
||
send_text(&mut socket, &done_msg.to_string()).await;
|
||
}
|
||
let total_turn_ms = turn_start.elapsed().as_millis() as u64;
|
||
shared
|
||
.metrics
|
||
.total_turn_ms_sum
|
||
.fetch_add(total_turn_ms, Ordering::Relaxed);
|
||
shared
|
||
.metrics
|
||
.total_turn_ms_count
|
||
.fetch_add(1, Ordering::Relaxed);
|
||
// Per-turn timing summary. recv + stt_post + conv_total ≈ total
|
||
// (small slack for mutex acquire, channel close, log formatting).
|
||
// llm_to_first_audio is a sub-phase of conv_total: it's the time
|
||
// from conv_fut start to the first PCM byte hitting the pump.
|
||
let stt_post_ms = conv_start_t
|
||
.saturating_duration_since(recv_done_t)
|
||
.as_millis() as u64;
|
||
tracing::info!(
|
||
"turn timing: recv={}ms stt_post={}ms llm_to_first_audio={}ms \
|
||
conv_total={}ms total={}ms",
|
||
recv_phase_ms,
|
||
stt_post_ms,
|
||
llm_to_first_audio_ms.unwrap_or(0),
|
||
conv_total_ms,
|
||
total_turn_ms,
|
||
);
|
||
// Stash any barge-in audio so the next turn picks up where the user started.
|
||
carry_over = barge_in_audio;
|
||
}
|
||
|
||
shared
|
||
.metrics
|
||
.connections_active
|
||
.fetch_sub(1, Ordering::Relaxed);
|
||
tracing::info!("WS connection closed");
|
||
}
|
||
|
||
async fn send_text(socket: &mut WebSocket, payload: &str) -> bool {
|
||
socket
|
||
.send(Message::Text(payload.to_string()))
|
||
.await
|
||
.is_ok()
|
||
}
|
||
|
||
/// Walks Word/EndWord events: pairs them, detokenizes each word via the
|
||
/// STT's sentencepiece tokenizer, and appends to `transcript_words`. Step
|
||
/// events (VAD prs) are ignored here — they're handled by the VAD loop.
|
||
async fn collect_transcript_events(
|
||
events: &[AsrEvent],
|
||
stt_lock: &Mutex<AsrEngine>,
|
||
pending: &mut Option<Vec<u32>>,
|
||
transcript_words: &mut Vec<String>,
|
||
) {
|
||
if events.is_empty() {
|
||
return;
|
||
}
|
||
let stt = stt_lock.lock().await;
|
||
let kyutai = match &*stt {
|
||
AsrEngine::Kyutai(s) => s,
|
||
// Batch backends (Moonshine / Whisper) don't go through this —
|
||
// events come from Kyutai's incremental Word/EndWord stream only.
|
||
AsrEngine::Moonshine(_) => return,
|
||
#[cfg(feature = "asr")]
|
||
AsrEngine::Whisper(_) => return,
|
||
};
|
||
for ev in events {
|
||
match ev {
|
||
AsrEvent::Word { tokens, .. } => *pending = Some(tokens.clone()),
|
||
AsrEvent::EndWord { .. } => {
|
||
if let Some(tokens) = pending.take() {
|
||
if let Some(w) = kyutai.decode_word_text(&tokens) {
|
||
let w = w.trim().to_string();
|
||
if !w.is_empty() {
|
||
transcript_words.push(w);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
AsrEvent::Step { .. } => {}
|
||
}
|
||
}
|
||
}
|