diff --git a/Cargo.toml b/Cargo.toml index 0d4f59c..1b85fcd 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -2,45 +2,56 @@ name = "iAs" version = "0.1.0" edition = "2024" +description = "微信 AI 智能助手 — 通过 DeepSeek LLM 实现 AI 自动回复" [dependencies] -# 异步运行时 -tokio = { version = "1.0", features = ["full"] } -# HTTP 客户端 -reqwest = { version = "0.12", features = ["json", "stream", "gzip"] } -# 序列化 -serde = { version = "1.0", features = ["derive"] } -serde_json = "1.0" -# 环境变量 -dotenvy = "0.15" -# CLI 参数解析 -clap = { version = "4", features = ["derive"] } -# 终端输入 -rustyline = "14.0" -# UUID -uuid = { version = "1", features = ["v4", "serde"] } -# 时间 -chrono = { version = "0.4", features = ["serde"] } -# 加密 -base64 = "0.22" -sha2 = "0.10" -hex = "0.4" -rand = "0.9" -# 日志 -tracing = "0.1" -tracing-subscriber = { version = "0.3", features = ["env-filter"] } -# 二维码终端显示 -qrcode = "0.14" -image = "0.25" -# 异步 trait -async-trait = "0.1" -# 流式处理 -futures-util = "0.3" -sqlx = { version = "0.8.6", features = ["runtime-tokio-native-tls", "postgres", "uuid", "chrono", "derive", "migrate"] } -tracing-appender = "0.2.5" -dialoguer = "0.12.0" -console = "0.16.3" -dirs = "6.0.0" -# Ed25519 JWT 签名（天气 API） -ed25519-dalek = { version = "2", features = ["pem"] } -scraper = "0.27.0" +# ── 异步运行时 ── +tokio = { version = "1.0", features = ["full"] } # 异步运行时核心 +# ── HTTP 客户端 ── +reqwest = { version = "0.12", features = ["json", "stream", "gzip"] } # HTTP 请求（微信 API + DeepSeek API + 工具调用） +# ── 序列化 ── +serde = { version = "1.0", features = ["derive"] } # 序列化/反序列化 +serde_json = "1.0" # JSON 处理 +# ── 环境变量加载 ── +dotenvy = "0.15" # .env 文件加载 +# ── CLI 参数解析 ── +clap = { version = "4", features = ["derive"] } # 命令行参数解析 +# ── 终端交互 ── +rustyline = "14.0" # REPL 行编辑 +dialoguer = "0.12.0" # 终端交互对话框 +console = "0.16.3" # 终端控制 +# ── UUID ── +uuid = { version = "1", features = ["v4", "serde"] } # 唯一标识符 +# ── 时间处理 ── +chrono = { version = "0.4", features = ["serde"] } # 日期时间 +# ── 加密与签名 ── +base64 = "0.22" # Base64 编码 +sha2 = "0.10" # SHA-256 哈希（审批码） +hex = "0.4" # 十六进制 +rand = "0.9" # 随机数 +ed25519-dalek = { version = "2", features = ["pem"] } # Ed25519 JWT 签名（和风天气 API） +# ── 日志 ── +tracing = "0.1" # 结构化日志 +tracing-subscriber = { version = "0.3", features = ["env-filter"] } # 日志输出 +tracing-appender = "0.2.5" # 文件日志（日滚） +# ── 二维码 ── +qrcode = "0.14" # 二维码生成 +image = "0.25" # 图片处理 +# ── 异步 trait ── +async-trait = "0.1" # 异步 trait 支持 +# ── 流式处理 ── +futures-util = "0.3" # 流式处理工具 +# ── 数据库 ── +sqlx = { version = "0.8.6", features = ["runtime-tokio-native-tls", "postgres", "uuid", "chrono", "derive", "migrate"] } # PostgreSQL +# ── 文件系统 ── +dirs = "6.0.0" # 系统目录路径 +# ── HTML 解析 ── +scraper = "0.27.0" # HTML 解析（fetch_page 工具） + +[[bin]] +name = "ias" +path = "src/main.rs" + +[[bin]] +name = "ias-auth" +path = "src/bin/ias-auth.rs" diff --git a/src/bin/ias-auth.rs b/src/bin/ias-auth.rs index d09b3f2..538224d 100644 --- a/src/bin/ias-auth.rs +++ b/src/bin/ias-auth.rs @@ -1,10 +1,18 @@ -//! ias-auth — 独立 Token 生成器 -//! -//! 专供 JSON 工具定义通过 `${cmd:ias-auth gen-qweather-jwt}` 调用。 -//! 编译快、无运行时依赖，独立于主程序 `ias`。 -//! -//! 用法: -//! ias-auth gen-qweather-jwt 生成和风天气 Ed25519 JWT Token +/// ## ias-auth —— 独立 Token 生成器（辅助二进制） +/// +/// 这是一个独立的二进制文件，不依赖主程序的任何模块。 +/// 专供 JSON 工具定义通过 `${cmd:ias-auth gen-qweather-jwt}` 调用， +/// 用于生成和风天气 API 的 Ed25519 JWT Token。 +/// +/// ### 设计原因 +/// 独立编译有两个好处： +/// 1. 编译速度快（没有主程序的庞大依赖链） +/// 2. 在 JSON 工具定义中通过 shell 命令调用的开销很小 +/// +/// ### 用法 +/// ```bash +/// ias-auth gen-qweather-jwt +/// ``` use std::process::exit; diff --git a/src/channel/mod.rs b/src/channel/mod.rs index c7f249c..fb4e487 100644 --- a/src/channel/mod.rs +++ b/src/channel/mod.rs @@ -1,11 +1,33 @@ -//! 渠道 / 通道基础类型 +//! ## 渠道/通道基础类型 //! -//! 定义消息的来源/目标渠道标识，以及从外部（微信等）收到的原始消息。 +//! 定义消息的来源和目标渠道标识，以及从外部轮询收到的原始消息。 +//! +//! ### 设计意图 +//! +//! `ChannelId` 是本项目中所有消息路由的核心标识符。 +//! 它由 `platform`（平台名，如 "wechat"）和 `user_id`（用户 ID）两部分组成， +//! 确保不同平台上的用户可以被唯一区分。 +//! +//! 未来如果要接入 Telegram、Discord 等新平台，只需要新增 `ChannelId` 的构造方法， +//! 消息队列的路由逻辑无需改动。 use serde::{Deserialize, Serialize}; use std::fmt; -/// 渠道标识 —— 区分消息来自哪个平台和哪个用户 +/// ## 渠道标识 —— 区分消息来自哪个平台和哪个用户 +/// +/// 这是消息队列和 daemon 中所有消息路由的基础。 +/// 每个 `PipelineMessage` 都携带一个 `ChannelId`，用来追踪消息的来龙去脉。 +/// +/// ### 字段 +/// * `platform` — 平台名称，如 `"wechat"`、`"telegram"`、`"console"` +/// * `user_id` — 在该平台上的用户标识 +/// +/// ### 示例 +/// ``` +/// let ch = ChannelId::wechat("wxid_123"); +/// assert_eq!(ch.to_string(), "wechat:wxid_123"); +/// ``` #[derive(Debug, Clone, Hash, Eq, PartialEq, Serialize, Deserialize)] pub struct ChannelId { /// 平台名称，如 "wechat"、"telegram"、"console" @@ -27,7 +49,15 @@ impl fmt::Display for ChannelId { } } -/// 从外部轮询收到的原始消息 +/// ## 从外部轮询收到的原始消息 +/// +/// 在 daemon 的主循环中，从 WeChat API 获取到的消息会被转换为这个结构， +/// 然后封装成 `PipelineMessage` 投递到消息队列。 +/// +/// * `channel` — 来源渠道（平台 + 用户） +/// * `text` — 消息文本 +/// * `message_id` — 消息 ID +/// * `context_token` — 微信的上下文令牌（用于多轮对话的上下文恢复） #[derive(Debug, Clone, Serialize, Deserialize)] #[allow(dead_code)] pub struct IncomingMessage { diff --git a/src/cli.rs b/src/cli.rs index da0ce6c..823e92e 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -1,15 +1,17 @@ use clap::{Parser, Subcommand}; -/// iAs — 微信 AI 智能助手 +/// ## iAs —— 微信 AI 智能助手 /// /// 支持扫码登录、长轮询收发消息、AI 自动回复（DeepSeek）、 /// 内置工具（天气/搜索/备忘录/日期时间等）、Token 用量统计。 /// -/// 快速开始: -/// ias login 扫码登录微信 -/// ias listen --llm 启动 AI 自动回复 -/// ias tool weather 北京 查询天气 -/// ias tool search "最新新闻" 联网搜索 +/// ### 快速开始 +/// ```bash +/// 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 { diff --git a/src/context/builder.rs b/src/context/builder.rs index 2b4fee8..f8b166e 100644 --- a/src/context/builder.rs +++ b/src/context/builder.rs @@ -4,13 +4,24 @@ use crate::llm::types::Message; use std::sync::Arc; use tokio::sync::Mutex; -/// 估算消息的 token 数 +/// ## 估算消息的 token 数 +/// +/// 粗略估算一条消息的 token 消耗，用于 token budget 管理。 +/// 公式：`文本 token + 8`（8 是消息格式的开销）。 +/// +/// 注意这只是估算，不是精确值，但在实践中足够做预算判断。 pub fn estimate_tokens(msg: &Message) -> u32 { let text = msg.content.as_str(); estimate_text_tokens(text) + 8 } -/// 估算纯文本的 token 数 +/// ## 估算纯文本的 token 数 +/// +/// 中文和英文的 token 消耗不同： +/// * 中文字符 ≈ 1.5 tokens/字（CJK 字符更密集） +/// * 英文字符 ≈ 0.25 tokens/字符 +/// +/// 这是一个粗略估算，用于在调用 API 之前判断上下文是否超 budget。 pub fn estimate_text_tokens(text: &str) -> u32 { if text.is_empty() { return 0; @@ -27,9 +38,17 @@ fn is_cjk(c: char) -> bool { || ('\u{f900}'..='\u{faff}').contains(&c) } -/// 构建给 LLM 的消息数组（带 token budget 管理） +/// ## 构建给 LLM 的消息数组（带 token budget 管理） /// -/// 返回消息数组和是否需要摘要的信息 +/// 这是 LLM 调用前的关键步骤 —— 从 `ChatSession` 中提取合适的消息， +/// 在 token 预算内尽可能保留更多上下文。 +/// +/// ### 构建策略 +/// 1. **系统提示词** — 始终包含 +/// 2. **溢出摘要** — 如果有最新的 overflow 摘要，加入（仍在 budget 内时） +/// 3. **近期消息** — 从 checkpoint 开始，从最新消息倒序添加，直到接近 budget +/// 4. **保护最近用户消息** — 确保最后一条用户消息始终在上下文中 +/// （即使超过 budget，也在 +2000 范围内强制包含） pub async fn build_context(session: &Arc>, system_prompt: &str) -> Vec { let s = session.lock().await; @@ -92,7 +111,18 @@ pub async fn build_context(session: &Arc>, system_prompt: &st messages } -/// 触发溢出摘要：压缩 checkpoint 到当前位置之间的新消息 +/// ## 触发溢出摘要 —— 压缩 checkpoint 到当前位置之间的新消息 +/// +/// 当上下文 token 超 budget 时调用。把 checkpoint 之后的消息交给 LLM 生成摘要， +/// 然后推进 checkpoint，让旧消息可以被释放（或截断）。 +/// +/// ### 流程 +/// 1. 获取 checkpoint 到当前末尾的消息（最多 40 条） +/// 2. 用 LLM summarizer 生成摘要（不可用时回退到简单截断） +/// 3. 将摘要存入 `summaries` 列表 +/// 4. 持久化到数据库 +/// 5. 推进 checkpoint 到当前末尾 +/// 6. 如果 checkpoint 超过 200，截断最早的消息（保留最近 100 条用于调试） pub async fn trigger_overflow_summary( session: &Arc>, summarizer: Option<&Summarizer>, @@ -153,7 +183,13 @@ pub async fn trigger_overflow_summary( true } -/// 触发空闲摘要（不注入，保存到 summaries 供 LLM 工具查询） +/// ## 触发空闲超时摘要（不注入上下文，只存档） +/// +/// 当距上条用户消息超过 12 小时时触发。 +/// 与溢出摘要不同，空闲摘要**不会**作为 system prompt 注入到下一次对话中， +/// 而是只保存到数据库，供 `read_summaries` 工具查询。 +/// +/// 生成摘要后会清空所有消息，checkpoint 重置为 0。 pub async fn trigger_idle_summary( session: &Arc>, summarizer: Option<&Summarizer>, diff --git a/src/context/tools.rs b/src/context/tools.rs index ff01866..146ea30 100644 --- a/src/context/tools.rs +++ b/src/context/tools.rs @@ -3,7 +3,20 @@ use std::collections::HashMap; use std::sync::Arc; use tokio::sync::Mutex; -/// 长期记忆管理器（内存 + PostgreSQL 双写，按用户隔离） +/// ## 长期记忆管理器（内存 + PostgreSQL 双写） +/// +/// 为每个用户维护独立的长期记忆列表，支持读写操作。 +/// +/// ### 存储策略 +/// * **内存缓存** — `HashMap>`，按 user_id 隔离 +/// * **数据库持久化** — 可选的 `PgPool`，写入时同时写到 PostgreSQL +/// +/// ### 数据流 +/// ```text +/// 启动时 load(user_id) → 从数据库加载到缓存 +/// 对话中 read_for() → 从缓存读取，返回 "1. ...\n2. ..." 格式 +/// 对话中 write_for() → 写入缓存 + 写入数据库 +/// ``` pub struct MemoryStore { pool: Option>, cache: Arc>>>, @@ -61,7 +74,10 @@ impl MemoryStore { } } -/// read_summaries 工具：读取历史摘要（内存 + 数据库） +/// ## `read_summaries` 工具 —— 读取历史摘要（内存 + 数据库） +/// +/// 返回当前 session 中的摘要列表 + 数据库中存档的摘要。 +/// 去重：数据库中的摘要如果与 session 中的摘要前 30 个字符相同，则跳过。 pub async fn read_summaries( session: &Arc>, ) -> String { diff --git a/src/context/types.rs b/src/context/types.rs index 911beba..ffc5a14 100644 --- a/src/context/types.rs +++ b/src/context/types.rs @@ -4,7 +4,11 @@ use serde::{Deserialize, Serialize}; use std::sync::Arc; use uuid::Uuid; -/// 摘要原因 +/// ## 摘要原因 +/// +/// 区分两种触发摘要的场景： +/// * `Overflow` — 上下文 token 总数超过预算（token budget），触发压缩 +/// * `Timeout` — 距上条用户消息超过 12 小时空闲，触发会话总结 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] pub enum SummaryReason { /// 上下文 token 超 budget 触发 @@ -13,7 +17,12 @@ pub enum SummaryReason { Timeout, } -/// 一条摘要记录 +/// ## 一条摘要记录 +/// +/// * `checkpoint` — 摘要对应的消息位置，此位置之前的消息已被压缩成摘要 +/// * `text` — 摘要文本 +/// * `reason` — 生成原因（Overflow / Timeout） +/// * `created_at` — 生成时间 #[derive(Debug, Clone)] pub struct SummaryEntry { /// 摘要对应的消息位置（checkpoint） @@ -26,7 +35,23 @@ pub struct SummaryEntry { pub created_at: DateTime, } -/// 聊天会话状态 +/// ## 聊天会话状态（核心数据结构） +/// +/// 这是整个项目中最重要的数据结构之一 —— 它维护了一次 LLM 对话的完整状态。 +/// +/// ### 关键设计 +/// * **消息历史** (`messages`) — 从 checkpoint 之后的消息保持完整 +/// * **摘要机制** (`summaries` + `checkpoint`) — checkpoint 之前的消息被压缩成摘要， +/// 用 token budget 确保上下文不会超出 LLM 的限制 +/// * **Token 预算** (`token_budget`) — 默认 28000，留 4000 给回复 +/// * **空闲超时** (`idle_timeout_secs`) — 默认 12 小时，超时后生成空闲摘要 +/// * **会话 ID** (`session_id`) — 每次新对话分配一个 Uuid，用于追踪工具调用往返 +/// +/// ### 消息生命周期 +/// 1. 用户消息到达 → `add_user()` → 存入 `messages` +/// 2. AI 回复 → `add_assistant()` → 存入 `messages` +/// 3. Token 超 budget → `trigger_overflow_summary()` → 压缩 checkpoint 前的消息 +/// 4. 12h 空闲 → `trigger_idle_summary()` → 生成摘要，清空消息 #[derive(Debug, Clone)] pub struct ChatSession { pub user_id: String, diff --git a/src/daemon.rs b/src/daemon.rs index 7f0ba07..5be3873 100644 --- a/src/daemon.rs +++ b/src/daemon.rs @@ -1,7 +1,43 @@ -//! Daemon — 常驻进程 +//! ## Daemon —— 常驻守护进程（核心循环） //! -//! 持有 WeChat 长连接 + 数据库。使用消息队列架构将消息处理分解为 -//! LLM / Tool / Send 三个消费者（tokio task），通过 QueueRunner 路由。 +//! 这是整个项目的核心入口循环。它的职责： +//! +//! 1. **持有 WeChat 长连接** — 通过 `WeChatClient` 长轮询接收微信消息 +//! 2. **持有数据库连接** — 可选 PostgreSQL，回退到文件存储 +//! 3. **管理审批流程** — `ApprovalManager` 处理高风险工具的确认码审批 +//! 4. **管理长期记忆** — `MemoryStore` 提供按用户隔离的记忆读写 +//! +//! ### 消息处理架构（三消费者 + 公平队列） +//! +//! daemon 将消息处理拆分为 3 个独立的 tokio async task（消费者）， +//! 通过 `MessageQueue` + `QueueRunner` 做公平轮转路由： +//! +//! ```text +//! receive_messages → enqueue PipelineMessage +//! │ +//! ▼ +//! ┌──────────────────┐ +//! │ MessageQueue │ ← 公平轮转（A1, B1, A2, A3） +//! └────────┬─────────┘ +//! │ dequeue +//! ▼ +//! ┌──────────────────┐ +//! │ QueueRunner │ ← 按 MessageKind 路由 +//! └──┬──────┬──────┬─┘ +//! │ │ │ +//! ▼ ▼ ▼ +//! LLM Tool Send +//! Cons. Cons. Cons. +//! ``` +//! +//! - **LLM Consumer** — 调用 DeepSeek API 做对话 + 工具循环 +//! - **Tool Consumer** — 执行内置工具，高风险工具先走审批 +//! - **Send Consumer** — 将 LLM 回复通过 WeChat API 发回给用户 +//! +//! ### 与旧架构的关系 +//! +//! 旧架构使用 `daemon + worker` 分离模式（UDS IPC），worker 已废弃。 +//! 当前是统一的单进程 tokio task 架构。 use crate::channel::ChannelId; use crate::context::MemoryStore; @@ -19,17 +55,29 @@ use tokio::sync::Mutex; use tracing::{error, info, warn}; use uuid::Uuid; -/// Daemon 上下文 —— 消费者共享的所有状态 +/// ## Daemon 上下文 —— 消费者之间共享的全部状态 +/// +/// 这个结构体被所有三个消费者（LLM / Tool / Send）通过 `Arc` 共享访问。 +/// 每个字段都只读（或内部有锁），所以不需要额外的 `Mutex` 包装。 struct DaemonCtx { + /// 数据库句柄（可选，没有时走文件存储） db: Option>, + /// WeChat API 客户端（内部用 Arc 保护 token/updates_buf） client: WeChatClient, + /// 本机器人的微信账号 ID account_id: String, + /// 审批管理器 —— 处理高风险工具的用户确认码流程 approval: Arc, + /// 长期记忆管理器 —— 按用户隔离的记忆读写 memory_store: Arc, + /// 注册给 LLM 的工具列表（JSON 格式，符合 function calling 规范） tools_list: Vec, + /// 系统提示词，覆盖默认值 system_prompt: String, + /// 使用的模型名称（如 `deepseek-v4-flash`） model: String, - /// 待审批: user_id → (原始文本, context_token, 工具名, tool_args, 创建时间) + /// 待审批的上下文：`user_id → (原始文本, context_token, 工具名, 工具参数, 创建时间)` + /// 用户在聊天中发送确认码时，从这里面查找匹配项 approval_ctx: Mutex, String, String, Instant)>>, } diff --git a/src/db/mod.rs b/src/db/mod.rs index e915a5f..886d7ff 100644 --- a/src/db/mod.rs +++ b/src/db/mod.rs @@ -1,10 +1,24 @@ +//! ## 数据库管理 +//! +//! ### 职责 +//! 1. 从 `DATABASE_URL` 环境变量连接 PostgreSQL +//! 2. 自动运行 `./migrations/` 下的 SQL 迁移 +//! 3. 提供连接池引用给各模块使用 +//! +//! ### 回退策略 +//! 数据库是可选的。如果数据库不可用，系统会优雅地回退到文件存储 +//! （`StateManager`），核心功能仍然可用。 + pub mod models; -use sqlx::PgPool; use sqlx::postgres::PgPoolOptions; +use sqlx::PgPool; use std::time::Duration; -/// 数据库管理器 +/// ## 数据库管理器 +/// +/// 封装了 PostgreSQL 连接池的创建和迁移管理。 +/// 通过 `Database::connect()` 创建，所有数据库操作都在 `models` 模块中。 pub struct Database { pool: PgPool, } diff --git a/src/db/models.rs b/src/db/models.rs index fa1c998..37a4346 100644 --- a/src/db/models.rs +++ b/src/db/models.rs @@ -1,3 +1,15 @@ +//! ## 数据库查询（models） +//! +//! 所有数据库操作集中在此模块，包括： +//! +//! ### 表操作 +//! * `app_state` — 认证状态的持久化（key-value 存储） +//! * `chat_records` — 聊天记录（inbound/outbound） +//! * `llm_usage` — Token 用量记录和统计 +//! * `session_summaries` — 会话摘要存档 +//! * `user_memories` — 用户长期记忆 +//! * `pending_approvals` — 待审批记录 + use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; use sqlx::PgPool; diff --git a/src/ipc.rs b/src/ipc.rs index 965dc45..eb349e0 100644 --- a/src/ipc.rs +++ b/src/ipc.rs @@ -1,10 +1,19 @@ -//! IPC 帧格式 — Unix Domain Socket 长度前缀帧 +//! ## IPC 帧格式 —— Unix Domain Socket 长度前缀帧（已废弃） //! -//! 帧格式: [4 字节 u32 BE: payload_len][N 字节: JSON payload] +//! 这是旧架构（daemon + worker 分离模式）的进程间通信协议。 +//! 新架构已改为统一的单进程 tokio task，此模块保留仅为编译参考。 //! -//! 消息类型定义: -//! Daemon → Worker: Task (携带用户消息、历史、记忆、摘要、环境变量) -//! Worker → Daemon: Chunk / Reply / NeedApproval / DbWrite / UsageReport / Bye +//! ### 帧格式 +//! ```text +//! [4 字节 u32 大端: payload_len][N 字节: JSON payload] +//! ``` +//! +//! ### 消息类型 +//! * Daemon → Worker: `TaskFrame`（用户消息 + 上下文 + 环境变量） +//! * Worker → Daemon: `OutputFrame`（Chunk / Reply / NeedApproval / DbWrite / UsageReport / Bye） +//! +//! ### 最大帧大小 +//! 10MB（防止恶意超大帧导致 OOM） use serde::{Deserialize, Serialize}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; diff --git a/src/llm/conversation.rs b/src/llm/conversation.rs index 852943c..87d7272 100644 --- a/src/llm/conversation.rs +++ b/src/llm/conversation.rs @@ -8,18 +8,39 @@ use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; use tokio::sync::Mutex; +/// ## 工具执行器 —— LLM 调用的工具回调 +/// +/// 这是一个函数别名：`Fn(工具名, JSON参数) → Future>`。 +/// +/// 当 LLM 在对话中发起工具调用时，Conversation 会调用这个 executor 来执行工具。 +/// +/// ### 两种执行模式 +/// 1. **同步模式** — executor 直接执行并返回结果字符串 +/// 2. **异步模式** — executor 将工具调用入队到消息队列，返回 `ASYNC_MARKER + 工具名` +/// 表示"工具已异步入队，结果稍后通过 ToolResult 消息回喂" pub type ToolExecutor = Arc< dyn Fn(&str, &str) -> Pin> + Send>> + Send + Sync, >; -/// 摘要生成器：接收要摘要的消息文本 → 返回 LLM 生成的摘要 +/// ## 摘要生成器 —— 用于生成对话摘要的回调 +/// +/// Conversation 在两种情况下会触发摘要： +/// 1. **溢出摘要** — 消息 token 总数接近 budget 时，压缩早期消息 +/// 2. **空闲超时摘要** — 距上条用户消息超过 12 小时时，总结本次会话 pub type Summarizer = Arc< dyn Fn(String) -> Pin> + Send>> + Send + Sync, >; -/// `chat_with_tools` 的返回结果 +/// ## `chat_with_tools` 的返回结果 +/// +/// * `reply` — LLM 最终的文本回复 +/// * `used_tools` — 本轮对话中是否调用了工具 +/// * `usage` — Token 用量统计（可选） +/// * `has_pending_async` — 是否有异步工具调用仍在等待结果 +/// 如果为 true，调用方应该保留 Conversation 实例，等待 ToolResult 回来后再调用 +/// `resume_tool_loop()` 恢复对话 #[derive(Debug)] pub struct ChatResult { pub reply: String, @@ -32,6 +53,27 @@ pub struct ChatResult { /// 异步工具标记 —— 工具执行器返回此前缀表示工具已异步入队 pub const ASYNC_MARKER: &str = "__ASYNC_PENDING__"; +/// ## Conversation —— LLM 对话管理器（核心类） +/// +/// 这是整个项目中"与 LLM 对话"的核心抽象，负责： +/// +/// ### 职责 +/// 1. **管理对话状态** — 持有 `ChatSession`（消息历史、摘要、token 预算） +/// 2. **流式对话** — 调用 provider 发起流式请求，逐 chunk 消费 +/// 3. **工具循环** — 自动处理 LLM 发起的工具调用，循环直到无工具或达到轮次上限 +/// 4. **异步工具** — 支持工具通过消息队列异步执行，结果回来后恢复对话 +/// 5. **上下文管理** — token 预算检查、溢出摘要、空闲超时摘要 +/// +/// ### 工具循环流程 +/// ```text +/// chat_with_tools(user_msg) +/// → run_tool_loop() +/// → 1. build_context() 构建上下文（含摘要、system prompt） +/// → 2. provider.chat_stream() 发起流式请求 +/// → 3. 消费流，收集文本 + tool_calls +/// → 4. 有 tool_calls？→ 执行 → 回到步骤 1（最多 5 轮） +/// → 5. 无 tool_calls？→ 返回最终回复 +/// ``` pub struct Conversation { config: ConversationConfig, provider: BoxedProvider, @@ -42,6 +84,11 @@ pub struct Conversation { } impl Conversation { + /// ## 创建 Conversation + /// + /// 1. 通过 `create_provider()` 工厂从配置创建 LLM 提供商 + /// 2. 初始化空的 `ChatSession`（消息历史在收到第一条用户消息后按需加载） + /// 3. tool_executor 初始为 None，需要调用 `set_tool_executor` 设置 pub fn new(config: ConversationConfig) -> Result { let provider = create_provider(&config)?; Ok(Self { @@ -77,7 +124,12 @@ impl Conversation { self.create_summarizer() } - /// 创建 LLM 摘要生成器（用当前 provider 做非流式调用） + /// ## 创建 LLM 摘要生成器 + /// + /// 用当前 LLM provider 创建一个非流式调用的摘要器。 + /// 把消息列表拼成 prompt 发给 LLM，返回生成的摘要文本。 + /// + /// 摘要 prompt 为："你是一个对话摘要助手，用中文输出简洁摘要。" pub fn create_summarizer(&self) -> Summarizer { let provider = self.provider.clone(); Arc::new(move |prompt: String| { @@ -110,7 +162,17 @@ impl Conversation { }) } - /// 单次对话（无工具） + /// ## 单次对话（无工具调用） + /// + /// 简单的流式对话：发送用户消息 → 获取 AI 回复。 + /// 不会触发工具循环，适用于不需要 function calling 的场景。 + /// + /// ### 流程 + /// 1. 检查是否空闲超时（距上条消息 > 12h），如果是则触发摘要 + /// 2. 将用户消息添加到 session + /// 3. 构建上下文（system prompt + 摘要 + 近期消息） + /// 4. 发起流式请求 + /// 5. 返回 `ChatHandle`（可逐 chunk 消费） pub async fn chat(&self, user_message: String) -> Result { // 空闲超时检查 { @@ -134,7 +196,18 @@ impl Conversation { }) } - /// 对话 + 工具循环（带上下文管理） + /// ## 对话 + 工具循环（主入口） + /// + /// 这是最常用的方法：发送用户消息，AI 可以选择回复文本或调用工具。 + /// 如果调用了工具，工具执行结果会被自动回喂给 AI，AI 再决定下一步操作。 + /// + /// ### 完整流程 + /// 1. 空闲超时检查 → 触发摘要 + /// 2. 添加用户消息到 session + /// 3. 进入 `run_tool_loop()` + /// - 构建上下文 → 调 LLM → 解析流 + /// - 有 tool_calls？→ 执行工具 → 继续循环（最多 5 轮） + /// - 无 tool_calls？→ 返回最终回复 pub async fn chat_with_tools(&self, user_message: String) -> Result { // ── 空闲超时检查 ── { diff --git a/src/llm/deepseek.rs b/src/llm/deepseek.rs index e503ad0..fc22d38 100644 --- a/src/llm/deepseek.rs +++ b/src/llm/deepseek.rs @@ -7,6 +7,27 @@ use reqwest::Client as HttpClient; use std::collections::BTreeMap; use tokio::sync::mpsc; +/// ## DeepSeek LLM 提供商 +/// +/// 实现了 `LlmProvider` trait，通过 DeepSeek API 提供流式对话能力。 +/// +/// ### 配置 +/// 从环境变量读取： +/// * `DEEPSEEK_API_KEY` — API 密钥（必需） +/// * `DEEPSEEK_BASE_URL` — API 地址（可选，默认 `https://api.deepseek.com/v1`） +/// * `DEEPSEEK_MODEL` — 模型名（可选，默认 `deepseek-v4-flash`） +/// +/// ### 请求格式 +/// 发送 POST 到 `{base_url}/chat/completions`，标准 OpenAI-compatible 格式： +/// ```json +/// { +/// "model": "deepseek-v4-flash", +/// "messages": [...], +/// "stream": true, +/// "tools": [...], // 可选，function calling +/// "thinking": {...} // 可选，DeepSeek 推理模式 +/// } +/// ``` pub struct DeepSeekProvider { http: HttpClient, api_key: String, diff --git a/src/llm/provider.rs b/src/llm/provider.rs index 74cecca..d74362c 100644 --- a/src/llm/provider.rs +++ b/src/llm/provider.rs @@ -8,7 +8,16 @@ use tokio::sync::mpsc; pub type StreamReceiver = mpsc::Receiver; pub type StreamSender = mpsc::Sender; -/// LLM 提供商抽象 +/// ## LLM 提供商抽象（trait） +/// +/// 这是 LLM 层的核心抽象接口，定义了"一个 LLM 提供商需要提供什么能力"。 +/// 目前唯一的实现是 `DeepSeekProvider`，但通过这个 trait 可以轻松添加 +/// 其他提供商（如 OpenAI、Claude 等）。 +/// +/// ### 关键设计 +/// * 使用 `async_trait` 宏支持异步 trait 方法 +/// * 返回 `StreamReceiver`（mpsc::Receiver），调用方可以按 chunk 消费流式结果 +/// * `Send + Sync` 保证可以在 tokio task 间安全共享 #[async_trait] pub trait LlmProvider: Send + Sync { /// 提供商名称 @@ -26,14 +35,23 @@ pub trait LlmProvider: Send + Sync { pub type BoxedProvider = Arc; -/// 从配置创建恰当的提供商 +/// ## `create_provider` — 工厂函数：从配置创建恰当的 LLM 提供商 +/// +/// 目前只返回 DeepSeek，后续可以在这里根据配置 switch 到不同提供商。 pub fn create_provider(_config: &ConversationConfig) -> Result { Ok(Arc::new(super::deepseek::DeepSeekProvider::new()?)) } // ─── 内部：SSE 解析工具 ─── -/// 流式块解析结果 +/// ## 流式块解析结果（内部枚举） +/// +/// 用于从 SSE 的 JSON delta 中解析出不同种类的块。 +/// * `Text` — 普通文本 delta +/// * `Reasoning` — 思考内容（DeepSeek reasoning_content） +/// * `ToolCallDelta` — 工具调用的增量信息（可能需要跨多个 chunk 拼接） +/// * `FinishReason` — 结束原因（stop / tool_calls 等） +/// * `Usage` — 用量统计（通常在最后一个 chunk 中） #[allow(dead_code)] pub(crate) enum ParsedChunk { Text(String), @@ -48,7 +66,20 @@ pub(crate) enum ParsedChunk { Usage(super::types::Usage), } -/// 从 JSON body 中解析 DeepSeek/OpenAI 流式 delta +/// ## 从 JSON body 中解析 DeepSeek/OpenAI 流式 delta +/// +/// ### 输入 +/// 一行 SSE 数据，格式为 `data: {json}`。 +/// +/// ### 处理逻辑 +/// 1. 跳过非 `data: ` 前缀的行和 `[DONE]` 行 +/// 2. 解析 JSON 为 `ChunkResponse`（内嵌多个 `choices`） +/// 3. 提取 usage（可能在任意一个 chunk 中） +/// 4. 对每个 choice，按优先级提取： +/// - tool_calls delta → `ToolCallDelta` +/// - reasoning_content → `Reasoning` +/// - content → `Text` +/// - finish_reason → `FinishReason` pub(crate) fn parse_chat_chunk(line: &str) -> Option { if !line.starts_with("data: ") { return None; diff --git a/src/llm/types.rs b/src/llm/types.rs index 43009fa..efaf532 100644 --- a/src/llm/types.rs +++ b/src/llm/types.rs @@ -1,7 +1,12 @@ use serde::{Deserialize, Serialize}; -// ─── 角色 ─── - +/// ## 对话角色（System / User / Assistant / Tool） +/// +/// 符合 OpenAI/DeepSeek 的 Chat Completion API 标准角色定义。 +/// * `System` — 系统提示词，设定 AI 的行为和限制 +/// * `User` — 用户消息 +/// * `Assistant` — AI 回复（可以是纯文本或带 tool_calls） +/// * `Tool` — 工具执行结果的回馈 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] #[serde(rename_all = "lowercase")] pub enum Role { @@ -11,8 +16,17 @@ pub enum Role { Tool, } -// ─── 消息 ─── - +/// ## LLM 消息 —— 符合 OpenAI/DeepSeek Chat Completion API 格式 +/// +/// 这是整个 LLM 交互的核心类型。`Vec` 构成一次对话的全部上下文， +/// 被序列化为 JSON 发送给 DeepSeek API。 +/// +/// ### 字段 +/// * `role` — 消息角色（system/user/assistant/tool） +/// * `content` — 消息内容（文本） +/// * `tool_call_id` — 如果是 tool 角色的结果消息，关联到原始 tool_call 的 ID +/// * `name` — 工具名（tool 角色时使用） +/// * `tool_calls` — assistant 发起工具调用时携带的调用列表 #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Message { pub role: Role, @@ -48,8 +62,12 @@ impl Message { } } -// ─── 工具调用 ─── - +/// ## 工具调用 —— LLM 发起的一个工具调用请求 +/// +/// 当 LLM 认为需要调用工具时，会在回复中携带这个结构。 +/// * `id` — 唯一标识这个调用（用于 tool result 回馈时匹配） +/// * `call_type` — 固定为 "function" +/// * `function` — 函数名 + JSON 格式的参数 #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ToolCall { pub id: String, @@ -64,8 +82,13 @@ pub struct ToolFunctionCall { pub arguments: String, // JSON string } -// ─── 流式响应块 ─── - +/// ## 流式响应块 —— 从 DeepSeek API 的 SSE 流中解析出的每个 delta +/// +/// ### 变体 +/// * `Text` — 文本片段 delta（逐步拼接得到完整回复） +/// * `Reasoning` — 思考/推理内容（DeepSeek 的 thinking 机制） +/// * `Done` — 完成信号，携带完整文本 + 可选的工具调用 + 用量统计 +/// * `Error` — 错误 #[derive(Debug, Clone)] #[allow(dead_code)] pub enum StreamChunk { @@ -85,8 +108,14 @@ pub enum StreamChunk { Error(String), } -// ─── Token 用量 ─── - +/// ## Token 用量统计 +/// +/// DeepSeek API 在流式响应的最后一个 chunk 中会附带用量信息。 +/// * `prompt_tokens` — 提示词消耗的 token 数 +/// * `completion_tokens` — 生成文本消耗的 token 数 +/// * `total_tokens` = prompt + completion +/// * `prompt_cache_hit_tokens` — 命中上下文的缓存 token 数（省钱） +/// * `prompt_cache_miss_tokens` — 未命中缓存的 token 数 #[derive(Debug, Clone, Serialize, Deserialize, Default)] pub struct Usage { pub prompt_tokens: u32, @@ -98,8 +127,15 @@ pub struct Usage { pub prompt_cache_miss_tokens: u32, } -// ─── 对话配置 ─── - +/// ## 对话配置 —— 创建 Conversation 时的参数 +/// +/// ### 关键字段 +/// * `system_prompt` — 系统提示词，决定 AI 的行为方式 +/// * `model` — 模型名称（如 `deepseek-v4-flash`） +/// * `temperature` — 生成随机性（0.0-1.0，越高越有创意） +/// * `max_tokens` — 最大生成 token 数 +/// * `thinking` — 是否启用 DeepSeek 的 thinking 模式 +/// * `tools` — OpenAI 格式的工具定义列表（function calling 用） #[derive(Debug, Clone)] pub struct ConversationConfig { pub system_prompt: String, diff --git a/src/logger.rs b/src/logger.rs index 6398855..8f2b328 100644 --- a/src/logger.rs +++ b/src/logger.rs @@ -3,6 +3,14 @@ use tracing_subscriber::layer::SubscriberExt; use tracing_subscriber::util::SubscriberInitExt; use tracing_subscriber::EnvFilter; +/// ## 日志初始化 +/// +/// ### 两种模式 +/// * **终端模式** — 彩色输出到 stderr，适合开发调试 +/// * **文件模式** (`with_file=true`) — 同时输出到终端和日滚文件 +/// `~/.ias/logs/ias.log.YYYY-MM-DD` +/// +/// 日志级别通过 `RUST_LOG` 环境变量控制（默认 `info`）。 pub fn init_logger(log_dir: Option<&str>, with_file: bool) { let env_filter = EnvFilter::try_from_default_env() .unwrap_or_else(|_| EnvFilter::new("info")); diff --git a/src/main.rs b/src/main.rs index 90cb96c..9bf5311 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,3 +1,34 @@ +//! ## iAs —— 微信 AI 智能助手（入口） +//! +//! ### 系统架构概览 +//! +//! ```text +//! CLI (main.rs → clap) +//! ├── ias login → 扫码登录，保存 token 到 DB 或文件 +//! ├── ias listen --llm → Daemon 模式（主运行模式） +//! │ └── daemon.rs: 长轮询 → MessageQueue → 3 个 tokio 消费者 +//! │ ├── LLM Consumer → Conversation.chat_with_tools() → DeepSeek API +//! │ ├── Tool Consumer → BuiltinRegistry.execute() / ApprovalManager +//! │ └── Send Consumer → WeChatClient.send_message() +//! ├── ias service → 等同 listen --llm 但启用文件日志 +//! ├── ias daemon → 守护进程入口（等同 listen --llm） +//! ├── ias tool → 直接调用内置工具（无需登录） +//! └── ias usage → Token 消耗统计查询 +//! ``` +//! +//! ### 数据流 +//! ```text +//! 微信消息 → 存 chat_records → 加载历史 + 记忆 → 构建 LLM 上下文 +//! → Conversation.chat_with_tools() → LLM 回复或调工具 +//! → 高风险工具走审批 → 工具结果回喂 LLM → 最终回复发送 +//! ``` +//! +//! ### 关键设计决策 +//! 1. **单进程架构** — 旧版 daemon+worker 分离模式已废弃，统一用 tokio task +//! 2. **公平队列** — MessageQueue 按用户轮转，防止一个用户刷屏饿死其他人 +//! 3. **可选数据库** — PostgreSQL 不可用时回退到文件存储 +//! 4. **两层元工具** — LLM 通过 query_capabilities / call_capability 间接调用工具 + mod channel; mod cli; mod context; diff --git a/src/queue/message_queue.rs b/src/queue/message_queue.rs index e3f8d23..f9b6ec2 100644 --- a/src/queue/message_queue.rs +++ b/src/queue/message_queue.rs @@ -1,14 +1,18 @@ -//! ## 多渠道消息队列 — MessageQueue +//! ## 多渠道公平轮转消息队列 (MessageQueue) //! -//! # 设计目标 +//! ### 设计目标 //! -//! 多用户场景下，保证公平性： -//! - A 发了 3 条消息、B 发了 1 条 → 处理顺序 A1, B1, A2, A3 -//! - 不会因为 A 消息多而饿死 B +//! 多用户场景下保证**公平性**：如果一个用户连续发了很多消息，不会让其他用户饿死。 +//! 出队顺序是**轮转（round-robin）**的： +//! - A 发了 3 条、B 发了 1 条 → 处理顺序 A1, B1, A2, A3 //! -//! 每个消息携带 user_id + correlation_id，保证回复能追溯到正确的用户。 +//! ### 消息关联追踪 //! -//! # 架构位置 +//! 每条消息有一个 `correlation_id`（Uuid），用来追踪一次用户请求的完整生命周期： +//! 用户消息 → LLM 处理 → 工具调用 → 工具结果回喂 → 最终回复发送。 +//! 这样即使多用户并发，回复也不会串。 +//! +//! ### 架构位置 //! //! ```text //! receive_messages → enqueue PipelineMessage diff --git a/src/queue/runner.rs b/src/queue/runner.rs index d952540..48567b3 100644 --- a/src/queue/runner.rs +++ b/src/queue/runner.rs @@ -1,9 +1,16 @@ -//! QueueRunner — 消息队列运行时 +//! ## QueueRunner —— 消息队列运行时（路由调度器） //! -//! 从公平轮转 MessageQueue 中出队消息，按 MessageKind 路由到对应的 +//! 从公平轮转 `MessageQueue` 中出队消息，然后按 `MessageKind` 路由到对应的 //! tokio mpsc 通道，由注册的消费者 task 处理。 //! -//! # 架构 +//! ### 角色定位 +//! +//! QueueRunner 是整个消息管线的**调度核心**： +//! 1. 持有 `MessageQueue`（公平轮转）+ 3 个 mpsc 通道 +//! 2. 在 `run()` 循环中不断出队 → 按类型路由 → 投递到正确的消费者 +//! 3. 支持优雅关闭（AtomicBool 信号） +//! +//! ### 架构 //! //! ```text //! ┌──────────────┐ diff --git a/src/scheduler.rs b/src/scheduler.rs index 423a0a8..3f99d49 100644 --- a/src/scheduler.rs +++ b/src/scheduler.rs @@ -3,7 +3,21 @@ use std::sync::Arc; use std::time::Duration; use tokio::time::interval; -/// 定时任务调度器 —— 轮询 PostgreSQL 中到期的 scheduled_tasks +/// ## 定时任务调度器 +/// +/// 轮询 PostgreSQL 中到期的 `scheduled_tasks` 表，执行任务后将结果通知用户。 +/// +/// ### 调度流程 +/// 1. 每 5 秒检查一次是否有到期任务（`next_run_at <= NOW()`） +/// 2. 使用 `FOR UPDATE SKIP LOCKED` 防止多个实例重复执行 +/// 3. 执行任务 shell 命令（120 秒超时） +/// 4. 更新 `next_run_at`（当前时间 + interval_seconds） +/// 5. 通过回调函数通知用户结果 +/// +/// ### 安全保障 +/// * `SKIP LOCKED` — 多实例部署时不会抢同一任务 +/// * `kill_on_drop` — 子进程在超时时会被杀掉 +/// * 120 秒超时 — 防止 shell 命令无限执行 pub struct Scheduler { pool: Option>, } diff --git a/src/state.rs b/src/state.rs index 2403c15..994fdf1 100644 --- a/src/state.rs +++ b/src/state.rs @@ -1,3 +1,14 @@ +//! ## 文件状态管理（PostgreSQL 不可用时的回退方案） +//! +//! ### 设计意图 +//! 在数据库不可用时，使用本地 JSON 文件持久化认证和运行时状态。 +//! 当数据库可用时，数据优先存数据库，并从文件迁移到数据库后清理文件。 +//! +//! ### 存储位置 +//! 默认 `~/.ias/data/weixin-ilink/` +//! * `auth.json` — 认证信息（token, account_id, base_url） +//! * `runtime.json` — 运行时状态（get_updates_buf） + use serde::{Deserialize, Serialize}; use std::path::PathBuf; @@ -15,7 +26,16 @@ pub struct RuntimeState { pub get_updates_buf: String, } -/// 状态管理器 +/// ## 状态管理器 —— 基于文件的状态持久化 +/// +/// 在数据库不可用时代替 DB 存储认证和运行时状态。 +/// +/// ### 文件格式 +/// * `auth.json` — `AuthState { token, account_id, base_url }` +/// * `runtime.json` — `RuntimeState { get_updates_buf }` +/// +/// ### 安全 +/// Unix 系统上 auth.json 的权限被设置为 600（仅所有者可读写） pub struct StateManager { state_dir: PathBuf, } diff --git a/src/tools/approval.rs b/src/tools/approval.rs index 4b49816..dfbaf95 100644 --- a/src/tools/approval.rs +++ b/src/tools/approval.rs @@ -6,7 +6,11 @@ use std::sync::Arc; use std::time::Instant; use tokio::sync::{Mutex, oneshot}; -/// 审批决策 +/// ## 审批决策 +/// +/// * `Approved` — 用户输入了正确的确认码 +/// * `Rejected` — 用户拒绝了（输入 0 或"取消"） +/// * `Expired` — 超时未确认（5 分钟有效） #[derive(Debug, Clone, PartialEq)] pub enum ApprovalDecision { Approved, @@ -24,7 +28,23 @@ struct PendingEntry { sender: oneshot::Sender, } -/// 审批管理器 +/// ## 审批管理器 +/// +/// 管理高风险工具的用户确认流程。 +/// +/// ### 流程 +/// 1. `create(user_id, skill_name)` → 生成 6 位随机确认码，哈希后存储 +/// 2. 向用户发送"⚠️ 操作确认\n\n技能：{name}\n确认码：{code}" +/// 3. 用户在微信聊天中输入确认码 +/// 4. `handle_reply(user_id, reply)` → 验证确认码哈希 +/// - 匹配 → Approved +/// - 不匹配 → 减少尝试次数（最多 3 次） +/// - 输入 0 或"取消" → Rejected +/// - 超时 → Expired（每 60 秒清理一次） +/// +/// ### 存储 +/// * 内存 — `HashMap>`（快速访问） +/// * 数据库 — `pending_approvals` 表（持久化审计） pub struct ApprovalManager { pool: Option>, /// user_id → 该用户的待审批队列 diff --git a/src/tools/builtin.rs b/src/tools/builtin.rs index 6ffb97d..2f1d2ee 100644 --- a/src/tools/builtin.rs +++ b/src/tools/builtin.rs @@ -1,6 +1,21 @@ use crate::tools::types::{RiskLevel, SkillResult}; -/// 内置工具注册表：名称 → 执行、spec 列表、风险判断 +/// ## 内置工具注册表 +/// +/// 连接 LLM 的元工具调用和实际 Rust 实现的工具函数。 +/// +/// ### 职责 +/// 1. `execute(name, args)` — 按名称路由到对应的 builtin 工具执行 +/// 2. `specs()` — 返回所有内置工具的 SkillSpec（给元工具 `query_capabilities` 使用） +/// 3. `is_high_risk(name)` — 判断工具是否需要审批 +/// +/// ### 当前内置工具 +/// * `get_current_datetime` — 获取当前北京时间 +/// * `manage_memos` — 管理备忘录（CRUD） +/// * `query_weather` — 查询天气（和风天气 API） +/// * `web_search` — 联网搜索（Tavily API） +/// * `fetch_page` — 抓取网页内容 +/// * `amap_*` — 高德地图系列（POI搜索、地理编码、路径规划等） pub struct BuiltinRegistry; impl BuiltinRegistry { diff --git a/src/tools/definitions.rs b/src/tools/definitions.rs index 83d3f52..ca9dcfd 100644 --- a/src/tools/definitions.rs +++ b/src/tools/definitions.rs @@ -1,16 +1,30 @@ -//! JSON 工具定义 — 像 curl + jq 一样描述一个工具 +//! ## JSON 工具定义格式 //! -//! 三要素：URL → 取节点 → 选字段 +//! 本模块定义了一种"声明式工具"格式 —— 只需要写一个 JSON 文件就能注册一个新工具， +//! 无需改 Rust 代码、重新编译。 +//! +//! ### 设计理念：像 curl + jq 一样描述工具 //! //! ```json //! { //! "name": "stocks", -//! "description": "股票行情", -//! "url": "https://api.example.com/stocks", +//! "description": "股票行情查询", +//! "url": "https://api.example.com/stocks/{{code}}", +//! "method": "GET", +//! "headers": { "Authorization": "${env:API_KEY}" }, //! "extract": "data.items", //! "fields": ["name", "price"] //! } //! ``` +//! +//! ### 三要素 +//! 1. **URL** — API 端点模板（支持 `{{var}}` 和 `${param:name:type}` 传参） +//! 2. **取节点** — `extract` 字段指定 JSON 点路径（类似 jq） +//! 3. **选字段** — `fields` 字段只保留需要的字段 +//! +//! ### 参数推导 +//! 自动从 URL/Body/Headers 中推导参数定义（无需手动写 parameters 字段）， +//! 也支持显式 `parameters` 声明以获得更精确的类型描述。 use std::collections::HashMap; @@ -60,7 +74,22 @@ impl<'de> serde::Deserialize<'de> for HttpMethod { } } -/// 一个 JSON 工具定义 +/// ## JSON 工具定义（ToolDef） +/// +/// 描述一个通过 HTTP API 调用的外部工具。 +/// +/// ### 字段说明 +/// * `name` — 工具名称，用作 `call_capability` 的 name +/// * `description` — 工具描述（给 LLM 看的） +/// * `url` — URL 模板，支持 `{{var}}` 和 `${param:name:type}` 两种传参 +/// * `method` — HTTP 方法（GET/POST/PUT/DELETE），默认 GET +/// * `headers` — HTTP 请求头，值可包含 `${env:KEY}`（环境变量）、 +/// `${cmd:command}`（命令输出）、`${param:name:type}`（参数值） +/// * `body` — POST/PUT 的 JSON body 模板，同样支持变量替换 +/// * `parameters` — 显式参数声明（可选，未声明时自动推导） +/// * `extract` — JSON 提取路径，如 "data.items" +/// * `fields` — 筛选要保留的字段 +/// * `timeout` — 超时秒数，默认 30 #[derive(Debug, Clone, serde::Deserialize)] pub struct ToolDef { /// 工具名称，用作 call_capability 的 name diff --git a/src/tools/mod.rs b/src/tools/mod.rs index a57e61f..218a41c 100644 --- a/src/tools/mod.rs +++ b/src/tools/mod.rs @@ -1,14 +1,45 @@ +//! ## 工具系统 —— 元工具 + 分离式内置工具 +//! +//! 本模块是 iAs 的"工具系统"，让 LLM 能够调用外部功能。 +//! +//! ### 架构设计（两层元工具） +//! +//! LLM 不直接感知每个具体工具，而是通过两个**元工具**间接调用： +//! +//! 1. **`query_capabilities`** — 列出所有可用工具的名称、描述、参数格式 +//! 2. **`call_capability`** — 按名称调用工具，参数通过 JSON 传递 +//! +//! 这样做的优势：新增工具时只需要在服务端注册，不需要修改 LLM 的 tool definitions。 +//! +//! ### 工具类型 +//! * **内置工具** (builtins/) — Rust 代码实现的工具（天气、搜索、高德地图等） +//! * **JSON 工具** (definitions.rs) — 从 `~/.ias/tools/*.json` 加载，由 ToolDef 定义 +//! * **上下文工具** — read_memories, write_memory, read_summaries（直接在 daemon 中处理） +//! +//! ### 工具执行流程 +//! ```text +//! LLM 请求 call_capability("query_weather", {"location":"北京"}) +//! → daemon 解析出 target_name="query_weather"，提取参数 +//! → 高风险？→ 走 ApprovalManager 审批流 +//! → 低风险？→ BuiltinRegistry::execute() 或 SubprocessRunner::execute() +//! → 结果返回 LLM 继续对话 + pub mod approval; pub mod builtin; pub mod builtins; +pub mod definitions; +pub mod subprocess; pub mod types; -/// 从 call_capability 的 `{name, prompt}` JSON 中提取目标工具的真实参数。 +/// ## 从 `call_capability` 的 `{name, prompt}` JSON 中提取目标工具的真实参数 /// -/// 逻辑： -/// - 如果 prompt 是 JSON 字符串 → parse 后合并 -/// - 如果 prompt 是 JSON 对象 → 直接合并 -/// - 顶层字段（非 name/prompt）也合并进去 +/// `call_capability` 的参数格式为 `{name: "工具名", prompt: "参数JSON字符串"}`。 +/// 这个函数负责从这种格式中提取出工具真正的参数。 +/// +/// ### 提取逻辑 +/// 1. 如果 `prompt` 是 JSON 字符串 → parse 后合并 +/// 2. 如果 `prompt` 是 JSON 对象 → 直接合并 +/// 3. 顶层字段（非 name/prompt）也合并进去（支持直接传参） pub fn unpack_call_params(cp: &serde_json::Value) -> serde_json::Value { let mut params = serde_json::Map::new(); diff --git a/src/tools/subprocess.rs b/src/tools/subprocess.rs index e6eb220..01cf047 100644 --- a/src/tools/subprocess.rs +++ b/src/tools/subprocess.rs @@ -1,14 +1,30 @@ -//! JSON 工具子进程执行器 +//! ## JSON 工具子进程执行器 //! -//! 从 ~/.ias/tools/*.json 加载 ToolDef，运行时通过 HTTP 执行。 -//! 无需编译，放 JSON 即用。 +//! ### 设计理念："放 JSON 即用" //! -//! 执行流程： -//! 模板替换 → HTTP 请求 → JSON 解析 -//! ├─ 失败 → 返回纯文本 -//! ├─ 提取 → json_select() -//! ├─ 筛选 → filter_fields() -//! └─ 格式化输出 +//! 从 `~/.ias/tools/*.json` 加载 `ToolDef`，运行时通过 HTTP 执行。 +//! 不需要修改 Rust 代码、不需要重新编译，放一个 JSON 文件就能注册新工具。 +//! +//! ### 执行流程 +//! +//! ```text +//! 1. 模板替换 +//! {{var}} → args.var（URL 编码） +//! ${param:name:type} → args.name +//! ${env:KEY} → 环境变量 +//! ${cmd:command} → shell 命令输出（有缓存，TTL=300s） +//! +//! 2. HTTP 请求（GET/POST/PUT/DELETE） +//! +//! 3. 返回处理 +//! ├─ 不是 JSON → 直接返回纯文本 +//! ├─ 有 extract → json_select() 按点路径取子节点 +//! ├─ 有 fields → filter_fields() 只保留指定字段 +//! └─ 格式化输出 → 可读文本 +//! ``` +//! +//! ### 缓存 +//! `${cmd:...}` 的结果会被缓存 300 秒，适合 JWT token 等有时效的值。 use crate::tools::definitions::{HttpMethod, ToolDef}; use crate::tools::types::{SkillResult, SkillSpec}; @@ -102,7 +118,14 @@ pub fn runner() -> &'static SubprocessRunner { &RUNNER } -/// 子进程执行器 +/// ## 子进程执行器 —— 管理所有从 JSON 文件加载的工具 +/// +/// ### 全局单例 +/// 通过 `LazyLock` 在首次访问时初始化，从 `~/ias/tools/*.json` 加载全部工具。 +/// +/// ### 运行时热加载 +/// `reload()` 方法可以重新扫描工具目录（目前未在运行时调用， +/// 但为后续热更新保留了接口）。 pub struct SubprocessRunner { tools: RwLock>, } diff --git a/src/tools/types.rs b/src/tools/types.rs index c0bb9a8..8181243 100644 --- a/src/tools/types.rs +++ b/src/tools/types.rs @@ -3,7 +3,10 @@ use std::future::Future; use std::pin::Pin; use std::sync::Arc; -// ─── 风险等级 ─── +/// ## 风险等级 —— 工具的安全级别 +/// +/// * `Low` — 低风险，直接执行（如天气查询、日期时间） +/// * `High` — 高风险，需要用户输入确认码审批（如写备忘录、发消息） #[derive(Debug, Clone, PartialEq, Deserialize)] #[serde(rename_all = "lowercase")] @@ -14,7 +17,14 @@ pub enum RiskLevel { impl RiskLevel {} -// ─── 工具元数据 ─── +/// ## 工具元数据（SkillSpec）—— 工具的自描述信息 +/// +/// 用于给 LLM 展示"有什么工具可用、怎么用"。 +/// * `name` — 工具名称 +/// * `description` — 工具描述 +/// * `risk_level` — 风险等级（Low/High） +/// * `parameters` — JSON Schema 格式的参数定义 +/// * `timeout_secs` — 执行超时时间 #[derive(Debug, Clone, Deserialize)] pub struct SkillSpec { @@ -36,7 +46,11 @@ fn default_timeout() -> u64 { 30 } -// ─── 工具执行结果 ─── +/// ## 工具执行结果 +/// +/// * `success` — 执行是否成功 +/// * `output` — 可读的输出文本（将返回给 LLM） +/// * `data` — 结构化数据（可选），用于后续处理 #[derive(Debug, Clone, Serialize)] pub struct SkillResult { @@ -89,7 +103,12 @@ impl SkillResult { } } -// ─── 执行上下文 ─── +/// ## 执行上下文 —— 工具执行时需要的环境信息 +/// +/// 当工具在 daemon 上下文中执行时，携带以下信息： +/// * `user_id` — 触发工具的用户 +/// * `approval_manager` — 审批管理器（高风险工具需要） +/// * `send_wechat` — 微信发送回调（用于向用户发送审批提示） pub type WechatSender = Arc< dyn Fn(&str, &str) -> Pin> + Send>> + Send + Sync, diff --git a/src/wechat/client.rs b/src/wechat/client.rs index 42c1084..1660646 100644 --- a/src/wechat/client.rs +++ b/src/wechat/client.rs @@ -6,7 +6,24 @@ use std::time::Duration; use tokio::sync::Mutex; use tracing::{debug, info, warn}; -/// 微信 iLink Bot API 客户端 +/// ## 微信 iLink Bot API 客户端 +/// +/// ### 职责 +/// 封装了与微信 iLink Bot API 的所有通信，包括： +/// * 扫码登录流程（获取二维码、轮询状态） +/// * 消息接收（长轮询 getUpdates） +/// * 消息发送（sendMessage） +/// * 生命周期管理（notifyStart/notifyStop） +/// +/// ### 线程安全 +/// `#[derive(Clone)]` + 内部 `Arc>` 保证所有方法可以在 tokio task 间共享。 +/// +/// ### API 端点 +/// * 二维码: `{base_url}/ilink/bot/get_bot_qrcode` +/// * 登录状态: `{base_url}/ilink/bot/get_qrcode_status` +/// * 收消息: `{base_url}/ilink/bot/getupdates` +/// * 发消息: `{base_url}/ilink/bot/sendmessage` +/// * 注册监听: `{base_url}/ilink/bot/msg/notifystart` #[derive(Clone)] pub struct WeChatClient { http: HttpClient, @@ -46,7 +63,10 @@ impl WeChatClient { // ─── 公共方法 ─── - /// 获取登录二维码链接 + /// ## 获取登录二维码链接 +/// + /// 调用 iLink API `/ilink/bot/get_bot_qrcode` 获取二维码。 + /// 二维码内容是一个 URL，浏览器打开后显示二维码图片。 pub async fn get_qrcode(&self) -> Result<(String, String), String> { let url = format!( "{}/ilink/bot/get_bot_qrcode?bot_type={}", @@ -68,7 +88,16 @@ impl WeChatClient { Ok((resp.qrcode, resp.qrcode_img_content)) } - /// 轮询扫码状态 + /// ## 轮询扫码状态（带 35 秒超时的长轮询） + /// + /// 调用 iLink API 检查用户是否已扫码。 + /// * `wait` — 继续等待 + /// * `scaned` — 已扫码，等待确认 + /// * `scaned_but_redirect` — IDC 重定向（切换到 redirect_host） + /// * `confirmed` — 登录成功 + /// * `expired` — 二维码过期 + /// + /// 网络错误不会导致退出，而是返回 `wait` 让调用者继续轮询。 pub async fn poll_qr_status( &self, qrcode: &str, @@ -358,6 +387,12 @@ impl WeChatClient { // ─── 内部方法 ─── + /// ## 公共请求头 +/// + /// 每个请求都携带以下头信息： + /// * `Content-Type: application/json` + /// * `iLink-App-ClientVersion` — 客户端版本号编码 + /// * `X-WECHAT-UIN` — 随机 uint32 → base64 编码的会话标识 fn common_headers() -> reqwest::header::HeaderMap { let mut headers = reqwest::header::HeaderMap::new(); headers.insert( diff --git a/src/wechat/types.rs b/src/wechat/types.rs index 0583a82..6e00ecc 100644 --- a/src/wechat/types.rs +++ b/src/wechat/types.rs @@ -116,7 +116,12 @@ pub struct RefMessage { pub title: Option, } -/// 消息条目（支持 text/image/voice/file/video） +/// ## 消息条目 —— 消息的具体内容 +/// +/// 支持 5 种消息类型：text(1), image(2), voice(3), file(4), video(5)。 +/// 每种类型有对应的 `*_item` 字段，其他字段保持 None。 +/// +/// `text()` 便捷方法用于快速构造文本消息（发送消息时使用）。 #[derive(Debug, Clone, Serialize, Deserialize)] pub struct MessageItem { #[serde(rename = "type", skip_serializing_if = "Option::is_none")] @@ -183,7 +188,18 @@ impl Default for MessageItem { } } -// ─── 微信消息 ─── +/// ## 微信消息结构 +/// +/// 对应 iLink API 返回的单个消息。所有字段都是 Option 的， +/// 因为不同消息类型携带不同字段。 +/// +/// ### 关键字段 +/// * `msg_type` — 消息类型：1=用户消息(TYPE_USER), 2=机器人消息(TYPE_BOT) +/// * `from_user_id` — 发送者微信 ID +/// * `to_user_id` — 接收者微信 ID +/// * `item_list` — 消息内容列表（支持 text/image/voice/file/video） +/// * `context_token` — 微信上下文令牌（用于多轮对话的上下文恢复） +/// * `group_id` — 群聊 ID（群消息时非空） #[derive(Debug, Clone, Serialize, Deserialize)] pub struct WeixinMessage { diff --git a/src/worker.rs b/src/worker.rs index 2ae8410..aeb3750 100644 --- a/src/worker.rs +++ b/src/worker.rs @@ -1,9 +1,15 @@ -#![deprecated = "Worker 进程模式已废弃，请使用 daemon 模式（统一 tokio task 架构）"] - -//! Worker — 无状态执行进程（已废弃） +//! ## Worker —— 无状态执行进程（已废弃） //! +//! ### 旧架构说明 +//! 在旧架构中，每条消息会 spawn 一个独立的 Worker 进程， +//! Worker 通过 UDS 连接到 daemon，处理完一条消息后退出。 +//! +//! 优势：Worker 代码热更新（编译后下一条消息自动用新版本） +//! 劣势：进程启动开销、IPC 序列化开销、状态管理复杂 +//! +//! ### 新架构 +//! 现在改为 daemon 内的 LLM/Tool/Send 三个消费者 tokio task。 //! 此模块保留供编译参考，不再被 daemon 使用。 -//! 新架构使用 daemon 内的 LLM/Tool/Send 三个消费者 tokio task。 //! 1. 连接 daemon 的 Unix Domain Socket //! 2. 读取 TaskFrame (用户消息 + 上下文 + 环境变量) //! 3. 注入环境变量，构建 Conversation，执行 LLM 对话 + 工具调用