Skip to main content

maple_proxy/
lib.rs

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