rustytorch/crates/models/rtx-csm/examples/converse_client.rs

//! WebSocket client for `converse_server`. Streams a WAV file as the
//! user turn, prints the transcript, saves the assistant response audio.
//!
//! Usage:
//! ```bash
//! # in another terminal: examples/converse_server with model loaded
//! cargo run -p rtx-csm --release --example converse_client -- \
//!   --url ws://127.0.0.1:18090/v1/converse \
//!   --in /tmp/asr_test.flac \
//!   --out /tmp/converse_response.wav
//! ```

use anyhow::{Context, Result};
use clap::Parser;
use futures_util::{SinkExt, StreamExt};
use rtx_csm::audio_io;
use std::path::PathBuf;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::protocol::Message;

const PCM_RATE: u32 = 24_000;
const FRAME_SAMPLES: usize = 4800; // 200 ms per frame

#[derive(Debug, Parser)]
struct Cli {
    #[arg(long, default_value = "ws://127.0.0.1:18090/v1/converse")]
    url: String,
    /// WAV/FLAC user-turn audio (any rate; resampled to 24 kHz mono).
    #[arg(long = "in")]
    input: PathBuf,
    /// Output WAV (assistant response).
    #[arg(long)]
    out: PathBuf,
    /// Bearer token sent in the Authorization header. Reads
    /// RTX_AUTH_TOKEN env var if not provided.
    #[arg(long)]
    auth_token: Option<String>,
    /// If set, after receiving N ms of assistant audio, inject a chunk
    /// of new audio (the same input WAV's first 200 ms by default) to
    /// simulate barge-in. Useful for testing 6c.3c.
    #[arg(long)]
    barge_in_after_ms: Option<u64>,
}

#[tokio::main]
async fn main() -> Result<()> {
    tracing_subscriber::fmt().init();
    let cli = Cli::parse();

    let samples = audio_io::load_mono_at_rate(&cli.input, PCM_RATE).context("load user audio")?;
    println!(
        "loaded {}: {} samples @ {} Hz ({:.2}s)",
        cli.input.display(),
        samples.len(),
        PCM_RATE,
        samples.len() as f32 / PCM_RATE as f32
    );

    let mut request = (&cli.url)
        .into_client_request()
        .with_context(|| format!("parse url {}", cli.url))?;
    if let Some(token) = cli
        .auth_token
        .or_else(|| std::env::var("RTX_AUTH_TOKEN").ok())
    {
        request.headers_mut().insert(
            "Authorization",
            format!("Bearer {token}").parse().context("auth header")?,
        );
    }
    let (mut ws, _resp) = tokio_tungstenite::connect_async(request)
        .await
        .with_context(|| format!("connect {}", cli.url))?;
    println!("connected to {}", cli.url);

    // Stream PCM in 200 ms frames. Server can begin transcribing as soon
    // as it has audio.
    for chunk in samples.chunks(FRAME_SAMPLES) {
        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());
        }
        ws.send(Message::Binary(buf.into())).await?;
    }
    ws.send(Message::Text("EOT".into())).await?;
    println!("sent {} samples + EOT", samples.len());

    // Read response until "done" event.
    let mut response_pcm: Vec<f32> = Vec::new();
    let t = std::time::Instant::now();
    let mut first_audio_ms: Option<u128> = None;
    let mut barge_in_fired = false;
    // Pre-encode 200 ms of the input as the barge-in payload (it's just
    // the first ~4800 samples, encoded as i16 LE bytes).
    let barge_payload: Vec<u8> = samples
        .iter()
        .take((PCM_RATE as f32 * 0.2) as usize)
        .flat_map(|s| {
            let v = (s.clamp(-1.0, 1.0) * i16::MAX as f32) as i16;
            v.to_le_bytes().to_vec()
        })
        .collect();

    while let Some(msg) = ws.next().await {
        let msg = msg?;
        match msg {
            Message::Binary(bytes) => {
                if first_audio_ms.is_none() {
                    first_audio_ms = Some(t.elapsed().as_millis());
                }
                for c in bytes.chunks_exact(2) {
                    let s = i16::from_le_bytes([c[0], c[1]]);
                    response_pcm.push(s as f32 / i16::MAX as f32);
                }
                // Optional: inject barge-in once we've heard N ms of audio.
                if let Some(after_ms) = cli.barge_in_after_ms {
                    if !barge_in_fired
                        && first_audio_ms
                            .is_some_and(|t0| t.elapsed().as_millis() - t0 > after_ms as u128)
                    {
                        println!(
                            "[barge-in] injecting {} bytes of audio after {} ms",
                            barge_payload.len(),
                            after_ms
                        );
                        ws.send(Message::Binary(barge_payload.clone().into()))
                            .await?;
                        barge_in_fired = true;
                    }
                }
            }
            Message::Text(t) => {
                let parsed: serde_json::Value =
                    serde_json::from_str(&t).context("parse server event")?;
                let event = parsed.get("event").and_then(|v| v.as_str()).unwrap_or("");
                match event {
                    "transcript" => {
                        let transcript = parsed.get("text").and_then(|v| v.as_str()).unwrap_or("");
                        println!("transcript: {transcript:?}");
                    }
                    "done" => {
                        let assistant = parsed
                            .get("assistant")
                            .and_then(|v| v.as_str())
                            .unwrap_or("");
                        println!("assistant: {assistant:?}");
                        break;
                    }
                    "barge_in" => {
                        println!("[server] barge_in event received — TTS cancelled");
                        // For this test we just exit; a real client would
                        // continue streaming the new turn's audio + EOT.
                        if cli.barge_in_after_ms.is_some() {
                            println!(
                                "[barge-in test] success — server correctly detected barge-in"
                            );
                            break;
                        }
                    }
                    "vad_eot" => {
                        println!("[server] vad_eot event received — auto end-of-turn");
                    }
                    "error" => {
                        let msg = parsed
                            .get("msg")
                            .and_then(|v| v.as_str())
                            .unwrap_or("(unknown)");
                        anyhow::bail!("server error: {msg}");
                    }
                    other => println!("(server event {other}: {t})"),
                }
            }
            Message::Close(_) => break,
            _ => {}
        }
    }

    audio_io::write_wav_24k_mono(&cli.out, &response_pcm)?;
    println!(
        "wrote {} samples = {:.2}s @ {} Hz to {}",
        response_pcm.len(),
        response_pcm.len() as f32 / PCM_RATE as f32,
        PCM_RATE,
        cli.out.display()
    );
    println!(
        "ttfa={:?}ms, total wall = {:.2}s",
        first_audio_ms,
        t.elapsed().as_secs_f32()
    );
    Ok(())
}