Crate cognate_axum 

Axum integration for Cognate.

Features

• CognateProvider — an Axum extractor that pulls an Arc<dyn Provider> out of shared application state via axum::extract::FromRef.
• UsageLayer / UsageHandle — Tower middleware that accumulates token usage across all requests.
• into_sse — convert a BoxStream<Result<Chunk>> directly into an Axum Sse response.

Wiring up the provider

use axum::{Router, routing::post};
use cognate_axum::CognateProvider;
use cognate_core::{Provider, Request};
use std::sync::Arc;

#[derive(Clone)]
struct AppState {
    provider: Arc<dyn Provider>,
}

impl axum::extract::FromRef<AppState> for Arc<dyn Provider> {
    fn from_ref(state: &AppState) -> Self {
        state.provider.clone()
    }
}

async fn chat(
    CognateProvider(provider): CognateProvider,
    axum::Json(req): axum::Json<Request>,
) -> String {
    provider.complete(req).await.unwrap().content().to_string()
}

#[tokio::main]
async fn main() {
    let provider: Arc<dyn Provider> = unimplemented!();
    let _app: Router = Router::new()
        .route("/chat", post(chat))
        .with_state(AppState { provider });
}

Structs

CognateProvider

Axum extractor that resolves an Arc<dyn Provider> from shared application state.

UsageHandle

A handle for reading token usage accumulated by UsageLayer.

UsageLayer

Tower [Layer] that wraps a Provider and records token usage.

Functions

into_sse

Convert a streaming provider response into an Axum Sse response.