crabllm_proxy/
lib.rs

1use axum::{
2    Json, Router,
3    extract::Request,
4    middleware,
5    response::Response,
6    routing::{get, post},
7};
8use crabllm_core::Storage;
9
10pub use auth::KeyName;
11pub use state::AppState;
12
13pub mod admin;
14pub mod auth;
15pub mod ext;
16mod handlers;
17mod state;
18pub mod storage;
19
20/// Middleware that tracks the number of in-flight API requests.
21/// For SSE streams, the gauge decrements when the response starts (not when the
22/// stream ends), so it undercounts long-lived streaming connections.
23async fn track_active_connections(request: Request, next: middleware::Next) -> Response {
24    metrics::gauge!("crabllm_active_connections").increment(1.0);
25    let response = next.run(request).await;
26    metrics::gauge!("crabllm_active_connections").decrement(1.0);
27    response
28}
29
30/// Build the Axum router with all API routes and admin routes.
31pub fn router<S: Storage + 'static>(state: AppState<S>, admin_routes: Vec<Router>) -> Router {
32    let mut app = Router::<AppState<S>>::new()
33        .route(
34            "/v1/chat/completions",
35            post(handlers::chat_completions::<S>),
36        )
37        .route("/v1/embeddings", post(handlers::embeddings::<S>))
38        .route(
39            "/v1/images/generations",
40            post(handlers::image_generations::<S>),
41        )
42        .route("/v1/audio/speech", post(handlers::audio_speech::<S>))
43        .route(
44            "/v1/audio/transcriptions",
45            post(handlers::audio_transcriptions::<S>),
46        )
47        .route("/v1/models", get(handlers::models::<S>))
48        .layer(middleware::from_fn_with_state(
49            state.clone(),
50            auth::auth::<S>,
51        ))
52        .layer(middleware::from_fn(track_active_connections))
53        .with_state(state);
54
55    // Health check — outside auth middleware so load balancers can probe it.
56    app = app.route(
57        "/health",
58        get(|| async { Json(serde_json::json!({"status": "ok"})) }),
59    );
60
61    // Merge extension-provided admin routes (stateless — extensions
62    // capture their own state via closures in the Router<()>).
63    for admin_router in admin_routes {
64        app = app.merge(admin_router);
65    }
66
67    app
68}
crabllm_proxy/lib.rs

crabllm_proxy/
lib.rs
