Integrate ComfyUI agents and video workflows
This commit is contained in:
+112
-23
@@ -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)
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user