Nvidia/BossAI
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: complete Rust backend UI compatibility

Browse Source
This commit is contained in:
boos
2026-05-18 16:52:54 +08:00
parent f2434e725d
commit f986fe0377
41 changed files with 3585 additions and 811 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/routes/comfyui_agent.rs
+145 -78
View File
@@ -3,12 +3,13 @@ use axum::{
extract::State,
http::HeaderMap,
response::sse::{Event, Sse},
routing::post,
routing::{get, post},
Json, Router,
};
use futures::stream::Stream;
use regex::Regex;
use reqwest::Client;
use serde::Deserialize;
use serde_json::{json, Value};
use std::convert::Infallible;
use tokio::sync::mpsc::{unbounded_channel, UnboundedSender};
@@ -58,19 +59,68 @@ emotion: 平静
- 如果没有用户明确要求，切勿自作主张输出全部内容
"#;
fn upstream_base() -> String {
std::env::var("BOOSAPI_UPSTREAM_BASE")
.unwrap_or_else(|_| "https://api.boos.lat/v1".to_string())
.trim_end_matches('/')
.to_string()
}
fn upstream_key() -> String {
std::env::var("BOOSAPI_UPSTREAM_KEY").unwrap_or_else(|_| "sk-boos4721".to_string())
}
pub fn routes() -> Router<AppState> {
Router::new().route("/api/comfyui-agent/chat", post(chat_handler))
Router::new()
.route("/api/comfyui-agent/status", get(status_handler))
.route("/api/comfyui-agent/chat", post(chat_handler))
}
#[derive(Debug, Deserialize)]
struct AgentChatRequest {
messages: Vec<AgentMessage>,
}
#[derive(Debug, Deserialize)]
struct AgentMessage {
role: String,
content: String,
}
fn upstream_config(state: &AppState) -> (String, String) {
let base = std::env::var("BOOSAPI_UPSTREAM_BASE")
.ok()
.map(|value| value.trim().trim_end_matches('/').to_string())
.filter(|value| !value.is_empty())
.unwrap_or_else(|| "https://api.boos.lat/v1".to_string());
let key = std::env::var("BOOSAPI_UPSTREAM_KEY")
.ok()
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
.unwrap_or_else(|| state.config.auth.proxy_token.clone());
(base, key)
}
fn authorize(headers: &HeaderMap, state: &AppState) -> bool {
let auth_header = headers
.get("authorization")
.and_then(|value| value.to_str().ok())
.unwrap_or_default();
let token = auth_header
.strip_prefix("Bearer ")
.or_else(|| auth_header.strip_prefix("bearer "))
.unwrap_or_default();
token == state.config.auth.auth_token
}
async fn status_handler(
State(state): State<AppState>,
headers: HeaderMap,
) -> Result<Json<Value>, (axum::http::StatusCode, Json<Value>)> {
if !authorize(&headers, &state) {
return Err((
axum::http::StatusCode::UNAUTHORIZED,
Json(json!({ "error": "Unauthorized" })),
));
}
let (upstream_base, _) = upstream_config(&state);
Ok(Json(json!({
"success": true,
"model": state.config.comfyui.agent_model,
"upstreamBase": upstream_base,
"comfyui": state.comfyui.status_response().await,
})))
}
#[derive(Debug, Default)]
@@ -79,7 +129,6 @@ struct ParsedSections {
scenes: Vec<SceneSection>,
tts_list: Vec<TtsSection>,
workflow_json: Option<String>,
text_content: String,
}
#[derive(Debug, Default)]
@@ -108,7 +157,6 @@ struct TtsSection {
fn parse_sections(full_text: &str) -> ParsedSections {
let mut sections = ParsedSections::default();
sections.text_content = full_text.to_string();
// Parse CHARACTER blocks
let char_re = Regex::new(r"(?s)【CHARACTER】\s*(.*?)【/CHARACTER】").unwrap();
@@ -118,21 +166,35 @@ fn parse_sections(full_text: &str) -> ParsedSections {
for line in block.lines() {
let line = line.trim();
let line_lower = line.to_lowercase();
if let Some(val) = line.strip_prefix("name:").or_else(|| line.strip_prefix("name :")) {
if let Some(val) = line
.strip_prefix("name:")
.or_else(|| line.strip_prefix("name :"))
{
char_section.name = val.trim().to_string();
} else if let Some(val) = line.strip_prefix("description:").or_else(|| line.strip_prefix("description :")) {
} else if let Some(val) = line
.strip_prefix("description:")
.or_else(|| line.strip_prefix("description :"))
{
char_section.description = val.trim().to_string();
} else if line_lower.starts_with("prompt_front:") || line_lower.starts_with("prompt_front :") {
let val = line.splitn(2, ':').nth(1).unwrap_or("").trim();
} else if line_lower.starts_with("prompt_front:")
|| line_lower.starts_with("prompt_front :")
{
let val = line.split_once(':').map(|x| x.1).unwrap_or("").trim();
char_section.prompt_front = val.to_string();
} else if line_lower.starts_with("prompt_back:") || line_lower.starts_with("prompt_back :") {
let val = line.splitn(2, ':').nth(1).unwrap_or("").trim();
} else if line_lower.starts_with("prompt_back:")
|| line_lower.starts_with("prompt_back :")
{
let val = line.split_once(':').map(|x| x.1).unwrap_or("").trim();
char_section.prompt_back = val.to_string();
} else if line_lower.starts_with("prompt_left:") || line_lower.starts_with("prompt_left :") {
let val = line.splitn(2, ':').nth(1).unwrap_or("").trim();
} else if line_lower.starts_with("prompt_left:")
|| line_lower.starts_with("prompt_left :")
{
let val = line.split_once(':').map(|x| x.1).unwrap_or("").trim();
char_section.prompt_left = val.to_string();
} else if line_lower.starts_with("prompt_right:") || line_lower.starts_with("prompt_right :") {
let val = line.splitn(2, ':').nth(1).unwrap_or("").trim();
} else if line_lower.starts_with("prompt_right:")
|| line_lower.starts_with("prompt_right :")
{
let val = line.split_once(':').map(|x| x.1).unwrap_or("").trim();
char_section.prompt_right = val.to_string();
}
}
@@ -154,12 +216,15 @@ fn parse_sections(full_text: &str) -> ParsedSections {
let line = line.trim();
let line_lower = line.to_lowercase();
if line_lower.starts_with("id:") || line_lower.starts_with("id :") {
let val = line.splitn(2, ':').nth(1).unwrap_or("").trim();
let val = line.split_once(':').map(|x| x.1).unwrap_or("").trim();
scene.id = val.parse().unwrap_or(0);
} else if let Some(val) = line.strip_prefix("description:").or_else(|| line.strip_prefix("description :")) {
} else if let Some(val) = line
.strip_prefix("description:")
.or_else(|| line.strip_prefix("description :"))
{
scene.description = val.trim().to_string();
} else if line_lower.starts_with("prompt:") || line_lower.starts_with("prompt :") {
let val = line.splitn(2, ':').nth(1).unwrap_or("").trim();
let val = line.split_once(':').map(|x| x.1).unwrap_or("").trim();
scene.prompt = val.to_string();
}
}
@@ -181,13 +246,13 @@ fn parse_sections(full_text: &str) -> ParsedSections {
let line = line.trim();
let line_lower = line.to_lowercase();
if line_lower.starts_with("character:") || line_lower.starts_with("character :") {
let val = line.splitn(2, ':').nth(1).unwrap_or("").trim();
let val = line.split_once(':').map(|x| x.1).unwrap_or("").trim();
tts.character = val.to_string();
} else if line_lower.starts_with("text:") || line_lower.starts_with("text :") {
let val = line.splitn(2, ':').nth(1).unwrap_or("").trim();
let val = line.split_once(':').map(|x| x.1).unwrap_or("").trim();
tts.text = val.to_string();
} else if line_lower.starts_with("emotion:") || line_lower.starts_with("emotion :") {
let val = line.splitn(2, ':').nth(1).unwrap_or("").trim();
let val = line.split_once(':').map(|x| x.1).unwrap_or("").trim();
tts.emotion = val.to_string();
}
}
@@ -206,13 +271,13 @@ fn parse_sections(full_text: &str) -> ParsedSections {
sections
}
async fn call_image_gen(prompt: &str) -> Result<String, String> {
async fn call_image_gen(api_base: &str, api_key: &str, prompt: &str) -> Result<String, String> {
let client = Client::new();
let url = format!("{}/images/generations", upstream_base());
let url = format!("{}/images/generations", api_base.trim_end_matches('/'));
let resp = client
.post(&url)
.bearer_auth(upstream_key())
.bearer_auth(api_key)
.json(&json!({
"model": "gpt-image-2",
"prompt": prompt,
@@ -243,13 +308,18 @@ async fn call_image_gen(prompt: &str) -> Result<String, String> {
.to_string())
}
async fn call_tts(text: &str, voice: &str) -> Result<Vec<u8>, String> {
async fn call_tts(
api_base: &str,
api_key: &str,
text: &str,
voice: &str,
) -> Result<Vec<u8>, String> {
let client = Client::new();
let url = format!("{}/audio/speech", upstream_base());
let url = format!("{}/audio/speech", api_base.trim_end_matches('/'));
let resp = client
.post(&url)
.bearer_auth(upstream_key())
.bearer_auth(api_key)
.json(&json!({
"model": "tts-1",
"input": text,
@@ -274,58 +344,51 @@ async fn call_tts(text: &str, voice: &str) -> Result<Vec<u8>, String> {
}
fn send_sse(tx: &UnboundedSender<Result<Event, Infallible>>, payload: Value) -> bool {
tx.send(Ok(Event::default().data(payload.to_string()))).is_ok()
tx.send(Ok(Event::default().data(payload.to_string())))
.is_ok()
}
async fn chat_handler(
State(state): State<AppState>,
headers: HeaderMap,
Json(body): Json<Value>,
Json(body): Json<AgentChatRequest>,
) -> Result<Sse<impl Stream<Item = Result<Event, Infallible>>>, (axum::http::StatusCode, Json<Value>)>
{
let auth_header = headers
.get("authorization")
.and_then(|value| value.to_str().ok())
.unwrap_or_default();
let token = auth_header
.strip_prefix("Bearer ")
.or_else(|| auth_header.strip_prefix("bearer "))
.unwrap_or_default();
if token != state.config.auth.auth_token {
if !authorize(&headers, &state) {
return Err((
axum::http::StatusCode::UNAUTHORIZED,
Json(json!({ "error": "Unauthorized" })),
));
}
let messages = body
.get("messages")
.and_then(Value::as_array)
.ok_or_else(|| {
(
axum::http::StatusCode::BAD_REQUEST,
Json(json!({ "error": "messages is required" })),
)
})?;
for message in messages {
if message.get("role").and_then(Value::as_str).is_none()
|| message.get("content").and_then(Value::as_str).is_none()
{
return Err((
axum::http::StatusCode::BAD_REQUEST,
Json(json!({ "error": "messages is required" })),
));
}
if body.messages.is_empty()
|| body
.messages
.iter()
.any(|message| message.role.trim().is_empty() || message.content.trim().is_empty())
{
return Err((
axum::http::StatusCode::BAD_REQUEST,
Json(json!({ "error": "messages is required" })),
));
}
let messages = messages.clone();
let messages: Vec<Value> = body
.messages
.into_iter()
.map(|message| {
json!({
"role": message.role,
"content": message.content,
})
})
.collect();
let (api_base, api_key) = upstream_config(&state);
let agent_model = state.config.comfyui.agent_model.clone();
let (tx, rx) = unbounded_channel::<Result<Event, Infallible>>();
tokio::spawn(async move {
let upstream_url = format!("{}/chat/completions", upstream_base().trim_end_matches('/'));
let upstream_url = format!("{}/chat/completions", api_base.trim_end_matches('/'));
let mut upstream_messages = Vec::with_capacity(messages.len() + 1);
upstream_messages.push(json!({
@@ -340,9 +403,9 @@ async fn chat_handler(
let response = match client
.post(&upstream_url)
.bearer_auth(upstream_key())
.bearer_auth(&api_key)
.json(&json!({
"model": "gpt-5.5",
"model": agent_model,
"messages": upstream_messages,
"temperature": 0.7,
"stream": true,
@@ -516,7 +579,7 @@ async fn chat_handler(
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
}
match call_image_gen(prompt).await {
match call_image_gen(&api_base, &api_key, prompt).await {
Ok(url) if !url.is_empty() => {
send_sse(
&tx,
@@ -543,11 +606,14 @@ async fn chat_handler(
// Generate images for each scene
for scene in &sections.scenes {
info!("comfyui-agent: generating scene image for scene {}", scene.id);
info!(
"comfyui-agent: generating scene image for scene {}",
scene.id
);
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
match call_image_gen(&scene.prompt).await {
match call_image_gen(&api_base, &api_key, &scene.prompt).await {
Ok(url) if !url.is_empty() => {
send_sse(
&tx,
@@ -586,11 +652,12 @@ async fn chat_handler(
_ => "alloy",
};
match call_tts(&tts.text, voice).await {
match call_tts(&api_base, &api_key, &tts.text, voice).await {
Ok(audio_bytes) => {
// For SSE, we'll send a data URL for small audio, or a placeholder
// Since SSE can't easily send binary, we store to file and send URL
let audio_path = format!("/tmp/comfyui_tts_{}_{}.mp3",
let audio_path = format!(
"/tmp/comfyui_tts_{}_{}.mp3",
tts.character.replace(' ', "_"),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)