boos
1482 lines
48 KiB
Rust
1482 lines
48 KiB
Rust
use crate::{auth::jwt::validate_jwt, db::models::downstream_api_key::DownstreamApiKey, AppState};
|
|
use axum::{
|
|
extract::{DefaultBodyLimit, State},
|
|
http::HeaderMap,
|
|
response::sse::{Event, Sse},
|
|
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};
|
|
use tokio_stream::{wrappers::UnboundedReceiverStream, StreamExt};
|
|
use tracing::info;
|
|
|
|
const SYSTEM_PROMPT: &str = r#"你是 BoosAPI 的 AI 视频生成助手，帮用户通过 AI 制作视频。
|
|
|
|
## 规则
|
|
1. 先用中文与用户讨论脚本构思、视觉风格、角色设定
|
|
2. 用户确认后（说"开始制作"或"生成"），再分步输出结构化内容
|
|
3. **分两步走：第一步输出角色+场景+TTS，第二步输出 WORKFLOW**（不要一次全出）
|
|
|
|
## 结构化标记格式
|
|
严格按照以下格式输出，系统会自动解析并生成图片/TTS：
|
|
|
|
【CHARACTER】
|
|
name: 林澈
|
|
description: 女宇航员，30岁，冷静坚毅，短发收在头盔内
|
|
prompt_front: cinematic full body front view of a female astronaut, white EVA suit, photorealistic
|
|
prompt_back: cinematic full body back view of a female astronaut, life support backpack
|
|
prompt_left: cinematic full body left side view of a female astronaut, side utility pouches
|
|
prompt_right: cinematic full body right side view of a female astronaut, helmet visor reflection
|
|
【/CHARACTER】
|
|
|
|
【SCENE】
|
|
id: 1
|
|
description: 飞船驾驶舱内景
|
|
prompt: interior of a sci-fi spacecraft cockpit, holographic displays, cold blue lighting, cinematic
|
|
【/SCENE】
|
|
|
|
【TTS】
|
|
character: 林澈
|
|
text: 我们准备出舱检查
|
|
emotion: 平静
|
|
【/TTS】
|
|
|
|
【WORKFLOW】
|
|
```json
|
|
{"comfyui_workflow": {}}
|
|
```
|
|
【/WORKFLOW】
|
|
|
|
## 关键要求
|
|
- prompt_* 用英文，同时包含主体 + 视角 + 细节 + 画质标签
|
|
- 不要一次输出太多，先输角色和场景，下次再输工作流
|
|
- 如果没有用户明确要求，切勿自作主张输出全部内容
|
|
"#;
|
|
|
|
const AGENT_CHAT_BODY_LIMIT_BYTES: usize = 64 * 1024 * 1024;
|
|
|
|
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))
|
|
.layer(DefaultBodyLimit::max(AGENT_CHAT_BODY_LIMIT_BYTES))
|
|
}
|
|
|
|
#[derive(Debug, Deserialize)]
|
|
struct AgentChatRequest {
|
|
messages: Vec<AgentMessage>,
|
|
model: Option<String>,
|
|
#[serde(rename = "imageModel")]
|
|
image_model: Option<String>,
|
|
}
|
|
|
|
#[derive(Debug, Deserialize)]
|
|
struct AgentMessage {
|
|
role: String,
|
|
content: String,
|
|
#[serde(default)]
|
|
attachments: Vec<AgentAttachment>,
|
|
}
|
|
|
|
#[derive(Debug, Deserialize)]
|
|
#[serde(rename_all = "camelCase")]
|
|
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>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
enum AgentMode {
|
|
/// Route through this server's own `/v1/*` using the caller's downstream key.
|
|
Local,
|
|
/// Legacy fallback: external upstream configured via `BOOSAPI_UPSTREAM_BASE` / `BOOSAPI_UPSTREAM_KEY`.
|
|
FallbackRemote,
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
struct AgentAuthContext {
|
|
user_id: Option<i64>,
|
|
role: String,
|
|
outbound_key: String,
|
|
key_name: String,
|
|
}
|
|
|
|
impl AgentMode {
|
|
fn as_str(self) -> &'static str {
|
|
match self {
|
|
AgentMode::Local => "local",
|
|
AgentMode::FallbackRemote => "fallback_remote",
|
|
}
|
|
}
|
|
}
|
|
|
|
fn detect_agent_mode() -> AgentMode {
|
|
let base_override = std::env::var("BOOSAPI_UPSTREAM_BASE")
|
|
.ok()
|
|
.map(|v| v.trim().to_string())
|
|
.filter(|v| !v.is_empty());
|
|
let key_override = std::env::var("BOOSAPI_UPSTREAM_KEY")
|
|
.ok()
|
|
.map(|v| v.trim().to_string())
|
|
.filter(|v| !v.is_empty());
|
|
if base_override.is_some() || key_override.is_some() {
|
|
AgentMode::FallbackRemote
|
|
} else {
|
|
AgentMode::Local
|
|
}
|
|
}
|
|
|
|
fn local_base(state: &AppState) -> String {
|
|
let host = if state.config.server.host == "0.0.0.0" || state.config.server.host == "::" {
|
|
"127.0.0.1".to_string()
|
|
} else {
|
|
state.config.server.host.clone()
|
|
};
|
|
format!("http://{}:{}/v1", host, state.config.server.port)
|
|
}
|
|
|
|
/// Resolve `(mode, api_base, api_key, key_name)` for the agent's outbound calls.
|
|
///
|
|
/// Defaults to `AgentMode::Local`, which targets this server's own `/v1`
|
|
/// endpoints with the current caller's downstream key — every call then runs
|
|
/// through the regular accounting / routing pipeline and is billed/logged
|
|
/// against that user key.
|
|
///
|
|
/// Setting `BOOSAPI_UPSTREAM_BASE` or `BOOSAPI_UPSTREAM_KEY` switches to
|
|
/// `FallbackRemote` for operators that still want to point straight at an
|
|
/// external aggregator.
|
|
fn upstream_config(
|
|
state: &AppState,
|
|
auth: &AgentAuthContext,
|
|
) -> (AgentMode, String, String, String) {
|
|
match detect_agent_mode() {
|
|
AgentMode::Local => {
|
|
let base = local_base(state);
|
|
(
|
|
AgentMode::Local,
|
|
base,
|
|
auth.outbound_key.clone(),
|
|
auth.key_name.clone(),
|
|
)
|
|
}
|
|
AgentMode::FallbackRemote => {
|
|
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(|| auth.outbound_key.clone());
|
|
let key_name = if std::env::var("BOOSAPI_UPSTREAM_KEY")
|
|
.ok()
|
|
.map(|value| !value.trim().is_empty())
|
|
.unwrap_or(false)
|
|
{
|
|
"external".to_string()
|
|
} else {
|
|
auth.key_name.clone()
|
|
};
|
|
(AgentMode::FallbackRemote, base, key, key_name)
|
|
}
|
|
}
|
|
}
|
|
|
|
fn extract_bearer_token(headers: &HeaderMap) -> Option<&str> {
|
|
let auth_header = headers
|
|
.get("authorization")
|
|
.and_then(|value| value.to_str().ok())
|
|
.unwrap_or_default();
|
|
|
|
auth_header
|
|
.strip_prefix("Bearer ")
|
|
.or_else(|| auth_header.strip_prefix("bearer "))
|
|
.map(str::trim)
|
|
.filter(|token| !token.is_empty())
|
|
}
|
|
|
|
async fn user_downstream_key(state: &AppState, user_id: i64) -> Result<DownstreamApiKey, String> {
|
|
sqlx::query_as::<_, DownstreamApiKey>(
|
|
"SELECT *
|
|
FROM downstream_api_keys
|
|
WHERE user_id = ?
|
|
AND COALESCE(enabled, 1) = 1
|
|
AND (
|
|
expires_at IS NULL
|
|
OR expires_at = ''
|
|
OR datetime(expires_at) > datetime('now')
|
|
)
|
|
ORDER BY updated_at DESC, created_at DESC, id DESC
|
|
LIMIT 1",
|
|
)
|
|
.bind(user_id)
|
|
.fetch_optional(&*state.db)
|
|
.await
|
|
.map_err(|err| err.to_string())?
|
|
.ok_or_else(|| "当前用户没有可用的 API Key，请先在用户中心创建或启用 API Key".to_string())
|
|
}
|
|
|
|
async fn resolve_agent_auth(
|
|
headers: &HeaderMap,
|
|
state: &AppState,
|
|
) -> Result<AgentAuthContext, (axum::http::StatusCode, Json<Value>)> {
|
|
let token = extract_bearer_token(headers).ok_or_else(|| {
|
|
(
|
|
axum::http::StatusCode::UNAUTHORIZED,
|
|
Json(json!({ "error": "Unauthorized" })),
|
|
)
|
|
})?;
|
|
|
|
let proxy_token = state.runtime_proxy_token();
|
|
if token == state.config.auth.auth_token || token == proxy_token {
|
|
return Ok(AgentAuthContext {
|
|
user_id: None,
|
|
role: "admin".to_string(),
|
|
outbound_key: proxy_token,
|
|
key_name: "global-proxy".to_string(),
|
|
});
|
|
}
|
|
|
|
let claims = validate_jwt(&state.config.auth.jwt_secret, token).map_err(|_| {
|
|
(
|
|
axum::http::StatusCode::UNAUTHORIZED,
|
|
Json(json!({ "error": "Unauthorized" })),
|
|
)
|
|
})?;
|
|
if claims.role != "admin" && claims.role != "user" {
|
|
return Err((
|
|
axum::http::StatusCode::FORBIDDEN,
|
|
Json(json!({ "error": "Forbidden" })),
|
|
));
|
|
}
|
|
|
|
let user_id = claims.sub.parse::<i64>().map_err(|_| {
|
|
(
|
|
axum::http::StatusCode::UNAUTHORIZED,
|
|
Json(json!({ "error": "Invalid user ID in token" })),
|
|
)
|
|
})?;
|
|
|
|
match user_downstream_key(state, user_id).await {
|
|
Ok(key) => Ok(AgentAuthContext {
|
|
user_id: Some(user_id),
|
|
role: claims.role,
|
|
outbound_key: key.key,
|
|
key_name: key.name,
|
|
}),
|
|
Err(_) if claims.role == "admin" => Ok(AgentAuthContext {
|
|
user_id: Some(user_id),
|
|
role: claims.role,
|
|
outbound_key: state.runtime_proxy_token(),
|
|
key_name: "global-proxy".to_string(),
|
|
}),
|
|
Err(message) => Err((
|
|
axum::http::StatusCode::FORBIDDEN,
|
|
Json(json!({ "error": message })),
|
|
)),
|
|
}
|
|
}
|
|
|
|
fn is_exact_model_pattern(value: &str) -> bool {
|
|
let normalized = value.trim();
|
|
!normalized.is_empty()
|
|
&& !normalized.to_ascii_lowercase().starts_with("re:")
|
|
&& !normalized.contains('*')
|
|
&& !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,
|
|
) -> Result<Json<Value>, (axum::http::StatusCode, Json<Value>)> {
|
|
let auth = resolve_agent_auth(&headers, &state).await?;
|
|
let (mode, upstream_base, _, key_name) = upstream_config(&state, &auth);
|
|
let upstream_label = match mode {
|
|
AgentMode::Local => "BoosAPI 本地路由".to_string(),
|
|
AgentMode::FallbackRemote => {
|
|
let parsed = reqwest::Url::parse(&upstream_base).ok();
|
|
parsed
|
|
.and_then(|url| url.host_str().map(ToOwned::to_owned))
|
|
.unwrap_or_else(|| "外部上游".to_string())
|
|
}
|
|
};
|
|
Ok(Json(json!({
|
|
"success": true,
|
|
"model": state.config.comfyui.agent_model,
|
|
"mode": mode.as_str(),
|
|
"keyName": key_name,
|
|
"userId": auth.user_id,
|
|
"role": auth.role,
|
|
"apiBase": upstream_base,
|
|
"upstreamBase": upstream_base,
|
|
"upstreamLabel": upstream_label,
|
|
"comfyui": state.comfyui.status_response().await,
|
|
})))
|
|
}
|
|
|
|
async fn 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_exact_model_pattern(value))
|
|
.filter(|value| seen.insert(value.clone()))
|
|
.collect();
|
|
|
|
Ok(Json(json!({
|
|
"success": true,
|
|
"models": models,
|
|
"defaultModel": state.config.comfyui.agent_model,
|
|
})))
|
|
}
|
|
|
|
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>,
|
|
scenes: Vec<SceneSection>,
|
|
tts_list: Vec<TtsSection>,
|
|
workflow_json: Option<String>,
|
|
}
|
|
|
|
#[derive(Debug, Default)]
|
|
struct CharacterSection {
|
|
name: String,
|
|
description: String,
|
|
prompt_front: String,
|
|
prompt_back: String,
|
|
prompt_left: String,
|
|
prompt_right: String,
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
struct SceneSection {
|
|
id: u32,
|
|
description: String,
|
|
prompt: String,
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
struct TtsSection {
|
|
character: String,
|
|
text: String,
|
|
emotion: String,
|
|
}
|
|
|
|
fn parse_sections(full_text: &str) -> ParsedSections {
|
|
let mut sections = ParsedSections::default();
|
|
|
|
// Parse CHARACTER blocks
|
|
let char_re = Regex::new(r"(?s)【CHARACTER】\s*(.*?)【/CHARACTER】").unwrap();
|
|
for cap in char_re.captures_iter(full_text) {
|
|
let block = cap.get(1).unwrap().as_str();
|
|
let mut char_section = CharacterSection::default();
|
|
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 :"))
|
|
{
|
|
char_section.name = val.trim().to_string();
|
|
} 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.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.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.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.split_once(':').map(|x| x.1).unwrap_or("").trim();
|
|
char_section.prompt_right = val.to_string();
|
|
}
|
|
}
|
|
if !char_section.name.is_empty() {
|
|
sections.characters.push(char_section);
|
|
}
|
|
}
|
|
|
|
// Parse SCENE blocks
|
|
let scene_re = Regex::new(r"(?s)【SCENE】\s*(.*?)【/SCENE】").unwrap();
|
|
for cap in scene_re.captures_iter(full_text) {
|
|
let block = cap.get(1).unwrap().as_str();
|
|
let mut scene = SceneSection {
|
|
id: 0,
|
|
description: String::new(),
|
|
prompt: String::new(),
|
|
};
|
|
for line in block.lines() {
|
|
let line = line.trim();
|
|
let line_lower = line.to_lowercase();
|
|
if line_lower.starts_with("id:") || line_lower.starts_with("id :") {
|
|
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 :"))
|
|
{
|
|
scene.description = val.trim().to_string();
|
|
} else if line_lower.starts_with("prompt:") || line_lower.starts_with("prompt :") {
|
|
let val = line.split_once(':').map(|x| x.1).unwrap_or("").trim();
|
|
scene.prompt = val.to_string();
|
|
}
|
|
}
|
|
if !scene.prompt.is_empty() {
|
|
sections.scenes.push(scene);
|
|
}
|
|
}
|
|
|
|
// Parse TTS blocks
|
|
let tts_re = Regex::new(r"(?s)【TTS】\s*(.*?)【/TTS】").unwrap();
|
|
for cap in tts_re.captures_iter(full_text) {
|
|
let block = cap.get(1).unwrap().as_str();
|
|
let mut tts = TtsSection {
|
|
character: String::new(),
|
|
text: String::new(),
|
|
emotion: String::new(),
|
|
};
|
|
for line in block.lines() {
|
|
let line = line.trim();
|
|
let line_lower = line.to_lowercase();
|
|
if line_lower.starts_with("character:") || line_lower.starts_with("character :") {
|
|
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.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.split_once(':').map(|x| x.1).unwrap_or("").trim();
|
|
tts.emotion = val.to_string();
|
|
}
|
|
}
|
|
if !tts.text.is_empty() {
|
|
sections.tts_list.push(tts);
|
|
}
|
|
}
|
|
|
|
// Parse WORKFLOW blocks
|
|
let wf_re = Regex::new(r"(?s)【WORKFLOW】\s*```(?:json)?\s*(.*?)```\s*【/WORKFLOW】").unwrap();
|
|
if let Some(cap) = wf_re.captures(full_text) {
|
|
let wf = cap.get(1).unwrap().as_str().trim().to_string();
|
|
sections.workflow_json = Some(wf);
|
|
}
|
|
|
|
sections
|
|
}
|
|
|
|
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('/'));
|
|
|
|
let resp = client
|
|
.post(&url)
|
|
.bearer_auth(api_key)
|
|
.json(&json!({
|
|
"model": image_model,
|
|
"prompt": prompt,
|
|
"n": 1,
|
|
"size": "1024x1024"
|
|
}))
|
|
.send()
|
|
.await
|
|
.map_err(|e| format!("image request failed: {}", e))?;
|
|
|
|
let status = resp.status();
|
|
let data: Value = resp
|
|
.json()
|
|
.await
|
|
.map_err(|e| format!("image response parse failed: {}", e))?;
|
|
|
|
if !status.is_success() {
|
|
return Err(format!("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(
|
|
api_base: &str,
|
|
api_key: &str,
|
|
text: &str,
|
|
voice: &str,
|
|
) -> Result<Vec<u8>, String> {
|
|
let client = Client::new();
|
|
let url = format!("{}/audio/speech", api_base.trim_end_matches('/'));
|
|
|
|
let resp = client
|
|
.post(&url)
|
|
.bearer_auth(api_key)
|
|
.json(&json!({
|
|
"model": "tts-1",
|
|
"input": text,
|
|
"voice": voice
|
|
}))
|
|
.send()
|
|
.await
|
|
.map_err(|e| format!("TTS request failed: {}", e))?;
|
|
|
|
let status = resp.status();
|
|
if !status.is_success() {
|
|
let body = resp.text().await.unwrap_or_default();
|
|
return Err(format!("TTS upstream returned {}: {}", status, body));
|
|
}
|
|
|
|
let bytes = resp
|
|
.bytes()
|
|
.await
|
|
.map_err(|e| format!("TTS response read failed: {}", e))?;
|
|
|
|
Ok(bytes.to_vec())
|
|
}
|
|
|
|
fn send_sse(tx: &UnboundedSender<Result<Event, Infallible>>, payload: Value) -> bool {
|
|
tx.send(Ok(Event::default().data(payload.to_string())))
|
|
.is_ok()
|
|
}
|
|
|
|
fn string_field(value: &Value, keys: &[&str]) -> Option<String> {
|
|
keys.iter()
|
|
.find_map(|key| value.get(*key).and_then(Value::as_str))
|
|
.map(ToOwned::to_owned)
|
|
.filter(|text| !text.is_empty())
|
|
}
|
|
|
|
fn extract_reasoning_delta(parsed: &Value) -> Option<String> {
|
|
if let Some(delta) = parsed
|
|
.get("choices")
|
|
.and_then(Value::as_array)
|
|
.and_then(|choices| choices.first())
|
|
.and_then(|choice| choice.get("delta"))
|
|
{
|
|
if let Some(text) = string_field(
|
|
delta,
|
|
&[
|
|
"reasoning_content",
|
|
"reasoning",
|
|
"thinking",
|
|
"thinking_delta",
|
|
"reasoning_delta",
|
|
],
|
|
) {
|
|
return Some(text);
|
|
}
|
|
|
|
if let Some(text) = delta
|
|
.get("reasoning")
|
|
.and_then(|value| value.get("content").or_else(|| value.get("text")))
|
|
.and_then(Value::as_str)
|
|
.map(ToOwned::to_owned)
|
|
.filter(|text| !text.is_empty())
|
|
{
|
|
return Some(text);
|
|
}
|
|
}
|
|
|
|
let event_type = parsed
|
|
.get("type")
|
|
.and_then(Value::as_str)
|
|
.unwrap_or_default();
|
|
match event_type {
|
|
"response.reasoning.delta"
|
|
| "response.reasoning_text.delta"
|
|
| "response.reasoning_summary_text.delta" => {
|
|
string_field(parsed, &["delta", "text", "content"])
|
|
}
|
|
"content_block_delta" => {
|
|
let delta = parsed.get("delta")?;
|
|
match delta
|
|
.get("type")
|
|
.and_then(Value::as_str)
|
|
.unwrap_or_default()
|
|
{
|
|
"thinking_delta" => string_field(delta, &["thinking", "text", "delta"]),
|
|
"signature_delta" => None,
|
|
other if other.contains("thinking") || other.contains("reasoning") => {
|
|
string_field(delta, &["thinking", "text", "content", "delta"])
|
|
}
|
|
_ => None,
|
|
}
|
|
}
|
|
other if other.contains("reasoning") || other.contains("thinking") => string_field(
|
|
parsed,
|
|
&["delta", "text", "content", "thinking", "reasoning"],
|
|
),
|
|
_ => None,
|
|
}
|
|
}
|
|
|
|
fn attachment_summary(attachment: &AgentAttachment) -> String {
|
|
let mime = attachment
|
|
.mime_type
|
|
.as_deref()
|
|
.map(str::trim)
|
|
.filter(|value| !value.is_empty())
|
|
.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
|
|
.as_deref()
|
|
.map(str::trim)
|
|
.filter(|value| !value.is_empty())
|
|
{
|
|
format!("{header}\n{text}")
|
|
} else {
|
|
header
|
|
}
|
|
}
|
|
|
|
fn agent_message_to_upstream(message: AgentMessage) -> Value {
|
|
let role = message.role;
|
|
let mut text = message.content;
|
|
let mut content_parts = Vec::new();
|
|
|
|
for attachment in &message.attachments {
|
|
let mime = attachment
|
|
.mime_type
|
|
.as_deref()
|
|
.unwrap_or_default()
|
|
.to_ascii_lowercase();
|
|
if mime.starts_with("image/") {
|
|
if let Some(data_url) = attachment
|
|
.data_url
|
|
.as_deref()
|
|
.map(str::trim)
|
|
.filter(|value| value.starts_with("data:image/"))
|
|
{
|
|
if !content_parts
|
|
.iter()
|
|
.any(|part: &Value| part["type"] == "text")
|
|
{
|
|
content_parts.push(json!({
|
|
"type": "text",
|
|
"text": text,
|
|
}));
|
|
text = String::new();
|
|
}
|
|
content_parts.push(json!({
|
|
"type": "image_url",
|
|
"image_url": { "url": data_url },
|
|
}));
|
|
continue;
|
|
}
|
|
}
|
|
|
|
text.push_str(&attachment_summary(attachment));
|
|
}
|
|
|
|
if !content_parts.is_empty() {
|
|
if !text.trim().is_empty() {
|
|
content_parts.push(json!({
|
|
"type": "text",
|
|
"text": text,
|
|
}));
|
|
}
|
|
json!({
|
|
"role": role,
|
|
"content": content_parts,
|
|
})
|
|
} else {
|
|
json!({
|
|
"role": role,
|
|
"content": text,
|
|
})
|
|
}
|
|
}
|
|
|
|
async fn chat_handler(
|
|
State(state): State<AppState>,
|
|
headers: HeaderMap,
|
|
Json(body): Json<AgentChatRequest>,
|
|
) -> Result<Sse<impl Stream<Item = Result<Event, Infallible>>>, (axum::http::StatusCode, Json<Value>)>
|
|
{
|
|
let auth = resolve_agent_auth(&headers, &state).await?;
|
|
|
|
if body.messages.is_empty()
|
|
|| body.messages.iter().any(|message| {
|
|
message.role.trim().is_empty()
|
|
|| (message.content.trim().is_empty() && message.attachments.is_empty())
|
|
})
|
|
{
|
|
return Err((
|
|
axum::http::StatusCode::BAD_REQUEST,
|
|
Json(json!({ "error": "messages is required" })),
|
|
));
|
|
}
|
|
|
|
let requested_model = body
|
|
.model
|
|
.as_deref()
|
|
.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
|
|
.into_iter()
|
|
.map(agent_message_to_upstream)
|
|
.collect();
|
|
let (mode, api_base, api_key, key_name) = upstream_config(&state, &auth);
|
|
info!(
|
|
mode = mode.as_str(),
|
|
upstream_base = %api_base,
|
|
key_name = %key_name,
|
|
user_id = ?auth.user_id,
|
|
"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 {
|
|
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!({
|
|
"role": "system",
|
|
"content": SYSTEM_PROMPT,
|
|
}));
|
|
upstream_messages.extend(messages);
|
|
|
|
let client = Client::new();
|
|
|
|
info!("comfyui-agent: sending chat request to upstream");
|
|
|
|
let response = match client
|
|
.post(&upstream_url)
|
|
.bearer_auth(&api_key)
|
|
.json(&json!({
|
|
"model": agent_model,
|
|
"messages": upstream_messages,
|
|
"temperature": 0.7,
|
|
"stream": true,
|
|
}))
|
|
.send()
|
|
.await
|
|
{
|
|
Ok(response) => response,
|
|
Err(e) => {
|
|
send_sse(
|
|
&tx,
|
|
json!({
|
|
"type": "text",
|
|
"content": format!("服务暂时不可用，请稍后重试。({})", e),
|
|
}),
|
|
);
|
|
send_sse(&tx, json!({ "type": "done" }));
|
|
return;
|
|
}
|
|
};
|
|
|
|
if !response.status().is_success() {
|
|
send_sse(
|
|
&tx,
|
|
json!({
|
|
"type": "text",
|
|
"content": "服务暂时不可用，请稍后重试",
|
|
}),
|
|
);
|
|
send_sse(&tx, json!({ "type": "done" }));
|
|
return;
|
|
}
|
|
|
|
let mut stream = response.bytes_stream();
|
|
let mut buffer: Vec<u8> = Vec::new();
|
|
let mut accumulated_text = String::new();
|
|
|
|
// Phase 1: Stream the LLM text response
|
|
loop {
|
|
let chunk = tokio::select! {
|
|
_ = tx.closed() => {
|
|
return;
|
|
}
|
|
chunk = stream.next() => chunk,
|
|
};
|
|
|
|
let Some(chunk) = chunk else {
|
|
break;
|
|
};
|
|
|
|
let chunk = match chunk {
|
|
Ok(chunk) => chunk,
|
|
Err(error) => {
|
|
send_sse(
|
|
&tx,
|
|
json!({
|
|
"type": "text",
|
|
"content": format!("\n\n[错误: {}]", error),
|
|
}),
|
|
);
|
|
break;
|
|
}
|
|
};
|
|
|
|
buffer.extend_from_slice(&chunk);
|
|
|
|
while let Some(pos) = buffer.iter().position(|byte| *byte == b'\n') {
|
|
let mut line_bytes: Vec<u8> = buffer.drain(..=pos).collect();
|
|
|
|
if line_bytes.last() == Some(&b'\n') {
|
|
line_bytes.pop();
|
|
}
|
|
if line_bytes.last() == Some(&b'\r') {
|
|
line_bytes.pop();
|
|
}
|
|
|
|
let line = String::from_utf8_lossy(&line_bytes);
|
|
let line = line.trim().to_string();
|
|
|
|
if !line.starts_with("data: ") {
|
|
continue;
|
|
}
|
|
|
|
let data = line.trim_start_matches("data: ").trim().to_string();
|
|
|
|
if data == "[DONE]" {
|
|
continue;
|
|
}
|
|
|
|
let Ok(parsed) = serde_json::from_str::<Value>(&data) else {
|
|
continue;
|
|
};
|
|
|
|
if let Some(reasoning) = extract_reasoning_delta(&parsed) {
|
|
if !send_sse(
|
|
&tx,
|
|
json!({
|
|
"type": "reasoning",
|
|
"content": reasoning,
|
|
}),
|
|
) {
|
|
return;
|
|
}
|
|
}
|
|
|
|
let delta = parsed
|
|
.get("choices")
|
|
.and_then(Value::as_array)
|
|
.and_then(|choices| choices.first())
|
|
.and_then(|choice| choice.get("delta"))
|
|
.cloned()
|
|
.unwrap_or(Value::Null);
|
|
|
|
let content = delta
|
|
.get("content")
|
|
.and_then(Value::as_str)
|
|
.unwrap_or_default()
|
|
.to_string();
|
|
|
|
if !content.is_empty() {
|
|
accumulated_text.push_str(&content);
|
|
if !send_sse(
|
|
&tx,
|
|
json!({
|
|
"type": "text",
|
|
"content": content,
|
|
}),
|
|
) {
|
|
return;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Phase 2: Parse structured sections from accumulated text
|
|
info!("comfyui-agent: text streaming complete, parsing sections");
|
|
let sections = parse_sections(&accumulated_text);
|
|
|
|
let has_structured_data = !sections.characters.is_empty()
|
|
|| !sections.scenes.is_empty()
|
|
|| !sections.tts_list.is_empty()
|
|
|| sections.workflow_json.is_some();
|
|
|
|
if !has_structured_data {
|
|
send_sse(&tx, json!({ "type": "done" }));
|
|
return;
|
|
}
|
|
|
|
// Send a processing indicator
|
|
send_sse(
|
|
&tx,
|
|
json!({
|
|
"type": "text",
|
|
"content": "\n\n🔄 正在处理结构化数据..."
|
|
}),
|
|
);
|
|
|
|
// Generate images for each character (4 angles each)
|
|
for (char_idx, character) in sections.characters.iter().enumerate() {
|
|
// Send character_prompts event
|
|
send_sse(
|
|
&tx,
|
|
json!({
|
|
"type": "character_prompts",
|
|
"character": character.name,
|
|
}),
|
|
);
|
|
|
|
let angles = [
|
|
("正面", &character.prompt_front, "front"),
|
|
("背面", &character.prompt_back, "back"),
|
|
("左侧", &character.prompt_left, "left"),
|
|
("右侧", &character.prompt_right, "right"),
|
|
];
|
|
|
|
for (angle_label, prompt, angle_key) in &angles {
|
|
if prompt.is_empty() {
|
|
continue;
|
|
}
|
|
info!(
|
|
"comfyui-agent: generating image for {} - {}",
|
|
character.name, angle_label
|
|
);
|
|
|
|
// Add a short delay between image gen requests to avoid rate limiting
|
|
if char_idx > 0 || *angle_key != "front" {
|
|
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
|
|
}
|
|
|
|
match call_image_gen(&api_base, &api_key, &image_model, prompt).await {
|
|
Ok(url) if !url.is_empty() => {
|
|
send_sse(
|
|
&tx,
|
|
json!({
|
|
"type": "image",
|
|
"url": url,
|
|
"character": character.name,
|
|
"angle": angle_key,
|
|
}),
|
|
);
|
|
}
|
|
_ => {
|
|
send_sse(
|
|
&tx,
|
|
json!({
|
|
"type": "text",
|
|
"content": format!("\n⚠️ {} [{}] 图片生成失败，请重试", character.name, angle_label),
|
|
}),
|
|
);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Generate images for each scene
|
|
for scene in §ions.scenes {
|
|
info!(
|
|
"comfyui-agent: generating scene image for scene {}",
|
|
scene.id
|
|
);
|
|
|
|
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
|
|
|
|
match call_image_gen(&api_base, &api_key, &image_model, &scene.prompt).await {
|
|
Ok(url) if !url.is_empty() => {
|
|
send_sse(
|
|
&tx,
|
|
json!({
|
|
"type": "scene",
|
|
"id": scene.id,
|
|
"description": scene.description,
|
|
"prompt": scene.prompt,
|
|
"url": url,
|
|
}),
|
|
);
|
|
}
|
|
_ => {
|
|
// Still send scene info even without image
|
|
send_sse(
|
|
&tx,
|
|
json!({
|
|
"type": "scene",
|
|
"id": scene.id,
|
|
"description": scene.description,
|
|
"prompt": scene.prompt,
|
|
}),
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Generate TTS for each dialogue
|
|
for tts in §ions.tts_list {
|
|
info!("comfyui-agent: generating TTS for {}", tts.character);
|
|
|
|
// Choose voice based on emotion
|
|
let voice = match tts.emotion.as_str() {
|
|
"激动" | "愤怒" => "nova",
|
|
"悲伤" | "恐惧" => "shimmer",
|
|
_ => "alloy",
|
|
};
|
|
|
|
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",
|
|
tts.character.replace(' ', "_"),
|
|
std::time::SystemTime::now()
|
|
.duration_since(std::time::UNIX_EPOCH)
|
|
.unwrap_or_default()
|
|
.as_millis()
|
|
);
|
|
if let Ok(()) = tokio::fs::write(&audio_path, &audio_bytes).await {
|
|
send_sse(
|
|
&tx,
|
|
json!({
|
|
"type": "tts",
|
|
"character": tts.character,
|
|
"text": tts.text,
|
|
"url": format!("/audio/{}", audio_path.trim_start_matches("/tmp/comfyui_tts_")),
|
|
}),
|
|
);
|
|
}
|
|
}
|
|
Err(e) => {
|
|
info!("comfyui-agent: TTS generation failed: {}", e);
|
|
send_sse(
|
|
&tx,
|
|
json!({
|
|
"type": "text",
|
|
"content": format!("\n⚠️ {} 的语音生成失败", tts.character),
|
|
}),
|
|
);
|
|
}
|
|
}
|
|
|
|
tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
|
|
}
|
|
|
|
// Send ComfyUI workflow if present
|
|
if let Some(wf_json) = §ions.workflow_json {
|
|
info!("comfyui-agent: sending workflow JSON");
|
|
send_sse(
|
|
&tx,
|
|
json!({
|
|
"type": "comfyui_workflow",
|
|
"content": wf_json,
|
|
}),
|
|
);
|
|
}
|
|
|
|
// Signal done with summary
|
|
let mut summary = String::from("\n\n✅ **处理完成！**");
|
|
if !sections.characters.is_empty() {
|
|
summary.push_str(&format!("\n- 角色: {} 个", sections.characters.len()));
|
|
}
|
|
if !sections.scenes.is_empty() {
|
|
summary.push_str(&format!("\n- 场景: {} 个", sections.scenes.len()));
|
|
}
|
|
if !sections.tts_list.is_empty() {
|
|
summary.push_str(&format!("\n- TTS: {} 段", sections.tts_list.len()));
|
|
}
|
|
if sections.workflow_json.is_some() {
|
|
summary.push_str("\n- ComfyUI 工作流: ✅ 已生成");
|
|
}
|
|
send_sse(&tx, json!({ "type": "text", "content": summary }));
|
|
send_sse(&tx, json!({ "type": "done" }));
|
|
});
|
|
|
|
Ok(Sse::new(UnboundedReceiverStream::new(rx)))
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::config::{AppConfig, AuthConfig, DbConfig, ProxyConfig, ServerConfig};
|
|
use crate::services::builtin_keys::BuiltinKeys;
|
|
use axum::http::HeaderValue;
|
|
use sqlx::sqlite::SqlitePoolOptions;
|
|
use tokio::sync::Mutex;
|
|
use wiremock::matchers::{header, method, path};
|
|
use wiremock::{Mock, MockServer, ResponseTemplate};
|
|
|
|
/// `BOOSAPI_UPSTREAM_*` is process-global. Serialize tests that read or
|
|
/// mutate it so cargo's parallel runner does not race their assertions.
|
|
/// Uses `tokio::sync::Mutex` so the guard can safely span `.await` points
|
|
/// (mock server setup, agent worker poll loop) without tripping
|
|
/// `clippy::await_holding_lock`.
|
|
static ENV_GUARD: Mutex<()> = Mutex::const_new(());
|
|
|
|
async fn make_state(host: &str, port: u16) -> AppState {
|
|
let pool = SqlitePoolOptions::new()
|
|
.max_connections(1)
|
|
.connect("sqlite::memory:")
|
|
.await
|
|
.unwrap();
|
|
AppState::new(
|
|
AppConfig {
|
|
server: ServerConfig {
|
|
host: host.to_string(),
|
|
port,
|
|
log_level: "info".into(),
|
|
cors_origins: vec![],
|
|
},
|
|
db: DbConfig {
|
|
url: "sqlite::memory:".into(),
|
|
max_connections: 1,
|
|
ssl: false,
|
|
runtime_db_url: None,
|
|
runtime_db_dialect: None,
|
|
runtime_db_ssl: None,
|
|
},
|
|
auth: AuthConfig {
|
|
jwt_secret: "secret".into(),
|
|
jwt_expiry_hours: 1,
|
|
auth_token: "auth".into(),
|
|
proxy_token: "proxy".into(),
|
|
},
|
|
proxy: ProxyConfig {
|
|
upstream_timeout_secs: 30,
|
|
first_byte_timeout_secs: 30,
|
|
max_retries: 3,
|
|
proxy_log_retention_days: 7,
|
|
template_proxy_url: None,
|
|
},
|
|
comfyui: Default::default(),
|
|
},
|
|
pool,
|
|
BuiltinKeys {
|
|
comfyui_agent: (1, "sk-comfyui-test".into()),
|
|
},
|
|
crate::services::runtime_auth::RuntimeAuthSettings::from_token_for_test("proxy"),
|
|
)
|
|
}
|
|
|
|
fn clear_upstream_env() {
|
|
std::env::remove_var("BOOSAPI_UPSTREAM_BASE");
|
|
std::env::remove_var("BOOSAPI_UPSTREAM_KEY");
|
|
}
|
|
|
|
fn auth_headers() -> HeaderMap {
|
|
let mut headers = HeaderMap::new();
|
|
headers.insert("Authorization", HeaderValue::from_static("Bearer auth"));
|
|
headers
|
|
}
|
|
|
|
async fn call_status(state: AppState) -> Value {
|
|
let result = status_handler(State(state), auth_headers()).await;
|
|
result.expect("status_handler returned error").0
|
|
}
|
|
|
|
#[test]
|
|
fn extracts_chat_reasoning_delta() {
|
|
let parsed = json!({
|
|
"choices": [{
|
|
"delta": {
|
|
"reasoning_content": "正在拆解镜头..."
|
|
}
|
|
}]
|
|
});
|
|
assert_eq!(
|
|
extract_reasoning_delta(&parsed).as_deref(),
|
|
Some("正在拆解镜头...")
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn extracts_responses_reasoning_delta() {
|
|
let parsed = json!({
|
|
"type": "response.reasoning_summary_text.delta",
|
|
"delta": "先分析角色和场景。"
|
|
});
|
|
assert_eq!(
|
|
extract_reasoning_delta(&parsed).as_deref(),
|
|
Some("先分析角色和场景。")
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn extracts_anthropic_thinking_delta() {
|
|
let parsed = json!({
|
|
"type": "content_block_delta",
|
|
"index": 0,
|
|
"delta": {
|
|
"type": "thinking_delta",
|
|
"thinking": "用户需要一个视频分镜。"
|
|
}
|
|
});
|
|
assert_eq!(
|
|
extract_reasoning_delta(&parsed).as_deref(),
|
|
Some("用户需要一个视频分镜。")
|
|
);
|
|
}
|
|
|
|
/// 3.5 — `/api/comfyui-agent/status` reflects the env state in one shot:
|
|
/// no env → local; whitespace/empty env → still local (treated as unset);
|
|
/// real env → fallback_remote with trimmed `upstreamBase`.
|
|
#[tokio::test]
|
|
async fn status_reflects_env_state_across_three_scenarios() {
|
|
let _g = ENV_GUARD.lock().await;
|
|
clear_upstream_env();
|
|
|
|
// Scenario A: env unset → local; 0.0.0.0 host normalizes to 127.0.0.1.
|
|
{
|
|
let state = make_state("0.0.0.0", 7860).await;
|
|
let body = call_status(state).await;
|
|
assert_eq!(body["mode"], "local");
|
|
assert_eq!(body["keyName"], "global-proxy");
|
|
assert_eq!(body["upstreamBase"], "http://127.0.0.1:7860/v1");
|
|
}
|
|
|
|
// Scenario B: env set to whitespace/empty → treated as unset → local.
|
|
std::env::set_var("BOOSAPI_UPSTREAM_BASE", " ");
|
|
std::env::set_var("BOOSAPI_UPSTREAM_KEY", "");
|
|
{
|
|
let state = make_state("127.0.0.1", 9000).await;
|
|
let body = call_status(state).await;
|
|
assert_eq!(body["mode"], "local");
|
|
assert_eq!(body["keyName"], "global-proxy");
|
|
assert_eq!(body["upstreamBase"], "http://127.0.0.1:9000/v1");
|
|
}
|
|
|
|
// Scenario C: env set → fallback_remote, base trimmed of trailing slash.
|
|
std::env::set_var("BOOSAPI_UPSTREAM_BASE", "https://api.example.com/v1/");
|
|
std::env::remove_var("BOOSAPI_UPSTREAM_KEY");
|
|
{
|
|
let state = make_state("0.0.0.0", 7860).await;
|
|
let body = call_status(state).await;
|
|
assert_eq!(body["mode"], "fallback_remote");
|
|
assert_eq!(body["keyName"], "global-proxy");
|
|
assert_eq!(body["upstreamBase"], "https://api.example.com/v1");
|
|
}
|
|
|
|
clear_upstream_env();
|
|
}
|
|
|
|
/// 3.6 — In Local mode the agent's chat call must hit the local
|
|
/// `/v1/chat/completions` endpoint with the current caller's downstream
|
|
/// key as the bearer token.
|
|
#[tokio::test]
|
|
async fn agent_local_mode_routes_chat_to_local_upstream() {
|
|
let _g = ENV_GUARD.lock().await;
|
|
clear_upstream_env();
|
|
|
|
let mock = MockServer::start().await;
|
|
let sse = "data: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\ndata: [DONE]\n\n";
|
|
Mock::given(method("POST"))
|
|
.and(path("/v1/chat/completions"))
|
|
.and(header("authorization", "Bearer proxy"))
|
|
.respond_with(
|
|
ResponseTemplate::new(200)
|
|
.insert_header("content-type", "text/event-stream")
|
|
.set_body_string(sse),
|
|
)
|
|
.mount(&mock)
|
|
.await;
|
|
|
|
let addr = mock.address();
|
|
let host = addr.ip().to_string();
|
|
let port = addr.port();
|
|
|
|
let state = make_state(&host, port).await;
|
|
|
|
let result = chat_handler(
|
|
State(state),
|
|
auth_headers(),
|
|
Json(AgentChatRequest {
|
|
messages: vec![AgentMessage {
|
|
role: "user".to_string(),
|
|
content: "hello".to_string(),
|
|
attachments: Vec::new(),
|
|
}],
|
|
model: None,
|
|
image_model: None,
|
|
}),
|
|
)
|
|
.await;
|
|
assert!(result.is_ok(), "chat_handler should accept the request");
|
|
// The Sse stream is dropped here; the worker still runs the upstream
|
|
// POST before its streaming loop notices `tx.closed()`.
|
|
|
|
// Wait up to ~2s for the spawned worker's POST to land at wiremock.
|
|
let mut hit = false;
|
|
for _ in 0..40 {
|
|
tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
|
|
let received = mock.received_requests().await.unwrap_or_default();
|
|
if received
|
|
.iter()
|
|
.any(|r| r.url.path() == "/v1/chat/completions")
|
|
{
|
|
hit = true;
|
|
break;
|
|
}
|
|
}
|
|
assert!(
|
|
hit,
|
|
"expected /v1/chat/completions to be hit by the local-mode agent"
|
|
);
|
|
}
|
|
}
|