llm_manager/
serve_api.rs

1use std::net::SocketAddr;
2use std::time::Instant;
3
4use axum::Json;
5use axum::Router;
6use axum::body::Body;
7use axum::extract::State;
8use axum::http::StatusCode;
9use axum::response::IntoResponse;
10use axum::routing::{get, post};
11use futures_util::{StreamExt, stream};
12use tower_http::cors::CorsLayer;
13use tower_http::trace::TraceLayer;
14use tracing::info;
15
16use reqwest::Client;
17
18#[derive(Clone)]
19pub struct ApiState {
20    pub server_url: String,
21    pub api_key: Option<String>,
22    pub model_name: String,
23    pub pid: u32,
24    pub start_time: Instant,
25    pub port: u16,
26    pub client: reqwest::Client,
27}
28
29fn extract_api_key(headers: &axum::http::HeaderMap) -> Option<String> {
30    headers
31        .get("Authorization")
32        .and_then(|v| v.to_str().ok())
33        .and_then(|v| v.strip_prefix("Bearer "))
34        .map(|s| s.to_string())
35}
36
37async fn auth_middleware(
38    State(state): State<ApiState>,
39    req: axum::extract::Request,
40    next: axum::middleware::Next,
41) -> axum::response::Response {
42    if let Some(expected) = &state.api_key {
43        let provided = extract_api_key(req.headers());
44        if provided.as_deref() != Some(expected) {
45            return (
46                StatusCode::UNAUTHORIZED,
47                Json(serde_json::json!({"error": "Unauthorized"})),
48            )
49                .into_response();
50        }
51    }
52    next.run(req).await
53}
54
55/// Proxy a request to the llama-server backend with SSE streaming support.
56/// Checks Content-Type: if text/event-stream, streams the body; otherwise buffers.
57async fn proxy_streaming(
58    State(state): State<ApiState>,
59    req: axum::extract::Request,
60) -> impl IntoResponse {
61    let path = req.uri().path().to_string();
62    let method = req.method().clone();
63    let headers = req.headers().clone();
64
65    let url = format!("{}{}", state.server_url, path);
66
67    // Convert request body to a stream for reqwest
68    let body_bytes = match axum::body::to_bytes(req.into_body(), 10 * 1024 * 1024).await {
69        Ok(b) => b,
70        Err(e) => {
71            info!("Failed to read request body for {}: {}", path, e);
72            return (
73                StatusCode::BAD_REQUEST,
74                Json(serde_json::json!({"error": format!("Failed to read request body: {}", e)})),
75            )
76                .into_response();
77        }
78    };
79    let body_stream = stream::iter(vec![Ok::<_, reqwest::Error>(body_bytes)]);
80
81    let mut request_builder = match method {
82        axum::http::Method::GET => state.client.get(&url),
83        axum::http::Method::POST => state.client.post(&url),
84        axum::http::Method::PUT => state.client.put(&url),
85        axum::http::Method::DELETE => state.client.delete(&url),
86        _ => {
87            return (
88                StatusCode::METHOD_NOT_ALLOWED,
89                Json(serde_json::json!({"error": "Method not supported"})),
90            )
91                .into_response();
92        }
93    };
94
95    const HOP_BY_HOP: &[&str] = &[
96        "connection",
97        "keep-alive",
98        "proxy-authenticate",
99        "proxy-authorization",
100        "te",
101        "trailer",
102        "transfer-encoding",
103        "upgrade",
104        "host",
105    ];
106    for (name, value) in headers.iter() {
107        let name_str = name.as_str();
108        if !HOP_BY_HOP.contains(&name_str) && name_str != "authorization" {
109            request_builder = request_builder.header(name, value);
110        }
111    }
112
113    let response = request_builder
114        .body(reqwest::Body::wrap_stream(body_stream))
115        .send()
116        .await;
117
118    match response {
119        Ok(resp) => {
120            let status = resp.status();
121            let headers = resp.headers().clone();
122            let is_sse = resp
123                .headers()
124                .get(axum::http::header::CONTENT_TYPE)
125                .and_then(|v| v.to_str().ok())
126                .map(|v| v.contains("text/event-stream"))
127                .unwrap_or(false);
128
129            if is_sse {
130                let mut response = axum::response::Response::new(Body::from_stream(
131                    resp.bytes_stream()
132                        .map(|result| result.map_err(std::io::Error::other)),
133                ));
134                *response.status_mut() = status;
135                for (name, value) in headers.iter() {
136                    response.headers_mut().insert(name, value.clone());
137                }
138                response
139            } else {
140                let bytes = match resp.bytes().await {
141                    Ok(b) => b,
142                    Err(e) => {
143                        info!("Failed to read response body for {}: {}", path, e);
144                        return (
145                            StatusCode::BAD_GATEWAY,
146                            Json(serde_json::json!({"error": format!("Failed to read backend response: {}", e)})),
147                        )
148                            .into_response();
149                    }
150                };
151                (status, headers, bytes).into_response()
152            }
153        }
154        Err(e) => {
155            info!("Proxy error for {}: {}", path, e);
156            (
157                StatusCode::BAD_GATEWAY,
158                Json(serde_json::json!({"error": format!("Backend unavailable: {}", e)})),
159            )
160                .into_response()
161        }
162    }
163}
164
165/// Simple health check endpoint - no auth, verifies backend
166async fn health(State(state): State<ApiState>) -> impl IntoResponse {
167    let resp = state
168        .client
169        .get(format!("{}/health", state.server_url))
170        .send()
171        .await;
172
173    match resp {
174        Ok(response) if response.status().is_success() => Json(serde_json::json!({
175            "status": "ok",
176            "backend": "healthy"
177        })),
178        Ok(_) => Json(serde_json::json!({
179            "status": "degraded",
180            "backend": "unreachable"
181        })),
182        Err(_) => Json(serde_json::json!({
183            "status": "degraded",
184            "backend": "unreachable"
185        })),
186    }
187}
188
189/// Custom status endpoint.
190async fn status(State(state): State<ApiState>) -> impl IntoResponse {
191    let uptime = state.start_time.elapsed();
192    let uptime_secs = uptime.as_secs();
193
194    // Try to get loaded models from llama-server
195    let loaded_models = match state
196        .client
197        .get(format!("{}/models", state.server_url))
198        .send()
199        .await
200    {
201        Ok(resp) if resp.status().is_success() => {
202            let json: serde_json::Value = match resp.json().await {
203                Ok(v) => v,
204                Err(_) => serde_json::json!([]),
205            };
206            json.get("data")
207                .and_then(|d| d.as_array())
208                .map(|a| a.len())
209                .unwrap_or(0)
210        }
211        _ => 0,
212    };
213
214    Json(serde_json::json!({
215        "status": "running",
216        "pid": state.pid,
217        "port": state.port,
218        "model": state.model_name,
219        "uptime_seconds": uptime_secs,
220        "loaded_models": loaded_models,
221    }))
222}
223
224pub async fn start_api_server(
225    addr: SocketAddr,
226    api_key: Option<String>,
227    server_port: u16,
228    model_name: String,
229    pid: u32,
230    mut shutdown_rx: tokio::sync::watch::Receiver<bool>,
231    host: String,
232    tls_config: Option<axum_server::tls_rustls::RustlsConfig>,
233) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
234    let bind = addr;
235    let start_time = Instant::now();
236    let client = Client::builder()
237        .pool_max_idle_per_host(20)
238        .timeout(std::time::Duration::from_secs(300))
239        .build()?;
240    let state = ApiState {
241        server_url: format!("http://127.0.0.1:{}", server_port),
242        api_key,
243        model_name,
244        pid,
245        start_time,
246        port: bind.port(),
247        client,
248    };
249
250    let cors = CorsLayer::new()
251        .allow_origin(tower_http::cors::Any)
252        .allow_methods([
253            axum::http::Method::GET,
254            axum::http::Method::POST,
255            axum::http::Method::PUT,
256            axum::http::Method::DELETE,
257            axum::http::Method::OPTIONS,
258        ])
259        .allow_headers([
260            axum::http::header::CONTENT_TYPE,
261            axum::http::header::AUTHORIZATION,
262        ]);
263
264    let api_key_clone = state.api_key.clone();
265    let protocol = if tls_config.is_some() {
266        "https"
267    } else {
268        "http"
269    };
270    info!(
271        "API server starting on {protocol}://{} (proxying to http://127.0.0.1:{})",
272        host, server_port
273    );
274    if api_key_clone.is_some() {
275        info!("API key authentication is ENABLED");
276    }
277
278    let app = Router::new()
279        .route("/health", get(health))
280        .route("/metrics", get(proxy_streaming))
281        .merge(
282            Router::new()
283                .route("/v1/chat/completions", post(proxy_streaming))
284                .route("/v1/completions", post(proxy_streaming))
285                .route("/v1/embeddings", post(proxy_streaming))
286                .route("/v1/models", get(proxy_streaming))
287                .route("/api/status", get(status))
288                .fallback(proxy_streaming)
289                .layer(cors)
290                .layer(TraceLayer::new_for_http())
291                .layer(axum::middleware::from_fn_with_state(
292                    state.clone(),
293                    auth_middleware,
294                )),
295        )
296        .with_state(state);
297
298    match tls_config {
299        Some(tls_cfg) => {
300            let tls_listener = axum_server::bind_rustls(bind, tls_cfg);
301            let shutdown_fut = async {
302                let _ = shutdown_rx.wait_for(|v| *v).await;
303            };
304            tokio::select! {
305                result = tls_listener.serve(app.into_make_service()) => result?,
306                _ = shutdown_fut => {},
307            };
308        }
309        None => {
310            axum::serve(tokio::net::TcpListener::bind(bind).await?, app)
311                .with_graceful_shutdown(async move {
312                    let _ = shutdown_rx.wait_for(|v| *v).await;
313                })
314                .await?;
315        }
316    }
317    Ok(())
318}
llm_manager/serve_api.rs

llm_manager/
serve_api.rs
