rtx-csm: Phase 6b — LlmClient trait + OpenAI-compatible streaming impl

Generic LLM client abstraction for the conversational stack:

- LlmClient trait with generate_stream(messages, config) -> TokenStream.
  Default generate() impl folds the stream for non-streaming callers.
- ChatMessage / Role / GenConfig types with sensible defaults.
- OpenAiCompatibleClient: HTTP impl with SSE streaming. Works against
  OpenAI, Z.AI, vLLM, llama.cpp's HTTP server, LiteLLM — any endpoint
  serving the Chat Completions schema.
- examples/llm_chat: demo CLI that prints token-by-token to stdout
  with TTFT + total-time + char-count metrics.

Promotes tokio + reqwest + futures-util to regular dependencies (no
longer dev-only) so the trait is part of the public library surface.
Adds async-trait + eventsource-stream for the SSE streaming.

3 unit tests (constructors, role serialization, default config); 82
lib tests total green.

Phase 6 progress:
- 6a Kyutai STT: integration scaffolded; output bridge needs Python
  reference diff (deferred)
- 6b LLM client: shipped (this commit)
- 6c session glue (axum WS + duplex audio loop): next
- 6d productionization: deferred

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

crates/models/rtx-csm/Cargo.toml

@@ -67,6 +67,15 @@ half = "2.3"
 rand = "0.8"
 bytemuck = { version = "1.14", features = ["derive"] }
 
+# Async + HTTP for the LlmClient abstraction (Phase 6b). Promoted from
+# dev-dependency to regular dependency so the trait is part of the public
+# library surface.
+tokio = { version = "1", features = ["macros", "rt-multi-thread", "sync"] }
+futures-util = "0.3"
+reqwest = { version = "0.12", default-features = false, features = ["json", "stream", "rustls-tls"] }
+async-trait = "0.1"
+eventsource-stream = "0.2"
+
 [dev-dependencies]
 clap = { version = "4.5", features = ["derive"] }
 tempfile = "3.0"
@@ -74,13 +83,12 @@ approx = "0.5"
 tracing-subscriber = "0.3"
 # For the TTS HTTP server example.
 axum = { version = "0.7", features = ["multipart"] }
+# tokio with extra features (signal handler) needed by tts_server.
 tokio = { version = "1", features = ["macros", "rt-multi-thread", "signal", "sync"] }
 tower = "0.5"
 tower-http = { version = "0.6", features = ["trace"] }
-# HTTP client for the tts_server_bench example.
+# Multipart support added on top of the public reqwest dep for tts_server_bench.
 reqwest = { version = "0.12", default-features = false, features = ["json", "multipart", "rustls-tls"] }
-# Streaming Stream trait for tts_server's /v1/tts_stream endpoint.
-futures-util = "0.3"
 
 [features]
 default = ["cpu"]
@@ -181,3 +189,7 @@ path = "examples/tts_server_bench.rs"
 [[example]]
 name = "stt_demo"
 path = "examples/stt_demo.rs"
+
+[[example]]
+name = "llm_chat"
+path = "examples/llm_chat.rs"

crates/models/rtx-csm/examples/llm_chat.rs

