Compare commits
10 Commits
8dc232ebab
...
f7d0781090
| Author | SHA1 | Date | |
|---|---|---|---|
 wunianxiao
|
f7d0781090 | ||
|
wunianxiao
|
19dc54f842 | ||
|
wunianxiao
|
a04447c7bc | ||
|
wunianxiao
|
74bb35bd26 | ||
|
wunianxiao
|
713cbb7cc3 | ||
|
wunianxiao
|
8ed608f01b | ||
|
wunianxiao
|
ee5c9b733f | ||
|
wunianxiao
|
079be45615 | ||
|
wunianxiao
|
23c7fbac62 | ||
|
wunianxiao
|
937556917c |
@@ -26,3 +26,5 @@ qweather/
|
||||
Thumbs.db
|
||||
|
||||
read.md
|
||||
.codegraph/
|
||||
.reasonix/
|
||||
|
||||
+445
@@ -0,0 +1,445 @@
|
||||
# iAs 代码调用图谱
|
||||
|
||||
> 微信 AI 智能助手 —— 完整架构与调用关系
|
||||
|
||||
---
|
||||
|
||||
## 1. 全局架构总览
|
||||
|
||||
```mermaid
|
||||
flowchart TB
|
||||
subgraph CLI["CLI 入口 (main.rs)"]
|
||||
A1["clap::Parser Cli 解析"]
|
||||
A2["cmd_login"]
|
||||
A3["cmd_listen"]
|
||||
A4["cmd_send"]
|
||||
A5["cmd_whoami"]
|
||||
A6["cmd_usage"]
|
||||
A7["cmd_tool"]
|
||||
A8["cmd_daemon"]
|
||||
end
|
||||
subgraph DAEMON["守护进程 (daemon.rs)"]
|
||||
B0["run()"]
|
||||
B1["spawn_consumers()"]
|
||||
B2["llm_consumer_loop"]
|
||||
B3["tool_consumer_loop"]
|
||||
B4["send_consumer_loop"]
|
||||
B5["run_llm_round / resume_llm_round"]
|
||||
B6["handle_chat_result"]
|
||||
B7["execute_tool"]
|
||||
end
|
||||
subgraph QUEUE["消息队列 (queue/)"]
|
||||
C1["MessageQueue 公平轮转"]
|
||||
C2["QueueRunner 路由调度"]
|
||||
C3["EnqueueHandle"]
|
||||
C4["ConsumerChannels"]
|
||||
end
|
||||
subgraph LLM["LLM 对话系统 (llm/)"]
|
||||
D1["Conversation 对话管理器"]
|
||||
D2["chat_with_tools() / run_tool_loop()"]
|
||||
D3["DeepSeekProvider"]
|
||||
D4["LlmProvider trait"]
|
||||
D5["StreamReceiver"]
|
||||
end
|
||||
subgraph CONTEXT["上下文管理 (context/)"]
|
||||
E1["ChatSession 消息历史+摘要"]
|
||||
E2["builder Token估算+上下文构建"]
|
||||
E3["MemoryStore 长期记忆"]
|
||||
E4["read_summaries"]
|
||||
end
|
||||
subgraph TOOLS["工具系统 (tools/)"]
|
||||
F1["BuiltinRegistry 工具注册表"]
|
||||
F2["ApprovalManager 审批管理"]
|
||||
F3["build_capability_guide"]
|
||||
F4["unpack_call_params"]
|
||||
subgraph BUILTINS["内置工具 (builtins/)"]
|
||||
G1["datetime"]
|
||||
G2["weather"]
|
||||
G3["web_search"]
|
||||
G4["fetch_page"]
|
||||
G5["memos"]
|
||||
G6["amap"]
|
||||
end
|
||||
end
|
||||
subgraph WECHAT["微信通道 (wechat/)"]
|
||||
H1["WeChatClient"]
|
||||
H2["receive_messages()"]
|
||||
H3["send_text()"]
|
||||
H4["login()"]
|
||||
H5["notify_start/stop"]
|
||||
end
|
||||
subgraph DB["数据库 (db/)"]
|
||||
I1["Database::connect()"]
|
||||
I2["models CRUD 操作"]
|
||||
end
|
||||
subgraph OTHER["其他模块"]
|
||||
J1["state.rs 文件状态回退"]
|
||||
J2["scheduler.rs 定时任务"]
|
||||
J3["logger.rs 日志初始化"]
|
||||
J4["channel/mod.rs ChannelId"]
|
||||
J5["ipc.rs UDS 协议(已废弃)"]
|
||||
J6["worker.rs 旧架构(已废弃)"]
|
||||
end
|
||||
|
||||
A1 --> A2 & A3 & A4 & A5 & A6 & A7 & A8
|
||||
A3 --> B0
|
||||
A8 --> B0
|
||||
B0 --> B1
|
||||
B1 --> B2 & B3 & B4
|
||||
B2 --> B5 & B6
|
||||
B3 --> B7
|
||||
B0 --> C1 & C2 & C3
|
||||
C2 --> C4
|
||||
C4 --> B2 & B3 & B4
|
||||
B2 --> D1
|
||||
D1 --> D2
|
||||
D2 --> D3
|
||||
D3 --> D4
|
||||
D4 --> D5
|
||||
D1 --> E1
|
||||
D2 --> E2
|
||||
B2 --> E3 & E4
|
||||
D2 --> F1
|
||||
B3 --> F1 & F2
|
||||
B7 --> F1 & F3 & F4
|
||||
F1 --> G1 & G2 & G3 & G4 & G5 & G6
|
||||
B4 --> H1
|
||||
H1 --> H2 & H3
|
||||
B0 --> I1
|
||||
B2 --> I2
|
||||
B4 --> I2
|
||||
A2 --> I2 & J1
|
||||
A3 --> J1
|
||||
B0 --> J1 & J2
|
||||
A1 --> J3
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 2. 消息处理全链路
|
||||
|
||||
```mermaid
|
||||
sequenceDiagram
|
||||
participant WX as 微信用户
|
||||
participant WC as WeChatClient
|
||||
participant MQ as MessageQueue
|
||||
participant QR as QueueRunner
|
||||
participant LC as LLM Consumer
|
||||
participant TC as Tool Consumer
|
||||
participant SC as Send Consumer
|
||||
participant LLM as DeepSeek API
|
||||
participant BT as BuiltinRegistry
|
||||
|
||||
WX->>WC: 发送消息
|
||||
WC->>MQ: enqueue(UserMessage)
|
||||
MQ->>QR: dequeue (公平轮转)
|
||||
QR->>LC: llm_tx.send(UserMessage)
|
||||
LC->>LC: 加载历史/记忆/摘要
|
||||
LC->>LLM: chat_stream(上下文)
|
||||
LLM-->>LC: 流式回复 + tool_calls
|
||||
|
||||
alt 有工具调用
|
||||
LC->>MQ: enqueue(ToolCall)
|
||||
MQ->>QR: dequeue
|
||||
QR->>TC: tool_tx.send(ToolCall)
|
||||
|
||||
alt 高风险工具
|
||||
TC->>TC: ApprovalManager.create()
|
||||
TC->>MQ: enqueue(LLMReply: 确认码)
|
||||
MQ->>QR: dequeue
|
||||
QR->>SC: send_tx.send(LLMReply)
|
||||
SC->>WC: send_text(确认码)
|
||||
WX-->>WC: 回复确认码
|
||||
WC->>MQ: enqueue(UserMessage: 确认码)
|
||||
MQ->>QR: dequeue
|
||||
QR->>LC: llm_tx.send(确认码)
|
||||
LC->>TC: ApprovalManager.handle_reply()
|
||||
TC-->>LC: Approved
|
||||
LC->>MQ: enqueue(ToolCall + approved=true)
|
||||
MQ->>QR: dequeue
|
||||
QR->>TC: tool_tx.send(ToolCall)
|
||||
end
|
||||
|
||||
TC->>BT: execute(工具名, 参数)
|
||||
BT-->>TC: SkillResult
|
||||
TC->>MQ: enqueue(ToolResult)
|
||||
MQ->>QR: dequeue
|
||||
QR->>LC: llm_tx.send(ToolResult)
|
||||
LC->>LLM: resume_tool_loop()
|
||||
LLM-->>LC: 最终回复
|
||||
end
|
||||
|
||||
LC->>MQ: enqueue(LLMReply)
|
||||
MQ->>QR: dequeue
|
||||
QR->>SC: send_tx.send(LLMReply)
|
||||
SC->>WC: send_text(回复)
|
||||
WC->>WX: 显示回复
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 3. 模块依赖关系
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
subgraph CORE["核心层"]
|
||||
MAIN["main.rs"]
|
||||
DAEMON["daemon.rs"]
|
||||
CLI["cli.rs"]
|
||||
end
|
||||
subgraph PIPELINE["消息管线"]
|
||||
QUEUE["queue/"]
|
||||
CHANNEL["channel/"]
|
||||
end
|
||||
subgraph AI["AI 层"]
|
||||
LLM["llm/"]
|
||||
CTX["context/"]
|
||||
end
|
||||
subgraph TOOL["工具层"]
|
||||
TOOLS["tools/"]
|
||||
BUILTINS["tools/builtins/"]
|
||||
APPROVAL["tools/approval.rs"]
|
||||
end
|
||||
subgraph IO["IO 层"]
|
||||
WECHAT["wechat/"]
|
||||
DB["db/"]
|
||||
STATE["state.rs"]
|
||||
SCHED["scheduler.rs"]
|
||||
end
|
||||
subgraph UTIL["工具层"]
|
||||
LOG["logger.rs"]
|
||||
IPC["ipc.rs"]
|
||||
WORKER["worker.rs"]
|
||||
end
|
||||
|
||||
MAIN --> CLI & DAEMON & LOG & STATE & DB & CTX & LLM & TOOLS & WECHAT & SCHED
|
||||
DAEMON --> QUEUE & CHANNEL & WECHAT & DB & STATE & LLM & CTX & TOOLS & SCHED
|
||||
QUEUE --> CHANNEL
|
||||
LLM --> CTX & TOOLS
|
||||
TOOLS --> BUILTINS & APPROVAL
|
||||
CTX --> DB
|
||||
TOOLS --> DB
|
||||
APPROVAL --> DB
|
||||
DB --> STATE
|
||||
WORKER --> IPC & LLM & TOOLS
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 4. 数据流（消息生命周期）
|
||||
|
||||
```mermaid
|
||||
flowchart TB
|
||||
START(["📱 用户发消息"]) --> A["WeChatClient receive_messages()"]
|
||||
A --> B["入库: chat_records (inbound)"]
|
||||
B --> C["PipelineMessage kind: UserMessage"]
|
||||
C --> D["MessageQueue 公平轮转"]
|
||||
D --> E["加载历史消息 (最近20条)"]
|
||||
E --> F["加载长期记忆 (MemoryStore)"]
|
||||
F --> G["加载摘要 (read_summaries)"]
|
||||
G --> H["构建上下文 (builder)"]
|
||||
H --> I["Conversation chat_with_tools()"]
|
||||
I --> J{"LLM 调用了工具？"}
|
||||
J -->|"是"| K["Tool Consumer 执行工具"]
|
||||
K --> L{"高风险？→ 审批流程"}
|
||||
L --> M["结果回喂 LLM"]
|
||||
M --> I
|
||||
J -->|"否"| N["LLM 生成最终回复"]
|
||||
N --> O["Send Consumer send_text()"]
|
||||
O --> P["入库: chat_records (outbound)"]
|
||||
P --> ENDD(["📱 用户收到回复"])
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 5. 关键函数调用关系
|
||||
|
||||
### 5.1 工具执行链路
|
||||
|
||||
```mermaid
|
||||
flowchart TB
|
||||
subgraph LLM_SIDE["LLM 侧"]
|
||||
A["LLM 调用 call_capability()"]
|
||||
B["ToolExecutor (闭包)"]
|
||||
C["返回 ASYNC_MARKER + 工具名"]
|
||||
end
|
||||
subgraph QUEUE_SIDE["队列侧"]
|
||||
D["PipelineMessage::tool_call()"]
|
||||
E["MessageQueue.enqueue()"]
|
||||
F["QueueRunner → Tool Consumer"]
|
||||
end
|
||||
subgraph TOOL_SIDE["工具侧"]
|
||||
G["tool_consumer_loop()"]
|
||||
H{"高风险？"}
|
||||
I["ApprovalManager 审批流程"]
|
||||
J["execute_tool()"]
|
||||
K{"工具类型？"}
|
||||
L["BuiltinRegistry.execute()"]
|
||||
M["上下文工具 直接处理"]
|
||||
N["unpack_call_params()"]
|
||||
end
|
||||
subgraph RESULT["结果回喂"]
|
||||
O["PipelineMessage::tool_result()"]
|
||||
P["LLM Consumer resume_llm_round()"]
|
||||
Q["Conversation resume_tool_loop()"]
|
||||
end
|
||||
|
||||
A --> B --> C --> D --> E --> F --> G
|
||||
G --> H
|
||||
H -->|"是"| I --> J
|
||||
H -->|"否"| J
|
||||
J --> K
|
||||
K -->|"内置工具"| L
|
||||
K -->|"上下文工具"| M
|
||||
K -->|"call_capability"| N --> K
|
||||
J --> O --> P --> Q
|
||||
```
|
||||
|
||||
### 5.2 审批流程
|
||||
|
||||
```mermaid
|
||||
flowchart TB
|
||||
A["LLM 请求高风险工具"] --> B["ApprovalManager.create(user_id, tool)"]
|
||||
B --> C["生成 6 位确认码 SHA-256 哈希存储"]
|
||||
B --> D["创建 oneshot channel"]
|
||||
C --> E["发送消息给用户 ⚠️ 操作确认"]
|
||||
D --> F["等待用户回复 超时 5 分钟"]
|
||||
F --> G{"用户回复"}
|
||||
G -->|"匹配确认码"| H["ApprovalDecision::Approved"]
|
||||
G -->|"输入 0/取消"| I["ApprovalDecision::Rejected"]
|
||||
G -->|"不匹配(最多3次)"| J["减少尝试次数"]
|
||||
G -->|"超时"| K["ApprovalDecision::Expired"]
|
||||
J --> F
|
||||
H --> L["重新入队 UserMessage + approved_tool"]
|
||||
I --> M["回复用户: 已取消"]
|
||||
K --> N["回复用户: 已超时"]
|
||||
```
|
||||
|
||||
### 5.3 上下文构建流程
|
||||
|
||||
```mermaid
|
||||
flowchart TB
|
||||
A["ChatSession (消息历史)"] --> B["estimate_tokens()"]
|
||||
B --> C{"Token 超 budget? (默认 28000)"}
|
||||
C -->|"是"| D{"checkpoint 之前有未摘要消息？"}
|
||||
D -->|"是"| E["trigger_overflow_summary() → 调用 LLM 生成摘要"]
|
||||
D -->|"否"| F["裁剪最早的消息 直到符合 budget"]
|
||||
C -->|"否"| G{"空闲超时? (12h 无消息)"}
|
||||
G -->|"是"| H["trigger_idle_summary() → 清空消息, 保留摘要"]
|
||||
G -->|"否"| I["build_context() → 组装最终上下文"]
|
||||
E --> I
|
||||
F --> I
|
||||
H --> I
|
||||
I --> J["发送给 LLM"]
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 6. 文件依赖矩阵
|
||||
|
||||
| 文件 | 依赖的关键模块 |
|
||||
|------|---------------|
|
||||
| `main.rs` | cli, daemon, db, llm, context, tools, wechat, state, scheduler, logger |
|
||||
| `daemon.rs` | queue, channel, wechat, db, state, llm, context, tools, scheduler |
|
||||
| `cli.rs` | clap (独立) |
|
||||
| `llm/conversation.rs` | context/builder, context/types, llm/provider, llm/types |
|
||||
| `llm/deepseek.rs` | llm/provider, llm/types |
|
||||
| `llm/provider.rs` | llm/types |
|
||||
| `llm/types.rs` | serde (独立) |
|
||||
| `context/types.rs` | llm/types |
|
||||
| `context/builder.rs` | context/types, llm/types |
|
||||
| `context/tools.rs` | db/models |
|
||||
| `queue/message_queue.rs` | channel/mod |
|
||||
| `queue/runner.rs` | queue/message_queue |
|
||||
| `tools/builtin.rs` | tools/types, tools/builtins/* |
|
||||
| `tools/approval.rs` | db/models |
|
||||
| `tools/mod.rs` | tools/types |
|
||||
| `wechat/client.rs` | wechat/types |
|
||||
| `wechat/types.rs` | serde (独立) |
|
||||
| `db/mod.rs` | db/models |
|
||||
| `db/models.rs` | sqlx (独立) |
|
||||
| `state.rs` | serde (独立) |
|
||||
| `scheduler.rs` | db/models |
|
||||
| `worker.rs` | ipc, llm, tools (已废弃) |
|
||||
| `ipc.rs` | serde (已废弃) |
|
||||
|
||||
---
|
||||
|
||||
## 7. 启动流程
|
||||
|
||||
```mermaid
|
||||
flowchart TB
|
||||
START(["ias listen --llm"]) --> A["main()"]
|
||||
A --> B["clap::Parser::parse()"]
|
||||
B --> C{"命令类型"}
|
||||
C -->|"Tool"| D["cmd_tool() 直接执行工具"]
|
||||
C -->|"Login"| E["cmd_login() 扫码登录"]
|
||||
C -->|"Listen/Service"| F["cmd_listen()"]
|
||||
|
||||
F --> G["加载 .env"]
|
||||
F --> H["Database::connect()"]
|
||||
H --> I{"数据库可用？"}
|
||||
I -->|"是"| J["创建 MemoryStore (DB 模式)"]
|
||||
I -->|"否"| K["创建 MemoryStore (文件模式)"]
|
||||
|
||||
F --> L["加载认证信息"]
|
||||
L --> M{"已登录？"}
|
||||
M -->|"否"| N["报错退出"]
|
||||
M -->|"是"| O["WeChatClient.set_auth()"]
|
||||
|
||||
O --> P["WeChatClient.notify_start()"]
|
||||
P --> Q["创建 ApprovalManager"]
|
||||
Q --> R["创建 Conversation + ToolExecutor"]
|
||||
R --> S["启动 Scheduler (定时任务)"]
|
||||
S --> T["启动 listen_loop()"]
|
||||
|
||||
T --> U["WeChatClient.receive_messages()"]
|
||||
U --> V{"收到消息？"}
|
||||
V -->|"是"| W["enqueue PipelineMessage"]
|
||||
W --> X["LLM/Tool/Send 消费者处理"]
|
||||
V -->|"否/超时"| U
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 8. 消费者架构（三消费者）
|
||||
|
||||
```mermaid
|
||||
flowchart TB
|
||||
subgraph QUEUE2["消息队列"]
|
||||
Q["MessageQueue 公平轮转"]
|
||||
R["QueueRunner 路由"]
|
||||
end
|
||||
subgraph LLM_C["LLM Consumer"]
|
||||
L1["接收: UserMessage / ToolResult / ScheduledTask"]
|
||||
L2["加载用户历史 + 记忆 + 摘要"]
|
||||
L3["创建/恢复 Conversation"]
|
||||
L4["chat_with_tools() / resume_tool_loop()"]
|
||||
L5{"有异步工具等待？"}
|
||||
L6["发送 LLMReply"]
|
||||
L7["保留 session"]
|
||||
L8["清理 session"]
|
||||
end
|
||||
subgraph TOOL_C["Tool Consumer"]
|
||||
T1["接收: ToolCall / ApprovalRequest"]
|
||||
T2{"高风险 + 未审批？"}
|
||||
T3["ApprovalManager.create()"]
|
||||
T4["发送确认码"]
|
||||
T5["execute_tool()"]
|
||||
T6["发送 ToolResult"]
|
||||
end
|
||||
subgraph SEND_C["Send Consumer"]
|
||||
S1["接收: LLMReply"]
|
||||
S2["WeChatClient.send_text()"]
|
||||
S3["入库: chat_records (outbound)"]
|
||||
end
|
||||
|
||||
Q --> R
|
||||
R -->|"UserMessage / ToolResult"| L1
|
||||
R -->|"ToolCall"| T1
|
||||
R -->|"LLMReply"| S1
|
||||
L1 --> L2 --> L3 --> L4
|
||||
L4 --> L5
|
||||
L5 -->|"是"| L7
|
||||
L5 -->|"否"| L6
|
||||
```
|
||||
+51
-40
@@ -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"
|
||||
|
||||
@@ -0,0 +1,10 @@
|
||||
[package]
|
||||
name = "ias-fetch"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
|
||||
[dependencies]
|
||||
chromiumoxide = { version = "0.7", default-features = false, features = ["tokio-runtime", "bytes"] }
|
||||
futures = "0.3"
|
||||
serde_json = "1.0"
|
||||
tokio = { version = "1.0", features = ["full"] }
|
||||
@@ -0,0 +1,143 @@
|
||||
//! ias-fetch — 通过 headless Chrome 提取网页正文
|
||||
//!
|
||||
//! 启动无头 Chrome 渲染页面 JS，取渲染后的纯文本输出到 stdout。
|
||||
//! 专供 JSON 工具 `${cmd:ias-fetch {{url}}}` 调用。
|
||||
//!
|
||||
//! 用法:
|
||||
//! ias-fetch <url> # 默认：取 body.innerText
|
||||
//! ias-fetch --raw <url> # 返回完整 HTML
|
||||
//! ias-fetch --timeout 15 <url> # 自定义超时（默认 15s）
|
||||
|
||||
use futures::StreamExt;
|
||||
use std::time::Duration;
|
||||
|
||||
use chromiumoxide::{Browser, BrowserConfig, Page};
|
||||
|
||||
const DEFAULT_TIMEOUT: u64 = 15;
|
||||
const MAX_OUTPUT_CHARS: usize = 10_000;
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
let args: Vec<String> = std::env::args().collect();
|
||||
|
||||
let mut raw_mode = false;
|
||||
let mut timeout_secs = DEFAULT_TIMEOUT;
|
||||
let mut url: Option<String> = None;
|
||||
|
||||
let mut i = 1;
|
||||
while i < args.len() {
|
||||
match args[i].as_str() {
|
||||
"--raw" => raw_mode = true,
|
||||
"--timeout" => {
|
||||
i += 1;
|
||||
timeout_secs = args.get(i).and_then(|s| s.parse().ok()).unwrap_or(DEFAULT_TIMEOUT);
|
||||
}
|
||||
s if s.starts_with("http://") || s.starts_with("https://") => {
|
||||
url = Some(s.to_string());
|
||||
}
|
||||
s if !s.starts_with('-') => url = Some(s.to_string()),
|
||||
_ => {}
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
|
||||
let url = match url {
|
||||
Some(u) => {
|
||||
if u.starts_with("http://") || u.starts_with("https://") {
|
||||
u
|
||||
} else {
|
||||
format!("https://{}", u)
|
||||
}
|
||||
}
|
||||
None => {
|
||||
eprintln!("用法: ias-fetch [--raw] [--timeout N] <url>");
|
||||
std::process::exit(1);
|
||||
}
|
||||
};
|
||||
|
||||
let timeout = Duration::from_secs(timeout_secs);
|
||||
|
||||
match tokio::time::timeout(timeout, fetch_page(&url, raw_mode)).await {
|
||||
Ok(Ok(output)) => {
|
||||
if output.is_empty() {
|
||||
eprintln!("页面无内容");
|
||||
std::process::exit(1);
|
||||
}
|
||||
println!("{}", output);
|
||||
}
|
||||
Ok(Err(e)) => {
|
||||
eprintln!("{}", e);
|
||||
std::process::exit(1);
|
||||
}
|
||||
Err(_) => {
|
||||
eprintln!("超时 ({}s): {}", timeout_secs, url);
|
||||
std::process::exit(1);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn fetch_page(url: &str, raw_mode: bool) -> Result<String, String> {
|
||||
let (mut browser, mut handler) = Browser::launch(
|
||||
BrowserConfig::builder()
|
||||
.build()
|
||||
.map_err(|e| format!("Chrome 配置失败: {}", e))?,
|
||||
)
|
||||
.await
|
||||
.map_err(|e| format!("启动 Chrome 失败: {}\n请确保已安装 chromium 或 google-chrome", e))?;
|
||||
|
||||
let handle = tokio::spawn(async move {
|
||||
while let Some(_event) = handler.next().await {
|
||||
// 持续消费 handler 事件
|
||||
}
|
||||
});
|
||||
|
||||
// new_page 接受 URL，自动导航并等待加载
|
||||
let page = browser
|
||||
.new_page(url)
|
||||
.await
|
||||
.map_err(|e| format!("导航到 {} 失败: {}", url, e))?;
|
||||
|
||||
// 额外等待 JS 渲染
|
||||
tokio::time::sleep(Duration::from_secs(2)).await;
|
||||
|
||||
// 提取内容
|
||||
let content = if raw_mode {
|
||||
page.content().await.map_err(|e| format!("获取 HTML 失败: {}", e))?
|
||||
} else {
|
||||
get_inner_text(&page).await?
|
||||
};
|
||||
|
||||
// 清理
|
||||
browser.close().await.ok();
|
||||
handle.await.ok();
|
||||
|
||||
// 截断
|
||||
if content.chars().count() > MAX_OUTPUT_CHARS {
|
||||
Ok(format!(
|
||||
"{}…（已截断，完整内容 {} 字符）",
|
||||
content.chars().take(MAX_OUTPUT_CHARS).collect::<String>(),
|
||||
content.chars().count()
|
||||
))
|
||||
} else {
|
||||
Ok(content)
|
||||
}
|
||||
}
|
||||
|
||||
/// 通过 evaluate JS 获取 body.innerText
|
||||
async fn get_inner_text(page: &Page) -> Result<String, String> {
|
||||
let result = page
|
||||
.evaluate("document.body?.innerText || document.documentElement?.innerText || ''")
|
||||
.await
|
||||
.map_err(|e| format!("提取正文失败: {}", e))?;
|
||||
|
||||
let value: serde_json::Value = result
|
||||
.into_value()
|
||||
.map_err(|e| format!("解析 evaluate 结果失败: {}", e))?;
|
||||
match value {
|
||||
serde_json::Value::String(s) => {
|
||||
let cleaned: Vec<&str> = s.lines().map(|l| l.trim()).filter(|l| !l.is_empty()).collect();
|
||||
Ok(cleaned.join("\n"))
|
||||
}
|
||||
_ => Ok(String::new()),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,112 @@
|
||||
/// ## 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;
|
||||
|
||||
fn main() {
|
||||
let args: Vec<String> = std::env::args().collect();
|
||||
if args.len() < 2 {
|
||||
eprintln!("用法: ias-auth <command>");
|
||||
eprintln!("命令:");
|
||||
eprintln!(" gen-qweather-jwt 生成和风天气 Ed25519 JWT Token");
|
||||
exit(1);
|
||||
}
|
||||
match args[1].as_str() {
|
||||
"gen-qweather-jwt" => gen_qweather_jwt(),
|
||||
_ => {
|
||||
eprintln!("未知命令: {}", args[1]);
|
||||
eprintln!("支持的命令: gen-qweather-jwt");
|
||||
exit(1);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn gen_qweather_jwt() {
|
||||
match generate_jwt() {
|
||||
Ok(token) => println!("{}", token),
|
||||
Err(e) => {
|
||||
eprintln!("{}", e);
|
||||
exit(1);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ─── 和风天气 Ed25519 JWT 生成 ───
|
||||
// 从 iAs 的 tools::builtins::weather 提取，保持独立
|
||||
|
||||
use base64::Engine;
|
||||
use chrono::Utc;
|
||||
use ed25519_dalek::pkcs8::DecodePrivateKey;
|
||||
use ed25519_dalek::{Signer, SigningKey};
|
||||
|
||||
fn base64url(data: &[u8]) -> String {
|
||||
base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(data)
|
||||
}
|
||||
|
||||
fn generate_jwt() -> Result<String, String> {
|
||||
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))?;
|
||||
|
||||
// 手动解析 PEM: 去掉头尾，base64 解码得到 DER
|
||||
let der = pem
|
||||
.lines()
|
||||
.filter(|l| !l.starts_with("-----"))
|
||||
.collect::<Vec<_>>()
|
||||
.join("");
|
||||
let der_bytes = base64::engine::general_purpose::STANDARD
|
||||
.decode(&der)
|
||||
.map_err(|e| format!("Base64 解码密钥失败: {e}"))?;
|
||||
|
||||
let signing_key = SigningKey::from_pkcs8_der(&der_bytes)
|
||||
.map_err(|e| format!("解析 Ed25519 密钥失败: {e}"))?;
|
||||
|
||||
let key_id = std::env::var("QWEATHER_JWT_KEY_ID").unwrap_or_default();
|
||||
let project_id = std::env::var("QWEATHER_JWT_PROJECT_ID").unwrap_or_default();
|
||||
let ttl: i64 = std::env::var("QWEATHER_JWT_TTL_SECONDS")
|
||||
.ok()
|
||||
.and_then(|s| s.parse().ok())
|
||||
.unwrap_or(3600);
|
||||
|
||||
let now = Utc::now().timestamp();
|
||||
|
||||
// Header
|
||||
let header = serde_json::json!({"alg": "EdDSA", "kid": key_id});
|
||||
let header_b64 = base64url(header.to_string().as_bytes());
|
||||
|
||||
// Payload
|
||||
let payload = serde_json::json!({
|
||||
"sub": project_id,
|
||||
"iat": now - 30,
|
||||
"exp": now + ttl,
|
||||
});
|
||||
let payload_b64 = base64url(payload.to_string().as_bytes());
|
||||
|
||||
// Sign
|
||||
let message = format!("{}.{}", header_b64, payload_b64);
|
||||
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())
|
||||
))
|
||||
}
|
||||
@@ -0,0 +1,68 @@
|
||||
//! ## 渠道/通道基础类型
|
||||
//!
|
||||
//! 定义消息的来源和目标渠道标识，以及从外部轮询收到的原始消息。
|
||||
//!
|
||||
//! ### 设计意图
|
||||
//!
|
||||
//! `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"
|
||||
pub platform: String,
|
||||
/// 在该平台上的用户 ID
|
||||
pub user_id: String,
|
||||
}
|
||||
|
||||
impl ChannelId {
|
||||
/// 快捷构造微信渠道
|
||||
pub fn wechat(user_id: impl Into<String>) -> Self {
|
||||
Self { platform: "wechat".into(), user_id: user_id.into() }
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Display for ChannelId {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
write!(f, "{}:{}", self.platform, self.user_id)
|
||||
}
|
||||
}
|
||||
|
||||
/// ## 从外部轮询收到的原始消息
|
||||
///
|
||||
/// 在 daemon 的主循环中，从 WeChat API 获取到的消息会被转换为这个结构，
|
||||
/// 然后封装成 `PipelineMessage` 投递到消息队列。
|
||||
///
|
||||
/// * `channel` — 来源渠道（平台 + 用户）
|
||||
/// * `text` — 消息文本
|
||||
/// * `message_id` — 消息 ID
|
||||
/// * `context_token` — 微信的上下文令牌（用于多轮对话的上下文恢复）
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[allow(dead_code)]
|
||||
pub struct IncomingMessage {
|
||||
pub channel: ChannelId,
|
||||
pub text: String,
|
||||
pub message_id: String,
|
||||
pub context_token: Option<String>,
|
||||
}
|
||||
+24
-7
@@ -1,15 +1,32 @@
|
||||
//! ## CLI 子命令定义 —— clap 参数解析
|
||||
//!
|
||||
//! 定义了 iAs 的所有命令行子命令和参数：
|
||||
//!
|
||||
//! | 命令 | 说明 |
|
||||
//! |------|------|
|
||||
//! | `ias login` | 扫码登录微信 |
|
||||
//! | `ias listen --llm` | 启动 AI 自动回复 |
|
||||
//! | `ias send --to <id> --text <msg>` | 手动发送消息 |
|
||||
//! | `ias whoami` | 查看登录状态 |
|
||||
//! | `ias usage` | Token 用量统计 |
|
||||
//! | `ias service` | 后台服务模式 |
|
||||
//! | `ias daemon` | 守护进程模式 |
|
||||
//! | `ias tool <tool>` | 调用内置工具 |
|
||||
|
||||
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 {
|
||||
@@ -115,7 +132,7 @@ pub enum Commands {
|
||||
sock: String,
|
||||
},
|
||||
|
||||
/// Worker 进程（由 daemon 自动 spawn，也可手动测试）
|
||||
/// [已废弃] Worker 进程（新架构使用 daemon 统一 tokio task）
|
||||
///
|
||||
/// 示例:
|
||||
/// ias worker --sock /tmp/ias_daemon.sock
|
||||
|
||||
+57
-9
@@ -1,16 +1,40 @@
|
||||
//! ## 上下文构建器 —— Token 估算 + 上下文组装 + 摘要触发
|
||||
//!
|
||||
//! 这是 LLM 调用前的关键步骤：
|
||||
//!
|
||||
//! 1. **Token 估算** — 中/英文分别估算 token 消耗
|
||||
//! 2. **上下文构建** — 从 ChatSession 中提取合适的消息，在 token budget 内保留最多上下文
|
||||
//! 3. **摘要触发** — 溢出摘要（token 超 budget）和空闲超时摘要（12h 无消息）
|
||||
//!
|
||||
//! ### Token 估算公式
|
||||
//! - 中文字符 ≈ 1.5 tokens/字
|
||||
//! - 英文字符 ≈ 0.25 tokens/字符
|
||||
//! - 每条消息 + 8 tokens 格式开销
|
||||
|
||||
use crate::context::types::ChatSession;
|
||||
use crate::llm::conversation::Summarizer;
|
||||
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 +51,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<Mutex<ChatSession>>, system_prompt: &str) -> Vec<Message> {
|
||||
let s = session.lock().await;
|
||||
|
||||
@@ -92,7 +124,18 @@ pub async fn build_context(session: &Arc<Mutex<ChatSession>>, 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<Mutex<ChatSession>>,
|
||||
summarizer: Option<&Summarizer>,
|
||||
@@ -153,7 +196,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<Mutex<ChatSession>>,
|
||||
summarizer: Option<&Summarizer>,
|
||||
@@ -194,8 +243,8 @@ pub async fn trigger_idle_summary(
|
||||
|
||||
/// 生成摘要：优先使用 LLM，不可用时回退到简单截断
|
||||
async fn generate_summary(messages: &[Message], summarizer: Option<&Summarizer>) -> String {
|
||||
if let Some(summarizer) = summarizer {
|
||||
if messages.len() >= 3 {
|
||||
if let Some(summarizer) = summarizer
|
||||
&& messages.len() >= 3 {
|
||||
let prompt = build_summary_prompt(messages);
|
||||
match summarizer(prompt).await {
|
||||
Ok(text) if !text.is_empty() => {
|
||||
@@ -205,7 +254,6 @@ async fn generate_summary(messages: &[Message], summarizer: Option<&Summarizer>)
|
||||
_ => tracing::warn!("LLM 摘要失败，回退到简单截断"),
|
||||
}
|
||||
}
|
||||
}
|
||||
summarize_messages(messages)
|
||||
}
|
||||
|
||||
|
||||
@@ -1,3 +1,16 @@
|
||||
//! ## 上下文管理 —— 对话状态 + Token 预算 + 长期记忆
|
||||
//!
|
||||
//! 管理 LLM 对话的上下文状态，包括：
|
||||
//!
|
||||
//! - `types` — `ChatSession` 核心数据结构（消息历史、摘要、checkpoint）
|
||||
//! - `builder` — Token 预算估算、上下文构建、摘要触发
|
||||
//! - `tools` — `MemoryStore` 长期记忆管理器
|
||||
//!
|
||||
//! ### 关键设计
|
||||
//! - **Checkpoint 机制**：checkpoint 之前的消息被压缩成摘要，之后的消息保持完整
|
||||
//! - **Token Budget**：默认 28000 tokens，留 4000 给回复
|
||||
//! - **双重摘要**：溢出摘要（token 超预算）和空闲超时摘要（12h 无消息）
|
||||
|
||||
pub mod builder;
|
||||
pub mod tools;
|
||||
pub mod types;
|
||||
|
||||
+32
-2
@@ -1,9 +1,36 @@
|
||||
//! ## 长期记忆管理器 + 摘要读取工具
|
||||
//!
|
||||
//! 提供两个核心功能：
|
||||
//!
|
||||
//! 1. **MemoryStore** — 按用户隔离的长期记忆存储（内存缓存 + PostgreSQL 双写）
|
||||
//! 2. **read_summaries** — 读取历史会话摘要（内存 + 数据库合并去重）
|
||||
//!
|
||||
//! ### 数据流
|
||||
//! ```text
|
||||
//! 启动时 load(user_id) → 从数据库加载到缓存
|
||||
//! 对话中 read_for() → 从缓存读取，返回 "1. ...\n2. ..." 格式
|
||||
//! 对话中 write_for() → 写入缓存 + 写入数据库
|
||||
//! ```
|
||||
|
||||
use sqlx::PgPool;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
/// 长期记忆管理器（内存 + PostgreSQL 双写，按用户隔离）
|
||||
/// ## 长期记忆管理器（内存 + PostgreSQL 双写）
|
||||
///
|
||||
/// 为每个用户维护独立的长期记忆列表，支持读写操作。
|
||||
///
|
||||
/// ### 存储策略
|
||||
/// * **内存缓存** — `HashMap<String, Vec<String>>`，按 user_id 隔离
|
||||
/// * **数据库持久化** — 可选的 `PgPool`，写入时同时写到 PostgreSQL
|
||||
///
|
||||
/// ### 数据流
|
||||
/// ```text
|
||||
/// 启动时 load(user_id) → 从数据库加载到缓存
|
||||
/// 对话中 read_for() → 从缓存读取，返回 "1. ...\n2. ..." 格式
|
||||
/// 对话中 write_for() → 写入缓存 + 写入数据库
|
||||
/// ```
|
||||
pub struct MemoryStore {
|
||||
pool: Option<Arc<PgPool>>,
|
||||
cache: Arc<Mutex<HashMap<String, Vec<String>>>>,
|
||||
@@ -61,7 +88,10 @@ impl MemoryStore {
|
||||
}
|
||||
}
|
||||
|
||||
/// read_summaries 工具：读取历史摘要（内存 + 数据库）
|
||||
/// ## `read_summaries` 工具 —— 读取历史摘要（内存 + 数据库）
|
||||
///
|
||||
/// 返回当前 session 中的摘要列表 + 数据库中存档的摘要。
|
||||
/// 去重：数据库中的摘要如果与 session 中的摘要前 30 个字符相同，则跳过。
|
||||
pub async fn read_summaries(
|
||||
session: &Arc<Mutex<super::types::ChatSession>>,
|
||||
) -> String {
|
||||
|
||||
+54
-3
@@ -1,9 +1,31 @@
|
||||
//! ## 聊天会话状态 —— 核心数据结构
|
||||
//!
|
||||
//! 这是整个项目中最重要的数据结构之一 —— 它维护了一次 LLM 对话的完整状态。
|
||||
//!
|
||||
//! ### 关键设计
|
||||
//! - **消息历史** (`messages`) — 从 checkpoint 之后的消息保持完整
|
||||
//! - **摘要机制** (`summaries` + `checkpoint`) — checkpoint 之前的消息被压缩成摘要
|
||||
//! - **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()` → 生成摘要，清空消息
|
||||
|
||||
use crate::llm::types::Message;
|
||||
use chrono::{DateTime, Utc};
|
||||
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 触发
|
||||
@@ -12,7 +34,12 @@ pub enum SummaryReason {
|
||||
Timeout,
|
||||
}
|
||||
|
||||
/// 一条摘要记录
|
||||
/// ## 一条摘要记录
|
||||
///
|
||||
/// * `checkpoint` — 摘要对应的消息位置，此位置之前的消息已被压缩成摘要
|
||||
/// * `text` — 摘要文本
|
||||
/// * `reason` — 生成原因（Overflow / Timeout）
|
||||
/// * `created_at` — 生成时间
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SummaryEntry {
|
||||
/// 摘要对应的消息位置（checkpoint）
|
||||
@@ -25,7 +52,23 @@ pub struct SummaryEntry {
|
||||
pub created_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
/// 聊天会话状态
|
||||
/// ## 聊天会话状态（核心数据结构）
|
||||
///
|
||||
/// 这是整个项目中最重要的数据结构之一 —— 它维护了一次 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,
|
||||
@@ -44,6 +87,8 @@ pub struct ChatSession {
|
||||
pub token_budget: u32,
|
||||
/// 12 小时空闲阈值（秒）
|
||||
pub idle_timeout_secs: i64,
|
||||
/// 会话 ID（用于追踪工具调用往返）
|
||||
pub session_id: Uuid,
|
||||
}
|
||||
|
||||
impl Default for ChatSession {
|
||||
@@ -58,6 +103,7 @@ impl Default for ChatSession {
|
||||
last_user_at: None,
|
||||
token_budget: 28000,
|
||||
idle_timeout_secs: 12 * 3600,
|
||||
session_id: Uuid::new_v4(),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -90,6 +136,11 @@ impl ChatSession {
|
||||
.push(Message::tool_result(tool_call_id, name, result));
|
||||
}
|
||||
|
||||
/// 添加一条系统说明消息（如已批准的工具）
|
||||
pub fn add_system_note(&mut self, note: &str) {
|
||||
self.messages.push(Message::system(note));
|
||||
}
|
||||
|
||||
/// 是否因空闲超过阈值需要摘要
|
||||
pub fn is_idle_timeout(&self) -> bool {
|
||||
if let Some(last) = self.last_user_at {
|
||||
|
||||
+596
-717
File diff suppressed because it is too large
Load Diff
+36
-2
@@ -1,10 +1,44 @@
|
||||
//! ## 数据库管理
|
||||
//!
|
||||
//! ### 职责
|
||||
//! 1. 从 `DATABASE_URL` 环境变量连接 PostgreSQL
|
||||
//! 2. 自动运行 `./migrations/` 下的 SQL 迁移
|
||||
//! 3. 提供连接池引用给各模块使用
|
||||
//!
|
||||
//! ### 回退策略
|
||||
//! 数据库是可选的。如果数据库不可用，系统会优雅地回退到文件存储
|
||||
//! （`StateManager`），核心功能仍然可用。
|
||||
|
||||
//! ## 数据库管理 —— PostgreSQL 连接 + 迁移
|
||||
//!
|
||||
//! ### 职责
|
||||
//! 1. 从 `DATABASE_URL` 环境变量连接 PostgreSQL
|
||||
//! 2. 自动运行 `./migrations/` 下的 SQL 迁移
|
||||
//! 3. 提供连接池引用给各模块使用
|
||||
//!
|
||||
//! ### 回退策略
|
||||
//! 数据库是可选的。如果数据库不可用，系统会优雅地回退到文件存储
|
||||
//! （`StateManager`），核心功能仍然可用。
|
||||
//!
|
||||
//! ### 表结构
|
||||
//! - `app_state` — 认证状态 KV 存储
|
||||
//! - `chat_records` — 收发消息记录
|
||||
//! - `llm_usage` — Token 用量统计
|
||||
//! - `user_memories` — 用户长期记忆
|
||||
//! - `pending_approvals` — 待审批记录
|
||||
//! - `session_summaries` — 会话摘要
|
||||
//! - `scheduled_tasks` — 定时任务
|
||||
|
||||
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,
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
|
||||
+15
-5
@@ -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};
|
||||
@@ -149,6 +158,7 @@ pub async fn recv_task(stream: &mut UnixStream) -> Result<TaskFrame, String> {
|
||||
}
|
||||
|
||||
/// 接收并解析为 OutputFrame
|
||||
#[allow(dead_code)]
|
||||
pub async fn recv_output(stream: &mut UnixStream) -> Result<OutputFrame, String> {
|
||||
let v = recv_frame(stream).await?;
|
||||
serde_json::from_value(v).map_err(|e| format!("解析 OutputFrame 失败: {}", e))
|
||||
|
||||
+160
-13
@@ -1,31 +1,116 @@
|
||||
//! ## 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？→ 返回最终回复
|
||||
//! ```
|
||||
|
||||
use crate::context::builder;
|
||||
use crate::context::types::ChatSession;
|
||||
use crate::llm::provider::{BoxedProvider, StreamReceiver, create_provider};
|
||||
use crate::llm::types::{ConversationConfig, Message, StreamChunk, ToolCall, Usage};
|
||||
use std::future::Future;
|
||||
use std::pin::Pin;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
/// ## 工具执行器 —— LLM 调用的工具回调
|
||||
///
|
||||
/// 这是一个函数别名：`Fn(工具名, JSON参数) → Future<Result<String>>`。
|
||||
///
|
||||
/// 当 LLM 在对话中发起工具调用时，Conversation 会调用这个 executor 来执行工具。
|
||||
///
|
||||
/// ### 两种执行模式
|
||||
/// 1. **同步模式** — executor 直接执行并返回结果字符串
|
||||
/// 2. **异步模式** — executor 将工具调用入队到消息队列，返回 `ASYNC_MARKER + 工具名`
|
||||
/// 表示"工具已异步入队，结果稍后通过 ToolResult 消息回喂"
|
||||
pub type ToolExecutor = Arc<
|
||||
dyn Fn(&str, &str) -> Pin<Box<dyn Future<Output = Result<String, String>> + Send>>
|
||||
+ Send
|
||||
+ Sync,
|
||||
>;
|
||||
|
||||
/// 摘要生成器：接收要摘要的消息文本 → 返回 LLM 生成的摘要
|
||||
/// ## 摘要生成器 —— 用于生成对话摘要的回调
|
||||
///
|
||||
/// Conversation 在两种情况下会触发摘要：
|
||||
/// 1. **溢出摘要** — 消息 token 总数接近 budget 时，压缩早期消息
|
||||
/// 2. **空闲超时摘要** — 距上条用户消息超过 12 小时时，总结本次会话
|
||||
pub type Summarizer = Arc<
|
||||
dyn Fn(String) -> Pin<Box<dyn Future<Output = Result<String, String>> + Send>> + Send + Sync,
|
||||
>;
|
||||
|
||||
/// ## `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,
|
||||
pub used_tools: bool,
|
||||
pub usage: Option<Usage>,
|
||||
/// 是否有异步工具调用仍在等待结果
|
||||
pub has_pending_async: bool,
|
||||
}
|
||||
|
||||
/// 异步工具标记 —— 工具执行器返回此前缀表示工具已异步入队
|
||||
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,
|
||||
session: Arc<Mutex<ChatSession>>,
|
||||
tool_executor: Option<ToolExecutor>,
|
||||
/// 是否有异步工具调用等待结果返回
|
||||
pending_async_tools: Arc<AtomicBool>,
|
||||
}
|
||||
|
||||
impl Conversation {
|
||||
/// ## 创建 Conversation
|
||||
///
|
||||
/// 1. 通过 `create_provider()` 工厂从配置创建 LLM 提供商
|
||||
/// 2. 初始化空的 `ChatSession`（消息历史在收到第一条用户消息后按需加载）
|
||||
/// 3. tool_executor 初始为 None，需要调用 `set_tool_executor` 设置
|
||||
pub fn new(config: ConversationConfig) -> Result<Self, String> {
|
||||
let provider = create_provider(&config)?;
|
||||
Ok(Self {
|
||||
@@ -33,12 +118,18 @@ impl Conversation {
|
||||
provider,
|
||||
session: Arc::new(Mutex::new(ChatSession::new())),
|
||||
tool_executor: None,
|
||||
pending_async_tools: Arc::new(AtomicBool::new(false)),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn set_tool_executor(&mut self, executor: ToolExecutor) {
|
||||
self.tool_executor = Some(executor);
|
||||
}
|
||||
/// 是否有异步工具调用正在等待结果
|
||||
#[allow(dead_code)]
|
||||
pub fn has_pending_async(&self) -> bool {
|
||||
self.pending_async_tools.load(Ordering::SeqCst)
|
||||
}
|
||||
pub fn session(&self) -> Arc<Mutex<ChatSession>> {
|
||||
Arc::clone(&self.session)
|
||||
}
|
||||
@@ -55,7 +146,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| {
|
||||
@@ -88,7 +184,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<ChatHandle, String> {
|
||||
// 空闲超时检查
|
||||
{
|
||||
@@ -112,11 +218,19 @@ impl Conversation {
|
||||
})
|
||||
}
|
||||
|
||||
/// 对话 + 工具循环（带上下文管理）
|
||||
pub async fn chat_with_tools(
|
||||
&self,
|
||||
user_message: String,
|
||||
) -> Result<(String, bool, Option<Usage>), String> {
|
||||
/// ## 对话 + 工具循环（主入口）
|
||||
///
|
||||
/// 这是最常用的方法：发送用户消息，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<ChatResult, String> {
|
||||
// ── 空闲超时检查 ──
|
||||
{
|
||||
let s = self.session.lock().await;
|
||||
@@ -128,7 +242,17 @@ impl Conversation {
|
||||
}
|
||||
|
||||
self.session.lock().await.add_user(user_message.clone());
|
||||
self.run_tool_loop().await
|
||||
}
|
||||
|
||||
/// 在异步工具结果返回后恢复工具循环（不追加用户消息）
|
||||
pub async fn resume_tool_loop(&self) -> Result<ChatResult, String> {
|
||||
self.pending_async_tools.store(false, Ordering::SeqCst);
|
||||
self.run_tool_loop().await
|
||||
}
|
||||
|
||||
/// 工具循环核心：调 LLM → 处理工具调用 → 循环直到无工具或需要异步等待
|
||||
async fn run_tool_loop(&self) -> Result<ChatResult, String> {
|
||||
let mut used_tools = false;
|
||||
let mut last_usage = None;
|
||||
let mut turn_count = 0u32;
|
||||
@@ -167,8 +291,8 @@ impl Conversation {
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(calls) = tool_calls {
|
||||
if !calls.is_empty() {
|
||||
if let Some(calls) = tool_calls
|
||||
&& !calls.is_empty() {
|
||||
used_tools = true;
|
||||
tracing::info!(
|
||||
"🔧 LLM 请求 {} 个工具: {}",
|
||||
@@ -185,6 +309,7 @@ impl Conversation {
|
||||
.await
|
||||
.add_assistant_tool_calls(calls.clone());
|
||||
|
||||
let mut has_async = false;
|
||||
for tc in &calls {
|
||||
let result = match &self.tool_executor {
|
||||
Some(exec) => {
|
||||
@@ -195,6 +320,14 @@ impl Conversation {
|
||||
}
|
||||
None => "工具执行器未配置".to_string(),
|
||||
};
|
||||
|
||||
if result.starts_with(ASYNC_MARKER) {
|
||||
has_async = true;
|
||||
tracing::info!("📦 {} → 异步入队", tc.function.name);
|
||||
// 不添加到 session，等待 ToolResult 异步回喂
|
||||
continue;
|
||||
}
|
||||
|
||||
tracing::info!("📦 {} → {:.150}", tc.function.name, result);
|
||||
self.session.lock().await.add_tool_result(
|
||||
&tc.id,
|
||||
@@ -202,9 +335,18 @@ impl Conversation {
|
||||
&result,
|
||||
);
|
||||
}
|
||||
|
||||
if has_async {
|
||||
self.pending_async_tools.store(true, Ordering::SeqCst);
|
||||
return Ok(ChatResult {
|
||||
reply: full_text,
|
||||
used_tools,
|
||||
usage: last_usage,
|
||||
has_pending_async: true,
|
||||
});
|
||||
}
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
// 无工具调用：记录回复，检查是否需要溢出摘要
|
||||
if !full_text.is_empty() {
|
||||
@@ -215,7 +357,7 @@ impl Conversation {
|
||||
{
|
||||
let s = self.session.lock().await;
|
||||
let recent = s.recent_messages();
|
||||
let estimated: u32 = recent.iter().map(|m| builder::estimate_tokens(m)).sum();
|
||||
let estimated: u32 = recent.iter().map(builder::estimate_tokens).sum();
|
||||
if estimated > s.token_budget {
|
||||
drop(s);
|
||||
builder::trigger_overflow_summary(&self.session, Some(&self.summarizer()))
|
||||
@@ -224,7 +366,12 @@ impl Conversation {
|
||||
}
|
||||
}
|
||||
|
||||
return Ok((full_text, used_tools, last_usage));
|
||||
return Ok(ChatResult {
|
||||
reply: full_text,
|
||||
used_tools,
|
||||
usage: last_usage,
|
||||
has_pending_async: false,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+50
-3
@@ -1,3 +1,30 @@
|
||||
//! ## 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": [...],
|
||||
//! "thinking": {...}
|
||||
//! }
|
||||
//! ```
|
||||
//!
|
||||
//! ### 流式处理
|
||||
//! 使用 `reqwest` 的 `bytes_stream()` 逐 chunk 读取 SSE 流，
|
||||
//! 通过 `parse_chat_chunk()` 解析每个 delta，通过 mpsc channel 发送给调用方。
|
||||
//! 工具调用的 delta 跨多个 chunk 拼接，在 Done 信号中一次性返回。
|
||||
|
||||
use crate::llm::provider::{
|
||||
LlmProvider, ParsedChunk, StreamReceiver, StreamSender, parse_chat_chunk,
|
||||
};
|
||||
@@ -7,6 +34,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,
|
||||
@@ -167,8 +215,8 @@ async fn stream_deepseek(
|
||||
}
|
||||
|
||||
// 检查最后的 buffer
|
||||
if !buf.trim().is_empty() {
|
||||
if let Some(parsed) = parse_chat_chunk(buf.trim()) {
|
||||
if !buf.trim().is_empty()
|
||||
&& let Some(parsed) = parse_chat_chunk(buf.trim()) {
|
||||
match parsed {
|
||||
ParsedChunk::Text(t) => full_text.push_str(&t),
|
||||
ParsedChunk::ToolCallDelta {
|
||||
@@ -195,7 +243,6 @@ async fn stream_deepseek(
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 构建 tool_calls
|
||||
let tool_calls: Option<Vec<ToolCall>> = if tool_call_builders.is_empty() {
|
||||
|
||||
+19
-1
@@ -1,7 +1,25 @@
|
||||
//! ## LLM 对话系统 —— DeepSeek API 集成
|
||||
//!
|
||||
//! 封装与 DeepSeek 大语言模型的交互：
|
||||
//!
|
||||
//! - `types` — 消息/角色/工具调用/用量等核心类型
|
||||
//! - `provider` — `LlmProvider` trait 抽象 + SSE 流式解析
|
||||
//! - `deepseek` — DeepSeek API 流式客户端实现
|
||||
//! - `conversation` — `Conversation` 对话管理器（工具循环、摘要、上下文管理）
|
||||
//!
|
||||
//! ### 架构
|
||||
//! ```text
|
||||
//! Conversation (对话管理器)
|
||||
//! ├── ChatSession (消息历史 + 摘要 + checkpoint)
|
||||
//! ├── LlmProvider (trait: chat_stream → StreamReceiver)
|
||||
//! │ └── DeepSeekProvider (实现: SSE 流式请求)
|
||||
//! └── ToolExecutor (回调: 执行 LLM 调用的工具)
|
||||
//! ```
|
||||
|
||||
pub mod conversation;
|
||||
pub mod deepseek;
|
||||
pub mod provider;
|
||||
pub mod types;
|
||||
|
||||
pub use conversation::{Conversation, ToolExecutor, DEFAULT_SYSTEM_PROMPT};
|
||||
pub use conversation::{ChatResult, Conversation, ToolExecutor, ASYNC_MARKER, DEFAULT_SYSTEM_PROMPT};
|
||||
pub use types::{ConversationConfig, Usage};
|
||||
|
||||
+65
-20
@@ -1,3 +1,21 @@
|
||||
//! ## LLM 提供商抽象（trait） + SSE 流式解析
|
||||
//!
|
||||
//! 定义了 LLM 层的核心抽象接口 `LlmProvider`，以及 SSE 流式响应的解析工具。
|
||||
//!
|
||||
//! ### 设计模式
|
||||
//! - `LlmProvider` trait — 所有 LLM 提供商需实现的接口
|
||||
//! - `create_provider()` — 工厂函数，根据配置创建对应的提供商实例
|
||||
//! - `parse_chat_chunk()` — SSE 流式 delta 解析器
|
||||
//!
|
||||
//! ### 流式架构
|
||||
//! ```text
|
||||
//! Provider.chat_stream() → mpsc::Receiver<StreamChunk>
|
||||
//! ├── StreamChunk::Text (delta 文本)
|
||||
//! ├── StreamChunk::Reasoning (思考内容)
|
||||
//! ├── StreamChunk::Done (完成信号 + 工具调用 + 用量)
|
||||
//! └── StreamChunk::Error (错误)
|
||||
//! ```
|
||||
|
||||
use crate::llm::types::{ConversationConfig, Message, StreamChunk};
|
||||
use async_trait::async_trait;
|
||||
use serde::Deserialize;
|
||||
@@ -8,7 +26,16 @@ use tokio::sync::mpsc;
|
||||
pub type StreamReceiver = mpsc::Receiver<StreamChunk>;
|
||||
pub type StreamSender = mpsc::Sender<StreamChunk>;
|
||||
|
||||
/// 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 +53,23 @@ pub trait LlmProvider: Send + Sync {
|
||||
|
||||
pub type BoxedProvider = Arc<dyn LlmProvider>;
|
||||
|
||||
/// 从配置创建恰当的提供商
|
||||
/// ## `create_provider` — 工厂函数：从配置创建恰当的 LLM 提供商
|
||||
///
|
||||
/// 目前只返回 DeepSeek，后续可以在这里根据配置 switch 到不同提供商。
|
||||
pub fn create_provider(_config: &ConversationConfig) -> Result<BoxedProvider, String> {
|
||||
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 +84,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<ParsedChunk> {
|
||||
if !line.starts_with("data: ") {
|
||||
return None;
|
||||
@@ -110,16 +159,16 @@ pub(crate) fn parse_chat_chunk(line: &str) -> Option<ParsedChunk> {
|
||||
};
|
||||
|
||||
// 提取 usage（可能在最后一个 chunk）
|
||||
if let Some(ref usage) = parsed.usage {
|
||||
if usage.total_tokens > 0 {
|
||||
if let Some(ref usage) = parsed.usage
|
||||
&& usage.total_tokens > 0 {
|
||||
return Some(ParsedChunk::Usage(usage.clone()));
|
||||
}
|
||||
}
|
||||
|
||||
for choice in parsed.choices {
|
||||
// SSE 流式响应中每个 chunk 只包含一个 choice 的 delta
|
||||
if let Some(choice) = parsed.choices.into_iter().next() {
|
||||
// 工具调用 delta
|
||||
if let Some(tool_calls) = &choice.delta.tool_calls {
|
||||
for tc in tool_calls {
|
||||
if let Some(tool_calls) = &choice.delta.tool_calls
|
||||
&& let Some(tc) = tool_calls.first() {
|
||||
let idx = tc.index.unwrap_or(0);
|
||||
let args = tc
|
||||
.function
|
||||
@@ -134,23 +183,19 @@ pub(crate) fn parse_chat_chunk(line: &str) -> Option<ParsedChunk> {
|
||||
arguments: args,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(reasoning) = &choice.delta.reasoning_content {
|
||||
if !reasoning.is_empty() {
|
||||
if let Some(reasoning) = &choice.delta.reasoning_content
|
||||
&& !reasoning.is_empty() {
|
||||
return Some(ParsedChunk::Reasoning(reasoning.clone()));
|
||||
}
|
||||
}
|
||||
if let Some(content) = &choice.delta.content {
|
||||
if !content.is_empty() {
|
||||
if let Some(content) = &choice.delta.content
|
||||
&& !content.is_empty() {
|
||||
return Some(ParsedChunk::Text(content.clone()));
|
||||
}
|
||||
}
|
||||
if let Some(reason) = &choice.finish_reason {
|
||||
if !reason.is_empty() {
|
||||
if let Some(reason) = &choice.finish_reason
|
||||
&& !reason.is_empty() {
|
||||
return Some(ParsedChunk::FinishReason(reason.clone()));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
None
|
||||
|
||||
+61
-12
@@ -1,7 +1,25 @@
|
||||
//! ## LLM 对话核心类型定义
|
||||
//!
|
||||
//! 定义了与 LLM API 交互的全部核心类型：
|
||||
//!
|
||||
//! - `Role` — 对话角色（System / User / Assistant / Tool）
|
||||
//! - `Message` — 单条消息（支持文本、工具调用、工具结果）
|
||||
//! - `ToolCall` — LLM 发起的工具调用请求
|
||||
//! - `StreamChunk` — 流式响应块（Text / Reasoning / Done / Error）
|
||||
//! - `Usage` — Token 用量统计
|
||||
//! - `ConversationConfig` — 对话配置参数
|
||||
//!
|
||||
//! 这些类型符合 OpenAI/DeepSeek 的 Chat Completion API 标准格式。
|
||||
|
||||
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 +29,17 @@ pub enum Role {
|
||||
Tool,
|
||||
}
|
||||
|
||||
// ─── 消息 ───
|
||||
|
||||
/// ## LLM 消息 —— 符合 OpenAI/DeepSeek Chat Completion API 格式
|
||||
///
|
||||
/// 这是整个 LLM 交互的核心类型。`Vec<Message>` 构成一次对话的全部上下文，
|
||||
/// 被序列化为 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 +75,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 +95,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 +121,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 +140,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,
|
||||
|
||||
+19
-3
@@ -1,8 +1,25 @@
|
||||
//! ## 日志初始化
|
||||
//!
|
||||
//! ### 两种模式
|
||||
//! - **终端模式** — 彩色输出到 stderr，适合开发调试
|
||||
//! - **文件模式** (`with_file=true`) — 同时输出到终端和日滚文件
|
||||
//! `~/.ias/logs/ias.log.YYYY-MM-DD`
|
||||
//!
|
||||
//! 日志级别通过 `RUST_LOG` 环境变量控制（默认 `info`）。
|
||||
|
||||
use std::path::PathBuf;
|
||||
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"));
|
||||
@@ -11,8 +28,8 @@ pub fn init_logger(log_dir: Option<&str>, with_file: bool) {
|
||||
.with_target(false)
|
||||
.with_ansi(true);
|
||||
|
||||
if with_file {
|
||||
if let Some(dir) = log_dir {
|
||||
if with_file
|
||||
&& let Some(dir) = log_dir {
|
||||
let dir_path = PathBuf::from(dir);
|
||||
std::fs::create_dir_all(&dir_path).ok();
|
||||
|
||||
@@ -29,7 +46,6 @@ pub fn init_logger(log_dir: Option<&str>, with_file: bool) {
|
||||
.init();
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
tracing_subscriber::registry()
|
||||
.with(env_filter)
|
||||
|
||||
+51
-25
@@ -1,3 +1,35 @@
|
||||
//! ## 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 <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;
|
||||
mod daemon;
|
||||
@@ -5,6 +37,7 @@ mod db;
|
||||
mod ipc;
|
||||
mod llm;
|
||||
mod logger;
|
||||
mod queue;
|
||||
mod scheduler;
|
||||
mod state;
|
||||
mod tools;
|
||||
@@ -15,11 +48,11 @@ use clap::Parser;
|
||||
use cli::{AmapAction, Cli, Commands, MemosAction, ToolCommand};
|
||||
use context::MemoryStore;
|
||||
use db::Database;
|
||||
use llm::{Conversation, ConversationConfig, DEFAULT_SYSTEM_PROMPT, ToolExecutor, Usage};
|
||||
use llm::{ChatResult, Conversation, ConversationConfig, DEFAULT_SYSTEM_PROMPT, ToolExecutor, Usage};
|
||||
use std::sync::Arc;
|
||||
use tools::approval::{ApprovalDecision, ApprovalManager};
|
||||
use tools::types::ExecutionContext;
|
||||
use tracing::{error, info};
|
||||
use tracing::{error, info, warn};
|
||||
use wechat::client::WeChatClient;
|
||||
|
||||
#[tokio::main]
|
||||
@@ -105,6 +138,8 @@ async fn main() {
|
||||
cmd_daemon(sock, &database, &file_state, &memory_store).await;
|
||||
}
|
||||
Commands::Worker { sock } => {
|
||||
warn!("Worker 模式已废弃，请使用 daemon 模式");
|
||||
#[allow(deprecated)]
|
||||
if let Err(e) = worker::run(&sock).await {
|
||||
error!("Worker 失败: {}", e);
|
||||
}
|
||||
@@ -336,7 +371,6 @@ async fn cmd_listen(
|
||||
.send_text(&uid, &msg, None)
|
||||
.await
|
||||
.map(|_| ())
|
||||
.map_err(|e| e)
|
||||
})
|
||||
})
|
||||
};
|
||||
@@ -501,8 +535,8 @@ async fn listen_loop(
|
||||
info!("收到消息 from={}: {}", from, text);
|
||||
|
||||
// === 入库：收到的消息 ===
|
||||
if let Some(db) = database {
|
||||
if let Err(e) = db::models::insert_chat_record(
|
||||
if let Some(db) = database
|
||||
&& let Err(e) = db::models::insert_chat_record(
|
||||
db.pool(),
|
||||
"inbound",
|
||||
from,
|
||||
@@ -516,7 +550,6 @@ async fn listen_loop(
|
||||
{
|
||||
error!("保存聊天记录失败: {}", e);
|
||||
}
|
||||
}
|
||||
|
||||
// Echo 模式
|
||||
if echo {
|
||||
@@ -524,8 +557,8 @@ async fn listen_loop(
|
||||
Ok(sent_id) => {
|
||||
info!("回显成功 msg_id={}", sent_id);
|
||||
// === 入库：发送的回显 ===
|
||||
if let Some(db) = database {
|
||||
if let Err(e) = db::models::insert_chat_record(
|
||||
if let Some(db) = database
|
||||
&& let Err(e) = db::models::insert_chat_record(
|
||||
db.pool(),
|
||||
"outbound",
|
||||
from,
|
||||
@@ -539,7 +572,6 @@ async fn listen_loop(
|
||||
{
|
||||
error!("保存聊天记录失败: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => error!("回显失败: {}", e),
|
||||
}
|
||||
@@ -577,8 +609,8 @@ async fn listen_loop(
|
||||
}
|
||||
|
||||
// LLM 回复（异步执行，避免审批阻塞主循环）
|
||||
if enable_llm {
|
||||
if let Some(conv) = &conversation {
|
||||
if enable_llm
|
||||
&& let Some(conv) = &conversation {
|
||||
let conv = conv.clone();
|
||||
let client = client.clone();
|
||||
let database = database.clone();
|
||||
@@ -604,8 +636,8 @@ async fn listen_loop(
|
||||
{
|
||||
Ok(sent_id) => {
|
||||
info!("LLM 回复成功 msg_id={}", sent_id);
|
||||
if let Some(db) = database {
|
||||
if let Err(e) = db::models::insert_chat_record(
|
||||
if let Some(db) = database
|
||||
&& let Err(e) = db::models::insert_chat_record(
|
||||
db.pool(),
|
||||
"outbound",
|
||||
&from_owned,
|
||||
@@ -619,13 +651,11 @@ async fn listen_loop(
|
||||
{
|
||||
error!("保存聊天记录失败: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => error!("发送 LLM 回复失败: {}", e),
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
@@ -665,7 +695,7 @@ async fn generate_reply_with_tools(
|
||||
user_text: String,
|
||||
db: &Option<Arc<Database>>,
|
||||
) -> Result<String, String> {
|
||||
let (reply, _used_tools, usage) = conv.chat_with_tools(user_text).await?;
|
||||
let ChatResult { reply, used_tools: _used_tools, usage, .. } = conv.chat_with_tools(user_text).await?;
|
||||
if let Some(u) = &usage {
|
||||
info!(
|
||||
"📊 Token: {} 总 / {} 提示 / {} 生成 | 缓存命中 {} | 缓存未命中 {}",
|
||||
@@ -687,8 +717,8 @@ async fn store_usage(
|
||||
provider: &str,
|
||||
u: &Usage,
|
||||
) {
|
||||
if let Some(db) = db {
|
||||
if let Err(e) = db::models::insert_llm_usage(
|
||||
if let Some(db) = db
|
||||
&& let Err(e) = db::models::insert_llm_usage(
|
||||
db.pool(),
|
||||
user_id,
|
||||
model,
|
||||
@@ -702,7 +732,6 @@ async fn store_usage(
|
||||
{
|
||||
tracing::error!("存储 LLM 用量失败: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn cmd_whoami(database: &Option<Arc<Database>>, file_state: &state::StateManager) {
|
||||
@@ -828,8 +857,8 @@ async fn cmd_send(
|
||||
println!("✅ 消息已发送, msg_id={}", msg_id);
|
||||
|
||||
// 入库：发送的消息
|
||||
if let Some(db) = database {
|
||||
if let Err(e) = db::models::insert_chat_record(
|
||||
if let Some(db) = database
|
||||
&& let Err(e) = db::models::insert_chat_record(
|
||||
db.pool(),
|
||||
"outbound",
|
||||
to,
|
||||
@@ -843,7 +872,6 @@ async fn cmd_send(
|
||||
{
|
||||
error!("保存聊天记录失败: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => error!("发送失败: {}", e),
|
||||
}
|
||||
@@ -863,9 +891,7 @@ async fn builtin_approve(ctx: &ExecutionContext, name: &str) -> Result<bool, Str
|
||||
name, code
|
||||
);
|
||||
if let Some(ref send) = ctx.send_wechat {
|
||||
if let Err(e) = send(&ctx.user_id, &msg).await {
|
||||
return Err(e);
|
||||
}
|
||||
send(&ctx.user_id, &msg).await?
|
||||
}
|
||||
|
||||
match tokio::time::timeout(std::time::Duration::from_secs(300), rx).await {
|
||||
|
||||
@@ -0,0 +1,582 @@
|
||||
//! ## 多渠道公平轮转消息队列 (MessageQueue)
|
||||
//!
|
||||
//! ### 设计目标
|
||||
//!
|
||||
//! 多用户场景下保证**公平性**：如果一个用户连续发了很多消息，不会让其他用户饿死。
|
||||
//! 出队顺序是**轮转（round-robin）**的：
|
||||
//! - A 发了 3 条、B 发了 1 条 → 处理顺序 A1, B1, A2, A3
|
||||
//!
|
||||
//! ### 消息关联追踪
|
||||
//!
|
||||
//! 每条消息有一个 `correlation_id`（Uuid），用来追踪一次用户请求的完整生命周期：
|
||||
//! 用户消息 → LLM 处理 → 工具调用 → 工具结果回喂 → 最终回复发送。
|
||||
//! 这样即使多用户并发，回复也不会串。
|
||||
//!
|
||||
//! ### 架构位置
|
||||
//!
|
||||
//! ```text
|
||||
//! receive_messages → enqueue PipelineMessage
|
||||
//! ┌─────────────────────────────┐
|
||||
//! │ MessageQueue │
|
||||
//! │ [User A] [User B] │ ← 每用户独立 VecDeque
|
||||
//! │ pending: [A, B, A] │ ← 公平轮转索引
|
||||
//! └─────────────────────────────┘
|
||||
//! QueueRunner::run() → dequeue → 按 kind 路由到 mpsc 通道
|
||||
//! ├─ llm_tx (UserMessage, ToolResult)
|
||||
//! ├─ tool_tx (ToolCall, ApprovalRequest)
|
||||
//! └─ send_tx (LLMReply)
|
||||
//! ```
|
||||
|
||||
use crate::channel::ChannelId;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::{HashMap, VecDeque};
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use tokio::sync::Notify;
|
||||
use uuid::Uuid;
|
||||
|
||||
// ─── 管线消息类型 ───
|
||||
|
||||
/// 消息种类 —— 区分不同消费路由
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub enum MessageKind {
|
||||
/// 从外部收到的用户消息 → 交给 LLM Consumer
|
||||
UserMessage {
|
||||
text: String,
|
||||
message_id: String,
|
||||
context_token: Option<String>,
|
||||
/// 审批通过后携带的已批准工具名
|
||||
approved_tool: Option<String>,
|
||||
},
|
||||
/// LLM 生成的回复文本 → 发给 Send Consumer
|
||||
LLMReply {
|
||||
text: String,
|
||||
context_token: Option<String>,
|
||||
},
|
||||
/// LLM 请求的工具调用 → 发给 Tool Consumer
|
||||
ToolCall {
|
||||
session_id: Uuid,
|
||||
tool_call_id: String,
|
||||
tool_name: String,
|
||||
arguments: String,
|
||||
context_token: Option<String>,
|
||||
/// 该工具调用是否已被用户批准（跳过审批流程）
|
||||
approved: bool,
|
||||
},
|
||||
/// 工具执行结果 → 回喂 LLM Consumer
|
||||
ToolResult {
|
||||
session_id: Uuid,
|
||||
tool_call_id: String,
|
||||
tool_name: String,
|
||||
result: String,
|
||||
context_token: Option<String>,
|
||||
},
|
||||
/// 高风险工具 → 需用户审批
|
||||
ApprovalRequest {
|
||||
session_id: Uuid,
|
||||
tool_name: String,
|
||||
tool_args: String,
|
||||
/// 审批后要追加到用户消息中的文本
|
||||
approval_prompt: String,
|
||||
},
|
||||
/// 定时任务到期 → 交给 LLM Consumer
|
||||
ScheduledTask {
|
||||
text: String,
|
||||
task_name: String,
|
||||
},
|
||||
}
|
||||
|
||||
/// 管线消息 —— 携带路由元数据的统一消息类型
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct PipelineMessage {
|
||||
/// 来源/目标渠道（平台 + 用户），全程不变
|
||||
pub channel: ChannelId,
|
||||
/// 关联 ID：追踪一次用户请求的完整链路（用户消息→LLM→工具→回喂→发送）
|
||||
pub correlation_id: Uuid,
|
||||
/// 消息种类
|
||||
pub kind: MessageKind,
|
||||
}
|
||||
|
||||
impl PipelineMessage {
|
||||
/// 便捷：渠道中的 user_id
|
||||
pub fn user_id(&self) -> &str { &self.channel.user_id }
|
||||
/// 便捷：渠道中的 platform
|
||||
#[allow(dead_code)]
|
||||
pub fn platform(&self) -> &str { &self.channel.platform }
|
||||
/// 快速构造一条用户消息
|
||||
#[allow(dead_code)]
|
||||
pub fn user_message(
|
||||
user_id: impl Into<String>,
|
||||
text: impl Into<String>,
|
||||
message_id: impl Into<String>,
|
||||
context_token: Option<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
channel: ChannelId::wechat(user_id),
|
||||
correlation_id: Uuid::new_v4(),
|
||||
kind: MessageKind::UserMessage {
|
||||
text: text.into(),
|
||||
message_id: message_id.into(),
|
||||
context_token,
|
||||
approved_tool: None,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/// 快速构造一条 LLM 回复
|
||||
pub fn llm_reply(
|
||||
channel: ChannelId,
|
||||
correlation_id: Uuid,
|
||||
text: impl Into<String>,
|
||||
context_token: Option<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
channel,
|
||||
correlation_id,
|
||||
kind: MessageKind::LLMReply {
|
||||
text: text.into(),
|
||||
context_token,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/// 快速构造一条工具调用
|
||||
pub fn tool_call(
|
||||
channel: ChannelId,
|
||||
correlation_id: Uuid,
|
||||
session_id: Uuid,
|
||||
tool_call_id: impl Into<String>,
|
||||
tool_name: impl Into<String>,
|
||||
arguments: impl Into<String>,
|
||||
context_token: Option<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
channel,
|
||||
correlation_id,
|
||||
kind: MessageKind::ToolCall {
|
||||
session_id,
|
||||
tool_call_id: tool_call_id.into(),
|
||||
tool_name: tool_name.into(),
|
||||
arguments: arguments.into(),
|
||||
context_token,
|
||||
approved: false,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/// 快速构造一条已批准的工具调用（跳过审批）
|
||||
pub fn approved_tool_call(
|
||||
channel: ChannelId,
|
||||
correlation_id: Uuid,
|
||||
session_id: Uuid,
|
||||
tool_call_id: impl Into<String>,
|
||||
tool_name: impl Into<String>,
|
||||
arguments: impl Into<String>,
|
||||
context_token: Option<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
channel,
|
||||
correlation_id,
|
||||
kind: MessageKind::ToolCall {
|
||||
session_id,
|
||||
tool_call_id: tool_call_id.into(),
|
||||
tool_name: tool_name.into(),
|
||||
arguments: arguments.into(),
|
||||
context_token,
|
||||
approved: true,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/// 快速构造一条工具结果
|
||||
pub fn tool_result(
|
||||
channel: ChannelId,
|
||||
correlation_id: Uuid,
|
||||
session_id: Uuid,
|
||||
tool_call_id: impl Into<String>,
|
||||
tool_name: impl Into<String>,
|
||||
result: impl Into<String>,
|
||||
context_token: Option<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
channel,
|
||||
correlation_id,
|
||||
kind: MessageKind::ToolResult {
|
||||
session_id,
|
||||
tool_call_id: tool_call_id.into(),
|
||||
tool_name: tool_name.into(),
|
||||
result: result.into(),
|
||||
context_token,
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ─── 公平轮转队列 ───
|
||||
|
||||
/// 多渠道公平轮转消息队列
|
||||
///
|
||||
/// 按 ChannelId（platform + user_id）分队列，轮转出队保证公平。
|
||||
/// 出队后由 QueueRunner 按 MessageKind 路由到对应 mpsc 通道。
|
||||
pub struct MessageQueue {
|
||||
/// 每渠道独立队列
|
||||
queues: HashMap<ChannelId, VecDeque<PipelineMessage>>,
|
||||
/// 有待处理消息的渠道（轮转顺序）
|
||||
pending: VecDeque<ChannelId>,
|
||||
/// 空→非空时触发 Notify 唤醒消费者
|
||||
signal: Arc<Notify>,
|
||||
/// 定时唤醒间隔（兜底用）
|
||||
poll_interval: Duration,
|
||||
}
|
||||
|
||||
impl MessageQueue {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
queues: HashMap::new(),
|
||||
pending: VecDeque::new(),
|
||||
signal: Arc::new(Notify::new()),
|
||||
poll_interval: Duration::from_millis(200),
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub fn with_poll_interval(mut self, d: Duration) -> Self {
|
||||
self.poll_interval = d;
|
||||
self
|
||||
}
|
||||
|
||||
/// 入队。如果该渠道之前为空 → 加入 pending 尾部 + 通知
|
||||
pub fn enqueue(&mut self, item: PipelineMessage) {
|
||||
let ch = item.channel.clone();
|
||||
let q = self.queues.entry(ch.clone()).or_default();
|
||||
let was_empty = q.is_empty();
|
||||
q.push_back(item);
|
||||
if was_empty {
|
||||
self.pending.push_back(ch);
|
||||
self.signal.notify_one();
|
||||
}
|
||||
}
|
||||
|
||||
/// 公平轮转出队
|
||||
///
|
||||
/// 每次调用切换到下一个有消息的渠道：
|
||||
/// - 从 pending 头部取出一个渠道
|
||||
/// - pop 其一条消息
|
||||
/// - 该渠道还有消息 → 放回 pending 尾部
|
||||
pub fn dequeue(&mut self) -> Option<PipelineMessage> {
|
||||
while let Some(ch) = self.pending.pop_front() {
|
||||
if let Some(q) = self.queues.get_mut(&ch)
|
||||
&& let Some(item) = q.pop_front() {
|
||||
if !q.is_empty() {
|
||||
self.pending.push_back(ch);
|
||||
} else {
|
||||
self.queues.remove(&ch);
|
||||
}
|
||||
return Some(item);
|
||||
}
|
||||
self.queues.remove(&ch);
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
pub fn signal(&self) -> Arc<Notify> {
|
||||
Arc::clone(&self.signal)
|
||||
}
|
||||
#[allow(dead_code)]
|
||||
pub fn poll_interval(&self) -> Duration {
|
||||
self.poll_interval
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub fn has_pending(&self, ch: &ChannelId) -> bool {
|
||||
self.queues.get(ch).is_some_and(|q| !q.is_empty())
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub fn pending_count(&self) -> usize {
|
||||
self.pending.len()
|
||||
}
|
||||
|
||||
/// 清空指定渠道的全部待处理消息
|
||||
#[allow(dead_code)]
|
||||
pub fn drain_channel(&mut self, ch: &ChannelId) -> Vec<PipelineMessage> {
|
||||
let drained = self
|
||||
.queues
|
||||
.remove(ch)
|
||||
.map(|mut q| q.drain(..).collect())
|
||||
.unwrap_or_default();
|
||||
self.pending.retain(|c| c != ch);
|
||||
drained
|
||||
}
|
||||
}
|
||||
|
||||
// ─── 消费者通道标识 ───
|
||||
|
||||
/// 消息路由目标
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum ConsumerTarget {
|
||||
/// LLM 消费者（处理 UserMessage, ToolResult）
|
||||
LLM,
|
||||
/// 工具消费者（处理 ToolCall, ApprovalRequest）
|
||||
Tool,
|
||||
/// 发送消费者（处理 LLMReply）
|
||||
Send,
|
||||
}
|
||||
|
||||
impl PipelineMessage {
|
||||
/// 判断这条消息应该路由到哪个消费者
|
||||
pub fn target(&self) -> ConsumerTarget {
|
||||
match self.kind {
|
||||
MessageKind::UserMessage { .. }
|
||||
| MessageKind::ToolResult { .. }
|
||||
| MessageKind::ScheduledTask { .. } => ConsumerTarget::LLM,
|
||||
MessageKind::ToolCall { .. } | MessageKind::ApprovalRequest { .. } => {
|
||||
ConsumerTarget::Tool
|
||||
}
|
||||
MessageKind::LLMReply { .. } => ConsumerTarget::Send,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ─── 测试 ───
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::queue::QueueRunner;
|
||||
|
||||
fn make_user_msg(user: &str, text: &str) -> PipelineMessage {
|
||||
PipelineMessage::user_message(user, text, "", None)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_target_routing() {
|
||||
let msg = PipelineMessage::user_message("u1", "hello", "m1", None);
|
||||
assert_eq!(msg.target(), ConsumerTarget::LLM);
|
||||
|
||||
let msg = PipelineMessage::llm_reply(ChannelId::wechat("u1"), Uuid::new_v4(), "hi", None);
|
||||
assert_eq!(msg.target(), ConsumerTarget::Send);
|
||||
|
||||
let msg = PipelineMessage::tool_call(
|
||||
ChannelId::wechat("u1"),
|
||||
Uuid::new_v4(),
|
||||
Uuid::new_v4(),
|
||||
"tc1",
|
||||
"query_weather",
|
||||
r#"{"location":"北京"}"#,
|
||||
None,
|
||||
);
|
||||
assert_eq!(msg.target(), ConsumerTarget::Tool);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_fair_round_robin() {
|
||||
let mut q = MessageQueue::new();
|
||||
q.enqueue(make_user_msg("A", "a1"));
|
||||
q.enqueue(make_user_msg("A", "a2"));
|
||||
q.enqueue(make_user_msg("A", "a3"));
|
||||
q.enqueue(make_user_msg("B", "b1"));
|
||||
assert_eq!(q.dequeue().unwrap().channel.user_id, "A");
|
||||
assert_eq!(q.dequeue().unwrap().channel.user_id, "B");
|
||||
assert_eq!(q.dequeue().unwrap().channel.user_id, "A");
|
||||
assert_eq!(q.dequeue().unwrap().channel.user_id, "A");
|
||||
assert!(q.dequeue().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_single_user() {
|
||||
let mut q = MessageQueue::new();
|
||||
q.enqueue(make_user_msg("X", "hi"));
|
||||
q.enqueue(make_user_msg("X", "hello"));
|
||||
assert_eq!(q.dequeue().unwrap().channel.user_id, "X");
|
||||
assert_eq!(q.dequeue().unwrap().channel.user_id, "X");
|
||||
assert!(q.dequeue().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_empty_queue() {
|
||||
let mut q = MessageQueue::new();
|
||||
assert!(q.dequeue().is_none());
|
||||
assert_eq!(q.pending_count(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_user_id_preserved() {
|
||||
let msg = PipelineMessage::user_message("wx_user_1", "天气?", "m1", None);
|
||||
assert_eq!(msg.channel.user_id, "wx_user_1");
|
||||
|
||||
let reply = PipelineMessage::llm_reply(ChannelId::wechat("wx_user_1"), msg.correlation_id, "晴天", None);
|
||||
assert_eq!(reply.channel.user_id, "wx_user_1");
|
||||
assert_eq!(reply.correlation_id, msg.correlation_id);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_mixed_item_types() {
|
||||
let mut q = MessageQueue::new();
|
||||
q.enqueue(PipelineMessage::user_message("U", "user msg", "1", None));
|
||||
q.enqueue(PipelineMessage {
|
||||
channel: ChannelId::wechat("U"),
|
||||
correlation_id: Uuid::new_v4(),
|
||||
kind: MessageKind::ScheduledTask {
|
||||
text: "task result".into(),
|
||||
task_name: "reminder".into(),
|
||||
},
|
||||
});
|
||||
assert!(matches!(
|
||||
q.dequeue().unwrap().kind,
|
||||
MessageKind::UserMessage { .. }
|
||||
));
|
||||
assert!(matches!(
|
||||
q.dequeue().unwrap().kind,
|
||||
MessageKind::ScheduledTask { .. }
|
||||
));
|
||||
}
|
||||
|
||||
// ─── 端到端关联测试 ───
|
||||
|
||||
/// 验证 user_id 从用户消息到 LLM 回复全程不变
|
||||
#[test]
|
||||
fn test_user_id_pipeline_preserved() {
|
||||
let user_id = "wx_user_1";
|
||||
|
||||
// 用户消息入队
|
||||
let user_msg = PipelineMessage::user_message(user_id, "天气怎么样?", "msg1", None);
|
||||
assert_eq!(user_msg.channel.user_id, user_id);
|
||||
let correlation_id = user_msg.correlation_id;
|
||||
|
||||
// LLM 回复（模拟 LLM Consumer 生成）
|
||||
let reply = PipelineMessage::llm_reply(ChannelId::wechat(user_id), correlation_id, "今天晴天", None);
|
||||
assert_eq!(reply.channel.user_id, user_id);
|
||||
assert_eq!(reply.correlation_id, correlation_id);
|
||||
|
||||
// 工具调用（模拟 LLM Consumer 生成）
|
||||
let tool_call = PipelineMessage::tool_call(
|
||||
ChannelId::wechat(user_id),
|
||||
correlation_id,
|
||||
Uuid::new_v4(),
|
||||
"tc1",
|
||||
"query_weather",
|
||||
r#"{"location":"北京"}"#,
|
||||
None,
|
||||
);
|
||||
assert_eq!(tool_call.channel.user_id, user_id);
|
||||
assert_eq!(tool_call.correlation_id, correlation_id);
|
||||
|
||||
// 工具结果回喂（模拟 Tool Consumer 生成）
|
||||
let tool_result = PipelineMessage::tool_result(
|
||||
ChannelId::wechat(user_id),
|
||||
correlation_id,
|
||||
Uuid::new_v4(),
|
||||
"tc1",
|
||||
"query_weather",
|
||||
"25°C",
|
||||
None,
|
||||
);
|
||||
assert_eq!(tool_result.channel.user_id, user_id);
|
||||
assert_eq!(tool_result.correlation_id, correlation_id);
|
||||
}
|
||||
|
||||
/// 验证多个用户的消息不会串
|
||||
#[test]
|
||||
fn test_multi_user_isolation() {
|
||||
let mut q = MessageQueue::new();
|
||||
|
||||
q.enqueue(PipelineMessage::user_message("u1", "你好", "1", None));
|
||||
q.enqueue(PipelineMessage::user_message("u2", "hello", "2", None));
|
||||
q.enqueue(PipelineMessage::user_message("u1", "天气?", "3", None));
|
||||
|
||||
// 按公平轮转出队，验证 user_id 正确
|
||||
let msg1 = q.dequeue().unwrap();
|
||||
assert_eq!(msg1.channel.user_id, "u1");
|
||||
let msg2 = q.dequeue().unwrap();
|
||||
assert_eq!(msg2.channel.user_id, "u2");
|
||||
let msg3 = q.dequeue().unwrap();
|
||||
assert_eq!(msg3.channel.user_id, "u1");
|
||||
}
|
||||
|
||||
/// 验证路由目标正确
|
||||
#[test]
|
||||
fn test_consumer_target_routing() {
|
||||
// UserMessage → LLM
|
||||
let m = PipelineMessage::user_message("u1", "hi", "1", None);
|
||||
assert_eq!(m.target(), ConsumerTarget::LLM);
|
||||
|
||||
// ToolResult → LLM
|
||||
let m = PipelineMessage::tool_result(ChannelId::wechat("u1"), Uuid::new_v4(), Uuid::new_v4(), "tc1", "test", "ok", None);
|
||||
assert_eq!(m.target(), ConsumerTarget::LLM);
|
||||
|
||||
// ScheduledTask → LLM
|
||||
let m = PipelineMessage {
|
||||
channel: ChannelId::wechat("u1"),
|
||||
correlation_id: Uuid::new_v4(),
|
||||
kind: MessageKind::ScheduledTask { text: "task".into(), task_name: "t1".into() },
|
||||
};
|
||||
assert_eq!(m.target(), ConsumerTarget::LLM);
|
||||
|
||||
// ToolCall → Tool
|
||||
let m = PipelineMessage::tool_call(ChannelId::wechat("u1"), Uuid::new_v4(), Uuid::new_v4(), "tc1", "weather", "{}", None);
|
||||
assert_eq!(m.target(), ConsumerTarget::Tool);
|
||||
|
||||
// ApprovalRequest → Tool
|
||||
let m = PipelineMessage {
|
||||
channel: ChannelId::wechat("u1"),
|
||||
correlation_id: Uuid::new_v4(),
|
||||
kind: MessageKind::ApprovalRequest {
|
||||
session_id: Uuid::new_v4(),
|
||||
tool_name: "high_risk_tool".into(),
|
||||
tool_args: "{}".into(),
|
||||
approval_prompt: "".into(),
|
||||
},
|
||||
};
|
||||
assert_eq!(m.target(), ConsumerTarget::Tool);
|
||||
|
||||
// LLMReply → Send
|
||||
let m = PipelineMessage::llm_reply(ChannelId::wechat("u1"), Uuid::new_v4(), "reply", None);
|
||||
assert_eq!(m.target(), ConsumerTarget::Send);
|
||||
}
|
||||
|
||||
/// 验证 QueueRunner 的路由逻辑：入队 → 路由到正确通道
|
||||
#[tokio::test]
|
||||
async fn test_queue_runner_routing() {
|
||||
let (runner, mut channels) = QueueRunner::new(16);
|
||||
let enqueue = runner.enqueue_handle();
|
||||
|
||||
// 启动 runner（在后台运行）
|
||||
let handle = tokio::spawn(async move {
|
||||
runner.run().await;
|
||||
});
|
||||
|
||||
// 入队一个 UserMessage → 应到 llm_rx
|
||||
let user_msg = PipelineMessage::user_message("u1", "hello", "m1", None);
|
||||
enqueue.enqueue(user_msg).await;
|
||||
|
||||
// 入队一个 LLMReply → 应到 send_rx
|
||||
let reply_msg = PipelineMessage::llm_reply(ChannelId::wechat("u1"), Uuid::new_v4(), "reply", None);
|
||||
enqueue.enqueue(reply_msg).await;
|
||||
|
||||
// 验证 llm_rx 收到 UserMessage
|
||||
let llm_msg = tokio::time::timeout(
|
||||
std::time::Duration::from_secs(2),
|
||||
channels.llm_rx.recv(),
|
||||
)
|
||||
.await
|
||||
.expect("llm_rx 超时")
|
||||
.expect("llm_rx 通道关闭");
|
||||
assert_eq!(llm_msg.channel.user_id, "u1");
|
||||
assert!(matches!(llm_msg.kind, MessageKind::UserMessage { .. }));
|
||||
|
||||
// 验证 send_rx 收到 LLMReply
|
||||
let send_msg = tokio::time::timeout(
|
||||
std::time::Duration::from_secs(2),
|
||||
channels.send_rx.recv(),
|
||||
)
|
||||
.await
|
||||
.expect("send_rx 超时")
|
||||
.expect("send_rx 通道关闭");
|
||||
assert_eq!(send_msg.channel.user_id, "u1");
|
||||
assert!(matches!(send_msg.kind, MessageKind::LLMReply { .. }));
|
||||
|
||||
// 关闭 runner
|
||||
handle.abort();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,19 @@
|
||||
//! ## 消息队列 —— 多渠道公平轮转 + 三消费者路由
|
||||
//!
|
||||
//! 这是 daemon 架构的消息管线核心：
|
||||
//!
|
||||
//! - `message_queue` — 公平轮转队列（按用户轮转出队）
|
||||
//! - `runner` — 队列运行时（出队 → 按 MessageKind 路由到对应 mpsc 通道）
|
||||
//!
|
||||
//! ### 消费者路由
|
||||
//! ```text
|
||||
//! MessageKind::UserMessage / ToolResult / ScheduledTask → LLM Consumer
|
||||
//! MessageKind::ToolCall / ApprovalRequest → Tool Consumer
|
||||
//! MessageKind::LLMReply → Send Consumer
|
||||
//! ```
|
||||
|
||||
pub mod message_queue;
|
||||
pub mod runner;
|
||||
|
||||
pub use message_queue::{ConsumerTarget, MessageKind, MessageQueue, PipelineMessage};
|
||||
pub use runner::{ConsumerChannels, EnqueueHandle, QueueRunner};
|
||||
@@ -0,0 +1,219 @@
|
||||
//! ## QueueRunner —— 消息队列运行时（路由调度器）
|
||||
//!
|
||||
//! 从公平轮转 `MessageQueue` 中出队消息，然后按 `MessageKind` 路由到对应的
|
||||
//! tokio mpsc 通道，由注册的消费者 task 处理。
|
||||
//!
|
||||
//! ### 角色定位
|
||||
//!
|
||||
//! QueueRunner 是整个消息管线的**调度核心**：
|
||||
//! 1. 持有 `MessageQueue`（公平轮转）+ 3 个 mpsc 通道
|
||||
//! 2. 在 `run()` 循环中不断出队 → 按类型路由 → 投递到正确的消费者
|
||||
//! 3. 支持优雅关闭（AtomicBool 信号）
|
||||
//!
|
||||
//! ### 架构
|
||||
//!
|
||||
//! ```text
|
||||
//! ┌──────────────┐
|
||||
//! enqueue ──────→ │ MessageQueue │ ← 公平轮转
|
||||
//! └──────┬───────┘
|
||||
//! │ dequeue
|
||||
//! ▼
|
||||
//! ┌──────────────┐
|
||||
//! │ QueueRunner │ ← 路由循环
|
||||
//! └──┬───┬───┬──┘
|
||||
//! ┌─┘ │ └─┐
|
||||
//! ▼ ▼ ▼
|
||||
//! llm_tx tool_tx send_tx
|
||||
//! │ │ │
|
||||
//! ┌────┘ ┌─┘ ┌──┘
|
||||
//! ▼ ▼ ▼
|
||||
//! LLMCons. ToolCons. SendCons.
|
||||
//! ```
|
||||
|
||||
use crate::queue::{ConsumerTarget, MessageQueue, PipelineMessage};
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::{mpsc, Mutex, Notify};
|
||||
use tracing::{error, info, warn};
|
||||
|
||||
/// 消费者通道集 —— 创建 QueueRunner 时获得接收端
|
||||
pub struct ConsumerChannels {
|
||||
pub llm_rx: mpsc::Receiver<PipelineMessage>,
|
||||
pub tool_rx: mpsc::Receiver<PipelineMessage>,
|
||||
pub send_rx: mpsc::Receiver<PipelineMessage>,
|
||||
}
|
||||
|
||||
/// 队列运行时 —— 负责从公平轮转队列出队并按类型路由
|
||||
pub struct QueueRunner {
|
||||
/// 公平轮转队列
|
||||
queue: Arc<Mutex<MessageQueue>>,
|
||||
/// LLM 消费者通道
|
||||
llm_tx: mpsc::Sender<PipelineMessage>,
|
||||
/// 工具消费者通道
|
||||
tool_tx: mpsc::Sender<PipelineMessage>,
|
||||
/// 发送消费者通道
|
||||
send_tx: mpsc::Sender<PipelineMessage>,
|
||||
/// 关闭信号
|
||||
shutdown: Arc<AtomicBool>,
|
||||
/// 队列信号（空→非空通知）
|
||||
signal: Arc<Notify>,
|
||||
/// 轮询间隔
|
||||
poll_interval: std::time::Duration,
|
||||
}
|
||||
|
||||
impl QueueRunner {
|
||||
/// 创建 QueueRunner 和对应的消费者通道
|
||||
///
|
||||
/// `channel_size`: 每个 mpsc 通道的缓冲区大小
|
||||
pub fn new(channel_size: usize) -> (Self, ConsumerChannels) {
|
||||
let (llm_tx, llm_rx) = mpsc::channel(channel_size);
|
||||
let (tool_tx, tool_rx) = mpsc::channel(channel_size);
|
||||
let (send_tx, send_rx) = mpsc::channel(channel_size);
|
||||
|
||||
let queue = MessageQueue::new();
|
||||
let signal = queue.signal();
|
||||
let queue = Arc::new(Mutex::new(queue));
|
||||
|
||||
let runner = Self {
|
||||
queue,
|
||||
llm_tx,
|
||||
tool_tx,
|
||||
send_tx,
|
||||
shutdown: Arc::new(AtomicBool::new(false)),
|
||||
signal,
|
||||
poll_interval: std::time::Duration::from_millis(200),
|
||||
};
|
||||
|
||||
let channels = ConsumerChannels {
|
||||
llm_rx,
|
||||
tool_rx,
|
||||
send_rx,
|
||||
};
|
||||
|
||||
(runner, channels)
|
||||
}
|
||||
|
||||
/// 获取队列的写句柄（用于外部入队）
|
||||
pub fn enqueue_handle(&self) -> EnqueueHandle {
|
||||
EnqueueHandle {
|
||||
queue: self.queue.clone(),
|
||||
shutdown: self.shutdown.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 获取关闭信号句柄
|
||||
pub fn shutdown_handle(&self) -> ShutdownHandle {
|
||||
ShutdownHandle {
|
||||
shutdown: self.shutdown.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 启动路由主循环（阻塞，应在 tokio::spawn 中运行）
|
||||
pub async fn run(&self) {
|
||||
info!("QueueRunner 已启动");
|
||||
loop {
|
||||
// 检查关闭信号
|
||||
if self.shutdown.load(Ordering::Relaxed) {
|
||||
info!("QueueRunner 收到关闭信号，停止路由");
|
||||
break;
|
||||
}
|
||||
|
||||
// 出队
|
||||
let item = {
|
||||
let mut q = self.queue.lock().await;
|
||||
q.dequeue()
|
||||
};
|
||||
|
||||
if let Some(msg) = item {
|
||||
// 预提取字段（避免 borrow + move 冲突）
|
||||
let target = msg.target();
|
||||
let user_id = msg.channel.user_id.clone();
|
||||
let correlation_id = msg.correlation_id;
|
||||
|
||||
let result = match target {
|
||||
ConsumerTarget::LLM => self.llm_tx.send(msg).await,
|
||||
ConsumerTarget::Tool => self.tool_tx.send(msg).await,
|
||||
ConsumerTarget::Send => self.send_tx.send(msg).await,
|
||||
};
|
||||
|
||||
if let Err(_e) = result {
|
||||
warn!(
|
||||
"路由消息失败 (target={:?}, user={}, corr={}): 通道已关闭，跳过",
|
||||
target, user_id, correlation_id
|
||||
);
|
||||
// 不中断整个路由循环，跳过该消息继续处理
|
||||
}
|
||||
} else {
|
||||
// 队列为空，等待通知或超时
|
||||
tokio::select! {
|
||||
_ = self.signal.notified() => {
|
||||
// 有新的消息入队，继续循环
|
||||
}
|
||||
_ = tokio::time::sleep(self.poll_interval) => {
|
||||
// 超时兜底，继续检查
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
info!("QueueRunner 已停止");
|
||||
}
|
||||
}
|
||||
|
||||
/// 入队句柄 —— 外部模块通过它向队列投放消息
|
||||
#[derive(Clone)]
|
||||
pub struct EnqueueHandle {
|
||||
queue: Arc<Mutex<MessageQueue>>,
|
||||
shutdown: Arc<AtomicBool>,
|
||||
}
|
||||
|
||||
impl EnqueueHandle {
|
||||
/// 向队列投放一条消息
|
||||
pub async fn enqueue(&self, msg: PipelineMessage) {
|
||||
if self.shutdown.load(Ordering::Relaxed) {
|
||||
warn!(
|
||||
"队列已关闭，丢弃消息: user={}",
|
||||
msg.user_id()
|
||||
);
|
||||
return;
|
||||
}
|
||||
self.queue.lock().await.enqueue(msg);
|
||||
}
|
||||
}
|
||||
|
||||
/// 关闭句柄
|
||||
#[derive(Clone)]
|
||||
pub struct ShutdownHandle {
|
||||
shutdown: Arc<AtomicBool>,
|
||||
}
|
||||
|
||||
impl ShutdownHandle {
|
||||
/// 发送关闭信号
|
||||
pub fn shutdown(&self) {
|
||||
self.shutdown.store(true, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
/// 等所有消费者处理完并关闭后，等待队列排空
|
||||
#[allow(dead_code)]
|
||||
pub async fn drain_and_shutdown(
|
||||
_enqueue: &EnqueueHandle,
|
||||
runner_shutdown: &ShutdownHandle,
|
||||
consumer_handles: Vec<tokio::task::JoinHandle<()>>,
|
||||
) {
|
||||
info!("开始关闭队列系统...");
|
||||
|
||||
// 1. 阻止新消息入队（设置关闭标志）
|
||||
runner_shutdown.shutdown();
|
||||
|
||||
// 2. 等待排空（给一个合理的时间）
|
||||
tokio::time::sleep(std::time::Duration::from_millis(500)).await;
|
||||
|
||||
// 3. 等待消费者 task 完成
|
||||
for handle in consumer_handles {
|
||||
if let Err(e) = handle.await {
|
||||
error!("消费者 task 结束异常: {:?}", e);
|
||||
}
|
||||
}
|
||||
|
||||
info!("队列系统已完全关闭");
|
||||
}
|
||||
+31
-1
@@ -1,9 +1,39 @@
|
||||
//! ## 定时任务调度器
|
||||
//!
|
||||
//! 轮询 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 命令无限执行
|
||||
|
||||
use sqlx::PgPool;
|
||||
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<Arc<PgPool>>,
|
||||
}
|
||||
|
||||
+21
-1
@@ -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,
|
||||
}
|
||||
|
||||
+44
-5
@@ -1,3 +1,23 @@
|
||||
//! ## 审批管理器 —— 高风险工具的用户确认流程
|
||||
//!
|
||||
//! 当 LLM 调用高风险工具时，需要用户通过微信输入确认码来批准。
|
||||
//!
|
||||
//! ### 完整流程
|
||||
//! 1. `create(user_id, skill_name)` → 生成 6 位随机确认码，SHA-256 哈希后存储
|
||||
//! 2. 向用户发送 "⚠️ 操作确认\n\n技能：{name}\n确认码：{code}"
|
||||
//! 3. 用户在微信聊天中输入确认码
|
||||
//! 4. `handle_reply(user_id, reply)` → 验证确认码哈希
|
||||
//! - 匹配 → `Approved`
|
||||
//! - 不匹配 → 减少尝试次数（最多 3 次）
|
||||
//! - 输入 0 或"取消" → `Rejected`
|
||||
//! - 超时 → `Expired`（每 60 秒清理一次）
|
||||
//!
|
||||
//! ### 安全设计
|
||||
//! - 确认码用 SHA-256 哈希后存储，原始码不持久化
|
||||
//! - 每次生成随机 6 位数字（100000-999999）
|
||||
//! - 最多 3 次尝试，超时 5 分钟
|
||||
//! - 同时支持内存和数据库双写
|
||||
|
||||
use rand::Rng;
|
||||
use sha2::{Digest, Sha256};
|
||||
use sqlx::PgPool;
|
||||
@@ -6,7 +26,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 +48,23 @@ struct PendingEntry {
|
||||
sender: oneshot::Sender<ApprovalDecision>,
|
||||
}
|
||||
|
||||
/// 审批管理器
|
||||
/// ## 审批管理器
|
||||
///
|
||||
/// 管理高风险工具的用户确认流程。
|
||||
///
|
||||
/// ### 流程
|
||||
/// 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<user_id, Vec<PendingEntry>>`（快速访问）
|
||||
/// * 数据库 — `pending_approvals` 表（持久化审计）
|
||||
pub struct ApprovalManager {
|
||||
pool: Option<Arc<PgPool>>,
|
||||
/// user_id → 该用户的待审批队列
|
||||
@@ -175,12 +215,11 @@ impl ApprovalManager {
|
||||
|
||||
let mut map = self.pending.lock().await;
|
||||
for (uid, i, _) in expired.iter().rev() {
|
||||
if let Some(entries) = map.get_mut(uid) {
|
||||
if *i < entries.len() {
|
||||
if let Some(entries) = map.get_mut(uid)
|
||||
&& *i < entries.len() {
|
||||
let entry = entries.remove(*i);
|
||||
let _ = entry.sender.send(ApprovalDecision::Expired);
|
||||
}
|
||||
}
|
||||
}
|
||||
map.retain(|_, v| !v.is_empty());
|
||||
}
|
||||
|
||||
+27
-1
@@ -1,6 +1,32 @@
|
||||
//! ## 内置工具注册表 —— 工具路由中心
|
||||
//!
|
||||
//! 连接 LLM 的元工具调用和实际 Rust 实现的工具函数。
|
||||
//! 所有内置工具在此统一注册，通过名称路由到对应实现。
|
||||
//!
|
||||
//! ### 职责
|
||||
//! 1. `execute(name, args)` — 按名称路由到对应的 builtin 工具执行
|
||||
//! 2. `specs()` — 返回所有内置工具的 SkillSpec（给元工具 `query_capabilities` 使用）
|
||||
//! 3. `is_high_risk(name)` — 判断工具是否需要审批
|
||||
//! 4. `is_builtin(name)` — 判断名称是否对应一个已注册的内置工具
|
||||
|
||||
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 {
|
||||
|
||||
@@ -1,3 +1,28 @@
|
||||
//! ## 高德地图综合工具（Rust 原生实现 — v5 API）
|
||||
//!
|
||||
//! 提供 6 个高德地图相关的工具：
|
||||
//!
|
||||
//! | 工具名 | 功能 | API 版本 |
|
||||
//! |--------|------|---------|
|
||||
//! | `amap_poi_search` | POI（地点）搜索 + 周边搜索 | v5 |
|
||||
//! | `amap_geocode` | 地理编码：地址 → 坐标 | v3 |
|
||||
//! | `amap_reverse_geocode` | 逆地理编码：坐标 → 地址 | v3 |
|
||||
//! | `amap_route_plan` | 路径规划（步行/驾车/骑行/公交） | v5 |
|
||||
//! | `amap_travel_plan` | 智能旅游规划（搜索兴趣点 + 规划路线） | v5 |
|
||||
//! | `amap_map_link` | 生成地图可视化链接 | - |
|
||||
//!
|
||||
//! ### 认证
|
||||
//! 环境变量 `AMAP_WEBSERVICE_KEY` 或 `AMAP_KEY`
|
||||
//! 获取 Key: https://lbs.amap.com/api/webservice/create-project-and-key
|
||||
//!
|
||||
//! ### API 端点
|
||||
//! - POI 文本搜索: `GET /v5/place/text`
|
||||
//! - POI 周边搜索: `GET /v5/place/around`
|
||||
//! - 地理编码: `GET /v3/geocode/geo`
|
||||
//! - 逆地理编码: `GET /v3/geocode/regeo`
|
||||
//! - 步行/驾车/骑行/公交路线: `GET /v5/direction/{type}`
|
||||
//! - 所有请求必须带 `appname=amap-lbs-skill`
|
||||
|
||||
// 高德地图综合工具（Rust 原生实现 — v5 API）
|
||||
//
|
||||
// API:
|
||||
|
||||
@@ -1,8 +1,21 @@
|
||||
//! 日期时间工具
|
||||
//! ## 日期时间工具 —— 获取当前北京时间
|
||||
//!
|
||||
//! 最简单的内置工具，无网络依赖，纯粹返回本地时间。
|
||||
//! 用于 LLM 需要知道当前时间时调用。
|
||||
//!
|
||||
//! ### 风险等级
|
||||
//! `Low` — 无需用户确认，直接执行
|
||||
//!
|
||||
//! ### 输出格式
|
||||
//! ```json
|
||||
//! {"datetime": "2026-06-10 14:30:00"}
|
||||
//! ```
|
||||
|
||||
use chrono::Local;
|
||||
|
||||
use super::super::types::{RiskLevel, SkillResult, SkillSpec};
|
||||
|
||||
/// 返回该工具的元数据描述（供 query_capabilities 元工具使用）
|
||||
pub fn spec() -> SkillSpec {
|
||||
SkillSpec {
|
||||
name: "get_current_datetime".into(),
|
||||
@@ -13,6 +26,10 @@ pub fn spec() -> SkillSpec {
|
||||
}
|
||||
}
|
||||
|
||||
/// 执行日期时间查询
|
||||
///
|
||||
/// 使用 chrono::Local 获取系统当前时间，格式化为 "YYYY-MM-DD HH:MM:SS"。
|
||||
/// 注意：系统时区需设置为 Asia/Shanghai（北京时间 UTC+8）。
|
||||
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())
|
||||
|
||||
@@ -1,9 +1,23 @@
|
||||
//! 网页内容提取工具
|
||||
//! ## 网页内容提取工具 —— 可读性算法提取正文
|
||||
//!
|
||||
//! 使用可读性算法自动识别正文区域：清除广告/导航/页脚噪声，
|
||||
//! 按文本密度评分，取最高分区域。同时支持站点自定义 CSS 选择器。
|
||||
//! 自动识别网页正文区域，清除广告、导航、页脚等噪声，
|
||||
//! 按文本密度评分，取最高分区域作为正文。
|
||||
//!
|
||||
//! 站点选择器配置: `~/.ias/site_selectors.json`
|
||||
//! ### 提取策略（三层降级）
|
||||
//! 1. **自定义选择器** — 从 `~/.ias/site_selectors.json` 读取按站点的 CSS 选择器
|
||||
//! 2. **可读性算法** — 清除噪声 → 对块级元素评分 → 取最高分
|
||||
//! 3. **回退** — 取 `<body>` 的全部文本
|
||||
//!
|
||||
//! ### 评分因子
|
||||
//! - 文本长度（基础分）
|
||||
//! - 逗号/句号数（句子结构丰富度）
|
||||
//! - `<p>` 标签数（段落丰富度）
|
||||
//! - 链接密度惩罚（链接太多 → 导航/索引）
|
||||
//! - 代码块惩罚（代码多 → 降分）
|
||||
//!
|
||||
//! ### 配置
|
||||
//! 站点选择器配置文件: `~/.ias/site_selectors.json`
|
||||
//! 格式: `{"example.com": ["article.main", "#content"]}`
|
||||
|
||||
use reqwest::Client as HttpClient;
|
||||
use scraper::{ElementRef, Html, Selector};
|
||||
@@ -11,19 +25,20 @@ use std::collections::HashMap;
|
||||
|
||||
use super::super::types::{RiskLevel, SkillResult, SkillSpec};
|
||||
|
||||
/// 站点选择器配置文件路径
|
||||
/// 站点选择器配置文件路径（相对于 home 目录）
|
||||
const CONFIG_PATH: &str = ".ias/site_selectors.json";
|
||||
|
||||
/// 最大输出字符数
|
||||
/// 最大输出字符数（超过则截断）
|
||||
const MAX_OUTPUT_CHARS: usize = 6000;
|
||||
|
||||
/// 需移除的 HTML 标签
|
||||
/// 需要移除的 HTML 标签列表（这些标签内的内容不可能是正文）
|
||||
const STRIP_TAGS: &[&str] = &[
|
||||
"script", "style", "noscript", "iframe", "svg",
|
||||
"nav", "header", "footer", "aside", "form",
|
||||
];
|
||||
|
||||
/// 需移除的 class/id 关键词
|
||||
/// 需要移除的 class/id 关键词列表
|
||||
/// 匹配元素的 class 或 id 属性，包含这些关键词的视为噪声元素
|
||||
const NOISE_PATTERNS: &[&str] = &[
|
||||
"ad", "ads", "advert", "banner",
|
||||
"sidebar", "side-bar", "widget",
|
||||
@@ -42,7 +57,7 @@ const NOISE_PATTERNS: &[&str] = &[
|
||||
"edit-section", "noprint", "thumb",
|
||||
];
|
||||
|
||||
/// 内容评分阈值：低于此分不取
|
||||
/// 内容评分阈值：低于此分的元素不视为正文候选
|
||||
const MIN_SCORE_THRESHOLD: f64 = 50.0;
|
||||
|
||||
// ─── 工具 spec ───
|
||||
@@ -136,20 +151,18 @@ fn extract_content(document: &Html, domain: &str) -> String {
|
||||
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 {
|
||||
if let Some(text) = try_selector(document, sel_str)
|
||||
&& text.len() > 50 {
|
||||
return text;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 2. 先用可读性算法
|
||||
if let Some(text) = readability_extract(document) {
|
||||
if text.len() > 50 {
|
||||
if let Some(text) = readability_extract(document)
|
||||
&& text.len() > 50 {
|
||||
return text;
|
||||
}
|
||||
}
|
||||
|
||||
// 3. 回退到 body
|
||||
if let Some(text) = try_selector(document, "body") {
|
||||
@@ -239,7 +252,7 @@ fn score_element(el: &ElementRef) -> f64 {
|
||||
score += commas * 3.0;
|
||||
|
||||
// 句号/问号/感叹号 → 完整句子
|
||||
let sentences = text.matches(|c| c == '.' || c == '。' || c == '?' || c == '？' || c == '!' || c == '！').count() as f64;
|
||||
let sentences = text.matches(['.', '。', '?', '？', '!', '！']).count() as f64;
|
||||
score += sentences * 2.0;
|
||||
|
||||
// 段落数（<p> 标签）
|
||||
@@ -280,11 +293,10 @@ fn extract_clean_text(el: &ElementRef) -> String {
|
||||
|
||||
for child in el.descendants() {
|
||||
// 跳过噪声元素
|
||||
if let Some(child_el) = ElementRef::wrap(child) {
|
||||
if is_noise(&child_el) {
|
||||
if let Some(child_el) = ElementRef::wrap(child)
|
||||
&& is_noise(&child_el) {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(t) = child.value().as_text() {
|
||||
let t = t.trim();
|
||||
@@ -311,12 +323,11 @@ fn extract_clean_text(el: &ElementRef) -> String {
|
||||
// ─── 辅助 ───
|
||||
|
||||
fn extract_title(document: &Html) -> String {
|
||||
if let Ok(sel) = Selector::parse("title") {
|
||||
if let Some(el) = document.select(&sel).next() {
|
||||
if let Ok(sel) = Selector::parse("title")
|
||||
&& let Some(el) = document.select(&sel).next() {
|
||||
let t: String = el.text().collect::<Vec<_>>().join(" ").trim().to_string();
|
||||
if !t.is_empty() { return t; }
|
||||
}
|
||||
}
|
||||
"无标题".to_string()
|
||||
}
|
||||
|
||||
|
||||
@@ -1,8 +1,25 @@
|
||||
//! 备忘录工具 — 文件存储，支持 add/list/delete
|
||||
//! ## 备忘录工具 —— 本地文件存储的简易备忘录
|
||||
//!
|
||||
//! 支持三种操作：
|
||||
//! - `add` — 添加一条备忘录（自动分配自增 ID）
|
||||
//! - `list` — 列出所有备忘录
|
||||
//! - `delete` — 按 ID 删除备忘录
|
||||
//!
|
||||
//! ### 存储
|
||||
//! 数据存储在 `.data/memos.json` 文件中，JSON 数组格式。
|
||||
//! 每条记录包含：`id`、`content`、`created_at`。
|
||||
//!
|
||||
//! ### 智能推断
|
||||
//! 当通过 `call_capability` 调用时，可以从 prompt 文本中推断操作类型：
|
||||
//! - 包含"添加"/"新增"/"add" → 自动转为 add 操作
|
||||
//! - 包含"删除"/"移除"/"delete" → 自动转为 delete 操作
|
||||
//! - 其他 → 默认 list 操作
|
||||
|
||||
use chrono::Local;
|
||||
|
||||
use super::super::types::{RiskLevel, SkillResult, SkillSpec};
|
||||
|
||||
/// 返回该工具的元数据描述（供 query_capabilities 元工具使用）
|
||||
pub fn spec() -> SkillSpec {
|
||||
SkillSpec {
|
||||
name: "manage_memos".into(),
|
||||
@@ -20,6 +37,10 @@ pub fn spec() -> SkillSpec {
|
||||
}
|
||||
}
|
||||
|
||||
/// 执行备忘录操作
|
||||
///
|
||||
/// 从 JSON 参数中解析 action/content/id，执行对应操作。
|
||||
/// 数据文件路径为 `.data/memos.json`，自动创建目录。
|
||||
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 在顶层
|
||||
|
||||
@@ -1,3 +1,19 @@
|
||||
//! ## 内置工具实现集合
|
||||
//!
|
||||
//! 每个子模块对应一个具体的工具实现：
|
||||
//!
|
||||
//! | 模块 | 工具名 | 功能 |
|
||||
//! |------|--------|------|
|
||||
//! | `datetime` | `get_current_datetime` | 获取当前北京时间 |
|
||||
//! | `weather` | `query_weather` | 和风天气查询（实时/预报/逐时） |
|
||||
//! | `web_search` | `web_search` | Tavily 联网搜索 |
|
||||
//! | `fetch_page` | `fetch_page` | 网页正文提取（可读性算法） |
|
||||
//! | `memos` | `manage_memos` | 备忘录增删查 |
|
||||
//! | `amap` | `amap_*` | 高德地图系列（POI/地理编码/路径规划/旅游规划） |
|
||||
//!
|
||||
//! 每个工具都提供 `spec()` 返回元数据，和 `execute()` 执行入口。
|
||||
//! 通过 `BuiltinRegistry` 统一注册和路由。
|
||||
|
||||
pub mod amap;
|
||||
pub mod datetime;
|
||||
pub mod fetch_page;
|
||||
|
||||
@@ -1,13 +1,25 @@
|
||||
// 和风天气查询工具（Rust 原生实现，替代 bash 脚本）
|
||||
//
|
||||
// API:
|
||||
// GeoAPI: GET /geo/v2/city/lookup?location=城市名 → city_id
|
||||
// Now: GET /v7/weather/now?location=city_id → 实时天气
|
||||
// Daily: GET /v7/weather/{3,7,10,15,30}d?location=id → 每日预报
|
||||
// Hourly: GET /v7/weather/{24,72,168}h?location=id → 逐时预报
|
||||
//
|
||||
// 认证: EdDSA JWT (Ed25519)
|
||||
// 缓存: city_id 内存缓存 1h TTL
|
||||
//! ## 和风天气查询工具（Rust 原生实现）
|
||||
//!
|
||||
//! 通过和风天气 API 查询实时天气、每日预报和逐小时预报。
|
||||
//!
|
||||
//! ### API 端点
|
||||
//! - 城市查询: `GET /geo/v2/city/lookup?location=城市名` → 获取 city_id
|
||||
//! - 实时天气: `GET /v7/weather/now?location=city_id`
|
||||
//! - 每日预报: `GET /v7/weather/{3,7,10,15,30}d?location=id`
|
||||
//! - 逐时预报: `GET /v7/weather/{24,72,168}h?location=id`
|
||||
//!
|
||||
//! ### 认证方式
|
||||
//! 使用 EdDSA (Ed25519) JWT 进行 API 认证。
|
||||
//! 密钥对存放在 `qweather/` 目录下。
|
||||
//!
|
||||
//! ### 环境变量
|
||||
//! - `QWEATHER_API_HOST` — API 地址（默认 `api.qweather.com`）
|
||||
//! - `QWEATHER_JWT_KEY_ID` — JWT Key ID
|
||||
//! - `QWEATHER_JWT_PROJECT_ID` — 项目 ID
|
||||
//! - `QWEATHER_JWT_PRIVATE_KEY_FILE` — Ed25519 私钥路径
|
||||
//!
|
||||
//! ### 缓存策略
|
||||
//! city_id 在内存中缓存 1 小时（TTL），减少重复查询。
|
||||
|
||||
use base64::{Engine as _, engine::general_purpose};
|
||||
use chrono::Utc;
|
||||
@@ -474,11 +486,9 @@ fn extract_days_from_prompt(prompt: &str) -> u32 {
|
||||
.rev()
|
||||
.collect::<String>()
|
||||
.into()
|
||||
{
|
||||
if let Ok(n) = num_str.parse::<u32>() {
|
||||
&& let Ok(n) = num_str.parse::<u32>() {
|
||||
return n.min(30).max(1);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
3 // 默认
|
||||
@@ -601,7 +611,7 @@ pub async fn execute(params: serde_json::Value) -> SkillResult {
|
||||
if let Some(arr) = obj.get_mut(*key).and_then(|v| v.as_array_mut()) {
|
||||
for item in arr.iter_mut() {
|
||||
if let Some(m) = item.as_object_mut() {
|
||||
m.retain(|_, v| !v.is_null() && v.as_str().map_or(true, |s| !s.is_empty()));
|
||||
m.retain(|_, v| !v.is_null() && v.as_str().is_none_or(|s| !s.is_empty()));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,15 +1,29 @@
|
||||
//! Web 搜索工具 — Tavily Search API
|
||||
//! ## Web 搜索工具 —— Tavily Search API
|
||||
//!
|
||||
//! API: POST https://api.tavily.com/search
|
||||
//! 认证: Authorization: Bearer <TAVILY_API_KEY>
|
||||
//! 文档: https://docs.tavily.com/documentation/api-reference/endpoint/search
|
||||
//! 通过 Tavily 搜索引擎获取互联网实时信息。
|
||||
//! Tavily 是一个专为 AI Agent 设计的搜索引擎，返回结构化搜索结果。
|
||||
//!
|
||||
//! ### API
|
||||
//! - 端点: `POST https://api.tavily.com/search`
|
||||
//! - 认证: `Authorization: Bearer <TAVILY_API_KEY>`
|
||||
//! - 文档: https://docs.tavily.com/documentation/api-reference/endpoint/search
|
||||
//!
|
||||
//! ### 功能特性
|
||||
//! - 支持 AI 摘要（include_answer）
|
||||
//! - 支持多种搜索深度（basic / advanced / fast / ultra-fast）
|
||||
//! - 支持按主题过滤（general / news / finance）
|
||||
//! - 支持按时间范围过滤
|
||||
//! - 支持限定/排除域名
|
||||
//! - 支持按国家优先搜索结果
|
||||
|
||||
use reqwest::Client as HttpClient;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use super::super::types::{RiskLevel, SkillResult, SkillSpec};
|
||||
|
||||
// ─── 请求 ───
|
||||
// ═══════════════════════════════════════════════
|
||||
// 请求结构体
|
||||
// ═══════════════════════════════════════════════
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct SearchRequest {
|
||||
@@ -73,7 +87,9 @@ struct SearchRequest {
|
||||
safe_search: Option<bool>, // Enterprise only
|
||||
}
|
||||
|
||||
// ─── 响应 ───
|
||||
// ═══════════════════════════════════════════════
|
||||
// 响应结构体
|
||||
// ═══════════════════════════════════════════════
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[allow(dead_code)]
|
||||
@@ -126,7 +142,9 @@ struct UsageInfo {
|
||||
credits: Option<u32>,
|
||||
}
|
||||
|
||||
// ─── 工具 spec ───
|
||||
// ═══════════════════════════════════════════════
|
||||
// 工具元数据
|
||||
// ═══════════════════════════════════════════════
|
||||
|
||||
pub fn spec() -> SkillSpec {
|
||||
SkillSpec {
|
||||
@@ -161,7 +179,9 @@ pub fn spec() -> SkillSpec {
|
||||
}
|
||||
}
|
||||
|
||||
// ─── 执行 ───
|
||||
// ═══════════════════════════════════════════════
|
||||
// 执行入口
|
||||
// ═══════════════════════════════════════════════
|
||||
|
||||
pub async fn execute(params: serde_json::Value) -> SkillResult {
|
||||
let api_key = match std::env::var("TAVILY_API_KEY") {
|
||||
@@ -305,11 +325,10 @@ pub async fn execute(params: serde_json::Value) -> SkillResult {
|
||||
let mut output = String::new();
|
||||
|
||||
// AI 摘要
|
||||
if let Some(ref answer) = data.answer {
|
||||
if !answer.is_empty() {
|
||||
if let Some(ref answer) = data.answer
|
||||
&& !answer.is_empty() {
|
||||
output.push_str(&format!("📝 {}\n\n", answer));
|
||||
}
|
||||
}
|
||||
|
||||
// 搜索结果
|
||||
if data.results.is_empty() {
|
||||
|
||||
+62
-9
@@ -1,14 +1,68 @@
|
||||
//! ## 工具系统 —— 元工具 + 分离式内置工具
|
||||
//!
|
||||
//! 本模块是 iAs 的"工具系统"，让 LLM 能够调用外部功能。
|
||||
//!
|
||||
//! ### 架构设计（两层元工具）
|
||||
//!
|
||||
//! LLM 不直接感知每个具体工具，而是通过两个**元工具**间接调用：
|
||||
//!
|
||||
//! 1. **`query_capabilities`** — 列出所有可用工具的名称、描述、参数格式
|
||||
//! 2. **`call_capability`** — 按名称调用工具，参数通过 JSON 传递
|
||||
//!
|
||||
//! 这样做的优势：新增工具时只需要在服务端注册，不需要修改 LLM 的 tool definitions。
|
||||
//!
|
||||
//! ### 工具类型
|
||||
//! * **内置工具** (builtins/) — Rust 代码实现的工具（天气、搜索、高德地图等）
|
||||
//! * **上下文工具** — 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 继续对话
|
||||
|
||||
//! ## 工具系统 —— 元工具 + 内置工具 + 审批管理
|
||||
//!
|
||||
//! 本模块是 iAs 的"工具系统"，让 LLM 能够调用外部功能。
|
||||
//!
|
||||
//! ### 架构设计（两层元工具）
|
||||
//!
|
||||
//! LLM 不直接感知每个具体工具，而是通过两个**元工具**间接调用：
|
||||
//!
|
||||
//! 1. **`query_capabilities`** — 列出所有可用工具的名称、描述、参数格式
|
||||
//! 2. **`call_capability`** — 按名称调用工具，参数通过 JSON 传递
|
||||
//!
|
||||
//! 这样做的优势：新增工具时只需要在服务端注册，不需要修改 LLM 的 tool definitions。
|
||||
//!
|
||||
//! ### 工具类型
|
||||
//! - **内置工具** (builtins/) — Rust 代码实现的工具（天气、搜索、高德地图等）
|
||||
//! - **上下文工具** — read_memories, write_memory, read_summaries（直接在 daemon 中处理）
|
||||
//!
|
||||
//! ### 工具执行流程
|
||||
//! ```text
|
||||
//! LLM 请求 call_capability("query_weather", {"location":"北京"})
|
||||
//! → daemon 解析出 target_name="query_weather"，提取参数
|
||||
//! → 高风险？→ 走 ApprovalManager 审批流
|
||||
//! → 低风险？→ BuiltinRegistry::execute()
|
||||
//! → 结果返回 LLM 继续对话
|
||||
//! ```
|
||||
|
||||
pub mod approval;
|
||||
pub mod builtin;
|
||||
pub mod builtins;
|
||||
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();
|
||||
|
||||
@@ -16,11 +70,10 @@ pub fn unpack_call_params(cp: &serde_json::Value) -> serde_json::Value {
|
||||
if let Some(prompt) = cp.get("prompt") {
|
||||
match prompt {
|
||||
serde_json::Value::String(s) => {
|
||||
if let Ok(v) = serde_json::from_str::<serde_json::Value>(s) {
|
||||
if let Some(obj) = v.as_object() {
|
||||
if let Ok(v) = serde_json::from_str::<serde_json::Value>(s)
|
||||
&& let Some(obj) = v.as_object() {
|
||||
params.extend(obj.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
serde_json::Value::Object(obj) => {
|
||||
params.extend(obj.clone());
|
||||
@@ -138,7 +191,7 @@ fn format_spec(lines: &mut Vec<String>, spec: &SkillSpec) {
|
||||
.parameters
|
||||
.get("properties")
|
||||
.and_then(|v| v.as_object())
|
||||
.map_or(true, |o| o.is_empty())
|
||||
.is_none_or(|o| o.is_empty())
|
||||
{
|
||||
lines.push(format!(
|
||||
" 参数: {}",
|
||||
|
||||
+33
-4
@@ -1,9 +1,22 @@
|
||||
//! ## 工具系统类型定义
|
||||
//!
|
||||
//! 定义了工具系统的核心类型：
|
||||
//!
|
||||
//! - `RiskLevel` — 工具的安全级别（Low/High）
|
||||
//! - `SkillSpec` — 工具的元数据描述（名称、参数、风险等级）
|
||||
//! - `SkillResult` — 工具执行结果
|
||||
//! - `ExecutionContext` — 工具执行时的上下文环境
|
||||
//! - `WechatSender` — 微信发送回调类型别名
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
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 +27,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 +56,11 @@ fn default_timeout() -> u64 {
|
||||
30
|
||||
}
|
||||
|
||||
// ─── 工具执行结果 ───
|
||||
/// ## 工具执行结果
|
||||
///
|
||||
/// * `success` — 执行是否成功
|
||||
/// * `output` — 可读的输出文本（将返回给 LLM）
|
||||
/// * `data` — 结构化数据（可选），用于后续处理
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct SkillResult {
|
||||
@@ -89,7 +113,12 @@ impl SkillResult {
|
||||
}
|
||||
}
|
||||
|
||||
// ─── 执行上下文 ───
|
||||
/// ## 执行上下文 —— 工具执行时需要的环境信息
|
||||
///
|
||||
/// 当工具在 daemon 上下文中执行时，携带以下信息：
|
||||
/// * `user_id` — 触发工具的用户
|
||||
/// * `approval_manager` — 审批管理器（高风险工具需要）
|
||||
/// * `send_wechat` — 微信发送回调（用于向用户发送审批提示）
|
||||
|
||||
pub type WechatSender = Arc<
|
||||
dyn Fn(&str, &str) -> Pin<Box<dyn Future<Output = Result<(), String>> + Send>> + Send + Sync,
|
||||
|
||||
+66
-9
@@ -1,3 +1,19 @@
|
||||
//! ## 微信 iLink Bot API 客户端
|
||||
//!
|
||||
//! 封装了与微信 iLink Bot API 的所有通信：
|
||||
//!
|
||||
//! ### 核心功能
|
||||
//! - **扫码登录** — `get_qrcode()` → `poll_qr_status()` → `login()`
|
||||
//! - **消息接收** — `receive_messages()`（长轮询）
|
||||
//! - **消息发送** — `send_text()`
|
||||
//! - **生命周期** — `notify_start()` / `notify_stop()`
|
||||
//!
|
||||
//! ### 线程安全
|
||||
//! `#[derive(Clone)]` + 内部 `Arc<Mutex<>>` 保证所有方法可以在 tokio task 间共享。
|
||||
//!
|
||||
//! ### 环境变量
|
||||
//! - `WEIXIN_BASE_URL` — API 地址（默认 `https://ilinkai.weixin.qq.com`）
|
||||
|
||||
use crate::wechat::types::*;
|
||||
use base64::Engine;
|
||||
use reqwest::Client as HttpClient;
|
||||
@@ -6,7 +22,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<Mutex<>>` 保证所有方法可以在 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,
|
||||
@@ -22,15 +55,21 @@ impl WeChatClient {
|
||||
const ILINK_APP_ID: &'static str = "";
|
||||
/// 默认 bot_type
|
||||
const DEFAULT_BOT_TYPE: &'static str = "3";
|
||||
/// 默认 API 地址
|
||||
/// 默认 API 地址（可通过 WEIXIN_BASE_URL 环境变量覆盖）
|
||||
const DEFAULT_BASE_URL: &'static str = "https://ilinkai.weixin.qq.com";
|
||||
/// 二维码固定 API 地址
|
||||
const FIXED_BASE_URL: &'static str = "https://ilinkai.weixin.qq.com";
|
||||
/// 二维码固定 API 地址（可通过 WEIXIN_BASE_URL 环境变量覆盖）
|
||||
fn fixed_base_url() -> String {
|
||||
std::env::var("WEIXIN_BASE_URL")
|
||||
.unwrap_or_else(|_| "https://ilinkai.weixin.qq.com".to_string())
|
||||
}
|
||||
/// 版本号编码：0x00MMNNPP
|
||||
const CLIENT_VERSION: u32 = 0x000100_00; // 1.0.0
|
||||
|
||||
pub fn new(base_url: Option<String>) -> Self {
|
||||
let base = base_url.unwrap_or_else(|| Self::DEFAULT_BASE_URL.to_string());
|
||||
// 优先级: 显式参数 > WEIXIN_BASE_URL 环境变量 > 默认值
|
||||
let base = base_url
|
||||
.or_else(|| std::env::var("WEIXIN_BASE_URL").ok())
|
||||
.unwrap_or_else(|| Self::DEFAULT_BASE_URL.to_string());
|
||||
|
||||
Self {
|
||||
http: HttpClient::builder()
|
||||
@@ -46,11 +85,14 @@ 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={}",
|
||||
Self::FIXED_BASE_URL,
|
||||
&Self::fixed_base_url(),
|
||||
Self::DEFAULT_BOT_TYPE
|
||||
);
|
||||
|
||||
@@ -68,7 +110,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,
|
||||
@@ -147,7 +198,7 @@ impl WeChatClient {
|
||||
|
||||
let deadline = tokio::time::Instant::now() + Duration::from_secs(timeout_secs);
|
||||
let mut scanned = false;
|
||||
let mut current_base = Self::FIXED_BASE_URL.to_string();
|
||||
let mut current_base = Self::fixed_base_url();
|
||||
|
||||
while tokio::time::Instant::now() < deadline {
|
||||
let status = self.poll_qr_status(&qrcode, ¤t_base).await?;
|
||||
@@ -358,6 +409,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(
|
||||
|
||||
@@ -1,2 +1,16 @@
|
||||
//! ## 微信 iLink Bot 通道
|
||||
//!
|
||||
//! 封装与微信 iLink Bot API 的所有通信：
|
||||
//!
|
||||
//! - `client` — HTTP 客户端（登录、长轮询收消息、发消息、生命周期管理）
|
||||
//! - `types` — API 协议类型定义（消息结构、请求/响应体）
|
||||
//!
|
||||
//! ### 核心流程
|
||||
//! 1. `login()` — 扫码登录，获取 token + account_id
|
||||
//! 2. `notify_start()` — 注册监听器
|
||||
//! 3. `receive_messages()` — 长轮询接收消息
|
||||
//! 4. `send_text()` — 发送文本回复
|
||||
//! 5. `notify_stop()` — 注销监听器
|
||||
|
||||
pub mod client;
|
||||
pub mod types;
|
||||
|
||||
+32
-39
@@ -1,3 +1,15 @@
|
||||
//! ## 微信 iLink Bot API 协议类型
|
||||
//!
|
||||
//! 定义了与微信 iLink Bot API 通信的全部数据结构：
|
||||
//!
|
||||
//! - `BaseInfo` — 公共请求元数据
|
||||
//! - `WeixinMessage` — 微信消息结构（支持 text/image/voice/file/video）
|
||||
//! - `MessageItem` — 消息条目（5 种类型）
|
||||
//! - `GetUpdatesReq/Resp` — 长轮询收消息的请求/响应
|
||||
//! - `SendMessageReq` — 发送消息请求
|
||||
//! - `QRCodeResponse` / `StatusResponse` — 扫码登录相关
|
||||
//! - `LoginResult` / `MessageEvent` — 业务层结果类型
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
// ─── 基础类型 ───
|
||||
@@ -116,8 +128,14 @@ pub struct RefMessage {
|
||||
pub title: Option<String>,
|
||||
}
|
||||
|
||||
/// 消息条目（支持 text/image/voice/file/video）
|
||||
/// ## 消息条目 —— 消息的具体内容
|
||||
///
|
||||
/// 支持 5 种消息类型：text(1), image(2), voice(3), file(4), video(5)。
|
||||
/// 每种类型有对应的 `*_item` 字段，其他字段保持 None。
|
||||
///
|
||||
/// `text()` 便捷方法用于快速构造文本消息（发送消息时使用）。
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[derive(Default)]
|
||||
pub struct MessageItem {
|
||||
#[serde(rename = "type", skip_serializing_if = "Option::is_none")]
|
||||
pub item_type: Option<i32>,
|
||||
@@ -165,27 +183,22 @@ impl MessageItem {
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for MessageItem {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
item_type: None,
|
||||
create_time_ms: None,
|
||||
update_time_ms: None,
|
||||
is_completed: None,
|
||||
msg_id: None,
|
||||
ref_msg: None,
|
||||
text_item: None,
|
||||
image_item: None,
|
||||
voice_item: None,
|
||||
file_item: None,
|
||||
video_item: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ─── 微信消息 ───
|
||||
/// ## 微信消息结构
|
||||
///
|
||||
/// 对应 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)]
|
||||
#[derive(Default)]
|
||||
pub struct WeixinMessage {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub seq: Option<i64>,
|
||||
@@ -217,26 +230,6 @@ pub struct WeixinMessage {
|
||||
pub context_token: Option<String>,
|
||||
}
|
||||
|
||||
impl Default for WeixinMessage {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
seq: None,
|
||||
message_id: None,
|
||||
from_user_id: None,
|
||||
to_user_id: None,
|
||||
client_id: None,
|
||||
create_time_ms: None,
|
||||
update_time_ms: None,
|
||||
delete_time_ms: None,
|
||||
session_id: None,
|
||||
group_id: None,
|
||||
msg_type: None,
|
||||
message_state: None,
|
||||
item_list: None,
|
||||
context_token: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl WeixinMessage {
|
||||
#[allow(dead_code)]
|
||||
|
||||
+15
-6
@@ -1,6 +1,15 @@
|
||||
//! Worker — 无状态执行进程
|
||||
//! ## Worker —— 无状态执行进程（已废弃）
|
||||
//!
|
||||
//! 每条消息 spawn 一个 worker:
|
||||
//! ### 旧架构说明
|
||||
//! 在旧架构中，每条消息会 spawn 一个独立的 Worker 进程，
|
||||
//! Worker 通过 UDS 连接到 daemon，处理完一条消息后退出。
|
||||
//!
|
||||
//! 优势：Worker 代码热更新（编译后下一条消息自动用新版本）
|
||||
//! 劣势：进程启动开销、IPC 序列化开销、状态管理复杂
|
||||
//!
|
||||
//! ### 新架构
|
||||
//! 现在改为 daemon 内的 LLM/Tool/Send 三个消费者 tokio task。
|
||||
//! 此模块保留供编译参考，不再被 daemon 使用。
|
||||
//! 1. 连接 daemon 的 Unix Domain Socket
|
||||
//! 2. 读取 TaskFrame (用户消息 + 上下文 + 环境变量)
|
||||
//! 3. 注入环境变量，构建 Conversation，执行 LLM 对话 + 工具调用
|
||||
@@ -103,7 +112,8 @@ pub async fn run(sock_path: &str) -> Result<(), String> {
|
||||
conv.set_tool_executor(executor);
|
||||
|
||||
// 7. 执行 LLM 对话（带审批检测）
|
||||
let (reply, _used_tools, usage) = conv
|
||||
use crate::llm::ChatResult;
|
||||
let ChatResult { reply, used_tools: _used_tools, usage, .. } = conv
|
||||
.chat_with_tools(task.msg.text.clone())
|
||||
.await
|
||||
.map_err(|e| format!("LLM 对话失败: {}", e))?;
|
||||
@@ -266,8 +276,8 @@ fn build_worker_executor(
|
||||
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()) {
|
||||
if cp.get("action").and_then(|v| v.as_str()) == Some("add")
|
||||
&& let Some(content) = cp.get("content").and_then(|v| v.as_str()) {
|
||||
shared
|
||||
.memories
|
||||
.lock()
|
||||
@@ -277,7 +287,6 @@ fn build_worker_executor(
|
||||
.push(content.to_string());
|
||||
shared.new_memories.lock().await.push(content.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
return Ok(result.output);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user