feat: add Axum route handlers for /v1/chat/completions, /health, /v1/models

openai-proxy/src/routes.rs

@@ -0,0 +1,120 @@
+use axum::{
+    body::Body,
+    extract::State,
+    http::{HeaderMap, HeaderValue, StatusCode},
+    response::{IntoResponse, Response},
+    Json,
+};
+use futures::StreamExt;
+use serde_json::{json, Value};
+use tracing::{error, info};
+
+use crate::proxy::{ProxyClient, ProxyResponse};
+use crate::types::ChatCompletionRequest;
+
+#[derive(Clone)]
+pub struct AppState {
+    pub proxy: ProxyClient,
+}
+
+pub async fn health() -> impl IntoResponse {
+    Json(json!({"status": "ok"}))
+}
+
+pub async fn chat_completions(
+    State(state): State<AppState>,
+    Json(req): Json<ChatCompletionRequest>,
+) -> Response {
+    let is_streaming = req.stream.unwrap_or(false);
+    info!(model = %req.model, streaming = is_streaming, "chat completion request");
+
+    match state.proxy.forward_request(req).await {
+        Ok(ProxyResponse::Buffered(bytes)) => {
+            let mut headers = HeaderMap::new();
+            headers.insert(
+                axum::http::header::CONTENT_TYPE,
+                HeaderValue::from_static("application/json"),
+            );
+            (StatusCode::OK, headers, bytes).into_response()
+        }
+        Ok(ProxyResponse::Streamed(upstream_resp)) => {
+            let status = upstream_resp.status();
+            let upstream_headers = upstream_resp.headers().clone();
+
+            let stream = upstream_resp.bytes_stream().map(|result| {
+                result.map_err(|e| {
+                    error!("stream error: {}", e);
+                    std::io::Error::new(std::io::ErrorKind::Other, e)
+                })
+            });
+
+            let body = Body::from_stream(stream);
+            let mut resp = Response::new(body);
+            *resp.status_mut() = status;
+
+            if let Some(ct) = upstream_headers.get(axum::http::header::CONTENT_TYPE) {
+                resp.headers_mut().insert(axum::http::header::CONTENT_TYPE, ct.clone());
+            }
+
+            resp
+        }
+        Err(e) => {
+            error!("proxy error: {}", e);
+            let body = Json(json!({
+                "error": {
+                    "message": e.to_string(),
+                    "type": "proxy_error",
+                    "code": "upstream_error"
+                }
+            }));
+            (StatusCode::BAD_GATEWAY, body).into_response()
+        }
+    }
+}
+
+pub async fn models(State(state): State<AppState>) -> Response {
+    match state.proxy.forward_get("/v1/models").await {
+        Ok(bytes) => {
+            let mut headers = HeaderMap::new();
+            headers.insert(
+                axum::http::header::CONTENT_TYPE,
+                HeaderValue::from_static("application/json"),
+            );
+            (StatusCode::OK, headers, bytes).into_response()
+        }
+        Err(e) => {
+            error!("models proxy error: {}", e);
+            (
+                StatusCode::BAD_GATEWAY,
+                Json(json!({"error": {"message": e.to_string()}})),
+            )
+                .into_response()
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use axum::{body::to_bytes, http::Request, routing::get};
+    use tower::util::ServiceExt;
+
+    #[tokio::test]
+    async fn test_health_returns_200() {
+        let app = axum::Router::new().route("/health", get(health));
+        let response = app
+            .oneshot(
+                Request::builder()
+                    .uri("/health")
+                    .body(Body::empty())
+                    .unwrap(),
+            )
+            .await
+            .unwrap();
+
+        assert_eq!(response.status(), StatusCode::OK);
+        let body = to_bytes(response.into_body(), usize::MAX).await.unwrap();
+        let json: Value = serde_json::from_slice(&body).unwrap();
+        assert_eq!(json["status"], "ok");
+    }
+}