use crate::middleware::{ChatStream, Middleware, RequestHandler};
use crate::types::ChatRequestIR;
use async_trait::async_trait;
use tokio_util::sync::CancellationToken;
use tracing::{info, instrument};
pub struct LoggingMiddleware;
impl LoggingMiddleware {
pub fn new() -> Self {
Self
}
}
impl Default for LoggingMiddleware {
fn default() -> Self {
Self::new()
}
}
#[async_trait]
impl Middleware for LoggingMiddleware {
#[instrument(skip(self, next, cancel), fields(
model = %request.model.alias,
provider = ?request.model.provider.endpoint.kind
))]
async fn handle(&self, request: ChatRequestIR, cancel: CancellationToken, next: &dyn RequestHandler) -> anyhow::Result<ChatStream> {
info!("Processing chat request");
let start_time = std::time::Instant::now();
let request_id = request.metadata.get("request_id").cloned().unwrap_or_else(|| "unknown".to_string());
let result = next.handle(request, cancel).await;
match &result {
Ok(_) => {
let duration = start_time.elapsed();
info!(request_id = %request_id, duration = ?duration, "Request started successfully");
}
Err(e) => {
let duration = start_time.elapsed();
info!(request_id = %request_id, duration = ?duration, error = %e, "Request failed to start");
}
}
result
}
}