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::{Provider, 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, P>(state: AppState<S, P>, admin_routes: Vec<Router>) -> Router
32where
33    S: Storage + 'static,
34    P: Provider + 'static,
35{
36    let mut app = Router::<AppState<S, P>>::new()
37        .route(
38            "/v1/chat/completions",
39            post(handlers::chat_completions::<S, P>),
40        )
41        .route("/v1/embeddings", post(handlers::embeddings::<S, P>))
42        .route(
43            "/v1/images/generations",
44            post(handlers::image_generations::<S, P>),
45        )
46        .route("/v1/audio/speech", post(handlers::audio_speech::<S, P>))
47        .route(
48            "/v1/audio/transcriptions",
49            post(handlers::audio_transcriptions::<S, P>),
50        )
51        .route("/v1/models", get(handlers::models::<S, P>))
52        .layer(middleware::from_fn_with_state(
53            state.clone(),
54            auth::auth::<S, P>,
55        ))
56        .layer(middleware::from_fn(track_active_connections))
57        .with_state(state);
58
59    // Health check — outside auth middleware so load balancers can probe it.
60    app = app.route(
61        "/health",
62        get(|| async { Json(serde_json::json!({"status": "ok"})) }),
63    );
64
65    // Merge extension-provided admin routes (stateless — extensions
66    // capture their own state via closures in the Router<()>).
67    for admin_router in admin_routes {
68        app = app.merge(admin_router);
69    }
70
71    app
72}
crabllm_proxy/lib.rs

crabllm_proxy/
lib.rs
