Stream a reply¶

Print tokens as they arrive¶

chat_stream returns a ChatStream, which is a futures_util::Stream. Bring StreamExt into scope for .next():

use futures_util::StreamExt;
use rtb_chat::{AiClient, ChatRequest, ChatStreamEvent, Message};
use std::io::Write;

async fn run(client: &AiClient) -> Result<(), rtb_chat::AiError> {
    let mut stream = client
        .chat_stream(ChatRequest {
            messages: vec![Message::user("Explain prompt caching in two sentences.")],
            ..ChatRequest::default()
        })
        .await?;

    while let Some(event) = stream.next().await {
        match event {
            ChatStreamEvent::Token(t) => {
                print!("{t}");
                let _ = std::io::stdout().flush();
            }
            ChatStreamEvent::ThinkingToken(_) => {}
            ChatStreamEvent::Done(_) => break,
            ChatStreamEvent::Error(e) => return Err(e),
            _ => {}
        }
    }
    println!();
    Ok(())
}

Three details that matter:

• Break on Done. The stream does not finish there. It keeps yielding until the connection closes, so a loop without the break sits waiting after the answer is complete.
• Keep the _ arm. ChatStreamEvent is #[non_exhaustive]; new variants can arrive in a minor release.
• Flush. Tokens are fragments, not lines, and a line-buffered stdout shows nothing until the first newline.

Show the model thinking¶

ThinkingToken carries reasoning tokens rather than the answer. On the Anthropic-direct path you get them by setting ChatRequest::thinking; on the genai path they appear when the model reasons by default, whether you asked or not.

Keep them separate from the answer:

use rtb_chat::ChatStreamEvent;
fn handle(event: ChatStreamEvent, answer: &mut String, thoughts: &mut String) {
    match event {
        ChatStreamEvent::Token(t) => answer.push_str(&t),
        ChatStreamEvent::ThinkingToken(t) => thoughts.push_str(&t),
        _ => {}
    }
}

Concatenating both into one buffer produces the model's reasoning followed by its answer, which is rarely what a user wants to read.

Handle a failure part-way through¶

Errors reach you in two different places, and code that only handles one of them will miss the other:

• Connection-time: a bad key, a 429, a refused connection. chat_stream itself returns Err, with the same status mapping as chat.
• Mid-stream: the provider sends an SSE error event, or the byte stream breaks. That arrives as ChatStreamEvent::Error inside the loop.

Nothing closes the stream after an Error event. Break out of the loop yourself.

Read the token counts from Done¶

Done carries the Usage the provider reported over the course of the stream:

use futures_util::StreamExt as _;
use rtb_chat::ChatStreamEvent;

async fn run(client: &rtb_chat::AiClient, req: rtb_chat::ChatRequest)
    -> Result<(), rtb_chat::AiError> {
    let mut stream = client.chat_stream(req).await?;
    while let Some(event) = stream.next().await {
        if let ChatStreamEvent::Done(usage) = event {
            println!("{} in, {} out", usage.input_tokens, usage.output_tokens);
            break;
        }
    }
    Ok(())
}

Which counts each backend fills in is in Streaming usage arrives on Done.

Move a stream onto a task¶

ChatStream is Send, so it can go into tokio::spawn. It is not Sync, so it cannot be shared between tasks behind a plain reference. Move it, do not borrow it.