Skip to main content

maple_proxy/
lib.rs

1mod config;
2mod proxy;
3
4pub use config::Config;
5pub use opensecret::Pcr0Environment;
6use proxy::{health_check, proxy_openai_request, ProxyState};
7
8use axum::{
9    extract::DefaultBodyLimit,
10    http::Method,
11    routing::{get, post},
12    Router,
13};
14use std::sync::Arc;
15use tower::ServiceBuilder;
16use tower_http::{
17    cors::{Any, CorsLayer},
18    trace::{DefaultMakeSpan, DefaultOnResponse, TraceLayer},
19};
20use tracing::Level;
21
22const MAX_PROXY_REQUEST_BODY_BYTES: usize = 50 * 1024 * 1024;
23
24/// Create the Axum application with the given configuration
25pub fn create_app(config: Config) -> Router {
26    let state = Arc::new(ProxyState::new(config.clone()));
27    create_app_with_state(config, state)
28}
29
30pub(crate) fn create_app_with_state(config: Config, state: Arc<ProxyState>) -> Router {
31    let mut app = Router::new()
32        // Health check endpoints
33        .route("/health", get(health_check))
34        .route("/", get(health_check))
35        // OpenAI-compatible endpoints
36        .route("/v1/models", get(proxy_openai_request))
37        .route("/v1/chat/completions", post(proxy_openai_request))
38        .route("/v1/embeddings", post(proxy_openai_request))
39        .with_state(state)
40        .layer(
41            ServiceBuilder::new()
42                .layer(DefaultBodyLimit::max(MAX_PROXY_REQUEST_BODY_BYTES))
43                .layer(
44                    TraceLayer::new_for_http()
45                        .make_span_with(DefaultMakeSpan::new().level(Level::INFO))
46                        .on_response(DefaultOnResponse::new().level(Level::INFO)),
47                ),
48        );
49
50    // Add CORS if enabled
51    if config.enable_cors {
52        app = app.layer(
53            CorsLayer::new()
54                .allow_origin(Any)
55                .allow_methods([Method::GET, Method::POST, Method::OPTIONS])
56                .allow_headers(Any),
57        );
58    }
59
60    app
61}