llm_manager/
serve_api.rs

1use std::net::SocketAddr;
2use std::sync::{Arc, RwLock, Mutex};
3use std::time::Instant;
4
5use axum::Json;
6use axum::Router;
7use axum::body::{Body, Bytes, to_bytes};
8use axum::extract::State;
9use axum::http::StatusCode;
10use axum::response::IntoResponse;
11use axum::routing::{get, post};
12use futures_util::StreamExt;
13use tower_http::cors::CorsLayer;
14use tower_http::trace::TraceLayer;
15use tracing::info;
16
17use reqwest::Client;
18
19use crate::backend::web_context;
20
21
22
23
24pub struct StatusCache {
25    pub models: usize,
26    pub cached_at: Instant,
27}
28
29#[derive(Clone)]
30pub struct ApiState {
31    pub server_url: String,
32    pub api_key: Option<String>,
33    pub model_name: String,
34    pub pid: u32,
35    pub start_time: Instant,
36    pub port: u16,
37    pub client: reqwest::Client,
38    pub status_cache: Arc<RwLock<StatusCache>>,
39    pub system_prompt_preset_name: String,
40    pub web_search_engine: String,
41    pub web_search_engine_url: String,
42    pub web_search_enabled: bool,
43    pub web_search_api_key: Option<String>,
44    pub log_callback: Arc<Mutex<Option<Box<dyn Fn(String) + Send + Sync>>>>,
45}
46
47fn extract_api_key(headers: &axum::http::HeaderMap) -> Option<String> {
48    headers
49        .get("Authorization")
50        .and_then(|v| v.to_str().ok())
51        .and_then(|v| v.strip_prefix("Bearer "))
52        .map(|s| s.to_string())
53}
54
55async fn auth_middleware(
56    State(state): State<ApiState>,
57    req: axum::extract::Request,
58    next: axum::middleware::Next,
59) -> axum::response::Response {
60    tracing::debug!("auth_middleware: api_key={:?}", state.api_key.as_deref());
61    if let Some(expected) = &state.api_key {
62        let provided = extract_api_key(req.headers());
63        let expected_bytes = expected.as_bytes();
64        let not_equal = if let Some(provided_str) = provided {
65            constant_time_not_eq(provided_str.as_bytes(), expected_bytes)
66        } else {
67            true
68        };
69        if not_equal {
70            tracing::debug!("auth_middleware: rejecting request, not_equal={}", not_equal);
71            return (
72                StatusCode::UNAUTHORIZED,
73                Json(serde_json::json!({"error": "Unauthorized"})),
74            )
75                .into_response();
76        }
77    }
78    next.run(req).await
79}
80
81/// Constant-time byte comparison: returns true if a != b.
82/// Always processes all bytes regardless of where the first difference occurs.
83fn constant_time_not_eq(a: &[u8], b: &[u8]) -> bool {
84    if a.len() != b.len() {
85        return true;
86    }
87    let mut result: u8 = 0;
88    for (x, y) in a.iter().zip(b.iter()) {
89        result |= x ^ y;
90    }
91    result != 0
92}
93
94/// Proxy a request to the llama-server backend with SSE streaming support.
95/// Checks Content-Type: if text/event-stream, streams the body; otherwise buffers.
96async fn proxy_streaming(
97    State(state): State<ApiState>,
98    req: axum::extract::Request,
99) -> impl IntoResponse {
100    let path = req.uri().path().to_string();
101    let method = req.method().clone();
102    let headers = req.headers().clone();
103
104    let url = format!("{}{}", state.server_url, path);
105
106    // For chat completions and completions, drain body and optionally inject web search
107    if (path == "/v1/chat/completions" || path == "/v1/completions") && method == axum::http::Method::POST {
108        info!("API: proxying {} {}", method, path);
109        {
110            let cb = state.log_callback.lock().unwrap();
111            if let Some(c) = cb.as_ref() {
112                c(format!("API: proxying {} {}", method, path));
113            }
114        }
115        let body_bytes = match to_bytes(req.into_body(), 10 * 1024 * 1024).await {
116            Ok(b) => b,
117            Err(e) => {
118                info!("Failed to collect request body for {}: {}", path, e);
119                return (
120                    StatusCode::BAD_GATEWAY,
121                    Json(serde_json::json!({"error": format!("Failed to read request body: {}", e)})),
122                )
123                    .into_response();
124            }
125        };
126
127        let body_bytes = body_bytes;
128        let mut request_json: serde_json::Value = match serde_json::from_slice(&body_bytes) {
129            Ok(j) => j,
130            Err(e) => {
131                info!("Failed to parse request JSON for {}: {}", path, e);
132                return (
133                    StatusCode::BAD_REQUEST,
134                    Json(serde_json::json!({"error": format!("Invalid JSON: {}", e)})),
135                )
136                    .into_response();
137            }
138        };
139
140        info!("API: web_search_enabled={}, preset='{}', engine='{}'", 
141              state.web_search_enabled, state.system_prompt_preset_name, state.web_search_engine);
142        {
143            let cb = state.log_callback.lock().unwrap();
144            if let Some(c) = cb.as_ref() {
145                c(format!("API: web_search_enabled={}, preset='{}', engine='{}'", 
146                    state.web_search_enabled, state.system_prompt_preset_name, state.web_search_engine));
147            }
148        }
149
150        let result = web_context::build_injected_prompt(
151            &state.system_prompt_preset_name,
152            &request_json,
153            state.web_search_enabled,
154            &state.web_search_engine,
155            &state.web_search_engine_url,
156            state.web_search_api_key.as_deref().unwrap_or(""),
157            &state.log_callback,
158        ).await;
159
160        info!("API: web search performed={}, content_len={}", result.performed, result.content.len());
161        {
162            let cb = state.log_callback.lock().unwrap();
163            if let Some(c) = cb.as_ref() {
164                c(format!("API: web search performed={}, content_len={}", result.performed, result.content.len()));
165            }
166        }
167        if result.performed && !result.content.is_empty() {
168            if let Some(obj) = request_json.as_object_mut() {
169                if let Some(messages) = obj.get_mut("messages").and_then(|m| m.as_array_mut()) {
170                    if let Some(last) = messages.last_mut() {
171                        if let Some(content_val) = last.get_mut("content") {
172                            *content_val = serde_json::Value::String(result.content);
173                        }
174                    }
175                }
176            }
177        }
178
179        let modified_body = request_json.clone();
180
181        if let Some(messages) = modified_body.get("messages").and_then(|m| m.as_array()) {
182            let last_content = messages.last()
183                .and_then(|m| m.get("content").and_then(|c| c.as_str()))
184                .unwrap_or("");
185            info!("Prompt to llama-server: {} messages, last content ({} chars):\n{}", 
186                  messages.len(), last_content.len(), last_content);
187        }
188
189        let body_stream = futures_util::stream::once(async move {
190            Ok::<Bytes, std::convert::Infallible>(Bytes::from(
191                serde_json::to_vec(&modified_body).unwrap_or(body_bytes.to_vec())
192            ))
193        });
194
195        let mut request_builder = state.client.post(&url);
196
197   const HOP_BY_HOP: &[&str] = &[
198        "connection",
199        "keep-alive",
200        "proxy-authenticate",
201        "proxy-authorization",
202        "te",
203        "trailer",
204        "transfer-encoding",
205        "upgrade",
206        "host",
207        "content-length",
208    ];
209    let mut filtered = axum::http::HeaderMap::new();
210    for (name, value) in headers.iter() {
211        let n = name.as_str();
212        if !HOP_BY_HOP.contains(&n) && n != "authorization" {
213            filtered.insert(name, value.clone());
214        }
215    }
216    request_builder = request_builder.headers(filtered);
217
218    let response = request_builder
219        .body(reqwest::Body::wrap_stream(body_stream))
220        .send()
221        .await;
222
223    let response = handle_response(response, &path).await;
224    return response.into_response();
225}
226
227    // Stream request body directly to backend (no drain to memory)
228    let body_stream = req.into_body().into_data_stream().map(|r| {
229        r.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, format!("{}", e)))
230    });
231
232    let mut request_builder = match method {
233        axum::http::Method::GET => state.client.get(&url),
234        axum::http::Method::POST => state.client.post(&url),
235        axum::http::Method::PUT => state.client.put(&url),
236        axum::http::Method::DELETE => state.client.delete(&url),
237        _ => {
238            return (
239                StatusCode::METHOD_NOT_ALLOWED,
240                Json(serde_json::json!({"error": "Method not supported"})),
241            )
242                .into_response();
243        }
244    };
245
246    const HOP_BY_HOP: &[&str] = &[
247        "connection",
248        "keep-alive",
249        "proxy-authenticate",
250        "proxy-authorization",
251        "te",
252        "trailer",
253        "transfer-encoding",
254        "upgrade",
255        "host",
256    ];
257    let mut filtered = axum::http::HeaderMap::new();
258    for (name, value) in headers.iter() {
259        let n = name.as_str();
260        if !HOP_BY_HOP.contains(&n) && n != "authorization" {
261            filtered.insert(name, value.clone());
262        }
263    }
264    request_builder = request_builder.headers(filtered);
265
266    let response = request_builder
267        .body(reqwest::Body::wrap_stream(body_stream))
268        .send()
269        .await;
270
271     let response = handle_response(response, &path).await;
272    return response.into_response();
273}
274
275async fn handle_response(
276    response: Result<reqwest::Response, reqwest::Error>,
277    path: &str,
278) -> impl IntoResponse {
279    match response {
280        Ok(resp) => {
281            let status = resp.status();
282            let headers = resp.headers().clone();
283            let is_sse = resp
284                .headers()
285                .get(axum::http::header::CONTENT_TYPE)
286                .and_then(|v| v.to_str().ok())
287                .map(|v| v.contains("text/event-stream"))
288                .unwrap_or(false);
289
290            if is_sse {
291                let mut response = axum::response::Response::new(Body::from_stream(
292                    resp.bytes_stream()
293                        .map(|result| result.map_err(std::io::Error::other)),
294                ));
295                *response.status_mut() = status;
296                for (name, value) in headers.iter() {
297                    response.headers_mut().insert(name, value.clone());
298                }
299                response
300            } else {
301                let bytes = match resp.bytes().await {
302                    Ok(b) => b,
303                    Err(e) => {
304                        info!("Failed to read response body for {}: {}", path, e);
305                        return (
306                            StatusCode::BAD_GATEWAY,
307                            Json(serde_json::json!({"error": format!("Failed to read backend response: {}", e)})),
308                        )
309                            .into_response();
310                    }
311                };
312                (status, headers, bytes).into_response()
313            }
314        }
315        Err(e) => {
316            info!("Proxy error for {}: {}", path, e);
317            (
318                StatusCode::BAD_GATEWAY,
319                Json(serde_json::json!({"error": format!("Backend unavailable: {}", e)})),
320            )
321                .into_response()
322        }
323    }
324}
325
326/// Simple health check endpoint - no auth, verifies backend
327async fn health(State(state): State<ApiState>) -> impl IntoResponse {
328    let resp = state
329        .client
330        .get(format!("{}/health", state.server_url))
331        .send()
332        .await;
333
334    match resp {
335        Ok(response) if response.status().is_success() => Json(serde_json::json!({
336            "status": "ok",
337            "backend": "healthy"
338        })),
339        Ok(_) => Json(serde_json::json!({
340            "status": "degraded",
341            "backend": "unreachable"
342        })),
343        Err(_) => Json(serde_json::json!({
344            "status": "degraded",
345            "backend": "unreachable"
346        })),
347    }
348}
349
350/// Custom status endpoint.
351const STATUS_CACHE_TTL: std::time::Duration = std::time::Duration::from_secs(5);
352
353#[axum::debug_handler]
354async fn status(State(state): State<ApiState>) -> impl IntoResponse {
355    let uptime = state.start_time.elapsed();
356    let uptime_secs = uptime.as_secs();
357
358    let loaded_models = {
359        let (is_stale, cached_models) = {
360            let cache = state.status_cache.read().unwrap();
361            (cache.cached_at.elapsed() >= STATUS_CACHE_TTL, cache.models)
362        };
363        if !is_stale {
364            cached_models
365        } else {
366            let count = match state
367                .client
368                .get(format!("{}/models", state.server_url))
369                .send()
370                .await
371            {
372                Ok(resp) if resp.status().is_success() => {
373                    let val: Option<serde_json::Value> = resp.json().await.ok();
374                    let data = val.as_ref().and_then(|v| v.get("data")).and_then(|d| d.as_array());
375                    let c = data.map(|a| a.len()).unwrap_or(0);
376                    let mut cache = state.status_cache.write().unwrap();
377                    cache.models = c;
378                    cache.cached_at = Instant::now();
379                    c
380                }
381                _ => {
382                    let mut cache = state.status_cache.write().unwrap();
383                    cache.models = 0;
384                    cache.cached_at = Instant::now();
385                    0
386                }
387            };
388            count
389        }
390    };
391
392    Json(serde_json::json!({
393        "status": "running",
394        "pid": state.pid,
395        "port": state.port,
396        "model": state.model_name,
397        "uptime_seconds": uptime_secs,
398        "loaded_models": loaded_models,
399    }))
400}
401
402  pub async fn start_api_server(
403    addr: SocketAddr,
404    api_key: Option<String>,
405    server_port: u16,
406    model_name: String,
407    pid: u32,
408    mut shutdown_rx: tokio::sync::watch::Receiver<bool>,
409    host: String,
410    tls_config: Option<axum_server::tls_rustls::RustlsConfig>,
411    system_prompt_preset_name: String,
412    web_search_engine: String,
413    web_search_engine_url: String,
414    web_search_enabled: bool,
415    web_search_api_key: Option<String>,
416    log_callback: Arc<Mutex<Option<Box<dyn Fn(String) + Send + Sync>>>>,
417) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
418    let bind = addr;
419    let start_time = Instant::now();
420    let client = Client::builder()
421        .pool_max_idle_per_host(20)
422        .timeout(std::time::Duration::from_secs(60))
423        .build()?;
424    let state = ApiState {
425        server_url: format!("http://{}:{}", host, server_port),
426        api_key,
427        model_name,
428        pid,
429        start_time,
430        port: bind.port(),
431        client,
432        status_cache: Arc::new(RwLock::new(StatusCache {
433            models: 0,
434            cached_at: Instant::now() - std::time::Duration::from_secs(10),
435        })),
436        system_prompt_preset_name,
437        web_search_engine,
438        web_search_engine_url,
439        web_search_enabled,
440        web_search_api_key,
441        log_callback,
442    };
443
444    let cors = CorsLayer::new()
445        .allow_origin(tower_http::cors::Any)
446        .allow_methods([
447            axum::http::Method::GET,
448            axum::http::Method::POST,
449            axum::http::Method::PUT,
450            axum::http::Method::DELETE,
451            axum::http::Method::OPTIONS,
452        ])
453        .allow_headers([
454            axum::http::header::CONTENT_TYPE,
455            axum::http::header::AUTHORIZATION,
456        ]);
457
458    let api_key_clone = state.api_key.clone();
459    let protocol = if tls_config.is_some() {
460        "https"
461    } else {
462        "http"
463    };
464    info!(
465        "API server starting on {protocol}://{} (proxying to http://127.0.0.1:{})",
466        host, server_port
467    );
468    if api_key_clone.is_some() {
469        info!("API key authentication is ENABLED");
470    }
471
472    let app = Router::new()
473        .route("/health", get(health))
474        .route("/metrics", get(proxy_streaming))
475        .merge(
476            Router::new()
477                .route("/v1/chat/completions", post(proxy_streaming))
478                .route("/v1/completions", post(proxy_streaming))
479                .route("/v1/embeddings", post(proxy_streaming))
480                .route("/v1/models", get(proxy_streaming))
481                .route("/api/status", get(status))
482                .fallback(proxy_streaming)
483                .layer(cors)
484                .layer(TraceLayer::new_for_http()),
485        )
486        .layer(axum::middleware::from_fn_with_state(
487            state.clone(),
488            auth_middleware,
489        ))
490        .with_state(state);
491
492    match tls_config {
493        Some(tls_cfg) => {
494            let tls_listener = axum_server::bind_rustls(bind, tls_cfg);
495            let shutdown_fut = async {
496                let _ = shutdown_rx.wait_for(|v| *v).await;
497            };
498            tokio::select! {
499                result = tls_listener.serve(app.into_make_service()) => result?,
500                _ = shutdown_fut => {},
501            };
502        }
503        None => {
504            axum::serve(tokio::net::TcpListener::bind(bind).await?, app)
505                .with_graceful_shutdown(async move {
506                    let _ = shutdown_rx.wait_for(|v| *v).await;
507                })
508                .await?;
509        }
510    }
511    Ok(())
512}
llm_manager/serve_api.rs

llm_manager/
serve_api.rs
