新增: AI 工具(download_file + http output_file + grep 窗口 + fetch_search 搜索 + get_app_config + generate_image)

download_file 跨平台 URL 到文件流式下载;http_request output_file 落盘 + 截断可配;
grep context_chars 大单行窗口;fetch_search DuckDuckGo 免 key 搜索;get_app_config
只读返当前配置(治 AI 查配置绕 PowerShell);generate_image SenseNova 图像生成。
This commit is contained in:
lxy
2026-08-02 02:21:32 +08:00
parent e4f7b432aa
commit b999f79fdf
17 changed files with 2862 additions and 56 deletions
+587
View File
@@ -0,0 +1,587 @@
//! generate_image AI 工具 — 调 provider 的 OpenAI 兼容图像生成端点生成图片并下载落地。
//!
//! 设计目标:让 LLM 调用图像生成模型(默认 SenseNova U1 Fast,也支持 DALL-E 等 OpenAI 兼容
//! 图像端点)生成图片并落到工作区磁盘。图像生成模型不是 chat 模型,走 `/v1/images/generations`
//! 端点(非 chat completions),故不能经主对话链路调用,必须用本工具显式触发。
//!
//! ## 与 http_request / fetch_url / download_file 的分工
//!
//! - `http_request`:结构化 API 调用,返回 body(截断 50KB),不落盘 —— 通用 HTTP,不感知图像端点。
//! - `fetch_url`:网页 → markdown 文档嗅探,不落盘。
//! - `download_file`:URL → 落盘文件(已知 URL)。
//! - `generate_image`:**调 provider 图像端点 + 自动下载落地**(封装图像生成的 provider 凭证
//! 解析 + 端点拼接 + 响应解析 + 图片下载全链路,LLM 只需给 prompt/model/output_path)。
//!
//! ## 风险:High
//!
//! ① 付费 API(图像生成按张计费);② 写磁盘(落盘 output_path);③ 外发网络(provider 端点
//! + 图片 URL 下载)。统一 High 须人工批准。
//!
//! ## 安全
//!
//! 1. **provider 凭证**:从 db.ai_providers 筛 `enabled && provider_type=="openai_compat"` 的
//! provider,经 `df_storage::secret::resolve_provider_secret` 解析 key(DB 优先 fallback keyring)。
//! anthropic provider 不参与(图像端点是 OpenAI 风格)。
//! 2. **端点拼接**:`build_images_url` 智能 base_url(`/v1`/`/v4` 后缀直接补 `/images/generations`,
//! 否则补 `/v1/images/generations`),对齐 `model_fetch_helpers::build_models_url` 思路。
//! 3. **provider 端点域名**:provider 已知域名(api.sensenova / api.openai / ...),非用户输入,
//! SSRF 风险低,故 provider POST 直接 reqwest::Client 不走 SSRF 防护。
//! 4. **图片 URL 下载**:图片 URL 来自 provider 响应,**可能被恶意 provider 篡改指向内网**
//! (provider 域名虽可信但响应内容不可信),故下载图片 URL 复用 SSRF 防护
//! (validate_url + resolve_and_check_host + build_client,与 download_file 同源)。
//! 5. **路径校验**:output_path 走 `resolve_workspace_path_with_allowed`(白名单 + symlink 逃逸
//! 拦截 + 路径遍历拒),与 write_file/download_file 同源;默认路径锚定首个持久授权目录。
//! 6. **大小上限**:50MB(图片够用),流式边下边累计,超限中止 + 删半成品(对齐 download_file)。
use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};
use base64::{engine::general_purpose::STANDARD, Engine as _};
use serde_json::{json, Value};
use tokio::io::AsyncWriteExt;
use tokio::sync::RwLock;
use df_storage::crud::AiProviderRepo;
use df_storage::db::Database;
use df_storage::models::AiProviderRecord;
use df_storage::secret::resolve_provider_secret;
use crate::commands::ai::http::{
build_client, execute_with_redirects, resolve_and_check_host, validate_url, MAX_REDIRECTS,
};
use crate::commands::ai::tool_registry::resolve_workspace_path_with_allowed;
use crate::state::AllowedDirs;
/// 下载图片大小硬上限(字节)。50MB —— 覆盖高分辨率图片(4K壁纸级 < 20MB),超限中止。
/// 图片 URL 来自 provider 响应可能被篡改指向超大文件,流式边下边累计防磁盘炸。
const MAX_IMAGE_BYTES: u64 = 50 * 1024 * 1024;
/// 生成图片 POST 请求超时(秒)。图像生成模型耗时较高(高分辨率 10-30s),给 120s 余量。
const GENERATE_TIMEOUT_SECS: u64 = 120;
/// generate_image 工具 handler 入口(供 tools/generate_image.rs register 调用)。
///
/// 参数:
/// - prompt: 必填,图片描述(中英文均可,空串报错)
/// - model: 默认 `sensenova-u1-fast`(SenseNova U1 信息图生成,走 /v1/images/generations)
/// - size: 可选,图片尺寸字符串(SenseNova U1 常用 `2752x1536`/`1536x2752`,DALL-E 常用 `1024x1024`)
/// - n: 默认 1,生成数量,clamp [1, 4](多张图取首张 url/b64 落地,余下仅返回 url 列表)
/// - provider_id: 可选,指定 provider(多 provider 时);不传则按 model 名匹配厂商 host
/// - output_path: 可选,落盘路径;不传则默认 `{首个持久授权目录}/generated_images/{model}-{时间戳}.png`
///
/// 返回 {path, url(或 null 若仅 b64), model, provider_id, bytes_written, elapsed_ms}
///
/// 注:handler 接收 db + allowed_dirs(经 declare_tool! 闭包捕获 tuple Arc 传入)。
/// 路径校验在本函数内完成(需 allowed_dirs 快照,与 http_request output_file 同款)。
pub(crate) async fn execute_generate_image(
args: Value,
db: &Arc<Database>,
allowed_dirs: &Arc<RwLock<AllowedDirs>>,
) -> anyhow::Result<Value> {
let started = Instant::now();
// ── 参数解析 ──
let prompt = args.get("prompt").and_then(|v| v.as_str())
.ok_or_else(|| anyhow::anyhow!("缺少 prompt 参数"))?
.trim()
.to_string();
if prompt.is_empty() {
anyhow::bail!("prompt 不能为空");
}
let model = args.get("model").and_then(|v| v.as_str())
.unwrap_or("sensenova-u1-fast")
.trim()
.to_string();
if model.is_empty() {
anyhow::bail!("model 不能为空");
}
let size = args.get("size").and_then(|v| v.as_str())
.map(|s| s.trim().to_string())
.filter(|s| !s.is_empty());
let n = args.get("n").and_then(|v| v.as_u64())
.unwrap_or(1)
.clamp(1, 4) as u32;
let provider_id = args.get("provider_id").and_then(|v| v.as_str())
.map(|s| s.trim().to_string())
.filter(|s| !s.is_empty());
let output_path_raw = args.get("output_path").and_then(|v| v.as_str())
.map(|s| s.trim().to_string())
.filter(|s| !s.is_empty());
// ── 选 provider(enabled && openai_compat,按 model 名匹配厂商 host 兜底首个) ──
let providers = AiProviderRepo::new(db).list_all().await?;
let provider = select_image_provider(&providers, &model, provider_id.as_deref())?;
let provider_id_resolved = provider.id.clone();
// ── 解析 api_key(DB 优先 fallback keyring,空则报错) ──
let api_key = resolve_provider_secret(&provider);
let api_key = api_key.trim();
if api_key.is_empty() {
anyhow::bail!(
"provider「{}」未配置 api_key,请在设置填入并保存",
provider.name
);
}
// ── 拼端点 URL ──
let endpoint = build_images_url(&provider.base_url);
// ── POST 请求(provider 域名非用户输入,SSRF 风险低,直接 reqwest) ──
let body = {
let mut m = serde_json::Map::new();
m.insert("model".into(), json!(model));
m.insert("prompt".into(), json!(prompt));
m.insert("n".into(), json!(n));
if let Some(sz) = &size {
m.insert("size".into(), json!(sz));
}
Value::Object(m)
};
let client = reqwest::Client::builder()
.timeout(Duration::from_secs(GENERATE_TIMEOUT_SECS))
.connect_timeout(Duration::from_secs(15))
.build()
.map_err(|e| anyhow::anyhow!("HTTP client 构建失败: {}", e))?;
let resp = client.post(&endpoint)
.header("Authorization", format!("Bearer {}", api_key))
.header("Content-Type", "application/json")
.json(&body)
.send()
.await
.map_err(|e| anyhow::anyhow!("调用图像生成端点失败 ({}): {}", endpoint, e))?;
let status = resp.status();
if !status.is_success() {
// 错误响应读 body 摘要(截断 500 chars)帮助定位
let err_text = resp.text().await.unwrap_or_default();
let snippet: String = err_text.chars().take(500).collect();
anyhow::bail!(
"图像生成端点返回 HTTP {}: {}",
status.as_u16(),
snippet
);
}
let resp_body: Value = resp.json().await
.map_err(|e| anyhow::anyhow!("解析图像生成响应 JSON 失败: {}", e))?;
// ── 解析响应:优先 data[0].url,次选 data[0].b64_json ──
let data0 = resp_body.get("data").and_then(|d| d.get(0))
.ok_or_else(|| {
let snippet = body_snippet(&resp_body);
anyhow::anyhow!("图像生成响应缺 data[0](响应摘要: {})", snippet)
})?;
let image_url_opt = data0.get("url").and_then(|u| u.as_str())
.filter(|s| !s.trim().is_empty())
.map(|s| s.trim().to_string());
let b64_opt = data0.get("b64_json").and_then(|b| b.as_str())
.filter(|s| !s.trim().is_empty())
.map(|s| s.trim().to_string());
if image_url_opt.is_none() && b64_opt.is_none() {
let snippet = body_snippet(&resp_body);
anyhow::bail!(
"图像生成响应 data[0] 既无 url 也无 b64_json(响应摘要: {})",
snippet
);
}
// ── 解析 output_path(LLM 传则走白名单校验;不传则默认 generated_images/ 时间戳名) ──
let snap = allowed_dirs.read().await.clone();
let final_path = if let Some(raw) = &output_path_raw {
resolve_workspace_path_with_allowed(raw, &snap)?
.to_str()
.ok_or_else(|| anyhow::anyhow!("output_path 含非法字符"))?
.to_string()
} else {
// 默认:首个持久授权目录/generated_images/{model}-{YYYYMMDD-HHMMSS}.png
let root = snap.first_persistent_dir()
.ok_or_else(|| anyhow::anyhow!(
"未传 output_path 且无持久授权目录,请先绑定项目目录或显式传 output_path"
))?;
let ts = chrono::Local::now().format("%Y%m%d-%H%M%S");
// model 含 / 等路径分隔符风险:替换为 - 防穿透到子目录
let safe_model = model.replace(['/', '\\', ':', '*', '?', '"', '<', '>', '|'], "-");
let fname = format!("{}-{}.png", safe_model, ts);
root.join("generated_images").join(fname)
.to_str()
.ok_or_else(|| anyhow::anyhow!("默认 output_path 含非法字符"))?
.to_string()
};
// ── 落盘:url 走 SSRF 防护下载流式写;b64 直接写解码后字节 ──
let written_bytes = if let Some(image_url) = &image_url_opt {
download_image_with_ssrf(image_url, &final_path).await?
} else {
// b64 路径:解码后直接写
let b64 = b64_opt.as_ref().unwrap();
let decoded = STANDARD.decode(b64)
.map_err(|e| anyhow::anyhow!("b64_json 解码失败: {}", e))?;
if (decoded.len() as u64) > MAX_IMAGE_BYTES {
anyhow::bail!(
"b64_json 解码后 {} 字节超过 {} 上限",
decoded.len(),
MAX_IMAGE_BYTES
);
}
write_atomic(&final_path, &decoded).await?;
decoded.len() as u64
};
let elapsed_ms = started.elapsed().as_millis() as u64;
Ok(json!({
"path": final_path,
"url": image_url_opt,
"model": model,
"provider_id": provider_id_resolved,
"bytes_written": written_bytes,
"elapsed_ms": elapsed_ms,
}))
}
/// 从 provider 列表筛选图像生成可用 provider。
///
/// 筛选条件:`enabled && provider_type == "openai_compat"`(图像端点是 OpenAI 风格,
/// anthropic provider 排除)。
///
/// 选择策略:
/// 1. LLM 传 `provider_id` → 按 id 精确匹配(找不到报错)。
/// 2. 否则按 model 名匹配厂商 host 提示:
/// - model 含 `sensenova` → 找 base_url 含 `sensenova`
/// - model 含 `dall-e` → 找 base_url 含 `openai`
/// - model 含 `google`/`imagen` → 找 base_url 含 `googleapis`
/// 3. host 命中 → 用首个匹配;host 未命中 → 兜底取筛选结果首个。
/// 4. 筛选后空 → 报错(无可用 openai_compat provider)。
fn select_image_provider(
providers: &[AiProviderRecord],
model: &str,
provider_id: Option<&str>,
) -> anyhow::Result<AiProviderRecord> {
// 筛 enabled + openai_compat
let candidates: Vec<&AiProviderRecord> = providers.iter()
.filter(|p| p.enabled && p.provider_type == "openai_compat")
.collect();
// 路径1:LLM 指定 provider_id,精确匹配
if let Some(pid) = provider_id {
let found = candidates.iter().find(|p| p.id == pid)
.ok_or_else(|| anyhow::anyhow!(
"未找到指定 provider(id={},且需 enabled + openai_compat),请检查 provider_id 参数",
pid
))?;
return Ok((*found).clone());
}
if candidates.is_empty() {
anyhow::bail!("无可用 openai_compat provider,请先在设置配置支持图像生成的 provider");
}
// 路径2:按 model 名匹配厂商 host
let model_lower = model.to_lowercase();
let host_hint: Option<&str> = if model_lower.contains("sensenova") {
Some("sensenova")
} else if model_lower.contains("dall-e") {
Some("openai")
} else if model_lower.contains("google") || model_lower.contains("imagen") {
Some("googleapis")
} else {
None
};
if let Some(hint) = host_hint {
let hit = candidates.iter()
.find(|p| p.base_url.to_lowercase().contains(hint));
if let Some(p) = hit {
return Ok((*p).clone());
}
}
// 路径3:兜底首个候选
Ok(candidates[0].clone())
}
/// 拼接图像生成端点 URL。
///
/// 规则(base_url 先去尾 `/`,对齐 `model_fetch_helpers::build_models_url` 思路,
/// 但目标是 `/images/generations`):
/// - 以 `/v1` 或 `/v4` 结尾 → 直接补 `/images/generations`(避免 `/v1/v1/images/...`)
/// - 否则 → 补 `/v1/images/generations`(最常见:用户填 `https://api.sensenova.cn`)
fn build_images_url(base_url: &str) -> String {
let url = base_url.trim_end_matches('/');
if url.ends_with("/v1") || url.ends_with("/v4") {
format!("{url}/images/generations")
} else {
format!("{url}/v1/images/generations")
}
}
/// 流式下载图片 URL 到 output_path,带 SSRF 防护 + 大小上限 + 原子写。
///
/// 图片 URL 来自 provider 响应可能被恶意 provider 篡改指向内网,故复用 SSRF 防护
/// (validate_url + resolve_and_check_host + build_client + execute_with_redirects)。
async fn download_image_with_ssrf(image_url: &str, output_path: &str) -> anyhow::Result<u64> {
// SSRF 三层校验:词法 URL → DNS resolve IP → 每跳重定向重复
let (_scheme, host, port) = validate_url(image_url)?;
resolve_and_check_host(&host, port).await?;
let client = build_client(Duration::from_secs(GENERATE_TIMEOUT_SECS))?;
let resp = execute_with_redirects(
&client,
reqwest::Method::GET,
image_url.to_string(),
&HashMap::new(), // 图片下载不带自定义头
&None,
MAX_REDIRECTS,
)
.await?;
if !resp.status().is_success() {
anyhow::bail!(
"下载图片失败:HTTP {} {}({})",
resp.status().as_u16(),
resp.status().canonical_reason().unwrap_or(""),
image_url
);
}
// 流式写文件(避免大图 OOM),边下边累计,超 MAX_IMAGE_BYTES 立即中止 + 删半成品。
// 父目录不存在自动创建。原子写:tmp → rename(同目录不跨卷)。
let target = std::path::Path::new(output_path);
if let Some(parent) = target.parent() {
tokio::fs::create_dir_all(parent).await
.map_err(|e| anyhow::anyhow!("创建目录失败: {}", e))?;
}
let tmp = format!("{}.tmp-img", output_path);
let mut file = tokio::fs::File::create(&tmp).await
.map_err(|e| anyhow::anyhow!("创建临时文件失败 ({}): {}", tmp, e))?;
let mut stream = resp.bytes_stream();
let mut written: u64 = 0;
let mut exceeded = false;
use futures::StreamExt;
while let Some(chunk) = stream.next().await {
let chunk = chunk.map_err(|e| anyhow::anyhow!("读取图片流失败: {}", e))?;
written += chunk.len() as u64;
if written > MAX_IMAGE_BYTES {
exceeded = true;
break;
}
file.write_all(&chunk).await
.map_err(|e| anyhow::anyhow!("写入图片失败: {}", e))?;
}
file.flush().await
.map_err(|e| anyhow::anyhow!("flush 图片文件失败: {}", e))?;
drop(file);
if exceeded {
let _ = tokio::fs::remove_file(&tmp).await;
anyhow::bail!(
"图片超过 {} 字节上限,已中止并清理临时文件",
MAX_IMAGE_BYTES
);
}
if let Err(e) = tokio::fs::rename(&tmp, output_path).await {
let _ = tokio::fs::remove_file(&tmp).await;
return Err(anyhow::anyhow!("原子替换失败: {}", e));
}
Ok(written)
}
/// 原子写字节到指定路径(tmp → rename,同目录不跨卷)。
/// b64 路径用此函数(已全部解码在内存,无需流式)。
async fn write_atomic(output_path: &str, data: &[u8]) -> anyhow::Result<()> {
let target = std::path::Path::new(output_path);
if let Some(parent) = target.parent() {
tokio::fs::create_dir_all(parent).await
.map_err(|e| anyhow::anyhow!("创建目录失败: {}", e))?;
}
let tmp = format!("{}.tmp-img", output_path);
tokio::fs::write(&tmp, data).await
.map_err(|e| {
let _ = std::fs::remove_file(&tmp);
anyhow::anyhow!("写入图片失败: {}", e)
})?;
tokio::fs::rename(&tmp, output_path).await
.map_err(|e| {
let _ = std::fs::remove_file(&tmp);
anyhow::anyhow!("原子替换失败: {}", e)
})?;
Ok(())
}
/// 响应 body 截断为最多 500 chars 的字符串(用于错误信息),供定位用。
fn body_snippet(v: &Value) -> String {
let s = v.to_string();
s.chars().take(500).collect()
}
// ============================================================
// 单元测试
//
// 覆盖:
// ① build_images_url:端点拼接规则(纯函数,零依赖)
// ② select_image_provider:provider 筛选逻辑(纯函数,内存构造 record)
// ③ handler 参数边界(prompt 缺失/空、provider_id 不存在)
// 真实 API 调用走 #[ignore](CI 无凭证/无网时跳过,本地手跑)。
// ============================================================
#[cfg(test)]
mod tests {
use super::*;
// ── build_images_url:端点拼接规则 ──
#[test]
fn build_images_url_v1_suffix() {
// /v1 结尾 → 直接补 /images/generations
assert_eq!(
build_images_url("https://api.openai.com/v1"),
"https://api.openai.com/v1/images/generations"
);
// 不应产生 /v1/v1/images/...
assert!(!build_images_url("https://api.openai.com/v1").contains("/v1/v1/"));
}
#[test]
fn build_images_url_v4_suffix() {
// GLM v4 风格 base
assert_eq!(
build_images_url("https://open.bigmodel.cn/api/paas/v4"),
"https://open.bigmodel.cn/api/paas/v4/images/generations"
);
}
#[test]
fn build_images_url_bare_base() {
// 用户只填根 → 补 /v1/images/generations
assert_eq!(
build_images_url("https://api.sensenova.cn"),
"https://api.sensenova.cn/v1/images/generations"
);
}
#[test]
fn build_images_url_strips_trailing_slash() {
assert_eq!(
build_images_url("https://api.x.com/v1/"),
"https://api.x.com/v1/images/generations"
);
}
// ── select_image_provider:筛选 + 匹配逻辑 ──
/// 构造测试用 AiProviderRecord(只填筛选关键字段)。
fn mk_provider(id: &str, enabled: bool, ptype: &str, base_url: &str) -> AiProviderRecord {
AiProviderRecord {
id: id.into(),
name: id.into(),
provider_type: ptype.into(),
api_key: String::new(),
base_url: base_url.into(),
default_model: String::new(),
models: None,
model_configs: Vec::new(),
is_default: false,
config: None,
created_at: "0".into(),
updated_at: "0".into(),
enabled,
weight: 50,
}
}
#[test]
fn select_provider_no_candidates_errors() {
// 全 disabled 或非 openai_compat → 报错
let providers = vec![
mk_provider("p1", false, "openai_compat", "https://x"),
mk_provider("p2", true, "anthropic", "https://y"),
];
let err = select_image_provider(&providers, "any-model", None).unwrap_err();
assert!(format!("{}", err).contains("无可用"));
}
#[test]
fn select_provider_provider_id_match() {
let providers = vec![
mk_provider("p1", true, "openai_compat", "https://api.sensenova.cn"),
mk_provider("p2", true, "openai_compat", "https://api.openai.com"),
];
let got = select_image_provider(&providers, "any", Some("p2")).unwrap();
assert_eq!(got.id, "p2");
}
#[test]
fn select_provider_provider_id_not_found_errors() {
let providers = vec![mk_provider("p1", true, "openai_compat", "https://x")];
let err = select_image_provider(&providers, "any", Some("nonexistent")).unwrap_err();
assert!(format!("{}", err).contains("未找到指定 provider"));
}
#[test]
fn select_provider_host_match_sensenova() {
// model 含 sensenova → 匹配 base_url 含 sensenova 的 provider
let providers = vec![
mk_provider("openai", true, "openai_compat", "https://api.openai.com"),
mk_provider("sense", true, "openai_compat", "https://api.sensenova.cn"),
];
let got = select_image_provider(&providers, "sensenova-u1-fast", None).unwrap();
assert_eq!(got.id, "sense");
}
#[test]
fn select_provider_host_match_dalle() {
// model 含 dall-e → 匹配 base_url 含 openai 的 provider
let providers = vec![
mk_provider("glm", true, "openai_compat", "https://open.bigmodel.cn"),
mk_provider("oai", true, "openai_compat", "https://api.openai.com"),
];
let got = select_image_provider(&providers, "dall-e-3", None).unwrap();
assert_eq!(got.id, "oai");
}
#[test]
fn select_provider_host_unmatched_falls_back_first() {
// model 名无厂商提示 → 兜底取筛选结果首个
let providers = vec![
mk_provider("first", true, "openai_compat", "https://api.unknown.com"),
mk_provider("second", true, "openai_compat", "https://api.other.com"),
];
let got = select_image_provider(&providers, "my-custom-model", None).unwrap();
assert_eq!(got.id, "first");
}
// ── handler 参数边界(走 db 需 in-memory db) ──
#[tokio::test]
async fn handler_missing_prompt_errors() {
// 缺 prompt → 参数解析阶段报错(不发请求)
let db = Arc::new(Database::open_in_memory().await.unwrap());
let allowed = Arc::new(RwLock::new(AllowedDirs::default()));
let args = json!({ "model": "sensenova-u1-fast" });
let err = execute_generate_image(args, &db, &allowed).await.unwrap_err();
assert!(format!("{}", err).contains("prompt"));
}
#[tokio::test]
async fn handler_empty_prompt_errors() {
let db = Arc::new(Database::open_in_memory().await.unwrap());
let allowed = Arc::new(RwLock::new(AllowedDirs::default()));
let args = json!({ "prompt": " " });
let err = execute_generate_image(args, &db, &allowed).await.unwrap_err();
assert!(format!("{}", err).contains("prompt 不能为空"));
}
#[tokio::test]
async fn handler_no_provider_errors() {
// 无 openai_compat provider → 报错(不发请求)
let db = Arc::new(Database::open_in_memory().await.unwrap());
let allowed = Arc::new(RwLock::new(AllowedDirs::default()));
let args = json!({ "prompt": "一只猫" });
let err = execute_generate_image(args, &db, &allowed).await.unwrap_err();
assert!(format!("{}", err).contains("无可用") || format!("{}", err).contains("provider"));
}
}