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

feat: add workflow studio and postgres support

Browse Source
This commit is contained in:
boos
2026-05-25 16:03:07 +08:00
parent 7cdb02ef1d
commit cb35d5fc6e
96 changed files with 12414 additions and 15548 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/routes/comfyui_agent.rs
+152 -8
View File
@@ -63,6 +63,7 @@ pub fn routes() -> Router<AppState> {
Router::new()
.route("/api/comfyui-agent/status", get(status_handler))
.route("/api/comfyui-agent/models", get(models_handler))
.route("/api/comfyui-agent/image-models", get(image_models_handler))
.route("/api/comfyui-agent/chat", post(chat_handler))
}
@@ -70,6 +71,8 @@ pub fn routes() -> Router<AppState> {
struct AgentChatRequest {
messages: Vec<AgentMessage>,
model: Option<String>,
#[serde(rename = "imageModel")]
image_model: Option<String>,
}
#[derive(Debug, Deserialize)]
@@ -85,6 +88,10 @@ struct AgentMessage {
struct AgentAttachment {
name: String,
mime_type: Option<String>,
#[serde(default, rename = "mediaKind")]
media_kind: Option<String>,
#[serde(default, rename = "referenceLabel")]
reference_label: Option<String>,
text: Option<String>,
data_url: Option<String>,
}
@@ -290,6 +297,52 @@ fn is_exact_model_pattern(value: &str) -> bool {
&& !normalized.contains('?')
}
fn is_image_model_pattern(value: &str) -> bool {
if !is_exact_model_pattern(value) {
return false;
}
let normalized = value.trim().to_ascii_lowercase();
let image_keywords = [
"gpt-image",
"dall-e",
"dalle",
"midjourney",
"recraft",
"kolors",
"cogview",
"seedream",
"sdxl",
"stable-diffusion",
"stablediffusion",
"pixart",
"playground",
"firefly",
"imagen",
"flux",
"image",
"img",
"canvas",
"painting",
"draw",
];
if image_keywords
.iter()
.any(|keyword| normalized.contains(keyword))
{
return true;
}
let chat_keywords = [
"claude", "chatgpt", "gpt-4", "gpt-3", "gpt-5", "o1", "o3", "o4", "gemini", "deepseek",
"qwen", "llama", "mistral", "command", "kimi", "doubao", "glm", "hunyuan", "spark",
"baichuan", "ernie", "wenxin", "grok",
];
!chat_keywords
.iter()
.any(|keyword| normalized.contains(keyword))
}
async fn status_handler(
State(state): State<AppState>,
headers: HeaderMap,
@@ -353,6 +406,47 @@ async fn models_handler(
})))
}
async fn image_models_handler(
State(state): State<AppState>,
headers: HeaderMap,
) -> Result<Json<Value>, (axum::http::StatusCode, Json<Value>)> {
let _auth = resolve_agent_auth(&headers, &state).await?;
let rows: Vec<String> = sqlx::query_scalar(
"SELECT model_pattern
FROM token_routes
WHERE COALESCE(enabled, 1) = 1
ORDER BY model_pattern COLLATE NOCASE",
)
.fetch_all(&*state.db)
.await
.map_err(|err| {
(
axum::http::StatusCode::INTERNAL_SERVER_ERROR,
Json(json!({ "error": err.to_string() })),
)
})?;
let mut seen = std::collections::HashSet::new();
let models: Vec<String> = rows
.into_iter()
.map(|value| value.trim().to_string())
.filter(|value| is_image_model_pattern(value))
.filter(|value| seen.insert(value.clone()))
.collect();
Ok(Json(json!({
"success": true,
"models": models,
"defaultModel": models
.iter()
.find(|model| model.to_ascii_lowercase().contains("image"))
.or_else(|| models.first())
.cloned()
.unwrap_or_else(|| "gpt-image-2".to_string()),
})))
}
#[derive(Debug, Default)]
struct ParsedSections {
characters: Vec<CharacterSection>,
@@ -501,7 +595,12 @@ fn parse_sections(full_text: &str) -> ParsedSections {
sections
}
async fn call_image_gen(api_base: &str, api_key: &str, prompt: &str) -> Result<String, String> {
async fn call_image_gen(
api_base: &str,
api_key: &str,
image_model: &str,
prompt: &str,
) -> Result<String, String> {
let client = Client::new();
let url = format!("{}/images/generations", api_base.trim_end_matches('/'));
@@ -509,7 +608,7 @@ async fn call_image_gen(api_base: &str, api_key: &str, prompt: &str) -> Result<S
.post(&url)
.bearer_auth(api_key)
.json(&json!({
"model": "gpt-image-2",
"model": image_model,
"prompt": prompt,
"n": 1,
"size": "1024x1024"
@@ -587,6 +686,19 @@ fn attachment_summary(attachment: &AgentAttachment) -> String {
.unwrap_or("application/octet-stream");
let name = attachment.name.trim();
let name = if name.is_empty() { "untitled" } else { name };
let reference = attachment
.reference_label
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("@素材");
let media_kind = attachment
.media_kind
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("file");
let header = format!("\n\n[素材引用: {reference}, 名称: {name}, 类型: {mime}, 分类: {media_kind}]");
if let Some(text) = attachment
.text
@@ -594,9 +706,9 @@ fn attachment_summary(attachment: &AgentAttachment) -> String {
.map(str::trim)
.filter(|value| !value.is_empty())
{
format!("\n\n[附件: {name}, 类型: {mime}]\n{text}")
format!("{header}\n{text}")
} else {
format!("\n\n[附件: {name}, 类型: {mime}]")
header
}
}
@@ -684,6 +796,12 @@ async fn chat_handler(
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
let requested_image_model = body
.image_model
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
let messages: Vec<Value> = body
.messages
@@ -699,6 +817,7 @@ async fn chat_handler(
"comfyui-agent: resolved upstream",
);
let agent_model = requested_model.unwrap_or_else(|| state.config.comfyui.agent_model.clone());
let image_model = requested_image_model.unwrap_or_else(|| "gpt-image-2".to_string());
let (tx, rx) = unbounded_channel::<Result<Event, Infallible>>();
tokio::spawn(async move {
@@ -813,12 +932,35 @@ async fn chat_handler(
continue;
};
let content = parsed
let delta = parsed
.get("choices")
.and_then(Value::as_array)
.and_then(|choices| choices.first())
.and_then(|choice| choice.get("delta"))
.and_then(|delta| delta.get("content"))
.cloned()
.unwrap_or(Value::Null);
let reasoning = delta
.get("reasoning_content")
.or_else(|| delta.get("reasoning"))
.and_then(Value::as_str)
.unwrap_or_default()
.to_string();
if !reasoning.is_empty() {
if !send_sse(
&tx,
json!({
"type": "reasoning",
"content": reasoning,
}),
) {
return;
}
}
let content = delta
.get("content")
.and_then(Value::as_str)
.unwrap_or_default()
.to_string();
@@ -893,7 +1035,7 @@ async fn chat_handler(
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
}
match call_image_gen(&api_base, &api_key, prompt).await {
match call_image_gen(&api_base, &api_key, &image_model, prompt).await {
Ok(url) if !url.is_empty() => {
send_sse(
&tx,
@@ -927,7 +1069,7 @@ async fn chat_handler(
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
match call_image_gen(&api_base, &api_key, &scene.prompt).await {
match call_image_gen(&api_base, &api_key, &image_model, &scene.prompt).await {
Ok(url) if !url.is_empty() => {
send_sse(
&tx,
@@ -1191,8 +1333,10 @@ mod tests {
messages: vec![AgentMessage {
role: "user".to_string(),
content: "hello".to_string(),
attachments: Vec::new(),
}],
model: None,
image_model: None,
}),
)
.await;