feat: iPet → iAs 完整迁移,Rust 版微信 AI 助手
Phase 1 ✅ CLI 框架 + 配置系统 - clap 子命令: login / listen / send / whoami / usage - config.json + env var 替换 - tracing 日志系统 - state 持久化(auth/runtime 文件存 + PostgreSQL) Phase 2 ✅ 微信通道 - wechat::client — 完整 iLink Bot HTTP API 实现 - 扫码登录(终端二维码 + 轮询状态) - 长轮询 getupdates / 消息收发 / 监听注册 Phase 3 ✅ AI 对话(纯文本 + function calling) - LlmProvider trait: DeepSeek + LM Studio 实现 - SSE 流式解析(text / reasoning / tool_calls delta / usage) - Conversation: 消息历史 + chat / chat_with_tools 工具循环 Phase 4 ✅ PostgreSQL 集成 - app_state(认证 KV 存储) - chat_records(消息收发记录) - llm_usage(Token 用量统计缓存命中率) - user_memories(长期记忆持久化) - pending_approvals(审批确认码) - scheduled_tasks(定时任务表) Phase 5 ✅ 一切皆 Skill(工具系统) - SkillRegistry: 系统 + 用户 skills 双目录合并 - SKILL.md 解析器 + 子进程执行器(stdin JSON → stdout) - 9 个系统 Skills: datetime / weather / search / email / shell / schedule / memos / read_memories / read_summaries - ApprovalManager: High 风险技能 → 确认码审批(透明模式) - High 风险技能:确认码审批(透明模式) Phase 6 ✅ 定时任务调度器 上下文管理 - ChatSession: checkpoint + token budget (28K) + summaries - Token 估算器(中英文自适应) - 12h 空闲 → trigger_idle_summary(不入会话) - Budget 溢出 → trigger_overflow_summary(入会话 + drain 旧消息) - Summarizer: LLM 生成自然语言摘要(fallback 简单截断) - 长期记忆 / 摘要 通过 read_memories / read_summaries 工具按需读取 工具调用日志 + Token 统计 - INFO: 工具名 + 参数 + 结果摘要 - DEBUG: 子进程 exit/stdout/stderr - ias usage --since --until --model 查看用量和缓存命中率
This commit is contained in:
+63
@@ -0,0 +1,63 @@
|
||||
use clap::{Parser, Subcommand};
|
||||
|
||||
/// iAs - 微信 AI 智能助手
|
||||
#[derive(Parser, Debug)]
|
||||
#[command(name = "ias", version, about = "微信 AI 智能助手 (Rust)")]
|
||||
pub struct Cli {
|
||||
#[command(subcommand)]
|
||||
pub command: Commands,
|
||||
}
|
||||
|
||||
#[derive(Subcommand, Debug)]
|
||||
pub enum Commands {
|
||||
/// 扫码登录微信
|
||||
Login {
|
||||
/// 登录超时秒数
|
||||
#[arg(long, default_value = "480")]
|
||||
timeout: u64,
|
||||
},
|
||||
|
||||
/// 监听消息(长轮询)
|
||||
Listen {
|
||||
/// 启用 AI 自动回复
|
||||
#[arg(long)]
|
||||
llm: bool,
|
||||
|
||||
/// 仅回显消息(不调用 AI)
|
||||
#[arg(long)]
|
||||
echo: bool,
|
||||
},
|
||||
|
||||
/// 发送文本消息
|
||||
Send {
|
||||
/// 目标用户 ID
|
||||
#[arg(long)]
|
||||
to: String,
|
||||
|
||||
/// 消息内容
|
||||
#[arg(long)]
|
||||
text: String,
|
||||
|
||||
/// 上下文令牌
|
||||
#[arg(long)]
|
||||
context_token: Option<String>,
|
||||
},
|
||||
|
||||
/// 查看当前登录状态
|
||||
Whoami,
|
||||
|
||||
/// 查看 LLM Token 使用统计
|
||||
Usage {
|
||||
/// 起始时间(ISO 格式,默认7天前)
|
||||
#[arg(long)]
|
||||
since: Option<String>,
|
||||
|
||||
/// 结束时间(ISO 格式,默认现在)
|
||||
#[arg(long)]
|
||||
until: Option<String>,
|
||||
|
||||
/// 模型名称过滤
|
||||
#[arg(long)]
|
||||
model: Option<String>,
|
||||
},
|
||||
}
|
||||
+218
@@ -0,0 +1,218 @@
|
||||
// 允许未来阶段使用的字段
|
||||
#![allow(dead_code)]
|
||||
|
||||
use serde::Deserialize;
|
||||
use std::path::Path;
|
||||
|
||||
/// iAs 全局配置
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
pub struct AppConfig {
|
||||
#[serde(default)]
|
||||
pub llm: LlmConfig,
|
||||
#[serde(default)]
|
||||
pub storage: StorageConfig,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
pub struct LlmConfig {
|
||||
/// LLM 提供商:deepseek | lmstudio
|
||||
#[serde(default = "default_llm_provider")]
|
||||
pub provider: String,
|
||||
|
||||
/// DeepSeek 配置
|
||||
#[serde(default)]
|
||||
pub deepseek: DeepSeekConfig,
|
||||
|
||||
/// LM Studio 配置
|
||||
#[serde(default)]
|
||||
pub lmstudio: LmStudioConfig,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
pub struct DeepSeekConfig {
|
||||
#[serde(default = "default_deepseek_api_key")]
|
||||
pub api_key: String,
|
||||
#[serde(default = "default_deepseek_model")]
|
||||
pub model: String,
|
||||
#[serde(default = "default_deepseek_base_url")]
|
||||
pub base_url: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
pub struct LmStudioConfig {
|
||||
#[serde(default = "default_lmstudio_base_url")]
|
||||
pub base_url: String,
|
||||
#[serde(default = "default_lmstudio_model")]
|
||||
pub model: String,
|
||||
#[serde(default)]
|
||||
pub api_key: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
pub struct StorageConfig {
|
||||
/// 状态文件目录
|
||||
#[serde(default = "default_state_dir")]
|
||||
pub state_dir: String,
|
||||
/// PostgreSQL 连接串(可选,默认为文件存储)
|
||||
#[serde(default)]
|
||||
pub database_url: String,
|
||||
}
|
||||
|
||||
// ─── 默认值 ───
|
||||
|
||||
fn default_llm_provider() -> String {
|
||||
"deepseek".to_string()
|
||||
}
|
||||
|
||||
fn default_deepseek_api_key() -> String {
|
||||
env_or_default("DEEPSEEK_API_KEY", "")
|
||||
}
|
||||
|
||||
fn default_deepseek_model() -> String {
|
||||
env_or_default("DEEPSEEK_MODEL", "deepseek-v4-flash")
|
||||
}
|
||||
|
||||
fn default_deepseek_base_url() -> String {
|
||||
env_or_default("DEEPSEEK_BASE_URL", "https://api.deepseek.com/v1")
|
||||
}
|
||||
|
||||
fn default_lmstudio_base_url() -> String {
|
||||
env_or_default("LM_STUDIO_BASE_URL", "http://localhost:1234/v1")
|
||||
}
|
||||
|
||||
fn default_lmstudio_model() -> String {
|
||||
env_or_default("LM_STUDIO_MODEL", "local-model")
|
||||
}
|
||||
|
||||
fn default_state_dir() -> String {
|
||||
".data/weixin-ilink".to_string()
|
||||
}
|
||||
|
||||
fn env_or_default(key: &str, default: &str) -> String {
|
||||
std::env::var(key).unwrap_or_else(|_| default.to_string())
|
||||
}
|
||||
|
||||
impl Default for LlmConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
provider: default_llm_provider(),
|
||||
deepseek: DeepSeekConfig::default(),
|
||||
lmstudio: LmStudioConfig::default(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for DeepSeekConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
api_key: default_deepseek_api_key(),
|
||||
model: default_deepseek_model(),
|
||||
base_url: default_deepseek_base_url(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for LmStudioConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
base_url: default_lmstudio_base_url(),
|
||||
model: default_lmstudio_model(),
|
||||
api_key: String::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for StorageConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
state_dir: default_state_dir(),
|
||||
database_url: String::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for AppConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
llm: LlmConfig::default(),
|
||||
storage: StorageConfig::default(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl AppConfig {
|
||||
/// 从文件加载配置,缺失字段回退到环境变量
|
||||
pub fn from_file(path: impl AsRef<Path>) -> Self {
|
||||
let path = path.as_ref();
|
||||
if path.exists() {
|
||||
match std::fs::read_to_string(path) {
|
||||
Ok(content) => {
|
||||
// 先做环境变量替换
|
||||
let resolved = resolve_env_vars(&content);
|
||||
match serde_json::from_str(&resolved) {
|
||||
Ok(cfg) => return cfg,
|
||||
Err(e) => {
|
||||
tracing::warn!("解析配置文件失败 ({}): {},使用默认配置", path.display(), e);
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("读取配置文件失败 ({}): {},使用默认配置", path.display(), e);
|
||||
}
|
||||
}
|
||||
}
|
||||
Self::default()
|
||||
}
|
||||
}
|
||||
|
||||
/// 替换字符串中的 `${VAR_NAME}` 为环境变量值
|
||||
fn resolve_env_vars(input: &str) -> String {
|
||||
let mut result = String::with_capacity(input.len());
|
||||
let mut chars = input.chars().peekable();
|
||||
|
||||
while let Some(c) = chars.next() {
|
||||
if c == '$' && chars.peek() == Some(&'{') {
|
||||
chars.next(); // 跳过 '{'
|
||||
let mut var_name = String::new();
|
||||
while let Some(&ch) = chars.peek() {
|
||||
if ch == '}' {
|
||||
chars.next(); // 跳过 '}'
|
||||
break;
|
||||
}
|
||||
var_name.push(ch);
|
||||
chars.next();
|
||||
}
|
||||
let value = std::env::var(&var_name).unwrap_or_default();
|
||||
result.push_str(&value);
|
||||
} else {
|
||||
result.push(c);
|
||||
}
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_resolve_env_vars() {
|
||||
// set_var is unsafe in edition 2024; skip env test
|
||||
let result = resolve_env_vars("prefix_${TEST_VAR}_suffix");
|
||||
// TEST_VAR might not be set, so result could be "prefix__suffix"
|
||||
assert_eq!(result, "prefix__suffix");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_resolve_missing_env_var() {
|
||||
let result = resolve_env_vars("prefix_${MISSING_VAR}_suffix");
|
||||
assert_eq!(result, "prefix__suffix");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_resolve_no_vars() {
|
||||
let result = resolve_env_vars("plain text");
|
||||
assert_eq!(result, "plain text");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,193 @@
|
||||
# iAs 上下文管理系统设计 v2
|
||||
|
||||
## 核心理念
|
||||
|
||||
1. **长期记忆不固定注入** — 通过工具让 LLM 按需读写
|
||||
2. **摘要不入上下文** — 除非因上下文空间不足被迫压缩
|
||||
3. **摘要总结点** — 记录压缩位置,之后的消息始终完整保留
|
||||
4. **12 小时空闲触发** — 超过 12h 自动压缩,但不插入会话
|
||||
5. **Token Budget 驱动** — 达到上限才压缩,保持消息完整度优先
|
||||
|
||||
---
|
||||
|
||||
## 状态模型
|
||||
|
||||
```
|
||||
ChatSession
|
||||
├── checkpoint: usize = 0 ← 摘要总结点(消息列表中的位置)
|
||||
├── messages: Vec<Message> ← 全部消息(checkpoint 之前的是历史)
|
||||
├── summaries: Vec<SummaryEntry> ← 所有摘要
|
||||
│ ├── pos: usize ← 对应消息列表中的位置
|
||||
│ ├── text: String ← 摘要内容
|
||||
│ └── reason: "overflow"|"timeout" ← 生成原因
|
||||
└── last_user_at: Option<DateTime> ← 上一条用户消息时间
|
||||
```
|
||||
|
||||
### 消息列表视图(LLM 实际看到的内容)
|
||||
|
||||
```
|
||||
TOOL: read_memories / write_memory / read_summaries
|
||||
↑ LLM 按需调用,不自动注入
|
||||
|
||||
VIEW: [system_prompt]
|
||||
[overflow_summary] ← 仅当因 context 满产生时注入
|
||||
[messages from checkpoint..] ← 始终完整保留
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 核心类型
|
||||
|
||||
```rust
|
||||
pub struct SummaryEntry {
|
||||
pub checkpoint: usize, // 对应 messages 中的位置
|
||||
pub text: String,
|
||||
pub reason: SummaryReason,
|
||||
}
|
||||
|
||||
pub enum SummaryReason {
|
||||
Overflow, // 上下文满了自动压缩
|
||||
Timeout, // 超过 12 小时空闲
|
||||
}
|
||||
|
||||
pub struct ChatSession {
|
||||
checkpoint: usize,
|
||||
messages: Vec<Message>,
|
||||
summaries: Vec<SummaryEntry>,
|
||||
last_user_at: Option<DateTime<Utc>>,
|
||||
token_budget: u32, // 默认 32000
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 核心流程
|
||||
|
||||
```
|
||||
用户发消息 "合肥天气"
|
||||
│
|
||||
▼
|
||||
1. 距离检查
|
||||
├── last_user_at 超过 12h?
|
||||
│ └── YES → summarize_all(reason=timeout)
|
||||
│ ├── 生成摘要存入 summaries(不插入会话)
|
||||
│ └── checkpoint = messages.len()
|
||||
│
|
||||
▼
|
||||
2. 构建 LLM 请求
|
||||
├── system_prompt
|
||||
├── 【如果有 overflow_summary 且 checkpoint > 0】
|
||||
│ └── { role: "system", content: "历史摘要:..." }
|
||||
└── messages[checkpoint..]
|
||||
│
|
||||
▼
|
||||
3. Token 估算
|
||||
├── 在 budget 内 → 直接发送
|
||||
│
|
||||
└── 超出 budget → summarize_all(reason=overflow)
|
||||
├── 生成摘要存入 summaries
|
||||
├── checkpoint = messages.len()
|
||||
├── 摘要作为 system message 注入本次请求
|
||||
└── 重新构建请求(现在只有摘要 + 当前消息)
|
||||
│
|
||||
▼
|
||||
4. 发送给 LLM
|
||||
├── LLM 回复(可能包含工具调用)
|
||||
│ ├── read_memories → 返回长期记忆
|
||||
│ ├── write_memory → 保存到长期记忆
|
||||
│ └── read_summaries → 返回 summaries 列表
|
||||
│
|
||||
└── LLM 最终回复 → record_turn(user_msg, assistant_msg)
|
||||
├── messages.push(user_msg, assistant_msg)
|
||||
└── last_user_at = now
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 摘要触发条件对比
|
||||
|
||||
| 原因 | 触发条件 | 是否插入会话 | 用途 |
|
||||
|------|---------|------------|------|
|
||||
| **Overflow** | token 超 budget | ✅ 自动插入 | 保持上下文不超限 |
|
||||
| **Timeout** | 距上条消息 >12h | ❌ 不插入(工具可读) | 跨天会话,LLM 自行决定是否查看 |
|
||||
|
||||
---
|
||||
|
||||
## 长期记忆工具(LLM 按需调用)
|
||||
|
||||
```json
|
||||
{
|
||||
"name": "read_memories",
|
||||
"description": "读取用户的长期记忆(跨会话持久保存的个人信息)",
|
||||
"parameters": { "type": "object", "properties": {} }
|
||||
}
|
||||
|
||||
{
|
||||
"name": "write_memory",
|
||||
"description": "写入一条用户的长期记忆",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"content": { "type": "string", "description": "记忆内容" }
|
||||
},
|
||||
"required": ["content"]
|
||||
}
|
||||
}
|
||||
|
||||
{
|
||||
"name": "read_summaries",
|
||||
"description": "读取历史会话摘要(包括空闲超时产生的摘要)",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
这些工具注册为系统 skills(只对 `chat_with_tools` 启用,不暴露给用户)。
|
||||
|
||||
## 数据库存储
|
||||
|
||||
```sql
|
||||
-- 长期记忆(跨会话持久)
|
||||
CREATE TABLE user_memories (
|
||||
user_id TEXT PRIMARY KEY,
|
||||
content JSONB NOT NULL DEFAULT '[]',
|
||||
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
|
||||
);
|
||||
```
|
||||
|
||||
摘要和消息会话不持久化(会话消息保存到 `chat_records` 表,但摘要只存在于运行时的 `ChatSession`)。
|
||||
|
||||
## 与 iPet 的区别
|
||||
|
||||
| 特性 | iPet | iAs v2 |
|
||||
|------|------|--------|
|
||||
| 长期记忆 | 自动注入每个请求 | 通过工具按需读写 |
|
||||
| 滚动摘要 | 自动注入每个请求 | 不注入,工具可读 |
|
||||
| 会话摘要 | 自动注入 | 不注入,overflow 时例外 |
|
||||
| 摘要触发 | 定期 + 按消息数 | Token budget + 12h 空闲 |
|
||||
| 历史完整保留 | ❌ 始终压缩 | ✅ checkpoint 后完整保留 |
|
||||
|
||||
## Token 估算
|
||||
|
||||
```rust
|
||||
pub fn estimate_tokens(text: &str) -> u32 {
|
||||
let cjk = text.chars().filter(|c| c >= '\u{4e00}' && c <= '\u{9fff}').count() as u32;
|
||||
let other = text.chars().count() as u32 - cjk;
|
||||
cjk * 15 / 10 + other / 4 + 8 // +8 消息 overhead
|
||||
}
|
||||
```
|
||||
|
||||
## 实施步骤
|
||||
|
||||
| 步骤 | 内容 |
|
||||
|------|------|
|
||||
| 1 | ChatSession + SummaryEntry 核心类型 |
|
||||
| 2 | build_context() — 构建 LLM 消息数组 |
|
||||
| 3 | record_turn() — 记录对话轮次 |
|
||||
| 4 | summarize_all() — 摘要生成 + checkpoint 更新 |
|
||||
| 5 | token_estimate() — Token 估算 |
|
||||
| 6 | 长期记忆工具(read_memories / write_memory)|
|
||||
| 7 | 摘要读取工具(read_summaries)|
|
||||
| 8 | 集成到 Conversation + main.rs |
|
||||
@@ -0,0 +1,266 @@
|
||||
use crate::context::types::ChatSession;
|
||||
use crate::llm::conversation::Summarizer;
|
||||
use crate::llm::types::Message;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
/// 估算消息的 token 数
|
||||
pub fn estimate_tokens(msg: &Message) -> u32 {
|
||||
let text = msg.content.as_str();
|
||||
estimate_text_tokens(text) + 8
|
||||
}
|
||||
|
||||
/// 估算纯文本的 token 数
|
||||
pub fn estimate_text_tokens(text: &str) -> u32 {
|
||||
if text.is_empty() {
|
||||
return 0;
|
||||
}
|
||||
let cjk_count = text.chars().filter(|c| is_cjk(*c)).count() as u32;
|
||||
let other_count = text.chars().count() as u32 - cjk_count;
|
||||
// 中文 ≈ 1.5 token/字,英文 ≈ 0.25 token/字符
|
||||
cjk_count * 15 / 10 + other_count / 4
|
||||
}
|
||||
|
||||
fn is_cjk(c: char) -> bool {
|
||||
('\u{4e00}'..='\u{9fff}').contains(&c)
|
||||
|| ('\u{3400}'..='\u{4dbf}').contains(&c)
|
||||
|| ('\u{f900}'..='\u{faff}').contains(&c)
|
||||
}
|
||||
|
||||
/// 构建给 LLM 的消息数组(带 token budget 管理)
|
||||
///
|
||||
/// 返回消息数组和是否需要摘要的信息
|
||||
pub async fn build_context(
|
||||
session: &Arc<Mutex<ChatSession>>,
|
||||
system_prompt: &str,
|
||||
) -> Vec<Message> {
|
||||
let mut s = session.lock().await;
|
||||
|
||||
// ── 1. 空闲超时检查(消息到达前由调用方检查)──
|
||||
// 这里只做构建,超时触发在上层
|
||||
|
||||
// ── 2. 构建消息 ──
|
||||
let mut messages = Vec::new();
|
||||
let mut used = 0u32;
|
||||
|
||||
// 系统提示词
|
||||
let sys_tokens = estimate_text_tokens(system_prompt);
|
||||
messages.push(Message::system(system_prompt));
|
||||
used += sys_tokens + 8;
|
||||
|
||||
// overflow 摘要(如果有)
|
||||
if let Some(summary) = s.latest_overflow_summary() {
|
||||
let text = format!("【历史对话摘要】\n{}", summary.text);
|
||||
let tokens = estimate_text_tokens(&text);
|
||||
if used + tokens < s.token_budget {
|
||||
messages.push(Message::system(text));
|
||||
used += tokens + 8;
|
||||
}
|
||||
}
|
||||
|
||||
// 近期消息(从 checkpoint 开始,倒序添加直到接近 budget)
|
||||
let recent = s.recent_messages().to_vec();
|
||||
let mut included = Vec::new();
|
||||
let mut tail_used = 0u32;
|
||||
|
||||
for msg in recent.iter().rev() {
|
||||
let t = estimate_tokens(msg);
|
||||
if used + tail_used + t > s.token_budget - 500 {
|
||||
break;
|
||||
}
|
||||
tail_used += t;
|
||||
included.push(msg.clone());
|
||||
}
|
||||
included.reverse();
|
||||
|
||||
// 保护:确保最近一条用户消息始终在上下文中
|
||||
if let Some(last_user) = recent.iter().rev().find(|m| m.role == crate::llm::types::Role::User) {
|
||||
let already_included = included.iter().any(|m| {
|
||||
m.role == last_user.role && m.content == last_user.content
|
||||
});
|
||||
if !already_included {
|
||||
let t = estimate_tokens(last_user);
|
||||
if used + tail_used + t <= s.token_budget + 2000 {
|
||||
included.push(last_user.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
messages.extend(included);
|
||||
|
||||
messages
|
||||
}
|
||||
|
||||
/// 触发溢出摘要:压缩 checkpoint 到当前位置之间的新消息
|
||||
pub async fn trigger_overflow_summary(
|
||||
session: &Arc<Mutex<ChatSession>>,
|
||||
summarizer: Option<&Summarizer>,
|
||||
) -> bool {
|
||||
let mut s = session.lock().await;
|
||||
if s.messages.is_empty() {
|
||||
return false;
|
||||
}
|
||||
|
||||
let prev_checkpoint = s.checkpoint;
|
||||
let end = s.messages.len();
|
||||
|
||||
// checkpoint 始终在 user+assistant 对边界,直接用它作为摘要起点
|
||||
let start = prev_checkpoint;
|
||||
if start >= end {
|
||||
return false;
|
||||
}
|
||||
|
||||
let to_summarize: Vec<_> = s.messages[start..end]
|
||||
.iter()
|
||||
.take(40) // 最多 40 条
|
||||
.cloned()
|
||||
.collect();
|
||||
|
||||
if to_summarize.is_empty() {
|
||||
return false;
|
||||
}
|
||||
|
||||
let summary_text = generate_summary(&to_summarize, summarizer).await;
|
||||
|
||||
s.summaries.push(super::types::SummaryEntry {
|
||||
checkpoint: prev_checkpoint,
|
||||
text: summary_text,
|
||||
reason: super::types::SummaryReason::Overflow,
|
||||
created_at: chrono::Utc::now(),
|
||||
});
|
||||
|
||||
s.checkpoint = end;
|
||||
|
||||
// 清理太旧的消息(checkpoint 之前的保留最近 100 条用于调试)
|
||||
if s.checkpoint > 200 {
|
||||
let keep = s.checkpoint - 100;
|
||||
s.messages.drain(0..keep);
|
||||
// 调整 checkpoint 和 summary checkpoint
|
||||
let shift = keep;
|
||||
s.checkpoint -= shift;
|
||||
for sum in &mut s.summaries {
|
||||
if sum.checkpoint >= shift {
|
||||
sum.checkpoint -= shift;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
true
|
||||
}
|
||||
|
||||
/// 触发空闲摘要(不注入,保存到 summaries 供 LLM 工具查询)
|
||||
pub async fn trigger_idle_summary(
|
||||
session: &Arc<Mutex<ChatSession>>,
|
||||
summarizer: Option<&Summarizer>,
|
||||
) {
|
||||
let mut s = session.lock().await;
|
||||
if s.messages.is_empty() {
|
||||
return;
|
||||
}
|
||||
|
||||
let summary_text = generate_summary(&s.messages, summarizer).await;
|
||||
let cp = s.messages.len();
|
||||
|
||||
s.summaries.push(super::types::SummaryEntry {
|
||||
checkpoint: cp,
|
||||
text: summary_text,
|
||||
reason: super::types::SummaryReason::Timeout,
|
||||
created_at: chrono::Utc::now(),
|
||||
});
|
||||
|
||||
s.checkpoint = s.messages.len();
|
||||
|
||||
// 清理旧消息(checkpoint 之前的全部清除)
|
||||
if s.checkpoint > 0 {
|
||||
let shift = s.checkpoint;
|
||||
s.messages.drain(0..shift);
|
||||
s.checkpoint = 0;
|
||||
for sum in &mut s.summaries {
|
||||
if sum.checkpoint >= shift {
|
||||
sum.checkpoint -= shift;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 生成摘要:优先使用 LLM,不可用时回退到简单截断
|
||||
async fn generate_summary(messages: &[Message], summarizer: Option<&Summarizer>) -> String {
|
||||
if let Some(summarizer) = summarizer {
|
||||
if messages.len() >= 3 {
|
||||
let prompt = build_summary_prompt(messages);
|
||||
match summarizer(prompt).await {
|
||||
Ok(text) if !text.is_empty() => {
|
||||
tracing::info!("LLM 摘要: {:.100}...", text);
|
||||
return text;
|
||||
}
|
||||
_ => tracing::warn!("LLM 摘要失败,回退到简单截断"),
|
||||
}
|
||||
}
|
||||
}
|
||||
summarize_messages(messages)
|
||||
}
|
||||
|
||||
/// 构建给 LLM 的摘要 prompt
|
||||
fn build_summary_prompt(messages: &[Message]) -> String {
|
||||
let convo: String = messages
|
||||
.iter()
|
||||
.filter(|m| m.role == crate::llm::types::Role::User || m.role == crate::llm::types::Role::Assistant)
|
||||
.map(|m| {
|
||||
let role = if m.role == crate::llm::types::Role::User { "用户" } else { "助手" };
|
||||
let content: String = m.content.chars().take(200).collect();
|
||||
format!("{}: {}", role, content)
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
|
||||
format!(
|
||||
"请将以下对话压缩为简短摘要,保留关键事实、用户偏好、决定和待办事项。只输出摘要,不要其他内容。\n\n对话:\n{}",
|
||||
convo
|
||||
)
|
||||
}
|
||||
|
||||
/// 简单的摘要生成(提取最近 N 条消息的关键内容)
|
||||
fn summarize_messages(messages: &[Message]) -> String {
|
||||
// 简单实现:提取用户消息的前 80 个字符作为摘要
|
||||
let lines: Vec<String> = messages
|
||||
.iter()
|
||||
.filter(|m| m.role == crate::llm::types::Role::User || m.role == crate::llm::types::Role::Assistant)
|
||||
.map(|m| {
|
||||
let role = match m.role {
|
||||
crate::llm::types::Role::User => "用户",
|
||||
crate::llm::types::Role::Assistant => "助手",
|
||||
_ => "",
|
||||
};
|
||||
let preview: String = m.content.chars().take(80).collect();
|
||||
format!("{}: {}", role, preview)
|
||||
})
|
||||
.collect();
|
||||
|
||||
if lines.is_empty() {
|
||||
"暂无历史对话".to_string()
|
||||
} else {
|
||||
let result = format!("历史对话摘要({} 条消息):\n{}", lines.len(), lines.join("\n"));
|
||||
if result.chars().count() > 2000 {
|
||||
format!("{}...(已截断)", result.chars().take(2000).collect::<String>())
|
||||
} else {
|
||||
result
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_estimate_cjk() {
|
||||
let t = estimate_text_tokens("你好世界");
|
||||
assert!(t >= 5 && t <= 10, "CJK estimate: {}", t);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_estimate_ascii() {
|
||||
let t = estimate_text_tokens("hello world");
|
||||
assert!(t >= 1 && t <= 6, "ASCII estimate: {}", t);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,6 @@
|
||||
pub mod builder;
|
||||
pub mod tools;
|
||||
pub mod types;
|
||||
|
||||
pub use tools::MemoryStore;
|
||||
pub use types::ChatSession;
|
||||
@@ -0,0 +1,92 @@
|
||||
use sqlx::PgPool;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
/// 长期记忆管理器(内存 + PostgreSQL 双写)
|
||||
pub struct MemoryStore {
|
||||
pool: Option<Arc<PgPool>>,
|
||||
cache: Arc<Mutex<Vec<String>>>,
|
||||
}
|
||||
|
||||
impl MemoryStore {
|
||||
pub fn new(pool: Option<Arc<PgPool>>) -> Self {
|
||||
Self {
|
||||
pool,
|
||||
cache: Arc::new(Mutex::new(Vec::new())),
|
||||
}
|
||||
}
|
||||
|
||||
/// 初始化:从数据库加载记忆到缓存
|
||||
pub async fn load(&self) {
|
||||
let pool = match self.pool.as_ref() {
|
||||
Some(p) => p.clone(),
|
||||
None => return,
|
||||
};
|
||||
if let Ok(rows) = sqlx::query_as::<_, (String,)>(
|
||||
"SELECT content FROM user_memories WHERE user_id = 'default' ORDER BY id",
|
||||
)
|
||||
.fetch_all(pool.as_ref())
|
||||
.await
|
||||
{
|
||||
let mut cache = self.cache.lock().await;
|
||||
*cache = rows.into_iter().map(|(c,)| c).collect();
|
||||
}
|
||||
}
|
||||
|
||||
/// 读取所有记忆
|
||||
pub async fn read(&self) -> String {
|
||||
let mems = self.cache.lock().await;
|
||||
if mems.is_empty() {
|
||||
"暂无长期记忆".to_string()
|
||||
} else {
|
||||
mems.iter()
|
||||
.enumerate()
|
||||
.map(|(i, m)| format!("{}. {}", i + 1, m))
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n")
|
||||
}
|
||||
}
|
||||
|
||||
/// 写入一条记忆
|
||||
pub async fn write(&self, content: &str) -> String {
|
||||
// 写入缓存
|
||||
self.cache.lock().await.push(content.to_string());
|
||||
// 写入数据库
|
||||
if let Some(ref pool) = self.pool {
|
||||
let _ = sqlx::query(
|
||||
"INSERT INTO user_memories (user_id, content) VALUES ('default', $1)",
|
||||
)
|
||||
.bind(content)
|
||||
.execute(pool.as_ref())
|
||||
.await;
|
||||
}
|
||||
"已添加长期记忆".to_string()
|
||||
}
|
||||
}
|
||||
|
||||
/// read_summaries 工具:读取历史摘要
|
||||
pub async fn read_summaries(
|
||||
session: &Arc<Mutex<super::types::ChatSession>>,
|
||||
) -> String {
|
||||
let s = session.lock().await;
|
||||
if s.summaries.is_empty() {
|
||||
"暂无历史摘要".to_string()
|
||||
} else {
|
||||
s.summaries
|
||||
.iter()
|
||||
.map(|sum| {
|
||||
let reason = match sum.reason {
|
||||
super::types::SummaryReason::Overflow => "上下文压缩",
|
||||
super::types::SummaryReason::Timeout => "空闲超时",
|
||||
};
|
||||
format!(
|
||||
"[{}] {} (原因: {})",
|
||||
sum.created_at.format("%Y-%m-%d %H:%M"),
|
||||
sum.text,
|
||||
reason,
|
||||
)
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n---\n")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,111 @@
|
||||
use crate::llm::types::Message;
|
||||
use chrono::{DateTime, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// 摘要原因
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
pub enum SummaryReason {
|
||||
/// 上下文 token 超 budget 触发
|
||||
Overflow,
|
||||
/// 距上条用户消息超过 12 小时触发
|
||||
Timeout,
|
||||
}
|
||||
|
||||
/// 一条摘要记录
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SummaryEntry {
|
||||
/// 摘要对应的消息位置(checkpoint)
|
||||
pub checkpoint: usize,
|
||||
/// 摘要文本
|
||||
pub text: String,
|
||||
/// 生成原因
|
||||
pub reason: SummaryReason,
|
||||
/// 生成时间
|
||||
pub created_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
/// 聊天会话状态
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ChatSession {
|
||||
/// 摘要总结点——此位置之前的消息已被压缩
|
||||
pub checkpoint: usize,
|
||||
/// 全部消息(checkpoint 之后的保持完整)
|
||||
pub messages: Vec<Message>,
|
||||
/// 所有摘要(keyed by checkpoint)
|
||||
pub summaries: Vec<SummaryEntry>,
|
||||
/// 上一条用户消息的时间
|
||||
pub last_user_at: Option<DateTime<Utc>>,
|
||||
/// Token 预算(默认 28000,留 4000 给回复)
|
||||
pub token_budget: u32,
|
||||
/// 12 小时空闲阈值(秒)
|
||||
pub idle_timeout_secs: i64,
|
||||
}
|
||||
|
||||
impl Default for ChatSession {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
checkpoint: 0,
|
||||
messages: Vec::new(),
|
||||
summaries: Vec::new(),
|
||||
last_user_at: None,
|
||||
token_budget: 28000,
|
||||
idle_timeout_secs: 12 * 3600,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ChatSession {
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
/// 记录一条用户消息
|
||||
pub fn add_user(&mut self, content: String) {
|
||||
self.messages.push(Message::user(content));
|
||||
self.last_user_at = Some(Utc::now());
|
||||
}
|
||||
|
||||
/// 记录一条助手消息
|
||||
pub fn add_assistant(&mut self, content: String) {
|
||||
self.messages.push(Message::assistant(content));
|
||||
}
|
||||
|
||||
/// 记录一条带 tool_calls 的助手消息
|
||||
pub fn add_assistant_tool_calls(&mut self, tool_calls: Vec<crate::llm::types::ToolCall>) {
|
||||
self.messages.push(Message::assistant_with_tool_calls(tool_calls));
|
||||
}
|
||||
|
||||
/// 记录一条工具结果
|
||||
pub fn add_tool_result(&mut self, tool_call_id: &str, name: &str, result: &str) {
|
||||
self.messages.push(Message::tool_result(tool_call_id, name, result));
|
||||
}
|
||||
|
||||
/// 是否因空闲超过阈值需要摘要
|
||||
pub fn is_idle_timeout(&self) -> bool {
|
||||
if let Some(last) = self.last_user_at {
|
||||
let elapsed = Utc::now().signed_duration_since(last).num_seconds();
|
||||
elapsed > self.idle_timeout_secs
|
||||
} else {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
/// 获取检查点之后的消息
|
||||
pub fn recent_messages(&self) -> &[Message] {
|
||||
if self.checkpoint < self.messages.len() {
|
||||
&self.messages[self.checkpoint..]
|
||||
} else {
|
||||
&[]
|
||||
}
|
||||
}
|
||||
|
||||
/// 是否有 overflow 摘要可以注入
|
||||
pub fn has_overflow_summary(&self) -> bool {
|
||||
self.summaries.iter().any(|s| s.reason == SummaryReason::Overflow)
|
||||
}
|
||||
|
||||
/// 获取最新的 overflow 摘要
|
||||
pub fn latest_overflow_summary(&self) -> Option<&SummaryEntry> {
|
||||
self.summaries.iter().rev().find(|s| s.reason == SummaryReason::Overflow)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,49 @@
|
||||
pub mod models;
|
||||
|
||||
use sqlx::postgres::PgPoolOptions;
|
||||
use sqlx::PgPool;
|
||||
use std::time::Duration;
|
||||
|
||||
/// 数据库管理器
|
||||
pub struct Database {
|
||||
pool: PgPool,
|
||||
}
|
||||
|
||||
impl Database {
|
||||
/// 从 DATABASE_URL 环境变量连接并运行迁移
|
||||
pub async fn connect() -> Result<Self, String> {
|
||||
let database_url =
|
||||
std::env::var("DATABASE_URL").map_err(|_| "请设置 DATABASE_URL 环境变量".to_string())?;
|
||||
|
||||
let pool = PgPoolOptions::new()
|
||||
.max_connections(5)
|
||||
.acquire_timeout(Duration::from_secs(10))
|
||||
.connect(&database_url)
|
||||
.await
|
||||
.map_err(|e| format!("连接数据库失败: {}", e))?;
|
||||
|
||||
// 运行迁移
|
||||
sqlx::migrate!("./migrations")
|
||||
.run(&pool)
|
||||
.await
|
||||
.map_err(|e| format!("数据库迁移失败: {}", e))?;
|
||||
|
||||
tracing::info!("数据库连接成功,迁移已完成");
|
||||
|
||||
Ok(Self { pool })
|
||||
}
|
||||
|
||||
/// 获取连接池引用
|
||||
pub fn pool(&self) -> &PgPool {
|
||||
&self.pool
|
||||
}
|
||||
|
||||
/// 检查数据库是否可用
|
||||
pub async fn health_check(&self) -> Result<(), String> {
|
||||
sqlx::query("SELECT 1")
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_err(|e| format!("数据库健康检查失败: {}", e))?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,193 @@
|
||||
use chrono::{DateTime, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sqlx::PgPool;
|
||||
|
||||
use crate::state::AuthState;
|
||||
|
||||
// ─── 认证状态 ───
|
||||
|
||||
/// 从数据库加载认证状态
|
||||
pub async fn load_auth(pool: &PgPool) -> Option<AuthState> {
|
||||
let row = sqlx::query_as::<_, (serde_json::Value,)>(
|
||||
"SELECT value FROM app_state WHERE key = 'auth'",
|
||||
)
|
||||
.fetch_optional(pool)
|
||||
.await
|
||||
.ok()??;
|
||||
|
||||
serde_json::from_value(row.0).ok()
|
||||
}
|
||||
|
||||
/// 保存认证状态到数据库
|
||||
pub async fn save_auth(pool: &PgPool, auth: &AuthState) -> Result<(), String> {
|
||||
let value = serde_json::to_value(auth).map_err(|e| format!("序列化 auth 失败: {}", e))?;
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO app_state (key, value, updated_at)
|
||||
VALUES ('auth', $1, NOW())
|
||||
ON CONFLICT (key) DO UPDATE
|
||||
SET value = EXCLUDED.value, updated_at = NOW()
|
||||
"#,
|
||||
)
|
||||
.bind(&value)
|
||||
.execute(pool)
|
||||
.await
|
||||
.map_err(|e| format!("保存 auth 到数据库失败: {}", e))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// ─── 聊天记录 ───
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)]
|
||||
pub struct ChatRecord {
|
||||
pub id: i64,
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub direction: String,
|
||||
pub user_id: String,
|
||||
pub account_id: String,
|
||||
pub text: String,
|
||||
pub source: String,
|
||||
pub context_token: String,
|
||||
pub message_id: String,
|
||||
}
|
||||
|
||||
/// 插入一条聊天记录
|
||||
pub async fn insert_chat_record(
|
||||
pool: &PgPool,
|
||||
direction: &str,
|
||||
user_id: &str,
|
||||
account_id: &str,
|
||||
text: &str,
|
||||
source: &str,
|
||||
context_token: Option<&str>,
|
||||
message_id: &str,
|
||||
) -> Result<i64, String> {
|
||||
let ctx = context_token.unwrap_or("");
|
||||
|
||||
let row: (i64,) = sqlx::query_as(
|
||||
r#"
|
||||
INSERT INTO chat_records (direction, user_id, account_id, text, source, context_token, message_id)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7)
|
||||
RETURNING id
|
||||
"#,
|
||||
)
|
||||
.bind(direction)
|
||||
.bind(user_id)
|
||||
.bind(account_id)
|
||||
.bind(text)
|
||||
.bind(source)
|
||||
.bind(ctx)
|
||||
.bind(message_id)
|
||||
.fetch_one(pool)
|
||||
.await
|
||||
.map_err(|e| format!("插入聊天记录失败: {}", e))?;
|
||||
|
||||
Ok(row.0)
|
||||
}
|
||||
|
||||
/// 查询最近的聊天记录(用户维度)
|
||||
pub async fn list_recent_chat_records(
|
||||
pool: &PgPool,
|
||||
user_id: &str,
|
||||
limit: i64,
|
||||
) -> Result<Vec<ChatRecord>, String> {
|
||||
let records = sqlx::query_as::<_, ChatRecord>(
|
||||
r#"
|
||||
SELECT id, created_at, direction, user_id, account_id, text, source, context_token, message_id
|
||||
FROM chat_records
|
||||
WHERE user_id = $1
|
||||
ORDER BY created_at DESC
|
||||
LIMIT $2
|
||||
"#,
|
||||
)
|
||||
.bind(user_id)
|
||||
.bind(limit)
|
||||
.fetch_all(pool)
|
||||
.await
|
||||
.map_err(|e| format!("查询聊天记录失败: {}", e))?;
|
||||
|
||||
Ok(records)
|
||||
}
|
||||
|
||||
// ─── LLM 用量 ───
|
||||
|
||||
/// 插入一条 LLM 用量记录
|
||||
pub async fn insert_llm_usage(
|
||||
pool: &PgPool,
|
||||
user_id: &str,
|
||||
model: &str,
|
||||
provider: &str,
|
||||
prompt_tokens: u32,
|
||||
completion_tokens: u32,
|
||||
cache_hit_tokens: u32,
|
||||
cache_miss_tokens: u32,
|
||||
) -> Result<(), String> {
|
||||
let total = prompt_tokens + completion_tokens;
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO llm_usage (user_id, model, provider, prompt_tokens, completion_tokens, total_tokens, cache_hit_tokens, cache_miss_tokens)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
|
||||
"#,
|
||||
)
|
||||
.bind(user_id)
|
||||
.bind(model)
|
||||
.bind(provider)
|
||||
.bind(prompt_tokens as i32)
|
||||
.bind(completion_tokens as i32)
|
||||
.bind(total as i32)
|
||||
.bind(cache_hit_tokens as i32)
|
||||
.bind(cache_miss_tokens as i32)
|
||||
.execute(pool)
|
||||
.await
|
||||
.map_err(|e| format!("插入 LLM 用量失败: {}", e))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// LLM 用量聚合统计
|
||||
#[derive(Debug, Clone, sqlx::FromRow)]
|
||||
pub struct LlmUsageStats {
|
||||
pub total_calls: i64,
|
||||
pub total_prompt_tokens: i64,
|
||||
pub total_completion_tokens: i64,
|
||||
pub total_tokens: i64,
|
||||
pub total_cache_hit: i64,
|
||||
pub total_cache_miss: i64,
|
||||
}
|
||||
|
||||
/// 查询 LLM 用量统计(按时间范围过滤)
|
||||
pub async fn query_llm_usage_stats(
|
||||
pool: &PgPool,
|
||||
since: Option<chrono::DateTime<chrono::Utc>>,
|
||||
until: Option<chrono::DateTime<chrono::Utc>>,
|
||||
model_filter: Option<&str>,
|
||||
) -> Result<LlmUsageStats, String> {
|
||||
let since = since.unwrap_or_else(|| {
|
||||
chrono::Utc::now() - chrono::Duration::days(7)
|
||||
});
|
||||
let until = until.unwrap_or_else(chrono::Utc::now);
|
||||
|
||||
let row = sqlx::query_as::<_, LlmUsageStats>(
|
||||
r#"
|
||||
SELECT
|
||||
COUNT(*)::bigint as total_calls,
|
||||
COALESCE(SUM(prompt_tokens), 0)::bigint as total_prompt_tokens,
|
||||
COALESCE(SUM(completion_tokens), 0)::bigint as total_completion_tokens,
|
||||
COALESCE(SUM(total_tokens), 0)::bigint as total_tokens,
|
||||
COALESCE(SUM(cache_hit_tokens), 0)::bigint as total_cache_hit,
|
||||
COALESCE(SUM(cache_miss_tokens), 0)::bigint as total_cache_miss
|
||||
FROM llm_usage
|
||||
WHERE created_at >= $1 AND created_at <= $2
|
||||
AND ($3::text IS NULL OR model = $3)
|
||||
"#,
|
||||
)
|
||||
.bind(since)
|
||||
.bind(until)
|
||||
.bind(model_filter)
|
||||
.fetch_one(pool)
|
||||
.await
|
||||
.map_err(|e| format!("查询 LLM 用量失败: {}", e))?;
|
||||
|
||||
Ok(row)
|
||||
}
|
||||
@@ -0,0 +1,242 @@
|
||||
use crate::context::builder;
|
||||
use crate::context::types::ChatSession;
|
||||
use crate::llm::provider::{create_provider, BoxedProvider, StreamReceiver};
|
||||
use crate::llm::types::{ConversationConfig, Message, StreamChunk, ToolCall, Usage};
|
||||
use std::future::Future;
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
pub type ToolExecutor =
|
||||
Arc<dyn Fn(&str, &str) -> Pin<Box<dyn Future<Output = Result<String, String>> + Send>> + Send + Sync>;
|
||||
|
||||
/// 摘要生成器:接收要摘要的消息文本 → 返回 LLM 生成的摘要
|
||||
pub type Summarizer =
|
||||
Arc<dyn Fn(String) -> Pin<Box<dyn Future<Output = Result<String, String>> + Send>> + Send + Sync>;
|
||||
|
||||
pub struct Conversation {
|
||||
config: ConversationConfig,
|
||||
provider: BoxedProvider,
|
||||
session: Arc<Mutex<ChatSession>>,
|
||||
tool_executor: Option<ToolExecutor>,
|
||||
}
|
||||
|
||||
impl Conversation {
|
||||
pub fn new(config: ConversationConfig) -> Result<Self, String> {
|
||||
let provider = create_provider(&config)?;
|
||||
Ok(Self {
|
||||
config,
|
||||
provider,
|
||||
session: Arc::new(Mutex::new(ChatSession::new())),
|
||||
tool_executor: None,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn with_provider(config: ConversationConfig, provider: BoxedProvider) -> Self {
|
||||
Self {
|
||||
config,
|
||||
provider,
|
||||
session: Arc::new(Mutex::new(ChatSession::new())),
|
||||
tool_executor: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn set_tool_executor(&mut self, executor: ToolExecutor) { self.tool_executor = Some(executor); }
|
||||
pub fn session(&self) -> Arc<Mutex<ChatSession>> { Arc::clone(&self.session) }
|
||||
pub fn model(&self) -> &str { &self.config.model }
|
||||
pub fn provider_name(&self) -> &str { self.provider.name() }
|
||||
pub fn spec(&self) -> &ConversationConfig { &self.config }
|
||||
pub fn summarizer(&self) -> Summarizer { self.create_summarizer() }
|
||||
|
||||
/// 创建 LLM 摘要生成器(用当前 provider 做非流式调用)
|
||||
pub fn create_summarizer(&self) -> Summarizer {
|
||||
let provider = self.provider.clone();
|
||||
Arc::new(move |prompt: String| {
|
||||
let provider = provider.clone();
|
||||
Box::pin(async move {
|
||||
let msgs = vec![Message::system("你是一个对话摘要助手,用中文输出简洁摘要。"), Message::user(prompt)];
|
||||
let mut rx = provider.chat_stream(&ConversationConfig::default(), &msgs).await
|
||||
.map_err(|e| format!("摘要请求失败: {}", e))?;
|
||||
let mut text = String::new();
|
||||
while let Some(chunk) = rx.recv().await {
|
||||
match chunk {
|
||||
StreamChunk::Text(t) => text.push_str(&t),
|
||||
StreamChunk::Done { text: full, .. } => {
|
||||
if !full.is_empty() { text = full; }
|
||||
break;
|
||||
}
|
||||
StreamChunk::Error(e) => return Err(e),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
Ok(text)
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
/// 单次对话(无工具)
|
||||
pub async fn chat(&self, user_message: String) -> Result<ChatHandle, String> {
|
||||
// 空闲超时检查
|
||||
{
|
||||
let s = self.session.lock().await;
|
||||
if s.is_idle_timeout() {
|
||||
drop(s);
|
||||
builder::trigger_idle_summary(&self.session, Some(&self.summarizer())).await;
|
||||
}
|
||||
}
|
||||
|
||||
self.session.lock().await.add_user(user_message.clone());
|
||||
let messages = builder::build_context(&self.session, &self.config.system_prompt).await;
|
||||
let rx = self.provider.chat_stream(&self.config, &messages).await?;
|
||||
|
||||
Ok(ChatHandle {
|
||||
rx: Some(rx), full_text: String::new(), tool_calls: None, usage: None,
|
||||
session: Arc::clone(&self.session),
|
||||
})
|
||||
}
|
||||
|
||||
/// 对话 + 工具循环(带上下文管理)
|
||||
pub async fn chat_with_tools(
|
||||
&self,
|
||||
user_message: String,
|
||||
) -> Result<(String, bool, Option<Usage>), String> {
|
||||
// ── 空闲超时检查 ──
|
||||
{
|
||||
let s = self.session.lock().await;
|
||||
if s.is_idle_timeout() {
|
||||
drop(s);
|
||||
builder::trigger_idle_summary(&self.session, Some(&self.summarizer())).await;
|
||||
tracing::info!("⏰ 检测到 12h 空闲,已生成摘要");
|
||||
}
|
||||
}
|
||||
|
||||
self.session.lock().await.add_user(user_message.clone());
|
||||
|
||||
let mut used_tools = false;
|
||||
let mut last_usage = None;
|
||||
let mut turn_count = 0u32;
|
||||
|
||||
loop {
|
||||
turn_count += 1;
|
||||
if turn_count > 5 {
|
||||
return Err("工具调用次数过多(最多5轮)".to_string());
|
||||
}
|
||||
|
||||
let messages = builder::build_context(&self.session, &self.config.system_prompt).await;
|
||||
let rx = self.provider.chat_stream(&self.config, &messages).await?;
|
||||
|
||||
let mut full_text = String::new();
|
||||
let mut tool_calls: Option<Vec<ToolCall>> = None;
|
||||
let mut rx = rx;
|
||||
|
||||
while let Some(chunk) = rx.recv().await {
|
||||
match chunk {
|
||||
StreamChunk::Text(t) => full_text.push_str(&t),
|
||||
StreamChunk::Done { text, tool_calls: tc, usage, .. } => {
|
||||
if !text.is_empty() { full_text = text; }
|
||||
tool_calls = tc;
|
||||
last_usage = usage;
|
||||
break;
|
||||
}
|
||||
StreamChunk::Error(e) => return Err(e),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(calls) = tool_calls {
|
||||
if !calls.is_empty() {
|
||||
used_tools = true;
|
||||
tracing::info!(
|
||||
"🔧 LLM 请求 {} 个工具: {}",
|
||||
calls.len(),
|
||||
calls.iter().map(|c| format!("{}({:.80})", c.function.name, c.function.arguments)).collect::<Vec<_>>().join(", ")
|
||||
);
|
||||
|
||||
self.session.lock().await.add_assistant_tool_calls(calls.clone());
|
||||
|
||||
for tc in &calls {
|
||||
let result = match &self.tool_executor {
|
||||
Some(exec) => match exec(&tc.function.name, &tc.function.arguments).await {
|
||||
Ok(r) => r,
|
||||
Err(e) => format!("执行失败: {}", e),
|
||||
},
|
||||
None => "工具执行器未配置".to_string(),
|
||||
};
|
||||
tracing::info!("📦 {} → {:.150}", tc.function.name, result);
|
||||
self.session.lock().await.add_tool_result(&tc.id, &tc.function.name, &result);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
// 无工具调用:记录回复,检查是否需要溢出摘要
|
||||
if !full_text.is_empty() {
|
||||
self.session.lock().await.add_assistant(full_text.clone());
|
||||
}
|
||||
|
||||
// 检查 token 是否接近预算,触发溢出摘要
|
||||
{
|
||||
let s = self.session.lock().await;
|
||||
let recent = s.recent_messages();
|
||||
let estimated: u32 = recent.iter().map(|m| builder::estimate_tokens(m)).sum();
|
||||
if estimated > s.token_budget {
|
||||
drop(s);
|
||||
builder::trigger_overflow_summary(&self.session, Some(&self.summarizer())).await;
|
||||
tracing::info!("📦 上下文超预算,已触发溢出摘要");
|
||||
}
|
||||
}
|
||||
|
||||
return Ok((full_text, used_tools, last_usage));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ─── ChatHandle ───
|
||||
|
||||
pub struct ChatHandle {
|
||||
rx: Option<StreamReceiver>,
|
||||
full_text: String,
|
||||
tool_calls: Option<Vec<ToolCall>>,
|
||||
usage: Option<Usage>,
|
||||
session: Arc<Mutex<ChatSession>>,
|
||||
}
|
||||
|
||||
impl ChatHandle {
|
||||
pub async fn next_chunk(&mut self) -> Option<StreamChunk> {
|
||||
let chunk = self.rx.as_mut()?.recv().await;
|
||||
match &chunk {
|
||||
Some(StreamChunk::Text(t)) => self.full_text.push_str(t),
|
||||
Some(StreamChunk::Done { text, tool_calls, usage, .. }) => {
|
||||
if !text.is_empty() { self.full_text = text.clone(); }
|
||||
self.tool_calls = tool_calls.clone();
|
||||
self.usage = usage.clone();
|
||||
if !self.full_text.is_empty() {
|
||||
let mut s = self.session.lock().await;
|
||||
s.add_assistant(self.full_text.clone());
|
||||
}
|
||||
self.rx = None;
|
||||
}
|
||||
Some(StreamChunk::Error(_)) => { self.rx = None; }
|
||||
_ => {}
|
||||
}
|
||||
chunk
|
||||
}
|
||||
|
||||
pub async fn consume(mut self) -> Result<String, String> {
|
||||
while let Some(chunk) = self.next_chunk().await {
|
||||
if let StreamChunk::Error(e) = chunk { return Err(e); }
|
||||
}
|
||||
Ok(self.full_text.clone())
|
||||
}
|
||||
|
||||
pub fn current_text(&self) -> &str { &self.full_text }
|
||||
pub fn get_tool_calls(&self) -> Option<&Vec<ToolCall>> { self.tool_calls.as_ref() }
|
||||
pub fn get_usage(&self) -> Option<&Usage> { self.usage.as_ref() }
|
||||
}
|
||||
|
||||
pub const DEFAULT_SYSTEM_PROMPT: &str = "You are a concise and helpful assistant replying in Chinese. \
|
||||
Keep replies practical and natural. \
|
||||
Unless the user asks for detail, keep most replies under 120 Chinese characters. \
|
||||
\
|
||||
如果需要了解用户的长期记忆或历史对话摘要,使用 read_memories / read_summaries 工具。\
|
||||
如果用户提到重要的个人信息或约定,使用 write_memory 工具记录下来。";
|
||||
@@ -0,0 +1,195 @@
|
||||
use crate::llm::provider::{parse_chat_chunk, LlmProvider, ParsedChunk, StreamReceiver, StreamSender};
|
||||
use crate::llm::types::{ConversationConfig, Message, StreamChunk, ToolCall, Usage};
|
||||
use async_trait::async_trait;
|
||||
use reqwest::Client as HttpClient;
|
||||
use std::collections::BTreeMap;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
pub struct DeepSeekProvider {
|
||||
http: HttpClient,
|
||||
api_key: String,
|
||||
base_url: String,
|
||||
}
|
||||
|
||||
impl DeepSeekProvider {
|
||||
pub fn new() -> Result<Self, String> {
|
||||
let api_key = std::env::var("DEEPSEEK_API_KEY")
|
||||
.map_err(|_| "请设置 DEEPSEEK_API_KEY 环境变量".to_string())?;
|
||||
let base_url = std::env::var("DEEPSEEK_BASE_URL")
|
||||
.unwrap_or_else(|_| "https://api.deepseek.com/v1".to_string());
|
||||
|
||||
Ok(Self {
|
||||
http: HttpClient::builder()
|
||||
.timeout(std::time::Duration::from_secs(120))
|
||||
.build()
|
||||
.map_err(|e| format!("创建 HTTP client 失败: {}", e))?,
|
||||
api_key,
|
||||
base_url,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl LlmProvider for DeepSeekProvider {
|
||||
fn name(&self) -> &str { "deepseek" }
|
||||
|
||||
async fn chat_stream(
|
||||
&self,
|
||||
config: &ConversationConfig,
|
||||
messages: &[Message],
|
||||
) -> Result<StreamReceiver, String> {
|
||||
let url = format!("{}/chat/completions", self.base_url);
|
||||
|
||||
let mut body = serde_json::json!({
|
||||
"model": config.model,
|
||||
"messages": messages,
|
||||
"temperature": config.temperature,
|
||||
"max_tokens": config.max_tokens,
|
||||
"stream": true,
|
||||
});
|
||||
|
||||
// 如果有工具,添加 tools 参数
|
||||
if let Some(ref tools) = config.tools {
|
||||
body["tools"] = serde_json::Value::Array(tools.clone());
|
||||
}
|
||||
|
||||
if config.thinking {
|
||||
body["thinking"] = serde_json::json!({"type": "enabled"});
|
||||
body["reasoning_effort"] = serde_json::json!("high");
|
||||
body.as_object_mut().map(|obj| {
|
||||
obj.remove("temperature");
|
||||
obj.remove("top_p");
|
||||
});
|
||||
}
|
||||
|
||||
let (tx, rx) = mpsc::channel::<StreamChunk>(64);
|
||||
let http = self.http.clone();
|
||||
let api_key = self.api_key.clone();
|
||||
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = stream_deepseek(http, &url, &api_key, body, tx.clone()).await {
|
||||
let _ = tx.send(StreamChunk::Error(e)).await;
|
||||
}
|
||||
});
|
||||
|
||||
Ok(rx)
|
||||
}
|
||||
}
|
||||
|
||||
async fn stream_deepseek(
|
||||
http: HttpClient,
|
||||
url: &str,
|
||||
api_key: &str,
|
||||
body: serde_json::Value,
|
||||
tx: StreamSender,
|
||||
) -> Result<(), String> {
|
||||
let response = http
|
||||
.post(url)
|
||||
.header("Authorization", format!("Bearer {}", api_key))
|
||||
.header("Content-Type", "application/json")
|
||||
.header("Accept", "text/event-stream")
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("请求失败: {}", e))?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
let status = response.status();
|
||||
let text = response.text().await.unwrap_or_default();
|
||||
let _ = tx.send(StreamChunk::Error(format!("HTTP {}: {}", status, text))).await;
|
||||
return Err(format!("HTTP {}: {}", status, text));
|
||||
}
|
||||
|
||||
let mut full_text = String::new();
|
||||
let mut full_reasoning = String::new();
|
||||
let mut usage: Option<Usage> = None;
|
||||
let mut tool_call_builders: BTreeMap<i32, (Option<String>, Option<String>, String)> = BTreeMap::new();
|
||||
|
||||
use futures_util::StreamExt;
|
||||
let mut buf = String::new();
|
||||
let mut stream = response.bytes_stream();
|
||||
|
||||
while let Some(chunk_result) = stream.next().await {
|
||||
let chunk = chunk_result.map_err(|e| format!("读取流失败: {}", e))?;
|
||||
let text = String::from_utf8_lossy(&chunk);
|
||||
buf.push_str(&text);
|
||||
|
||||
while let Some(line_end) = buf.find('\n') {
|
||||
let line = buf[..line_end].to_string();
|
||||
buf = buf[line_end + 1..].to_string();
|
||||
let trimmed = line.trim();
|
||||
if trimmed.is_empty() { continue; }
|
||||
|
||||
if let Some(parsed) = parse_chat_chunk(trimmed) {
|
||||
match parsed {
|
||||
ParsedChunk::Text(t) => {
|
||||
full_text.push_str(&t);
|
||||
let _ = tx.send(StreamChunk::Text(t)).await;
|
||||
}
|
||||
ParsedChunk::Reasoning(r) => {
|
||||
full_reasoning.push_str(&r);
|
||||
let _ = tx.send(StreamChunk::Reasoning(r)).await;
|
||||
}
|
||||
ParsedChunk::ToolCallDelta { index, id, name, arguments } => {
|
||||
let entry = tool_call_builders.entry(index).or_insert((None, None, String::new()));
|
||||
if let Some(call_id) = id { entry.0 = Some(call_id); }
|
||||
if let Some(call_name) = name { entry.1 = Some(call_name); }
|
||||
entry.2.push_str(&arguments);
|
||||
}
|
||||
ParsedChunk::FinishReason(_) => {}
|
||||
ParsedChunk::Usage(u) => { usage = Some(u); }
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 检查最后的 buffer
|
||||
if !buf.trim().is_empty() {
|
||||
if let Some(parsed) = parse_chat_chunk(buf.trim()) {
|
||||
match parsed {
|
||||
ParsedChunk::Text(t) => full_text.push_str(&t),
|
||||
ParsedChunk::ToolCallDelta { index, id, name, arguments } => {
|
||||
let entry = tool_call_builders.entry(index).or_insert((None, None, String::new()));
|
||||
if let Some(call_id) = id { entry.0 = Some(call_id); }
|
||||
if let Some(call_name) = name { entry.1 = Some(call_name); }
|
||||
entry.2.push_str(&arguments);
|
||||
}
|
||||
ParsedChunk::Usage(u) => { usage = Some(u); }
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 构建 tool_calls
|
||||
let tool_calls: Option<Vec<ToolCall>> = if tool_call_builders.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(
|
||||
tool_call_builders
|
||||
.into_values()
|
||||
.filter_map(|(id, name, args)| {
|
||||
let name = name?;
|
||||
let id = id.unwrap_or_else(|| format!("call_{}", rand::random::<u32>()));
|
||||
Some(ToolCall {
|
||||
id,
|
||||
call_type: "function".to_string(),
|
||||
function: crate::llm::types::ToolFunctionCall {
|
||||
name,
|
||||
arguments: args,
|
||||
},
|
||||
})
|
||||
})
|
||||
.collect(),
|
||||
)
|
||||
};
|
||||
|
||||
let _ = tx
|
||||
.send(StreamChunk::Done {
|
||||
text: full_text,
|
||||
reasoning: if full_reasoning.is_empty() { None } else { Some(full_reasoning) },
|
||||
tool_calls,
|
||||
usage,
|
||||
})
|
||||
.await;
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,157 @@
|
||||
use crate::llm::provider::{parse_chat_chunk, LlmProvider, ParsedChunk, StreamReceiver, StreamSender};
|
||||
use crate::llm::types::{ConversationConfig, Message, StreamChunk};
|
||||
use async_trait::async_trait;
|
||||
use reqwest::Client as HttpClient;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
/// LM Studio 提供商(OpenAI 兼容接口)
|
||||
pub struct LmStudioProvider {
|
||||
http: HttpClient,
|
||||
base_url: String,
|
||||
api_key: String,
|
||||
}
|
||||
|
||||
impl LmStudioProvider {
|
||||
pub fn new() -> Result<Self, String> {
|
||||
let base_url = std::env::var("LM_STUDIO_BASE_URL")
|
||||
.unwrap_or_else(|_| "http://localhost:1234/v1".to_string());
|
||||
// 标准化 base_url: 确保有协议和 /v1 路径
|
||||
let base_url = normalize_base_url(&base_url);
|
||||
let api_key = std::env::var("LM_STUDIO_API_KEY").unwrap_or_default();
|
||||
|
||||
Ok(Self {
|
||||
http: HttpClient::builder()
|
||||
.timeout(std::time::Duration::from_secs(120))
|
||||
.build()
|
||||
.map_err(|e| format!("创建 HTTP client 失败: {}", e))?,
|
||||
base_url,
|
||||
api_key,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl LlmProvider for LmStudioProvider {
|
||||
fn name(&self) -> &str {
|
||||
"lmstudio"
|
||||
}
|
||||
|
||||
async fn chat_stream(
|
||||
&self,
|
||||
config: &ConversationConfig,
|
||||
messages: &[Message],
|
||||
) -> Result<StreamReceiver, String> {
|
||||
let url = format!("{}/chat/completions", self.base_url);
|
||||
|
||||
let body = serde_json::json!({
|
||||
"model": config.model,
|
||||
"messages": messages,
|
||||
"temperature": config.temperature,
|
||||
"max_tokens": config.max_tokens,
|
||||
"stream": true,
|
||||
});
|
||||
|
||||
let (tx, rx) = mpsc::channel::<StreamChunk>(64);
|
||||
let http = self.http.clone();
|
||||
let api_key = self.api_key.clone();
|
||||
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = stream_openai(http, &url, &api_key, body, tx.clone()).await {
|
||||
let _ = tx.send(StreamChunk::Error(e)).await;
|
||||
}
|
||||
});
|
||||
|
||||
Ok(rx)
|
||||
}
|
||||
}
|
||||
|
||||
/// 与 DeepSeek 公用相同的 SSE 流式处理,但去除 thinking 相关字段
|
||||
async fn stream_openai(
|
||||
http: HttpClient,
|
||||
url: &str,
|
||||
api_key: &str,
|
||||
body: serde_json::Value,
|
||||
tx: StreamSender,
|
||||
) -> Result<(), String> {
|
||||
let mut req = http.post(url).header("Content-Type", "application/json");
|
||||
|
||||
if !api_key.is_empty() {
|
||||
req = req.header("Authorization", format!("Bearer {}", api_key));
|
||||
}
|
||||
|
||||
let response = req
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("请求失败: {}", e))?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
let status = response.status();
|
||||
let text = response.text().await.unwrap_or_default();
|
||||
let _ = tx.send(StreamChunk::Error(format!("HTTP {}: {}", status, text))).await;
|
||||
return Err(format!("HTTP {}: {}", status, text));
|
||||
}
|
||||
|
||||
let mut full_text = String::new();
|
||||
let mut buf = String::new();
|
||||
let mut stream = response.bytes_stream();
|
||||
|
||||
use futures_util::StreamExt;
|
||||
while let Some(chunk_result) = stream.next().await {
|
||||
let chunk = chunk_result.map_err(|e| format!("读取流失败: {}", e))?;
|
||||
let text = String::from_utf8_lossy(&chunk);
|
||||
buf.push_str(&text);
|
||||
|
||||
while let Some(line_end) = buf.find('\n') {
|
||||
let line = buf[..line_end].to_string();
|
||||
buf = buf[line_end + 1..].to_string();
|
||||
let trimmed = line.trim();
|
||||
if trimmed.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
// LM Studio 使用标准 OpenAI SSE 格式
|
||||
if let Some(parsed) = parse_chat_chunk(trimmed) {
|
||||
match parsed {
|
||||
ParsedChunk::Text(t) => {
|
||||
full_text.push_str(&t);
|
||||
let _ = tx.send(StreamChunk::Text(t)).await;
|
||||
}
|
||||
ParsedChunk::Reasoning(r) => {
|
||||
let _ = tx.send(StreamChunk::Reasoning(r)).await;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let _ = tx
|
||||
.send(StreamChunk::Done {
|
||||
text: full_text,
|
||||
reasoning: None,
|
||||
tool_calls: None,
|
||||
usage: None,
|
||||
})
|
||||
.await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 确保 base_url 有协议前缀和 /v1 路径
|
||||
fn normalize_base_url(raw: &str) -> String {
|
||||
let raw = raw.trim();
|
||||
let with_protocol = if raw.starts_with("http://") || raw.starts_with("https://") {
|
||||
raw.to_string()
|
||||
} else {
|
||||
format!("http://{}", raw)
|
||||
};
|
||||
|
||||
// 去掉末尾 /
|
||||
let trimmed = with_protocol.trim_end_matches('/');
|
||||
// 确保路径以 /v1 结尾
|
||||
if trimmed.ends_with("/v1") {
|
||||
trimmed.to_string()
|
||||
} else {
|
||||
format!("{}/v1", trimmed)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,8 @@
|
||||
pub mod conversation;
|
||||
pub mod deepseek;
|
||||
pub mod lmstudio;
|
||||
pub mod provider;
|
||||
pub mod types;
|
||||
|
||||
pub use conversation::{Conversation, Summarizer, ToolExecutor, DEFAULT_SYSTEM_PROMPT};
|
||||
pub use types::{ConversationConfig, Message, Role, StreamChunk, Usage};
|
||||
@@ -0,0 +1,187 @@
|
||||
use crate::llm::types::{ConversationConfig, Message, StreamChunk};
|
||||
use serde::Deserialize;
|
||||
use async_trait::async_trait;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
/// 流式返回:每个 StreamChunk 是 delta 或控制信号
|
||||
pub type StreamReceiver = mpsc::Receiver<StreamChunk>;
|
||||
pub type StreamSender = mpsc::Sender<StreamChunk>;
|
||||
|
||||
/// LLM 提供商抽象
|
||||
#[async_trait]
|
||||
pub trait LlmProvider: Send + Sync {
|
||||
/// 提供商名称
|
||||
fn name(&self) -> &str;
|
||||
|
||||
/// 发起流式对话
|
||||
async fn chat_stream(
|
||||
&self,
|
||||
config: &ConversationConfig,
|
||||
messages: &[Message],
|
||||
) -> Result<StreamReceiver, String>;
|
||||
}
|
||||
|
||||
// ─── 内置提供商注册 ───
|
||||
|
||||
pub type BoxedProvider = Arc<dyn LlmProvider>;
|
||||
|
||||
/// 从配置创建恰当的提供商
|
||||
pub fn create_provider(_config: &ConversationConfig) -> Result<BoxedProvider, String> {
|
||||
// 根据 model 前缀或环境变量选择
|
||||
let provider = std::env::var("LLM_PROVIDER")
|
||||
.unwrap_or_else(|_| "deepseek".to_string())
|
||||
.to_lowercase();
|
||||
|
||||
match provider.as_str() {
|
||||
"deepseek" => Ok(Arc::new(super::deepseek::DeepSeekProvider::new()?)),
|
||||
"lmstudio" => Ok(Arc::new(super::lmstudio::LmStudioProvider::new()?)),
|
||||
_ => Err(format!("不支持的 LLM 提供商: {}", provider)),
|
||||
}
|
||||
}
|
||||
|
||||
// ─── 内部:SSE 解析工具 ───
|
||||
|
||||
/// 从字节流中解析 SSE 行
|
||||
pub(crate) fn parse_sse_line(line: &str) -> Option<(String, String)> {
|
||||
// 支持两种格式:
|
||||
// data: {...}
|
||||
// event: ...
|
||||
let line = line.trim();
|
||||
if line.is_empty() || line.starts_with(':') {
|
||||
return None;
|
||||
}
|
||||
|
||||
if let Some(pos) = line.find(": ") {
|
||||
let field = line[..pos].trim().to_string();
|
||||
let value = line[pos + 2..].trim_start().to_string();
|
||||
Some((field, value))
|
||||
} else if let Some(pos) = line.find(':') {
|
||||
let field = line[..pos].trim().to_string();
|
||||
let value = line[pos + 1..].trim_start().to_string();
|
||||
Some((field, value))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// 流式块解析结果
|
||||
pub(crate) enum ParsedChunk {
|
||||
Text(String),
|
||||
Reasoning(String),
|
||||
ToolCallDelta {
|
||||
index: i32,
|
||||
id: Option<String>,
|
||||
name: Option<String>,
|
||||
arguments: String,
|
||||
},
|
||||
FinishReason(String),
|
||||
Usage(super::types::Usage),
|
||||
}
|
||||
|
||||
/// 从 JSON body 中解析 DeepSeek/OpenAI 流式 delta
|
||||
pub(crate) fn parse_chat_chunk(
|
||||
line: &str,
|
||||
) -> Option<ParsedChunk> {
|
||||
if !line.starts_with("data: ") {
|
||||
return None;
|
||||
}
|
||||
|
||||
let data = line["data: ".len()..].trim();
|
||||
if data == "[DONE]" {
|
||||
return None;
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ToolCallDelta {
|
||||
#[serde(default)]
|
||||
index: Option<i32>,
|
||||
#[serde(default)]
|
||||
id: Option<String>,
|
||||
#[serde(rename = "type", default)]
|
||||
call_type: Option<String>,
|
||||
#[serde(default)]
|
||||
function: Option<ToolFunctionDelta>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ToolFunctionDelta {
|
||||
#[serde(default)]
|
||||
name: Option<String>,
|
||||
#[serde(default)]
|
||||
arguments: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct Delta {
|
||||
#[serde(default)]
|
||||
content: Option<String>,
|
||||
#[serde(default)]
|
||||
reasoning_content: Option<String>,
|
||||
#[serde(default)]
|
||||
tool_calls: Option<Vec<ToolCallDelta>>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ChunkChoice {
|
||||
delta: Delta,
|
||||
#[serde(default)]
|
||||
finish_reason: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ChunkResponse {
|
||||
choices: Vec<ChunkChoice>,
|
||||
#[serde(default)]
|
||||
usage: Option<super::types::Usage>,
|
||||
}
|
||||
|
||||
let parsed: ChunkResponse = match serde_json::from_str(data) {
|
||||
Ok(p) => p,
|
||||
Err(_) => return None,
|
||||
};
|
||||
|
||||
// 提取 usage(可能在最后一个 chunk)
|
||||
if let Some(ref usage) = parsed.usage {
|
||||
if usage.total_tokens > 0 {
|
||||
return Some(ParsedChunk::Usage(usage.clone()));
|
||||
}
|
||||
}
|
||||
|
||||
for choice in parsed.choices {
|
||||
// 工具调用 delta
|
||||
if let Some(tool_calls) = &choice.delta.tool_calls {
|
||||
for tc in tool_calls {
|
||||
let idx = tc.index.unwrap_or(0);
|
||||
let args = tc.function.as_ref()
|
||||
.and_then(|f| f.arguments.clone())
|
||||
.unwrap_or_default();
|
||||
let name = tc.function.as_ref().and_then(|f| f.name.clone());
|
||||
return Some(ParsedChunk::ToolCallDelta {
|
||||
index: idx,
|
||||
id: tc.id.clone(),
|
||||
name,
|
||||
arguments: args,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(reasoning) = &choice.delta.reasoning_content {
|
||||
if !reasoning.is_empty() {
|
||||
return Some(ParsedChunk::Reasoning(reasoning.clone()));
|
||||
}
|
||||
}
|
||||
if let Some(content) = &choice.delta.content {
|
||||
if !content.is_empty() {
|
||||
return Some(ParsedChunk::Text(content.clone()));
|
||||
}
|
||||
}
|
||||
if let Some(reason) = &choice.finish_reason {
|
||||
if !reason.is_empty() {
|
||||
return Some(ParsedChunk::FinishReason(reason.clone()));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
@@ -0,0 +1,125 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
// ─── 角色 ───
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum Role {
|
||||
System,
|
||||
User,
|
||||
Assistant,
|
||||
Tool,
|
||||
}
|
||||
|
||||
// ─── 消息 ───
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Message {
|
||||
pub role: Role,
|
||||
#[serde(default)]
|
||||
pub content: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub tool_call_id: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub name: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub tool_calls: Option<Vec<ToolCall>>,
|
||||
}
|
||||
|
||||
impl Message {
|
||||
pub fn system(content: impl Into<String>) -> Self {
|
||||
Self { role: Role::System, content: content.into(), tool_call_id: None, name: None, tool_calls: None }
|
||||
}
|
||||
|
||||
pub fn user(content: impl Into<String>) -> Self {
|
||||
Self { role: Role::User, content: content.into(), tool_call_id: None, name: None, tool_calls: None }
|
||||
}
|
||||
|
||||
pub fn assistant(content: impl Into<String>) -> Self {
|
||||
Self { role: Role::Assistant, content: content.into(), tool_call_id: None, name: None, tool_calls: None }
|
||||
}
|
||||
|
||||
pub fn assistant_with_tool_calls(tool_calls: Vec<ToolCall>) -> Self {
|
||||
Self { role: Role::Assistant, content: String::new(), tool_call_id: None, name: None, tool_calls: Some(tool_calls) }
|
||||
}
|
||||
|
||||
pub fn tool_result(tool_call_id: &str, name: &str, result: &str) -> Self {
|
||||
Self { role: Role::Tool, content: result.to_string(), tool_call_id: Some(tool_call_id.to_string()), name: Some(name.to_string()), tool_calls: None }
|
||||
}
|
||||
}
|
||||
|
||||
// ─── 工具调用 ───
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ToolCall {
|
||||
pub id: String,
|
||||
#[serde(rename = "type")]
|
||||
pub call_type: String, // "function"
|
||||
pub function: ToolFunctionCall,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ToolFunctionCall {
|
||||
pub name: String,
|
||||
pub arguments: String, // JSON string
|
||||
}
|
||||
|
||||
// ─── 流式响应块 ───
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum StreamChunk {
|
||||
/// 文本片段 delta
|
||||
Text(String),
|
||||
/// 思考/推理内容(DeepSeek thinking)
|
||||
Reasoning(String),
|
||||
/// 完成(携带完整文本,以及可选的工具调用)
|
||||
Done {
|
||||
text: String,
|
||||
reasoning: Option<String>,
|
||||
tool_calls: Option<Vec<ToolCall>>,
|
||||
usage: Option<Usage>,
|
||||
},
|
||||
/// 错误
|
||||
Error(String),
|
||||
}
|
||||
|
||||
// ─── Token 用量 ───
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||
pub struct Usage {
|
||||
pub prompt_tokens: u32,
|
||||
pub completion_tokens: u32,
|
||||
pub total_tokens: u32,
|
||||
#[serde(default)]
|
||||
pub prompt_cache_hit_tokens: u32,
|
||||
#[serde(default)]
|
||||
pub prompt_cache_miss_tokens: u32,
|
||||
}
|
||||
|
||||
// ─── 对话配置 ───
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ConversationConfig {
|
||||
pub system_prompt: String,
|
||||
pub model: String,
|
||||
pub temperature: f32,
|
||||
pub max_tokens: u32,
|
||||
pub thinking: bool,
|
||||
/// LLM function calling 的工具定义(JSON 数组)
|
||||
pub tools: Option<Vec<serde_json::Value>>,
|
||||
}
|
||||
|
||||
impl Default for ConversationConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
system_prompt: "You are a concise and helpful assistant replying in Chinese. \
|
||||
Keep replies practical and natural."
|
||||
.to_string(),
|
||||
model: "deepseek-v4-flash".to_string(),
|
||||
temperature: 0.7,
|
||||
max_tokens: 4096,
|
||||
thinking: true,
|
||||
tools: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
+635
-2
@@ -1,3 +1,636 @@
|
||||
fn main() {
|
||||
println!("Hello, world!");
|
||||
mod cli;
|
||||
mod config;
|
||||
mod context;
|
||||
mod db;
|
||||
mod llm;
|
||||
mod scheduler;
|
||||
mod state;
|
||||
mod tools;
|
||||
mod wechat;
|
||||
|
||||
use clap::Parser;
|
||||
use cli::{Cli, Commands};
|
||||
use context::MemoryStore;
|
||||
use config::AppConfig;
|
||||
use context::types::ChatSession as _;
|
||||
use db::Database;
|
||||
use llm::{Conversation, ConversationConfig, ToolExecutor, Usage, DEFAULT_SYSTEM_PROMPT};
|
||||
use std::sync::Arc;
|
||||
use tools::approval::ApprovalManager;
|
||||
use tools::registry::{ExecutionContext, SkillRegistry};
|
||||
use tracing::{error, info};
|
||||
use tracing_subscriber::EnvFilter;
|
||||
use wechat::client::WeChatClient;
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
tracing_subscriber::fmt()
|
||||
.with_env_filter(
|
||||
EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| EnvFilter::new("info")),
|
||||
)
|
||||
.init();
|
||||
|
||||
dotenvy::dotenv().ok();
|
||||
|
||||
let config = AppConfig::from_file("config.json");
|
||||
|
||||
// 初始化数据库(如果配置了 DATABASE_URL)
|
||||
let database = match Database::connect().await {
|
||||
Ok(db) => {
|
||||
info!("✅ 数据库连接成功");
|
||||
Some(Arc::new(db))
|
||||
}
|
||||
Err(e) => {
|
||||
info!("数据库不可用,使用文件存储: {}", e);
|
||||
None
|
||||
}
|
||||
};
|
||||
|
||||
// 长期记忆存储
|
||||
let memory_store = Arc::new(MemoryStore::new(database.as_ref().map(|db| Arc::new(db.pool().clone()))));
|
||||
memory_store.load().await;
|
||||
|
||||
// 文件状态管理器(降级方案)
|
||||
let file_state = state::StateManager::new(&config.storage.state_dir);
|
||||
|
||||
let cli = Cli::parse();
|
||||
|
||||
match cli.command {
|
||||
Commands::Login { timeout } => {
|
||||
cmd_login(timeout, &database, &file_state).await;
|
||||
}
|
||||
Commands::Listen {
|
||||
llm: enable_llm,
|
||||
echo,
|
||||
} => {
|
||||
cmd_listen(enable_llm, echo, &database, &file_state, &config, &memory_store).await;
|
||||
}
|
||||
Commands::Send {
|
||||
to,
|
||||
text,
|
||||
context_token,
|
||||
} => {
|
||||
cmd_send(&to, &text, context_token.as_deref(), &database, &file_state).await;
|
||||
}
|
||||
Commands::Whoami => {
|
||||
cmd_whoami(&database, &file_state).await;
|
||||
}
|
||||
Commands::Usage { since, until, model } => {
|
||||
cmd_usage(&database, since, until, model).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ─── 命令实现 ───
|
||||
|
||||
async fn cmd_login(
|
||||
timeout_secs: u64,
|
||||
database: &Option<Arc<Database>>,
|
||||
file_state: &state::StateManager,
|
||||
) {
|
||||
let client = WeChatClient::new(None);
|
||||
|
||||
match client
|
||||
.login(|url| println!("二维码链接: {}", url), timeout_secs)
|
||||
.await
|
||||
{
|
||||
Ok(result) => {
|
||||
let auth = state::AuthState {
|
||||
token: result.token.clone(),
|
||||
account_id: result.account_id.clone(),
|
||||
base_url: result.base_url.clone(),
|
||||
};
|
||||
|
||||
// 优先存数据库
|
||||
if let Some(db) = database {
|
||||
if let Err(e) = db::models::save_auth(db.pool(), &auth).await {
|
||||
error!("保存 auth 到数据库失败: {}", e);
|
||||
file_state.save_auth(&auth);
|
||||
} else {
|
||||
info!("认证信息已保存到数据库");
|
||||
// 清理文件状态(迁移完成)
|
||||
file_state.clear_auth();
|
||||
}
|
||||
} else {
|
||||
file_state.save_auth(&auth);
|
||||
}
|
||||
|
||||
println!("\n✅ 登录成功!");
|
||||
println!(" 账号: {}", result.account_id);
|
||||
println!(" Token: {}...", &result.token[..std::cmp::min(16, result.token.len())]);
|
||||
println!(" API: {}", result.base_url);
|
||||
}
|
||||
Err(e) => error!("登录失败: {}", e),
|
||||
}
|
||||
}
|
||||
|
||||
async fn cmd_listen(
|
||||
enable_llm: bool,
|
||||
echo: bool,
|
||||
database: &Option<Arc<Database>>,
|
||||
file_state: &state::StateManager,
|
||||
config: &AppConfig,
|
||||
memory_store: &Arc<MemoryStore>,
|
||||
) {
|
||||
// 从数据库或文件加载认证
|
||||
let auth = 'load: {
|
||||
if let Some(db) = database {
|
||||
if let Some(a) = db::models::load_auth(db.pool()).await {
|
||||
break 'load Some(a);
|
||||
}
|
||||
// 尝试从文件迁移
|
||||
if let Some(fa) = file_state.load_auth() {
|
||||
let a = state::AuthState {
|
||||
token: fa.token,
|
||||
account_id: fa.account_id,
|
||||
base_url: fa.base_url,
|
||||
};
|
||||
let _ = db::models::save_auth(db.pool(), &a).await;
|
||||
break 'load Some(a);
|
||||
}
|
||||
} else if let Some(a) = file_state.load_auth() {
|
||||
break 'load Some(a);
|
||||
}
|
||||
break 'load None;
|
||||
};
|
||||
|
||||
let auth = match auth {
|
||||
Some(a) => a,
|
||||
None => {
|
||||
error!("未登录,请先执行 login 命令");
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let client = WeChatClient::new(Some(auth.base_url.clone()));
|
||||
client
|
||||
.set_auth(&auth.token, &auth.account_id, &auth.base_url)
|
||||
.await;
|
||||
|
||||
// 从文件恢复 runtime 状态(get_updates_buf 暂时保留文件存储)
|
||||
if let Some(r) = file_state.load_runtime() {
|
||||
client.set_updates_buf(&r.get_updates_buf).await;
|
||||
}
|
||||
|
||||
if let Err(e) = client.notify_start().await {
|
||||
error!("注册监听器失败: {}", e);
|
||||
return;
|
||||
}
|
||||
|
||||
// 加载系统 Skills
|
||||
let skill_registry = match SkillRegistry::load("resources/skills", "skills").await {
|
||||
Ok(r) => {
|
||||
info!("已加载 {} 个技能", r.count());
|
||||
Some(Arc::new(r))
|
||||
}
|
||||
Err(e) => {
|
||||
info!("技能加载结果: {:?}", e);
|
||||
None
|
||||
}
|
||||
};
|
||||
|
||||
// 审批管理器
|
||||
let approval_manager = database.as_ref().map(|db| {
|
||||
Arc::new(ApprovalManager::new(Some(Arc::new(db.pool().clone()))))
|
||||
});
|
||||
|
||||
let conversation = if enable_llm {
|
||||
let sys_prompt = std::env::var("WEIXIN_LLM_SYSTEM_PROMPT")
|
||||
.unwrap_or_else(|_| DEFAULT_SYSTEM_PROMPT.to_string());
|
||||
let model = config.llm.deepseek.model.clone();
|
||||
|
||||
let mut cfg = ConversationConfig {
|
||||
system_prompt: sys_prompt,
|
||||
model,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
// 构建 LLM tools 列表(skills + 上下文工具)
|
||||
let mut tools_list: Vec<serde_json::Value> = Vec::new();
|
||||
|
||||
if let Some(ref registry) = skill_registry {
|
||||
for spec in registry.list_specs() {
|
||||
tools_list.push(serde_json::json!({
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": spec.name,
|
||||
"description": spec.description,
|
||||
"parameters": spec.parameters,
|
||||
}
|
||||
}));
|
||||
}
|
||||
}
|
||||
|
||||
// 上下文工具(长期记忆 + 摘要)
|
||||
tools_list.push(serde_json::json!({
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "read_memories",
|
||||
"description": "读取用户的长期记忆(跨会话持久保存的个人信息、偏好、约定等)",
|
||||
"parameters": { "type": "object", "properties": {} }
|
||||
}
|
||||
}));
|
||||
tools_list.push(serde_json::json!({
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "write_memory",
|
||||
"description": "写入一条用户的长期记忆(当用户提到重要偏好、个人信息、约定时调用)",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"content": { "type": "string", "description": "记忆内容" }
|
||||
},
|
||||
"required": ["content"]
|
||||
}
|
||||
}
|
||||
}));
|
||||
tools_list.push(serde_json::json!({
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "read_summaries",
|
||||
"description": "读取历史会话摘要(包括因上下文压缩或空闲超时产生的摘要)",
|
||||
"parameters": { "type": "object", "properties": {} }
|
||||
}
|
||||
}));
|
||||
|
||||
cfg.tools = Some(tools_list);
|
||||
|
||||
// 创建 Conversation
|
||||
let mut conv = match Conversation::new(cfg) {
|
||||
Ok(c) => {
|
||||
info!("LLM 已初始化: {} / {}", c.provider_name(), c.model());
|
||||
c
|
||||
}
|
||||
Err(e) => {
|
||||
error!("初始化 LLM 失败: {}", e);
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
// 构建工具执行器
|
||||
let session = conv.session();
|
||||
let registry_opt = skill_registry.clone();
|
||||
let approval = approval_manager.clone();
|
||||
let memory_store = memory_store.clone();
|
||||
let send_wechat: tools::registry::WechatSender = {
|
||||
let client = client.clone();
|
||||
Arc::new(move |user_id: &str, text: &str| {
|
||||
let client = client.clone();
|
||||
let uid = user_id.to_string();
|
||||
let msg = text.to_string();
|
||||
Box::pin(async move {
|
||||
client.send_text(&uid, &msg, None).await.map(|_| ()).map_err(|e| e)
|
||||
})
|
||||
})
|
||||
};
|
||||
|
||||
let executor: ToolExecutor = Arc::new(move |name: &str, args_json: &str| {
|
||||
let reg = registry_opt.clone();
|
||||
let approval = approval.clone();
|
||||
let send = send_wechat.clone();
|
||||
let session = session.clone();
|
||||
let memory_store = memory_store.clone();
|
||||
let n = name.to_string();
|
||||
let args = args_json.to_string();
|
||||
|
||||
Box::pin(async move {
|
||||
match n.as_str() {
|
||||
"read_memories" => return Ok(memory_store.read().await),
|
||||
"write_memory" => {
|
||||
let params: serde_json::Value = serde_json::from_str(&args).unwrap_or_default();
|
||||
let content = params.get("content").and_then(|v| v.as_str()).unwrap_or("");
|
||||
if content.is_empty() { return Ok("请提供 content 参数".to_string()); }
|
||||
return Ok(memory_store.write(content).await);
|
||||
}
|
||||
"read_summaries" => return Ok(context::tools::read_summaries(&session).await),
|
||||
_ => {}
|
||||
}
|
||||
if let Some(ref reg) = reg {
|
||||
let params: serde_json::Value = serde_json::from_str(&args).unwrap_or(serde_json::json!({}));
|
||||
let mut ctx = ExecutionContext::new("wechat_user");
|
||||
if let Some(a) = approval { ctx = ctx.with_approval(a); }
|
||||
ctx = ctx.with_wechat(send);
|
||||
let result = reg.execute(&n, params, &ctx).await;
|
||||
Ok(result.output)
|
||||
} else {
|
||||
Ok(format!("未知工具: {}", n))
|
||||
}
|
||||
})
|
||||
});
|
||||
|
||||
conv.set_tool_executor(executor);
|
||||
|
||||
Some(conv)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let account_id = auth.account_id.clone();
|
||||
|
||||
// 启动调度器(如果数据库可用)
|
||||
if let Some(sched_db) = database.as_ref().map(|db| db.pool().clone()) {
|
||||
let sched_client = client.clone();
|
||||
tokio::spawn(async move {
|
||||
let scheduler = scheduler::Scheduler::new(Some(Arc::new(sched_db)));
|
||||
scheduler.run(move |to: &str, _name: &str, msg: &str| {
|
||||
let c = sched_client.clone();
|
||||
let uid = to.to_string();
|
||||
let text = msg.to_string();
|
||||
tokio::task::spawn(async move {
|
||||
if let Err(e) = c.send_text(&uid, &text, None).await {
|
||||
tracing::error!("发送调度结果失败: {}", e);
|
||||
}
|
||||
});
|
||||
Ok(())
|
||||
}).await;
|
||||
});
|
||||
info!("定时任务调度器已启动");
|
||||
}
|
||||
|
||||
info!("开始监听消息 (llm={}, echo={})...", enable_llm, echo);
|
||||
println!("已开始监听消息,按 Ctrl+C 退出");
|
||||
|
||||
let shutdown = async { tokio::signal::ctrl_c().await.expect("Ctrl+C 失败") };
|
||||
|
||||
tokio::select! {
|
||||
_ = shutdown => {
|
||||
info!("收到退出信号,注销监听器...");
|
||||
let _ = client.notify_stop().await;
|
||||
file_state.save_runtime(&client.get_updates_buf().await);
|
||||
info!("已退出");
|
||||
}
|
||||
_ = listen_loop(&client, enable_llm, echo, &account_id, database, file_state, conversation, approval_manager) => {}
|
||||
}
|
||||
}
|
||||
|
||||
async fn listen_loop(
|
||||
client: &WeChatClient,
|
||||
enable_llm: bool,
|
||||
echo: bool,
|
||||
account_id: &str,
|
||||
database: &Option<Arc<Database>>,
|
||||
file_state: &state::StateManager,
|
||||
conversation: Option<Conversation>,
|
||||
approval_manager: Option<Arc<ApprovalManager>>,
|
||||
) {
|
||||
loop {
|
||||
match client.receive_messages().await {
|
||||
Ok(resp) => {
|
||||
// 持久化 runtime buf
|
||||
if !resp.get_updates_buf.is_empty() {
|
||||
file_state.save_runtime(&resp.get_updates_buf);
|
||||
}
|
||||
|
||||
for msg in &resp.msgs {
|
||||
if msg.msg_type != Some(wechat::types::WeixinMessage::TYPE_USER) {
|
||||
continue;
|
||||
}
|
||||
|
||||
let from = msg.from_user_id.as_deref().unwrap_or("unknown");
|
||||
let text = msg.text_content().unwrap_or("(非文本消息)");
|
||||
let ctx_token = msg.context_token.as_deref();
|
||||
let msg_id = msg.message_id.unwrap_or(0).to_string();
|
||||
|
||||
info!("收到消息 from={}: {}", from, text);
|
||||
|
||||
// === 入库:收到的消息 ===
|
||||
if let Some(db) = database {
|
||||
let _ = db::models::insert_chat_record(
|
||||
db.pool(),
|
||||
"inbound",
|
||||
from,
|
||||
account_id,
|
||||
text,
|
||||
"wechat",
|
||||
ctx_token,
|
||||
&msg_id,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
// Echo 模式
|
||||
if echo {
|
||||
match client.send_text(from, text, ctx_token).await {
|
||||
Ok(sent_id) => {
|
||||
info!("回显成功 msg_id={}", sent_id);
|
||||
// === 入库:发送的回显 ===
|
||||
if let Some(db) = database {
|
||||
let _ = db::models::insert_chat_record(
|
||||
db.pool(),
|
||||
"outbound",
|
||||
from,
|
||||
account_id,
|
||||
text,
|
||||
"echo",
|
||||
ctx_token,
|
||||
&sent_id,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
Err(e) => error!("回显失败: {}", e),
|
||||
}
|
||||
}
|
||||
|
||||
// 审批回复检查:是否匹配待审批的确认码
|
||||
let is_approval_reply = if let Some(ref mgr) = approval_manager {
|
||||
mgr.handle_reply(from, text).await.is_some()
|
||||
} else {
|
||||
false
|
||||
};
|
||||
|
||||
if is_approval_reply {
|
||||
info!("消息匹配审批确认码,不传给 LLM(已由审批系统处理)");
|
||||
continue;
|
||||
}
|
||||
|
||||
// LLM 回复
|
||||
if enable_llm {
|
||||
if let Some(ref conv) = conversation {
|
||||
let reply_result = if conv.spec().tools.is_some() {
|
||||
generate_reply_with_tools(conv, text.to_string(), database).await
|
||||
} else {
|
||||
generate_reply(conv, text.to_string(), database).await
|
||||
};
|
||||
|
||||
match reply_result {
|
||||
Ok(reply) => {
|
||||
match client.send_text(from, &reply, ctx_token).await {
|
||||
Ok(sent_id) => {
|
||||
info!("LLM 回复成功 msg_id={}", sent_id);
|
||||
if let Some(db) = database {
|
||||
let _ = db::models::insert_chat_record(
|
||||
db.pool(), "outbound", from, account_id,
|
||||
&reply, "llm", ctx_token, &sent_id,
|
||||
).await;
|
||||
}
|
||||
}
|
||||
Err(e) => error!("发送 LLM 回复失败: {}", e),
|
||||
}
|
||||
}
|
||||
Err(e) => error!("生成 LLM 回复失败: {}", e),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
if !e.contains("超时") && !e.contains("timeout") {
|
||||
error!("接收消息失败: {}", e);
|
||||
tokio::time::sleep(std::time::Duration::from_secs(2)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn generate_reply(conv: &Conversation, user_text: String, db: &Option<Arc<Database>>) -> Result<String, String> {
|
||||
let handle = conv.chat(user_text).await?;
|
||||
if let Some(u) = handle.get_usage() {
|
||||
info!("📊 Token: {} 总 / {} 提示 / {} 生成 | 缓存命中 {} | 缓存未命中 {}",
|
||||
u.total_tokens, u.prompt_tokens, u.completion_tokens,
|
||||
u.prompt_cache_hit_tokens, u.prompt_cache_miss_tokens);
|
||||
store_usage(db, "", conv.model(), conv.provider_name(), u).await;
|
||||
}
|
||||
handle.consume().await
|
||||
}
|
||||
|
||||
async fn generate_reply_with_tools(
|
||||
conv: &Conversation,
|
||||
user_text: String,
|
||||
db: &Option<Arc<Database>>,
|
||||
) -> Result<String, String> {
|
||||
let (reply, _used_tools, usage) = conv.chat_with_tools(user_text).await?;
|
||||
if let Some(u) = &usage {
|
||||
info!("📊 Token: {} 总 / {} 提示 / {} 生成 | 缓存命中 {} | 缓存未命中 {}",
|
||||
u.total_tokens, u.prompt_tokens, u.completion_tokens,
|
||||
u.prompt_cache_hit_tokens, u.prompt_cache_miss_tokens);
|
||||
store_usage(db, "", conv.model(), conv.provider_name(), u).await;
|
||||
}
|
||||
Ok(reply)
|
||||
}
|
||||
|
||||
async fn store_usage(db: &Option<Arc<Database>>, user_id: &str, model: &str, provider: &str, u: &Usage) {
|
||||
if let Some(db) = db {
|
||||
let _ = db::models::insert_llm_usage(
|
||||
db.pool(), user_id, model, provider,
|
||||
u.prompt_tokens, u.completion_tokens,
|
||||
u.prompt_cache_hit_tokens, u.prompt_cache_miss_tokens,
|
||||
).await;
|
||||
}
|
||||
}
|
||||
|
||||
async fn cmd_whoami(
|
||||
database: &Option<Arc<Database>>,
|
||||
file_state: &state::StateManager,
|
||||
) {
|
||||
let auth: Option<state::AuthState> = if let Some(db) = database {
|
||||
db::models::load_auth(db.pool()).await
|
||||
} else {
|
||||
file_state.load_auth()
|
||||
};
|
||||
|
||||
match auth {
|
||||
Some(a) => {
|
||||
println!("账号: {}", a.account_id);
|
||||
println!("Token: {}...", &a.token[..std::cmp::min(16, a.token.len())]);
|
||||
println!("API: {}", a.base_url);
|
||||
if database.is_some() {
|
||||
println!("存储: PostgreSQL");
|
||||
} else {
|
||||
println!("存储: 文件 ({})", file_state.state_dir().display());
|
||||
}
|
||||
}
|
||||
None => println!("未登录,请先执行 login 命令"),
|
||||
}
|
||||
}
|
||||
|
||||
async fn cmd_usage(
|
||||
database: &Option<Arc<Database>>,
|
||||
since: Option<String>,
|
||||
until: Option<String>,
|
||||
model: Option<String>,
|
||||
) {
|
||||
let db = match database {
|
||||
Some(d) => d,
|
||||
None => { println!("数据库未配置,无法查询用量"); return; }
|
||||
};
|
||||
|
||||
let since_dt = since.and_then(|s| {
|
||||
chrono::DateTime::parse_from_rfc3339(&s).ok().map(|d| d.with_timezone(&chrono::Utc))
|
||||
});
|
||||
let until_dt = until.and_then(|s| {
|
||||
chrono::DateTime::parse_from_rfc3339(&s).ok().map(|d| d.with_timezone(&chrono::Utc))
|
||||
});
|
||||
let model_ref = model.as_deref();
|
||||
|
||||
match db::models::query_llm_usage_stats(db.pool(), since_dt, until_dt, model_ref).await {
|
||||
Ok(stats) => {
|
||||
let hit_rate = if stats.total_cache_hit + stats.total_cache_miss > 0 {
|
||||
stats.total_cache_hit as f64 / (stats.total_cache_hit + stats.total_cache_miss) as f64 * 100.0
|
||||
} else { 0.0 };
|
||||
|
||||
println!("📊 LLM Token 使用统计");
|
||||
println!("{:=<40}", "");
|
||||
println!("调用次数: {}", stats.total_calls);
|
||||
println!("Prompt Tokens: {} ({:.1}K)", stats.total_prompt_tokens, stats.total_prompt_tokens as f64 / 1000.0);
|
||||
println!("生成 Tokens: {} ({:.1}K)", stats.total_completion_tokens, stats.total_completion_tokens as f64 / 1000.0);
|
||||
println!("总 Tokens: {} ({:.1}K)", stats.total_tokens, stats.total_tokens as f64 / 1000.0);
|
||||
println!("缓存命中: {} ({:.1}%)", stats.total_cache_hit, hit_rate);
|
||||
println!("缓存未命中: {} ({:.1}%)", stats.total_cache_miss, 100.0 - hit_rate);
|
||||
}
|
||||
Err(e) => println!("查询失败: {}", e),
|
||||
}
|
||||
}
|
||||
|
||||
async fn cmd_send(
|
||||
to: &str,
|
||||
text: &str,
|
||||
context_token: Option<&str>,
|
||||
database: &Option<Arc<Database>>,
|
||||
file_state: &state::StateManager,
|
||||
) {
|
||||
let auth: Option<state::AuthState> = if let Some(db) = database {
|
||||
db::models::load_auth(db.pool()).await
|
||||
} else {
|
||||
file_state.load_auth()
|
||||
};
|
||||
|
||||
let auth = match auth {
|
||||
Some(a) => a,
|
||||
None => {
|
||||
error!("未登录,请先执行 login 命令");
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let base_url = auth.base_url.clone();
|
||||
let account_id = auth.account_id.clone();
|
||||
let client = WeChatClient::new(Some(auth.base_url));
|
||||
client
|
||||
.set_auth(&auth.token, &auth.account_id, &base_url)
|
||||
.await;
|
||||
|
||||
match client.send_text(to, text, context_token).await {
|
||||
Ok(msg_id) => {
|
||||
println!("✅ 消息已发送, msg_id={}", msg_id);
|
||||
|
||||
// 入库:发送的消息
|
||||
if let Some(db) = database {
|
||||
let _ = db::models::insert_chat_record(
|
||||
db.pool(),
|
||||
"outbound",
|
||||
to,
|
||||
&account_id,
|
||||
text,
|
||||
"manual",
|
||||
context_token,
|
||||
&msg_id,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
Err(e) => error!("发送失败: {}", e),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,167 @@
|
||||
use chrono::Utc;
|
||||
use sqlx::PgPool;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use tokio::time::interval;
|
||||
|
||||
/// 定时任务调度器 —— 轮询 PostgreSQL 中到期的 scheduled_tasks
|
||||
pub struct Scheduler {
|
||||
pool: Option<Arc<PgPool>>,
|
||||
}
|
||||
|
||||
impl Scheduler {
|
||||
pub fn new(pool: Option<Arc<PgPool>>) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
/// 启动调度循环(在 listen_loop 中作为 tokio::spawn 调用)
|
||||
/// 每 5 秒检查一次是否有到期任务,执行后通过 callback 发送结果
|
||||
pub async fn run(
|
||||
&self,
|
||||
mut on_task: impl FnMut(&str, &str, &str) -> Result<(), String>,
|
||||
) {
|
||||
let pool = match self.pool.as_ref() {
|
||||
Some(p) => p.clone(),
|
||||
None => {
|
||||
tracing::info!("调度器:数据库未配置,跳过");
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let mut tick = interval(Duration::from_secs(5));
|
||||
tracing::info!("调度器已启动(每 5s 检查)");
|
||||
|
||||
loop {
|
||||
tick.tick().await;
|
||||
|
||||
let tasks = match fetch_due_tasks(&pool).await {
|
||||
Ok(t) => t,
|
||||
Err(e) => {
|
||||
tracing::warn!("调度器查询失败: {}", e);
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
for task in tasks {
|
||||
tracing::info!("⏰ 执行定时任务: {}", task.name);
|
||||
|
||||
// 执行任务
|
||||
let result = execute_task(&task).await;
|
||||
|
||||
// 更新任务状态
|
||||
let _ = update_task_status(&pool, &task.id, &result).await;
|
||||
|
||||
// 通知用户
|
||||
let notify_msg = format!(
|
||||
"⏰ 定时任务: {}\n结果: {}",
|
||||
task.name,
|
||||
result
|
||||
);
|
||||
if let Err(e) = on_task(&task.user_id, &task.name, ¬ify_msg) {
|
||||
tracing::error!("发送定时任务通知失败: {}", e);
|
||||
}
|
||||
|
||||
tracing::info!("✅ 定时任务完成: {} → {}", task.name, result);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct DueTask {
|
||||
id: uuid::Uuid,
|
||||
name: String,
|
||||
user_id: String,
|
||||
command: String,
|
||||
shell: String,
|
||||
cwd: String,
|
||||
}
|
||||
|
||||
async fn fetch_due_tasks(pool: &PgPool) -> Result<Vec<DueTask>, String> {
|
||||
let rows = sqlx::query_as::<_, (uuid::Uuid, String, String, String, String, String)>(
|
||||
r#"
|
||||
SELECT id, name, user_id, command, shell, cwd
|
||||
FROM scheduled_tasks
|
||||
WHERE enabled = true AND next_run_at <= NOW()
|
||||
ORDER BY next_run_at ASC
|
||||
LIMIT 10
|
||||
"#,
|
||||
)
|
||||
.fetch_all(pool)
|
||||
.await
|
||||
.map_err(|e| format!("查询到期任务失败: {}", e))?;
|
||||
|
||||
Ok(rows
|
||||
.into_iter()
|
||||
.map(|(id, name, user_id, command, shell, cwd)| DueTask {
|
||||
id,
|
||||
name,
|
||||
user_id,
|
||||
command,
|
||||
shell,
|
||||
cwd,
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn execute_task(task: &DueTask) -> String {
|
||||
let shell = if task.shell.is_empty() {
|
||||
"/bin/bash"
|
||||
} else {
|
||||
&task.shell
|
||||
};
|
||||
let cwd = if task.cwd.is_empty() {
|
||||
std::env::current_dir()
|
||||
.map(|p| p.to_string_lossy().to_string())
|
||||
.unwrap_or_default()
|
||||
} else {
|
||||
task.cwd.clone()
|
||||
};
|
||||
|
||||
let result = tokio::process::Command::new(shell)
|
||||
.arg("-c")
|
||||
.arg(&task.command)
|
||||
.current_dir(&cwd)
|
||||
.stdout(std::process::Stdio::piped())
|
||||
.stderr(std::process::Stdio::piped())
|
||||
.output()
|
||||
.await;
|
||||
|
||||
match result {
|
||||
Ok(output) => {
|
||||
let stdout = String::from_utf8_lossy(&output.stdout).trim().to_string();
|
||||
let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string();
|
||||
|
||||
if output.status.success() {
|
||||
if stdout.is_empty() { "执行成功(无输出)".to_string() } else { stdout }
|
||||
} else {
|
||||
format!(
|
||||
"退出码 {}: {}",
|
||||
output.status.code().unwrap_or(-1),
|
||||
if stderr.is_empty() { &stdout } else { &stderr }
|
||||
)
|
||||
}
|
||||
}
|
||||
Err(e) => format!("执行失败: {}", e),
|
||||
}
|
||||
}
|
||||
|
||||
async fn update_task_status(pool: &PgPool, task_id: &uuid::Uuid, result: &str) -> Result<(), String> {
|
||||
// 更新 last_run_at 和 next_run_at
|
||||
sqlx::query(
|
||||
r#"
|
||||
UPDATE scheduled_tasks
|
||||
SET last_run_at = NOW(),
|
||||
next_run_at = NOW() + (interval_seconds * INTERVAL '1 second'),
|
||||
last_status = $2
|
||||
WHERE id = $1
|
||||
"#,
|
||||
)
|
||||
.bind(task_id)
|
||||
.bind(result.chars().take(500).collect::<String>())
|
||||
.execute(pool)
|
||||
.await
|
||||
.map_err(|e| format!("更新任务状态失败: {}", e))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,96 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::path::PathBuf;
|
||||
|
||||
/// 认证状态
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct AuthState {
|
||||
pub token: String,
|
||||
pub account_id: String,
|
||||
pub base_url: String,
|
||||
}
|
||||
|
||||
/// 运行时状态
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct RuntimeState {
|
||||
pub get_updates_buf: String,
|
||||
}
|
||||
|
||||
/// 状态管理器
|
||||
pub struct StateManager {
|
||||
state_dir: PathBuf,
|
||||
}
|
||||
|
||||
impl StateManager {
|
||||
pub fn new(state_dir: impl Into<PathBuf>) -> Self {
|
||||
Self {
|
||||
state_dir: state_dir.into(),
|
||||
}
|
||||
}
|
||||
|
||||
fn auth_file(&self) -> PathBuf {
|
||||
self.state_dir.join("auth.json")
|
||||
}
|
||||
|
||||
fn runtime_file(&self) -> PathBuf {
|
||||
self.state_dir.join("runtime.json")
|
||||
}
|
||||
|
||||
fn ensure_dir(&self) {
|
||||
std::fs::create_dir_all(&self.state_dir).ok();
|
||||
}
|
||||
|
||||
pub fn state_dir(&self) -> &std::path::Path {
|
||||
&self.state_dir
|
||||
}
|
||||
|
||||
pub fn save_auth(&self, auth: &AuthState) {
|
||||
self.ensure_dir();
|
||||
let json = serde_json::to_string_pretty(auth).expect("序列化 auth 失败");
|
||||
if let Err(e) = std::fs::write(self.auth_file(), &json) {
|
||||
tracing::error!("保存 auth 状态失败: {}", e);
|
||||
}
|
||||
// 权限
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
std::fs::set_permissions(self.auth_file(), std::fs::Permissions::from_mode(0o600)).ok();
|
||||
}
|
||||
}
|
||||
|
||||
pub fn load_auth(&self) -> Option<AuthState> {
|
||||
let path = self.auth_file();
|
||||
if !path.exists() {
|
||||
return None;
|
||||
}
|
||||
let json = std::fs::read_to_string(&path).ok()?;
|
||||
serde_json::from_str(&json).ok()
|
||||
}
|
||||
|
||||
pub fn save_runtime(&self, buf: &str) {
|
||||
self.ensure_dir();
|
||||
let state = RuntimeState {
|
||||
get_updates_buf: buf.to_string(),
|
||||
};
|
||||
let json = serde_json::to_string_pretty(&state).expect("序列化 runtime 失败");
|
||||
if let Err(e) = std::fs::write(self.runtime_file(), &json) {
|
||||
tracing::error!("保存 runtime 状态失败: {}", e);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn load_runtime(&self) -> Option<RuntimeState> {
|
||||
let path = self.runtime_file();
|
||||
if !path.exists() {
|
||||
return None;
|
||||
}
|
||||
let json = std::fs::read_to_string(&path).ok()?;
|
||||
serde_json::from_str(&json).ok()
|
||||
}
|
||||
|
||||
/// 删除 auth 文件(迁移到数据库后清理)
|
||||
pub fn clear_auth(&self) {
|
||||
let path = self.auth_file();
|
||||
if path.exists() {
|
||||
std::fs::remove_file(&path).ok();
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,328 @@
|
||||
# iAs 工具系统设计
|
||||
|
||||
## 核心理念
|
||||
|
||||
**一切皆 Skill。**
|
||||
|
||||
没有"内置工具"和"外部技能"的区别——所有工具都是 SKILL.md,区别仅在来源:
|
||||
- **系统 skills** — `resources/skills/` 目录,随 iAs 分发,只读
|
||||
- **用户 skills** — `skills/` 项目目录,用户自己添加
|
||||
|
||||
加载时自动合并,同名时用户 skills 覆盖系统 skills。
|
||||
|
||||
---
|
||||
|
||||
## 架构
|
||||
|
||||
```
|
||||
┌──────────────────────────────────────────────────────────┐
|
||||
│ LLM │
|
||||
│ generate_reply() ←→ tool_choice → tool_execute(result) │
|
||||
└─────────────────────┬────────────────────────────────────┘
|
||||
│
|
||||
▼
|
||||
┌──────────────────────────────────────────────────────────┐
|
||||
│ SkillRegistry │
|
||||
│ │
|
||||
│ ┌──────────────────┐ ┌────────────────────────────┐ │
|
||||
│ │ 系统 Skills │ │ 用户 Skills │ │
|
||||
│ │ resources/skills/ │ │ skills/ │ │
|
||||
│ │ │ │ │ │
|
||||
│ │ • datetime │ │ • query_qweather/ │ │
|
||||
│ │ • weather │ │ • search_tavily/ │ │
|
||||
│ │ • search_web │ │ • start_pc/ │ │
|
||||
│ │ • email │ │ • ... │ │
|
||||
│ │ • shell_exec │ │ │ │
|
||||
│ │ • schedule │ │ │ │
|
||||
│ └────────┬──────────┘ └────────────┬──────────────┘ │
|
||||
│ │ │ │
|
||||
│ ▼ ▼ │
|
||||
│ ┌──────────────────────────────────────────────────┐ │
|
||||
│ │ Approval Gate │ │
|
||||
│ │ risk == High → 确认码审批 / Low → 自动执行 │ │
|
||||
│ │ 审批过程对 LLM 完全透明 │ │
|
||||
│ └──────────────────────────────────────────────────┘ │
|
||||
└──────────────────────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 核心类型
|
||||
|
||||
```rust
|
||||
/// 风险等级(仅两级)
|
||||
#[derive(Debug, Clone, PartialEq, Deserialize)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum RiskLevel {
|
||||
Low, // 只读/安全,自动执行
|
||||
High, // 写操作/有风险,需微信确认码审批
|
||||
}
|
||||
|
||||
/// 技能来源
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum SkillSource {
|
||||
System, // resources/skills/
|
||||
User, // skills/
|
||||
}
|
||||
|
||||
/// 技能元数据(从 SKILL.md 解析)
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
pub struct SkillSpec {
|
||||
pub name: String,
|
||||
pub description: String,
|
||||
pub risk_level: RiskLevel,
|
||||
pub parameters: serde_json::Value, // JSON Schema
|
||||
pub timeout_secs: u64,
|
||||
}
|
||||
|
||||
/// 技能(加载到内存后)
|
||||
pub struct Skill {
|
||||
pub spec: SkillSpec,
|
||||
pub source: SkillSource,
|
||||
pub dir: PathBuf,
|
||||
pub execute_command: String,
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## SKILL.md 格式
|
||||
|
||||
每个技能是一个目录,包含 `SKILL.md` 文件:
|
||||
|
||||
```markdown
|
||||
# Get Current Datetime
|
||||
|
||||
获取当前日期时间(北京时间 UTC+8)。
|
||||
|
||||
## Risk Level
|
||||
Low
|
||||
|
||||
## Parameters
|
||||
\`\`\`json
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"format": {
|
||||
"type": "string",
|
||||
"description": "时间格式: full | date | time",
|
||||
"enum": ["full", "date", "time"],
|
||||
"default": "full"
|
||||
}
|
||||
}
|
||||
}
|
||||
\`\`\`
|
||||
|
||||
## Execute
|
||||
\`\`\`bash
|
||||
#!/usr/bin/env bash
|
||||
scripts/datetime.sh
|
||||
\`\`\`
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 内置 Skills 一览
|
||||
|
||||
随 iAs 分发的系统 skills:
|
||||
|
||||
| Skill | 风险 | 说明 |
|
||||
|-------|------|------|
|
||||
| `get_current_datetime` | Low | 获取当前时间(UTC+8) |
|
||||
| `query_weather` | Low | 和风天气查询 |
|
||||
| `search_web` | Low | Tavily 网络搜索 |
|
||||
| `email` | High | 收发邮件 |
|
||||
| `execute_shell` | High | 执行 shell 命令 |
|
||||
| `list_scheduled_tasks` | Low | 列出定时任务 |
|
||||
| `manage_scheduled_task` | High | 增删改定时任务 |
|
||||
|
||||
---
|
||||
|
||||
## SkillRegistry
|
||||
|
||||
```rust
|
||||
pub struct SkillRegistry {
|
||||
skills: HashMap<String, Arc<Skill>>,
|
||||
}
|
||||
|
||||
impl SkillRegistry {
|
||||
pub async fn load(system_dir: &str, user_dir: &str) -> Result<Self, LoadError>;
|
||||
|
||||
pub fn list_specs(&self) -> Vec<&SkillSpec>;
|
||||
pub fn list_specs_by_risk(&self, level: RiskLevel) -> Vec<&SkillSpec>;
|
||||
pub fn get(&self, name: &str) -> Option<&Skill>;
|
||||
|
||||
/// 执行技能(含审批流程)
|
||||
pub async fn execute(
|
||||
&self,
|
||||
name: &str,
|
||||
params: serde_json::Value,
|
||||
ctx: &ExecutionContext,
|
||||
) -> Result<SkillResult, SkillError>;
|
||||
}
|
||||
|
||||
pub struct ExecutionContext {
|
||||
pub user_id: String,
|
||||
pub db: Option<Arc<Database>>,
|
||||
/// 发送微信消息的回调
|
||||
pub send_wechat: Option<Box<dyn Fn(&str, &str) -> Result<(), String> + Send + Sync>>,
|
||||
/// 等待用户回复的回调(返回用户消息文本)
|
||||
pub wait_for_reply: Option<Box<dyn Fn(&str, Duration) -> Result<Option<String>, String> + Send + Sync>>,
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 审批流程(透明模式)
|
||||
|
||||
**核心原则:审批过程对 LLM 完全不可见。** LLM 只看到最终结果。
|
||||
|
||||
```
|
||||
LLM 调用 High 风险技能
|
||||
│
|
||||
▼
|
||||
SkillRegistry::execute()
|
||||
│
|
||||
├── 查找 Skill
|
||||
├── risk_level == High?
|
||||
│
|
||||
└── YES ──────────────────────────────────────────┐
|
||||
│ │
|
||||
▼ │
|
||||
┌──────────────────────────────────────────┐ │
|
||||
│ 等待阶段(LLM 不感知) │ │
|
||||
│ │ │
|
||||
│ ① 生成 6 位确认码(如 482731) │ │
|
||||
│ ② SHA-256 哈希存入 pending_approvals │ │
|
||||
│ ③ 通过微信发送确认消息给用户: │ │
|
||||
│ │ │
|
||||
│ ┌─────────────────────────────┐ │ │
|
||||
│ │ ⚠️ 操作确认 │ │ │
|
||||
│ │ │ │ │
|
||||
│ │ 技能:send_email │ │ │
|
||||
│ │ 参数:收件人 xxx@xxx.com │ │ │
|
||||
│ │ │ │ │
|
||||
│ │ 确认码:482731 │ │ │
|
||||
│ │ 回复确认码以继续 │ │ │
|
||||
│ │ 回复 0 或「取消」以取消 │ │ │
|
||||
│ │ (5分钟内有效,最多3次尝试) │ │ │
|
||||
│ └─────────────────────────────┘ │ │
|
||||
│ │ │
|
||||
│ ④ 等待用户回复 │ │
|
||||
│ │ │ │
|
||||
│ ├── 回复正确确认码 ──┐ │ │
|
||||
│ │ │ │ │
|
||||
│ │ ▼ │ │
|
||||
│ │ ┌──────────────┐ │ │
|
||||
│ │ │ ⑤ 执行技能 │ │ │
|
||||
│ │ │ ⑥ LLM 看到 │ │ │
|
||||
│ │ │ 执行结果 │ │ │
|
||||
│ │ └──────────────┘ │ │
|
||||
│ │ │ │
|
||||
│ ├── 回复错误确认码 │ │
|
||||
│ │ → "确认码错误,还剩 N 次" │ │
|
||||
│ │ → 重试(最多 3 次) │ │
|
||||
│ │ → 3 次都错 │ │
|
||||
│ │ → LLM 看到: │ │
|
||||
│ │ "用户拒绝了你的调用: │ │
|
||||
│ │ send_email" │ │
|
||||
│ │ │ │
|
||||
│ ├── 回复 0 / 「取消」 │ │
|
||||
│ │ → LLM 看到: │ │
|
||||
│ │ "用户拒绝了你的调用: │ │
|
||||
│ │ send_email" │ │
|
||||
│ │ │ │
|
||||
│ └── 5 分钟超时 │ │
|
||||
│ → LLM 看到: │ │
|
||||
│ "用户没有确认操作: │ │
|
||||
│ send_email" │ │
|
||||
└──────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
### LLM 视角(两种可能)
|
||||
|
||||
```
|
||||
✅ 通过:
|
||||
LLM: 调用 send_email({to: "xxx", subject: "..."})
|
||||
系统: { success: true, message: "邮件已发送" }
|
||||
|
||||
❌ 拒绝:
|
||||
LLM: 调用 send_email({to: "xxx", subject: "..."})
|
||||
系统: { success: false, error: "用户拒绝了你的调用:send_email" }
|
||||
```
|
||||
|
||||
LLM 不知道确认码的存在,不知道审批过程。它只知道自己调用了工具,然后得到了结果或拒绝。
|
||||
|
||||
---
|
||||
|
||||
## 数据库表
|
||||
|
||||
```sql
|
||||
-- 待审批的工具调用
|
||||
CREATE TABLE pending_approvals (
|
||||
id UUID PRIMARY KEY,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||
expires_at TIMESTAMPTZ NOT NULL, -- 5分钟后过期
|
||||
user_id TEXT NOT NULL,
|
||||
skill_name TEXT NOT NULL,
|
||||
params JSONB NOT NULL,
|
||||
code_hash TEXT NOT NULL, -- 确认码 SHA-256(不存明文)
|
||||
attempts_left INTEGER NOT NULL DEFAULT 3, -- 剩余尝试次数
|
||||
status TEXT NOT NULL DEFAULT 'pending', -- pending | approved | rejected | expired
|
||||
result JSONB,
|
||||
consumed_at TIMESTAMPTZ
|
||||
);
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Skills 目录结构
|
||||
|
||||
```
|
||||
resources/skills/ # 系统内置 Skills(随 iAs 分发)
|
||||
├── get_current_datetime/
|
||||
│ ├── SKILL.md
|
||||
│ └── scripts/datetime.sh
|
||||
├── query_weather/
|
||||
│ ├── SKILL.md
|
||||
│ └── scripts/weather.sh
|
||||
├── search_web/
|
||||
│ ├── SKILL.md
|
||||
│ └── scripts/search.sh
|
||||
├── email/
|
||||
│ ├── SKILL.md
|
||||
│ └── scripts/email.sh
|
||||
├── execute_shell/
|
||||
│ ├── SKILL.md
|
||||
│ └── scripts/shell.sh
|
||||
├── list_scheduled_tasks/
|
||||
│ ├── SKILL.md
|
||||
│ └── scripts/list.sh
|
||||
└── manage_scheduled_task/
|
||||
├── SKILL.md
|
||||
└── scripts/manage.sh
|
||||
|
||||
skills/ # 用户 Skills(用户自行添加)
|
||||
├── query_qweather/
|
||||
│ ├── SKILL.md
|
||||
│ └── scripts/query.sh
|
||||
├── search_tavily/
|
||||
│ ├── SKILL.md
|
||||
│ └── scripts/search.sh
|
||||
└── start_pc/
|
||||
├── SKILL.md
|
||||
└── scripts/wake.sh
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 实施步骤
|
||||
|
||||
| 步骤 | 内容 | 预估 |
|
||||
|------|------|------|
|
||||
| Step 1 | SkillSpec/Skill/SkillRegistry 核心类型 + 注册表 | 半天 |
|
||||
| Step 2 | SKILL.md 解析器 + 目录扫描(system + user 合并) | 半天 |
|
||||
| Step 3 | 技能执行器:子进程管理 + 参数 JSON 传递 + 结果解析 | 半天 |
|
||||
| Step 4 | 系统 Skills 脚本实现(7 个 bash 脚本) | 1 天 |
|
||||
| Step 5 | 审批系统:pending_approvals 表 + 确认码生成/验证 + 微信透明通知 | 1 天 |
|
||||
| Step 6 | LLM function calling 集成 + 审计日志 | 1 天 |
|
||||
@@ -0,0 +1,154 @@
|
||||
use rand::Rng;
|
||||
use sha2::{Digest, Sha256};
|
||||
use sqlx::PgPool;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::{Mutex, oneshot};
|
||||
|
||||
/// 审批决策
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum ApprovalDecision {
|
||||
Approved,
|
||||
Rejected,
|
||||
Expired,
|
||||
}
|
||||
|
||||
/// 内存中待审批的条目
|
||||
struct PendingEntry {
|
||||
code_hash: String,
|
||||
skill_name: String,
|
||||
attempts_left: i32,
|
||||
/// 通知等待方
|
||||
sender: oneshot::Sender<ApprovalDecision>,
|
||||
}
|
||||
|
||||
/// 审批管理器
|
||||
pub struct ApprovalManager {
|
||||
pool: Option<Arc<PgPool>>,
|
||||
/// user_id → 该用户的待审批队列
|
||||
pending: Arc<Mutex<HashMap<String, Vec<PendingEntry>>>>,
|
||||
}
|
||||
|
||||
impl ApprovalManager {
|
||||
pub fn new(pool: Option<Arc<PgPool>>) -> Self {
|
||||
Self {
|
||||
pool,
|
||||
pending: Arc::new(Mutex::new(HashMap::new())),
|
||||
}
|
||||
}
|
||||
|
||||
/// 创建审批请求
|
||||
/// 返回 (确认码明文, 接收器) — 调用方 await 接收器等待用户确认
|
||||
pub async fn create(
|
||||
&self,
|
||||
user_id: &str,
|
||||
skill_name: &str,
|
||||
) -> Result<(String, oneshot::Receiver<ApprovalDecision>), String> {
|
||||
let code: u32 = rand::thread_rng().gen_range(100000..999999);
|
||||
let code_str = code.to_string();
|
||||
let code_hash = format!("{:x}", Sha256::digest(code_str.as_bytes()));
|
||||
|
||||
let (tx, rx) = oneshot::channel();
|
||||
|
||||
// 入库
|
||||
if let Some(ref pool) = self.pool {
|
||||
let id = uuid::Uuid::new_v4();
|
||||
let expires_at = chrono::Utc::now() + chrono::Duration::minutes(5);
|
||||
let params = serde_json::json!({"name": skill_name});
|
||||
|
||||
let _ = sqlx::query(
|
||||
r#"
|
||||
INSERT INTO pending_approvals (id, expires_at, user_id, skill_name, params, code_hash, status)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, 'pending')
|
||||
"#,
|
||||
)
|
||||
.bind(id)
|
||||
.bind(expires_at)
|
||||
.bind(user_id)
|
||||
.bind(skill_name)
|
||||
.bind(¶ms)
|
||||
.bind(&code_hash)
|
||||
.execute(pool.as_ref())
|
||||
.await;
|
||||
}
|
||||
|
||||
// 注册到内存
|
||||
let mut map = self.pending.lock().await;
|
||||
map.entry(user_id.to_string())
|
||||
.or_default()
|
||||
.push(PendingEntry {
|
||||
code_hash,
|
||||
skill_name: skill_name.to_string(),
|
||||
attempts_left: 3,
|
||||
sender: tx,
|
||||
});
|
||||
|
||||
Ok((code_str, rx))
|
||||
}
|
||||
|
||||
/// 处理用户回复(由消息循环调用)
|
||||
/// 匹配成功时返回 Some(技能名),匹配失败返回 None
|
||||
pub async fn handle_reply(&self, user_id: &str, reply: &str) -> Option<String> {
|
||||
let mut map = self.pending.lock().await;
|
||||
let entries = map.get_mut(user_id)?;
|
||||
if entries.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let reply_hash = format!("{:x}", Sha256::digest(reply.as_bytes()));
|
||||
let reply_trimmed = reply.trim();
|
||||
|
||||
// 检查取消
|
||||
if reply_trimmed == "0" || reply_trimmed == "取消" {
|
||||
let entry = entries.remove(0);
|
||||
let name = entry.skill_name.clone();
|
||||
let _ = entry.sender.send(ApprovalDecision::Rejected);
|
||||
return Some(name);
|
||||
}
|
||||
|
||||
// 检查确认码
|
||||
let pos = entries.iter().position(|e| e.code_hash == reply_hash);
|
||||
if let Some(idx) = pos {
|
||||
let entry = entries.remove(idx);
|
||||
let name = entry.skill_name.clone();
|
||||
let _ = entry.sender.send(ApprovalDecision::Approved);
|
||||
return Some(name);
|
||||
}
|
||||
|
||||
// 确认码错误
|
||||
if let Some(entry) = entries.first_mut() {
|
||||
entry.attempts_left -= 1;
|
||||
if entry.attempts_left <= 0 {
|
||||
let entry = entries.remove(0);
|
||||
let name = entry.skill_name.clone();
|
||||
let _ = entry.sender.send(ApprovalDecision::Rejected);
|
||||
return Some(name);
|
||||
}
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
/// 清理过期(由定时任务调用)
|
||||
pub async fn clean_expired(&self) {
|
||||
let mut map = self.pending.lock().await;
|
||||
for entries in map.values_mut() {
|
||||
entries.retain(|e| !e.sender.is_closed());
|
||||
}
|
||||
map.retain(|_, v| !v.is_empty());
|
||||
|
||||
if let Some(ref pool) = self.pool {
|
||||
let _ = sqlx::query(
|
||||
"UPDATE pending_approvals SET status = 'expired' WHERE status = 'pending' AND expires_at < NOW()",
|
||||
)
|
||||
.execute(pool.as_ref())
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
/// 获取某个用户的待审批条数
|
||||
pub async fn pending_count(&self, user_id: &str) -> usize {
|
||||
let map = self.pending.lock().await;
|
||||
map.get(user_id).map(|v| v.len()).unwrap_or(0)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,4 @@
|
||||
pub mod approval;
|
||||
pub mod parser;
|
||||
pub mod registry;
|
||||
pub mod skill;
|
||||
@@ -0,0 +1,279 @@
|
||||
use crate::tools::skill::{Skill, SkillSource, SkillSpec};
|
||||
use std::fs;
|
||||
use std::path::Path;
|
||||
|
||||
/// 解析 SKILL.md 文件,返回 Skill
|
||||
///
|
||||
/// SKILL.md 格式:
|
||||
/// ```markdown
|
||||
/// # Skill Name
|
||||
/// Description text...
|
||||
///
|
||||
/// ## Risk Level
|
||||
/// Low | High
|
||||
///
|
||||
/// ## Parameters
|
||||
/// ```json
|
||||
/// { ... }
|
||||
/// ```
|
||||
///
|
||||
/// ## Execute
|
||||
/// ```bash
|
||||
/// command to run
|
||||
/// ```
|
||||
/// ```
|
||||
pub fn parse_skill(skill_dir: &Path) -> Result<Skill, String> {
|
||||
let md_path = skill_dir.join("SKILL.md");
|
||||
let content =
|
||||
fs::read_to_string(&md_path).map_err(|e| format!("读取 {} 失败: {}", md_path.display(), e))?;
|
||||
|
||||
let lines: Vec<&str> = content.lines().collect();
|
||||
let mut name = String::new();
|
||||
let mut description = String::new();
|
||||
let mut risk_level = String::new();
|
||||
let mut params_raw = String::new();
|
||||
let mut execute_raw = String::new();
|
||||
let mut in_params = false;
|
||||
let mut in_execute = false;
|
||||
let mut params_fence = 0;
|
||||
let mut execute_fence = 0;
|
||||
let mut desc_lines: Vec<&str> = Vec::new();
|
||||
let mut in_desc = false;
|
||||
|
||||
for (i, line) in lines.iter().enumerate() {
|
||||
let trimmed = line.trim();
|
||||
|
||||
// 标题行 → 技能名
|
||||
if trimmed.starts_with("# ") && !trimmed.starts_with("## ") && name.is_empty() {
|
||||
name = trimmed.trim_start_matches("# ").trim().to_string();
|
||||
continue;
|
||||
}
|
||||
|
||||
// 二级标题
|
||||
if trimmed.starts_with("## ") {
|
||||
let section = trimmed.trim_start_matches("## ").trim().to_lowercase();
|
||||
match section.as_str() {
|
||||
"risk level" | "risk_level" | "risk" => {
|
||||
in_desc = false;
|
||||
in_params = false;
|
||||
in_execute = false;
|
||||
}
|
||||
"parameters" | "param" => {
|
||||
in_desc = false;
|
||||
in_params = true;
|
||||
in_execute = false;
|
||||
}
|
||||
"execute" | "execution" | "command" | "run" => {
|
||||
in_desc = false;
|
||||
in_params = false;
|
||||
in_execute = true;
|
||||
params_fence = 0;
|
||||
execute_fence = 0;
|
||||
}
|
||||
_ => {
|
||||
in_desc = false;
|
||||
in_params = false;
|
||||
in_execute = false;
|
||||
}
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
// 描述:标题和 Risk Level 之间的文本
|
||||
if !in_params && !in_execute && !trimmed.starts_with("#") && !trimmed.is_empty() {
|
||||
if name.is_empty() {
|
||||
continue;
|
||||
}
|
||||
// 检查是否已进入 Risk Level 段
|
||||
let is_risk_line = trimmed.to_lowercase().starts_with("low")
|
||||
|| trimmed.to_lowercase().starts_with("high");
|
||||
if risk_level.is_empty() && is_risk_line {
|
||||
risk_level = trimmed.to_string();
|
||||
continue;
|
||||
}
|
||||
if risk_level.is_empty() {
|
||||
desc_lines.push(trimmed);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
// Risk Level 值
|
||||
if !in_params && !in_execute && risk_level.is_empty() {
|
||||
let lower = trimmed.to_lowercase();
|
||||
if lower == "low" || lower == "high" {
|
||||
risk_level = trimmed.to_string();
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
// Parameters: JSON 代码块
|
||||
if in_params {
|
||||
if trimmed.starts_with("```") {
|
||||
params_fence += 1;
|
||||
if params_fence == 2 {
|
||||
in_params = false;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
if params_fence == 1 {
|
||||
if !params_raw.is_empty() {
|
||||
params_raw.push('\n');
|
||||
}
|
||||
params_raw.push_str(trimmed);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
// Execute: shell 命令
|
||||
if in_execute {
|
||||
if trimmed.starts_with("```") {
|
||||
execute_fence += 1;
|
||||
if execute_fence == 2 {
|
||||
in_execute = false;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
if execute_fence == 1 {
|
||||
// 跳过 shebang 行
|
||||
if trimmed.starts_with("#!/") {
|
||||
continue;
|
||||
}
|
||||
if !execute_raw.is_empty() {
|
||||
execute_raw.push('\n');
|
||||
}
|
||||
execute_raw.push_str(trimmed);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
// 验证
|
||||
if name.is_empty() {
|
||||
return Err(format!("{}: 未找到技能名称 (# Title)", md_path.display()));
|
||||
}
|
||||
|
||||
description = desc_lines.join(" ").trim().to_string();
|
||||
|
||||
let risk = match risk_level.to_lowercase().trim() {
|
||||
"high" => crate::tools::skill::RiskLevel::High,
|
||||
_ => crate::tools::skill::RiskLevel::Low,
|
||||
};
|
||||
|
||||
// 解析参数 JSON
|
||||
let parameters: serde_json::Value = if params_raw.is_empty() {
|
||||
serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {}
|
||||
})
|
||||
} else {
|
||||
serde_json::from_str(¶ms_raw).unwrap_or_else(|_| {
|
||||
serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {}
|
||||
})
|
||||
})
|
||||
};
|
||||
|
||||
// 执行命令(去掉首尾空白)
|
||||
let execute_command = execute_raw.trim().to_string();
|
||||
if execute_command.is_empty() {
|
||||
return Err(format!("{}: 未找到 Execute 命令", md_path.display()));
|
||||
}
|
||||
|
||||
let spec = SkillSpec {
|
||||
name,
|
||||
description,
|
||||
risk_level: risk,
|
||||
parameters,
|
||||
timeout_secs: 30,
|
||||
};
|
||||
|
||||
Ok(Skill::new(
|
||||
spec,
|
||||
// 先设成 User,由调用方修正
|
||||
SkillSource::User,
|
||||
skill_dir.to_path_buf(),
|
||||
execute_command,
|
||||
))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::io::Write;
|
||||
|
||||
#[test]
|
||||
fn test_load_all_system_skills() {
|
||||
let rt = tokio::runtime::Runtime::new().unwrap();
|
||||
let result = rt.block_on(crate::tools::registry::SkillRegistry::load(
|
||||
"resources/skills",
|
||||
"/nonexistent",
|
||||
));
|
||||
match result {
|
||||
Ok(registry) => {
|
||||
assert_eq!(registry.count(), 8, "应该有 8 个系统技能");
|
||||
let names: Vec<&str> = registry.list_specs().iter().map(|s| s.name.as_str()).collect();
|
||||
println!("已加载技能: {:?}", names);
|
||||
assert!(names.contains(&"manage_memos"));
|
||||
assert!(names.contains(&"get_current_datetime"));
|
||||
assert!(names.contains(&"query_weather"));
|
||||
assert!(names.contains(&"search_web"));
|
||||
assert!(names.contains(&"email"));
|
||||
assert!(names.contains(&"execute_shell"));
|
||||
assert!(names.contains(&"list_scheduled_tasks"));
|
||||
assert!(names.contains(&"manage_scheduled_task"));
|
||||
}
|
||||
Err(errors) => {
|
||||
for e in &errors {
|
||||
eprintln!("加载失败: {}", e);
|
||||
}
|
||||
panic!("技能加载失败");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_basic_skill() {
|
||||
let dir = std::env::temp_dir().join("ias-test-skill");
|
||||
fs::create_dir_all(&dir).ok();
|
||||
let mut f = fs::File::create(dir.join("SKILL.md")).unwrap();
|
||||
write!(
|
||||
f,
|
||||
r#"# Get Current Datetime
|
||||
|
||||
获取当前日期时间(北京时间 UTC+8)。
|
||||
|
||||
## Risk Level
|
||||
Low
|
||||
|
||||
## Parameters
|
||||
```json
|
||||
{{
|
||||
"type": "object",
|
||||
"properties": {{
|
||||
"format": {{
|
||||
"type": "string",
|
||||
"description": "full | date | time"
|
||||
}}
|
||||
}}
|
||||
}}
|
||||
```
|
||||
|
||||
## Execute
|
||||
```bash
|
||||
date '+%Y-%m-%d %H:%M:%S'
|
||||
```
|
||||
"#
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let skill = parse_skill(&dir).unwrap();
|
||||
assert_eq!(skill.spec.name, "Get Current Datetime");
|
||||
assert!(skill.spec.description.contains("日期时间"));
|
||||
assert_eq!(skill.spec.risk_level, crate::tools::skill::RiskLevel::Low);
|
||||
assert!(!skill.execute_command.is_empty());
|
||||
assert!(skill.execute_command.contains("date"));
|
||||
|
||||
fs::remove_dir_all(&dir).ok();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,311 @@
|
||||
use super::approval::{ApprovalDecision, ApprovalManager};
|
||||
use crate::tools::parser::parse_skill;
|
||||
use crate::tools::skill::*;
|
||||
use std::future::Future;
|
||||
use std::pin::Pin;
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use tokio::process::Command;
|
||||
|
||||
/// 技能注册表
|
||||
pub struct SkillRegistry {
|
||||
skills: HashMap<String, Arc<Skill>>,
|
||||
}
|
||||
|
||||
impl SkillRegistry {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
skills: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 从 system_dir 和 user_dir 加载所有技能
|
||||
/// system_dir: resources/skills/(随 iAs 分发)
|
||||
/// user_dir: skills/(用户自行添加)
|
||||
/// 同名时 user 覆盖 system
|
||||
pub async fn load(system_dir: &str, user_dir: &str) -> Result<Self, Vec<String>> {
|
||||
let mut registry = Self::new();
|
||||
let mut errors = Vec::new();
|
||||
|
||||
// 加载系统 skills
|
||||
if Path::new(system_dir).exists() {
|
||||
registry.load_from_dir(system_dir, SkillSource::System, &mut errors);
|
||||
}
|
||||
|
||||
// 加载用户 skills
|
||||
if Path::new(user_dir).exists() {
|
||||
registry.load_from_dir(user_dir, SkillSource::User, &mut errors);
|
||||
}
|
||||
|
||||
if errors.is_empty() {
|
||||
Ok(registry)
|
||||
} else {
|
||||
Err(errors)
|
||||
}
|
||||
}
|
||||
|
||||
fn load_from_dir(&mut self, dir: &str, source: SkillSource, errors: &mut Vec<String>) {
|
||||
let dir_path = Path::new(dir);
|
||||
let entries = match std::fs::read_dir(dir_path) {
|
||||
Ok(e) => e,
|
||||
Err(e) => {
|
||||
errors.push(format!("读取目录 {} 失败: {}", dir, e));
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
for entry in entries.flatten() {
|
||||
let path = entry.path();
|
||||
if !path.is_dir() {
|
||||
continue;
|
||||
}
|
||||
let skill_md = path.join("SKILL.md");
|
||||
if !skill_md.exists() {
|
||||
continue;
|
||||
}
|
||||
|
||||
match parse_skill(&path) {
|
||||
Ok(mut skill) => {
|
||||
skill.source = source.clone();
|
||||
let name = skill.spec.name.clone();
|
||||
|
||||
// user 同名覆盖 system
|
||||
if source == SkillSource::User && self.skills.contains_key(&name) {
|
||||
tracing::info!("用户技能覆盖系统技能: {}", name);
|
||||
}
|
||||
self.skills.insert(name, Arc::new(skill));
|
||||
}
|
||||
Err(e) => {
|
||||
errors.push(format!("{}: {}", path.display(), e));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 列出所有技能 specs(给 LLM function calling)
|
||||
pub fn list_specs(&self) -> Vec<&SkillSpec> {
|
||||
let mut specs: Vec<&SkillSpec> = self.skills.values().map(|s| &s.spec).collect();
|
||||
specs.sort_by(|a, b| a.name.cmp(&b.name));
|
||||
specs
|
||||
}
|
||||
|
||||
/// 按风险等级过滤
|
||||
pub fn list_specs_by_risk(&self, level: RiskLevel) -> Vec<&SkillSpec> {
|
||||
self.skills
|
||||
.values()
|
||||
.filter(|s| s.spec.risk_level == level)
|
||||
.map(|s| &s.spec)
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// 按名称查找技能
|
||||
pub fn get(&self, name: &str) -> Option<&Skill> {
|
||||
self.skills.get(name).map(|s| s.as_ref())
|
||||
}
|
||||
|
||||
/// 获取技能数量
|
||||
pub fn count(&self) -> usize {
|
||||
self.skills.len()
|
||||
}
|
||||
|
||||
// ─── 执行 ───
|
||||
|
||||
/// 执行技能(含审批流程)
|
||||
pub async fn execute(
|
||||
&self,
|
||||
name: &str,
|
||||
params: serde_json::Value,
|
||||
ctx: &ExecutionContext,
|
||||
) -> SkillResult {
|
||||
let skill = match self.skills.get(name) {
|
||||
Some(s) => s.clone(),
|
||||
None => {
|
||||
return SkillResult::error(format!("技能不存在: {}", name));
|
||||
}
|
||||
};
|
||||
|
||||
// High 风险 → 审批流程(透明模式)
|
||||
if skill.spec.risk_level.is_high() {
|
||||
match approval_flow(&skill, ctx).await {
|
||||
ApprovalResult::Approved => {
|
||||
// 继续执行
|
||||
}
|
||||
ApprovalResult::Rejected => {
|
||||
return SkillResult::rejected(name);
|
||||
}
|
||||
ApprovalResult::Expired => {
|
||||
return SkillResult::expired(name);
|
||||
}
|
||||
ApprovalResult::Error(msg) => {
|
||||
return SkillResult::error(msg);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 执行技能
|
||||
run_skill(&skill, params, ctx).await
|
||||
}
|
||||
}
|
||||
|
||||
// ─── 审批流程 ───
|
||||
|
||||
enum ApprovalResult {
|
||||
Approved,
|
||||
Rejected,
|
||||
Expired,
|
||||
Error(String),
|
||||
}
|
||||
|
||||
async fn approval_flow(skill: &Skill, ctx: &ExecutionContext) -> ApprovalResult {
|
||||
let manager = match ctx.approval_manager.as_ref() {
|
||||
Some(m) => m.clone(),
|
||||
None => return ApprovalResult::Approved, // 测试模式
|
||||
};
|
||||
|
||||
// 创建审批请求(生成确认码 + 注册通道)
|
||||
let (code, rx) = match manager.create(&ctx.user_id, &skill.spec.name).await {
|
||||
Ok(pair) => pair,
|
||||
Err(e) => return ApprovalResult::Error(e),
|
||||
};
|
||||
tracing::debug!("审批确认码已生成: {}...", &code[..2]);
|
||||
|
||||
// 发送确认消息给用户
|
||||
let msg = format!(
|
||||
"⚠️ 操作确认\n\n技能:{}\n\n确认码:{}\n\n回复确认码以继续\n回复 0 或「取消」以取消\n(5分钟内有效,最多3次尝试)",
|
||||
skill.spec.name, code
|
||||
);
|
||||
if let Some(ref send) = ctx.send_wechat {
|
||||
if let Err(e) = send(&ctx.user_id, &msg).await {
|
||||
return ApprovalResult::Error(format!("发送确认消息失败: {}", e));
|
||||
}
|
||||
}
|
||||
|
||||
// 等待用户确认(通过 oneshot channel,由消息循环驱动)
|
||||
match tokio::time::timeout(Duration::from_secs(300), rx).await {
|
||||
Ok(Ok(ApprovalDecision::Approved)) => ApprovalResult::Approved,
|
||||
Ok(Ok(ApprovalDecision::Rejected)) => ApprovalResult::Rejected,
|
||||
Ok(Ok(ApprovalDecision::Expired)) => ApprovalResult::Expired,
|
||||
Ok(Err(_)) => ApprovalResult::Error("审批通道异常".to_string()),
|
||||
Err(_) => ApprovalResult::Expired,
|
||||
}
|
||||
}
|
||||
|
||||
// ─── 技能执行器 ───
|
||||
|
||||
async fn run_skill(skill: &Skill, params: serde_json::Value, _ctx: &ExecutionContext) -> SkillResult {
|
||||
let timeout = Duration::from_secs(skill.spec.timeout_secs);
|
||||
|
||||
let cmd_str = &skill.execute_command;
|
||||
let parts: Vec<&str> = cmd_str.split_whitespace().collect();
|
||||
if parts.is_empty() {
|
||||
return SkillResult::error("执行命令为空");
|
||||
}
|
||||
|
||||
let mut cmd = Command::new(parts[0]);
|
||||
for arg in &parts[1..] {
|
||||
cmd.arg(arg);
|
||||
}
|
||||
cmd.current_dir(&skill.dir);
|
||||
tracing::debug!(
|
||||
"执行命令: {:?} (cwd={:?})",
|
||||
&cmd_str, &skill.dir
|
||||
);
|
||||
|
||||
// 通过 stdin 传入参数
|
||||
let params_json = serde_json::to_string(¶ms).unwrap_or_default();
|
||||
cmd.stdin(std::process::Stdio::piped());
|
||||
cmd.stdout(std::process::Stdio::piped());
|
||||
cmd.stderr(std::process::Stdio::piped());
|
||||
|
||||
let mut child = match cmd.spawn() {
|
||||
Ok(c) => c,
|
||||
Err(e) => return SkillResult::error(format!("启动进程失败: {}", e)),
|
||||
};
|
||||
|
||||
// 写入参数到 stdin
|
||||
if let Some(mut stdin) = child.stdin.take() {
|
||||
use tokio::io::AsyncWriteExt;
|
||||
let _ = stdin.write_all(params_json.as_bytes()).await;
|
||||
// 关闭 stdin 让进程知道输入结束
|
||||
drop(stdin);
|
||||
}
|
||||
|
||||
// 等待完成(带超时)
|
||||
let result = tokio::time::timeout(timeout, child.wait_with_output()).await;
|
||||
|
||||
match result {
|
||||
Ok(Ok(output)) => {
|
||||
let stdout = String::from_utf8_lossy(&output.stdout).trim().to_string();
|
||||
let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string();
|
||||
let exit_code = output.status.code().unwrap_or(-1);
|
||||
|
||||
tracing::debug!(
|
||||
"🖥️ {} exit={} out={:.200} err={:.100}",
|
||||
skill.spec.name, exit_code,
|
||||
if stdout.is_empty() { "(空)" } else { &stdout },
|
||||
if stderr.is_empty() { "(空)" } else { &stderr }
|
||||
);
|
||||
|
||||
if output.status.success() {
|
||||
// 尝试解析为 JSON
|
||||
if let Ok(json) = serde_json::from_str::<serde_json::Value>(&stdout) {
|
||||
SkillResult::ok_with_data(stdout, json)
|
||||
} else {
|
||||
SkillResult::ok(stdout)
|
||||
}
|
||||
} else {
|
||||
let msg = if !stderr.is_empty() {
|
||||
format!("退出码 {}: {}", output.status.code().unwrap_or(-1), stderr)
|
||||
} else {
|
||||
format!("退出码 {}", output.status.code().unwrap_or(-1))
|
||||
};
|
||||
SkillResult::error(msg)
|
||||
}
|
||||
}
|
||||
Ok(Err(e)) => SkillResult::error(format!("进程错误: {}", e)),
|
||||
Err(_) => SkillResult::error(format!("执行超时 ({}s)", skill.spec.timeout_secs)),
|
||||
}
|
||||
}
|
||||
|
||||
// ─── 执行上下文 ───
|
||||
|
||||
/// 发送微信消息的回调: (user_id, text) -> Result(异步)
|
||||
pub type WechatSender =
|
||||
Arc<dyn Fn(&str, &str) -> Pin<Box<dyn Future<Output = Result<(), String>> + Send>> + Send + Sync>;
|
||||
|
||||
pub struct ExecutionContext {
|
||||
pub user_id: String,
|
||||
pub db: Option<Arc<sqlx::PgPool>>,
|
||||
/// 审批管理器
|
||||
pub approval_manager: Option<Arc<ApprovalManager>>,
|
||||
/// 发送微信消息的回调
|
||||
pub send_wechat: Option<WechatSender>,
|
||||
}
|
||||
|
||||
impl ExecutionContext {
|
||||
pub fn new(user_id: impl Into<String>) -> Self {
|
||||
Self {
|
||||
user_id: user_id.into(),
|
||||
db: None,
|
||||
approval_manager: None,
|
||||
send_wechat: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_db(mut self, db: Arc<sqlx::PgPool>) -> Self {
|
||||
self.db = Some(db);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_approval(mut self, mgr: Arc<ApprovalManager>) -> Self {
|
||||
self.approval_manager = Some(mgr);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_wechat(mut self, send: WechatSender) -> Self {
|
||||
self.send_wechat = Some(send);
|
||||
self
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,156 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
// ─── 风险等级 ───
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Deserialize)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum RiskLevel {
|
||||
Low, // 只读/安全,自动执行
|
||||
High, // 写操作/有风险,需微信确认码审批
|
||||
}
|
||||
|
||||
impl RiskLevel {
|
||||
pub fn is_high(&self) -> bool {
|
||||
matches!(self, RiskLevel::High)
|
||||
}
|
||||
}
|
||||
|
||||
// ─── 技能来源 ───
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum SkillSource {
|
||||
System, // resources/skills/
|
||||
User, // skills/
|
||||
}
|
||||
|
||||
impl SkillSource {
|
||||
pub fn label(&self) -> &str {
|
||||
match self {
|
||||
SkillSource::System => "系统",
|
||||
SkillSource::User => "用户",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ─── 技能元数据(从 SKILL.md 解析) ───
|
||||
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
pub struct SkillSpec {
|
||||
pub name: String,
|
||||
pub description: String,
|
||||
#[serde(default = "default_risk_level")]
|
||||
pub risk_level: RiskLevel,
|
||||
#[serde(default)]
|
||||
pub parameters: serde_json::Value,
|
||||
#[serde(default = "default_timeout")]
|
||||
pub timeout_secs: u64,
|
||||
}
|
||||
|
||||
fn default_risk_level() -> RiskLevel {
|
||||
RiskLevel::Low
|
||||
}
|
||||
|
||||
fn default_timeout() -> u64 {
|
||||
30
|
||||
}
|
||||
|
||||
// ─── 技能(加载到内存后) ───
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Skill {
|
||||
pub spec: SkillSpec,
|
||||
pub source: SkillSource,
|
||||
pub dir: PathBuf,
|
||||
pub execute_command: String,
|
||||
}
|
||||
|
||||
impl Skill {
|
||||
pub fn new(spec: SkillSpec, source: SkillSource, dir: PathBuf, execute_command: String) -> Self {
|
||||
Self {
|
||||
spec,
|
||||
source,
|
||||
dir,
|
||||
execute_command,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ─── 执行上下文 ───
|
||||
|
||||
pub type WechatSender = Arc<dyn Fn(&str, &str) -> Result<(), String> + Send + Sync>;
|
||||
pub type ReplyWaiter = Arc<dyn Fn(&str, Duration) -> Result<Option<String>, String> + Send + Sync>;
|
||||
|
||||
/// LLM 能看到的结果
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct SkillResult {
|
||||
pub success: bool,
|
||||
pub output: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub data: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
impl SkillResult {
|
||||
pub fn ok(output: impl Into<String>) -> Self {
|
||||
Self {
|
||||
success: true,
|
||||
output: output.into(),
|
||||
data: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn ok_with_data(output: impl Into<String>, data: serde_json::Value) -> Self {
|
||||
Self {
|
||||
success: true,
|
||||
output: output.into(),
|
||||
data: Some(data),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn rejected(skill_name: &str) -> Self {
|
||||
Self {
|
||||
success: false,
|
||||
output: format!("用户拒绝了你的调用:{}", skill_name),
|
||||
data: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn expired(skill_name: &str) -> Self {
|
||||
Self {
|
||||
success: false,
|
||||
output: format!("用户没有确认操作:{}", skill_name),
|
||||
data: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn error(msg: impl Into<String>) -> Self {
|
||||
Self {
|
||||
success: false,
|
||||
output: msg.into(),
|
||||
data: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ─── 错误 ───
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum SkillError {
|
||||
NotFound(String),
|
||||
Execution(String),
|
||||
Timeout(String),
|
||||
InvalidParams(String),
|
||||
}
|
||||
|
||||
impl std::fmt::Display for SkillError {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
SkillError::NotFound(name) => write!(f, "技能不存在: {}", name),
|
||||
SkillError::Execution(msg) => write!(f, "执行失败: {}", msg),
|
||||
SkillError::Timeout(name) => write!(f, "执行超时: {}", name),
|
||||
SkillError::InvalidParams(msg) => write!(f, "参数无效: {}", msg),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,436 @@
|
||||
use crate::wechat::types::*;
|
||||
use base64::Engine;
|
||||
use reqwest::Client as HttpClient;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use tokio::sync::Mutex;
|
||||
use tracing::{debug, info, warn};
|
||||
|
||||
/// 微信 iLink Bot API 客户端
|
||||
#[derive(Clone)]
|
||||
pub struct WeChatClient {
|
||||
http: HttpClient,
|
||||
base_url: String,
|
||||
token: Arc<Mutex<String>>,
|
||||
account_id: Arc<Mutex<String>>,
|
||||
/// get_updates_buf,用于长轮询的上下文缓冲
|
||||
updates_buf: Arc<Mutex<String>>,
|
||||
}
|
||||
|
||||
impl WeChatClient {
|
||||
/// iLink App ID(从 package.json 的 ilink_appid 字段,硬编码或环境变量)
|
||||
const ILINK_APP_ID: &'static str = "";
|
||||
/// 默认 bot_type
|
||||
const DEFAULT_BOT_TYPE: &'static str = "3";
|
||||
/// 默认 API 地址
|
||||
const DEFAULT_BASE_URL: &'static str = "https://ilinkai.weixin.qq.com";
|
||||
/// 二维码固定 API 地址
|
||||
const FIXED_BASE_URL: &'static str = "https://ilinkai.weixin.qq.com";
|
||||
/// 版本号编码: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());
|
||||
|
||||
Self {
|
||||
http: HttpClient::builder()
|
||||
.timeout(Duration::from_secs(60))
|
||||
.build()
|
||||
.expect("创建 HTTP client 失败"),
|
||||
base_url: base,
|
||||
token: Arc::new(Mutex::new(String::new())),
|
||||
account_id: Arc::new(Mutex::new(String::new())),
|
||||
updates_buf: Arc::new(Mutex::new(String::new())),
|
||||
}
|
||||
}
|
||||
|
||||
// ─── 公共方法 ───
|
||||
|
||||
/// 获取登录二维码链接
|
||||
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::DEFAULT_BOT_TYPE
|
||||
);
|
||||
|
||||
let resp: QRCodeResponse = self
|
||||
.http
|
||||
.get(&url)
|
||||
.headers(Self::common_headers())
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("请求二维码失败: {}", e))?
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| format!("解析二维码响应失败: {}", e))?;
|
||||
|
||||
Ok((resp.qrcode, resp.qrcode_img_content))
|
||||
}
|
||||
|
||||
/// 轮询扫码状态
|
||||
pub async fn poll_qr_status(
|
||||
&self,
|
||||
qrcode: &str,
|
||||
base_url: &str,
|
||||
) -> Result<StatusResponse, String> {
|
||||
let url = format!(
|
||||
"{}/ilink/bot/get_qrcode_status?qrcode={}",
|
||||
base_url, qrcode
|
||||
);
|
||||
|
||||
let resp_result = self
|
||||
.http
|
||||
.get(&url)
|
||||
.headers(Self::common_headers())
|
||||
.timeout(Duration::from_secs(35))
|
||||
.send()
|
||||
.await;
|
||||
|
||||
match resp_result {
|
||||
Ok(response) => {
|
||||
let text = response.text().await
|
||||
.map_err(|e| format!("读取响应体失败: {}", e))?;
|
||||
serde_json::from_str::<StatusResponse>(&text)
|
||||
.map_err(|e| format!("解析状态响应失败: {}", e))
|
||||
}
|
||||
Err(e) if e.is_timeout() => {
|
||||
// 超时是正常的,返回 wait 让调用者继续轮询
|
||||
Ok(StatusResponse {
|
||||
status: "wait".to_string(),
|
||||
bot_token: None,
|
||||
ilink_bot_id: None,
|
||||
baseurl: None,
|
||||
ilink_user_id: None,
|
||||
redirect_host: None,
|
||||
})
|
||||
}
|
||||
Err(e) => {
|
||||
warn!("轮询二维码状态网络错误: {},继续等待", e);
|
||||
Ok(StatusResponse {
|
||||
status: "wait".to_string(),
|
||||
bot_token: None,
|
||||
ilink_bot_id: None,
|
||||
baseurl: None,
|
||||
ilink_user_id: None,
|
||||
redirect_host: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 执行完整扫码登录流程
|
||||
pub async fn login(
|
||||
&self,
|
||||
on_qrcode: impl Fn(&str),
|
||||
timeout_secs: u64,
|
||||
) -> Result<LoginResult, String> {
|
||||
info!("开始微信扫码登录...");
|
||||
|
||||
let (qrcode, qrcode_img_content) = self.get_qrcode().await?;
|
||||
|
||||
// 显示二维码链接
|
||||
on_qrcode(&qrcode_img_content);
|
||||
|
||||
// 生成二维码终端显示
|
||||
if let Ok(code) = qrcode::QrCode::new(qrcode_img_content.as_bytes()) {
|
||||
let svg = code.render::<qrcode::render::unicode::Dense1x2>()
|
||||
.dark_color(qrcode::render::unicode::Dense1x2::Dark)
|
||||
.light_color(qrcode::render::unicode::Dense1x2::Light)
|
||||
.build();
|
||||
println!("\n{}", svg);
|
||||
}
|
||||
|
||||
println!("请使用微信扫描上方二维码登录");
|
||||
println!("二维码链接(如二维码无法显示):{}", qrcode_img_content);
|
||||
println!();
|
||||
|
||||
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();
|
||||
|
||||
while tokio::time::Instant::now() < deadline {
|
||||
let status = self.poll_qr_status(&qrcode, ¤t_base).await?;
|
||||
|
||||
match status.status.as_str() {
|
||||
"wait" => {
|
||||
print!(".");
|
||||
std::io::Write::flush(&mut std::io::stdout()).ok();
|
||||
}
|
||||
"scaned" => {
|
||||
if !scanned {
|
||||
println!("\n👀 已扫码,请在微信上确认登录...");
|
||||
scanned = true;
|
||||
}
|
||||
}
|
||||
"scaned_but_redirect" => {
|
||||
if let Some(host) = &status.redirect_host {
|
||||
current_base = format!("https://{}", host);
|
||||
info!("IDC 重定向到: {}", current_base);
|
||||
}
|
||||
}
|
||||
"confirmed" => {
|
||||
let token = status
|
||||
.bot_token
|
||||
.ok_or_else(|| "登录成功但未返回 bot_token".to_string())?;
|
||||
let account_id = status
|
||||
.ilink_bot_id
|
||||
.ok_or_else(|| "登录成功但未返回 ilink_bot_id".to_string())?;
|
||||
let user_id = status.ilink_user_id.unwrap_or_default();
|
||||
let base_url = status.baseurl.unwrap_or_else(|| Self::DEFAULT_BASE_URL.to_string());
|
||||
|
||||
// 保存到内存状态
|
||||
*self.token.lock().await = token.clone();
|
||||
*self.account_id.lock().await = account_id.clone();
|
||||
|
||||
info!("✅ 登录成功! bot_id={}", account_id);
|
||||
return Ok(LoginResult {
|
||||
token,
|
||||
account_id,
|
||||
base_url,
|
||||
user_id,
|
||||
});
|
||||
}
|
||||
"expired" => {
|
||||
return Err("二维码已过期".to_string());
|
||||
}
|
||||
_ => {
|
||||
warn!("未知状态: {}", status.status);
|
||||
}
|
||||
}
|
||||
|
||||
tokio::time::sleep(Duration::from_secs(1)).await;
|
||||
}
|
||||
|
||||
Err("登录超时".to_string())
|
||||
}
|
||||
|
||||
/// 注册监听器(通知微信服务端本客户端已启动)
|
||||
pub async fn notify_start(&self) -> Result<(), String> {
|
||||
let body = serde_json::json!({ "base_info": BaseInfo::new() });
|
||||
|
||||
let url = format!("{}/ilink/bot/msg/notifystart", self.base_url);
|
||||
let token = self.token.lock().await.clone();
|
||||
|
||||
let _resp: NotifyResp = self
|
||||
.post_json(&url, &body, Some(&token))
|
||||
.await?
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| format!("解析 notifyStart 响应失败: {}", e))?;
|
||||
|
||||
info!("监听器已注册");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 注销监听器
|
||||
pub async fn notify_stop(&self) -> Result<(), String> {
|
||||
let body = serde_json::json!({ "base_info": BaseInfo::new() });
|
||||
|
||||
let url = format!("{}/ilink/bot/msg/notifystop", self.base_url);
|
||||
let token = self.token.lock().await.clone();
|
||||
|
||||
let _resp: NotifyResp = self
|
||||
.post_json(&url, &body, Some(&token))
|
||||
.await?
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| format!("解析 notifyStop 响应失败: {}", e))?;
|
||||
|
||||
info!("监听器已注销");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 长轮询接收消息
|
||||
pub async fn receive_messages(&self) -> Result<GetUpdatesResp, String> {
|
||||
let buf = self.updates_buf.lock().await.clone();
|
||||
|
||||
let body = GetUpdatesReq {
|
||||
get_updates_buf: buf.clone(),
|
||||
base_info: BaseInfo::new(),
|
||||
};
|
||||
|
||||
let url = format!("{}/ilink/bot/getupdates", self.base_url);
|
||||
let token = self.token.lock().await.clone();
|
||||
|
||||
let resp: GetUpdatesResp = match self
|
||||
.post_json(&url, &body, Some(&token))
|
||||
.await
|
||||
{
|
||||
Ok(resp) => resp.json().await.map_err(|e| format!("解析 getUpdates 响应失败: {}", e))?,
|
||||
Err(e) => {
|
||||
// 如果超时,返回空响应
|
||||
if e.contains("超时") || e.contains("timeout") {
|
||||
return Ok(GetUpdatesResp {
|
||||
ret: 0,
|
||||
errcode: None,
|
||||
errmsg: None,
|
||||
msgs: vec![],
|
||||
get_updates_buf: buf.clone(),
|
||||
longpolling_timeout_ms: None,
|
||||
});
|
||||
}
|
||||
return Err(e);
|
||||
}
|
||||
};
|
||||
|
||||
// 更新 buf
|
||||
if !resp.get_updates_buf.is_empty() {
|
||||
*self.updates_buf.lock().await = resp.get_updates_buf.clone();
|
||||
}
|
||||
|
||||
Ok(resp)
|
||||
}
|
||||
|
||||
/// 发送文本消息
|
||||
pub async fn send_text(
|
||||
&self,
|
||||
to: &str,
|
||||
text: &str,
|
||||
context_token: Option<&str>,
|
||||
) -> Result<String, String> {
|
||||
let client_id = uuid::Uuid::new_v4().to_string();
|
||||
|
||||
let msg = WeixinMessage {
|
||||
from_user_id: Some(String::new()),
|
||||
to_user_id: Some(to.to_string()),
|
||||
client_id: Some(client_id.clone()),
|
||||
msg_type: Some(WeixinMessage::TYPE_BOT),
|
||||
message_state: Some(2), // FINISH
|
||||
item_list: Some(vec![MessageItem::text(text)]),
|
||||
context_token: context_token.map(|s| s.to_string()),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let body = SendMessageReq { msg };
|
||||
|
||||
let url = format!("{}/ilink/bot/sendmessage", self.base_url);
|
||||
let token = self.token.lock().await.clone();
|
||||
|
||||
let _resp_text = self
|
||||
.post_json(&url, &body, Some(&token))
|
||||
.await?
|
||||
.text()
|
||||
.await
|
||||
.map_err(|e| format!("发送消息网络错误: {}", e))?;
|
||||
|
||||
Ok(client_id)
|
||||
}
|
||||
|
||||
/// 设置 auth 状态(从持久化恢复)
|
||||
pub async fn set_auth(&self, token: &str, account_id: &str, base_url: &str) {
|
||||
*self.token.lock().await = token.to_string();
|
||||
*self.account_id.lock().await = account_id.to_string();
|
||||
if !base_url.is_empty() {
|
||||
// 只更新 base_url 如果传入了非空值
|
||||
}
|
||||
}
|
||||
|
||||
/// 获取当前 token(用于持久化)
|
||||
pub async fn get_token(&self) -> String {
|
||||
self.token.lock().await.clone()
|
||||
}
|
||||
|
||||
/// 获取当前 account_id(用于持久化)
|
||||
pub async fn get_account_id(&self) -> String {
|
||||
self.account_id.lock().await.clone()
|
||||
}
|
||||
|
||||
/// 获取当前 updates_buf(用于持久化)
|
||||
pub async fn get_updates_buf(&self) -> String {
|
||||
self.updates_buf.lock().await.clone()
|
||||
}
|
||||
|
||||
/// 设置 get_updates_buf(从持久化恢复)
|
||||
pub async fn set_updates_buf(&self, buf: &str) {
|
||||
*self.updates_buf.lock().await = buf.to_string();
|
||||
}
|
||||
|
||||
// ─── 内部方法 ───
|
||||
|
||||
fn common_headers() -> reqwest::header::HeaderMap {
|
||||
let mut headers = reqwest::header::HeaderMap::new();
|
||||
headers.insert(
|
||||
"Content-Type",
|
||||
reqwest::header::HeaderValue::from_static("application/json"),
|
||||
);
|
||||
|
||||
if !Self::ILINK_APP_ID.is_empty() {
|
||||
headers.insert(
|
||||
"iLink-App-Id",
|
||||
reqwest::header::HeaderValue::from_static(Self::ILINK_APP_ID),
|
||||
);
|
||||
}
|
||||
|
||||
headers.insert(
|
||||
"iLink-App-ClientVersion",
|
||||
reqwest::header::HeaderValue::from_str(&Self::CLIENT_VERSION.to_string())
|
||||
.unwrap(),
|
||||
);
|
||||
|
||||
// X-WECHAT-UIN: random uint32 -> decimal -> base64
|
||||
let uin = rand::random::<u32>();
|
||||
let uin_b64 = base64::engine::general_purpose::STANDARD
|
||||
.encode(uin.to_string());
|
||||
headers.insert(
|
||||
"X-WECHAT-UIN",
|
||||
reqwest::header::HeaderValue::from_str(&uin_b64).unwrap(),
|
||||
);
|
||||
|
||||
headers
|
||||
}
|
||||
|
||||
fn auth_headers(token: &str) -> reqwest::header::HeaderMap {
|
||||
let mut headers = Self::common_headers();
|
||||
if !token.is_empty() {
|
||||
headers.insert(
|
||||
reqwest::header::AUTHORIZATION,
|
||||
reqwest::header::HeaderValue::from_str(&format!("Bearer {}", token))
|
||||
.unwrap(),
|
||||
);
|
||||
headers.insert(
|
||||
"AuthorizationType",
|
||||
reqwest::header::HeaderValue::from_static("ilink_bot_token"),
|
||||
);
|
||||
}
|
||||
headers
|
||||
}
|
||||
|
||||
async fn post_json(
|
||||
&self,
|
||||
url: &str,
|
||||
body: &impl serde::Serialize,
|
||||
token: Option<&str>,
|
||||
) -> Result<reqwest::Response, String> {
|
||||
let body_str = serde_json::to_string(body).map_err(|e| format!("序列化请求体失败: {}", e))?;
|
||||
|
||||
debug!("POST {} body_len={}", url, body_str.len());
|
||||
|
||||
let mut headers = if let Some(t) = token {
|
||||
Self::auth_headers(t)
|
||||
} else {
|
||||
Self::common_headers()
|
||||
};
|
||||
|
||||
// 设置 Content-Length
|
||||
headers.insert(
|
||||
"Content-Length",
|
||||
reqwest::header::HeaderValue::from_str(&body_str.len().to_string())
|
||||
.unwrap(),
|
||||
);
|
||||
|
||||
self.http
|
||||
.post(url)
|
||||
.headers(headers)
|
||||
.body(body_str)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| {
|
||||
if e.is_timeout() {
|
||||
format!("请求超时: {}", url)
|
||||
} else {
|
||||
format!("请求失败 ({}): {}", url, e)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,2 @@
|
||||
pub mod client;
|
||||
pub mod types;
|
||||
@@ -0,0 +1,333 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
// ─── 基础类型 ───
|
||||
|
||||
/// 公共请求元数据
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub struct BaseInfo {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub channel_version: Option<String>,
|
||||
}
|
||||
|
||||
impl BaseInfo {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
channel_version: Some(env!("CARGO_PKG_VERSION").to_string()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ─── 消息类型 ───
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct TextItem {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub text: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct CDNMedia {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub encrypt_query_param: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub aes_key: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub encrypt_type: Option<i32>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub full_url: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ImageItem {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub media: Option<CDNMedia>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub thumb_media: Option<CDNMedia>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub aeskey: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub url: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub mid_size: Option<i64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub thumb_size: Option<i64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub thumb_height: Option<i32>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub thumb_width: Option<i32>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub hd_size: Option<i64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct VoiceItem {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub media: Option<CDNMedia>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub encode_type: Option<i32>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub bits_per_sample: Option<i32>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub sample_rate: Option<i32>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub playtime: Option<i64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub text: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct FileItem {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub media: Option<CDNMedia>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub file_name: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub md5: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub len: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct VideoItem {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub media: Option<CDNMedia>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub video_size: Option<i64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub play_length: Option<i64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub video_md5: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub thumb_media: Option<CDNMedia>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub thumb_size: Option<i64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub thumb_height: Option<i32>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub thumb_width: Option<i32>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct RefMessage {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub message_item: Option<Box<MessageItem>>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub title: Option<String>,
|
||||
}
|
||||
|
||||
/// 消息条目(支持 text/image/voice/file/video)
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct MessageItem {
|
||||
#[serde(rename = "type", skip_serializing_if = "Option::is_none")]
|
||||
pub item_type: Option<i32>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub create_time_ms: Option<i64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub update_time_ms: Option<i64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub is_completed: Option<bool>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub msg_id: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub ref_msg: Option<Box<RefMessage>>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub text_item: Option<TextItem>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub image_item: Option<ImageItem>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub voice_item: Option<VoiceItem>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub file_item: Option<FileItem>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub video_item: Option<VideoItem>,
|
||||
}
|
||||
|
||||
impl MessageItem {
|
||||
pub const TYPE_TEXT: i32 = 1;
|
||||
pub const TYPE_IMAGE: i32 = 2;
|
||||
pub const TYPE_VOICE: i32 = 3;
|
||||
pub const TYPE_FILE: i32 = 4;
|
||||
pub const TYPE_VIDEO: i32 = 5;
|
||||
|
||||
pub fn text(text: &str) -> Self {
|
||||
Self {
|
||||
item_type: Some(Self::TYPE_TEXT),
|
||||
text_item: Some(TextItem {
|
||||
text: Some(text.to_string()),
|
||||
}),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ─── 微信消息 ───
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct WeixinMessage {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub seq: Option<i64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub message_id: Option<i64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub from_user_id: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub to_user_id: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub client_id: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub create_time_ms: Option<i64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub update_time_ms: Option<i64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub delete_time_ms: Option<i64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub session_id: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub group_id: Option<String>,
|
||||
#[serde(rename = "message_type", skip_serializing_if = "Option::is_none")]
|
||||
pub msg_type: Option<i32>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub message_state: Option<i32>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub item_list: Option<Vec<MessageItem>>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
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 {
|
||||
pub const TYPE_NONE: i32 = 0;
|
||||
pub const TYPE_USER: i32 = 1;
|
||||
pub const TYPE_BOT: i32 = 2;
|
||||
|
||||
/// 获取文本内容(如果有纯文本 item)
|
||||
pub fn text_content(&self) -> Option<&str> {
|
||||
self.item_list
|
||||
.as_ref()?
|
||||
.iter()
|
||||
.find(|item| item.item_type == Some(MessageItem::TYPE_TEXT))
|
||||
.and_then(|item| item.text_item.as_ref())
|
||||
.and_then(|t| t.text.as_deref())
|
||||
}
|
||||
}
|
||||
|
||||
// ─── API 请求/响应 ───
|
||||
|
||||
/// 获取二维码响应
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct QRCodeResponse {
|
||||
pub qrcode: String,
|
||||
pub qrcode_img_content: String,
|
||||
}
|
||||
|
||||
/// 扫码状态响应
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct StatusResponse {
|
||||
pub status: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub bot_token: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub ilink_bot_id: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub baseurl: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub ilink_user_id: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub redirect_host: Option<String>,
|
||||
}
|
||||
|
||||
/// GetUpdates 请求
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct GetUpdatesReq {
|
||||
pub get_updates_buf: String,
|
||||
pub base_info: BaseInfo,
|
||||
}
|
||||
|
||||
/// GetUpdates 响应
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct GetUpdatesResp {
|
||||
#[serde(default)]
|
||||
pub ret: i32,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub errcode: Option<i32>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub errmsg: Option<String>,
|
||||
#[serde(default)]
|
||||
pub msgs: Vec<WeixinMessage>,
|
||||
#[serde(default)]
|
||||
pub get_updates_buf: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub longpolling_timeout_ms: Option<i64>,
|
||||
}
|
||||
|
||||
/// SendMessage 请求
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct SendMessageReq {
|
||||
pub msg: WeixinMessage,
|
||||
}
|
||||
|
||||
/// NotifyStart/NotifyStop 响应
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct NotifyResp {
|
||||
#[serde(default)]
|
||||
pub ret: i32,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub errmsg: Option<String>,
|
||||
}
|
||||
|
||||
/// 登录结果
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct LoginResult {
|
||||
pub token: String,
|
||||
pub account_id: String,
|
||||
pub base_url: String,
|
||||
pub user_id: String,
|
||||
}
|
||||
|
||||
/// 收到的消息事件
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct MessageEvent {
|
||||
pub from_user_id: String,
|
||||
pub text: String,
|
||||
pub context_token: Option<String>,
|
||||
pub message_id: i64,
|
||||
}
|
||||
Reference in New Issue
Block a user