@@ -0,0 +1,88 @@
+//! Tiny demo of the LlmClient streaming abstraction.
+//!
+//! Streams the assistant response token-by-token to stdout. Demonstrates
+//! the OpenAI-compatible client; works with OpenAI, Z.AI, vLLM, llama.cpp,
+//! or any Chat Completions endpoint.
+//!
+//! Usage (Z.AI defaults):
+//! ```bash
+//! export OPENAI_API_KEY=$Z_AI_API_KEY
+//! cargo run -p rtx-csm --release --example llm_chat -- \
+//!     --base "https://api.z.ai/api/coding/paas/v4" \
+//!     --model glm-4.6 \
+//!     --prompt "In one short sentence, why is rust good for ML?"
+//! ```
+
+use anyhow::Result;
+use clap::Parser;
+use futures_util::StreamExt;
+use rtx_csm::llm_client::{ChatMessage, GenConfig, LlmClient, OpenAiCompatibleClient};
+use std::io::Write;
+
+#[derive(Debug, Parser)]
+#[command(name = "llm_chat")]
+struct Cli {
+    /// Base URL of the OpenAI-compatible endpoint. Default OpenAI public.
+    #[arg(long, default_value = "https://api.openai.com/v1")]
+    base: String,
+    /// Model name.
+    #[arg(long, default_value = "gpt-4o-mini")]
+    model: String,
+    /// API key. Reads OPENAI_API_KEY from env if not set.
+    #[arg(long)]
+    api_key: Option<String>,
+    /// User prompt.
+    #[arg(long, default_value = "Say hello in one short sentence.")]
+    prompt: String,
+    /// Optional system prompt.
+    #[arg(long, default_value = "You are a concise assistant.")]
+    system: String,
+    #[arg(long, default_value_t = 0.7)]
+    temperature: f32,
+    #[arg(long, default_value_t = 256)]
+    max_tokens: u32,
+}
+
+#[tokio::main]
+async fn main() -> Result<()> {
+    let cli = Cli::parse();
+    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"))?;
+    let client = OpenAiCompatibleClient::new(&cli.base, api_key, &cli.model);
+    let messages = vec![
+        ChatMessage::system(&cli.system),
+        ChatMessage::user(&cli.prompt),
+    ];
+    let config = GenConfig {
+        max_tokens: Some(cli.max_tokens),
+        temperature: cli.temperature,
+        ..GenConfig::default()
+    };
+    println!("== {} via {} ==", cli.model, cli.base);
+    print!("user: {}\nassistant: ", cli.prompt);
+    std::io::stdout().flush().ok();
+
+    let t = std::time::Instant::now();
+    let mut first_token_ms: Option<u128> = None;
+    let mut stream = client.generate_stream(messages, config).await?;
+    let mut total_chars = 0;
+    while let Some(chunk) = stream.next().await {
+        let chunk = chunk?;
+        if first_token_ms.is_none() {
+            first_token_ms = Some(t.elapsed().as_millis());
+        }
+        total_chars += chunk.len();
+        print!("{chunk}");
+        std::io::stdout().flush().ok();
+    }
+    println!();
+    println!(
+        "[ttft={}ms total={}ms chars={}]",
+        first_token_ms.unwrap_or(0),
+        t.elapsed().as_millis(),
+        total_chars
+    );
+    Ok(())
+}

crates/models/rtx-csm/src/lib.rs

@@ -13,6 +13,7 @@ pub mod csm_quantized;
 pub mod error;
 pub mod generator;
 pub mod hub;
+pub mod llm_client;
 pub mod longform;
 pub mod lora;
 pub mod mimi;

crates/models/rtx-csm/src/llm_client.rs

