forked from boos/BoosAPI
feat: Rust refactor — full backend rewrite (Phases 1-6)
Phase 1-3: Scaffold + types + auth (JWT, Argon2, API keys) Phase 4a: Admin API (sites/accounts/tokens/downstream-keys/settings) Phase 5: Proxy core (/v1/chat, /v1/images, /v1/models) Phase 6: ComfyUI proxy routes migrated - 80+ source files, all cargo check + 18 tests green - Drop-in compatible with Node.js API surface
This commit is contained in:
@@ -0,0 +1,190 @@
|
||||
// Proxy routes — OpenAI-compatible /v1/ endpoints
|
||||
//
|
||||
// These routes handle:
|
||||
// POST /v1/chat/completions — streaming and non-streaming
|
||||
// POST /v1/images/generations — image generation
|
||||
// GET /v1/models — list available models
|
||||
|
||||
use crate::{
|
||||
auth::api_key::{ authorize_downstream_token, DownstreamAuthResult },
|
||||
error::{ AppError, AppResult },
|
||||
proxy::executors::{ self, UpstreamRoute },
|
||||
AppState,
|
||||
};
|
||||
use axum::{
|
||||
extract::State,
|
||||
http::HeaderMap,
|
||||
response::{ IntoResponse, Response, Sse },
|
||||
routing::{ get, post },
|
||||
Json, Router,
|
||||
};
|
||||
use futures::stream;
|
||||
use futures::StreamExt;
|
||||
use serde_json::{ json, Value };
|
||||
use std::convert::Infallible;
|
||||
use tracing::info;
|
||||
|
||||
// ── Auth types ────────────────────────────────────────────────────────
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct ProxyAuth {
|
||||
pub token: String,
|
||||
pub key_id: Option<i64>,
|
||||
pub source: DownstreamSource,
|
||||
}
|
||||
|
||||
#[derive(Clone, PartialEq)]
|
||||
pub enum DownstreamSource {
|
||||
Managed { key_id: i64 },
|
||||
Global,
|
||||
}
|
||||
|
||||
// ── Route definitions ─────────────────────────────────────────────────
|
||||
|
||||
/// Build the /v1/ proxy router
|
||||
pub fn routes() -> Router<AppState> {
|
||||
Router::new()
|
||||
.route("/v1/chat/completions", post(chat_completions))
|
||||
.route("/v1/images/generations", post(images_generations))
|
||||
.route("/v1/models", get(list_models))
|
||||
}
|
||||
|
||||
// ── Handlers ───────────────────────────────────────────────────────────
|
||||
|
||||
/// POST /v1/chat/completions
|
||||
async fn chat_completions(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
Json(body): Json<Value>,
|
||||
) -> Result<Response, AppError> {
|
||||
info!("POST /v1/chat/completions");
|
||||
|
||||
check_auth(&state, &headers).await?;
|
||||
|
||||
let model = body
|
||||
.get("model")
|
||||
.and_then(|m| m.as_str())
|
||||
.ok_or_else(|| AppError::BadRequest("model is required".into()))?
|
||||
.to_string();
|
||||
|
||||
let route = resolve_route(&state, &model).await?;
|
||||
let stream = body.get("stream").and_then(|s| s.as_bool()).unwrap_or(false);
|
||||
|
||||
if stream {
|
||||
let chunks = executors::execute_chat_completions_stream(state, route, body).await?;
|
||||
let sse = Sse::new(stream::iter(chunks).map(|chunk| {
|
||||
Ok::<_, Infallible>(axum::response::sse::Event::default().data(chunk))
|
||||
}));
|
||||
Ok(sse.into_response())
|
||||
} else {
|
||||
let result = executors::execute_chat_completions(state, route, body).await?;
|
||||
Ok(Json(result).into_response())
|
||||
}
|
||||
}
|
||||
|
||||
/// POST /v1/images/generations
|
||||
async fn images_generations(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
Json(body): Json<Value>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
info!("POST /v1/images/generations");
|
||||
|
||||
check_auth(&state, &headers).await?;
|
||||
|
||||
let model = body
|
||||
.get("model")
|
||||
.and_then(|m| m.as_str())
|
||||
.ok_or_else(|| AppError::BadRequest("model is required".into()))?
|
||||
.to_string();
|
||||
|
||||
let route = resolve_route(&state, &model).await?;
|
||||
let result = executors::execute_image_generation(state, route, body).await?;
|
||||
Ok(Json(result))
|
||||
}
|
||||
|
||||
/// GET /v1/models
|
||||
async fn list_models(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
) -> AppResult<Json<Value>> {
|
||||
info!("GET /v1/models");
|
||||
|
||||
check_auth(&state, &headers).await?;
|
||||
|
||||
let rows: Vec<(String, String)> = sqlx::query_as(
|
||||
"SELECT tr.name, tr.model FROM token_routes tr WHERE tr.enabled = 1 ORDER BY tr.name"
|
||||
)
|
||||
.fetch_all(&*state.db)
|
||||
.await?;
|
||||
|
||||
let models: Vec<Value> = rows
|
||||
.into_iter()
|
||||
.map(|(name, model)| {
|
||||
json!({
|
||||
"id": model,
|
||||
"object": "model",
|
||||
"created": 0,
|
||||
"owned_by": name,
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
Ok(Json(json!({
|
||||
"object": "list",
|
||||
"data": models,
|
||||
})))
|
||||
}
|
||||
|
||||
// ── Auth helper ───────────────────────────────────────────────────────
|
||||
|
||||
async fn check_auth(state: &AppState, headers: &HeaderMap) -> Result<(), AppError> {
|
||||
let auth = headers
|
||||
.get("Authorization")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.and_then(|v| v.strip_prefix("Bearer "))
|
||||
.and_then(|v| v.trim().split_whitespace().next())
|
||||
.ok_or_else(|| AppError::Unauthorized("Missing Authorization header".into()))?;
|
||||
|
||||
let result = authorize_downstream_token(auth, &state.db).await
|
||||
.map_err(|e| AppError::Unauthorized(e))?;
|
||||
|
||||
match result {
|
||||
DownstreamAuthResult::Managed { key, .. } => {
|
||||
if !key.enabled.unwrap_or(true) {
|
||||
return Err(AppError::Forbidden("API key is disabled".into()));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
DownstreamAuthResult::Global { .. } => Ok(()),
|
||||
}
|
||||
}
|
||||
|
||||
// ── Route Resolution ──────────────────────────────────────────────────
|
||||
|
||||
/// Resolve upstream endpoint + API key for a given model.
|
||||
async fn resolve_route(state: &AppState, model: &str) -> AppResult<UpstreamRoute> {
|
||||
let row: Option<(String, String)> = sqlx::query_as(
|
||||
"SELECT s.base_url, COALESCE(a.api_key, s.default_key, '') as api_key
|
||||
FROM token_routes tr
|
||||
JOIN sites s ON s.id = tr.site_id
|
||||
LEFT JOIN accounts a ON a.id = tr.account_id AND a.site_id = s.id
|
||||
WHERE tr.enabled = 1 AND tr.model = ?
|
||||
LIMIT 1"
|
||||
)
|
||||
.bind(model)
|
||||
.fetch_optional(&*state.db)
|
||||
.await?;
|
||||
|
||||
match row {
|
||||
Some((base_url, api_key)) => Ok(UpstreamRoute {
|
||||
base_url,
|
||||
api_key,
|
||||
model: model.to_string(),
|
||||
}),
|
||||
None => Err(AppError::BadRequest(format!(
|
||||
"No route configured for model '{}'",
|
||||
model
|
||||
))),
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user