diff --git a/Cargo.lock b/Cargo.lock index a58c3a8..a060e13 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -567,6 +567,29 @@ dependencies = [ "typenum", ] +[[package]] +name = "cssparser" +version = "0.37.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8c9cdaae01d5ed7882b04d795e7f752f46ff52d2fa3b50a20d28c464510bba98" +dependencies = [ + "cssparser-macros", + "dtoa-short", + "itoa", + "phf", + "smallvec", +] + +[[package]] +name = "cssparser-macros" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10a2a99df6e410a8ff4245aa2006499ea662245f967cc7c0a38c83ef8eb44dbf" +dependencies = [ + "quote", + "syn", +] + [[package]] name = "curve25519-dalek" version = "4.1.3" @@ -614,6 +637,27 @@ dependencies = [ "powerfmt", ] +[[package]] +name = "derive_more" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d751e9e49156b02b44f9c1815bcb94b984cdcc4396ecc32521c739452808b134" +dependencies = [ + "derive_more-impl", +] + +[[package]] +name = "derive_more-impl" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "799a97264921d8623a957f6c3b9011f3b5492f557bbb7a5a19b7fa6d06ba8dcb" +dependencies = [ + "proc-macro2", + "quote", + "rustc_version", + "syn", +] + [[package]] name = "dialoguer" version = "0.12.0" @@ -676,6 +720,21 @@ version = "0.15.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1aaf95b3e5c8f23aa320147307562d361db0ae0d51242340f558153b4eb2439b" +[[package]] +name = "dtoa" +version = "1.0.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c3cf4824e2d5f025c7b531afcb2325364084a16806f6d47fbc1f5fbd9960590" + +[[package]] +name = "dtoa-short" +version = "0.3.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd1511a7b6a56299bd043a9c167a6d2bfb37bf84a6dfceaba651168adfb43c87" +dependencies = [ + "dtoa", +] + [[package]] name = "ed25519" version = "2.2.3" @@ -700,6 +759,12 @@ dependencies = [ "zeroize", ] +[[package]] +name = "ego-tree" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b04dc5a38e4f151a79d9f2451ae6037fb6eaf5cba34771f44781f80e508498e3" + [[package]] name = "either" version = "1.16.0" @@ -1003,6 +1068,15 @@ dependencies = [ "version_check", ] +[[package]] +name = "getopts" +version = "0.2.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfe4fbac503b8d1f88e6676011885f34b7174f46e59956bba534ba83abded4df" +dependencies = [ + "unicode-width 0.2.2", +] + [[package]] name = "getrandom" version = "0.2.17" @@ -1144,6 +1218,16 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "html5ever" +version = "0.39.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "46a1761807faccc9a19e86944bbf40610014066306f96edcdedc2fb714bcb7b8" +dependencies = [ + "log", + "markup5ever", +] + [[package]] name = "http" version = "1.4.1" @@ -1280,6 +1364,7 @@ dependencies = [ "rand 0.9.4", "reqwest", "rustyline", + "scraper", "serde", "serde_json", "sha2", @@ -1637,6 +1722,17 @@ dependencies = [ "imgref", ] +[[package]] +name = "markup5ever" +version = "0.39.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7122d987ec5f704ee56f6e5b41a7d93722e9aae27ae07cafa4036c4d3f9757de" +dependencies = [ + "log", + "tendril", + "web_atoms", +] + [[package]] name = "matchers" version = "0.2.0" @@ -1987,6 +2083,59 @@ version = "2.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" +[[package]] +name = "phf" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c1562dc717473dbaa4c1f85a36410e03c047b2e7df7f45ee938fbef64ae7fadf" +dependencies = [ + "phf_macros", + "phf_shared", + "serde", +] + +[[package]] +name = "phf_codegen" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "49aa7f9d80421bca176ca8dbfebe668cc7a2684708594ec9f3c0db0805d5d6e1" +dependencies = [ + "phf_generator", + "phf_shared", +] + +[[package]] +name = "phf_generator" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "135ace3a761e564ec88c03a77317a7c6b80bb7f7135ef2544dbe054243b89737" +dependencies = [ + "fastrand", + "phf_shared", +] + +[[package]] +name = "phf_macros" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "812f032b54b1e759ccd5f8b6677695d5268c588701effba24601f6932f8269ef" +dependencies = [ + "phf_generator", + "phf_shared", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "phf_shared" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e57fef6bc5981e38c2ce2d63bfa546861309f875b8a75f092d1d54ae2d64f266" +dependencies = [ + "siphasher", +] + [[package]] name = "pin-project-lite" version = "0.2.17" @@ -2063,6 +2212,12 @@ dependencies = [ "zerocopy", ] +[[package]] +name = "precomputed-hash" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "925383efa346730478fb4838dbe9137d2a47675ad789c546d150a6e1dd4ab31c" + [[package]] name = "prettyplease" version = "0.2.37" @@ -2420,6 +2575,12 @@ dependencies = [ "zeroize", ] +[[package]] +name = "rustc-hash" +version = "2.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94300abf3f1ae2e2b8ffb7b58043de3d399c73fa6f4b73826402a5c457614dbe" + [[package]] name = "rustc_version" version = "0.4.1" @@ -2524,6 +2685,21 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" +[[package]] +name = "scraper" +version = "0.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bdd0be4d296f048bfb06dd01bbc80ef789ddd2e55583e8d2e6b804942abfabc2" +dependencies = [ + "cssparser", + "ego-tree", + "getopts", + "html5ever", + "precomputed-hash", + "selectors", + "tendril", +] + [[package]] name = "security-framework" version = "3.7.0" @@ -2547,6 +2723,25 @@ dependencies = [ "libc", ] +[[package]] +name = "selectors" +version = "0.38.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8adfa1c298912827b8a28b223b3b874357397ae706e6190acd9bf28cee99114d" +dependencies = [ + "bitflags", + "cssparser", + "derive_more", + "log", + "new_debug_unreachable", + "phf", + "phf_codegen", + "precomputed-hash", + "rustc-hash", + "servo_arc", + "smallvec", +] + [[package]] name = "semver" version = "1.0.28" @@ -2608,6 +2803,15 @@ dependencies = [ "serde", ] +[[package]] +name = "servo_arc" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "170fb83ab34de17dc69aa7c67482b22218ddb85da56546f9bd6b929e32a05930" +dependencies = [ + "stable_deref_trait", +] + [[package]] name = "sha1" version = "0.10.6" @@ -2686,6 +2890,12 @@ dependencies = [ "quote", ] +[[package]] +name = "siphasher" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ee5873ec9cce0195efcb7a4e9507a04cd49aec9c83d0389df45b1ef7ba2e649" + [[package]] name = "slab" version = "0.4.12" @@ -2933,6 +3143,30 @@ version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" +[[package]] +name = "string_cache" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a18596f8c785a729f2819c0f6a7eae6ebeebdfffbfe4214ae6b087f690e31901" +dependencies = [ + "new_debug_unreachable", + "parking_lot", + "phf_shared", + "precomputed-hash", +] + +[[package]] +name = "string_cache_codegen" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "585635e46db231059f76c5849798146164652513eb9e8ab2685939dd90f29b69" +dependencies = [ + "phf_generator", + "phf_shared", + "proc-macro2", + "quote", +] + [[package]] name = "stringprep" version = "0.1.5" @@ -3027,6 +3261,16 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "tendril" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4790fc369d5a530f4b544b094e31388b9b3a37c0f4652ade4505945f5660d24" +dependencies = [ + "new_debug_unreachable", + "utf-8", +] + [[package]] name = "thiserror" version = "2.0.18" @@ -3404,6 +3648,12 @@ dependencies = [ "serde", ] +[[package]] +name = "utf-8" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09cc8ee72d2a9becf2f2febe0205bbed8fc6615b7cb429ad062dc7b7ddd036a9" + [[package]] name = "utf8_iter" version = "1.0.4" @@ -3608,6 +3858,18 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "web_atoms" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7cff6eef815df1834fd250e3a2ff436044d82a9f1bc1980ca1dbdf07effc538" +dependencies = [ + "phf", + "phf_codegen", + "string_cache", + "string_cache_codegen", +] + [[package]] name = "weezl" version = "0.1.12" diff --git a/Cargo.toml b/Cargo.toml index b3bf394..0d4f59c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -43,3 +43,4 @@ console = "0.16.3" dirs = "6.0.0" # Ed25519 JWT 签名(天气 API) ed25519-dalek = { version = "2", features = ["pem"] } +scraper = "0.27.0" diff --git a/src/cli.rs b/src/cli.rs index 9280560..35b3b8f 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -2,8 +2,14 @@ use clap::{Parser, Subcommand}; /// iAs — 微信 AI 智能助手 /// -/// 支持扫码登录、长轮询收发消息、AI 自动回复(DeepSeek/LM Studio)、 -/// 工具调用(天气/搜索/邮件/Shell/定时任务等)、Token 用量统计。 +/// 支持扫码登录、长轮询收发消息、AI 自动回复(DeepSeek)、 +/// 内置工具(天气/搜索/备忘录/日期时间等)、Token 用量统计。 +/// +/// 快速开始: +/// ias login 扫码登录微信 +/// ias listen --llm 启动 AI 自动回复 +/// ias tool weather 北京 查询天气 +/// ias tool search "最新新闻" 联网搜索 #[derive(Parser, Debug)] #[command(name = "ias", version, about, long_about = None)] pub struct Cli { @@ -15,8 +21,12 @@ pub struct Cli { pub enum Commands { /// 扫码登录微信 /// - /// 在终端显示二维码,扫描后完成登录。 + /// 在终端显示登录二维码,扫描后完成登录。 /// 认证信息存入 PostgreSQL 或本地文件。 + /// + /// 示例: + /// ias login + /// ias login --timeout 600 Login { /// 登录超时秒数(默认 480s) #[arg(long, default_value = "480")] @@ -26,62 +36,55 @@ pub enum Commands { /// 监听微信消息(长轮询) /// /// 启动微信消息监听,接收并处理消息。 - /// 使用 --llm 启用 AI 自动回复,--echo 原样回显。 /// /// 示例: /// ias listen # 仅监听,不做回复 /// ias listen --llm # AI 自动回复 /// ias listen --echo # 原样回显 Listen { - /// 启用 AI 自动回复(DeepSeek/LM Studio) - #[arg(long, help = "启用 AI 自动回复")] + /// 启用 AI 自动回复 + #[arg(long)] llm: bool, - /// 仅回显消息(不调用 AI) - #[arg(long, help = "原样回显收到的消息")] + /// 原样回显收到的消息 + #[arg(long)] echo: bool, }, /// 发送文本消息 /// - /// 向指定微信用户发送文本。 - /// /// 示例: - /// ias send --to --text "你好" + /// ias send --to wxid_xxx --text "你好" Send { /// 目标用户 ID - #[arg(long, help = "微信用户 ID")] + #[arg(long)] to: String, /// 消息内容 - #[arg(long, help = "文本消息内容")] + #[arg(long)] text: String, /// 上下文令牌 - #[arg(long, help = "微信上下文令牌")] + #[arg(long)] context_token: Option, }, /// 查看当前登录状态 - /// - /// 显示已登录的账号、Token 前缀、API 地址及存储方式。 Whoami, /// 查看 LLM Token 使用统计 /// - /// 查询 LLM 调用的 Token 消耗、缓存命中率。 - /// 默认查询最近 7 天,支持时间范围和模型过滤。 + /// 默认查询最近 7 天。 /// /// 示例: /// ias usage - /// ias usage --since 2026-06-01T00:00:00Z - /// ias usage --model deepseek-v4-flash + /// ias usage --since 2026-06-01T00:00:00Z --model deepseek-v4-flash Usage { - /// 起始时间(RFC3339 格式,默认 7 天前) - #[arg(long, help = "起始时间(如 2026-06-01T00:00:00Z)")] + /// 起始时间(RFC3339 格式) + #[arg(long)] since: Option, - /// 结束时间(RFC3339 格式,默认现在) + /// 结束时间(RFC3339 格式) #[arg(long)] until: Option, @@ -90,13 +93,228 @@ pub enum Commands { model: Option, }, - /// 后台服务模式 + /// 后台服务模式(等同 listen --llm + 文件日志) /// - /// 等同 listen --llm,同时日志写入 ~/.ias/logs/ 滚动文件。 - /// 配合 systemd 实现开机自启和进程守护。 + /// 配合 systemd 实现开机自启。 /// /// 示例: - /// ias service # 前台运行 - /// sudo systemctl enable ias --now # systemd 管理 + /// ias service Service, + + /// 守护进程模式(daemon + worker 分离架构) + /// + /// Daemon 持有 WeChat 长连接和数据库,每条消息 spawn 独立 worker。 + /// Worker 代码更新后编译即可,下一条消息自动用新版本。 + /// + /// 示例: + /// ias daemon + /// ias daemon --sock /tmp/ias.sock + Daemon { + /// Unix Domain Socket 路径 + #[arg(long, default_value = "/tmp/ias_daemon.sock")] + sock: String, + }, + + /// Worker 进程(由 daemon 自动 spawn,也可手动测试) + /// + /// 示例: + /// ias worker --sock /tmp/ias_daemon.sock + Worker { + /// Unix Domain Socket 路径 + #[arg(long, default_value = "/tmp/ias_daemon.sock")] + sock: String, + }, + + /// 调用内置工具(无需登录,独立运行) + /// + /// 可直接在终端调用天气、搜索、备忘录、日期时间、高德地图等工具。 + #[command( + subcommand, + after_help = "示例:\n ias tool datetime\n ias tool weather 北京\n ias tool weather 上海 --days 7\n ias tool search \"Rust 2024 edition\"\n ias tool fetch https://example.com/article\n ias tool memos list\n ias tool memos add \"明天下午2点开会\"\n ias tool memos delete 3\n ias tool amap poi-search --keywords 肯德基 --city 北京\n ias tool amap geocode --address 西直门\n ias tool amap route-plan --type walking --origin 116.397,39.909 --destination 116.427,39.903\n ias tool amap travel-plan --city 北京 --interests 景点,美食 --route-type walking" + )] + Tool(ToolCommand), +} + +// ─── 工具子命令 ─── + +#[derive(Subcommand, Debug, Clone)] +pub enum ToolCommand { + /// 获取当前日期时间(北京时间 UTC+8) + Datetime, + + /// 管理备忘录 + #[command(subcommand)] + Memos(MemosAction), + + /// 查询指定城市的天气 + /// + /// 示例: + /// ias tool weather 北京 + /// ias tool weather 上海 --days 7 + Weather { + /// 城市名称 + location: String, + + /// 查询天数(1-7,默认 3) + #[arg(long, default_value = "3")] + days: u32, + }, + + /// 联网搜索最新信息(需配置 TAVILY_API_KEY) + /// + /// 示例: + /// ias tool search "最新 AI 新闻" + /// ias tool search "Rust async" --max-results 10 --no-answer + /// ias tool search "比特币" --topic finance --time-range week + Search { + /// 搜索关键词 + query: String, + + /// 最大结果数(默认 5) + #[arg(long, default_value = "5")] + max_results: u32, + + /// 不包含 AI 生成的摘要 + #[arg(long)] + no_answer: bool, + + /// 搜索主题: general / news / finance + #[arg(long)] + topic: Option, + + /// 发布时间过滤: day / week / month / year + #[arg(long)] + time_range: Option, + }, + + /// 抓取网页并提取正文内容 + Fetch { + /// 网页 URL + url: String, + }, + + /// 高德地图工具(POI搜索、地理编码、路径规划、旅游规划) + #[command(subcommand)] + Amap(AmapAction), +} + +#[derive(Subcommand, Debug, Clone)] +pub enum AmapAction { + /// POI(地点)搜索 + PoiSearch { + /// 搜索关键词 + #[arg(long)] + keywords: String, + + /// 城市名称 + #[arg(long)] + city: Option, + + /// 中心点坐标(经度,纬度),用于周边搜索 + #[arg(long)] + location: Option, + + /// 搜索半径(米),默认 1000 + #[arg(long)] + radius: Option, + + /// 页码,默认 1 + #[arg(long, default_value = "1")] + page: u32, + + /// 每页数量,默认 10 + #[arg(long, default_value = "10")] + offset: u32, + }, + + /// 地理编码:地址 → 坐标 + Geocode { + /// 地址名称 + #[arg(long)] + address: String, + + /// 城市名称(可选) + #[arg(long)] + city: Option, + }, + + /// 逆地理编码:坐标 → 地址 + ReverseGeocode { + /// 坐标(经度,纬度) + #[arg(long)] + location: String, + }, + + /// 路径规划 + RoutePlan { + /// 出行方式:walking/driving/riding/transit + #[arg(long)] + r#type: String, + + /// 起点坐标(经度,纬度) + #[arg(long)] + origin: String, + + /// 终点坐标(经度,纬度) + #[arg(long)] + destination: String, + + /// 城市名称(公交方式必填) + #[arg(long)] + city: Option, + + /// 途经点(驾车可选,多个用;分隔) + #[arg(long)] + waypoints: Option, + + /// 策略(驾车默认10躲避拥堵,公交默认0最快捷) + #[arg(long)] + strategy: Option, + }, + + /// 智能旅游规划 + TravelPlan { + /// 城市名称 + #[arg(long)] + city: String, + + /// 兴趣点关键词(逗号分隔),如 景点,美食,酒店 + #[arg(long, default_value = "景点,美食")] + interests: String, + + /// 路线类型:walking/driving/riding/transit + #[arg(long, default_value = "walking")] + route_type: String, + }, + + /// 生成地图可视化链接 + MapLink { + /// 地图数据 JSON 数组 + #[arg(long)] + data: String, + }, +} + +#[derive(Subcommand, Debug, Clone)] +pub enum MemosAction { + /// 列出所有备忘录 + List, + + /// 添加备忘录 + /// + /// 示例: + /// ias tool memos add "明天下午2点开会" + Add { + /// 备忘录内容 + content: String, + }, + + /// 删除指定 ID 的备忘录 + /// + /// 示例: + /// ias tool memos delete 3 + Delete { + /// 备忘录 ID + id: i64, + }, } diff --git a/src/context/tools.rs b/src/context/tools.rs index 3589f9e..ff01866 100644 --- a/src/context/tools.rs +++ b/src/context/tools.rs @@ -82,7 +82,8 @@ pub async fn read_summaries( } if let Some(pool) = &s.db_pool { - if let Ok(db_summaries) = crate::db::models::load_summaries(pool, 5).await { + let uid = if s.current_user_id.is_empty() { "default" } else { &s.current_user_id }; + if let Ok(db_summaries) = crate::db::models::load_summaries(pool, uid, 5).await { for text in db_summaries { if !lines.iter().any(|l| l.contains(&text.chars().take(30).collect::())) { lines.push(format!("[数据库存档] {}", text)); diff --git a/src/context/types.rs b/src/context/types.rs index 033272f..b05c5c9 100644 --- a/src/context/types.rs +++ b/src/context/types.rs @@ -117,6 +117,20 @@ impl ChatSession { .find(|s| s.reason == SummaryReason::Overflow) } + /// 从预载的历史记录初始化会话(Worker 模式) + pub fn load_from_history(&mut self, history: &[crate::ipc::HistoryEntry]) { + for entry in history { + match entry.role.as_str() { + "user" => self.add_user(entry.content.clone()), + "assistant" => self.add_assistant(entry.content.clone()), + _ => {} + } + } + if !history.is_empty() { + self.last_user_at = Some(chrono::Utc::now()); + } + } + /// 从数据库加载最近的聊天记录到会话中 pub async fn load_recent_messages(&mut self, limit: i64) { let pool = match &self.db_pool { @@ -166,10 +180,15 @@ impl ChatSession { /// 保存摘要到数据库 pub async fn save_summary_to_db(&self, text: &str, reason: &str, message_count: i32) { if let Some(pool) = &self.db_pool { + let uid = if self.current_user_id.is_empty() { + if self.user_id.is_empty() { "default" } else { &self.user_id } + } else { + &self.current_user_id + }; let _ = sqlx::query( "INSERT INTO session_summaries (user_id, summary_text, reason, message_count) VALUES ($1, $2, $3, $4)", ) - .bind(&self.user_id) + .bind(uid) .bind(text) .bind(reason) .bind(message_count) diff --git a/src/daemon.rs b/src/daemon.rs new file mode 100644 index 0000000..5c3661d --- /dev/null +++ b/src/daemon.rs @@ -0,0 +1,927 @@ +//! Daemon — 常驻进程 +//! +//! 持有 WeChat 长连接 + 数据库。每收到一条消息 spawn 一个 worker 进程, +//! 通过 Unix Domain Socket 通信。Worker 代码更新后编译即可,下一条消息 +//! 自动使用新版本。 + +use crate::context::MemoryStore; +use crate::db::Database; +use crate::ipc::{HistoryEntry, OutputFrame, TaskFrame, TaskMessage, recv_output, send_frame}; +use crate::state::StateManager; +use crate::tools::approval::{ApprovalDecision, ApprovalManager}; +use crate::wechat::client::WeChatClient; +use std::collections::{HashMap, VecDeque}; +use std::sync::Arc; +use std::time::Instant; +use tokio::net::{UnixListener, UnixStream}; +use tokio::process::Command; +use tokio::sync::Mutex; +use tracing::{error, info, warn}; + +/// 消息队列:每个用户最多一个活跃 worker,等待队列 FIFO +struct MessageQueue { + /// 当前是否有活跃 worker + active: HashMap, + /// 每个用户的待处理消息队列 + pending: HashMap>, + /// 等待 worker 连接的用户顺序 + waiting: VecDeque, +} + +struct PendingMessage { + from: String, + text: String, + account_id: String, + context_token: Option, + #[allow(dead_code)] + message_id: String, + approved_tool: Option, +} + +impl MessageQueue { + fn new() -> Self { + Self { + active: HashMap::new(), + pending: HashMap::new(), + waiting: VecDeque::new(), + } + } + + fn enqueue(&mut self, user_id: &str, msg: PendingMessage) { + self.pending + .entry(user_id.to_string()) + .or_default() + .push_back(msg); + if !self.active.get(user_id).unwrap_or(&false) && !self.waiting.contains(&user_id.to_string()) { + self.waiting.push_back(user_id.to_string()); + } + } + + /// 取出下一个等待处理的用户(daemon accept worker 后调用) + fn pop_waiting(&mut self) -> Option { + while let Some(uid) = self.waiting.pop_front() { + // 跳过已经活跃的 + if *self.active.get(&uid).unwrap_or(&false) { + continue; + } + // 检查是否有待处理消息 + if self.pending.get(&uid).map_or(true, |q| q.is_empty()) { + continue; + } + self.active.insert(uid.clone(), true); + return Some(uid); + } + None + } + + fn pop_message(&mut self, user_id: &str) -> Option { + self.pending.get_mut(user_id)?.pop_front() + } + + fn deactivate(&mut self, user_id: &str) { + self.active.insert(user_id.to_string(), false); + } + + fn has_pending(&self, user_id: &str) -> bool { + self.pending + .get(user_id) + .map_or(false, |q| !q.is_empty()) + } +} + +/// Daemon 入口 +pub async fn run( + sock_path: String, + database: &Option>, + file_state: &StateManager, + memory_store: &Arc, +) { + // 1. 加载认证 + let auth = match load_auth(database, file_state).await { + Some(a) => a, + None => { + error!("未登录,请先执行 ias login"); + return; + } + }; + + // 2. 创建 WeChat 客户端 + let mut client = WeChatClient::new(Some(auth.base_url.clone())); + client + .set_auth(&auth.token, &auth.account_id, &auth.base_url) + .await; + + // 恢复 runtime buf + if let Some(r) = file_state.load_runtime() { + client.set_updates_buf(&r.get_updates_buf).await; + } + + // 3. 注册监听器 + if let Err(e) = client.notify_start().await { + error!("注册监听器失败: {}", e); + return; + } + + let account_id = auth.account_id.clone(); + + // 4. 绑定 Unix Domain Socket + let _ = std::fs::remove_file(&sock_path); + let listener = match UnixListener::bind(&sock_path) { + Ok(l) => l, + Err(e) => { + error!("绑定 Unix Socket 失败 ({}): {}", sock_path, e); + return; + } + }; + + info!("Daemon 已启动: sock={}", sock_path); + + // 5. 审批管理器 + let approval_manager = Arc::new(ApprovalManager::new( + database.as_ref().map(|db| Arc::new(db.pool().clone())), + )); + + // 6. 消息队列 + let queue = Arc::new(Mutex::new(MessageQueue::new())); + + // 7. 构建工具列表(传给 worker) + let tools_list = build_tools_list(); + + // 8. 环境变量映射 + let env_map = build_env_map(); + + // 9. 系统提示词 + let system_prompt = + std::env::var("WEIXIN_LLM_SYSTEM_PROMPT").unwrap_or_else(|_| String::new()); + let model = + std::env::var("DEEPSEEK_MODEL").unwrap_or_else(|_| "deepseek-v4-flash".to_string()); + + // 10. 共享数据 + let shared = Arc::new(DaemonShared { + db: database.clone(), + client: client.clone(), + account_id: account_id.clone(), + approval: approval_manager.clone(), + memory_store: memory_store.clone(), + sock_path: sock_path.clone(), + tools_list: tools_list.clone(), + env_map, + system_prompt, + model, + approval_ctx: Mutex::new(HashMap::new()), + }); + + // 11. 审批过期清理 + let approval_clean = approval_manager.clone(); + let shared_for_clean = shared.clone(); + tokio::spawn(async move { + loop { + tokio::time::sleep(std::time::Duration::from_secs(60)).await; + approval_clean.clean_expired().await; + // 清理过期的审批上下文(超过 5 分钟) + let now = Instant::now(); + shared_for_clean.approval_ctx.lock().await + .retain(|_, (_, _, _, created)| now.duration_since(*created).as_secs() < 300); + } + }); + + // 12. 调度器 + if let Some(sched_db) = database.as_ref().map(|db| db.pool().clone()) { + let sched_client = client.clone(); + tokio::spawn(async move { + let scheduler = crate::scheduler::Scheduler::new(Some(Arc::new(sched_db))); + scheduler + .run(move |to: &str, _name: &str, msg: &str| { + let c = sched_client.clone(); + let uid = to.to_string(); + let text = msg.to_string(); + tokio::task::spawn(async move { + if let Err(e) = c.send_text(&uid, &text, None).await { + tracing::error!("发送调度结果失败: {}", e); + } + }); + Ok(()) + }) + .await; + }); + info!("定时任务调度器已启动"); + } + + info!("开始监听消息..."); + + // 13. 主循环 + loop { + tokio::select! { + // WeChat 消息 + msg_result = client.receive_messages() => { + match msg_result { + Ok(resp) => { + // 持久化 runtime buf + if !resp.get_updates_buf.is_empty() { + file_state.save_runtime(&resp.get_updates_buf); + } + + for msg in &resp.msgs { + if msg.msg_type != Some(crate::wechat::types::WeixinMessage::TYPE_USER) { + continue; + } + let from = msg.from_user_id.as_deref().unwrap_or("unknown"); + let text = msg.text_content().unwrap_or("(非文本消息)"); + let ctx_token = msg.context_token.as_deref(); + let msg_id = msg.message_id.map(|id| id.to_string()).unwrap_or_default(); + + info!("收到消息 from={}: {}", from, text); + + // 入库:收到的消息 + if let Some(db) = &shared.db { + if let Err(e) = crate::db::models::insert_chat_record( + db.pool(), + "inbound", + from, + &shared.account_id, + text, + "wechat", + ctx_token, + &msg_id, + ) + .await + { + error!("保存聊天记录失败: {}", e); + } + } + + // 审批回复检查 + if let Some((skill_name, decision)) = shared.approval.handle_reply(from, text).await { + match decision { + ApprovalDecision::Approved => { + info!("审批通过: {}", skill_name); + let original = shared.approval_ctx.lock().await.remove(from); + if let Some((orig_text, orig_ctx, _tool, _ts)) = original { + let mut q = queue.lock().await; + q.enqueue(from, PendingMessage { + from: from.to_string(), + text: orig_text, + account_id: shared.account_id.clone(), + context_token: orig_ctx, + message_id: String::new(), + approved_tool: Some(skill_name), + }); + drop(q); + spawn_worker_for_user(from, &queue, &shared).await; + } + } + ApprovalDecision::Rejected => { + info!("审批被拒绝: {}", skill_name); + shared.approval_ctx.lock().await.remove(from); + } + ApprovalDecision::Expired => { + info!("审批已过期: {}", skill_name); + shared.approval_ctx.lock().await.remove(from); + } + } + continue; + } + + // 消息入队 + { + let mut q = queue.lock().await; + q.enqueue(from, PendingMessage { + from: from.to_string(), + text: text.to_string(), + account_id: shared.account_id.clone(), + context_token: ctx_token.map(String::from), + message_id: msg_id, + approved_tool: None, + }); + drop(q); + } + + // 尝试 spawn worker + spawn_worker_for_user(from, &queue, &shared).await; + } + } + Err(e) => { + if !e.contains("超时") && !e.contains("timeout") { + error!("接收消息失败: {}", e); + tokio::time::sleep(std::time::Duration::from_secs(2)).await; + } + } + } + } + + // Worker 连接 + accept_result = listener.accept() => { + match accept_result { + Ok((stream, _)) => { + let q = queue.clone(); + let shr = shared.clone(); + tokio::spawn(async move { + handle_worker(stream, q, shr).await; + }); + } + Err(e) => { + error!("接受 Worker 连接失败: {}", e); + } + } + } + } + } +} + +/// 共享数据 +struct DaemonShared { + db: Option>, + client: WeChatClient, + account_id: String, + approval: Arc, + memory_store: Arc, + sock_path: String, + tools_list: Vec, + env_map: HashMap, + system_prompt: String, + model: String, + /// 待审批消息: user_id → (原始消息文本, context_token, 审批工具名, 创建时间) + approval_ctx: Mutex, String, Instant)>>, +} + +/// 处理一个 Worker 连接 +async fn handle_worker( + mut stream: UnixStream, + queue: Arc>, + shared: Arc, +) { + // 1. 取出等待的用户 + let user_id = { + let mut q = queue.lock().await; + q.pop_waiting() + }; + + let user_id = match user_id { + Some(uid) => uid, + None => { + warn!("Worker 连接但没有等待的用户"); + return; + } + }; + + // 2. 取出该用户的一条消息 + let pending_msg = { + let mut q = queue.lock().await; + q.pop_message(&user_id) + }; + + let pending_msg = match pending_msg { + Some(m) => m, + None => { + warn!("Worker 连接但用户 {} 无待处理消息", user_id); + queue.lock().await.deactivate(&user_id); + return; + } + }; + + info!( + "Worker 处理: user={} text={:.60}", + user_id, pending_msg.text + ); + + // 3. 加载用户记忆 + shared.memory_store.load(&user_id).await; + + // 4. 预载历史消息 + let history = load_history(&shared.db, &user_id).await; + let memories = load_memories(&shared.memory_store, &user_id).await; + let summaries = load_summaries(&shared.db, &user_id).await; + + // 5. 构建 TaskFrame + let approved_tool_info = pending_msg.approved_tool.clone(); + let task = TaskFrame { + user_id: user_id.clone(), + msg: TaskMessage { + from: pending_msg.from, + text: pending_msg.text.clone(), + account_id: pending_msg.account_id, + context_token: pending_msg.context_token.clone(), + }, + history, + memories, + summaries, + approved_tool: approved_tool_info.clone(), + env: shared.env_map.clone(), + tools: shared.tools_list.clone(), + system_prompt: shared.system_prompt.clone(), + model: shared.model.clone(), + }; + + // 6. 发送 task 帧 + if let Err(e) = send_frame(&mut stream, &task).await { + error!("发送 task 帧失败: {}", e); + queue.lock().await.deactivate(&user_id); + return; + } + + // 7. 逐帧处理 worker 输出 + let mut reply_text = String::new(); + let mut chunk_buffer = String::new(); + let ctx_token = task.msg.context_token.clone(); + + loop { + match recv_output(&mut stream).await { + Ok(frame) => match frame { + OutputFrame::Chunk { text } => { + chunk_buffer.push_str(&text); + } + OutputFrame::Reply { text } => { + reply_text = text; + } + OutputFrame::RequestApproval { tool, reason } => { + info!("Worker 请求审批: tool={} reason={}", tool, reason); + // 创建审批记录 + match shared.approval.create(&user_id, &tool).await { + Ok((code, _rx)) => { + let msg = format!( + "⚠️ 操作确认\n\n技能:{}\n\n确认码:{}\n\n回复确认码以继续\n回复 0 或「取消」以取消\n(5分钟内有效)", + tool, code + ); + // 保存原始消息上下文以便审批后重新处理 + shared.approval_ctx.lock().await.insert( + user_id.clone(), + (task.msg.text.clone(), ctx_token.clone(), tool.clone(), Instant::now()), + ); + // 发送确认消息到微信 + if let Err(e) = shared.client.send_text(&user_id, &msg, ctx_token.as_deref()).await { + error!("发送审批消息失败: {}", e); + } + } + Err(e) => { + error!("创建审批失败: {}", e); + let _ = shared.client.send_text(&user_id, &format!("审批创建失败: {}", e), ctx_token.as_deref()).await; + } + } + } + OutputFrame::NeedApproval { tool, .. } => { + warn!("收到已废弃的 NeedApproval 帧 (tool={})", tool); + } + OutputFrame::DbWrite { table, row } => { + handle_db_write(&shared.db, &shared.memory_store, &table, &row, &user_id, &shared.account_id).await; + } + OutputFrame::UsageReport { + model, + provider, + prompt_tokens, + completion_tokens, + total_tokens: _, + cache_hit_tokens, + cache_miss_tokens, + user_id: uid, + } => { + if let Some(db) = &shared.db { + if let Err(e) = crate::db::models::insert_llm_usage( + db.pool(), + &uid, + &model, + &provider, + prompt_tokens, + completion_tokens, + cache_hit_tokens, + cache_miss_tokens, + ) + .await + { + error!("存储 LLM 用量失败: {}", e); + } + } + } + OutputFrame::Bye => break, + }, + Err(e) => { + error!("读取 Worker 输出帧失败: {}", e); + break; + } + } + } + + // 8. 发送最终回复 + if !reply_text.is_empty() { + match shared + .client + .send_text(&user_id, &reply_text, ctx_token.as_deref()) + .await + { + Ok(sent_id) => { + info!("回复成功 user={} msg_id={}", user_id, sent_id); + // 入库:发送的回复 + if let Some(db) = &shared.db { + if let Err(e) = crate::db::models::insert_chat_record( + db.pool(), + "outbound", + &user_id, + &shared.account_id, + &reply_text, + "llm", + ctx_token.as_deref(), + &sent_id, + ) + .await + { + error!("保存聊天记录失败: {}", e); + } + } + } + Err(e) => error!("发送回复失败: {}", e), + } + } else if !chunk_buffer.is_empty() { + // 有 chunk 但没有 reply — 用 chunk 作为回复 + match shared + .client + .send_text(&user_id, &chunk_buffer, ctx_token.as_deref()) + .await + { + Ok(sent_id) => { + info!("回复成功 (chunk) user={} msg_id={}", user_id, sent_id); + if let Some(db) = &shared.db { + if let Err(e) = crate::db::models::insert_chat_record( + db.pool(), + "outbound", + &user_id, + &shared.account_id, + &chunk_buffer, + "llm", + ctx_token.as_deref(), + &sent_id, + ) + .await + { + error!("保存聊天记录失败: {}", e); + } + } + } + Err(e) => error!("发送回复失败: {}", e), + } + } + + // 9. 清理并继续处理该用户的下一消息 + { + let mut q = queue.lock().await; + q.deactivate(&user_id); + let has_more = q.has_pending(&user_id); + drop(q); + if has_more { + spawn_worker_for_user(&user_id, &queue, &shared).await; + } + } +} + +/// 为用户 spawn worker 进程 +async fn spawn_worker_for_user( + user_id: &str, + queue: &Arc>, + shared: &Arc, +) { + // 检查是否已有活跃 worker + { + let q = queue.lock().await; + if *q.active.get(user_id).unwrap_or(&false) { + return; // 已有 worker 在处理 + } + } + + let sock = shared.sock_path.clone(); + let exe = match std::env::current_exe() { + Ok(p) => p, + Err(e) => { + error!("获取当前可执行文件路径失败: {}", e); + return; + } + }; + + info!("Spawning worker for user={}", user_id); + + match Command::new(&exe) + .arg("worker") + .arg("--sock") + .arg(&sock) + .spawn() + { + Ok(mut child) => { + // 后台等待 worker 退出 + tokio::spawn(async move { + match child.wait().await { + Ok(status) => { + if !status.success() { + warn!("Worker 退出码: {:?}", status.code()); + } + } + Err(e) => error!("等待 Worker 退出失败: {}", e), + } + }); + } + Err(e) => { + error!("Spawn worker 失败: {}", e); + queue.lock().await.deactivate(user_id); + } + } +} + +// ─── 辅助函数 ─── + +async fn load_auth( + database: &Option>, + file_state: &StateManager, +) -> Option { + if let Some(db) = database { + if let Some(a) = crate::db::models::load_auth(db.pool()).await { + return Some(a); + } + // 尝试从文件迁移 + if let Some(fa) = file_state.load_auth() { + let a = crate::state::AuthState { + token: fa.token, + account_id: fa.account_id, + base_url: fa.base_url, + }; + if let Err(e) = crate::db::models::save_auth(db.pool(), &a).await { + error!("保存认证信息失败: {}", e); + } + return Some(a); + } + } else if let Some(a) = file_state.load_auth() { + return Some(a); + } + None +} + +async fn load_history(db: &Option>, user_id: &str) -> Vec { + let Some(db) = db else { + return vec![]; + }; + match crate::db::models::list_recent_chat_records(db.pool(), user_id, 20).await { + Ok(records) => records + .into_iter() + .rev() // 时间顺序 + .map(|r| HistoryEntry { + role: if r.direction == "inbound" { + "user".into() + } else { + "assistant".into() + }, + content: r.text, + }) + .collect(), + Err(e) => { + warn!("加载历史消息失败: {}", e); + vec![] + } + } +} + +async fn load_memories(memory_store: &MemoryStore, user_id: &str) -> Vec { + let text = memory_store.read_for(user_id).await; + if text == "暂无长期记忆" { + vec![] + } else { + text.lines().map(|l| l.trim().to_string()).filter(|l| !l.is_empty()).collect() + } +} + +async fn load_summaries(db: &Option>, user_id: &str) -> Vec { + let Some(db) = db else { + return vec![]; + }; + match crate::db::models::load_summaries(db.pool(), user_id, 5).await { + Ok(s) => s, + Err(e) => { + warn!("加载摘要失败: {}", e); + vec![] + } + } +} + +async fn handle_db_write( + db: &Option>, + memory_store: &MemoryStore, + table: &str, + row: &serde_json::Value, + user_id: &str, + account_id: &str, +) { + // user_memories 不依赖 DB(memory_store.write_for 自己处理持久化) + if table == "user_memories" { + let content = row.get("content").and_then(|v| v.as_str()).unwrap_or(""); + if !content.is_empty() { + memory_store.write_for(user_id, content).await; + } + return; + } + + let Some(db) = db else { + return; + }; + match table { + "chat_records" => { + let direction = row.get("direction").and_then(|v| v.as_str()).unwrap_or("outbound"); + let text = row.get("text").and_then(|v| v.as_str()).unwrap_or(""); + let source = row.get("source").and_then(|v| v.as_str()).unwrap_or("llm"); + let ctx_token = row.get("context_token").and_then(|v| v.as_str()); + let msg_id = row.get("message_id").and_then(|v| v.as_str()).unwrap_or(""); + if let Err(e) = crate::db::models::insert_chat_record( + db.pool(), + direction, + user_id, + account_id, + text, + source, + ctx_token, + msg_id, + ) + .await + { + error!("DbWrite chat_records 失败: {}", e); + } + } + "llm_usage" => { + let model = row.get("model").and_then(|v| v.as_str()).unwrap_or(""); + let provider = row.get("provider").and_then(|v| v.as_str()).unwrap_or(""); + let prompt_tokens = row.get("prompt_tokens").and_then(|v| v.as_u64()).unwrap_or(0) as u32; + let completion_tokens = row.get("completion_tokens").and_then(|v| v.as_u64()).unwrap_or(0) as u32; + let cache_hit = row.get("cache_hit_tokens").and_then(|v| v.as_u64()).unwrap_or(0) as u32; + let cache_miss = row.get("cache_miss_tokens").and_then(|v| v.as_u64()).unwrap_or(0) as u32; + if let Err(e) = crate::db::models::insert_llm_usage( + db.pool(), + user_id, + model, + provider, + prompt_tokens, + completion_tokens, + cache_hit, + cache_miss, + ) + .await + { + error!("DbWrite llm_usage 失败: {}", e); + } + } + _ => warn!("未知 DbWrite table: {}", table), + } +} + +fn build_tools_list() -> Vec { + let mut list = Vec::new(); + + // query_capabilities + list.push(serde_json::json!({ + "type": "function", + "function": { + "name": "query_capabilities", + "description": "列出所有可用工具的详细说明、参数格式和调用示例。在需要了解有哪些工具可用时首先调用此工具。", + "parameters": { "type": "object", "properties": {} } + } + })); + + // call_capability + list.push(serde_json::json!({ + "type": "function", + "function": { + "name": "call_capability", + "description": "调用一个已注册的工具。name 为工具名(来自 query_capabilities),工具所需的具体参数应直接放在 JSON 中。", + "parameters": { + "type": "object", + "properties": { + "name": { "type": "string", "description": "工具名称" }, + "prompt": { "type": "string", "description": "json格式参数" } + }, + "required": ["name"] + } + } + })); + + // 上下文工具 + list.push(serde_json::json!({ + "type": "function", + "function": { + "name": "read_memories", + "description": "读取用户的长期记忆(偏好、个人信息、约定)", + "parameters": { "type": "object", "properties": {} } + } + })); + list.push(serde_json::json!({ + "type": "function", + "function": { + "name": "write_memory", + "description": "记录用户的重要信息或偏好", + "parameters": { + "type": "object", + "properties": { + "content": { "type": "string", "description": "记忆内容" } + }, + "required": ["content"] + } + } + })); + list.push(serde_json::json!({ + "type": "function", + "function": { + "name": "read_summaries", + "description": "读取历史会话摘要", + "parameters": { "type": "object", "properties": {} } + } + })); + + // 业务工具(与 BuiltinRegistry 保持一致) + list.push(serde_json::json!({ + "type": "function", + "function": { + "name": "get_current_datetime", + "description": "获取当前日期时间(北京时间 UTC+8)", + "parameters": { "type": "object", "properties": {"format": {"type": "string"}} } + } + })); + list.push(serde_json::json!({ + "type": "function", + "function": { + "name": "manage_memos", + "description": "管理备忘录:添加/列出/删除", + "parameters": { + "type": "object", + "properties": { + "action": {"type": "string", "enum": ["add", "list", "delete"]}, + "content": {"type": "string"}, + "id": {"type": "integer"} + } + } + } + })); + list.push(serde_json::json!({ + "type": "function", + "function": { + "name": "query_weather", + "description": "查询指定城市的天气信息", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string", "description": "城市名称"}, + "days": {"type": "integer", "description": "查询天数(1-7)"} + }, + "required": ["location"] + } + } + })); + list.push(serde_json::json!({ + "type": "function", + "function": { + "name": "web_search", + "description": "联网搜索最新信息。通过 Tavily API 搜索互联网,获取实时、准确的结果。适用于需要最新资讯、事实查询的场景。", + "parameters": { + "type": "object", + "properties": { + "query": {"type": "string", "description": "搜索关键词"}, + "max_results": {"type": "integer", "description": "最大结果数(1-20,默认5)"}, + "include_answer": {"type": "boolean", "description": "是否包含 AI 生成的摘要回答"}, + "search_depth": {"type": "string", "enum": ["basic", "advanced", "fast", "ultra-fast"], "description": "搜索深度"}, + "topic": {"type": "string", "enum": ["general", "news", "finance"], "description": "搜索主题"}, + "time_range": {"type": "string", "enum": ["day", "week", "month", "year"], "description": "按发布时间过滤"}, + "start_date": {"type": "string", "description": "起始日期 YYYY-MM-DD"}, + "end_date": {"type": "string", "description": "结束日期 YYYY-MM-DD"}, + "include_domains": {"type": "array", "items": {"type": "string"}, "description": "限定域名列表"}, + "exclude_domains": {"type": "array", "items": {"type": "string"}, "description": "排除域名列表"}, + "country": {"type": "string", "description": "优先指定国家"} + }, + "required": ["query"] + } + } + })); + list.push(serde_json::json!({ + "type": "function", + "function": { + "name": "fetch_page", + "description": "抓取网页并提取正文内容。传入 URL 即可阅读文章正文。可通过 ~/.ias/site_selectors.json 配置按站点用 CSS 选择器精确定位内容区域。", + "parameters": { + "type": "object", + "properties": { + "url": {"type": "string", "description": "网页 URL"} + }, + "required": ["url"] + } + } + })); + + list +} + +fn build_env_map() -> HashMap { + let mut map = HashMap::new(); + for var in &[ + "DEEPSEEK_API_KEY", + "DEEPSEEK_MODEL", + "DEEPSEEK_BASE_URL", + "QWEATHER_API_HOST", + "QWEATHER_JWT_KEY_ID", + "QWEATHER_JWT_PROJECT_ID", + "QWEATHER_JWT_PRIVATE_KEY_FILE", + "TAVILY_API_KEY", + ] { + if let Ok(val) = std::env::var(var) { + map.insert(var.to_string(), val); + } + } + map +} diff --git a/src/db/models.rs b/src/db/models.rs index 896aa26..fa1c998 100644 --- a/src/db/models.rs +++ b/src/db/models.rs @@ -219,11 +219,12 @@ pub async fn save_summary( Ok(()) } -/// 加载最近的摘要(用于 read_summaries 工具) -pub async fn load_summaries(pool: &PgPool, limit: i64) -> Result, String> { +/// 加载最近的摘要(用于 read_summaries 工具),按 user_id 过滤 +pub async fn load_summaries(pool: &PgPool, user_id: &str, limit: i64) -> Result, String> { let rows = sqlx::query_as::<_, (String,)>( - "SELECT summary_text FROM session_summaries ORDER BY created_at DESC LIMIT $1", + "SELECT summary_text FROM session_summaries WHERE user_id = $1 ORDER BY created_at DESC LIMIT $2", ) + .bind(user_id) .bind(limit) .fetch_all(pool) .await diff --git a/src/ipc.rs b/src/ipc.rs new file mode 100644 index 0000000..492bc28 --- /dev/null +++ b/src/ipc.rs @@ -0,0 +1,198 @@ +//! IPC 帧格式 — Unix Domain Socket 长度前缀帧 +//! +//! 帧格式: [4 字节 u32 BE: payload_len][N 字节: JSON payload] +//! +//! 消息类型定义: +//! Daemon → Worker: Task (携带用户消息、历史、记忆、摘要、环境变量) +//! Worker → Daemon: Chunk / Reply / NeedApproval / DbWrite / UsageReport / Bye + +use serde::{Deserialize, Serialize}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::UnixStream; + +// ─── 帧消息类型 ─── + +/// Daemon → Worker: 任务描述 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TaskFrame { + pub user_id: String, + pub msg: TaskMessage, + pub history: Vec, + pub memories: Vec, + pub summaries: Vec, + #[serde(default)] + pub approved_tool: Option, + pub env: std::collections::HashMap, + pub tools: Vec, + pub system_prompt: String, + pub model: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TaskMessage { + pub from: String, + pub text: String, + pub account_id: String, + #[serde(default)] + pub context_token: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct HistoryEntry { + pub role: String, + pub content: String, +} + +/// Worker → Daemon: 输出帧 +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(tag = "type")] +pub enum OutputFrame { + /// 流式文本块 + #[serde(rename = "chunk")] + Chunk { text: String }, + /// 最终回复(完整文本) + #[serde(rename = "reply")] + Reply { text: String }, + /// 需要审批 + #[serde(rename = "need_approval")] + NeedApproval { + tool: String, + code: String, + message: String, + }, + /// Worker 请求 daemon 创建审批(daemon 生成确认码) + #[serde(rename = "request_approval")] + RequestApproval { + tool: String, + reason: String, + }, + /// 数据库写入 + #[serde(rename = "db_write")] + DbWrite { + table: String, + row: serde_json::Value, + }, + /// Token 用量报告 + #[serde(rename = "usage")] + UsageReport { + model: String, + provider: String, + prompt_tokens: u32, + completion_tokens: u32, + total_tokens: u32, + cache_hit_tokens: u32, + cache_miss_tokens: u32, + user_id: String, + }, + /// Worker 已完成,即将退出 + #[serde(rename = "bye")] + Bye, +} + +// ─── 帧读写 ─── + +/// 发送一个 JSON 帧(长度前缀 + JSON) +pub async fn send_frame(stream: &mut UnixStream, msg: &impl Serialize) -> Result<(), String> { + let payload = serde_json::to_vec(msg).map_err(|e| format!("序列化帧失败: {}", e))?; + let len = payload.len() as u32; + let mut header = len.to_be_bytes().to_vec(); + header.extend_from_slice(&payload); + + stream + .write_all(&header) + .await + .map_err(|e| format!("发送帧失败: {}", e))?; + Ok(()) +} + +/// 最大帧大小:10MB(防止 OOM) +const MAX_FRAME_SIZE: usize = 10 * 1024 * 1024; + +/// 接收一个 JSON 帧 +pub async fn recv_frame(stream: &mut UnixStream) -> Result { + // 读 4 字节长度头 + let mut header = [0u8; 4]; + stream + .read_exact(&mut header) + .await + .map_err(|e| format!("读取帧头失败: {}", e))?; + let len = u32::from_be_bytes(header) as usize; + + // 防御:拒绝超大帧 + if len > MAX_FRAME_SIZE { + return Err(format!("帧过大: {} bytes (max {})", len, MAX_FRAME_SIZE)); + } + + // 读 payload + let mut payload = vec![0u8; len]; + stream + .read_exact(&mut payload) + .await + .map_err(|e| format!("读取帧体失败 (len={}): {}", len, e))?; + + serde_json::from_slice(&payload).map_err(|e| { + format!( + "解析帧失败: {} — payload({}): {:.200}", + e, + len, + String::from_utf8_lossy(&payload[..std::cmp::min(len, 200)]) + ) + }) +} + +// ─── 便捷方法: 按类型收发 ─── + +/// 接收并解析为 TaskFrame +pub async fn recv_task(stream: &mut UnixStream) -> Result { + let v = recv_frame(stream).await?; + serde_json::from_value(v).map_err(|e| format!("解析 TaskFrame 失败: {}", e)) +} + +/// 接收并解析为 OutputFrame +pub async fn recv_output(stream: &mut UnixStream) -> Result { + let v = recv_frame(stream).await?; + serde_json::from_value(v).map_err(|e| format!("解析 OutputFrame 失败: {}", e)) +} + +// ─── 测试 ─── + +#[cfg(test)] +mod tests { + use super::*; + use tokio::net::UnixListener; + + #[tokio::test] + async fn test_frame_roundtrip() { + let sock_path = "/tmp/ias_ipc_test.sock"; + let _ = std::fs::remove_file(sock_path); + + let listener = UnixListener::bind(sock_path).unwrap(); + + let server = tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + let v = recv_frame(&mut stream).await.unwrap(); + assert_eq!(v["type"], "chunk"); + assert_eq!(v["text"], "hello"); + + let output = OutputFrame::Reply { + text: "world".into(), + }; + send_frame(&mut stream, &output).await.unwrap(); + }); + + let mut client = UnixStream::connect(sock_path).await.unwrap(); + let chunk = OutputFrame::Chunk { + text: "hello".into(), + }; + send_frame(&mut client, &chunk).await.unwrap(); + + let resp = recv_output(&mut client).await.unwrap(); + match resp { + OutputFrame::Reply { text } => assert_eq!(text, "world"), + _ => panic!("expected reply"), + } + + server.await.unwrap(); + let _ = std::fs::remove_file(sock_path); + } +} diff --git a/src/llm/conversation.rs b/src/llm/conversation.rs index fa7efa8..618dc85 100644 --- a/src/llm/conversation.rs +++ b/src/llm/conversation.rs @@ -288,6 +288,12 @@ pub const DEFAULT_SYSTEM_PROMPT: &str = "You are a concise and helpful assistant Keep replies practical and natural. \ Unless the user asks for detail, keep most replies under 120 Chinese characters. \ \ - 可用的核心工具:query_capabilities(查看所有可用工具)、call_capability(按名调用工具)。\ + 可用的核心工具:query_capabilities(查看所有可用工具及优先级)、call_capability(按名调用工具)。\ 可使用 read_memories / write_memory 管理用户长期记忆,read_summaries 查看历史摘要。\ - 调用工具前先用 query_capabilities 确认工具有哪些。"; + \ + 工具选择优先级(重要):\ + 1. 🗺️ 旅游/路线/地点查询 → amap_* 系列(amap_poi_search, amap_route_plan, amap_travel_plan 等)\ + 2. 🌤️ 天气查询 → query_weather\ + 3. 🔍 以上都不匹配 → web_search(Tavily 联网搜索)\ + \ + 调用工具前先用 query_capabilities 确认有哪些工具及其参数格式。"; diff --git a/src/main.rs b/src/main.rs index 9305f0c..f0f1ba0 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,15 +1,18 @@ mod cli; mod context; +mod daemon; mod db; +mod ipc; mod llm; mod logger; mod scheduler; mod state; mod tools; mod wechat; +mod worker; use clap::Parser; -use cli::{Cli, Commands}; +use cli::{AmapAction, Cli, Commands, MemosAction, ToolCommand}; use context::MemoryStore; use db::Database; use llm::{Conversation, ConversationConfig, DEFAULT_SYSTEM_PROMPT, ToolExecutor, Usage}; @@ -42,6 +45,17 @@ async fn main() { } let file_state = state::StateManager::new(".data/weixin-ilink"); + + // 工具命令无需登录/数据库,提前处理 + let tool_cmd = match &cli.command { + Commands::Tool(cmd) => Some(cmd), + _ => None, + }; + if let Some(cmd) = tool_cmd { + cmd_tool(cmd.clone()).await; + return; + } + let database = match Database::connect().await { Ok(db) => { info!("✅ 数据库连接成功"); @@ -87,11 +101,29 @@ async fn main() { cmd_usage(&database, since, until, model).await; } Commands::Service => cmd_listen(true, false, &database, &file_state, &memory_store).await, + Commands::Daemon { sock } => { + cmd_daemon(sock, &database, &file_state, &memory_store).await; + } + Commands::Worker { sock } => { + if let Err(e) = worker::run(&sock).await { + error!("Worker 失败: {}", e); + } + } + Commands::Tool(..) => unreachable!("Tool 命令已提前处理"), } } // ─── 命令实现 ─── +async fn cmd_daemon( + sock_path: String, + database: &Option>, + file_state: &state::StateManager, + memory_store: &Arc, +) { + daemon::run(sock_path, database, file_state, memory_store).await; +} + async fn cmd_login( timeout_secs: u64, database: &Option>, @@ -339,10 +371,11 @@ async fn cmd_listen( // query_capabilities: 列出所有内置工具 if n == "query_capabilities" { let specs = tools::builtin::BuiltinRegistry::specs(); - return Ok(build_capability_list(&specs)); + return Ok(tools::build_capability_guide(&specs)); } - // 确定目标工具名和参数(call_capability 提取 name + 透传完整参数) + // 确定目标工具名和参数 + // call_capability: 从 {name, prompt} 中提取 name,从 prompt 中解包真实参数 let (target_name, target_args) = if n == "call_capability" { let cp: serde_json::Value = serde_json::from_str(&args).unwrap_or_default(); let name = cp @@ -350,7 +383,8 @@ async fn cmd_listen( .and_then(|v| v.as_str()) .map(|s| s.to_string()) .unwrap_or_default(); - (name, serde_json::to_string(&cp).unwrap_or_default()) + let unpacked = tools::unpack_call_params(&cp); + (name, serde_json::to_string(&unpacked).unwrap_or_default()) } else { (n.clone(), args.clone()) }; @@ -841,30 +875,143 @@ async fn builtin_approve(ctx: &ExecutionContext, name: &str) -> Result String { - let mut lines = vec!["📋 可用工具:".to_string()]; - for spec in specs { - let risk = if spec.risk_level == tools::types::RiskLevel::High { - " ⚠️需确认" - } else { - "" - }; - lines.push(format!( - "\n🔹 {} {} — {}", - spec.name, risk, spec.description - )); - if !spec - .parameters - .get("properties") - .and_then(|v| v.as_object()) - .map_or(true, |o| o.is_empty()) - { - lines.push(format!( - " 参数: {}", - serde_json::to_string(&spec.parameters).unwrap_or_default() - )); + +// ─── 工具 CLI 命令 ─── + +async fn cmd_tool(cmd: ToolCommand) { + match cmd { + ToolCommand::Datetime => { + let result = tools::builtins::datetime::execute(); + println!("{}", result.output); + } + ToolCommand::Memos(action) => { + let args = match action { + MemosAction::List => serde_json::json!({"action": "list"}), + MemosAction::Add { content } => { + serde_json::json!({"action": "add", "content": content}) + } + MemosAction::Delete { id } => { + serde_json::json!({"action": "delete", "id": id}) + } + }; + let args_str = serde_json::to_string(&args).unwrap_or_default(); + let result = tools::builtins::memos::execute(&args_str).await; + println!("{}", result.output); + } + ToolCommand::Weather { location, days } => { + let params = serde_json::json!({"location": location, "days": days}); + let result = tools::builtins::weather::execute(params).await; + println!("{}", result.output); + } + ToolCommand::Search { + query, + max_results, + no_answer, + topic, + time_range, + } => { + let mut params = serde_json::json!({ + "query": query, + "max_results": max_results, + }); + if no_answer { + params["include_answer"] = serde_json::json!(false); + } + if let Some(t) = topic { + params["topic"] = serde_json::json!(t); + } + if let Some(tr) = time_range { + params["time_range"] = serde_json::json!(tr); + } + let result = tools::builtins::web_search::execute(params).await; + println!("{}", result.output); + } + ToolCommand::Fetch { url } => { + let params = serde_json::json!({"url": url}); + let result = tools::builtins::fetch_page::execute(params).await; + println!("{}", result.output); + } + ToolCommand::Amap(action) => { + cmd_amap(action).await; + } + } +} + +async fn cmd_amap(action: AmapAction) { + match action { + AmapAction::PoiSearch { + keywords, + city, + location, + radius, + page, + offset, + } => { + let mut params = serde_json::json!({ + "keywords": keywords, + "page": page, + "offset": offset, + }); + if let Some(c) = city { params["city"] = serde_json::json!(c); } + if let Some(l) = location { params["location"] = serde_json::json!(l); } + if let Some(r) = radius { params["radius"] = serde_json::json!(r); } + let result = tools::builtins::amap::execute("amap_poi_search", params).await; + if let Some(r) = result { println!("{}", r.output); } + } + AmapAction::Geocode { address, city } => { + let mut params = serde_json::json!({"address": address}); + if let Some(c) = city { params["city"] = serde_json::json!(c); } + let result = tools::builtins::amap::execute("amap_geocode", params).await; + if let Some(r) = result { println!("{}", r.output); } + } + AmapAction::ReverseGeocode { location } => { + let params = serde_json::json!({"location": location}); + let result = tools::builtins::amap::execute("amap_reverse_geocode", params).await; + if let Some(r) = result { println!("{}", r.output); } + } + AmapAction::RoutePlan { + r#type, + origin, + destination, + city, + waypoints, + strategy, + } => { + let mut params = serde_json::json!({ + "type": r#type, + "origin": origin, + "destination": destination, + }); + if let Some(c) = city { params["city"] = serde_json::json!(c); } + if let Some(w) = waypoints { params["waypoints"] = serde_json::json!(w); } + if let Some(s) = strategy { params["strategy"] = serde_json::json!(s); } + let result = tools::builtins::amap::execute("amap_route_plan", params).await; + if let Some(r) = result { println!("{}", r.output); } + } + AmapAction::TravelPlan { + city, + interests, + route_type, + } => { + let interests_arr: Vec = interests + .split(',') + .map(|s| serde_json::json!(s.trim())) + .collect(); + let params = serde_json::json!({ + "city": city, + "interests": interests_arr, + "route_type": route_type, + }); + let result = tools::builtins::amap::execute("amap_travel_plan", params).await; + if let Some(r) = result { println!("{}", r.output); } + } + AmapAction::MapLink { data } => { + let map_data: serde_json::Value = + serde_json::from_str(&data).unwrap_or(serde_json::json!([])); + let params = serde_json::json!({"map_data": map_data}); + let result = tools::builtins::amap::execute("amap_map_link", params).await; + if let Some(r) = result { println!("{}", r.output); } } } - lines.join("\n") } diff --git a/src/tools/approval.rs b/src/tools/approval.rs index bdb5d85..4b49816 100644 --- a/src/tools/approval.rs +++ b/src/tools/approval.rs @@ -3,6 +3,7 @@ use sha2::{Digest, Sha256}; use sqlx::PgPool; use std::collections::HashMap; use std::sync::Arc; +use std::time::Instant; use tokio::sync::{Mutex, oneshot}; /// 审批决策 @@ -18,6 +19,7 @@ struct PendingEntry { code_hash: String, skill_name: String, attempts_left: i32, + expires_at: Instant, /// 通知等待方 sender: oneshot::Sender, } @@ -80,6 +82,7 @@ impl ApprovalManager { code_hash, skill_name: skill_name.to_string(), attempts_left: 3, + expires_at: Instant::now() + std::time::Duration::from_secs(300), sender: tx, }); @@ -87,8 +90,8 @@ impl ApprovalManager { } /// 处理用户回复(由消息循环调用) - /// 匹配成功时返回 Some(技能名),匹配失败返回 None - pub async fn handle_reply(&self, user_id: &str, reply: &str) -> Option { + /// 返回 (技能名, 决策);匹配失败返回 None + pub async fn handle_reply(&self, user_id: &str, reply: &str) -> Option<(String, ApprovalDecision)> { let (result, name, code_hash) = { let mut map = self.pending.lock().await; let entries = map.get_mut(user_id)?; @@ -141,17 +144,18 @@ impl ApprovalManager { .await; } - Some(name) + Some((name, result)) } - /// 清理过期审批(通知等待方 + 更新 DB) + /// 清理过期审批(按时间判定,通知等待方 + 更新 DB) pub async fn clean_expired(&self) { + let now = Instant::now(); let mut expired = Vec::new(); { let map = self.pending.lock().await; for (uid, entries) in map.iter() { for (i, e) in entries.iter().enumerate() { - if e.sender.is_closed() { + if now > e.expires_at { expired.push((uid.clone(), i, e.code_hash.clone())); } } diff --git a/src/tools/builtin.rs b/src/tools/builtin.rs index 04ff1ef..6ffb97d 100644 --- a/src/tools/builtin.rs +++ b/src/tools/builtin.rs @@ -1,17 +1,32 @@ use crate::tools::types::{RiskLevel, SkillResult}; -use chrono::Local; -/// 内置工具执行器:统一处理参数校验、执行、错误返回 +/// 内置工具注册表:名称 → 执行、spec 列表、风险判断 pub struct BuiltinRegistry; impl BuiltinRegistry { pub async fn execute(name: &str, args_json: &str) -> Option { match name { - "get_current_datetime" => Some(execute_datetime()), - "manage_memos" => Some(handle_memos(args_json).await), + "get_current_datetime" => Some(super::builtins::datetime::execute()), + "manage_memos" => Some(super::builtins::memos::execute(args_json).await), "query_weather" => { - let params: serde_json::Value = serde_json::from_str(args_json).unwrap_or_default(); - Some(super::weather::execute(params).await) + let params: serde_json::Value = + serde_json::from_str(args_json).unwrap_or_default(); + Some(super::builtins::weather::execute(params).await) + } + "web_search" => { + let params: serde_json::Value = + serde_json::from_str(args_json).unwrap_or_default(); + Some(super::builtins::web_search::execute(params).await) + } + "fetch_page" => { + let params: serde_json::Value = + serde_json::from_str(args_json).unwrap_or_default(); + Some(super::builtins::fetch_page::execute(params).await) + } + n if n.starts_with("amap_") => { + let params: serde_json::Value = + serde_json::from_str(args_json).unwrap_or_default(); + super::builtins::amap::execute(n, params).await } _ => None, } @@ -19,23 +34,15 @@ impl BuiltinRegistry { /// 返回所有内置工具的 spec 列表 pub fn specs() -> Vec { - vec![ - super::types::SkillSpec { - name: "get_current_datetime".into(), - description: "获取当前日期时间(北京时间 UTC+8)".into(), - risk_level: RiskLevel::Low, - parameters: serde_json::json!({"type":"object","properties":{"format":{"type":"string"}}}), - timeout_secs: 5, - }, - super::types::SkillSpec { - name: "manage_memos".into(), - description: "管理备忘录:添加/列出/删除。高风险操作(add/delete)需确认".into(), - risk_level: RiskLevel::Low, - parameters: serde_json::json!({"type":"object","properties":{"action":{"type":"string","enum":["add","list","delete"]},"content":{"type":"string"},"id":{"type":"integer"}}}), - timeout_secs: 10, - }, - super::weather::spec(), - ] + let mut specs = vec![ + super::builtins::datetime::spec(), + super::builtins::memos::spec(), + super::builtins::weather::spec(), + super::builtins::web_search::spec(), + super::builtins::fetch_page::spec(), + ]; + specs.extend(super::builtins::amap::specs()); + specs } pub fn is_high_risk(name: &str) -> bool { @@ -48,101 +55,3 @@ impl BuiltinRegistry { Self::specs().iter().any(|s| s.name == name) } } - -fn execute_datetime() -> SkillResult { - let now = Local::now().format("%Y-%m-%d %H:%M:%S").to_string(); - SkillResult::ok(serde_json::json!({"datetime": now}).to_string()) -} - -async fn handle_memos(args_json: &str) -> SkillResult { - let mut params: serde_json::Value = serde_json::from_str(args_json).unwrap_or_default(); - // call_capability 透传完整参数,action/content 在顶层 - // 如果顶层没有 action,尝试从 prompt 推断 - if params.get("action").is_none() { - let prompt = params.get("prompt").and_then(|v| v.as_str()).unwrap_or(""); - let action = if prompt.contains("添加") || prompt.contains("新增") || prompt.contains("add") - { - "add" - } else if prompt.contains("删除") || prompt.contains("移除") || prompt.contains("delete") - { - "delete" - } else { - "list" - }; - let content = params.get("prompt").and_then(|v| v.as_str()).unwrap_or(""); - let memo_id: i32 = content - .chars() - .filter(|c| c.is_ascii_digit()) - .collect::() - .parse() - .unwrap_or(0); - params = serde_json::json!({"action": action, "content": content, "id": memo_id}); - } - let action = params - .get("action") - .and_then(|v| v.as_str()) - .unwrap_or("list"); - let path = ".data/memos.json"; - if let Err(e) = std::fs::create_dir_all(".data") { - return SkillResult::error(format!("创建目录失败: {}", e)); - } - - let mut data: Vec = if std::path::Path::new(path).exists() { - serde_json::from_str(&std::fs::read_to_string(path).unwrap_or_default()).unwrap_or_default() - } else { - vec![] - }; - - match action { - "add" => { - let content = params.get("content").and_then(|v| v.as_str()).unwrap_or(""); - if content.is_empty() { - return SkillResult::error("请提供 content 参数"); - } - let next_id = data - .iter() - .filter_map(|m| m.get("id").and_then(|v| v.as_i64())) - .max() - .unwrap_or(0) - + 1; - data.push(serde_json::json!({ - "id": next_id, "content": content, - "created_at": Local::now().format("%Y-%m-%d %H:%M").to_string() - })); - if let Err(e) = std::fs::write(path, serde_json::to_string_pretty(&data).unwrap()) { - return SkillResult::error(format!("写入备忘录失败: {}", e)); - } - SkillResult::ok(format!("已添加备忘录 #{}", next_id)) - } - "list" => { - if data.is_empty() { - return SkillResult::ok("暂无备忘录".to_string()); - } - let lines: Vec = data - .iter() - .map(|m| { - format!( - "#{} [{}] {}", - m["id"], - m.get("created_at").and_then(|v| v.as_str()).unwrap_or("?"), - m["content"].as_str().unwrap_or("?") - ) - }) - .collect(); - SkillResult::ok(lines.join("\n")) - } - "delete" => { - let id = params.get("id").and_then(|v| v.as_i64()).unwrap_or(0); - let before = data.len(); - data.retain(|m| m.get("id").and_then(|v| v.as_i64()) != Some(id)); - if data.len() == before { - return SkillResult::error(format!("未找到备忘录 #{}", id)); - } - if let Err(e) = std::fs::write(path, serde_json::to_string_pretty(&data).unwrap()) { - return SkillResult::error(format!("写入备忘录失败: {}", e)); - } - SkillResult::ok(format!("已删除备忘录 #{}", id)) - } - _ => SkillResult::error("未知操作,支持: add, list, delete"), - } -} diff --git a/src/tools/builtins/amap.rs b/src/tools/builtins/amap.rs new file mode 100644 index 0000000..6798d5e --- /dev/null +++ b/src/tools/builtins/amap.rs @@ -0,0 +1,1116 @@ +// 高德地图综合工具(Rust 原生实现) +// +// API: +// POI 搜索: GET /v5/place/text → poi 列表 +// POI 周边搜索: GET /v5/place/around → poi 列表 +// 地理编码: GET /v3/geocode/geo → 坐标 +// 逆地理编码: GET /v3/geocode/regeo → 地址 +// 步行路线: GET /v3/direction/walking → 路线 +// 驾车路线: GET /v3/direction/driving → 路线 +// 骑行路线: GET /v4/direction/bicycling → 路线 +// 公交路线: GET /v3/direction/transit/integrated → 路线 +// +// 认证: AMAP_WEBSERVICE_KEY 或 AMAP_KEY 环境变量 +// 所有请求必须带 appname=amap-lbs-skill + +use reqwest::Client; +use std::time::Duration; + +use crate::tools::types::{RiskLevel, SkillResult, SkillSpec}; + +// ═══════════════════════════════════════════════ +// AmapClient — 共享的高德 API 客户端 +// ═══════════════════════════════════════════════ + +struct AmapClient { + http: Client, + key: String, +} + +impl AmapClient { + fn new() -> Result { + let key = std::env::var("AMAP_WEBSERVICE_KEY") + .or_else(|_| std::env::var("AMAP_KEY")) + .map_err(|_| { + "未设置高德 API Key,请设置环境变量 AMAP_WEBSERVICE_KEY 或 AMAP_KEY。\ + 获取 Key: https://lbs.amap.com/api/webservice/create-project-and-key" + .to_string() + })?; + + Ok(Self { + http: Client::builder() + .timeout(Duration::from_secs(15)) + .build() + .map_err(|e| format!("创建 HTTP client 失败: {}", e))?, + key, + }) + } + + /// POI 文本搜索 + async fn poi_text_search( + &self, + keywords: &str, + city: Option<&str>, + page: u32, + offset: u32, + ) -> Result { + let mut params = vec![ + ("key", self.key.as_str()), + ("keywords", keywords), + ("appname", "amap-lbs-skill"), + ]; + let region = city.unwrap_or(""); + params.push(("region", region)); + let page_str = page.to_string(); + params.push(("page", &page_str)); + let offset_str = offset.to_string(); + params.push(("offset", &offset_str)); + + let resp: serde_json::Value = self + .http + .get("https://restapi.amap.com/v5/place/text") + .query(¶ms) + .send() + .await + .map_err(|e| format!("POI 搜索请求失败: {}", e))? + .json() + .await + .map_err(|e| format!("解析 POI 搜索响应失败: {}", e))?; + + if resp["status"].as_str() != Some("1") { + return Err(format!( + "POI 搜索失败: {}", + resp["info"].as_str().unwrap_or("未知错误") + )); + } + Ok(resp) + } + + /// POI 周边搜索 + async fn poi_around( + &self, + keywords: &str, + location: &str, + radius: u32, + page: u32, + offset: u32, + ) -> Result { + let page_str = page.to_string(); + let offset_str = offset.to_string(); + let radius_str = radius.to_string(); + + let resp: serde_json::Value = self + .http + .get("https://restapi.amap.com/v5/place/around") + .query(&[ + ("key", self.key.as_str()), + ("keywords", keywords), + ("location", location), + ("radius", &radius_str), + ("page", &page_str), + ("offset", &offset_str), + ("appname", "amap-lbs-skill"), + ]) + .send() + .await + .map_err(|e| format!("周边搜索请求失败: {}", e))? + .json() + .await + .map_err(|e| format!("解析周边搜索响应失败: {}", e))?; + + if resp["status"].as_str() != Some("1") { + return Err(format!( + "周边搜索失败: {}", + resp["info"].as_str().unwrap_or("未知错误") + )); + } + Ok(resp) + } + + /// 地理编码:地址 → 坐标 + async fn geocode(&self, address: &str, city: Option<&str>) -> Result { + let mut params = vec![ + ("key", self.key.as_str()), + ("address", address), + ("output", "JSON"), + ("appname", "amap-lbs-skill"), + ]; + let city_val; + if let Some(c) = city { + city_val = c.to_string(); + params.push(("city", &city_val)); + } + + let resp: serde_json::Value = self + .http + .get("https://restapi.amap.com/v3/geocode/geo") + .query(¶ms) + .send() + .await + .map_err(|e| format!("地理编码请求失败: {}", e))? + .json() + .await + .map_err(|e| format!("解析地理编码响应失败: {}", e))?; + + if resp["status"].as_str() != Some("1") { + return Err(format!( + "地理编码失败: {}", + resp["info"].as_str().unwrap_or("未知错误") + )); + } + Ok(resp) + } + + /// 逆地理编码:坐标 → 地址 + async fn regeocode(&self, location: &str) -> Result { + let resp: serde_json::Value = self + .http + .get("https://restapi.amap.com/v3/geocode/regeo") + .query(&[ + ("key", self.key.as_str()), + ("location", location), + ("output", "JSON"), + ("appname", "amap-lbs-skill"), + ]) + .send() + .await + .map_err(|e| format!("逆地理编码请求失败: {}", e))? + .json() + .await + .map_err(|e| format!("解析逆地理编码响应失败: {}", e))?; + + if resp["status"].as_str() != Some("1") { + return Err(format!( + "逆地理编码失败: {}", + resp["info"].as_str().unwrap_or("未知错误") + )); + } + Ok(resp) + } + + /// 步行路线规划 + async fn walking_route( + &self, + origin: &str, + destination: &str, + ) -> Result { + let resp: serde_json::Value = self + .http + .get("https://restapi.amap.com/v3/direction/walking") + .query(&[ + ("key", self.key.as_str()), + ("origin", origin), + ("destination", destination), + ("appname", "amap-lbs-skill"), + ]) + .send() + .await + .map_err(|e| format!("步行路线请求失败: {}", e))? + .json() + .await + .map_err(|e| format!("解析步行路线响应失败: {}", e))?; + + if resp["status"].as_str() != Some("1") { + return Err(format!( + "步行路线失败: {}", + resp["info"].as_str().unwrap_or("未知错误") + )); + } + Ok(resp) + } + + /// 驾车路线规划 + async fn driving_route( + &self, + origin: &str, + destination: &str, + waypoints: Option<&str>, + strategy: u32, + ) -> Result { + let strategy_s = strategy.to_string(); + let mut params = vec![ + ("key", self.key.as_str()), + ("origin", origin), + ("destination", destination), + ("strategy", &strategy_s), + ("extensions", "base"), + ("appname", "amap-lbs-skill"), + ]; + if let Some(wp) = waypoints { + params.push(("waypoints", wp)); + } + + let resp: serde_json::Value = self + .http + .get("https://restapi.amap.com/v3/direction/driving") + .query(¶ms) + .send() + .await + .map_err(|e| format!("驾车路线请求失败: {}", e))? + .json() + .await + .map_err(|e| format!("解析驾车路线响应失败: {}", e))?; + + if resp["status"].as_str() != Some("1") { + return Err(format!( + "驾车路线失败: {}", + resp["info"].as_str().unwrap_or("未知错误") + )); + } + Ok(resp) + } + + /// 骑行路线规划 + async fn riding_route( + &self, + origin: &str, + destination: &str, + ) -> Result { + let resp: serde_json::Value = self + .http + .get("https://restapi.amap.com/v4/direction/bicycling") + .query(&[ + ("key", self.key.as_str()), + ("origin", origin), + ("destination", destination), + ("appname", "amap-lbs-skill"), + ]) + .send() + .await + .map_err(|e| format!("骑行路线请求失败: {}", e))? + .json() + .await + .map_err(|e| format!("解析骑行路线响应失败: {}", e))?; + + if resp["errcode"].as_i64() != Some(0) { + return Err(format!( + "骑行路线失败: {}", + resp["errmsg"].as_str().unwrap_or("未知错误") + )); + } + Ok(resp) + } + + /// 公交路线规划 + async fn transit_route( + &self, + origin: &str, + destination: &str, + city: &str, + strategy: u32, + ) -> Result { + let resp: serde_json::Value = self + .http + .get("https://restapi.amap.com/v3/direction/transit/integrated") + .query(&[ + ("key", self.key.as_str()), + ("origin", origin), + ("destination", destination), + ("city", city), + ("strategy", &strategy.to_string()), + ("appname", "amap-lbs-skill"), + ]) + .send() + .await + .map_err(|e| format!("公交路线请求失败: {}", e))? + .json() + .await + .map_err(|e| format!("解析公交路线响应失败: {}", e))?; + + if resp["status"].as_str() != Some("1") { + return Err(format!( + "公交路线失败: {}", + resp["info"].as_str().unwrap_or("未知错误") + )); + } + Ok(resp) + } +} + +// ═══════════════════════════════════════════════ +// 1. amap_poi_search — POI 搜索 +// ═══════════════════════════════════════════════ + +pub fn poi_search_spec() -> SkillSpec { + SkillSpec { + name: "amap_poi_search".into(), + description: "高德地图 POI(地点)搜索,支持关键词搜索、城市限定、周边搜索。\ + 参数 keywords 为搜索关键词,city 为城市名(可选),location 为中心点坐标'经度,纬度'(周边搜索时必填),\ + radius 为搜索半径米(默认1000),page 为页码(默认1),offset 为每页数量(默认10,最大25)" + .into(), + risk_level: RiskLevel::Low, + parameters: serde_json::json!({ + "type": "object", + "properties": { + "keywords": {"type": "string", "description": "搜索关键词,如 肯德基、酒店、景点"}, + "city": {"type": "string", "description": "城市名称或城市编码,如 北京、上海"}, + "location": {"type": "string", "description": "中心点坐标,格式:经度,纬度。用于周边搜索"}, + "radius": {"type": "integer", "description": "搜索半径(米),默认1000"}, + "page": {"type": "integer", "description": "页码,默认1"}, + "offset": {"type": "integer", "description": "每页数量,默认10,最大25"} + }, + "required": ["keywords"] + }), + timeout_secs: 15, + } +} + +pub async fn poi_search_execute(params: serde_json::Value) -> SkillResult { + let keywords = params["keywords"] + .as_str() + .unwrap_or_else(|| { + tracing::warn!("amap_poi_search: 缺少 keywords"); + "" + }); + if keywords.is_empty() { + return SkillResult::error("参数错误:缺少 keywords(搜索关键词)"); + } + + let city = params["city"].as_str(); + let location = params["location"].as_str(); + let radius = params["radius"].as_u64().unwrap_or(1000) as u32; + let page = params["page"].as_u64().unwrap_or(1) as u32; + let offset = params["offset"].as_u64().unwrap_or(10).min(25) as u32; + + let client = match AmapClient::new() { + Ok(c) => c, + Err(e) => return SkillResult::error(e), + }; + + // 有 location → 周边搜索,否则 → 文本搜索 + let result = if let Some(loc) = location { + tracing::info!("📍 周边搜索: {} @ {} (radius={}m)", keywords, loc, radius); + client + .poi_around(keywords, loc, radius, page, offset) + .await + } else { + tracing::info!("🔍 POI 搜索: {} city={:?}", keywords, city); + client + .poi_text_search(keywords, city, page, offset) + .await + }; + + match result { + Ok(data) => { + let count = data["count"].as_str().unwrap_or("?"); + let pois = &data["pois"]; + let poi_count = pois.as_array().map(|a| a.len()).unwrap_or(0); + let text = format!( + "✅ 搜索「{keywords}」共 {count} 条结果,当前页显示 {poi_count} 条", + ); + SkillResult::ok_with_data(text, data) + } + Err(e) => SkillResult::error(e), + } +} + +// ═══════════════════════════════════════════════ +// 2. amap_geocode — 地理编码 +// ═══════════════════════════════════════════════ + +pub fn geocode_spec() -> SkillSpec { + SkillSpec { + name: "amap_geocode".into(), + description: + "高德地图地理编码:将地址转换为经纬度坐标。参数 address 为地址名称,city 为城市名(可选)" + .into(), + risk_level: RiskLevel::Low, + parameters: serde_json::json!({ + "type": "object", + "properties": { + "address": {"type": "string", "description": "地址名称,如 西直门、北京南站、天安门"}, + "city": {"type": "string", "description": "城市名称,可选,用于缩小搜索范围"} + }, + "required": ["address"] + }), + timeout_secs: 10, + } +} + +pub async fn geocode_execute(params: serde_json::Value) -> SkillResult { + let address = params["address"].as_str().unwrap_or(""); + if address.is_empty() { + return SkillResult::error("参数错误:缺少 address(地址名称)"); + } + let city = params["city"].as_str(); + + let client = match AmapClient::new() { + Ok(c) => c, + Err(e) => return SkillResult::error(e), + }; + + tracing::info!("📍 地理编码: {} city={:?}", address, city); + + match client.geocode(address, city).await { + Ok(data) => { + let geocodes = &data["geocodes"]; + if let Some(arr) = geocodes.as_array() { + if arr.is_empty() { + return SkillResult::error(format!("未找到地址: {address}")); + } + let loc = arr[0]["location"].as_str().unwrap_or("未知"); + let formatted = arr[0]["formatted_address"].as_str().unwrap_or(address); + let text = format!("📍 {formatted} 坐标: {loc}"); + let out = serde_json::json!({ + "address": formatted, + "location": loc, + "top_result": arr[0] + }); + SkillResult::ok_with_data(text, out) + } else { + SkillResult::error(format!("未找到地址: {address}")) + } + } + Err(e) => SkillResult::error(e), + } +} + +// ═══════════════════════════════════════════════ +// 3. amap_reverse_geocode — 逆地理编码 +// ═══════════════════════════════════════════════ + +pub fn reverse_geocode_spec() -> SkillSpec { + SkillSpec { + name: "amap_reverse_geocode".into(), + description: + "高德地图逆地理编码:将经纬度坐标转换为地址描述。参数 location 为坐标,格式'经度,纬度'" + .into(), + risk_level: RiskLevel::Low, + parameters: serde_json::json!({ + "type": "object", + "properties": { + "location": {"type": "string", "description": "经纬度坐标,格式:经度,纬度"} + }, + "required": ["location"] + }), + timeout_secs: 10, + } +} + +pub async fn reverse_geocode_execute(params: serde_json::Value) -> SkillResult { + let location = params["location"].as_str().unwrap_or(""); + if location.is_empty() { + return SkillResult::error("参数错误:缺少 location(经纬度坐标)"); + } + + let client = match AmapClient::new() { + Ok(c) => c, + Err(e) => return SkillResult::error(e), + }; + + tracing::info!("📍 逆地理编码: {}", location); + + match client.regeocode(location).await { + Ok(data) => { + let regeocode = &data["regeocode"]; + let addr = regeocode["formatted_address"] + .as_str() + .unwrap_or("未知地址"); + let text = format!("📍 坐标 {location} → {addr}"); + SkillResult::ok_with_data(text, data) + } + Err(e) => SkillResult::error(e), + } +} + +// ═══════════════════════════════════════════════ +// 4. amap_route_plan — 路径规划 +// ═══════════════════════════════════════════════ + +pub fn route_plan_spec() -> SkillSpec { + SkillSpec { + name: "amap_route_plan".into(), + description: "高德地图路径规划,支持步行(walking)、驾车(driving)、骑行(riding)、公交(transit)四种出行方式。\ + 参数 origin 为起点坐标'经度,纬度',destination 为终点坐标'经度,纬度',\ + type 为出行方式,city 为城市名(公交必填),waypoints 为途经点(驾车可选,多个用;分隔)" + .into(), + risk_level: RiskLevel::Low, + parameters: serde_json::json!({ + "type": "object", + "properties": { + "type": {"type": "string", "description": "出行方式:walking/driving/riding/transit", "enum": ["walking", "driving", "riding", "transit"]}, + "origin": {"type": "string", "description": "起点坐标,格式:经度,纬度"}, + "destination": {"type": "string", "description": "终点坐标,格式:经度,纬度"}, + "city": {"type": "string", "description": "城市名称(公交方式必填)"}, + "waypoints": {"type": "string", "description": "途经点坐标,多个用;分隔(仅驾车方式)"}, + "strategy": {"type": "integer", "description": "策略:驾车默认10(躲避拥堵),公交默认0(最快捷)"} + }, + "required": ["type", "origin", "destination"] + }), + timeout_secs: 15, + } +} + +pub async fn route_plan_execute(params: serde_json::Value) -> SkillResult { + let route_type = params["type"].as_str().unwrap_or(""); + let origin = params["origin"].as_str().unwrap_or(""); + let destination = params["destination"].as_str().unwrap_or(""); + + if route_type.is_empty() || origin.is_empty() || destination.is_empty() { + return SkillResult::error( + "参数错误:缺少 type/origin/destination。格式示例:\ + {\"type\":\"walking\",\"origin\":\"116.397428,39.90923\",\"destination\":\"116.427281,39.903719\"}", + ); + } + + let client = match AmapClient::new() { + Ok(c) => c, + Err(e) => return SkillResult::error(e), + }; + + tracing::info!( + "🛣️ 路径规划: type={}, {} → {}", + route_type, + origin, + destination + ); + + match route_type { + "walking" => { + match client.walking_route(origin, destination).await { + Ok(data) => { + let summary = route_summary(&data, "walking"); + SkillResult::ok_with_data(summary, data) + } + Err(e) => SkillResult::error(e), + } + } + "driving" => { + let waypoints = params["waypoints"].as_str(); + let strategy = params["strategy"].as_u64().unwrap_or(10) as u32; + match client + .driving_route(origin, destination, waypoints, strategy) + .await + { + Ok(data) => { + let summary = route_summary(&data, "driving"); + SkillResult::ok_with_data(summary, data) + } + Err(e) => SkillResult::error(e), + } + } + "riding" => { + match client.riding_route(origin, destination).await { + Ok(data) => { + let summary = route_summary(&data, "riding"); + SkillResult::ok_with_data(summary, data) + } + Err(e) => SkillResult::error(e), + } + } + "transit" => { + let city = params["city"].as_str().unwrap_or(""); + if city.is_empty() { + return SkillResult::error("公交路线需要提供 city 参数(城市名称)"); + } + let strategy = params["strategy"].as_u64().unwrap_or(0) as u32; + match client + .transit_route(origin, destination, city, strategy) + .await + { + Ok(data) => { + let summary = route_summary(&data, "transit"); + SkillResult::ok_with_data(summary, data) + } + Err(e) => SkillResult::error(e), + } + } + other => SkillResult::error(format!( + "不支持的出行方式: {other},支持 walking/driving/riding/transit" + )), + } +} + +/// 提取路线摘要信息 +fn route_summary(data: &serde_json::Value, route_type: &str) -> String { + match route_type { + "walking" => { + if let Some(path) = data["route"]["paths"] + .as_array() + .and_then(|a| a.first()) + { + let dist = path["distance"] + .as_str() + .and_then(|s| s.parse::().ok()) + .unwrap_or(0.0); + let dur = path["duration"] + .as_str() + .and_then(|s| s.parse::().ok()) + .unwrap_or(0.0); + format!( + "🚶 步行路线:距离 {:.2} 公里,预计 {} 分钟", + dist / 1000.0, + (dur / 60.0) as u32 + ) + } else { + "🚶 步行路线规划完成".into() + } + } + "driving" => { + if let Some(path) = data["route"]["paths"] + .as_array() + .and_then(|a| a.first()) + { + let dist = path["distance"] + .as_str() + .and_then(|s| s.parse::().ok()) + .unwrap_or(0.0); + let dur = path["duration"] + .as_str() + .and_then(|s| s.parse::().ok()) + .unwrap_or(0.0); + let tolls = path["tolls"] + .as_str() + .and_then(|s| s.parse::().ok()) + .unwrap_or(0.0); + format!( + "🚗 驾车路线:距离 {:.2} 公里,预计 {} 分钟,过路费 {} 元", + dist / 1000.0, + (dur / 60.0) as u32, + tolls + ) + } else { + "🚗 驾车路线规划完成".into() + } + } + "riding" => { + if let Some(path) = data["data"]["paths"] + .as_array() + .and_then(|a| a.first()) + { + let dist = path["distance"].as_f64().unwrap_or(0.0); + let dur = path["duration"].as_f64().unwrap_or(0.0); + format!( + "🚴 骑行路线:距离 {:.2} 公里,预计 {} 分钟", + dist / 1000.0, + (dur / 60.0) as u32 + ) + } else { + "🚴 骑行路线规划完成".into() + } + } + "transit" => { + let count = data["route"]["transits"] + .as_array() + .map(|a| a.len()) + .unwrap_or(0); + if let Some(t) = data["route"]["transits"] + .as_array() + .and_then(|a| a.first()) + { + let dur = t["duration"] + .as_str() + .and_then(|s| s.parse::().ok()) + .unwrap_or(0.0); + let cost = t["cost"].as_str().and_then(|s| s.parse::().ok()); + let wdist = t["walking_distance"] + .as_str() + .and_then(|s| s.parse::().ok()) + .unwrap_or(0.0); + let mut text = format!( + "🚌 公交路线:共 {count} 个方案,推荐方案预计 {} 分钟", + (dur / 60.0) as u32 + ); + if let Some(c) = cost { + text.push_str(&format!(",费用 {c} 元")); + } + text.push_str(&format!( + ",步行 {:.0} 米", + wdist + )); + text + } else { + format!("🚌 公交路线规划完成,共 {count} 个方案") + } + } + _ => "路线规划完成".into(), + } +} + +// ═══════════════════════════════════════════════ +// 5. amap_travel_plan — 旅游规划 +// ═══════════════════════════════════════════════ + +pub fn travel_plan_spec() -> SkillSpec { + SkillSpec { + name: "amap_travel_plan".into(), + description: "高德地图智能旅游规划:自动搜索城市内的兴趣点并规划游览路线。\ + 参数 city 为城市名(必填),interests 为兴趣点关键词数组如['景点','美食'],\ + route_type 为路线类型 walking/driving/riding/transit(默认walking)" + .into(), + risk_level: RiskLevel::Low, + parameters: serde_json::json!({ + "type": "object", + "properties": { + "city": {"type": "string", "description": "城市名称,如 北京、杭州、上海"}, + "interests": {"type": "array", "items": {"type": "string"}, "description": "兴趣点关键词,如 ['景点','美食','酒店']"}, + "route_type": {"type": "string", "description": "路线类型:walking/driving/riding/transit", "enum": ["walking", "driving", "riding", "transit"]} + }, + "required": ["city"] + }), + timeout_secs: 60, + } +} + +pub async fn travel_plan_execute(params: serde_json::Value) -> SkillResult { + let city = params["city"].as_str().unwrap_or(""); + if city.is_empty() { + return SkillResult::error("参数错误:缺少 city(城市名称)"); + } + + let interests: Vec<&str> = params["interests"] + .as_array() + .map(|a| { + a.iter() + .filter_map(|v| v.as_str()) + .collect::>() + }) + .unwrap_or_else(|| vec!["景点", "美食"]); + + let route_type = params["route_type"].as_str().unwrap_or("walking"); + let valid_types = ["walking", "driving", "riding", "transit"]; + if !valid_types.contains(&route_type) { + return SkillResult::error(format!( + "无效的路线类型: {route_type},支持 walking/driving/riding/transit" + )); + } + + let client = match AmapClient::new() { + Ok(c) => c, + Err(e) => return SkillResult::error(e), + }; + + tracing::info!( + "🗺️ 旅游规划: city={}, interests={:?}, route={}", + city, + interests, + route_type + ); + + let mut all_pois: Vec = Vec::new(); + let mut map_tasks: Vec = Vec::new(); + + // 第一步:搜索各类兴趣点 + for interest in &interests { + tracing::info!("📍 搜索 {}...", interest); + + match client + .poi_text_search(interest, Some(city), 1, 5) + .await + { + Ok(data) => { + if let Some(pois) = data["pois"].as_array() { + for poi in pois { + let name = poi["name"].as_str().unwrap_or("未知"); + let address = poi["address"].as_str().unwrap_or(""); + let loc_str = poi["location"].as_str().unwrap_or("0,0"); + let parts: Vec<&str> = loc_str.split(',').collect(); + let lng = parts + .first() + .and_then(|s| s.parse::().ok()) + .unwrap_or(0.0); + let lat = parts + .get(1) + .and_then(|s| s.parse::().ok()) + .unwrap_or(0.0); + + map_tasks.push(serde_json::json!({ + "type": "poi", + "lnglat": [lng, lat], + "sort": interest, + "text": name, + "remark": address + })); + } + all_pois.extend(pois.clone()); + tracing::info!(" 找到 {} 个 {}", pois.len(), interest); + } + } + Err(e) => { + tracing::warn!(" 搜索 {} 失败: {}", interest, e); + } + } + } + + if all_pois.is_empty() { + return SkillResult::error(format!("在 {city} 未找到相关兴趣点")); + } + + // 第二步:规划 POI 之间的路线 + if all_pois.len() >= 2 { + for i in 0..all_pois.len() - 1 { + let start = &all_pois[i]; + let end = &all_pois[i + 1]; + let start_loc = start["location"].as_str().unwrap_or("0,0"); + let end_loc = end["location"].as_str().unwrap_or("0,0"); + + let start_parts: Vec = start_loc + .split(',') + .filter_map(|s| s.parse().ok()) + .collect(); + let end_parts: Vec = end_loc + .split(',') + .filter_map(|s| s.parse().ok()) + .collect(); + + let start_name = start["name"].as_str().unwrap_or("?"); + let end_name = end["name"].as_str().unwrap_or("?"); + + let mut route_task = serde_json::json!({ + "type": "route", + "routeType": route_type, + "start": start_parts, + "end": end_parts, + "remark": format!("从 {start_name} 到 {end_name}") + }); + + if route_type == "transit" { + route_task["city"] = serde_json::json!(city); + } + + map_tasks.push(route_task); + } + } + + // 第三步:生成地图可视化链接 + let map_link = generate_map_link(&map_tasks); + + let route_type_cn = match route_type { + "walking" => "步行", + "driving" => "驾车", + "riding" => "骑行", + "transit" => "公交", + _ => route_type, + }; + + let poi_names: Vec<&str> = all_pois + .iter() + .filter_map(|p| p["name"].as_str()) + .collect(); + + let text = format!( + "🗺️ {city} 旅游规划完成({route_type_cn})\n\ + 📍 推荐地点({} 个):{} \n\ + 🔗 地图链接:{map_link}", + all_pois.len(), + poi_names.join(" → ") + ); + + let output = serde_json::json!({ + "city": city, + "pois": all_pois, + "map_tasks": map_tasks, + "map_link": map_link, + "route_type": route_type, + }); + + SkillResult::ok_with_data(text, output) +} + +// ═══════════════════════════════════════════════ +// 6. amap_map_link — 生成地图可视化链接 +// ═══════════════════════════════════════════════ + +pub fn map_link_spec() -> SkillSpec { + SkillSpec { + name: "amap_map_link".into(), + description: "生成高德地图可视化链接,支持 POI 和路线展示。\ + 参数 map_data 为数组,每项包含:type('poi'|'route')、lnglat/start/end、sort/text/remark 等字段" + .into(), + risk_level: RiskLevel::Low, + parameters: serde_json::json!({ + "type": "object", + "properties": { + "map_data": {"type": "array", "description": "地图数据数组,每项格式:{\"type\":\"poi\",\"lnglat\":[lng,lat],\"sort\":\"分类\",\"text\":\"名称\",\"remark\":\"备注\"} 或 {\"type\":\"route\",\"routeType\":\"walking\",\"start\":[lng,lat],\"end\":[lng,lat],\"remark\":\"描述\"}"} + }, + "required": ["map_data"] + }), + timeout_secs: 5, + } +} + +pub async fn map_link_execute(params: serde_json::Value) -> SkillResult { + let map_data = match params.get("map_data") { + Some(serde_json::Value::Array(arr)) => arr.clone(), + Some(_) => return SkillResult::error("map_data 必须是数组"), + None => return SkillResult::error("缺少 map_data 参数"), + }; + + if map_data.is_empty() { + return SkillResult::error("map_data 不能为空"); + } + + let link = generate_map_link(&map_data); + let text = format!("🗺️ 地图可视化链接:\n{link}"); + let output = serde_json::json!({ + "map_link": link, + "item_count": map_data.len() + }); + SkillResult::ok_with_data(text, output) +} + +/// 生成高德地图可视化链接 +fn generate_map_link(map_task_data: &[serde_json::Value]) -> String { + let base_url = "https://a.amap.com/jsapi_demo_show/static/openclaw/travel_plan.html"; + let data_str = serde_json::to_string(map_task_data).unwrap_or_default(); + let encoded = urlencoding(&data_str); // 手动 URL 编码(只编码特殊字符) + format!("{base_url}?data={encoded}") +} + +/// 简易 URL 编码(仅编码必要的特殊字符) +fn urlencoding(s: &str) -> String { + s.chars() + .map(|c| match c { + ' ' => "%20".to_string(), + ':' => "%3A".to_string(), + '/' => "%2F".to_string(), + '?' => "%3F".to_string(), + '#' => "%23".to_string(), + '[' => "%5B".to_string(), + ']' => "%5D".to_string(), + '@' => "%40".to_string(), + '!' => "%21".to_string(), + '$' => "%24".to_string(), + '&' => "%26".to_string(), + '\'' => "%27".to_string(), + '(' => "%28".to_string(), + ')' => "%29".to_string(), + '*' => "%2A".to_string(), + '+' => "%2B".to_string(), + ',' => "%2C".to_string(), + ';' => "%3B".to_string(), + '=' => "%3D".to_string(), + '%' => "%25".to_string(), + '{' => "%7B".to_string(), + '}' => "%7D".to_string(), + '"' => "%22".to_string(), + '\\' => "%5C".to_string(), + other => other.to_string(), + }) + .collect() +} + +// ═══════════════════════════════════════════════ +// 统一分发入口 +// ═══════════════════════════════════════════════ + +/// 所有 amap 工具的 spec 列表 +pub fn specs() -> Vec { + vec![ + poi_search_spec(), + geocode_spec(), + reverse_geocode_spec(), + route_plan_spec(), + travel_plan_spec(), + map_link_spec(), + ] +} + +/// 根据名称分发执行 +pub async fn execute(name: &str, params: serde_json::Value) -> Option { + match name { + "amap_poi_search" => Some(poi_search_execute(params).await), + "amap_geocode" => Some(geocode_execute(params).await), + "amap_reverse_geocode" => Some(reverse_geocode_execute(params).await), + "amap_route_plan" => Some(route_plan_execute(params).await), + "amap_travel_plan" => Some(travel_plan_execute(params).await), + "amap_map_link" => Some(map_link_execute(params).await), + _ => None, + } +} + +// ═══════════════════════════════════════════════ +// 测试 +// ═══════════════════════════════════════════════ + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_url_encoding() { + let encoded = urlencoding("https://example.com/path?q=hello world"); + // 不应该有冒号、斜杠等特殊字符被漏掉 + assert!(!encoded.contains("://")); + assert!(encoded.contains("%3A%2F%2F")); + assert!(encoded.contains("%20")); + } + + #[test] + fn test_generate_map_link() { + let data = vec![ + serde_json::json!({ + "type": "poi", + "lnglat": [116.397428, 39.90923], + "sort": "风景名胜", + "text": "故宫博物院", + "remark": "明清皇家宫殿" + }), + serde_json::json!({ + "type": "route", + "routeType": "walking", + "start": [116.397428, 39.90923], + "end": [116.427281, 39.903719], + "remark": "步行路线" + }), + ]; + let link = generate_map_link(&data); + assert!(link.starts_with("https://a.amap.com/jsapi_demo_show/static/openclaw/travel_plan.html?data=")); + assert!(link.contains("walking")); + } + + #[test] + fn test_route_summary_walking() { + let data = serde_json::json!({ + "route": { + "paths": [{ + "distance": "1500", + "duration": "1200" + }] + } + }); + let summary = route_summary(&data, "walking"); + assert!(summary.contains("1.50")); + assert!(summary.contains("20")); + } + + #[test] + fn test_specs_count() { + let s = specs(); + assert_eq!(s.len(), 6); + // 验证每个都有非空 name 和 description + for spec in &s { + assert!(!spec.name.is_empty()); + assert!(!spec.description.is_empty()); + } + } + + /// 端到端测试(需要 AMAP_WEBSERVICE_KEY 环境变量) + #[tokio::test] + async fn test_poi_search_integration() { + if std::env::var("AMAP_WEBSERVICE_KEY").is_err() + && std::env::var("AMAP_KEY").is_err() + { + eprintln!("跳过集成测试: 未设置 AMAP_WEBSERVICE_KEY"); + return; + } + let result = + poi_search_execute(serde_json::json!({"keywords": "肯德基", "city": "北京"})).await; + assert!(result.success, "POI 搜索失败: {}", result.output); + } + + #[tokio::test] + async fn test_geocode_integration() { + if std::env::var("AMAP_WEBSERVICE_KEY").is_err() + && std::env::var("AMAP_KEY").is_err() + { + eprintln!("跳过集成测试: 未设置 AMAP_WEBSERVICE_KEY"); + return; + } + let result = + geocode_execute(serde_json::json!({"address": "天安门"})).await; + assert!(result.success, "地理编码失败: {}", result.output); + let data = result.data.unwrap(); + assert!(data["location"].as_str().is_some()); + } +} diff --git a/src/tools/builtins/datetime.rs b/src/tools/builtins/datetime.rs new file mode 100644 index 0000000..3eed213 --- /dev/null +++ b/src/tools/builtins/datetime.rs @@ -0,0 +1,19 @@ +//! 日期时间工具 +use chrono::Local; + +use super::super::types::{RiskLevel, SkillResult, SkillSpec}; + +pub fn spec() -> SkillSpec { + SkillSpec { + name: "get_current_datetime".into(), + description: "获取当前日期时间(北京时间 UTC+8)".into(), + risk_level: RiskLevel::Low, + parameters: serde_json::json!({"type":"object","properties":{"format":{"type":"string"}}}), + timeout_secs: 5, + } +} + +pub fn execute() -> SkillResult { + let now = Local::now().format("%Y-%m-%d %H:%M:%S").to_string(); + SkillResult::ok(serde_json::json!({"datetime": now}).to_string()) +} diff --git a/src/tools/builtins/fetch_page.rs b/src/tools/builtins/fetch_page.rs new file mode 100644 index 0000000..09b9de8 --- /dev/null +++ b/src/tools/builtins/fetch_page.rs @@ -0,0 +1,348 @@ +//! 网页内容提取工具 +//! +//! 使用可读性算法自动识别正文区域:清除广告/导航/页脚噪声, +//! 按文本密度评分,取最高分区域。同时支持站点自定义 CSS 选择器。 +//! +//! 站点选择器配置: `~/.ias/site_selectors.json` + +use reqwest::Client as HttpClient; +use scraper::{ElementRef, Html, Selector}; +use std::collections::HashMap; + +use super::super::types::{RiskLevel, SkillResult, SkillSpec}; + +/// 站点选择器配置文件路径 +const CONFIG_PATH: &str = ".ias/site_selectors.json"; + +/// 最大输出字符数 +const MAX_OUTPUT_CHARS: usize = 6000; + +/// 需移除的 HTML 标签 +const STRIP_TAGS: &[&str] = &[ + "script", "style", "noscript", "iframe", "svg", + "nav", "header", "footer", "aside", "form", +]; + +/// 需移除的 class/id 关键词 +const NOISE_PATTERNS: &[&str] = &[ + "ad", "ads", "advert", "banner", + "sidebar", "side-bar", "widget", + "footer", "header", "nav", "menu", "navbar", + "comment", "comments", + "social", "share", "sharing", + "related", "recommend", "recommended", + "subscribe", "newsletter", "signup", + "cookie", "consent", "popup", "modal", + "breadcrumb", "pagination", + "print", "disclaimer", + "toc", "table-of-contents", + "metadata", "infobox", "navbox", + "reference", "references", "reflist", + "catlinks", "mw-jump-link", "mw-editsection", + "edit-section", "noprint", "thumb", +]; + +/// 内容评分阈值:低于此分不取 +const MIN_SCORE_THRESHOLD: f64 = 50.0; + +// ─── 工具 spec ─── + +pub fn spec() -> SkillSpec { + SkillSpec { + name: "fetch_page".into(), + description: "抓取网页并自动提取正文内容。自动过滤广告、导航、页脚等噪声,适用于阅读文章、获取信息。可通过 ~/.ias/site_selectors.json 配置按站点的 CSS 选择器精确定位。".into(), + risk_level: RiskLevel::Low, + parameters: serde_json::json!({ + "type": "object", + "properties": { + "url": { "type": "string", "description": "网页 URL" } + }, + "required": ["url"] + }), + timeout_secs: 30, + } +} + +// ─── 执行 ─── + +pub async fn execute(params: serde_json::Value) -> SkillResult { + let url_str = params + .get("url") + .and_then(|v| v.as_str()) + .or_else(|| params.get("prompt").and_then(|v| v.as_str())) + .unwrap_or(""); + + if url_str.is_empty() { + return SkillResult::error("请提供 url 参数"); + } + + let url = match reqwest::Url::parse(url_str) { + Ok(u) => u, + Err(e) => return SkillResult::error(format!("无效的 URL: {}", e)), + }; + let domain = url.host_str().unwrap_or("unknown"); + + // 抓取页面 + let http = match HttpClient::builder() + .timeout(std::time::Duration::from_secs(20)) + .user_agent("Mozilla/5.0 (compatible; iAs-Bot/1.0)") + .build() + { + Ok(c) => c, + Err(e) => return SkillResult::error(format!("创建 HTTP 客户端失败: {}", e)), + }; + + let resp = match http.get(url.clone()).send().await { + Ok(r) => r, + Err(e) => { + if e.is_timeout() { return SkillResult::error("请求超时"); } + return SkillResult::error(format!("请求失败: {}", e)); + } + }; + + if !resp.status().is_success() { + return SkillResult::error(format!("HTTP {}", resp.status())); + } + + let html_text = match resp.text().await { + Ok(t) => t, + Err(e) => return SkillResult::error(format!("读取响应失败: {}", e)), + }; + + let document = Html::parse_document(&html_text); + let title = extract_title(&document); + + // ── 提取正文 ── + let content = extract_content(&document, domain); + + let display = if content.chars().count() > MAX_OUTPUT_CHARS { + let truncated: String = content.chars().take(MAX_OUTPUT_CHARS).collect(); + format!("{}\n\n... (内容已截断)", truncated) + } else { + content + }; + + let mut output = String::new(); + output.push_str(&format!("📄 {}\n\n🔗 {}\n\n{}", title, url_str, display)); + + SkillResult::ok(output) +} + +// ─── 核心:可读性提取 ─── + +/// 提取主要内容 +fn extract_content(document: &Html, domain: &str) -> String { + // 1. 先尝试自定义选择器 + let config = load_config(); + if let Some(selectors) = config.get(domain) { + for sel_str in selectors { + if let Some(text) = try_selector(document, sel_str) { + if text.len() > 50 { + return text; + } + } + } + } + + // 2. 先用可读性算法 + if let Some(text) = readability_extract(document) { + if text.len() > 50 { + return text; + } + } + + // 3. 回退到 body + if let Some(text) = try_selector(document, "body") { + return text; + } + + "无法提取内容".to_string() +} + +/// 可读性算法:清除噪声 → 评分 → 取最高分 +fn readability_extract(document: &Html) -> Option { + let body = document.root_element(); + + // 直接对原始文档评分(跳过噪声元素) + let candidates = score_elements(&body); + + // 取最高分 + candidates + .into_iter() + .max_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal)) + .map(|(el, _score)| { + let text = extract_clean_text(&el); + // 合并空白行 + let lines: Vec<&str> = text.lines().map(|l| l.trim()).filter(|l| !l.is_empty()).collect(); + lines.join("\n") + }) + .filter(|t| t.len() > 50) +} + +/// 评分所有候选元素 +fn score_elements<'a>(root: &ElementRef<'a>) -> Vec<(ElementRef<'a>, f64)> { + let mut candidates: Vec<(ElementRef, f64)> = Vec::new(); + + // 收集所有块级元素 + let block_tags = [ + "div", "article", "section", "main", "p", + ]; + + for tag in &block_tags { + if let Ok(sel) = Selector::parse(tag) { + for el in root.select(&sel) { + // 跳过噪声 + if is_noise(&el) { + continue; + } + let score = score_element(&el); + if score > MIN_SCORE_THRESHOLD { + candidates.push((el, score)); + } + } + } + } + + candidates +} + +/// 判断元素是否为噪声 +fn is_noise(el: &ElementRef) -> bool { + let tag = el.value().name().to_lowercase(); + if STRIP_TAGS.contains(&tag.as_str()) { + return true; + } + + let class = el.value().attr("class").unwrap_or("").to_lowercase(); + let id = el.value().attr("id").unwrap_or("").to_lowercase(); + let combined = format!("{} {}", class, id); + + for pattern in NOISE_PATTERNS { + if combined.contains(pattern) { + return true; + } + } + + false +} + +/// 评分单个元素 +fn score_element(el: &ElementRef) -> f64 { + let text = el.text().collect::>().join(""); + let text_len = text.trim().len() as f64; + + // 基础分 = 文本长度 + let mut score = text_len; + + // 逗号数 → 句子结构丰富度 + let commas = text.matches(',').count() as f64; + score += commas * 3.0; + + // 句号/问号/感叹号 → 完整句子 + let sentences = text.matches(|c| c == '.' || c == '。' || c == '?' || c == '?' || c == '!' || c == '!').count() as f64; + score += sentences * 2.0; + + // 段落数(

