boos/BoosAPI
2
0
Fork 2
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: merge opus comfyui and login updates

Browse Source
This commit is contained in:
boos
2026-05-21 16:32:04 +02:00
parent c16479da7c
commit 9c84759498
105 changed files with 20160 additions and 1034 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/routes/proxy/mod.rs
+309 -13
View File
@@ -17,7 +17,7 @@ use crate::{
AppState,
};
use axum::{
extract::State,
extract::{Multipart, State},
http::HeaderMap,
response::{IntoResponse, Response, Sse},
routing::{get, post},
@@ -25,7 +25,7 @@ use axum::{
};
use futures::{stream, FutureExt, StreamExt};
use serde_json::{json, Value};
use std::{convert::Infallible, sync::Arc};
use std::{convert::Infallible, sync::Arc, time::Instant};
use tracing::info;
// ── Auth types ────────────────────────────────────────────────────────
@@ -58,6 +58,7 @@ pub fn routes() -> Router<AppState> {
.route("/v1/audio/translations", post(audio_translations))
.route("/v1/audio/speech", post(audio_speech))
.route("/v1/images/generations", post(images_generations))
.route("/v1/images/edits", post(images_edits))
.route("/v1/videos", post(videos_handler))
.route("/v1/models", get(list_models))
// ── ComfyUI shim aliases ──────────────────────────────────────
@@ -70,6 +71,10 @@ pub fn routes() -> Router<AppState> {
"/api/comfyui-agent/proxy/openai/images/generations",
post(images_generations),
)
.route(
"/api/comfyui-agent/proxy/openai/images/edits",
post(images_edits),
)
.route(
"/api/comfyui-agent/proxy/openai/v1/responses",
post(responses_handler),
@@ -321,6 +326,98 @@ async fn images_generations(
Ok(Json(result))
}
/// POST /v1/images/edits
///
/// `multipart/form-data` request: one or more image binaries
/// (`image`, `image[]`, `mask`) plus text fields (`prompt`, `model`,
/// `n`, `size`, `response_format`, `user`, `quality`, …). Mirrors
/// /images/generations for auth + routing + billing, but rebuilds a
/// fresh upstream multipart form rather than streaming the raw body —
/// that lets the route override `model` with the conductor-resolved
/// upstream model and skip frontend-only fields.
async fn images_edits(
State(state): State<AppState>,
headers: HeaderMap,
mut multipart: Multipart,
) -> AppResult<Json<Value>> {
info!("POST /v1/images/edits");
let key_id = check_auth(&state, &headers).await?;
let mut fields: Vec<executors::MultipartField> = Vec::new();
let mut model: Option<String> = None;
while let Some(field) = multipart
.next_field()
.await
.map_err(|e| AppError::BadRequest(format!("invalid multipart body: {e}")))?
{
let name = field
.name()
.ok_or_else(|| AppError::BadRequest("multipart field missing name".into()))?
.to_string();
let filename = field.file_name().map(str::to_string);
let content_type = field.content_type().map(str::to_string);
// OpenAI's spec lists exactly which fields carry binaries; everything
// else is a text field. We tag a field as binary when it has a
// filename _or_ a non-text content-type — covers both well-formed
// SDKs and minimal curl uploads.
let is_file = filename.is_some()
|| content_type
.as_deref()
.map(|ct| !ct.starts_with("text/") && ct != "application/x-www-form-urlencoded")
.unwrap_or(false);
if is_file {
let bytes = field
.bytes()
.await
.map_err(|e| AppError::BadRequest(format!("failed to read upload: {e}")))?
.to_vec();
fields.push(executors::MultipartField {
name,
bytes: Some(bytes),
filename,
content_type,
text: None,
});
} else {
let text = field
.text()
.await
.map_err(|e| AppError::BadRequest(format!("failed to read field: {e}")))?;
if name == "model" {
model = Some(text.clone());
}
fields.push(executors::MultipartField {
name,
bytes: None,
filename: None,
content_type: None,
text: Some(text),
});
}
}
let model = model.ok_or_else(|| AppError::BadRequest("model is required".into()))?;
let route = resolve_route(&state, &model).await?;
// Override the multipart `model` field with the conductor-resolved
// upstream model name (route may map "dall-e-3" → upstream-specific id).
if route.model != model {
for field in fields.iter_mut() {
if field.name == "model" && field.bytes.is_none() {
field.text = Some(route.model.clone());
}
}
}
let result = executors::execute_image_edit(state.clone(), route, fields).await?;
record_billing(&state, key_id, &model, 0, 0, 0.0).await;
Ok(Json(result))
}
/// GET /v1/models
async fn list_models(State(state): State<AppState>, headers: HeaderMap) -> AppResult<Json<Value>> {
info!("GET /v1/models");
@@ -401,6 +498,91 @@ async fn record_billing(
}
}
async fn record_proxy_log_success(
state: &AppState,
key_id: Option<i64>,
selected: &SelectedChannel,
requested_model: &str,
latency_ms: Option<i64>,
prompt_tokens: i64,
completion_tokens: i64,
cost: f64,
attempts: usize,
is_stream: bool,
) {
if let Err(err) = sqlx::query(
"INSERT INTO proxy_logs (
route_id, channel_id, account_id, downstream_api_key_id,
model_requested, model_actual, status, http_status,
latency_ms, prompt_tokens, completion_tokens, total_tokens, estimated_cost,
retry_count, is_stream
) VALUES (?, ?, ?, ?, ?, ?, 'success', 200, ?, ?, ?, ?, ?, ?, ?)",
)
.bind(selected.route_id)
.bind(selected.channel_id)
.bind(selected.account_id)
.bind(key_id)
.bind(requested_model)
.bind(&selected.actual_model)
.bind(latency_ms)
.bind(prompt_tokens)
.bind(completion_tokens)
.bind(prompt_tokens + completion_tokens)
.bind(cost)
.bind(attempts.saturating_sub(1) as i64)
.bind(if is_stream { 1 } else { 0 })
.execute(&*state.db)
.await
{
tracing::warn!(
%err,
model = requested_model,
channel_id = selected.channel_id,
"failed to record proxy_logs success entry",
);
}
}
async fn record_proxy_log_failure(
state: &AppState,
key_id: Option<i64>,
selected: Option<&SelectedChannel>,
requested_model: &str,
status: Option<u16>,
latency_ms: Option<i64>,
error_message: Option<&str>,
attempts: usize,
is_stream: bool,
) {
if let Err(err) = sqlx::query(
"INSERT INTO proxy_logs (
route_id, channel_id, account_id, downstream_api_key_id,
model_requested, model_actual, status, http_status,
latency_ms, error_message, retry_count, is_stream
) VALUES (?, ?, ?, ?, ?, ?, 'error', ?, ?, ?, ?, ?)",
)
.bind(selected.and_then(|item| item.route_id))
.bind(selected.map(|item| item.channel_id))
.bind(selected.and_then(|item| item.account_id))
.bind(key_id)
.bind(requested_model)
.bind(selected.map(|item| item.actual_model.as_str()))
.bind(status.map(i64::from))
.bind(latency_ms)
.bind(error_message.unwrap_or("proxy request failed"))
.bind(attempts.saturating_sub(1) as i64)
.bind(if is_stream { 1 } else { 0 })
.execute(&*state.db)
.await
{
tracing::warn!(
%err,
model = requested_model,
"failed to record proxy_logs failure entry",
);
}
}
/// Pull `usage` block out of an OpenAI-shaped response. Returns
/// (prompt_tokens, completion_tokens). Missing fields default to 0.
fn extract_usage(response: &Value) -> (i64, i64) {
@@ -419,6 +601,10 @@ fn extract_usage(response: &Value) -> (i64, i64) {
(prompt, completion)
}
fn elapsed_ms_since(started_at: Instant) -> i64 {
started_at.elapsed().as_millis().min(i64::MAX as u128) as i64
}
// ── Route Resolution ──────────────────────────────────────────────────
/// Resolve upstream endpoint + API key for a given model.
@@ -491,6 +677,7 @@ async fn execute_chat_completions_with_conductor(
model: String,
body: Value,
) -> AppResult<Value> {
let started_at = Instant::now();
let conductor = build_conductor(&state);
let billing_state = state.clone();
let billing_model = model.clone();
@@ -527,10 +714,55 @@ async fn execute_chat_completions_with_conductor(
})
.await;
let response = conductor_result_to_app_result(result)?;
let (prompt, completion) = extract_usage(&response);
record_billing(&billing_state, key_id, &billing_model, prompt, completion, 0.0).await;
Ok(response)
match result {
ExecuteResult::Success {
selected,
response,
attempts,
} => {
let (prompt, completion) = extract_usage(&response);
record_proxy_log_success(
&billing_state,
key_id,
&selected,
&billing_model,
Some(elapsed_ms_since(started_at)),
prompt,
completion,
0.0,
attempts,
false,
)
.await;
record_billing(&billing_state, key_id, &billing_model, prompt, completion, 0.0).await;
Ok(response)
}
ExecuteResult::Failure {
reason,
selected,
status,
raw_error_text,
attempts,
} => {
record_proxy_log_failure(
&billing_state,
key_id,
selected.as_ref(),
&billing_model,
status,
Some(elapsed_ms_since(started_at)),
raw_error_text.as_deref(),
attempts,
false,
)
.await;
Err(AppError::UpstreamError(format!(
"Proxy conductor failed after {attempts} attempts ({reason:?}, status: {:?}): {}",
status,
raw_error_text.unwrap_or_else(|| "no upstream error text".to_string())
)))
}
}
}
async fn execute_chat_completions_stream_with_conductor(
@@ -539,6 +771,7 @@ async fn execute_chat_completions_stream_with_conductor(
model: String,
body: Value,
) -> AppResult<Vec<String>> {
let started_at = Instant::now();
let conductor = build_conductor(&state);
let billing_state = state.clone();
let billing_model = model.clone();
@@ -575,19 +808,50 @@ async fn execute_chat_completions_stream_with_conductor(
})
.await;
let chunks = match conductor_result_to_app_result(result)? {
Value::Array(chunks) => chunks
let (selected, attempts, chunks) = match result {
ExecuteResult::Success {
selected,
response,
attempts,
} => match response {
Value::Array(chunks) => (selected, attempts, chunks
.into_iter()
.map(|chunk| {
chunk.as_str().map(str::to_string).ok_or_else(|| {
AppError::UpstreamError("Invalid streaming response from conductor".into())
})
})
.collect::<AppResult<Vec<String>>>()?,
_ => {
return Err(AppError::UpstreamError(
"Invalid streaming response from conductor".into(),
))
.collect::<AppResult<Vec<String>>>()?),
_ => {
return Err(AppError::UpstreamError(
"Invalid streaming response from conductor".into(),
))
}
},
ExecuteResult::Failure {
reason,
selected,
status,
raw_error_text,
attempts,
} => {
record_proxy_log_failure(
&billing_state,
key_id,
selected.as_ref(),
&billing_model,
status,
Some(elapsed_ms_since(started_at)),
raw_error_text.as_deref(),
attempts,
true,
)
.await;
return Err(AppError::UpstreamError(format!(
"Proxy conductor failed after {attempts} attempts ({reason:?}, status: {:?}): {}",
status,
raw_error_text.unwrap_or_else(|| "no upstream error text".to_string())
)));
}
};
@@ -611,6 +875,19 @@ async fn execute_chat_completions_stream_with_conductor(
}
})
.unwrap_or((0, 0));
record_proxy_log_success(
&billing_state,
key_id,
&selected,
&billing_model,
Some(elapsed_ms_since(started_at)),
prompt,
completion,
0.0,
attempts,
true,
)
.await;
record_billing(&billing_state, key_id, &billing_model, prompt, completion, 0.0).await;
Ok(chunks)
}
@@ -881,6 +1158,25 @@ mod tests {
.execute(&pool)
.await
.unwrap();
// The channel-selection query joins through site_api_endpoints to
// pick a per-site override URL. We don't need any rows — the
// COALESCE falls back to sites.url when this table is empty —
// but the table itself must exist or the SELECT errors out and
// the conductor sees `None` (NoChannel) instead of an attempt
// failure that bumps fail_count.
sqlx::query(
"CREATE TABLE site_api_endpoints (
id INTEGER PRIMARY KEY,
site_id INTEGER NOT NULL,
url TEXT NOT NULL,
enabled INTEGER DEFAULT 1,
sort_order INTEGER DEFAULT 0,
cooldown_until TEXT
)",
)
.execute(&pool)
.await
.unwrap();
sqlx::query("INSERT INTO downstream_api_keys (id, name, key, enabled) VALUES (1, 'test', 'sk-test', 1)")
.execute(&pool)