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

//! End-to-end LLM → CSM TTS pipeline demo.
//!
//! Streams an LLM response token-by-token, flushes each completed
//! sentence to CSM, writes the concatenated audio to a single WAV.
//!
//! Two modes:
//!   --mock          : use a hardcoded mock LLM (no API key needed; useful
//!                     for local end-to-end testing)
//!   --base/--model  : use a real OpenAI-compatible endpoint
//!
//! Usage (mock):
//! ```
//! cargo run -p rtx-csm --release --features metal --example converse -- \
//!     --mock --speaker 0 --out /tmp/converse_out.wav
//! ```
//!
//! Usage (live LLM):
//! ```
//! export OPENAI_API_KEY=sk-...
//! cargo run -p rtx-csm --release --features metal --example converse -- \
//!     --base "https://api.openai.com/v1" --model gpt-4o-mini \
//!     --prompt "Tell me a two-sentence story about a sleepy turtle." \
//!     --out /tmp/converse_out.wav
//! ```

use anyhow::{Context, Result};
use async_trait::async_trait;
use clap::Parser;
use futures_util::{stream, Stream};
use rtx_csm::{
    audio_io,
    converse::{Converse, ConverseOptions, FlushPolicy, Utterance},
    llm_client::{ChatMessage, GenConfig, LlmClient, OpenAiCompatibleClient, TokenStream},
    Generator,
};
use rtx_csm::error::Result as CsmResult;
use std::path::PathBuf;
use std::pin::Pin;

#[derive(Debug, Parser)]
#[command(name = "converse")]
struct Cli {
    /// Use the hardcoded mock LLM instead of a real API.
    #[arg(long)]
    mock: bool,
    /// LLM base URL (e.g. https://api.openai.com/v1).
    #[arg(long, default_value = "https://api.openai.com/v1")]
    base: String,
    #[arg(long, default_value = "gpt-4o-mini")]
    model: String,
    #[arg(long)]
    api_key: Option<String>,
    #[arg(long, default_value = "Tell me a two-sentence story about a sleepy turtle.")]
    prompt: String,
    #[arg(long, default_value = "You are a concise storyteller.")]
    system: String,
    #[arg(long, default_value_t = 0)]
    speaker: u32,
    #[arg(long, default_value_t = 0.7)]
    temperature: f32,
    #[arg(long, default_value_t = 256)]
    max_tokens: u32,
    /// Eager flush (every comma) for lower first-audio latency.
    #[arg(long)]
    eager: bool,
    /// Output WAV path (concatenated full assistant response).
    #[arg(long)]
    out: PathBuf,
    #[arg(long)]
    cpu: bool,
}

/// Mock LLM: emits a fixed response chunked into small token-like pieces
/// so the streaming pipeline is exercised without a network call.
struct MockLlm {
    chunks: Vec<&'static str>,
}

impl MockLlm {
    fn new() -> Self {
        Self {
            chunks: vec![
                "Once ", "upon ", "a ", "time ", "there ", "was ", "a ", "very ", "sleepy ",
                "turtle. ", "She ", "yawned ", "loudly ", "and ", "fell ", "asleep ", "in ",
                "the ", "warm ", "sun.",
            ],
        }
    }
}

#[async_trait]
impl LlmClient for MockLlm {
    async fn generate_stream(
        &self,
        _messages: Vec<ChatMessage>,
        _config: GenConfig,
    ) -> CsmResult<TokenStream> {
        let chunks: Vec<CsmResult<String>> =
            self.chunks.iter().map(|s| Ok(s.to_string())).collect();
        let s: Pin<Box<dyn Stream<Item = CsmResult<String>> + Send>> =
            Box::pin(stream::iter(chunks));
        Ok(s)
    }
}

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

    let device = if cli.cpu {
        candle_core::Device::Cpu
    } else {
        Generator::default_device()?
    };
    println!("device: {device:?}");
    let mut generator = Generator::load_csm_1b(&device)?;
    println!(
        "loaded CSM-1B (sr={} Hz)",
        generator.config.sample_rate
    );

    let messages = vec![
        ChatMessage::system(&cli.system),
        ChatMessage::user(&cli.prompt),
    ];
    let gen_cfg = GenConfig {
        max_tokens: Some(cli.max_tokens),
        temperature: cli.temperature,
        ..GenConfig::default()
    };
    let opts = ConverseOptions {
        speaker: cli.speaker,
        flush: if cli.eager { FlushPolicy::Eager } else { FlushPolicy::Punctuation },
        ..ConverseOptions::default()
    };

    let mut all: Vec<Utterance> = Vec::new();
    let t = std::time::Instant::now();
    let full_text = if cli.mock {
        let llm = MockLlm::new();
        let mut conv = Converse::new(&llm, &mut generator);
        conv.run(messages, gen_cfg, opts, |u| {
            println!(
                "  utt ({} samples, tts {}ms): {:?}",
                u.audio.len(),
                u.tts_latency_ms,
                u.text
            );
            all.push(u.clone());
            Ok(())
        })
        .await?
    } else {
        let api_key = cli
            .api_key
            .or_else(|| std::env::var("OPENAI_API_KEY").ok())
            .ok_or_else(|| anyhow::anyhow!("set --api-key or OPENAI_API_KEY (or use --mock)"))?;
        let llm = OpenAiCompatibleClient::new(&cli.base, api_key, &cli.model);
        let mut conv = Converse::new(&llm, &mut generator);
        conv.run(messages, gen_cfg, opts, |u| {
            println!(
                "  utt ({} samples, tts {}ms): {:?}",
                u.audio.len(),
                u.tts_latency_ms,
                u.text
            );
            all.push(u.clone());
            Ok(())
        })
        .await?
    };

    println!(
        "\n== full assistant response ({} chars, {} sentences) ==",
        full_text.len(),
        all.len()
    );
    println!("{full_text}");

    // Concatenate and write.
    let mut concat: Vec<f32> = Vec::new();
    for u in &all {
        concat.extend_from_slice(&u.audio);
    }
    audio_io::write_wav_24k_mono(&cli.out, &concat).context("write wav")?;
    println!(
        "\nwrote {} ({} samples = {:.2}s @ 24 kHz, total {:.2}s elapsed)",
        cli.out.display(),
        concat.len(),
        concat.len() as f32 / 24000.0,
        t.elapsed().as_secs_f32()
    );

    Ok(())
}