osobhandClaude Opus 4.7
385858a3ba
Adds `--watermark-silentcipher` flag to converse_server. Sesame's actual
production watermarker is now a drop-in option in the conversation
pipeline alongside the existing AudioSeal flags.
Usage:
--watermark-silentcipher hf # download from sony/silentcipher
--watermark-silentcipher /path/to/dir # local checkpoint dir
--watermark-message 0xCAFE # 16-bit message (also drives AudioSeal)
Mutex with `--watermark-generator` / `--watermark-detector` (AudioSeal):
the Generator only carries one watermarker. Both are wrapped with
`ResampledWatermarker(24 kHz <-> 16 kHz)` for the CSM TTS path.
End-to-end production test (Q8 + Kyutai + mock LLM + SilentCipher 0xCAFE):
client TTFA: 7888 ms
total wall: 21860 ms
assistant audio: 12.16 s @ 24 kHz, written to /tmp/converse_silent_response.wav
re-detect: confidence 0.7614, payload 0x0000CAFE (PASS)
The 0.76 confidence (vs 1.00 in the standalone CLI test) is expected —
the assistant audio went through 24->16->24 resample plus stream-
encode-decode, all of which add noise. Still well above the 0.7
threshold we use for `Option<u16> -> Some/None` mapping in the
Watermarker trait impl.
A/B vs AudioSeal on the same /tmp/asr_test.flac (10.43 s @ 24 kHz):
AudioSeal SilentCipher
Embed timing not in CLI 988 ms (0.10x rt)
Detect timing not in CLI 1493 ms (0.14x rt)
Bit accuracy 16/16 bits 15/15 codes
Confidence 1.0000 1.0000
Message 0xCAFE 0xCAFE (decimal 51966)
Both bit-perfect. AudioSeal carries 16 bits, SilentCipher carries up
to ~24 bits per patch (15 base-3 codes). For our use (16-bit job_id
or message hash), either fits.
Production recommendation: ship SilentCipher for literal Sesame
parity AND the structural advantages (smaller model, identical bit
accuracy, confidence-based threshold). AudioSeal stays available for
callers who want the per-sample presence map (which SilentCipher
doesn't provide).
Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
1601 lines
66 KiB
Rust
1601 lines
66 KiB
Rust
//! Rust Unmute MVP — WebSocket conversational server.
|
||
//!
|
||
//! 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::{
|
||
extract::{
|
||
ws::{Message, WebSocket},
|
||
State, WebSocketUpgrade,
|
||
},
|
||
http::{HeaderMap, StatusCode},
|
||
response::IntoResponse,
|
||
routing::get,
|
||
Router,
|
||
};
|
||
use clap::Parser;
|
||
use futures_util::{stream, StreamExt};
|
||
use rtx_csm::{
|
||
converse::{Converse, ConverseOptions, FlushPolicy},
|
||
error::Result as CsmResult,
|
||
llm_client::{ChatMessage, GenConfig, LlmClient, OpenAiCompatibleClient, TokenStream},
|
||
stt::{AsrEvent, Stt, ASR_DELAY_FRAMES, SAMPLE_RATE as STT_SR},
|
||
GenerateOptions, Generator,
|
||
};
|
||
use std::net::SocketAddr;
|
||
use std::path::PathBuf;
|
||
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
|
||
use std::sync::Arc;
|
||
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,
|
||
|
||
/// 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`).
|
||
/// Loaded into the FP backbone before the conversation loop starts so
|
||
/// every assistant utterance uses the cloned voice. Cannot combine
|
||
/// with `--quantized-gguf` yet (LoRA-on-Q8 path exists in `generate`
|
||
/// but the server hasn't been audited for that combo).
|
||
#[arg(long)]
|
||
lora: Option<PathBuf>,
|
||
/// LoRA rank — must match training. Default 8.
|
||
#[arg(long, default_value_t = 8)]
|
||
lora_rank: usize,
|
||
/// LoRA alpha — must match training. Default 16.
|
||
#[arg(long, default_value_t = 16.0)]
|
||
lora_alpha: f32,
|
||
|
||
/// 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,
|
||
}
|
||
|
||
/// 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>>,
|
||
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();
|
||
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() {
|
||
if cli.lora.is_some() {
|
||
anyhow::bail!(
|
||
"--lora and --quantized-gguf cannot combine in converse_server yet \
|
||
(the LoRA-on-Q8 path works in `generate` but the server path \
|
||
hasn't been audited; load FP for voice clone)."
|
||
);
|
||
}
|
||
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 lora_cfg = rtx_csm::lora::LoraConfig {
|
||
rank: cli.lora_rank,
|
||
alpha: cli.lora_alpha,
|
||
..rtx_csm::lora::LoraConfig::default()
|
||
};
|
||
let vm = candle_nn::VarMap::new();
|
||
generator.model.inner.add_lora_to_backbone(&lora_cfg, &vm)?;
|
||
rtx_csm::training::load_lora_adapter(&vm, lora_path, &device)?;
|
||
generator.model.inner.refresh_lora(&vm)?;
|
||
tracing::info!(
|
||
"loaded LoRA adapter from {} (rank={} alpha={})",
|
||
lora_path.display(),
|
||
cli.lora_rank,
|
||
cli.lora_alpha,
|
||
);
|
||
}
|
||
|
||
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()
|
||
);
|
||
}
|
||
|
||
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
|
||
},
|
||
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",
|
||
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,
|
||
);
|
||
(
|
||
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;
|
||
}
|
||
|
||
// -- 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,
|
||
..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>>();
|
||
let history_clone = 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 { .. } => {}
|
||
}
|
||
}
|
||
}
|