Integrate ComfyUI agents and video workflows

This commit is contained in:
boos
2026-05-22 04:46:15 +08:00
parent 3c12738138
commit c1074ddf71
35 changed files with 2527 additions and 213 deletions
+112 -23
View File
@@ -17,7 +17,7 @@ use crate::{
AppState,
};
use axum::{
extract::{Multipart, State},
extract::{Multipart, Query, State},
http::HeaderMap,
response::{IntoResponse, Response, Sse},
routing::{get, post},
@@ -60,6 +60,7 @@ pub fn routes() -> Router<AppState> {
.route("/v1/images/generations", post(images_generations))
.route("/v1/images/edits", post(images_edits))
.route("/v1/videos", post(videos_handler))
.route("/v1/videos/{id}", get(video_task_handler))
.route("/v1/models", get(list_models))
// ── ComfyUI shim aliases ──────────────────────────────────────
// ComfyUI's OpenAI nodes hit /proxy/openai/... paths (see
@@ -115,8 +116,7 @@ async fn responses_handler(
.unwrap_or(false);
if stream {
let chunks =
execute_responses_stream_with_conductor(state, key_id, model, body).await?;
let chunks = execute_responses_stream_with_conductor(state, key_id, model, body).await?;
let sse =
Sse::new(stream::iter(chunks).map(|chunk| {
Ok::<_, Infallible>(axum::response::sse::Event::default().data(chunk))
@@ -276,8 +276,7 @@ async fn audio_translations(
let route = resolve_route(&state, &model).await?;
let result =
executors::execute_proxy_request(state.clone(), route, "/audio/translations", body)
.await?;
executors::execute_proxy_request(state.clone(), route, "/audio/translations", body).await?;
record_billing(&state, key_id, &model, 0, 0, 0.0).await;
Ok(Json(result))
}
@@ -299,7 +298,33 @@ async fn videos_handler(
.to_string();
let route = resolve_route(&state, &model).await?;
let result = executors::execute_proxy_request(state.clone(), route, "/videos", body).await?;
let result = executors::execute_video_generation(state.clone(), route, body).await?;
record_billing(&state, key_id, &model, 0, 0, 0.0).await;
Ok(Json(result))
}
/// GET /v1/videos/{id} — inspect video generation task.
async fn video_task_handler(
State(state): State<AppState>,
headers: HeaderMap,
axum::extract::Path(id): axum::extract::Path<String>,
Query(query): Query<std::collections::HashMap<String, String>>,
) -> AppResult<Json<Value>> {
info!("GET /v1/videos/{id}");
let key_id = check_auth(&state, &headers).await?;
let model = query
.get("model")
.map(String::as_str)
.unwrap_or("default")
.to_string();
let route = resolve_route(&state, &model).await?;
let suffix = if route.base_url.contains("ark.cn-beijing.volces.com") {
format!("/contents/generations/tasks/{id}")
} else {
format!("/videos/{id}")
};
let result = executors::execute_proxy_get(route, &suffix).await?;
record_billing(&state, key_id, &model, 0, 0, 0.0).await;
Ok(Json(result))
}
@@ -426,7 +451,22 @@ async fn list_models(State(state): State<AppState>, headers: HeaderMap) -> AppRe
let _ = check_auth(&state, &headers).await?;
let rows: Vec<(String, String)> = sqlx::query_as(
"SELECT tr.model_pattern, COALESCE(tr.model_mapping, tr.model_pattern) FROM token_routes tr WHERE tr.enabled = 1 ORDER BY tr.id",
"SELECT DISTINCT
tr.model_pattern,
COALESCE(NULLIF(TRIM(tr.model_mapping), ''), tr.model_pattern)
FROM token_routes tr
JOIN route_channels rc ON rc.route_id = tr.id
JOIN accounts a ON a.id = rc.account_id
JOIN sites s ON s.id = a.site_id
LEFT JOIN account_tokens at ON at.id = rc.token_id AND at.enabled = 1
WHERE tr.enabled = 1
AND rc.enabled = 1
AND a.status = 'active'
AND s.status = 'active'
AND TRIM(COALESCE(s.url, '')) <> ''
AND (rc.cooldown_until IS NULL OR rc.cooldown_until <= datetime('now'))
AND (rc.token_id IS NULL OR at.id IS NOT NULL)
ORDER BY tr.id",
)
.fetch_all(&*state.db)
.await?;
@@ -671,6 +711,17 @@ fn build_conductor(state: &AppState) -> ProxyConductor {
}
}
fn upstream_status_from_error(err: &AppError) -> Option<u16> {
let message = err.to_string();
let marker = "Upstream returned ";
let start = message.find(marker)? + marker.len();
message
.get(start..)?
.split_whitespace()
.next()
.and_then(|value| value.parse::<u16>().ok())
}
async fn execute_chat_completions_with_conductor(
state: AppState,
key_id: Option<i64>,
@@ -703,7 +754,7 @@ async fn execute_chat_completions_with_conductor(
},
Err(err) => AttemptResult::Failure {
action: crate::proxy::conductor::AttemptFailureAction::Failover,
status: None,
status: upstream_status_from_error(&err),
raw_error_text: Some(err.to_string()),
},
}
@@ -734,7 +785,15 @@ async fn execute_chat_completions_with_conductor(
false,
)
.await;
record_billing(&billing_state, key_id, &billing_model, prompt, completion, 0.0).await;
record_billing(
&billing_state,
key_id,
&billing_model,
prompt,
completion,
0.0,
)
.await;
Ok(response)
}
ExecuteResult::Failure {
@@ -797,7 +856,7 @@ async fn execute_chat_completions_stream_with_conductor(
},
Err(err) => AttemptResult::Failure {
action: crate::proxy::conductor::AttemptFailureAction::Failover,
status: None,
status: upstream_status_from_error(&err),
raw_error_text: Some(err.to_string()),
},
}
@@ -814,14 +873,20 @@ async fn execute_chat_completions_stream_with_conductor(
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>>>()?),
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(),
@@ -888,7 +953,15 @@ async fn execute_chat_completions_stream_with_conductor(
true,
)
.await;
record_billing(&billing_state, key_id, &billing_model, prompt, completion, 0.0).await;
record_billing(
&billing_state,
key_id,
&billing_model,
prompt,
completion,
0.0,
)
.await;
Ok(chunks)
}
@@ -941,7 +1014,7 @@ async fn execute_responses_with_conductor(
},
Err(err) => AttemptResult::Failure {
action: crate::proxy::conductor::AttemptFailureAction::Failover,
status: None,
status: upstream_status_from_error(&err),
raw_error_text: Some(err.to_string()),
},
}
@@ -954,7 +1027,15 @@ async fn execute_responses_with_conductor(
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;
record_billing(
&billing_state,
key_id,
&billing_model,
prompt,
completion,
0.0,
)
.await;
Ok(response)
}
@@ -988,7 +1069,7 @@ async fn execute_responses_stream_with_conductor(
},
Err(err) => AttemptResult::Failure {
action: crate::proxy::conductor::AttemptFailureAction::Failover,
status: None,
status: upstream_status_from_error(&err),
raw_error_text: Some(err.to_string()),
},
}
@@ -1037,7 +1118,15 @@ async fn execute_responses_stream_with_conductor(
}
})
.unwrap_or((0, 0));
record_billing(&billing_state, key_id, &billing_model, prompt, completion, 0.0).await;
record_billing(
&billing_state,
key_id,
&billing_model,
prompt,
completion,
0.0,
)
.await;
Ok(chunks)
}