Middleware

Middleware lets you run your own async code around a chat run. chat-rs gives you two hooks, the before strategy and the after strategy, plus a retry strategy for handling failures.

Before and after strategies

A strategy is a closure that receives the Messages (which it may mutate synchronously) and the most recent Metadata, and returns a future for any async work. The before strategy runs once before the loop starts; the after strategy runs once after it completes successfully.

use chat_rs::{ChatBuilder, openai::OpenAIBuilder};
use chat_rs::types::messages::content;
use chat_rs::parts;

let client = OpenAIBuilder::new().with_model("gpt-4o").build();

let chat = ChatBuilder::new()
    .with_model(client)
    .with_before_strategy(Box::new(|messages, _metadata| {
        // Synchronous edits to `messages` are fine here.
        messages.push(content::from_system(parts!["Keep answers short."]));
        Box::pin(async {})
    }))
    .with_after_strategy(Box::new(|_messages, metadata| {
        let tokens = metadata.and_then(|m| m.usage.total_tokens);
        Box::pin(async move {
            println!("run finished, tokens used: {tokens:?}");
        })
    }))
    .build();

When they run

Both hooks run on the completion path, the embeddings path, and the streaming paths, so the same instrumentation applies no matter how you call the model.

The shape of a strategy

pub type CallbackStrategy =
    Box<dyn Fn(&mut Messages, Option<&Metadata>) -> CallbackFuture + Send + Sync>;

The returned future is 'static, so it cannot borrow messages across an await. Do any reading or mutation of messages synchronously, then move the owned values you need into the async block (as tokens is moved above).

Retrieval and external services (RAG)

Because a before strategy can mutate Messages and reach the outside world, it is a natural place to bring in extra information. Retrieval-augmented generation fits here: look up relevant documents and inject them as context before the model runs.

.with_before_strategy(Box::new(|messages, _metadata| {
    // Fetch relevant context, then add it to the conversation.
    let docs = my_vector_store.search_blocking("recent question");
    for doc in docs {
        messages.push(content::from_system(parts![doc.text]));
    }
    Box::pin(async {})
}))

The lookup happens synchronously here, before the returned future, which keeps it inside the 'static constraint. If your retrieval is itself async, run it before you call complete and pass the results in, or keep a blocking handle to your store for use inside the hook.

The same hook is a clean home for anything external: a feature flag check, a policy or moderation pass, a call to another model, or pulling state from a device or service. The after strategy is the mirror image, good for persisting the transcript or recording usage once the run finishes.