boos
636 lines
18 KiB
Rust
636 lines
18 KiB
Rust
use axum::{
|
||
body::Body,
|
||
extract::State,
|
||
http::{header, HeaderMap, Response, StatusCode},
|
||
response::{
|
||
sse::{Event, Sse},
|
||
IntoResponse,
|
||
},
|
||
routing::{get, post},
|
||
Json, Router,
|
||
};
|
||
use futures::stream;
|
||
use serde_json::{json, Value};
|
||
use std::convert::Infallible;
|
||
use tracing::info;
|
||
|
||
use crate::{
|
||
error::{AppError, AppResult},
|
||
AppState,
|
||
};
|
||
|
||
const SCRIPT_ANALYSIS_PROMPT: &str = r#"你是一个 AI 视频生成助手。分析用户提供的脚本，提取以下结构化信息，以 JSON 格式返回：
|
||
|
||
{
|
||
"title": "视频标题",
|
||
"characters": [
|
||
{
|
||
"name": "角色名",
|
||
"gender": "男/女",
|
||
"age": "年龄描述",
|
||
"appearance": "外貌详细描述（发型、脸型、服装、体型等）",
|
||
"personality": "性格特征"
|
||
}
|
||
],
|
||
"scenes": [
|
||
{
|
||
"id": 1,
|
||
"description": "场景描述",
|
||
"environment": "环境详细描述（时间、地点、天气、光线等）",
|
||
"characters": ["出现的角色名"],
|
||
"action": "场景中发生的动作",
|
||
"dialogue": [
|
||
{"character": "角色名", "line": "台词", "emotion": "情感语气"}
|
||
]
|
||
}
|
||
],
|
||
"style": "整体视觉风格（动画/写实/像素等）"
|
||
}
|
||
|
||
只返回 JSON，不要其他文字。"#;
|
||
|
||
const CHARACTER_PROMPT_GEN: &str = r#"根据角色描述生成 4 组稳定的图像生成提示词（正面、背面、左侧面、右侧面），用于 AI 图像生成模型。
|
||
每组提示词应包含：人物姿态、服装、光照、背景。
|
||
返回 JSON 格式：
|
||
{
|
||
"front": "正面提示词",
|
||
"back": "背面提示词",
|
||
"left": "左侧面提示词",
|
||
"right": "右侧面提示词"
|
||
}
|
||
只返回 JSON。"#;
|
||
|
||
const SCENE_PROMPT_GEN: &str = r#"根据场景描述生成图像生成提示词，用于 AI 图像生成模型创建场景背景。
|
||
包含：环境、光照、色调、氛围。
|
||
返回 JSON：
|
||
{
|
||
"prompt": "场景生成提示词",
|
||
"negative_prompt": "负面提示词"
|
||
}
|
||
只返回 JSON。"#;
|
||
|
||
pub fn routes() -> Router<AppState> {
|
||
Router::new()
|
||
.route("/api/video-agent/status", get(status))
|
||
.route("/api/video-agent/start", post(start_comfyui))
|
||
.route("/api/video-agent/stop", post(stop_comfyui))
|
||
.route("/api/video-agent/restart", post(restart_comfyui))
|
||
.route("/api/video-agent/chat", post(chat))
|
||
.route("/api/video-agent/analyze", post(analyze))
|
||
.route(
|
||
"/api/video-agent/character-prompts",
|
||
post(character_prompts),
|
||
)
|
||
.route(
|
||
"/api/video-agent/generate-character",
|
||
post(generate_character),
|
||
)
|
||
.route("/api/video-agent/scene-prompt", post(scene_prompt))
|
||
.route("/api/video-agent/generate-scene", post(generate_scene))
|
||
.route("/api/video-agent/generate-tts", post(generate_tts))
|
||
.route("/api/video-agent/full-workflow", post(full_workflow))
|
||
}
|
||
|
||
async fn status(State(state): State<AppState>, headers: HeaderMap) -> AppResult<Json<Value>> {
|
||
check_auth(&headers, &state)?;
|
||
Ok(Json(external_comfyui_status(&state).await))
|
||
}
|
||
|
||
async fn start_comfyui(
|
||
State(state): State<AppState>,
|
||
headers: HeaderMap,
|
||
) -> AppResult<Json<Value>> {
|
||
check_auth(&headers, &state)?;
|
||
Ok(Json(json!({
|
||
"success": true,
|
||
"message": "内置 Docker ComfyUI 工作台已移除；请配置外部 ComfyUI 或使用 /workflows 工作流画布。",
|
||
"comfyui": external_comfyui_status_value(&state).await
|
||
})))
|
||
}
|
||
|
||
async fn stop_comfyui(State(state): State<AppState>, headers: HeaderMap) -> AppResult<Json<Value>> {
|
||
check_auth(&headers, &state)?;
|
||
Ok(Json(json!({
|
||
"success": true,
|
||
"message": "内置 Docker ComfyUI 工作台已移除；BoosAPI 不再管理 ComfyUI 进程生命周期。",
|
||
"comfyui": external_comfyui_status_value(&state).await
|
||
})))
|
||
}
|
||
|
||
async fn restart_comfyui(
|
||
State(state): State<AppState>,
|
||
headers: HeaderMap,
|
||
) -> AppResult<Json<Value>> {
|
||
check_auth(&headers, &state)?;
|
||
Ok(Json(json!({
|
||
"success": true,
|
||
"message": "内置 Docker ComfyUI 工作台已移除；请在外部 ComfyUI 服务侧重启。",
|
||
"comfyui": external_comfyui_status_value(&state).await
|
||
})))
|
||
}
|
||
|
||
async fn chat(
|
||
State(state): State<AppState>,
|
||
headers: HeaderMap,
|
||
Json(body): Json<Value>,
|
||
) -> AppResult<Sse<impl futures::Stream<Item = Result<Event, Infallible>>>> {
|
||
check_auth(&headers, &state)?;
|
||
|
||
let prompt = body
|
||
.get("message")
|
||
.and_then(Value::as_str)
|
||
.or_else(|| {
|
||
body.get("messages")
|
||
.and_then(Value::as_array)
|
||
.and_then(|messages| messages.last())
|
||
.and_then(|message| message.get("content"))
|
||
.and_then(Value::as_str)
|
||
})
|
||
.unwrap_or("生成一个短视频工作流");
|
||
|
||
let workflow = build_basic_workflow(prompt);
|
||
let events = vec![
|
||
json!({
|
||
"type": "text",
|
||
"content": "已收到需求，正在生成 ComfyUI workflow..."
|
||
}),
|
||
json!({
|
||
"type": "workflow_generated",
|
||
"workflow": workflow
|
||
}),
|
||
json!({
|
||
"type": "done"
|
||
}),
|
||
];
|
||
|
||
Ok(Sse::new(stream::iter(events.into_iter().map(|event| {
|
||
Ok(Event::default().data(event.to_string()))
|
||
}))))
|
||
}
|
||
|
||
fn local_api_base(state: &AppState) -> String {
|
||
state.comfyui.model_backend_url()
|
||
}
|
||
|
||
async fn external_comfyui_status(state: &AppState) -> Value {
|
||
json!({
|
||
"success": true,
|
||
"comfyui": external_comfyui_status_value(state).await,
|
||
})
|
||
}
|
||
|
||
async fn external_comfyui_status_value(state: &AppState) -> Value {
|
||
let status = state.comfyui.status_response().await;
|
||
let reachable = state.comfyui.health_check().await;
|
||
json!({
|
||
"status": if reachable { "running" } else { "unavailable" },
|
||
"pid": null,
|
||
"uptime_secs": 0,
|
||
"mode": "external",
|
||
"model_backend": status.model_backend,
|
||
"base_url": status.base_url,
|
||
"version": status.version,
|
||
"last_error": if reachable {
|
||
Value::Null
|
||
} else {
|
||
json!("外部 ComfyUI 不可达；内置 Docker ComfyUI 工作台已移除，请配置 comfyui.base_url 或使用 /workflows 工作流画布。")
|
||
},
|
||
"dockerManaged": false,
|
||
})
|
||
}
|
||
|
||
fn check_auth(headers: &HeaderMap, state: &AppState) -> AppResult<()> {
|
||
let expected = format!("Bearer {}", state.config.auth.auth_token);
|
||
|
||
let actual = headers
|
||
.get(header::AUTHORIZATION)
|
||
.and_then(|v| v.to_str().ok())
|
||
.unwrap_or("");
|
||
|
||
if actual != expected {
|
||
return Err(AppError::BadRequest("unauthorized".to_string()));
|
||
}
|
||
|
||
Ok(())
|
||
}
|
||
|
||
fn build_basic_workflow(prompt: &str) -> Value {
|
||
json!({
|
||
"boosapi": {
|
||
"kind": "video-agent-workflow",
|
||
"version": 1,
|
||
"prompt": prompt,
|
||
},
|
||
"nodes": {},
|
||
})
|
||
}
|
||
|
||
fn required_str<'a>(body: &'a Value, key: &str) -> AppResult<&'a str> {
|
||
body.get(key)
|
||
.and_then(Value::as_str)
|
||
.filter(|s| !s.trim().is_empty())
|
||
.ok_or_else(|| AppError::BadRequest(format!("{key} is required")))
|
||
}
|
||
|
||
fn required_value<'a>(body: &'a Value, key: &str) -> AppResult<&'a Value> {
|
||
body.get(key)
|
||
.filter(|v| !v.is_null())
|
||
.ok_or_else(|| AppError::BadRequest(format!("{key} is required")))
|
||
}
|
||
|
||
fn strip_json_markers(s: &str) -> String {
|
||
let mut text = s.trim().to_string();
|
||
|
||
if text.starts_with("```json") {
|
||
text = text.trim_start_matches("```json").to_string();
|
||
} else if text.starts_with("```") {
|
||
text = text.trim_start_matches("```").to_string();
|
||
}
|
||
|
||
if text.ends_with("```") {
|
||
text = text.trim_end_matches("```").to_string();
|
||
}
|
||
|
||
text.trim().to_string()
|
||
}
|
||
|
||
fn parse_llm_json(raw: &str) -> AppResult<Value> {
|
||
let cleaned = strip_json_markers(raw);
|
||
serde_json::from_str::<Value>(&cleaned).map_err(|e| {
|
||
AppError::Internal(anyhow::anyhow!(
|
||
"failed to parse LLM JSON response: {}; response: {}",
|
||
e,
|
||
cleaned
|
||
))
|
||
})
|
||
}
|
||
|
||
async fn call_llm(state: &AppState, prompt: &str, system: &str) -> AppResult<String> {
|
||
let client = reqwest::Client::new();
|
||
let url = format!("{}/chat/completions", local_api_base(state));
|
||
|
||
info!("video-agent calling BoosAPI LLM endpoint: {}", url);
|
||
|
||
let resp = client
|
||
.post(url)
|
||
.bearer_auth(state.runtime_proxy_token())
|
||
.json(&json!({
|
||
"model": &state.config.comfyui.agent_model,
|
||
"messages": [
|
||
{ "role": "system", "content": system },
|
||
{ "role": "user", "content": prompt }
|
||
],
|
||
"temperature": 0.7
|
||
}))
|
||
.send()
|
||
.await
|
||
.map_err(|e| AppError::Internal(anyhow::anyhow!("LLM request failed: {}", e)))?;
|
||
|
||
let status = resp.status();
|
||
let data: Value = resp
|
||
.json()
|
||
.await
|
||
.map_err(|e| AppError::Internal(anyhow::anyhow!("LLM response parse failed: {}", e)))?;
|
||
|
||
if !status.is_success() {
|
||
return Err(AppError::Internal(anyhow::anyhow!(
|
||
"LLM upstream returned {}: {}",
|
||
status,
|
||
data
|
||
)));
|
||
}
|
||
|
||
Ok(data
|
||
.get("choices")
|
||
.and_then(Value::as_array)
|
||
.and_then(|choices| choices.first())
|
||
.and_then(|choice| choice.get("message"))
|
||
.and_then(|message| message.get("content"))
|
||
.and_then(Value::as_str)
|
||
.unwrap_or("")
|
||
.to_string())
|
||
}
|
||
|
||
async fn call_image_gen(state: &AppState, prompt: &str) -> AppResult<String> {
|
||
let client = reqwest::Client::new();
|
||
let url = format!("{}/images/generations", local_api_base(state));
|
||
|
||
info!("video-agent calling BoosAPI image endpoint: {}", url);
|
||
|
||
let resp = client
|
||
.post(url)
|
||
.bearer_auth(state.runtime_proxy_token())
|
||
.json(&json!({
|
||
"model": "gpt-image-2",
|
||
"prompt": prompt,
|
||
"n": 1,
|
||
"size": "1024x1024"
|
||
}))
|
||
.send()
|
||
.await
|
||
.map_err(|e| AppError::Internal(anyhow::anyhow!("image request failed: {}", e)))?;
|
||
|
||
let status = resp.status();
|
||
let data: Value = resp
|
||
.json()
|
||
.await
|
||
.map_err(|e| AppError::Internal(anyhow::anyhow!("image response parse failed: {}", e)))?;
|
||
|
||
if !status.is_success() {
|
||
return Err(AppError::Internal(anyhow::anyhow!(
|
||
"image upstream returned {}: {}",
|
||
status,
|
||
data
|
||
)));
|
||
}
|
||
|
||
Ok(data
|
||
.get("data")
|
||
.and_then(Value::as_array)
|
||
.and_then(|items| items.first())
|
||
.and_then(|item| item.get("url"))
|
||
.and_then(Value::as_str)
|
||
.unwrap_or("")
|
||
.to_string())
|
||
}
|
||
|
||
async fn call_tts(state: &AppState, text: &str, voice: &str) -> AppResult<Vec<u8>> {
|
||
let client = reqwest::Client::new();
|
||
let url = format!("{}/audio/speech", local_api_base(state));
|
||
|
||
info!("video-agent calling BoosAPI TTS endpoint: {}", url);
|
||
|
||
let resp = client
|
||
.post(url)
|
||
.bearer_auth(state.runtime_proxy_token())
|
||
.json(&json!({
|
||
"model": "tts-1",
|
||
"input": text,
|
||
"voice": voice
|
||
}))
|
||
.send()
|
||
.await
|
||
.map_err(|e| AppError::Internal(anyhow::anyhow!("TTS request failed: {}", e)))?;
|
||
|
||
let status = resp.status();
|
||
|
||
if !status.is_success() {
|
||
let body = resp.text().await.unwrap_or_default();
|
||
return Err(AppError::Internal(anyhow::anyhow!(
|
||
"TTS upstream returned {}: {}",
|
||
status,
|
||
body
|
||
)));
|
||
}
|
||
|
||
let bytes = resp
|
||
.bytes()
|
||
.await
|
||
.map_err(|e| AppError::Internal(anyhow::anyhow!("TTS response read failed: {}", e)))?;
|
||
|
||
Ok(bytes.to_vec())
|
||
}
|
||
|
||
async fn analyze(
|
||
State(state): State<AppState>,
|
||
headers: HeaderMap,
|
||
Json(body): Json<Value>,
|
||
) -> AppResult<Json<Value>> {
|
||
check_auth(&headers, &state)?;
|
||
|
||
let script = required_str(&body, "script")?;
|
||
|
||
let result = call_llm(&state, script, SCRIPT_ANALYSIS_PROMPT).await?;
|
||
let parsed = parse_llm_json(&result)?;
|
||
|
||
Ok(Json(json!({
|
||
"success": true,
|
||
"data": parsed
|
||
})))
|
||
}
|
||
|
||
async fn character_prompts(
|
||
State(state): State<AppState>,
|
||
headers: HeaderMap,
|
||
Json(body): Json<Value>,
|
||
) -> AppResult<Json<Value>> {
|
||
check_auth(&headers, &state)?;
|
||
|
||
let character = required_value(&body, "character")?;
|
||
|
||
let name = character
|
||
.get("name")
|
||
.and_then(Value::as_str)
|
||
.unwrap_or_default();
|
||
let appearance = character
|
||
.get("appearance")
|
||
.and_then(Value::as_str)
|
||
.unwrap_or_default();
|
||
let gender = character
|
||
.get("gender")
|
||
.and_then(Value::as_str)
|
||
.unwrap_or_default();
|
||
let age = character
|
||
.get("age")
|
||
.and_then(Value::as_str)
|
||
.unwrap_or_default();
|
||
|
||
let desc = format!("{name}: {appearance}, {gender}, {age}");
|
||
let result = call_llm(&state, &desc, CHARACTER_PROMPT_GEN).await?;
|
||
let prompts = parse_llm_json(&result)?;
|
||
|
||
Ok(Json(json!({
|
||
"success": true,
|
||
"data": prompts
|
||
})))
|
||
}
|
||
|
||
async fn generate_character(
|
||
State(state): State<AppState>,
|
||
headers: HeaderMap,
|
||
Json(body): Json<Value>,
|
||
) -> AppResult<Json<Value>> {
|
||
check_auth(&headers, &state)?;
|
||
|
||
let prompts = required_value(&body, "prompts")?;
|
||
let prompt_obj = prompts
|
||
.as_object()
|
||
.ok_or_else(|| AppError::BadRequest("prompts must be an object".to_string()))?;
|
||
|
||
let mut results = serde_json::Map::new();
|
||
|
||
for (angle, prompt_value) in prompt_obj {
|
||
let prompt = prompt_value.as_str().ok_or_else(|| {
|
||
AppError::BadRequest(format!("prompt for angle '{angle}' must be a string"))
|
||
})?;
|
||
|
||
let url = call_image_gen(&state, prompt).await?;
|
||
results.insert(angle.clone(), Value::String(url));
|
||
}
|
||
|
||
Ok(Json(json!({
|
||
"success": true,
|
||
"data": Value::Object(results)
|
||
})))
|
||
}
|
||
|
||
async fn scene_prompt(
|
||
State(state): State<AppState>,
|
||
headers: HeaderMap,
|
||
Json(body): Json<Value>,
|
||
) -> AppResult<Json<Value>> {
|
||
check_auth(&headers, &state)?;
|
||
|
||
let scene = required_value(&body, "scene")?;
|
||
|
||
let environment = scene
|
||
.get("environment")
|
||
.and_then(Value::as_str)
|
||
.unwrap_or_default();
|
||
let description = scene
|
||
.get("description")
|
||
.and_then(Value::as_str)
|
||
.unwrap_or_default();
|
||
|
||
let desc = format!("{environment}. {description}");
|
||
let result = call_llm(&state, &desc, SCENE_PROMPT_GEN).await?;
|
||
let prompt = parse_llm_json(&result)?;
|
||
|
||
Ok(Json(json!({
|
||
"success": true,
|
||
"data": prompt
|
||
})))
|
||
}
|
||
|
||
async fn generate_scene(
|
||
State(state): State<AppState>,
|
||
headers: HeaderMap,
|
||
Json(body): Json<Value>,
|
||
) -> AppResult<Json<Value>> {
|
||
check_auth(&headers, &state)?;
|
||
|
||
let prompt = required_str(&body, "prompt")?;
|
||
let url = call_image_gen(&state, prompt).await?;
|
||
|
||
Ok(Json(json!({
|
||
"success": true,
|
||
"data": {
|
||
"url": url
|
||
}
|
||
})))
|
||
}
|
||
|
||
async fn generate_tts(
|
||
State(state): State<AppState>,
|
||
headers: HeaderMap,
|
||
Json(body): Json<Value>,
|
||
) -> AppResult<Response<Body>> {
|
||
check_auth(&headers, &state)?;
|
||
|
||
let text = required_str(&body, "text")?;
|
||
let voice = body
|
||
.get("voice")
|
||
.and_then(Value::as_str)
|
||
.filter(|s| !s.trim().is_empty())
|
||
.unwrap_or("alloy");
|
||
|
||
let audio = call_tts(&state, text, voice).await?;
|
||
|
||
Ok((
|
||
StatusCode::OK,
|
||
[(header::CONTENT_TYPE, "audio/mpeg")],
|
||
audio,
|
||
)
|
||
.into_response())
|
||
}
|
||
|
||
async fn full_workflow(
|
||
State(state): State<AppState>,
|
||
headers: HeaderMap,
|
||
Json(body): Json<Value>,
|
||
) -> AppResult<Json<Value>> {
|
||
check_auth(&headers, &state)?;
|
||
|
||
let script = required_str(&body, "script")?;
|
||
let requested_style = body.get("style").and_then(Value::as_str);
|
||
|
||
let analysis_raw = call_llm(&state, script, SCRIPT_ANALYSIS_PROMPT).await?;
|
||
let analysis = parse_llm_json(&analysis_raw)?;
|
||
|
||
let mut characters = Vec::new();
|
||
|
||
if let Some(chars) = analysis.get("characters").and_then(Value::as_array) {
|
||
for character in chars {
|
||
let name = character
|
||
.get("name")
|
||
.and_then(Value::as_str)
|
||
.unwrap_or_default();
|
||
let appearance = character
|
||
.get("appearance")
|
||
.and_then(Value::as_str)
|
||
.unwrap_or_default();
|
||
let gender = character
|
||
.get("gender")
|
||
.and_then(Value::as_str)
|
||
.unwrap_or_default();
|
||
let age = character
|
||
.get("age")
|
||
.and_then(Value::as_str)
|
||
.unwrap_or_default();
|
||
|
||
let desc = format!("{name}: {appearance}, {gender}, {age}");
|
||
let prompts_raw = call_llm(&state, &desc, CHARACTER_PROMPT_GEN).await?;
|
||
let prompts = parse_llm_json(&prompts_raw)?;
|
||
|
||
let mut item = character.clone();
|
||
if let Some(obj) = item.as_object_mut() {
|
||
obj.insert("imagePrompts".to_string(), prompts);
|
||
}
|
||
|
||
characters.push(item);
|
||
}
|
||
}
|
||
|
||
let mut scenes = Vec::new();
|
||
|
||
if let Some(scene_items) = analysis.get("scenes").and_then(Value::as_array) {
|
||
for scene in scene_items {
|
||
let environment = scene
|
||
.get("environment")
|
||
.and_then(Value::as_str)
|
||
.unwrap_or_default();
|
||
let description = scene
|
||
.get("description")
|
||
.and_then(Value::as_str)
|
||
.unwrap_or_default();
|
||
|
||
let desc = format!("{environment}. {description}");
|
||
let scene_prompt_raw = call_llm(&state, &desc, SCENE_PROMPT_GEN).await?;
|
||
let image_prompt = parse_llm_json(&scene_prompt_raw)?;
|
||
|
||
let mut item = scene.clone();
|
||
if let Some(obj) = item.as_object_mut() {
|
||
obj.insert("imagePrompt".to_string(), image_prompt);
|
||
}
|
||
|
||
scenes.push(item);
|
||
}
|
||
}
|
||
|
||
let style = analysis
|
||
.get("style")
|
||
.cloned()
|
||
.or_else(|| requested_style.map(|s| Value::String(s.to_string())))
|
||
.unwrap_or(Value::Null);
|
||
|
||
Ok(Json(json!({
|
||
"success": true,
|
||
"data": {
|
||
"title": analysis.get("title").cloned().unwrap_or(Value::Null),
|
||
"style": style,
|
||
"characters": characters,
|
||
"scenes": scenes
|
||
}
|
||
})))
|
||
}
|