@@ -0,0 +1,297 @@
+//! Generic LLM client abstraction for the Rust Unmute conversational stack.
+//!
+//! Provides a uniform [`LlmClient`] trait with `generate_stream` returning
+//! a stream of text tokens. Implementations:
+//!
+//!   - [`OpenAiCompatibleClient`] — works with OpenAI's Chat Completions
+//!     API and any compatible endpoint (Z.AI, vLLM, llama.cpp's HTTP server,
+//!     LiteLLM, etc.). Streams via Server-Sent Events.
+//!
+//! This is the bridge layer in the STT → LLM → TTS conversational pipeline.
+//! Token-level streaming is critical: TTS can start speaking the assistant
+//! response as soon as the first token arrives, instead of waiting for the
+//! full LLM completion.
+//!
+//! ## Example
+//!
+//! ```no_run
+//! use rtx_csm::llm_client::{ChatMessage, GenConfig, LlmClient, OpenAiCompatibleClient, Role};
+//! use futures_util::StreamExt;
+//!
+//! # async fn run() -> anyhow::Result<()> {
+//! let client = OpenAiCompatibleClient::new(
+//!     "https://api.openai.com/v1",
+//!     std::env::var("OPENAI_API_KEY")?,
+//!     "gpt-4o-mini",
+//! );
+//! let messages = vec![
+//!     ChatMessage::system("You are a concise assistant."),
+//!     ChatMessage::user("Say hello in one word."),
+//! ];
+//! let mut stream = client.generate_stream(messages, GenConfig::default()).await?;
+//! while let Some(tok) = stream.next().await {
+//!     print!("{}", tok?);
+//! }
+//! # Ok(()) }
+//! ```
+
+use crate::error::{CsmError, Result};
+use async_trait::async_trait;
+use eventsource_stream::Eventsource;
+use futures_util::stream::{Stream, StreamExt};
+use serde::{Deserialize, Serialize};
+use std::pin::Pin;
+
+/// Boxed stream of text tokens. Each token is one chunk of the assistant
+/// response — typically a sub-word from the LLM tokenizer. Concatenate to
+/// rebuild the full message; pipe to TTS as they arrive for streaming UX.
+pub type TokenStream = Pin<Box<dyn Stream<Item = Result<String>> + Send>>;
+
+#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
+#[serde(rename_all = "lowercase")]
+pub enum Role {
+    System,
+    User,
+    Assistant,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ChatMessage {
+    pub role: Role,
+    pub content: String,
+}
+
+impl ChatMessage {
+    pub fn system(content: impl Into<String>) -> Self {
+        Self {
+            role: Role::System,
+            content: content.into(),
+        }
+    }
+    pub fn user(content: impl Into<String>) -> Self {
+        Self {
+            role: Role::User,
+            content: content.into(),
+        }
+    }
+    pub fn assistant(content: impl Into<String>) -> Self {
+        Self {
+            role: Role::Assistant,
+            content: content.into(),
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct GenConfig {
+    /// Maximum tokens in the response.
+    pub max_tokens: Option<u32>,
+    /// Sampling temperature; 0.0 = greedy.
+    pub temperature: f32,
+    /// Top-p nucleus cutoff. 1.0 = disabled.
+    pub top_p: f32,
+    /// Optional stop sequences.
+    pub stop: Vec<String>,
+}
+
+impl Default for GenConfig {
+    fn default() -> Self {
+        Self {
+            max_tokens: Some(512),
+            temperature: 0.7,
+            top_p: 1.0,
+            stop: Vec::new(),
+        }
+    }
+}
+
+#[async_trait]
+pub trait LlmClient: Send + Sync {
+    /// Generate a streaming response. The returned stream yields text
+    /// chunks as the model produces them. Each chunk is typically one or
+    /// a few tokens; concatenate for the full response.
+    async fn generate_stream(
+        &self,
+        messages: Vec<ChatMessage>,
+        config: GenConfig,
+    ) -> Result<TokenStream>;
+
+    /// Convenience: collect the full response. Default impl folds the
+    /// stream into a single String. Implementations are free to override
+    /// for non-streaming endpoints.
+    async fn generate(
+        &self,
+        messages: Vec<ChatMessage>,
+        config: GenConfig,
+    ) -> Result<String> {
+        let mut stream = self.generate_stream(messages, config).await?;
+        let mut out = String::new();
+        while let Some(chunk) = stream.next().await {
+            out.push_str(&chunk?);
+        }
+        Ok(out)
+    }
+}
+
+// -- OpenAI-compatible Chat Completions impl ------------------------------
+
+/// HTTP client for OpenAI-compatible Chat Completions endpoints.
+/// Tested patterns:
+///   - OpenAI: base="https://api.openai.com/v1", model="gpt-4o-mini"
+///   - Z.AI:   base="https://api.z.ai/api/coding/paas/v4", model="glm-4.6"
+///   - vLLM:   base="http://localhost:8000/v1", model="<served-model-name>"
+///   - llama.cpp: base="http://localhost:8080/v1", model="<arbitrary>"
+#[derive(Debug, Clone)]
+pub struct OpenAiCompatibleClient {
+    base_url: String,
+    api_key: String,
+    model: String,
+    http: reqwest::Client,
+}
+
+impl OpenAiCompatibleClient {
+    pub fn new(
+        base_url: impl Into<String>,
+        api_key: impl Into<String>,
+        model: impl Into<String>,
+    ) -> Self {
+        Self {
+            base_url: base_url.into(),
+            api_key: api_key.into(),
+            model: model.into(),
+            http: reqwest::Client::builder()
+                .timeout(std::time::Duration::from_secs(120))
+                .build()
+                .expect("reqwest client init"),
+        }
+    }
+}
+
+// -- Wire types: OpenAI Chat Completions request/response -----------------
+
+#[derive(Serialize)]
+struct ChatRequest<'a> {
+    model: &'a str,
+    messages: &'a [ChatMessage],
+    stream: bool,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    max_tokens: Option<u32>,
+    temperature: f32,
+    top_p: f32,
+    #[serde(skip_serializing_if = "Vec::is_empty")]
+    stop: Vec<String>,
+}
+
+#[derive(Deserialize)]
+struct ChatStreamEvent {
+    choices: Vec<ChatStreamChoice>,
+}
+
+#[derive(Deserialize)]
+struct ChatStreamChoice {
+    delta: ChatStreamDelta,
+    /// Captured for completeness; OpenAI sets to `"stop"`/`"length"` on
+    /// the last event. We don't currently surface it to callers.
+    #[serde(default)]
+    #[allow(dead_code)]
+    finish_reason: Option<String>,
+}
+
+#[derive(Deserialize, Default)]
+struct ChatStreamDelta {
+    #[serde(default)]
+    content: Option<String>,
+}
+
+#[async_trait]
+impl LlmClient for OpenAiCompatibleClient {
+    async fn generate_stream(
+        &self,
+        messages: Vec<ChatMessage>,
+        config: GenConfig,
+    ) -> Result<TokenStream> {
+        let url = format!("{}/chat/completions", self.base_url.trim_end_matches('/'));
+        let req = ChatRequest {
+            model: &self.model,
+            messages: &messages,
+            stream: true,
+            max_tokens: config.max_tokens,
+            temperature: config.temperature,
+            top_p: config.top_p,
+            stop: config.stop.clone(),
+        };
+        let response = self
+            .http
+            .post(&url)
+            .bearer_auth(&self.api_key)
+            .header("accept", "text/event-stream")
+            .json(&req)
+            .send()
+            .await
+            .map_err(|e| CsmError::Config(format!("LLM request: {e}")))?;
+        if !response.status().is_success() {
+            let status = response.status();
+            let body = response.text().await.unwrap_or_default();
+            return Err(CsmError::Config(format!(
+                "LLM HTTP {status}: {body}"
+            )));
+        }
+
+        // SSE stream: each event has data: <json>\n\n. Final event is data: [DONE].
+        let bytes_stream = response.bytes_stream();
+        let event_stream = bytes_stream.eventsource();
+        let token_stream = event_stream.filter_map(|ev| async move {
+            match ev {
+                Err(e) => Some(Err(CsmError::Config(format!("SSE: {e}")))),
+                Ok(event) => {
+                    let data = event.data;
+                    if data.trim() == "[DONE]" {
+                        None
+                    } else {
+                        match serde_json::from_str::<ChatStreamEvent>(&data) {
+                            Ok(parsed) => parsed
+                                .choices
+                                .into_iter()
+                                .next()
+                                .and_then(|c| c.delta.content)
+                                .filter(|s| !s.is_empty())
+                                .map(Ok),
+                            Err(e) => Some(Err(CsmError::Config(format!(
+                                "SSE parse: {e} body={data}"
+                            )))),
+                        }
+                    }
+                }
+            }
+        });
+        Ok(Box::pin(token_stream))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn chat_message_constructors() {
+        assert_eq!(ChatMessage::system("s").role, Role::System);
+        assert_eq!(ChatMessage::user("u").role, Role::User);
+        assert_eq!(ChatMessage::assistant("a").role, Role::Assistant);
+    }
+
+    #[test]
+    fn role_serializes_lowercase() {
+        let m = ChatMessage::user("hi");
+        let json = serde_json::to_string(&m).unwrap();
+        assert!(json.contains("\"role\":\"user\""), "got {json}");
+    }
+
+    #[test]
+    fn gen_config_defaults_sane() {
+        let c = GenConfig::default();
+        assert_eq!(c.max_tokens, Some(512));
+        assert!((c.temperature - 0.7).abs() < 1e-6);
+        assert_eq!(c.top_p, 1.0);
+        assert!(c.stop.is_empty());
+    }
+}