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]>
89 lines
2.8 KiB
Rust
89 lines
2.8 KiB
Rust
//! 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(())
|
|
}
|