标签) + if let Ok(p_sel) = Selector::parse("p") { + let p_count = el.select(&p_sel).count() as f64; + score += p_count * 5.0; + } + + // 链接密度惩罚:链接文字占总文字比例越高,越像导航/索引 + let link_text: String = el + .select(&Selector::parse("a").unwrap_or_else(|_| unreachable!())) + .flat_map(|a| a.text()) + .collect(); + let link_len = link_text.trim().len() as f64; + if text_len > 0.0 { + let link_ratio = link_len / text_len; + if link_ratio > 0.8 { + score *= 0.05; // 几乎全是链接 → 导航/目录 + } else if link_ratio > 0.5 { + score *= 0.15; + } else if link_ratio > 0.3 { + score *= 0.4; + } + } + + // 代码块惩罚(代码块多 → 技术文档,保留但降分) + if let Ok(code_sel) = Selector::parse("pre, code") { + let code_count = el.select(&code_sel).count() as f64; + score -= code_count * 10.0; + } + + score +} + +/// 提取元素内纯文本(过滤噪声标签) +fn extract_clean_text(el: &ElementRef) -> String { + let mut text = String::new(); + + for child in el.descendants() { + // 跳过噪声元素 + if let Some(child_el) = ElementRef::wrap(child) { + if is_noise(&child_el) { + continue; + } + } + + if let Some(t) = child.value().as_text() { + let t = t.trim(); + if !t.is_empty() { + if !text.is_empty() && !text.ends_with('\n') && !text.ends_with(' ') { + text.push(' '); + } + text.push_str(t); + } + } + + // 块级标签换行 + if let Some(el_ref) = ElementRef::wrap(child) { + let tag = el_ref.value().name(); + if matches!(tag, "br" | "p" | "div" | "li" | "h1" | "h2" | "h3" | "h4" | "h5" | "h6" | "tr" | "article" | "section" | "blockquote") { + text.push('\n'); + } + } + } + + text +} + +// ─── 辅助 ─── + +fn extract_title(document: &Html) -> String { + if let Ok(sel) = Selector::parse("title") { + if let Some(el) = document.select(&sel).next() { + let t: String = el.text().collect::>().join(" ").trim().to_string(); + if !t.is_empty() { return t; } + } + } + "无标题".to_string() +} + +fn try_selector(document: &Html, sel_str: &str) -> Option { + let sel = Selector::parse(sel_str).ok()?; + let mut text = String::new(); + for element in document.select(&sel) { + text.push_str(&extract_clean_text(&element)); + text.push('\n'); + } + let trimmed = text.trim().to_string(); + if trimmed.is_empty() { None } else { Some(trimmed) } +} + +fn load_config() -> HashMap> { + let home = match dirs::home_dir() { + Some(h) => h, + None => return HashMap::new(), + }; + let path = home.join(CONFIG_PATH); + if !path.exists() { + return HashMap::new(); + } + let content = match std::fs::read_to_string(&path) { + Ok(c) => c, + Err(_) => return HashMap::new(), + }; + serde_json::from_str(&content).unwrap_or_default() +} diff --git a/src/tools/builtins/memos.rs b/src/tools/builtins/memos.rs new file mode 100644 index 0000000..10de92f --- /dev/null +++ b/src/tools/builtins/memos.rs @@ -0,0 +1,114 @@ +//! 备忘录工具 — 文件存储,支持 add/list/delete +use chrono::Local; + +use super::super::types::{RiskLevel, SkillResult, SkillSpec}; + +pub fn spec() -> SkillSpec { + SkillSpec { + name: "manage_memos".into(), + description: "管理备忘录:添加/列出/删除".into(), + risk_level: RiskLevel::Low, + parameters: serde_json::json!({ + "type": "object", + "properties": { + "action": {"type": "string", "enum": ["add", "list", "delete"]}, + "content": {"type": "string"}, + "id": {"type": "integer"} + } + }), + timeout_secs: 10, + } +} + +pub async fn execute(args_json: &str) -> SkillResult { + let mut params: serde_json::Value = serde_json::from_str(args_json).unwrap_or_default(); + // call_capability 透传完整参数,action/content 在顶层 + // 如果顶层没有 action,尝试从 prompt 推断 + if params.get("action").is_none() { + let prompt = params.get("prompt").and_then(|v| v.as_str()).unwrap_or(""); + let action = if prompt.contains("添加") || prompt.contains("新增") || prompt.contains("add") + { + "add" + } else if prompt.contains("删除") || prompt.contains("移除") || prompt.contains("delete") + { + "delete" + } else { + "list" + }; + let content = params.get("prompt").and_then(|v| v.as_str()).unwrap_or(""); + let memo_id: i32 = content + .chars() + .filter(|c| c.is_ascii_digit()) + .collect::() + .parse() + .unwrap_or(0); + params = serde_json::json!({"action": action, "content": content, "id": memo_id}); + } + let action = params + .get("action") + .and_then(|v| v.as_str()) + .unwrap_or("list"); + let path = ".data/memos.json"; + if let Err(e) = std::fs::create_dir_all(".data") { + return SkillResult::error(format!("创建目录失败: {}", e)); + } + + let mut data: Vec = if std::path::Path::new(path).exists() { + serde_json::from_str(&std::fs::read_to_string(path).unwrap_or_default()).unwrap_or_default() + } else { + vec![] + }; + + match action { + "add" => { + let content = params.get("content").and_then(|v| v.as_str()).unwrap_or(""); + if content.is_empty() { + return SkillResult::error("请提供 content 参数"); + } + let next_id = data + .iter() + .filter_map(|m| m.get("id").and_then(|v| v.as_i64())) + .max() + .unwrap_or(0) + + 1; + data.push(serde_json::json!({ + "id": next_id, "content": content, + "created_at": Local::now().format("%Y-%m-%d %H:%M").to_string() + })); + if let Err(e) = std::fs::write(path, serde_json::to_string_pretty(&data).unwrap()) { + return SkillResult::error(format!("写入备忘录失败: {}", e)); + } + SkillResult::ok(format!("已添加备忘录 #{}", next_id)) + } + "list" => { + if data.is_empty() { + return SkillResult::ok("暂无备忘录".to_string()); + } + let lines: Vec = data + .iter() + .map(|m| { + format!( + "#{} [{}] {}", + m["id"], + m.get("created_at").and_then(|v| v.as_str()).unwrap_or("?"), + m["content"].as_str().unwrap_or("?") + ) + }) + .collect(); + SkillResult::ok(lines.join("\n")) + } + "delete" => { + let id = params.get("id").and_then(|v| v.as_i64()).unwrap_or(0); + let before = data.len(); + data.retain(|m| m.get("id").and_then(|v| v.as_i64()) != Some(id)); + if data.len() == before { + return SkillResult::error(format!("未找到备忘录 #{}", id)); + } + if let Err(e) = std::fs::write(path, serde_json::to_string_pretty(&data).unwrap()) { + return SkillResult::error(format!("写入备忘录失败: {}", e)); + } + SkillResult::ok(format!("已删除备忘录 #{}", id)) + } + _ => SkillResult::error("未知操作,支持: add, list, delete"), + } +} diff --git a/src/tools/builtins/mod.rs b/src/tools/builtins/mod.rs new file mode 100644 index 0000000..4b86d71 --- /dev/null +++ b/src/tools/builtins/mod.rs @@ -0,0 +1,6 @@ +pub mod amap; +pub mod datetime; +pub mod fetch_page; +pub mod memos; +pub mod weather; +pub mod web_search; diff --git a/src/tools/weather.rs b/src/tools/builtins/weather.rs similarity index 91% rename from src/tools/weather.rs rename to src/tools/builtins/weather.rs index d291941..0265d2e 100644 --- a/src/tools/weather.rs +++ b/src/tools/builtins/weather.rs @@ -30,8 +30,8 @@ fn base64url(data: &[u8]) -> String { } fn generate_jwt() -> Result { - let key_file = - std::env::var("QWEATHER_JWT_PRIVATE_KEY_FILE").unwrap_or_else(|_| "qweather/ed25519-private.pem".into()); + let key_file = std::env::var("QWEATHER_JWT_PRIVATE_KEY_FILE") + .unwrap_or_else(|_| "qweather/ed25519-private.pem".into()); let pem = std::fs::read_to_string(&key_file) .map_err(|e| format!("读取密钥文件 {} 失败: {}", key_file, e))?; @@ -72,10 +72,16 @@ fn generate_jwt() -> Result { // Sign let message = format!("{}.{}", header_b64, payload_b64); - let signature = signing_key.try_sign(message.as_bytes()) + let signature = signing_key + .try_sign(message.as_bytes()) .map_err(|e| format!("JWT 签名失败: {e}"))?; - Ok(format!("{}.{}.{}", header_b64, payload_b64, base64url(&signature.to_bytes()))) + Ok(format!( + "{}.{}.{}", + header_b64, + payload_b64, + base64url(&signature.to_bytes()) + )) } // ═══════════════════════════════════════════════ @@ -226,8 +232,8 @@ struct QWeatherClient { impl QWeatherClient { async fn new() -> Result { let jwt = generate_jwt()?; - let api_host = std::env::var("QWEATHER_API_HOST") - .unwrap_or_else(|_| "api.qweather.com".into()); + let api_host = + std::env::var("QWEATHER_API_HOST").unwrap_or_else(|_| "api.qweather.com".into()); Ok(Self { http: Client::builder() @@ -310,11 +316,7 @@ impl QWeatherClient { } /// 每日预报 - async fn daily_forecast( - &self, - city_id: &str, - days: u32, - ) -> Result { + async fn daily_forecast(&self, city_id: &str, days: u32) -> Result { let d = match days { 0..=3 => "3d", 4..=7 => "7d", @@ -418,7 +420,9 @@ impl QWeatherClient { pub fn spec() -> SkillSpec { SkillSpec { name: "query_weather".into(), - description: "查询指定城市的天气信息,包括实时天气、未来几天预报、逐小时预报。参数 location 为城市名".into(), + description: + "查询指定城市的天气信息,包括实时天气、未来几天预报、逐小时预报。参数 location 为城市名" + .into(), risk_level: RiskLevel::Low, parameters: serde_json::json!({ "type": "object", @@ -437,7 +441,7 @@ fn extract_params(params: &serde_json::Value) -> (&str, u32, u32) { let location = params .get("location") .and_then(|v| v.as_str()) - .unwrap_or("北京"); + .unwrap_or("合肥"); let days = params .get("days") @@ -489,18 +493,32 @@ fn extract_city_from_prompt(prompt: &str) -> &str { // 跳过前缀 let mut i = 0; - for prefix in &["帮我查一下", "帮我看看", "帮我查", "帮我", "查一下", "查询", "看看"] { + for prefix in &[ + "帮我查一下", + "帮我看看", + "帮我查", + "帮我", + "查一下", + "查询", + "看看", + ] { if prompt.starts_with(prefix) { i = prefix.chars().count(); - while i < chars.len() && chars[i] as u32 <= 0x4e00 { i += 1; } + while i < chars.len() && chars[i] as u32 <= 0x4e00 { + i += 1; + } break; } } // 跳过非中文 - while i < chars.len() && chars[i] as u32 <= 0x4e00 { i += 1; } + while i < chars.len() && chars[i] as u32 <= 0x4e00 { + i += 1; + } // 取至多4个连续中文 let start = i; - while i < chars.len() && chars[i] as u32 > 0x4e00 && i - start < 4 { i += 1; } + while i < chars.len() && chars[i] as u32 > 0x4e00 && i - start < 4 { + i += 1; + } if i - start >= 2 { &prompt[byte_pos(start)..byte_pos(i)] } else { @@ -509,23 +527,14 @@ fn extract_city_from_prompt(prompt: &str) -> &str { } pub async fn execute(params: serde_json::Value) -> SkillResult { - // call_capability 透传但没有 structured 参数 → 告知正确格式 - if params.get("location").is_none() - && params.get("name").and_then(|v| v.as_str()) == Some("query_weather") - { - let spec = spec(); - return SkillResult::error(format!( - "参数格式错误。请将 query_weather 所需的参数直接放在 call_capability 的 prompt 字段中,格式:\ - prompt: {{\"location\": \"城市名\", \"days\": 3, \"hours\": 0}}。\ - 工具说明:{}。参数 schema:{}", - spec.description, - serde_json::to_string(&spec.parameters).unwrap_or_default() - )); - } - let (location, days, hours) = extract_params(¶ms); - tracing::info!("🌤️ 查询天气: location={}, days={}, hours={}", location, days, hours); + tracing::info!( + "🌤️ 查询天气: location={}, days={}, hours={}", + location, + days, + hours + ); let client = match QWeatherClient::new().await { Ok(c) => c, @@ -541,9 +550,16 @@ pub async fn execute(params: serde_json::Value) -> SkillResult { // 并发查询 let now_fut = client.weather_now(&city_id); let daily_fut = client.daily_forecast(&city_id, days); - let hourly_fut = client.hourly_forecast(&city_id, hours); - let (now_result, daily_result, hourly_result) = tokio::join!(now_fut, daily_fut, hourly_fut); + // hours == 0 时跳过逐时预报,节省 API 调用 + let (now_result, daily_result, hourly_result) = if hours > 0 { + let hourly_fut = client.hourly_forecast(&city_id, hours); + let (n, d, h) = tokio::join!(now_fut, daily_fut, hourly_fut); + (n, d, h) + } else { + let (n, d) = tokio::join!(now_fut, daily_fut); + (n, d, Err("hours=0, skipped".into())) + }; let now = match now_result { Ok(n) => n, @@ -626,7 +642,10 @@ mod tests { #[test] fn test_extract_city_from_prompt() { // 取前缀后的连续中文给 GeoAPI 模糊搜索(多取不影响结果) - assert_eq!(extract_city_from_prompt("查询合肥未来7天的天气预报"), "合肥未来"); + assert_eq!( + extract_city_from_prompt("查询合肥未来7天的天气预报"), + "合肥未来" + ); assert_eq!(extract_city_from_prompt("北京今天天气"), "北京今天"); assert_eq!(extract_city_from_prompt("上海"), "上海"); assert_eq!(extract_city_from_prompt("帮我查一下广州"), "广州"); diff --git a/src/tools/builtins/web_search.rs b/src/tools/builtins/web_search.rs new file mode 100644 index 0000000..e940696 --- /dev/null +++ b/src/tools/builtins/web_search.rs @@ -0,0 +1,362 @@ +//! Web 搜索工具 — Tavily Search API +//! +//! API: POST https://api.tavily.com/search +//! 认证: Authorization: Bearer +//! 文档: https://docs.tavily.com/documentation/api-reference/endpoint/search + +use reqwest::Client as HttpClient; +use serde::{Deserialize, Serialize}; + +use super::super::types::{RiskLevel, SkillResult, SkillSpec}; + +// ─── 请求 ─── + +#[derive(Debug, Serialize)] +struct SearchRequest { + query: String, + + #[serde(skip_serializing_if = "Option::is_none")] + search_depth: Option, // basic | advanced | fast | ultra-fast + + #[serde(skip_serializing_if = "Option::is_none")] + chunks_per_source: Option, // 1-3, 仅 search_depth=advanced + + #[serde(skip_serializing_if = "Option::is_none")] + max_results: Option, // 0-20 + + #[serde(skip_serializing_if = "Option::is_none")] + topic: Option, // general | news | finance + + #[serde(skip_serializing_if = "Option::is_none")] + time_range: Option, // day | week | month | year | d | w | m | y + + #[serde(skip_serializing_if = "Option::is_none")] + start_date: Option, // YYYY-MM-DD + + #[serde(skip_serializing_if = "Option::is_none")] + end_date: Option, // YYYY-MM-DD + + #[serde(skip_serializing_if = "Option::is_none")] + include_answer: Option, // false | true(=basic) | "basic" | "advanced" + + #[serde(skip_serializing_if = "Option::is_none")] + include_raw_content: Option, // false | true(=markdown) | "markdown" | "text" + + #[serde(skip_serializing_if = "Option::is_none")] + include_images: Option, + + #[serde(skip_serializing_if = "Option::is_none")] + include_image_descriptions: Option, + + #[serde(skip_serializing_if = "Option::is_none")] + include_favicon: Option, + + #[serde(skip_serializing_if = "Option::is_none")] + include_domains: Option>, + + #[serde(skip_serializing_if = "Option::is_none")] + exclude_domains: Option>, + + #[serde(skip_serializing_if = "Option::is_none")] + country: Option, + + #[serde(skip_serializing_if = "Option::is_none")] + auto_parameters: Option, // 默认 false + + #[serde(skip_serializing_if = "Option::is_none")] + exact_match: Option, + + #[serde(skip_serializing_if = "Option::is_none")] + include_usage: Option, + + #[serde(skip_serializing_if = "Option::is_none")] + safe_search: Option, // Enterprise only +} + +// ─── 响应 ─── + +#[derive(Debug, Deserialize)] +struct SearchResponse { + query: String, + + #[serde(default)] + answer: Option, + + #[serde(default)] + images: Vec, + + #[serde(default)] + results: Vec, + + #[serde(default)] + response_time: serde_json::Value, + + #[serde(default)] + auto_parameters: Option, + + #[serde(default)] + usage: Option, + + #[serde(default)] + request_id: Option, +} + +#[derive(Debug, Deserialize)] +struct SearchResult { + title: String, + url: String, + content: String, + #[serde(default)] + score: Option, + #[serde(default)] + raw_content: Option, + #[serde(default)] + favicon: Option, + #[serde(default)] + images: Vec, + #[serde(default)] + published_date: Option, +} + +#[derive(Debug, Deserialize)] +struct UsageInfo { + #[serde(default)] + credits: Option, +} + +// ─── 工具 spec ─── + +pub fn spec() -> SkillSpec { + SkillSpec { + name: "web_search".into(), + description: + "联网搜索最新信息。通过 Tavily API 搜索互联网,获取实时、准确的结果。\ + 适用于需要最新资讯、事实查询的场景。".into(), + risk_level: RiskLevel::Low, + parameters: serde_json::json!({ + "type": "object", + "properties": { + "query": {"type": "string", "description": "搜索关键词"}, + "max_results": {"type": "integer", "description": "最大结果数 1-20(默认5)"}, + "include_answer": {"type": "boolean", "description": "是否包含 AI 摘要(默认 true)"}, + "search_depth": {"type": "string", "enum": ["basic","advanced","fast","ultra-fast"], + "description": "basic 均衡 / advanced 深度(2 credits) / fast 快速 / ultra-fast 极速"}, + "topic": {"type": "string", "enum": ["general","news","finance"], + "description": "general 通用 / news 新闻 / finance 财经"}, + "time_range": {"type": "string", "enum": ["day","week","month","year"], + "description": "按发布时间过滤(仅 topic=news/finance 时有效)"}, + "start_date": {"type": "string", "description": "起始日期 YYYY-MM-DD"}, + "end_date": {"type": "string", "description": "结束日期 YYYY-MM-DD"}, + "include_domains": {"type": "array", "items": {"type": "string"}, + "description": "限定搜索域名列表"}, + "exclude_domains": {"type": "array", "items": {"type": "string"}, + "description": "排除搜索域名列表"}, + "country": {"type": "string", "description": "优先指定国家的搜索结果(仅 topic=general)"} + }, + "required": ["query"] + }), + timeout_secs: 30, + } +} + +// ─── 执行 ─── + +pub async fn execute(params: serde_json::Value) -> SkillResult { + let api_key = match std::env::var("TAVILY_API_KEY") { + Ok(k) if !k.is_empty() => k, + _ => return SkillResult::error("未配置 TAVILY_API_KEY 环境变量"), + }; + + // —— 解析参数 —— + + let query = params + .get("query") + .and_then(|v| v.as_str()) + .or_else(|| params.get("prompt").and_then(|v| v.as_str())) + .unwrap_or(""); + if query.is_empty() { + return SkillResult::error("请提供 query 参数"); + } + + let max_results = params + .get("max_results") + .and_then(|v| v.as_u64()) + .map(|n| n.min(20) as u32); + + let include_answer = params + .get("include_answer") + .and_then(|v| v.as_bool()) + .unwrap_or(true); + + let search_depth = params + .get("search_depth") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + + let topic = params + .get("topic") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + + let time_range = params + .get("time_range") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + + let start_date = params + .get("start_date") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + + let end_date = params + .get("end_date") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + + let include_domains = params.get("include_domains").and_then(|v| { + v.as_array().map(|arr| { + arr.iter() + .filter_map(|item| item.as_str().map(String::from)) + .collect() + }) + }); + + let exclude_domains = params.get("exclude_domains").and_then(|v| { + v.as_array().map(|arr| { + arr.iter() + .filter_map(|item| item.as_str().map(String::from)) + .collect() + }) + }); + + let country = params + .get("country") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + + // days 参数转换为 time_range(便捷兼容) + let time_range = time_range.or_else(|| { + params + .get("days") + .and_then(|v| v.as_u64()) + .map(|d| format!("{}d", d)) + }); + + // —— 发起请求 —— + + let http = match HttpClient::builder() + .timeout(std::time::Duration::from_secs(30)) + .build() + { + Ok(c) => c, + Err(e) => return SkillResult::error(format!("创建 HTTP 客户端失败: {}", e)), + }; + + let body = SearchRequest { + query: query.to_string(), + search_depth, + chunks_per_source: None, + max_results: Some(max_results.unwrap_or(5)), + topic, + time_range, + start_date, + end_date, + include_answer: if include_answer { + Some(serde_json::Value::Bool(true)) + } else { + None + }, + include_raw_content: None, + include_images: None, + include_image_descriptions: None, + include_favicon: None, + include_domains, + exclude_domains, + country, + auto_parameters: None, + exact_match: None, + include_usage: None, + safe_search: None, + }; + + match http + .post("https://api.tavily.com/search") + .header("Authorization", format!("Bearer {}", api_key)) + .header("Content-Type", "application/json") + .json(&body) + .send() + .await + { + Ok(resp) => { + if !resp.status().is_success() { + let status = resp.status(); + let text = resp.text().await.unwrap_or_default(); + return SkillResult::error(format!( + "Tavily API 错误 ({}): {:.300}", + status, text + )); + } + + let resp_text = resp.text().await.unwrap_or_default(); + match serde_json::from_str::(&resp_text) { + Ok(data) => { + let mut output = String::new(); + + // AI 摘要 + if let Some(ref answer) = data.answer { + if !answer.is_empty() { + output.push_str(&format!("📝 {}\n\n", answer)); + } + } + + // 搜索结果 + if data.results.is_empty() { + output.push_str("未找到相关结果。"); + } else { + for (i, r) in data.results.iter().enumerate() { + output.push_str(&format!( + "{}. **{}**\n 🔗 {}\n {}\n\n", + i + 1, + r.title, + r.url, + r.content + )); + } + } + + // 用量信息 + let credits = data + .usage + .as_ref() + .and_then(|u| u.credits) + .map(|c| format!("{} credits", c)) + .unwrap_or_default(); + let rt = data.response_time; + output.push_str(&format!( + "⏱ {} | {} 条结果{}", + rt, + data.results.len(), + if credits.is_empty() { + String::new() + } else { + format!(" | {}", credits) + } + )); + + SkillResult::ok(output) + } + Err(e) => SkillResult::error(format!( + "解析 Tavily 响应失败: {} — 原始: {:.300}", + e, resp_text + )), + } + } + Err(e) => { + if e.is_timeout() { + SkillResult::error("Tavily 搜索超时") + } else { + SkillResult::error(format!("Tavily 请求失败: {}", e)) + } + } + } +} diff --git a/src/tools/mod.rs b/src/tools/mod.rs index a950927..a57e61f 100644 --- a/src/tools/mod.rs +++ b/src/tools/mod.rs @@ -1,4 +1,148 @@ pub mod approval; pub mod builtin; +pub mod builtins; pub mod types; -pub mod weather; + +/// 从 call_capability 的 `{name, prompt}` JSON 中提取目标工具的真实参数。 +/// +/// 逻辑: +/// - 如果 prompt 是 JSON 字符串 → parse 后合并 +/// - 如果 prompt 是 JSON 对象 → 直接合并 +/// - 顶层字段(非 name/prompt)也合并进去 +pub fn unpack_call_params(cp: &serde_json::Value) -> serde_json::Value { + let mut params = serde_json::Map::new(); + + // 1. 提取 prompt 字段中的参数 + if let Some(prompt) = cp.get("prompt") { + match prompt { + serde_json::Value::String(s) => { + if let Ok(v) = serde_json::from_str::(s) { + if let Some(obj) = v.as_object() { + params.extend(obj.clone()); + } + } + } + serde_json::Value::Object(obj) => { + params.extend(obj.clone()); + } + _ => {} + } + } + + // 2. 合并顶层字段(排除 name、prompt) + if let Some(obj) = cp.as_object() { + for (k, v) in obj { + if k == "name" || k == "prompt" { + continue; + } + params.insert(k.clone(), v.clone()); + } + } + + serde_json::Value::Object(params) +} + +#[cfg(test)] +mod tests { + #[test] + fn test_unpack_call_params() { + use super::unpack_call_params; + + // prompt 是 JSON 字符串 + let cp = serde_json::json!({"name": "amap_poi_search", "prompt": "{\"keywords\":\"肯德基\",\"city\":\"北京\"}"}); + let result = unpack_call_params(&cp); + assert_eq!(result["keywords"], "肯德基"); + assert_eq!(result["city"], "北京"); + assert!(result.get("name").is_none()); + assert!(result.get("prompt").is_none()); + + // prompt 是 JSON 对象 + let cp = serde_json::json!({"name": "query_weather", "prompt": {"location": "上海", "days": 7}}); + let result = unpack_call_params(&cp); + assert_eq!(result["location"], "上海"); + assert_eq!(result["days"], 7); + + // 仅有 name,无 prompt + let cp = serde_json::json!({"name": "get_current_datetime"}); + let result = unpack_call_params(&cp); + assert!(result.as_object().unwrap().is_empty()); + + // 顶层字段合并 + let cp = serde_json::json!({"name": "call_capability", "prompt": "{\"a\": 1}", "extra": 42}); + let result = unpack_call_params(&cp); + assert_eq!(result["a"], 1); + assert_eq!(result["extra"], 42); + } +} + +use types::SkillSpec; + +/// 构建工具调用指南(含优先级说明) +/// 返回给 LLM 的工具列表 + 使用建议 +pub fn build_capability_guide(specs: &[SkillSpec]) -> String { + let mut lines = vec![ + "📋 可用工具及调用优先级指南:".to_string(), + String::new(), + "## 🗺️ 高德地图工具(amap_*)— 旅游、路线、地点查询优先使用:".to_string(), + ]; + + // 先列出 amap 工具 + let amap_specs: Vec<&SkillSpec> = specs.iter().filter(|s| s.name.starts_with("amap_")).collect(); + if amap_specs.is_empty() { + lines.push(" (未配置高德 API Key,暂不可用)".to_string()); + } else { + for spec in &amap_specs { + format_spec(&mut lines, spec); + } + } + lines.push(String::new()); + lines.push(" ▸ 旅游规划、路径规划、地点/周边搜索、地理编码 → 优先使用 amap_* 系列工具".to_string()); + + lines.push(String::new()); + lines.push("## 🌤️ 天气工具 — 天气相关查询:".to_string()); + for spec in specs.iter().filter(|s| s.name == "query_weather") { + format_spec(&mut lines, spec); + } + lines.push(String::new()); + lines.push(" ▸ 任何天气相关问题(实时天气、天气预报、逐小时预报等)→ 使用 query_weather".to_string()); + + lines.push(String::new()); + lines.push("## 🔍 其他工具:".to_string()); + for spec in specs.iter().filter(|s| !s.name.starts_with("amap_") && s.name != "query_weather") { + format_spec(&mut lines, spec); + } + + lines.push(String::new()); + lines.push("## ⚡ 工具选择优先级:".to_string()); + lines.push(" 1. 旅游/路线/地点/地址查询 → amap_* 系列".to_string()); + lines.push(" 2. 天气查询 → query_weather".to_string()); + lines.push(" 3. 联网搜索/网页抓取 → web_search / fetch_page".to_string()); + lines.push(" 4. 以上都不匹配的需求 → web_search 搜索最新信息".to_string()); + lines.push(String::new()); + lines.push("📌 注意:每次调用工具前请确认它确实是当前最合适的工具。".to_string()); + + lines.join("\n") +} + +fn format_spec(lines: &mut Vec, spec: &SkillSpec) { + let risk = if spec.risk_level == crate::tools::types::RiskLevel::High { + " ⚠️需确认" + } else { + "" + }; + lines.push(format!( + "\n🔹 {} {} — {}", + spec.name, risk, spec.description + )); + if !spec + .parameters + .get("properties") + .and_then(|v| v.as_object()) + .map_or(true, |o| o.is_empty()) + { + lines.push(format!( + " 参数: {}", + serde_json::to_string(&spec.parameters).unwrap_or_default() + )); + } +} diff --git a/src/tools/types.rs b/src/tools/types.rs index b6de0f9..954a0d1 100644 --- a/src/tools/types.rs +++ b/src/tools/types.rs @@ -12,11 +12,7 @@ pub enum RiskLevel { High, } -impl RiskLevel { - pub fn is_high(&self) -> bool { - matches!(self, RiskLevel::High) - } -} +impl RiskLevel {} // ─── 工具元数据 ─── diff --git a/src/wechat/client.rs b/src/wechat/client.rs index cee0d4e..42c1084 100644 --- a/src/wechat/client.rs +++ b/src/wechat/client.rs @@ -74,10 +74,7 @@ impl WeChatClient { qrcode: &str, base_url: &str, ) -> Result { - let url = format!( - "{}/ilink/bot/get_qrcode_status?qrcode={}", - base_url, qrcode - ); + let url = format!("{}/ilink/bot/get_qrcode_status?qrcode={}", base_url, qrcode); let resp_result = self .http @@ -89,7 +86,9 @@ impl WeChatClient { match resp_result { Ok(response) => { - let text = response.text().await + let text = response + .text() + .await .map_err(|e| format!("读取响应体失败: {}", e))?; serde_json::from_str::(&text) .map_err(|e| format!("解析状态响应失败: {}", e)) @@ -134,7 +133,8 @@ impl WeChatClient { // 生成二维码终端显示 if let Ok(code) = qrcode::QrCode::new(qrcode_img_content.as_bytes()) { - let svg = code.render::() + let svg = code + .render::() .dark_color(qrcode::render::unicode::Dense1x2::Dark) .light_color(qrcode::render::unicode::Dense1x2::Light) .build(); @@ -177,7 +177,9 @@ impl WeChatClient { .ilink_bot_id .ok_or_else(|| "登录成功但未返回 ilink_bot_id".to_string())?; let user_id = status.ilink_user_id.unwrap_or_default(); - let base_url = status.baseurl.unwrap_or_else(|| Self::DEFAULT_BASE_URL.to_string()); + let base_url = status + .baseurl + .unwrap_or_else(|| Self::DEFAULT_BASE_URL.to_string()); // 保存到内存状态 *self.token.lock().await = token.clone(); @@ -220,7 +222,10 @@ impl WeChatClient { .map_err(|e| format!("解析 notifyStart 响应失败: {}", e))?; if _resp.ret != 0 { - return Err(format!("notifyStart 失败: ret={} errmsg={:?}", _resp.ret, _resp.errmsg)); + return Err(format!( + "notifyStart 失败: ret={} errmsg={:?}", + _resp.ret, _resp.errmsg + )); } info!("监听器已注册"); @@ -257,15 +262,19 @@ impl WeChatClient { let url = format!("{}/ilink/bot/getupdates", self.base_url); let token = self.token.lock().await.clone(); - let resp: GetUpdatesResp = match self - .post_json(&url, &body, Some(&token)) - .await - { + let resp: GetUpdatesResp = match self.post_json(&url, &body, Some(&token)).await { Ok(resp) => { - let resp: GetUpdatesResp = resp.json().await.map_err(|e| format!("解析 getUpdates 响应失败: {}", e))?; + let resp: GetUpdatesResp = resp + .json() + .await + .map_err(|e| format!("解析 getUpdates 响应失败: {}", e))?; if resp.ret != 0 || resp.errcode.unwrap_or(0) != 0 { - tracing::warn!("getUpdates 返回错误: ret={} errcode={:?} errmsg={:?}", - resp.ret, resp.errcode, resp.errmsg); + tracing::warn!( + "getUpdates 返回错误: ret={} errcode={:?} errmsg={:?}", + resp.ret, + resp.errcode, + resp.errmsg + ); } resp } @@ -337,16 +346,6 @@ impl WeChatClient { } } - /// 获取当前 token(用于持久化) - pub async fn get_token(&self) -> String { - self.token.lock().await.clone() - } - - /// 获取当前 account_id(用于持久化) - pub async fn get_account_id(&self) -> String { - self.account_id.lock().await.clone() - } - /// 获取当前 updates_buf(用于持久化) pub async fn get_updates_buf(&self) -> String { self.updates_buf.lock().await.clone() @@ -375,14 +374,12 @@ impl WeChatClient { headers.insert( "iLink-App-ClientVersion", - reqwest::header::HeaderValue::from_str(&Self::CLIENT_VERSION.to_string()) - .unwrap(), + reqwest::header::HeaderValue::from_str(&Self::CLIENT_VERSION.to_string()).unwrap(), ); // X-WECHAT-UIN: random uint32 -> decimal -> base64 let uin = rand::random::(); - let uin_b64 = base64::engine::general_purpose::STANDARD - .encode(uin.to_string()); + let uin_b64 = base64::engine::general_purpose::STANDARD.encode(uin.to_string()); headers.insert( "X-WECHAT-UIN", reqwest::header::HeaderValue::from_str(&uin_b64).unwrap(), @@ -396,8 +393,7 @@ impl WeChatClient { if !token.is_empty() { headers.insert( reqwest::header::AUTHORIZATION, - reqwest::header::HeaderValue::from_str(&format!("Bearer {}", token)) - .unwrap(), + reqwest::header::HeaderValue::from_str(&format!("Bearer {}", token)).unwrap(), ); headers.insert( "AuthorizationType", @@ -413,7 +409,8 @@ impl WeChatClient { body: &impl serde::Serialize, token: Option<&str>, ) -> Result { - let body_str = serde_json::to_string(body).map_err(|e| format!("序列化请求体失败: {}", e))?; + let body_str = + serde_json::to_string(body).map_err(|e| format!("序列化请求体失败: {}", e))?; debug!("POST {} body_len={}", url, body_str.len()); @@ -426,8 +423,7 @@ impl WeChatClient { // 设置 Content-Length headers.insert( "Content-Length", - reqwest::header::HeaderValue::from_str(&body_str.len().to_string()) - .unwrap(), + reqwest::header::HeaderValue::from_str(&body_str.len().to_string()).unwrap(), ); self.http diff --git a/src/worker.rs b/src/worker.rs new file mode 100644 index 0000000..f7b4059 --- /dev/null +++ b/src/worker.rs @@ -0,0 +1,291 @@ +//! Worker — 无状态执行进程 +//! +//! 每条消息 spawn 一个 worker: +//! 1. 连接 daemon 的 Unix Domain Socket +//! 2. 读取 TaskFrame (用户消息 + 上下文 + 环境变量) +//! 3. 注入环境变量,构建 Conversation,执行 LLM 对话 + 工具调用 +//! 4. 逐帧输出结果 (usage, reply, bye) +//! 5. 高风险工具遇审批 → 发送 RequestApproval → 退出,由 daemon 重新 spawn + +use crate::ipc::{OutputFrame, TaskFrame, recv_task, send_frame}; +use crate::llm::{Conversation, ConversationConfig, DEFAULT_SYSTEM_PROMPT}; +use crate::tools::types::SkillResult; +use std::collections::HashMap; +use std::sync::Arc; +use tokio::net::UnixStream; +use tracing::info; + +/// Worker 与工具执行器之间的共享状态 +struct WorkerShared { + /// 待发送的审批请求 (tool, reason) + pending_approval: tokio::sync::Mutex>, + /// 本次对话中已由用户批准的工具名(避免重复审批) + approved_tool: tokio::sync::Mutex>, + /// 本地记忆缓存 + memories: tokio::sync::Mutex>>, + /// 本次写入的新记忆(退出前发送 DbWrite 帧同步到 daemon) + new_memories: tokio::sync::Mutex>, + /// 预载的摘要 + summaries: Vec, + /// 用户 ID + user_id: String, +} + +/// Worker 入口: 连接 daemon,处理一条消息后退出 +pub async fn run(sock_path: &str) -> Result<(), String> { + let mut stream = UnixStream::connect(sock_path) + .await + .map_err(|e| format!("连接 daemon 失败 ({}): {}", sock_path, e))?; + + info!("Worker 已连接 daemon"); + + // 1. 读 task + let task: TaskFrame = recv_task(&mut stream).await?; + info!( + "Worker 收到任务: user={} msg={:.60}", + task.user_id, task.msg.text + ); + + // 2. 注入环境变量 + for (k, v) in &task.env { + unsafe { + std::env::set_var(k, v); + } + } + + // 3. 构建 ConversationConfig + let system_prompt = if task.system_prompt.is_empty() { + DEFAULT_SYSTEM_PROMPT.to_string() + } else { + task.system_prompt.clone() + }; + let model = if task.model.is_empty() { + std::env::var("DEEPSEEK_MODEL").unwrap_or_else(|_| "deepseek-v4-flash".to_string()) + } else { + task.model.clone() + }; + + let cfg = ConversationConfig { + system_prompt, + model, + tools: Some(task.tools.clone()), + ..Default::default() + }; + + // 4. 构建 Conversation 并加载上下文 + let mut conv = + Conversation::new(cfg).map_err(|e| format!("初始化 Conversation 失败: {}", e))?; + info!("Worker LLM: {}/{}", conv.provider_name(), conv.model()); + + { + let session = conv.session(); + let mut s = session.lock().await; + s.load_from_history(&task.history); + s.current_user_id = task.user_id.clone(); + } + + // 5. 共享状态 + let approved_tool = task.approved_tool.clone().filter(|t| !t.is_empty()); + let shared = Arc::new(WorkerShared { + pending_approval: tokio::sync::Mutex::new(None), + approved_tool: tokio::sync::Mutex::new(approved_tool), + memories: tokio::sync::Mutex::new(HashMap::from([( + task.user_id.clone(), + task.memories.clone(), + )])), + new_memories: tokio::sync::Mutex::new(Vec::new()), + summaries: task.summaries.clone(), + user_id: task.user_id.clone(), + }); + + // 6. 构建工具执行器 + let executor = build_worker_executor(conv.session(), shared.clone()); + conv.set_tool_executor(executor); + + // 7. 执行 LLM 对话(带审批检测) + let (reply, _used_tools, usage) = conv + .chat_with_tools(task.msg.text.clone()) + .await + .map_err(|e| format!("LLM 对话失败: {}", e))?; + + // 8. 检查是否有审批请求(由工具执行器设置) + let approval_req = shared.pending_approval.lock().await.take(); + if let Some((tool, reason)) = approval_req { + info!("Worker 发送审批请求: tool={}", tool); + let _ = send_frame(&mut stream, &OutputFrame::RequestApproval { tool, reason }).await; + let _ = send_frame(&mut stream, &OutputFrame::Bye).await; + return Ok(()); + } + + // 9. 发送结果帧 + // 先同步新写入的记忆到 daemon + let new_mems = shared.new_memories.lock().await; + for content in new_mems.iter() { + let _ = send_frame( + &mut stream, + &OutputFrame::DbWrite { + table: "user_memories".into(), + row: serde_json::json!({"content": content}), + }, + ) + .await; + } + drop(new_mems); + + if let Some(u) = &usage { + let _ = send_frame( + &mut stream, + &OutputFrame::UsageReport { + model: conv.model().to_string(), + provider: conv.provider_name().to_string(), + prompt_tokens: u.prompt_tokens, + completion_tokens: u.completion_tokens, + total_tokens: u.total_tokens, + cache_hit_tokens: u.prompt_cache_hit_tokens, + cache_miss_tokens: u.prompt_cache_miss_tokens, + user_id: task.user_id.clone(), + }, + ) + .await; + } + + let _ = send_frame( + &mut stream, + &OutputFrame::Reply { + text: reply.clone(), + }, + ) + .await; + + let _ = send_frame(&mut stream, &OutputFrame::Bye).await; + + info!( + "Worker 完成: user={} reply_len={}", + task.user_id, + reply.len() + ); + Ok(()) +} + +/// 构建 Worker 侧的工具执行器(无 DB 连接) +fn build_worker_executor( + _session: Arc>, + shared: Arc, +) -> crate::llm::conversation::ToolExecutor { + Arc::new(move |name: &str, args_json: &str| { + let shared = shared.clone(); + let n = name.to_string(); + let args = args_json.to_string(); + + Box::pin(async move { + match n.as_str() { + "read_memories" => { + let cache = shared.memories.lock().await; + let mems = cache.get(&shared.user_id); + match mems { + Some(m) if !m.is_empty() => { + let lines: Vec = m + .iter() + .enumerate() + .map(|(i, v)| format!("{}. {}", i + 1, v)) + .collect(); + Ok(lines.join("\n")) + } + _ => Ok("暂无长期记忆".to_string()), + } + } + "write_memory" => { + let params: serde_json::Value = serde_json::from_str(&args).unwrap_or_default(); + let content = params.get("content").and_then(|v| v.as_str()).unwrap_or(""); + if content.is_empty() { + return Ok("请提供 content 参数".to_string()); + } + shared + .memories + .lock() + .await + .entry(shared.user_id.clone()) + .or_default() + .push(content.to_string()); + shared.new_memories.lock().await.push(content.to_string()); + Ok(format!("已添加长期记忆: {}", content)) + } + "read_summaries" => { + if shared.summaries.is_empty() { + Ok("暂无历史摘要".to_string()) + } else { + Ok(shared + .summaries + .iter() + .enumerate() + .map(|(i, s)| format!("{}. {}", i + 1, s)) + .collect::>() + .join("\n---\n")) + } + } + "query_capabilities" => { + let specs = crate::tools::builtin::BuiltinRegistry::specs(); + Ok(crate::tools::build_capability_guide(&specs)) + } + // call_capability 或直接调用 + _ => { + let (target_name, target_args) = if n == "call_capability" { + let cp: serde_json::Value = serde_json::from_str(&args).unwrap_or_default(); + let name = cp + .get("name") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let unpacked = crate::tools::unpack_call_params(&cp); + (name, serde_json::to_string(&unpacked).unwrap_or_default()) + } else { + (n.clone(), args.clone()) + }; + + if target_name.is_empty() { + return Ok(format!("未知工具: {}", n)); + } + + // 高风险工具 → 请求审批(除非已被用户批准) + if crate::tools::builtin::BuiltinRegistry::is_high_risk(&target_name) { + let is_approved = shared.approved_tool.lock().await.as_deref() == Some(&target_name); + if !is_approved { + let reason = format!("LLM 请求执行高风险工具: {}", target_name); + *shared.pending_approval.lock().await = Some((target_name.clone(), reason)); + return Ok(SkillResult::rejected(&target_name).output); + } + // 批准过 → 执行一次后清除(同一对话内同一工具不会再被拦截) + *shared.approved_tool.lock().await = None; + } + + // 执行内置工具 + if let Some(result) = + crate::tools::builtin::BuiltinRegistry::execute(&target_name, &target_args) + .await + { + if target_name == "manage_memos" { + let cp: serde_json::Value = + serde_json::from_str(&target_args).unwrap_or_default(); + if cp.get("action").and_then(|v| v.as_str()) == Some("add") { + if let Some(content) = cp.get("content").and_then(|v| v.as_str()) { + shared + .memories + .lock() + .await + .entry(shared.user_id.clone()) + .or_default() + .push(content.to_string()); + shared.new_memories.lock().await.push(content.to_string()); + } + } + } + return Ok(result.output); + } + + Ok(format!("未知工具: {}", target_name)) + } + } + }) + }) +} +