wunianxiao/iAs
1
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

清理废弃代码，重构长期记忆系统

Browse Source 
核心变更：
- 删除 worker.rs + ipc.rs（废弃进程架构）
- 删除 ias listen 路径（cmd_listen / listen_loop / ToolExecutor）
- 合并 ias service → ias daemon --log-file
- 修复 main.rs call_capability 不处理上下文工具的 bug

MemoryStore 重构：
- write_for 空内容防御 + 精确去重
- 缓存结构 Vec<String> → Vec<(i32, String)>（用数据库 id 标识）
- 新增 delete_for / update_for 方法 + LLM 工具定义
- load_if_needed 避免每次消息查 DB
- 缓存 TTL（1h） + 每用户 200 条上限
- read_for limit 参数（默认 30 条）

daemon.rs：
- execute_tool 统一为递归派发，消除 call_capability 重复路由

净减少约 1400 行代码，0 warning，20 测试通过
This commit is contained in:
无念晓
2026-06-16 11:56:36 +08:00
parent 5a54368437
commit 1d78567102
19 changed files with 1060 additions and 1895 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/cli.rs
+18 -44
View File
@@ -50,22 +50,20 @@ pub enum Commands {
timeout: u64,
},
/// 监听微信消息（长轮询）
///
/// 启动微信消息监听，接收并处理消息。
/// 守护进程模式（常驻后台，持有 WeChat 长连接和数据库）
///
/// 示例:
/// ias listen # 仅监听，不做回复
/// ias listen --llm # AI 自动回复
/// ias listen --echo # 原样回显
Listen {
/// 启用 AI 自动回复
#[arg(long)]
llm: bool,
/// ias daemon
/// ias daemon --sock /tmp/ias.sock
/// ias daemon --log-file # 启用文件日志（配合 systemd）
Daemon {
/// Unix Domain Socket 路径
#[arg(long, default_value = "/tmp/ias_daemon.sock")]
sock: String,
/// 原样回显收到的消息
/// 启用文件日志（默认输出到 stderr）
#[arg(long)]
echo: bool,
log_file: bool,
},
/// 发送文本消息
@@ -110,38 +108,6 @@ pub enum Commands {
model: Option<String>,
},
/// 后台服务模式（等同 listen --llm + 文件日志）
///
/// 配合 systemd 实现开机自启。
///
/// 示例:
/// ias service
Service,
/// 守护进程模式（daemon + worker 分离架构）
///
/// Daemon 持有 WeChat 长连接和数据库，每条消息 spawn 独立 worker。
/// Worker 代码更新后编译即可，下一条消息自动用新版本。
///
/// 示例:
/// ias daemon
/// ias daemon --sock /tmp/ias.sock
Daemon {
/// Unix Domain Socket 路径
#[arg(long, default_value = "/tmp/ias_daemon.sock")]
sock: String,
},
/// [已废弃] Worker 进程（新架构使用 daemon 统一 tokio task）
///
/// 示例:
/// ias worker --sock /tmp/ias_daemon.sock
Worker {
/// Unix Domain Socket 路径
#[arg(long, default_value = "/tmp/ias_daemon.sock")]
sock: String,
},
/// 管理定时任务（增删改查）
///
/// 需要 PostgreSQL 数据库。
@@ -155,6 +121,14 @@ pub enum Commands {
#[command(subcommand)]
ScheduledTask(ScheduledTaskAction),
/// 列出所有可用的 LLM 工具
///
/// 显示 LLM 可以调用的所有工具名称、描述、参数格式和风险等级。
///
/// 示例:
/// ias tools
Tools,
/// 调用内置工具（无需登录，独立运行）
///
/// 可直接在终端调用天气、搜索、备忘录、日期时间、高德地图等工具。
src/context/mod.rs
+1
View File
@@ -16,3 +16,4 @@ pub mod tools;
pub mod types;
pub use tools::MemoryStore;
pub use types::HistoryEntry;
src/context/tools.rs
+213 -48
View File
@@ -8,32 +8,54 @@
//! ### 数据流
//! ```text
//! 启动时 load(user_id) → 从数据库加载到缓存
//! 对话中 read_for() → 从缓存读取，返回 "1. ...\n2. ..." 格式
//! 对话中 write_for() → 写入缓存 + 写入数据库
//! 对话中 read_for() → 从缓存读取，返回 "[id] ..." 格式
//! 对话中 write_for() → 写入缓存 + 写入数据库，返回 id
//! 对话中 delete_for() / update_for() → 按 id 删除或更新
//! ```
use sqlx::PgPool;
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Instant;
use tokio::sync::Mutex;
/// 缓存条目：数据库 id + 记忆内容
type CacheEntry = (i32, String);
/// 单用户缓存状态
struct UserCache {
memories: Vec<CacheEntry>,
last_access: Instant,
}
/// 缓存 TTL：1 小时无活动则清理
const CACHE_TTL_SECS: u64 = 3600;
/// 每用户记忆上限（防止单用户撑爆）
const MAX_MEMORIES_PER_USER: usize = 200;
/// 默认返回给 LLM 的记忆条数
const DEFAULT_MEMORY_LIMIT: usize = 30;
/// ## 长期记忆管理器（内存 + PostgreSQL 双写）
///
/// 为每个用户维护独立的长期记忆列表，支持读写操作。
/// 为每个用户维护独立的长期记忆列表，支持读写删改操作。
///
/// ### 存储策略
/// * **内存缓存** — `HashMap<String, Vec<String>>`，按 user_id 隔离
/// * **内存缓存** — `HashMap<String, UserCache>`，按 user_id 隔离
/// * **数据库持久化** — 可选的 `PgPool`，写入时同时写到 PostgreSQL
///
/// ### 数据流
/// ```text
/// 启动时 load(user_id) → 从数据库加载到缓存
/// 对话中 read_for() → 从缓存读取，返回 "1. ...\n2. ..." 格式
/// 对话中 write_for() → 写入缓存 + 写入数据库
/// 对话中 read_for() → 从缓存读取，返回 "[id] ..." 格式
/// 对话中 write_for() → INSERT RETURNING id，缓存存 (id, content)
/// 对话中 delete_for() → 删缓存 + DELETE DB
/// 对话中 update_for() → 更新缓存 + UPDATE DB
/// ```
pub struct MemoryStore {
pool: Option<Arc<PgPool>>,
cache: Arc<Mutex<HashMap<String, Vec<String>>>>,
cache: Arc<Mutex<HashMap<String, UserCache>>>,
}
impl MemoryStore {
@@ -44,83 +66,226 @@ impl MemoryStore {
}
}
/// 更新用户访问时间
async fn touch(&self, uid: &str) {
let mut cache = self.cache.lock().await;
if let Some(uc) = cache.get_mut(uid) {
uc.last_access = Instant::now();
}
}
/// 清理过期缓存（超过 1h 无活动的用户）
async fn evict_expired(&self) {
let now = Instant::now();
let mut cache = self.cache.lock().await;
cache.retain(|_, uc| now.duration_since(uc.last_access).as_secs() < CACHE_TTL_SECS);
}
/// 条件加载：仅在缓存未命中时查 DB
pub async fn load_if_needed(&self, user_id: &str) {
// 先清理过期缓存
self.evict_expired().await;
let uid = if user_id.is_empty() { "default" } else { user_id };
{
let cache = self.cache.lock().await;
if cache.contains_key(uid) {
return; // 缓存已有，跳过 DB 查询
}
}
self.load(user_id).await;
}
pub async fn load(&self, user_id: &str) {
let pool = match self.pool.as_ref() {
Some(p) => p.clone(),
None => return,
};
let uid = if user_id.is_empty() { "default" } else { user_id };
if let Ok(rows) = sqlx::query_as::<_, (String,)>(
"SELECT content FROM user_memories WHERE user_id = $1 ORDER BY id",
if let Ok(rows) = sqlx::query_as::<_, (i32, String)>(
"SELECT id, content FROM user_memories WHERE user_id = $1 ORDER BY id",
)
.bind(uid)
.fetch_all(pool.as_ref())
.await
{
let mems: Vec<String> = rows.into_iter().map(|(c,)| c).collect();
self.cache.lock().await.insert(uid.to_string(), mems);
self.cache.lock().await.insert(
uid.to_string(),
UserCache {
memories: rows,
last_access: Instant::now(),
},
);
}
}
pub async fn read_for(&self, user_id: &str) -> String {
pub async fn read_for(&self, user_id: &str, limit: Option<usize>) -> String {
let uid = if user_id.is_empty() { "default" } else { user_id };
self.touch(uid).await;
let max = limit.unwrap_or(DEFAULT_MEMORY_LIMIT);
let cache = self.cache.lock().await;
match cache.get(uid) {
Some(m) if !m.is_empty() => m.iter().enumerate()
.map(|(i, v)| format!("{}. {}", i + 1, v))
Some(uc) if !uc.memories.is_empty() => {
let mems = &uc.memories;
let start = if mems.len() > max { mems.len() - max } else { 0 };
mems[start..]
.iter()
.map(|(id, content)| format!("[{}] {}", id, content))
.collect::<Vec<_>>()
.join("\n"),
.join("\n")
}
_ => "暂无长期记忆".to_string(),
}
}
pub async fn write_for(&self, user_id: &str, content: &str) -> String {
let uid = if user_id.is_empty() { "default" } else { user_id };
self.cache.lock().await.entry(uid.to_string()).or_default().push(content.to_string());
if let Some(ref pool) = self.pool {
let _ = sqlx::query("INSERT INTO user_memories (user_id, content) VALUES ($1, $2)")
// 空内容防御
if content.trim().is_empty() {
return "记忆内容不能为空".to_string();
}
// 精确去重
{
let cache = self.cache.lock().await;
if let Some(uc) = cache.get(uid) {
if uc.memories.iter().any(|(_, c)| c == content) {
return format!("该记忆已存在: {}", content);
}
}
}
// 写入数据库（获取 id）
let new_id = if let Some(ref pool) = self.pool {
match sqlx::query_as::<_, (i32,)>(
"INSERT INTO user_memories (user_id, content) VALUES ($1, $2) RETURNING id",
)
.bind(uid)
.bind(content)
.fetch_one(pool.as_ref())
.await
{
Ok((id,)) => Some(id),
Err(_) => None,
}
} else {
None
};
// 无 DB 时生成临时负数 id
let id = match new_id {
Some(id) => id,
None => {
let cache = self.cache.lock().await;
-(cache
.get(uid)
.map(|uc| uc.memories.len() as i32 + 1)
.unwrap_or(1))
}
};
// 写入缓存
{
let mut cache = self.cache.lock().await;
let uc = cache
.entry(uid.to_string())
.or_insert_with(|| UserCache {
memories: Vec::new(),
last_access: Instant::now(),
});
uc.last_access = Instant::now();
uc.memories.push((id, content.to_string()));
// 容量限制：超限删最旧的（id 最小的）
while uc.memories.len() > MAX_MEMORIES_PER_USER {
if let Some(pos) = uc
.memories
.iter()
.enumerate()
.min_by_key(|(_, (id, _))| *id)
.map(|(i, _)| i)
{
uc.memories.remove(pos);
}
}
}
format!("已添加长期记忆 (id={}): {}", id, content)
}
/// 删除指定 id 的记忆
pub async fn delete_for(&self, user_id: &str, memory_id: i32) -> String {
let uid = if user_id.is_empty() { "default" } else { user_id };
let removed = {
let mut cache = self.cache.lock().await;
if let Some(uc) = cache.get_mut(uid) {
let old_len = uc.memories.len();
uc.memories.retain(|(id, _)| *id != memory_id);
uc.last_access = Instant::now();
old_len != uc.memories.len()
} else {
false
}
};
if !removed {
return format!("未找到 id={} 的记忆", memory_id);
}
if let Some(ref pool) = self.pool {
let _ = sqlx::query("DELETE FROM user_memories WHERE id = $1 AND user_id = $2")
.bind(memory_id)
.bind(uid)
.execute(pool.as_ref())
.await;
}
"已添加长期记忆".to_string()
}
format!("已删除记忆 (id={})", memory_id)
}
/// ## `read_summaries` 工具 —— 读取历史摘要（内存 + 数据库）
///
/// 返回当前 session 中的摘要列表 + 数据库中存档的摘要。
/// 去重：数据库中的摘要如果与 session 中的摘要前 30 个字符相同，则跳过。
pub async fn read_summaries(
session: &Arc<Mutex<super::types::ChatSession>>,
) -> String {
let s = session.lock().await;
let mut lines: Vec<String> = Vec::new();
/// 更新指定 id 的记忆内容
pub async fn update_for(&self, user_id: &str, memory_id: i32, content: &str) -> String {
let uid = if user_id.is_empty() { "default" } else { user_id };
for sum in &s.summaries {
let reason = match sum.reason {
super::types::SummaryReason::Overflow => "上下文压缩",
super::types::SummaryReason::Timeout => "空闲超时",
if content.trim().is_empty() {
return "记忆内容不能为空".to_string();
}
let found = {
let mut cache = self.cache.lock().await;
if let Some(uc) = cache.get_mut(uid) {
if let Some(entry) = uc.memories.iter_mut().find(|(id, _)| *id == memory_id) {
entry.1 = content.to_string();
uc.last_access = Instant::now();
true
} else {
false
}
} else {
false
}
};
lines.push(format!(
"[{}] {} (原因: {})",
sum.created_at.format("%Y-%m-%d %H:%M"),
sum.text,
reason,
));
if !found {
return format!("未找到 id={} 的记忆", memory_id);
}
if let Some(pool) = &s.db_pool {
let uid = if s.current_user_id.is_empty() { "default" } else { &s.current_user_id };
if let Ok(db_summaries) = crate::db::models::load_summaries(pool, uid, 5).await {
for text in db_summaries {
if !lines.iter().any(|l| l.contains(&text.chars().take(30).collect::<String>())) {
lines.push(format!("[数据库存档] {}", text));
}
if let Some(ref pool) = self.pool {
let _ = sqlx::query(
"UPDATE user_memories SET content = $1 WHERE id = $2 AND user_id = $3",
)
.bind(content)
.bind(memory_id)
.bind(uid)
.execute(pool.as_ref())
.await;
}
format!("已更新记忆 (id={}): {}", memory_id, content)
}
}
if lines.is_empty() { "暂无历史摘要".to_string() } else { lines.join("\n---\n") }
}
src/context/types.rs
+10 -46
View File
@@ -21,6 +21,13 @@ use serde::{Deserialize, Serialize};
use std::sync::Arc;
use uuid::Uuid;
/// 历史对话条目（用于从数据库加载历史消息）
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HistoryEntry {
pub role: String,
pub content: String,
}
/// ## 摘要原因
///
/// 区分两种触发摘要的场景：
@@ -49,6 +56,7 @@ pub struct SummaryEntry {
/// 生成原因
pub reason: SummaryReason,
/// 生成时间
#[allow(dead_code)]
pub created_at: DateTime<Utc>,
}
@@ -168,8 +176,8 @@ impl ChatSession {
.find(|s| s.reason == SummaryReason::Overflow)
}
/// 从预载的历史记录初始化会话（Worker 模式）
pub fn load_from_history(&mut self, history: &[crate::ipc::HistoryEntry]) {
/// 从预载的历史记录初始化会话
pub fn load_from_history(&mut self, history: &[HistoryEntry]) {
for entry in history {
match entry.role.as_str() {
"user" => self.add_user(entry.content.clone()),
@@ -183,50 +191,6 @@ impl ChatSession {
}
/// 从数据库加载最近的聊天记录到会话中
pub async fn load_recent_messages(&mut self, limit: i64) {
let pool = match &self.db_pool {
Some(p) => p.clone(),
None => return,
};
let uid = self.current_user_id.clone();
let rows = if uid.is_empty() {
sqlx::query_as::<_, (String, String, String, String)>(
"SELECT direction, user_id, text, message_id FROM chat_records ORDER BY created_at DESC LIMIT $1",
)
.bind(limit)
.fetch_all(pool.as_ref())
.await
} else {
sqlx::query_as::<_, (String, String, String, String)>(
"SELECT direction, user_id, text, message_id FROM chat_records WHERE user_id = $1 ORDER BY created_at DESC LIMIT $2",
)
.bind(&uid)
.bind(limit)
.fetch_all(pool.as_ref())
.await
};
if let Ok(rows) = rows {
let msgs: Vec<Message> = rows
.into_iter()
.rev() // 恢复时间顺序
.map(|(dir, _uid, text, _mid)| {
if dir == "inbound" {
Message::user(text)
} else {
Message::assistant(text)
}
})
.collect();
if !msgs.is_empty() {
tracing::info!("从数据库恢复 {} 条历史消息", msgs.len());
self.messages = msgs;
self.checkpoint = 0;
self.last_user_at = Some(Utc::now());
}
}
}
/// 保存摘要到数据库（委托给 db::models::save_summary）
pub async fn save_summary_to_db(&self, text: &str, reason: &str, message_count: i32) {
src/daemon.rs
+64 -54
View File
@@ -34,15 +34,14 @@
//! - **Tool Consumer** — 执行内置工具，高风险工具先走审批
//! - **Send Consumer** — 将 LLM 回复通过 WeChat API 发回给用户
//!
//! ### 与旧架构的关系
//! ### 架构
//!
//! 旧架构使用 `daemon + worker` 分离模式（UDS IPC），worker 已废弃。
//! 当前是统一的单进程 tokio task 架构。
//! 采用统一的单进程 tokio task 架构。旧版 daemon+worker 分离模式（UDS IPC）已在 v2 中移除。
use crate::channel::ChannelId;
use crate::context::MemoryStore;
use crate::db::Database;
use crate::ipc::HistoryEntry;
use crate::context::HistoryEntry;
use crate::llm::{ChatResult, Conversation, ConversationConfig, ToolExecutor, DEFAULT_SYSTEM_PROMPT, ASYNC_MARKER};
use crate::queue::{ConsumerChannels, EnqueueHandle, MessageKind, PipelineMessage, QueueRunner};
use crate::state::StateManager;
@@ -100,7 +99,7 @@ pub async fn run(
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; }
let account_id = auth.account_id.clone();
info!("Daemon 已启动 (无 worker 进程模式)");
info!("Daemon 已启动");
// 4-6. 审批管理器、工具列表、模型配置
let approval_manager = Arc::new(ApprovalManager::new(
@@ -315,7 +314,7 @@ async fn llm_consumer_loop(
info!("LLM Consumer: user={} text={:.60}", user_id, text);
// 加载数据
ctx.memory_store.load(&user_id).await;
ctx.memory_store.load_if_needed(&user_id).await;
let history = load_history(&ctx.db, &user_id).await;
let sys_prompt = if ctx.system_prompt.is_empty() {
@@ -595,15 +594,19 @@ async fn tool_consumer_loop(
// ─── 统一工具执行 ───
/// 所有工具的统一执行入口 — 仅此一处
async fn execute_tool(name: &str, arguments: &str, ctx: &DaemonCtx, user_id: &str) -> String {
///
/// `call_capability` 通过递归派发到此函数，需 boxing 避免无限大小。
fn execute_tool<'a>(
name: &'a str,
arguments: &'a str,
ctx: &'a DaemonCtx,
user_id: &'a str,
) -> std::pin::Pin<std::boxed::Box<dyn std::future::Future<Output = String> + Send + 'a>> {
Box::pin(async move {
// 上下文工具
match name {
"read_memories" => {
let mems = load_memories(&ctx.memory_store, user_id).await;
if mems.is_empty() { return "暂无长期记忆".to_string(); }
let lines: Vec<String> = mems.iter().enumerate()
.map(|(i, v)| format!("{}. {}", i + 1, v)).collect();
return lines.join("\n");
return ctx.memory_store.read_for(user_id, None).await;
}
"write_memory" => {
let params: serde_json::Value = serde_json::from_str(arguments).unwrap_or_default();
@@ -612,6 +615,20 @@ async fn execute_tool(name: &str, arguments: &str, ctx: &DaemonCtx, user_id: &st
ctx.memory_store.write_for(user_id, content).await;
return format!("已添加长期记忆: {}", content);
}
"delete_memory" => {
let params: serde_json::Value = serde_json::from_str(arguments).unwrap_or_default();
let memory_id = params.get("memory_id").and_then(|v| v.as_i64()).unwrap_or(-1) as i32;
if memory_id <= 0 { return "请提供有效的 memory_id 参数".to_string(); }
return ctx.memory_store.delete_for(user_id, memory_id).await;
}
"update_memory" => {
let params: serde_json::Value = serde_json::from_str(arguments).unwrap_or_default();
let memory_id = params.get("memory_id").and_then(|v| v.as_i64()).unwrap_or(-1) as i32;
let content = params.get("content").and_then(|v| v.as_str()).unwrap_or("");
if memory_id <= 0 { return "请提供有效的 memory_id 参数".to_string(); }
if content.is_empty() { return "请提供 content 参数".to_string(); }
return ctx.memory_store.update_for(user_id, memory_id, content).await;
}
"read_summaries" => {
let sums = load_summaries(&ctx.db, user_id).await;
if sums.is_empty() { return "暂无历史摘要".to_string(); }
@@ -624,48 +641,20 @@ async fn execute_tool(name: &str, arguments: &str, ctx: &DaemonCtx, user_id: &st
return crate::tools::build_capability_guide(&specs);
}
"call_capability" => {
// 解包 {name, prompt} → 提取目标工具名和参数，直接派发
// 解包 {name, prompt} → 提取目标工具名和参数，递归派发
let cp: serde_json::Value = serde_json::from_str(arguments).unwrap_or_default();
let target_name = cp.get("name").and_then(|v| v.as_str()).unwrap_or("").to_string();
if target_name.is_empty() {
return "请提供 name 参数".to_string();
}
// 防止嵌套 call_capability
if target_name == "call_capability" {
return "不支持嵌套调用 call_capability".to_string();
}
let target_args = serde_json::to_string(&crate::tools::unpack_call_params(&cp))
.unwrap_or_default();
// 重新派发到 execute_tool（通过顶层 match 分支避免 async 递归）
// 上下文工具
match target_name.as_str() {
"read_memories" => {
let mems = load_memories(&ctx.memory_store, user_id).await;
if mems.is_empty() { return "暂无长期记忆".to_string(); }
return mems.iter().enumerate()
.map(|(i, v)| format!("{}. {}", i + 1, v)).collect::<Vec<_>>().join("\n");
}
"write_memory" => {
let params: serde_json::Value = serde_json::from_str(&target_args).unwrap_or_default();
let content = params.get("content").and_then(|v| v.as_str()).unwrap_or("");
if content.is_empty() { return "请提供 content 参数".to_string(); }
ctx.memory_store.write_for(user_id, content).await;
return format!("已添加长期记忆: {}", content);
}
"read_summaries" => {
let sums = load_summaries(&ctx.db, user_id).await;
if sums.is_empty() { return "暂无历史摘要".to_string(); }
return sums.iter().enumerate()
.map(|(i, s)| format!("{}. {}", i + 1, s))
.collect::<Vec<_>>().join("\n---\n");
}
"query_capabilities" => {
let specs = crate::tools::specs();
return crate::tools::build_capability_guide(&specs);
}
_ => {}
}
// 内置工具注册表
match crate::tools::execute(&target_name, &target_args).await {
Some(r) => return r.output,
None => return format!("未知工具: {}", target_name),
}
// 递归派发
return execute_tool(&target_name, &target_args, ctx, user_id).await;
}
_ => {}
}
@@ -675,6 +664,7 @@ async fn execute_tool(name: &str, arguments: &str, ctx: &DaemonCtx, user_id: &st
Some(r) => r.output,
None => format!("未知工具: {}", name),
}
})
}
// ─── Send Consumer ───
@@ -738,13 +728,6 @@ async fn load_history(db: &Option<Arc<Database>>, user_id: &str) -> Vec<HistoryE
}
}
async fn load_memories(memory_store: &MemoryStore, user_id: &str) -> Vec<String> {
let text = memory_store.read_for(user_id).await;
// read_for() 返回 "暂无长期记忆" 表示无记忆，否则返回 "1. ...\n2. ..." 格式
if text.starts_with("暂无") { vec![] }
else { text.lines().map(|l| l.trim().to_string()).filter(|l| !l.is_empty()).collect() }
}
async fn load_summaries(db: &Option<Arc<Database>>, user_id: &str) -> Vec<String> {
let Some(db) = db else { return vec![]; };
match crate::db::models::load_summaries(db.pool(), user_id, 5).await {
@@ -794,6 +777,33 @@ fn build_tools_list() -> Vec<serde_json::Value> {
"parameters": { "type": "object", "properties": {} }
}
}));
list.push(serde_json::json!({
"type": "function", "function": {
"name": "delete_memory",
"description": "删除指定 id 的长期记忆。id 来自 read_memories 返回的 [id] 前缀。",
"parameters": {
"type": "object",
"properties": {
"memory_id": { "type": "integer", "description": "要删除的记忆 id" }
},
"required": ["memory_id"]
}
}
}));
list.push(serde_json::json!({
"type": "function", "function": {
"name": "update_memory",
"description": "更新指定 id 的长期记忆内容。id 来自 read_memories 返回的 [id] 前缀。",
"parameters": {
"type": "object",
"properties": {
"memory_id": { "type": "integer", "description": "要更新的记忆 id" },
"content": { "type": "string", "description": "新的记忆内容" }
},
"required": ["memory_id", "content"]
}
}
}));
let specs = crate::tools::specs();
for spec in &specs {
src/db/models.rs
-35
View File
@@ -126,41 +126,6 @@ pub async fn list_recent_chat_records(
// ─── LLM 用量 ───
/// 插入一条 LLM 用量记录
/// kind/entry 为 iPet 遗留兼容字段，写入默认值 ("llm_call" / null)
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, kind, entry)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
"#,
)
.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)
.bind("llm_call")
.bind(serde_json::Value::Null)
.execute(pool)
.await
.map_err(|e| format!("插入 LLM 用量失败: {}", e))?;
Ok(())
}
/// LLM 用量聚合统计
#[derive(Debug, Clone, sqlx::FromRow)]
pub struct LlmUsageStats {
src/ipc.rs
-208
View File
@@ -1,208 +0,0 @@
//! ## IPC 帧格式 —— Unix Domain Socket 长度前缀帧（已废弃）
//!
//! 这是旧架构（daemon + worker 分离模式）的进程间通信协议。
//! 新架构已改为统一的单进程 tokio task，此模块保留仅为编译参考。
//!
//! ### 帧格式
//! ```text
//! [4 字节 u32 大端: payload_len][N 字节: JSON payload]
//! ```
//!
//! ### 消息类型
//! * Daemon → Worker: `TaskFrame`（用户消息 + 上下文 + 环境变量）
//! * Worker → Daemon: `OutputFrame`（Chunk / Reply / NeedApproval / DbWrite / UsageReport / Bye）
//!
//! ### 最大帧大小
//! 10MB（防止恶意超大帧导致 OOM）
use serde::{Deserialize, Serialize};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::UnixStream;
// ─── 帧消息类型 ───
/// Daemon → Worker: 任务描述
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TaskFrame {
pub user_id: String,
pub msg: TaskMessage,
pub history: Vec<HistoryEntry>,
pub memories: Vec<String>,
pub summaries: Vec<String>,
#[serde(default)]
pub approved_tool: Option<String>,
pub env: std::collections::HashMap<String, String>,
pub tools: Vec<serde_json::Value>,
pub system_prompt: String,
pub model: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TaskMessage {
pub from: String,
pub text: String,
pub account_id: String,
#[serde(default)]
pub context_token: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HistoryEntry {
pub role: String,
pub content: String,
}
/// Worker → Daemon: 输出帧
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type")]
pub enum OutputFrame {
/// 流式文本块
#[serde(rename = "chunk")]
Chunk { text: String },
/// 最终回复（完整文本）
#[serde(rename = "reply")]
Reply { text: String },
/// 需要审批
#[serde(rename = "need_approval")]
NeedApproval {
tool: String,
code: String,
message: String,
},
/// Worker 请求 daemon 创建审批（daemon 生成确认码）
#[serde(rename = "request_approval")]
RequestApproval {
tool: String,
reason: String,
},
/// 数据库写入
#[serde(rename = "db_write")]
DbWrite {
table: String,
row: serde_json::Value,
},
/// Token 用量报告
#[serde(rename = "usage")]
UsageReport {
model: String,
provider: String,
prompt_tokens: u32,
completion_tokens: u32,
total_tokens: u32,
cache_hit_tokens: u32,
cache_miss_tokens: u32,
user_id: String,
},
/// Worker 已完成，即将退出
#[serde(rename = "bye")]
Bye,
}
// ─── 帧读写 ───
/// 发送一个 JSON 帧（长度前缀 + JSON）
pub async fn send_frame(stream: &mut UnixStream, msg: &impl Serialize) -> Result<(), String> {
let payload = serde_json::to_vec(msg).map_err(|e| format!("序列化帧失败: {}", e))?;
let len = payload.len() as u32;
let mut header = len.to_be_bytes().to_vec();
header.extend_from_slice(&payload);
stream
.write_all(&header)
.await
.map_err(|e| format!("发送帧失败: {}", e))?;
Ok(())
}
/// 最大帧大小：10MB（防止 OOM）
const MAX_FRAME_SIZE: usize = 10 * 1024 * 1024;
/// 接收一个 JSON 帧
pub async fn recv_frame(stream: &mut UnixStream) -> Result<serde_json::Value, String> {
// 读 4 字节长度头
let mut header = [0u8; 4];
stream
.read_exact(&mut header)
.await
.map_err(|e| format!("读取帧头失败: {}", e))?;
let len = u32::from_be_bytes(header) as usize;
// 防御：拒绝超大帧
if len > MAX_FRAME_SIZE {
return Err(format!("帧过大: {} bytes (max {})", len, MAX_FRAME_SIZE));
}
// 读 payload
let mut payload = vec![0u8; len];
stream
.read_exact(&mut payload)
.await
.map_err(|e| format!("读取帧体失败 (len={}): {}", len, e))?;
serde_json::from_slice(&payload).map_err(|e| {
format!(
"解析帧失败: {} — payload({}): {:.200}",
e,
len,
String::from_utf8_lossy(&payload[..std::cmp::min(len, 200)])
)
})
}
// ─── 便捷方法: 按类型收发 ───
/// 接收并解析为 TaskFrame
pub async fn recv_task(stream: &mut UnixStream) -> Result<TaskFrame, String> {
let v = recv_frame(stream).await?;
serde_json::from_value(v).map_err(|e| format!("解析 TaskFrame 失败: {}", e))
}
/// 接收并解析为 OutputFrame
#[allow(dead_code)]
pub async fn recv_output(stream: &mut UnixStream) -> Result<OutputFrame, String> {
let v = recv_frame(stream).await?;
serde_json::from_value(v).map_err(|e| format!("解析 OutputFrame 失败: {}", e))
}
// ─── 测试 ───
#[cfg(test)]
mod tests {
use super::*;
use tokio::net::UnixListener;
#[tokio::test]
async fn test_frame_roundtrip() {
let sock_path = "/tmp/ias_ipc_test.sock";
let _ = std::fs::remove_file(sock_path);
let listener = UnixListener::bind(sock_path).unwrap();
let server = tokio::spawn(async move {
let (mut stream, _) = listener.accept().await.unwrap();
let v = recv_frame(&mut stream).await.unwrap();
assert_eq!(v["type"], "chunk");
assert_eq!(v["text"], "hello");
let output = OutputFrame::Reply {
text: "world".into(),
};
send_frame(&mut stream, &output).await.unwrap();
});
let mut client = UnixStream::connect(sock_path).await.unwrap();
let chunk = OutputFrame::Chunk {
text: "hello".into(),
};
send_frame(&mut client, &chunk).await.unwrap();
let resp = recv_output(&mut client).await.unwrap();
match resp {
OutputFrame::Reply { text } => assert_eq!(text, "world"),
_ => panic!("expected reply"),
}
server.await.unwrap();
let _ = std::fs::remove_file(sock_path);
}
}
src/llm/conversation.rs
+8
View File
@@ -66,7 +66,9 @@ pub type Summarizer = Arc<
#[derive(Debug)]
pub struct ChatResult {
pub reply: String,
#[allow(dead_code)]
pub used_tools: bool,
#[allow(dead_code)]
pub usage: Option<Usage>,
/// 是否有异步工具调用仍在等待结果
pub has_pending_async: bool,
@@ -133,12 +135,15 @@ impl Conversation {
pub fn session(&self) -> Arc<Mutex<ChatSession>> {
Arc::clone(&self.session)
}
#[allow(dead_code)]
pub fn model(&self) -> &str {
&self.config.model
}
#[allow(dead_code)]
pub fn provider_name(&self) -> &str {
self.provider.name()
}
#[allow(dead_code)]
pub fn spec(&self) -> &ConversationConfig {
&self.config
}
@@ -195,6 +200,7 @@ impl Conversation {
/// 3. 构建上下文（system prompt + 摘要 + 近期消息）
/// 4. 发起流式请求
/// 5. 返回 `ChatHandle`（可逐 chunk 消费）
#[allow(dead_code)]
pub async fn chat(&self, user_message: String) -> Result<ChatHandle, String> {
// 空闲超时检查
{
@@ -378,6 +384,7 @@ impl Conversation {
// ─── ChatHandle ───
#[allow(dead_code)]
pub struct ChatHandle {
rx: Option<StreamReceiver>,
full_text: String,
@@ -386,6 +393,7 @@ pub struct ChatHandle {
session: Arc<Mutex<ChatSession>>,
}
#[allow(dead_code)]
impl ChatHandle {
pub async fn next_chunk(&mut self) -> Option<StreamChunk> {
let chunk = self.rx.as_mut()?.recv().await;
src/llm/mod.rs
+1 -1
View File
@@ -22,4 +22,4 @@ pub mod provider;
pub mod types;
pub use conversation::{ChatResult, Conversation, ToolExecutor, ASYNC_MARKER, DEFAULT_SYSTEM_PROMPT};
pub use types::{ConversationConfig, Usage};
pub use types::ConversationConfig;
src/main.rs
+83 -614
View File
@@ -5,13 +5,12 @@
//! ```text
//! CLI (main.rs → clap)
//! ├── ias login → 扫码登录，保存 token 到 DB 或文件
//! ├── ias listen --llm → Daemon 模式（主运行模式）
//! ├── ias daemon → 守护进程模式（主运行模式）
//! │ └── daemon.rs: 长轮询 → MessageQueue → 3 个 tokio 消费者
//! │ ├── LLM Consumer → Conversation.chat_with_tools() → DeepSeek API
//! │ ├── Tool Consumer → tools::execute() / ApprovalManager
//! │ └── Send Consumer → WeChatClient.send_message()
//! ├── ias service → 等同 listen --llm 但启用文件日志
//! ├── ias daemon → 守护进程入口（等同 listen --llm）
//! ├── ias daemon --log-file → 等同 daemon 但启用文件日志
//! ├── ias tool <tool> → 直接调用内置工具（无需登录）
//! └── ias usage → Token 消耗统计查询
//! ```
@@ -24,7 +23,7 @@
//! ```
//!
//! ### 关键设计决策
//! 1. **单进程架构** — 旧版 daemon+worker 分离模式已废弃，统一用 tokio task
//! 1. **单进程架构** — 统一用 tokio task
//! 2. **公平队列** — MessageQueue 按用户轮转，防止一个用户刷屏饿死其他人
//! 3. **可选数据库** — PostgreSQL 不可用时回退到文件存储
//! 4. **两层元工具** — LLM 通过 query_capabilities / call_capability 间接调用工具
@@ -34,7 +33,6 @@ mod cli;
mod context;
mod daemon;
mod db;
mod ipc;
mod llm;
mod logger;
mod queue;
@@ -42,29 +40,25 @@ mod scheduler;
mod state;
mod tools;
mod wechat;
mod worker;
use clap::Parser;
use cli::{AmapAction, Cli, Commands, MemosAction, ScheduledTaskAction, ToolCommand};
use context::MemoryStore;
use db::Database;
use llm::{ChatResult, Conversation, ConversationConfig, DEFAULT_SYSTEM_PROMPT, ToolExecutor, Usage};
use std::sync::Arc;
use tools::approval::{ApprovalDecision, ApprovalManager};
use tools::types::ExecutionContext;
use tracing::{error, info, warn};
use tracing::{error, info};
use wechat::client::WeChatClient;
#[tokio::main]
async fn main() {
let cli = Cli::parse();
let is_service = matches!(cli.command, Commands::Service);
let log_file = matches!(&cli.command, Commands::Daemon { log_file: true, .. });
let log_dir = dirs::home_dir()
.unwrap_or_else(|| std::path::PathBuf::from("."))
.join(".ias")
.join("logs");
let log_dir_str = log_dir.to_string_lossy().to_string();
logger::init_logger(Some(&log_dir_str), is_service);
logger::init_logger(Some(&log_dir_str), log_file);
// .env 加载：先当前目录，再 ~/.ias/.env
if std::path::Path::new(".env").exists() {
@@ -110,12 +104,6 @@ async fn main() {
Commands::Login { timeout } => {
cmd_login(timeout, &database, &file_state).await;
}
Commands::Listen {
llm: enable_llm,
echo,
} => {
cmd_listen(enable_llm, echo, &database, &file_state, &memory_store).await;
}
Commands::Send {
to,
text,
@@ -133,20 +121,15 @@ async fn main() {
} => {
cmd_usage(&database, since, until, model).await;
}
Commands::Service => cmd_listen(true, false, &database, &file_state, &memory_store).await,
Commands::Daemon { sock } => {
Commands::Daemon { sock, .. } => {
cmd_daemon(sock, &database, &file_state, &memory_store).await;
}
Commands::Worker { sock } => {
warn!("Worker 模式已废弃，请使用 daemon 模式");
#[allow(deprecated)]
if let Err(e) = worker::run(&sock).await {
error!("Worker 失败: {}", e);
}
}
Commands::ScheduledTask(action) => {
cmd_scheduled_task(action, &database).await;
}
Commands::Tools => {
cmd_tools().await;
}
Commands::Tool(..) => unreachable!("Tool 命令已提前处理"),
}
}
@@ -206,536 +189,10 @@ async fn cmd_login(
}
}
async fn cmd_listen(
enable_llm: bool,
echo: bool,
database: &Option<Arc<Database>>,
file_state: &state::StateManager,
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,
};
if let Err(e) = db::models::save_auth(db.pool(), &a).await {
error!("保存认证信息失败: {}", e);
}
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 mut 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;
}
let approval_manager = Arc::new(ApprovalManager::new(
database.as_ref().map(|db| 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 =
std::env::var("DEEPSEEK_MODEL").unwrap_or_else(|_| "deepseek-v4-flash".to_string());
let mut cfg = ConversationConfig {
system_prompt: sys_prompt,
model,
..Default::default()
};
// 各项能力通过这两个元工具实现：查询能力 + 调用能力
// 具体工具描述通过 query_capabilities 返回
let mut tools_list: Vec<serde_json::Value> = Vec::new();
// query_capabilities：列出所有可用工具及其说明
tools_list.push(serde_json::json!({
"type": "function",
"function": {
"name": "query_capabilities",
"description": "列出所有可用工具的详细说明、参数格式和调用示例。在需要了解有哪些工具可用时首先调用此工具。",
"parameters": { "type": "object", "properties": {} }
}
}));
// call_capability：按名称调用任意工具
tools_list.push(serde_json::json!({
"type": "function",
"function": {
"name": "call_capability",
"description": "调用一个已注册的工具。name 为工具名（来自 query_capabilities），工具所需的具体参数（如 location、days）应直接放在 JSON 中。",
"parameters": {
"type": "object",
"properties": {
"name": { "type": "string", "description": "工具名称" },
"prompt": { "type": "string", "description": "json格式参数，来自 query_capabilities 中的对应工具的描述。" }
},
"required": ["name"]
}
}
}));
// 上下文工具（仍然直接暴露给 LLM）
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) => {
// 设置 db_pool（启动时不加载消息，等第一条消息到达后再按需恢复）
if let Some(db) = database {
let session = c.session();
let mut s = session.lock().await;
s.db_pool = Some(Arc::new(db.pool().clone()));
}
info!("LLM 已初始化: {} / {}", c.provider_name(), c.model());
c
}
Err(e) => {
error!("初始化 LLM 失败: {}", e);
return;
}
};
// 构建工具执行器
let session = conv.session();
let approval = approval_manager.clone();
let memory_store = memory_store.clone();
let send_wechat: tools::types::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(|_| ())
})
})
};
let executor: ToolExecutor = Arc::new(move |name: &str, args_json: &str| {
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 {
let user_id = session.lock().await.current_user_id.clone();
// 上下文工具直接处理
if n == "read_memories" {
return Ok(memory_store.read_for(&user_id).await);
}
if n == "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_for(&user_id, content).await);
}
if n == "read_summaries" {
return Ok(context::tools::read_summaries(&session).await);
}
// query_capabilities: 列出所有内置工具
if n == "query_capabilities" {
let specs = tools::specs();
return Ok(tools::build_capability_guide(&specs));
}
// 确定目标工具名和参数
// call_capability: 从 {name, prompt} 中提取 name，从 prompt 中解包真实参数
let (target_name, target_args) = if n == "call_capability" {
let cp: serde_json::Value = serde_json::from_str(&args).unwrap_or_default();
let name = cp
.get("name")
.and_then(|v| v.as_str())
.map(|s| s.to_string())
.unwrap_or_default();
let unpacked = tools::unpack_call_params(&cp);
(name, serde_json::to_string(&unpacked).unwrap_or_default())
} else {
(n.clone(), args.clone())
};
// 内置工具
if tools::is_builtin(&target_name) {
if tools::is_high_risk(&target_name) {
let mut ctx = ExecutionContext::new(&user_id);
ctx = ctx.with_approval(approval.clone());
ctx = ctx.with_wechat(send.clone());
let approved = builtin_approve(&ctx, &target_name).await;
match approved {
Ok(true) => {}
Ok(false) => return Ok(SkillResult::rejected(&target_name).output),
Err(e) => return Ok(e),
}
}
if let Some(result) =
tools::execute(&target_name, &target_args).await
{
return Ok(result.output);
}
}
Ok(format!("未知工具: {}", target_name))
})
});
conv.set_tool_executor(executor);
Some(Arc::new(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!("定时任务调度器已启动");
// 审批定期清理
let approval_clean = approval_manager.clone();
tokio::spawn(async move {
loop {
tokio::time::sleep(std::time::Duration::from_secs(30)).await;
approval_clean.clean_expired().await;
}
});
}
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, memory_store) => {}
}
}
async fn listen_loop(
client: &WeChatClient,
enable_llm: bool,
echo: bool,
account_id: &str,
database: &Option<Arc<Database>>,
file_state: &state::StateManager,
conversation: Option<Arc<Conversation>>,
approval_manager: &ApprovalManager,
memory_store: &Arc<MemoryStore>,
) {
let last_user_id = Arc::new(tokio::sync::Mutex::new(String::new()));
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.map(|id| id.to_string()).unwrap_or_default();
info!("收到消息 from={}: {}", from, text);
// === 入库：收到的消息 ===
if let Some(db) = database
&& let Err(e) = db::models::insert_chat_record(
db.pool(),
"inbound",
from,
account_id,
text,
"wechat",
ctx_token,
&msg_id,
)
.await
{
error!("保存聊天记录失败: {}", e);
}
// 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 Err(e) = db::models::insert_chat_record(
db.pool(),
"outbound",
from,
account_id,
text,
"echo",
ctx_token,
&sent_id,
)
.await
{
error!("保存聊天记录失败: {}", e);
}
}
Err(e) => error!("回显失败: {}", e),
}
}
// 审批回复检查
let is_approval_reply =
approval_manager.handle_reply(from, text).await.is_some();
if is_approval_reply {
info!("消息匹配审批确认码，不传给 LLM");
continue;
}
// 用户切换：清空上下文，加载新用户历史
{
let mut last_user = last_user_id.lock().await;
if *last_user != *from {
info!("用户切换: {} → {}", *last_user, from);
if let Some(conv) = &conversation {
let session = conv.session();
let mut s = session.lock().await;
s.messages.clear();
s.checkpoint = 0;
s.summaries.clear();
s.last_user_at = None;
s.current_user_id = from.to_string();
// 加载该用户最近 20 条历史
s.load_recent_messages(20).await;
drop(s);
memory_store.load(from).await;
}
}
*last_user = from.to_string();
}
// LLM 回复（异步执行，避免审批阻塞主循环）
if enable_llm
&& let Some(conv) = &conversation {
let conv = conv.clone();
let client = client.clone();
let database = database.clone();
let from_owned = from.to_string();
let text_owned = text.to_string();
let ctx_owned = ctx_token.map(String::from);
let aid = account_id.to_string();
tokio::spawn(async move {
let reply_result = if conv.spec().tools.is_some() {
generate_reply_with_tools(&conv, text_owned, &database).await
} else {
generate_reply(&conv, text_owned, &database).await
};
let reply = match reply_result {
Ok(r) if !r.trim().is_empty() => r,
Ok(_) => "抱歉，生成回复为空，请重试。".to_string(),
Err(e) => format!("处理消息时出错：{:.200}", e),
};
match client
.send_text(&from_owned, &reply, ctx_owned.as_deref())
.await
{
Ok(sent_id) => {
info!("LLM 回复成功 msg_id={}", sent_id);
if let Some(db) = database
&& let Err(e) = db::models::insert_chat_record(
db.pool(),
"outbound",
&from_owned,
&aid,
&reply,
"llm",
ctx_owned.as_deref(),
&sent_id,
)
.await
{
error!("保存聊天记录失败: {}", 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 mut handle = conv.chat(user_text).await?;
// 先消费流（usage 在 Done chunk 中才可用）
let reply = handle.consume().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;
}
Ok(reply)
}
async fn generate_reply_with_tools(
conv: &Conversation,
user_text: String,
db: &Option<Arc<Database>>,
) -> Result<String, String> {
let ChatResult { reply, used_tools: _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 Err(e) = 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
{
tracing::error!("存储 LLM 用量失败: {}", e);
}
}
async fn cmd_whoami(database: &Option<Arc<Database>>, file_state: &state::StateManager) {
let auth: Option<state::AuthState> = if let Some(db) = database {
@@ -880,28 +337,6 @@ async fn cmd_send(
}
}
use tools::types::SkillResult;
async fn builtin_approve(ctx: &ExecutionContext, name: &str) -> Result<bool, String> {
let manager = match ctx.approval_manager.as_ref() {
Some(m) => m.clone(),
None => return Ok(true),
};
let (code, rx) = manager.create(&ctx.user_id, name).await?;
let msg = format!(
"⚠️ 操作确认\n\n技能：{}\n\n确认码：{}\n\n回复确认码以继续\n回复 0 或「取消」以取消\n(5分钟内有效，最多3次尝试)",
name, code
);
if let Some(ref send) = ctx.send_wechat {
send(&ctx.user_id, &msg).await?
}
match tokio::time::timeout(std::time::Duration::from_secs(300), rx).await {
Ok(Ok(ApprovalDecision::Approved)) => Ok(true),
_ => Ok(false),
}
}
// ─── 定时任务 CLI 命令 ───
@@ -1054,6 +489,47 @@ async fn cmd_scheduled_task(action: ScheduledTaskAction, database: &Option<Arc<D
}
}
// ─── 工具列表 CLI 命令 ───
async fn cmd_tools() {
let specs = tools::specs();
if specs.is_empty() {
println!("📋 没有可用的工具");
return;
}
println!("📋 LLM 可调用的工具列表（共 {} 个）:\n", specs.len());
println!("{:=<80}", "");
for spec in &specs {
let risk = match spec.risk_level {
tools::types::RiskLevel::Low => "低",
tools::types::RiskLevel::High => "高 ⚠️",
};
println!("🔹 {}", spec.name);
println!(" 描述: {}", spec.description);
println!(" 风险: {}", risk);
println!(" 超时: {}秒", spec.timeout_secs);
// 显示参数
if let Some(props) = spec.parameters.get("properties").and_then(|v| v.as_object()) {
if !props.is_empty() {
println!(" 参数:");
for (name, info) in props {
let desc = info.get("description").and_then(|v| v.as_str()).unwrap_or("");
let required = spec.parameters.get("required")
.and_then(|v| v.as_array())
.map(|arr| arr.iter().any(|r| r.as_str() == Some(name)))
.unwrap_or(false);
let req_mark = if required { "*必填" } else { "可选" };
println!(" - {} ({}) {}", name, req_mark, desc);
}
}
}
println!("{:=<80}", "");
}
}
// ─── 工具 CLI 命令 ───
async fn cmd_tool(cmd: ToolCommand) {
@@ -1063,23 +539,28 @@ async fn cmd_tool(cmd: ToolCommand) {
println!("{}", result.output);
}
ToolCommand::Memos(action) => {
let args = match action {
MemosAction::List => serde_json::json!({"action": "list"}),
let result = match action {
MemosAction::List => {
tools::builtins::memos::list_execute(serde_json::json!({})).await
}
MemosAction::Add { content } => {
serde_json::json!({"action": "add", "content": content})
tools::builtins::memos::add_execute(serde_json::json!({"content": content})).await
}
MemosAction::Delete { id } => {
serde_json::json!({"action": "delete", "id": id})
tools::builtins::memos::delete_execute(serde_json::json!({"id": id})).await
}
};
let args_str = serde_json::to_string(&args).unwrap_or_default();
let result = tools::builtins::memos::execute(&args_str).await;
println!("{}", result.output);
}
ToolCommand::Weather { location, days } => {
let params = serde_json::json!({"location": location, "days": days});
let result = tools::builtins::weather::execute(params).await;
println!("{}", result.output);
let now_params = serde_json::json!({"location": location});
let fc_params = serde_json::json!({"location": location, "days": days});
let now = tools::builtins::weather::now_execute(now_params).await;
let fc = tools::builtins::weather::forecast_execute(fc_params).await;
println!("=== 实时天气 ===");
println!("{}", now.output);
println!("\n=== 天气预报 ===");
println!("{}", fc.output);
}
ToolCommand::Search {
query,
@@ -1139,40 +620,34 @@ async fn cmd_amap(action: AmapAction) {
if let Some(r) = radius {
params["radius"] = serde_json::json!(r);
}
let result = tools::builtins::amap::execute("amap_poi_search", params).await;
if let Some(r) = result {
let result = tools::builtins::amap::poi_search_execute(params).await;
println!(
"{}\n{}",
r.output,
serde_json::to_string(&r.data).unwrap_or_default()
result.output,
serde_json::to_string(&result.data).unwrap_or_default()
);
}
}
AmapAction::Geocode { address, city } => {
let mut params = serde_json::json!({"address": address});
if let Some(c) = city {
params["city"] = serde_json::json!(c);
}
let result = tools::builtins::amap::execute("amap_geocode", params).await;
if let Some(r) = result {
let result = tools::builtins::amap::geocode_execute(params).await;
println!(
"{}\n{}",
r.output,
serde_json::to_string(&r.data).unwrap_or_default()
result.output,
serde_json::to_string(&result.data).unwrap_or_default()
);
}
}
AmapAction::ReverseGeocode { location } => {
let params = serde_json::json!({"location": location});
let result = tools::builtins::amap::execute("amap_reverse_geocode", params).await;
if let Some(r) = result {
let result = tools::builtins::amap::reverse_geocode_execute(params).await;
println!(
"{}\n{}",
r.output,
serde_json::to_string(&r.data).unwrap_or_default()
result.output,
serde_json::to_string(&result.data).unwrap_or_default()
);
}
}
AmapAction::RoutePlan {
r#type,
origin,
@@ -1195,15 +670,13 @@ async fn cmd_amap(action: AmapAction) {
if let Some(s) = strategy {
params["strategy"] = serde_json::json!(s);
}
let result = tools::builtins::amap::execute("amap_route_plan", params).await;
if let Some(r) = result {
let result = tools::builtins::amap::route_plan_execute(params).await;
println!(
"{}\n{}",
r.output,
serde_json::to_string(&r.data).unwrap_or_default()
result.output,
serde_json::to_string(&result.data).unwrap_or_default()
);
}
}
AmapAction::TravelPlan {
city,
interests,
@@ -1218,27 +691,23 @@ async fn cmd_amap(action: AmapAction) {
"interests": interests_arr,
"route_type": route_type,
});
let result = tools::builtins::amap::execute("amap_travel_plan", params).await;
if let Some(r) = result {
let result = tools::builtins::amap::travel_plan_execute(params).await;
println!(
"{}\n{}",
r.output,
serde_json::to_string(&r.data).unwrap_or_default()
result.output,
serde_json::to_string(&result.data).unwrap_or_default()
);
}
}
AmapAction::MapLink { data } => {
let map_data: serde_json::Value =
serde_json::from_str(&data).unwrap_or(serde_json::json!([]));
let params = serde_json::json!({"map_data": map_data});
let result = tools::builtins::amap::execute("amap_map_link", params).await;
if let Some(r) = result {
let result = tools::builtins::amap::map_link_execute(params).await;
println!(
"{}\n{}",
r.output,
serde_json::to_string(&r.data).unwrap_or_default()
result.output,
serde_json::to_string(&result.data).unwrap_or_default()
);
}
}
}
}
src/tools/builtins/amap.rs
+2
View File
@@ -904,6 +904,7 @@ fn urlencoding(s: &str) -> String {
// 统一分发入口
// ═══════════════════════════════════════════════
#[allow(dead_code)]
pub fn specs() -> Vec<SkillSpec> {
vec![
poi_search_spec(),
@@ -915,6 +916,7 @@ pub fn specs() -> Vec<SkillSpec> {
]
}
#[allow(dead_code)]
pub async fn execute(name: &str, params: serde_json::Value) -> Option<SkillResult> {
match name {
"amap_poi_search" => Some(poi_search_execute(params).await),
src/tools/builtins/memos.rs
+149 -94
View File
@@ -1,109 +1,56 @@
//! ## 备忘录工具 —— 本地文件存储的简易备忘录
//!
//! 支持三种操作：
//! - `add` — 添加一条备忘录（自动分配自增 ID）
//! - `list` — 列出所有备忘录
//! - `delete` — 按 ID 删除备忘录
//! 拆分为 4 个独立工具：
//! - `memos_list` — 列出所有备忘录
//! - `memos_add` — 添加备忘录
//! - `memos_get` — 查看指定备忘录
//! - `memos_delete` — 按 ID 删除备忘录
//!
//! ### 存储
//! 数据存储在 `.data/memos.json` 文件中，JSON 数组格式。
//! 每条记录包含：`id`、`content`、`created_at`。
//!
//! ### 智能推断
//! 当通过 `call_capability` 调用时，可以从 prompt 文本中推断操作类型：
//! - 包含"添加"/"新增"/"add" → 自动转为 add 操作
//! - 包含"删除"/"移除"/"delete" → 自动转为 delete 操作
//! - 其他 → 默认 list 操作
use chrono::Local;
use super::super::types::{RiskLevel, SkillResult, SkillSpec};
/// 返回该工具的元数据描述（供 query_capabilities 元工具使用）
pub fn spec() -> SkillSpec {
SkillSpec {
name: "manage_memos".into(),
description: "管理备忘录：添加/列出/删除".into(),
risk_level: RiskLevel::Low,
parameters: serde_json::json!({
"type": "object",
"properties": {
"action": {"type": "string", "enum": ["add", "list", "delete"]},
"content": {"type": "string"},
"id": {"type": "integer"}
}
}),
timeout_secs: 10,
}
}
/// 数据文件路径
const DATA_PATH: &str = ".data/memos.json";
/// 执行备忘录操作
///
/// 从 JSON 参数中解析 action/content/id，执行对应操作。
/// 数据文件路径为 `.data/memos.json`，自动创建目录。
pub async fn execute(args_json: &str) -> SkillResult {
let mut params: serde_json::Value = serde_json::from_str(args_json).unwrap_or_default();
// call_capability 透传完整参数，action/content 在顶层
// 如果顶层没有 action，尝试从 prompt 推断
if params.get("action").is_none() {
let prompt = params.get("prompt").and_then(|v| v.as_str()).unwrap_or("");
let action = if prompt.contains("添加") || prompt.contains("新增") || prompt.contains("add")
{
"add"
} else if prompt.contains("删除") || prompt.contains("移除") || prompt.contains("delete")
{
"delete"
} else {
"list"
};
let content = params.get("prompt").and_then(|v| v.as_str()).unwrap_or("");
let memo_id: i32 = content
.chars()
.filter(|c| c.is_ascii_digit())
.collect::<String>()
.parse()
.unwrap_or(0);
params = serde_json::json!({"action": action, "content": content, "id": memo_id});
}
let action = params
.get("action")
.and_then(|v| v.as_str())
.unwrap_or("list");
let path = ".data/memos.json";
if let Err(e) = std::fs::create_dir_all(".data") {
return SkillResult::error(format!("创建目录失败: {}", e));
}
let mut data: Vec<serde_json::Value> = if std::path::Path::new(path).exists() {
serde_json::from_str(&std::fs::read_to_string(path).unwrap_or_default()).unwrap_or_default()
/// 读取所有备忘录
fn load_all() -> Vec<serde_json::Value> {
if std::path::Path::new(DATA_PATH).exists() {
serde_json::from_str(&std::fs::read_to_string(DATA_PATH).unwrap_or_default()).unwrap_or_default()
} else {
vec![]
};
}
}
match action {
"add" => {
let content = params.get("content").and_then(|v| v.as_str()).unwrap_or("");
if content.is_empty() {
return SkillResult::error("请提供 content 参数");
/// 写入所有备忘录
fn save_all(data: &[serde_json::Value]) -> Result<(), String> {
std::fs::create_dir_all(".data").map_err(|e| format!("创建目录失败: {}", e))?;
std::fs::write(DATA_PATH, serde_json::to_string_pretty(data).unwrap())
.map_err(|e| format!("写入备忘录失败: {}", e))
}
let next_id = data
.iter()
.filter_map(|m| m.get("id").and_then(|v| v.as_i64()))
.max()
.unwrap_or(0)
+ 1;
data.push(serde_json::json!({
"id": next_id, "content": content,
"created_at": Local::now().format("%Y-%m-%d %H:%M").to_string()
}));
if let Err(e) = std::fs::write(path, serde_json::to_string_pretty(&data).unwrap()) {
return SkillResult::error(format!("写入备忘录失败: {}", e));
// ═══════════════════════════════════════════════
// 1. memos_list — 列出所有备忘录
// ═══════════════════════════════════════════════
pub fn list_spec() -> SkillSpec {
SkillSpec {
name: "memos_list".into(),
description: "列出所有备忘录，显示每条备忘录的编号、内容和创建时间".into(),
risk_level: RiskLevel::Low,
parameters: serde_json::json!({"type": "object", "properties": {}}),
timeout_secs: 5,
}
SkillResult::ok(format!("已添加备忘录 #{}", next_id))
}
"list" => {
pub async fn list_execute(_params: serde_json::Value) -> SkillResult {
let data = load_all();
if data.is_empty() {
return SkillResult::ok("暂无备忘录".to_string());
return SkillResult::ok("暂无备忘录。你可以通过 memos_add 添加一条备忘录。");
}
let lines: Vec<String> = data
.iter()
@@ -118,18 +65,126 @@ pub async fn execute(args_json: &str) -> SkillResult {
.collect();
SkillResult::ok(lines.join("\n"))
}
"delete" => {
// ═══════════════════════════════════════════════
// 2. memos_add — 添加备忘录
// ═══════════════════════════════════════════════
pub fn add_spec() -> SkillSpec {
SkillSpec {
name: "memos_add".into(),
description: "添加一条备忘录，自动分配编号。参数 content 为备忘录内容".into(),
risk_level: RiskLevel::Low,
parameters: serde_json::json!({
"type": "object",
"properties": {
"content": {"type": "string", "description": "备忘录内容，如 明天下午2点开会"}
},
"required": ["content"]
}),
timeout_secs: 5,
}
}
pub async fn add_execute(params: serde_json::Value) -> SkillResult {
let content = params.get("content").and_then(|v| v.as_str()).unwrap_or("");
if content.is_empty() {
return SkillResult::error("请提供 content（备忘录内容）");
}
let mut data = load_all();
let next_id = data
.iter()
.filter_map(|m| m.get("id").and_then(|v| v.as_i64()))
.max()
.unwrap_or(0)
+ 1;
data.push(serde_json::json!({
"id": next_id,
"content": content,
"created_at": Local::now().format("%Y-%m-%d %H:%M").to_string()
}));
if let Err(e) = save_all(&data) {
return SkillResult::error(e);
}
SkillResult::ok(format!("✅ 已添加备忘录 #{}", next_id))
}
// ═══════════════════════════════════════════════
// 3. memos_get — 查看指定备忘录
// ═══════════════════════════════════════════════
pub fn get_spec() -> SkillSpec {
SkillSpec {
name: "memos_get".into(),
description: "查看指定编号的备忘录详情。参数 id 为备忘录编号".into(),
risk_level: RiskLevel::Low,
parameters: serde_json::json!({
"type": "object",
"properties": {
"id": {"type": "integer", "description": "备忘录编号"}
},
"required": ["id"]
}),
timeout_secs: 5,
}
}
pub async fn get_execute(params: serde_json::Value) -> SkillResult {
let id = params.get("id").and_then(|v| v.as_i64()).unwrap_or(0);
if id <= 0 {
return SkillResult::error("请提供有效的 id（备忘录编号）");
}
let data = load_all();
match data.iter().find(|m| m.get("id").and_then(|v| v.as_i64()) == Some(id)) {
Some(memo) => {
let output = format!(
"#{} [{}]\n{}",
memo["id"],
memo.get("created_at").and_then(|v| v.as_str()).unwrap_or("?"),
memo["content"].as_str().unwrap_or("?")
);
SkillResult::ok(output)
}
None => SkillResult::error(format!("未找到备忘录 #{}", id)),
}
}
// ═══════════════════════════════════════════════
// 4. memos_delete — 删除备忘录
// ═══════════════════════════════════════════════
pub fn delete_spec() -> SkillSpec {
SkillSpec {
name: "memos_delete".into(),
description: "删除指定编号的备忘录。参数 id 为备忘录编号".into(),
risk_level: RiskLevel::Low,
parameters: serde_json::json!({
"type": "object",
"properties": {
"id": {"type": "integer", "description": "要删除的备忘录编号"}
},
"required": ["id"]
}),
timeout_secs: 5,
}
}
pub async fn delete_execute(params: serde_json::Value) -> SkillResult {
let id = params.get("id").and_then(|v| v.as_i64()).unwrap_or(0);
if id <= 0 {
return SkillResult::error("请提供有效的 id（备忘录编号）");
}
let mut data = load_all();
let before = data.len();
data.retain(|m| m.get("id").and_then(|v| v.as_i64()) != Some(id));
if data.len() == before {
return SkillResult::error(format!("未找到备忘录 #{}", id));
}
if let Err(e) = std::fs::write(path, serde_json::to_string_pretty(&data).unwrap()) {
return SkillResult::error(format!("写入备忘录失败: {}", e));
}
SkillResult::ok(format!("已删除备忘录 #{}", id))
}
_ => SkillResult::error("未知操作，支持: add, list, delete"),
if let Err(e) = save_all(&data) {
return SkillResult::error(e);
}
SkillResult::ok(format!("✅ 已删除备忘录 #{}", id))
}
src/tools/builtins/mod.rs
+3 -2
View File
@@ -5,10 +5,10 @@
//! | 模块 | 工具名 | 功能 |
//! |------|--------|------|
//! | `datetime` | `get_current_datetime` | 获取当前北京时间 |
//! | `weather` | `query_weather` | 和风天气查询（实时/预报/逐时） |
//! | `weather` | `weather_now` / `weather_forecast` / `weather_hourly` | 和风天气查询（实时/预报/逐时） |
//! | `web_search` | `web_search` | Tavily 联网搜索 |
//! | `fetch_page` | `fetch_page` | 网页正文提取（可读性算法） |
//! | `memos` | `manage_memos` | 备忘录增删查 |
//! | `memos` | `memos_list` / `memos_add` / `memos_get` / `memos_delete` | 备忘录增删查 |
//! | `amap` | `amap_*` | 高德地图系列（POI/地理编码/路径规划/旅游规划） |
//!
//! 每个工具都提供 `spec()` 返回元数据，和 `execute()` 执行入口。
@@ -16,6 +16,7 @@
pub mod amap;
pub mod datetime;
// pub mod email;
pub mod fetch_page;
pub mod memos;
pub mod scheduled_task;
src/tools/builtins/scheduled_task.rs
+149 -109
View File
@@ -1,7 +1,11 @@
//! ## 定时任务管理工具 —— 增删改查定时任务
//!
//! 通过 PostgreSQL 数据库管理 `scheduled_tasks` 表。
//! 支持 list / add / update / delete / toggle 五种操作。
//! 拆分为 5 个独立工具：
//! - `scheduled_task_list` — 列出所有定时任务
//! - `scheduled_task_add` — 添加定时任务
//! - `scheduled_task_update` — 更新定时任务
//! - `scheduled_task_delete` — 删除定时任务
//! - `scheduled_task_toggle` — 启用/禁用定时任务
//!
//! ### 任务类型
//! - `model` — 模型任务，LLM 可以自由增删改查
@@ -18,85 +22,8 @@ static POOL: LazyLock<Option<sqlx::PgPool>> = LazyLock::new(|| {
Some(pool)
});
/// 返回该工具的元数据描述
pub fn spec() -> SkillSpec {
SkillSpec {
name: "manage_scheduled_tasks".into(),
description: "管理定时任务：列出/添加/更新/删除/启用禁用。可以管理模型任务（task_type=model），系统任务（task_type=system）只能查看不能修改。需要 PostgreSQL 数据库。".into(),
risk_level: RiskLevel::Low,
parameters: serde_json::json!({
"type": "object",
"properties": {
"action": {
"type": "string",
"enum": ["list", "add", "update", "delete", "toggle"],
"description": "操作类型"
},
"name": {
"type": "string",
"description": "任务名称（add/update 时使用）"
},
"user": {
"type": "string",
"description": "所属用户 ID（add 时使用）"
},
"command": {
"type": "string",
"description": "要执行的 shell 命令（add/update 时使用）"
},
"interval": {
"type": "integer",
"description": "执行间隔（秒，add/update 时使用）"
},
"description": {
"type": "string",
"description": "任务描述（add/update 时使用，可选）"
},
"shell": {
"type": "string",
"description": "Shell 解释器，默认 /bin/bash（add/update 时使用，可选）"
},
"cwd": {
"type": "string",
"description": "工作目录（add/update 时使用，可选）"
},
"id": {
"type": "string",
"description": "任务 UUID（update/delete/toggle 时使用）"
},
"task_type": {
"type": "string",
"enum": ["model", "system"],
"description": "任务类型：model（模型任务，默认）或 system（系统任务，仅 add 时使用，可选）"
}
},
"required": ["action"]
}),
timeout_secs: 10,
}
}
/// 执行定时任务管理操作
pub async fn execute(args_json: &str) -> SkillResult {
let pool = match POOL.as_ref() {
Some(p) => p,
None => return SkillResult::error("数据库未配置，请设置 DATABASE_URL 环境变量"),
};
let params: serde_json::Value = serde_json::from_str(args_json).unwrap_or_default();
let action = params
.get("action")
.and_then(|v| v.as_str())
.unwrap_or("list");
match action {
"list" => cmd_list(pool).await,
"add" => cmd_add(pool, &params).await,
"update" => cmd_update(pool, &params).await,
"delete" => cmd_delete(pool, &params).await,
"toggle" => cmd_toggle(pool, &params).await,
_ => SkillResult::error(format!("未知操作: {}，支持: list, add, update, delete, toggle", action)),
}
fn pool() -> Option<&'static sqlx::PgPool> {
POOL.as_ref()
}
/// 检查任务是否为 system 类型（不可修改）
@@ -115,14 +42,32 @@ async fn check_system_task(pool: &sqlx::PgPool, id: &uuid::Uuid) -> Result<bool,
}
}
async fn cmd_list(pool: &sqlx::PgPool) -> SkillResult {
// ═══════════════════════════════════════════════
// 1. scheduled_task_list — 列出所有定时任务
// ═══════════════════════════════════════════════
pub fn list_spec() -> SkillSpec {
SkillSpec {
name: "scheduled_task_list".into(),
description: "列出所有定时任务（名称、ID、命令、间隔、状态、上次/下次执行时间）。系统任务标记 [系统] 标签。需要 PostgreSQL 数据库。".into(),
risk_level: RiskLevel::Low,
parameters: serde_json::json!({"type": "object", "properties": {}}),
timeout_secs: 10,
}
}
pub async fn list_execute(_params: serde_json::Value) -> SkillResult {
let pool = match pool() {
Some(p) => p,
None => return SkillResult::error("数据库未配置，请设置 DATABASE_URL 环境变量"),
};
let tasks = match crate::db::models::list_scheduled_tasks(pool).await {
Ok(t) => t,
Err(e) => return SkillResult::error(format!("查询失败: {}", e)),
};
if tasks.is_empty() {
return SkillResult::ok("暂无定时任务。你可以通过「添加定时任务」来创建，例如每天早8点播报天气。");
return SkillResult::ok("暂无定时任务。你可以通过 scheduled_task_add 来创建，例如每天早8点播报天气。");
}
let mut lines = vec![format!("📋 共有 {} 个定时任务:\n", tasks.len())];
@@ -158,25 +103,48 @@ async fn cmd_list(pool: &sqlx::PgPool) -> SkillResult {
SkillResult::ok(lines.join("\n"))
}
async fn cmd_add(pool: &sqlx::PgPool, params: &serde_json::Value) -> SkillResult {
// ═══════════════════════════════════════════════
// 2. scheduled_task_add — 添加定时任务
// ═══════════════════════════════════════════════
pub fn add_spec() -> SkillSpec {
SkillSpec {
name: "scheduled_task_add".into(),
description: "添加一个定时任务，按指定间隔自动执行 shell 命令。参数 name 为任务名称，user 为所属用户 ID，command 为要执行的命令，interval 为执行间隔（秒）。可选参数：description（描述）、shell（Shell解释器，默认/bin/bash）、cwd（工作目录）、task_type（任务类型，默认model）。需要 PostgreSQL 数据库。".into(),
risk_level: RiskLevel::Low,
parameters: serde_json::json!({
"type": "object",
"properties": {
"name": {"type": "string", "description": "任务名称（唯一）"},
"user": {"type": "string", "description": "所属用户 ID"},
"command": {"type": "string", "description": "要执行的 shell 命令"},
"interval": {"type": "integer", "description": "执行间隔（秒），如 86400 为每天"},
"description": {"type": "string", "description": "任务描述（可选）"},
"shell": {"type": "string", "description": "Shell 解释器，默认 /bin/bash（可选）"},
"cwd": {"type": "string", "description": "工作目录（可选）"},
"task_type": {"type": "string", "description": "任务类型：model（默认）或 system（可选）"}
},
"required": ["name", "user", "command", "interval"]
}),
timeout_secs: 10,
}
}
pub async fn add_execute(params: serde_json::Value) -> SkillResult {
let pool = match pool() {
Some(p) => p,
None => return SkillResult::error("数据库未配置，请设置 DATABASE_URL 环境变量"),
};
let name = params.get("name").and_then(|v| v.as_str()).unwrap_or("");
let user = params.get("user").and_then(|v| v.as_str()).unwrap_or("");
let command = params.get("command").and_then(|v| v.as_str()).unwrap_or("");
let interval = params.get("interval").and_then(|v| v.as_i64()).unwrap_or(0) as i32;
let task_type = params.get("task_type").and_then(|v| v.as_str()).unwrap_or("model");
if name.is_empty() {
return SkillResult::error("请提供 name（任务名称）");
}
if user.is_empty() {
return SkillResult::error("请提供 user（所属用户 ID）");
}
if command.is_empty() {
return SkillResult::error("请提供 command（要执行的命令）");
}
if interval <= 0 {
return SkillResult::error("请提供 interval（执行间隔，大于0秒）");
}
if name.is_empty() { return SkillResult::error("请提供 name（任务名称）"); }
if user.is_empty() { return SkillResult::error("请提供 user（所属用户 ID）"); }
if command.is_empty() { return SkillResult::error("请提供 command（要执行的命令）"); }
if interval <= 0 { return SkillResult::error("请提供 interval（执行间隔，大于0秒）"); }
if task_type != "model" && task_type != "system" {
return SkillResult::error("task_type 必须是 model 或 system");
}
@@ -194,7 +162,37 @@ async fn cmd_add(pool: &sqlx::PgPool, params: &serde_json::Value) -> SkillResult
}
}
async fn cmd_update(pool: &sqlx::PgPool, params: &serde_json::Value) -> SkillResult {
// ═══════════════════════════════════════════════
// 3. scheduled_task_update — 更新定时任务
// ═══════════════════════════════════════════════
pub fn update_spec() -> SkillSpec {
SkillSpec {
name: "scheduled_task_update".into(),
description: "更新定时任务的字段（只更新传入的字段）。参数 id 为任务 UUID（必填），可选字段：name/command/interval/description/shell/cwd。系统任务不可修改。需要 PostgreSQL 数据库。".into(),
risk_level: RiskLevel::Low,
parameters: serde_json::json!({
"type": "object",
"properties": {
"id": {"type": "string", "description": "任务 UUID"},
"name": {"type": "string", "description": "新任务名称（可选）"},
"command": {"type": "string", "description": "新命令（可选）"},
"interval": {"type": "integer", "description": "新执行间隔秒数（可选）"},
"description": {"type": "string", "description": "新描述（可选）"},
"shell": {"type": "string", "description": "新 Shell 解释器（可选）"},
"cwd": {"type": "string", "description": "新工作目录（可选）"}
},
"required": ["id"]
}),
timeout_secs: 10,
}
}
pub async fn update_execute(params: serde_json::Value) -> SkillResult {
let pool = match pool() {
Some(p) => p,
None => return SkillResult::error("数据库未配置，请设置 DATABASE_URL 环境变量"),
};
let id_str = params.get("id").and_then(|v| v.as_str()).unwrap_or("");
let id = match uuid::Uuid::parse_str(id_str) {
Ok(u) => u,
@@ -204,7 +202,7 @@ async fn cmd_update(pool: &sqlx::PgPool, params: &serde_json::Value) -> SkillRes
// 检查是否为系统任务
match check_system_task(pool, &id).await {
Ok(true) => return SkillResult::error("❌ 系统任务不可修改。系统任务由管理员配置，如需变更请联系管理员。"),
Ok(false) => {} // model 任务，可以修改
Ok(false) => {}
Err(e) => return SkillResult::error(e),
}
@@ -221,28 +219,47 @@ async fn cmd_update(pool: &sqlx::PgPool, params: &serde_json::Value) -> SkillRes
return SkillResult::error("请至少提供要更新的字段（name/command/interval/description/shell/cwd）");
}
match crate::db::models::update_scheduled_task(
pool, &id, name, command, interval, description, shell, cwd,
)
.await
{
match crate::db::models::update_scheduled_task(pool, &id, name, command, interval, description, shell, cwd).await {
Ok(true) => SkillResult::ok("✅ 定时任务已更新"),
Ok(false) => SkillResult::error("未找到该任务或无字段更新"),
Err(e) => SkillResult::error(format!("更新失败: {}", e)),
}
}
async fn cmd_delete(pool: &sqlx::PgPool, params: &serde_json::Value) -> SkillResult {
// ═══════════════════════════════════════════════
// 4. scheduled_task_delete — 删除定时任务
// ═══════════════════════════════════════════════
pub fn delete_spec() -> SkillSpec {
SkillSpec {
name: "scheduled_task_delete".into(),
description: "删除指定 UUID 的定时任务。系统任务不可删除。需要 PostgreSQL 数据库。".into(),
risk_level: RiskLevel::Low,
parameters: serde_json::json!({
"type": "object",
"properties": {
"id": {"type": "string", "description": "要删除的任务 UUID"}
},
"required": ["id"]
}),
timeout_secs: 10,
}
}
pub async fn delete_execute(params: serde_json::Value) -> SkillResult {
let pool = match pool() {
Some(p) => p,
None => return SkillResult::error("数据库未配置，请设置 DATABASE_URL 环境变量"),
};
let id_str = params.get("id").and_then(|v| v.as_str()).unwrap_or("");
let id = match uuid::Uuid::parse_str(id_str) {
Ok(u) => u,
Err(_) => return SkillResult::error("请提供有效的 id（任务 UUID）"),
};
// 检查是否为系统任务
match check_system_task(pool, &id).await {
Ok(true) => return SkillResult::error("❌ 系统任务不可删除。系统任务由管理员配置，如需变更请联系管理员。"),
Ok(false) => {} // model 任务，可以删除
Ok(false) => {}
Err(e) => return SkillResult::error(e),
}
@@ -253,17 +270,40 @@ async fn cmd_delete(pool: &sqlx::PgPool, params: &serde_json::Value) -> SkillRes
}
}
async fn cmd_toggle(pool: &sqlx::PgPool, params: &serde_json::Value) -> SkillResult {
// ═══════════════════════════════════════════════
// 5. scheduled_task_toggle — 启用/禁用定时任务
// ═══════════════════════════════════════════════
pub fn toggle_spec() -> SkillSpec {
SkillSpec {
name: "scheduled_task_toggle".into(),
description: "切换定时任务的启用/禁用状态。系统任务不可切换。需要 PostgreSQL 数据库。".into(),
risk_level: RiskLevel::Low,
parameters: serde_json::json!({
"type": "object",
"properties": {
"id": {"type": "string", "description": "任务 UUID"}
},
"required": ["id"]
}),
timeout_secs: 10,
}
}
pub async fn toggle_execute(params: serde_json::Value) -> SkillResult {
let pool = match pool() {
Some(p) => p,
None => return SkillResult::error("数据库未配置，请设置 DATABASE_URL 环境变量"),
};
let id_str = params.get("id").and_then(|v| v.as_str()).unwrap_or("");
let id = match uuid::Uuid::parse_str(id_str) {
Ok(u) => u,
Err(_) => return SkillResult::error("请提供有效的 id（任务 UUID）"),
};
// 检查是否为系统任务
match check_system_task(pool, &id).await {
Ok(true) => return SkillResult::error("❌ 系统任务不可启用/禁用。系统任务由管理员配置，如需变更请联系管理员。"),
Ok(false) => {} // model 任务，可以切换
Ok(false) => {}
Err(e) => return SkillResult::error(e),
}
src/tools/builtins/weather.rs
+137 -204
View File
@@ -42,23 +42,20 @@ fn base64url(data: &[u8]) -> String {
}
fn generate_jwt() -> Result<String, String> {
let key_file = std::env::var("QWEATHER_JWT_PRIVATE_KEY_FILE")
.unwrap_or_else(|_| "qweather/ed25519-private.pem".into());
let key_file = std::env::var("QWEATHER_JWT_PRIVATE_KEY_FILE").unwrap_or_else(|_| {
dirs::home_dir()
.unwrap_or_else(|| std::path::PathBuf::from("."))
.join(".ias")
.join("qweather")
.join("ed25519-private.pem")
.to_string_lossy()
.to_string()
});
let pem = std::fs::read_to_string(&key_file)
.map_err(|e| format!("读取密钥文件 {} 失败: {}", key_file, e))?;
// 手动解析 PEM: 去掉头尾，base64 解码得到 DER，再用 from_pkcs8_der
let der = pem
.lines()
.filter(|l| !l.starts_with("-----"))
.collect::<Vec<_>>()
.join("");
let der_bytes = general_purpose::STANDARD
.decode(&der)
.map_err(|e| format!("Base64 解码密钥失败: {e}"))?;
let signing_key = SigningKey::from_pkcs8_der(&der_bytes)
let signing_key = SigningKey::from_pkcs8_pem(&pem)
.map_err(|e| format!("解析 Ed25519 密钥失败: {e}"))?;
let key_id = std::env::var("QWEATHER_JWT_KEY_ID").unwrap_or_default();
@@ -426,200 +423,153 @@ impl QWeatherClient {
}
// ═══════════════════════════════════════════════
// 公开接口
// 公开接口 — 3 个独立工具
// ═══════════════════════════════════════════════
pub fn spec() -> SkillSpec {
/// 1. weather_now — 实时天气
pub fn now_spec() -> SkillSpec {
SkillSpec {
name: "query_weather".into(),
description:
"查询指定城市的天气信息，包括实时天气、未来几天预报、逐小时预报。参数 location 为城市名"
.into(),
name: "weather_now".into(),
description: "查询指定城市的实时天气（温度、体感温度、天气状况、风向风力、湿度、降水量、气压、能见度）。参数 location 为城市名".into(),
risk_level: RiskLevel::Low,
parameters: serde_json::json!({
"type": "object",
"properties": {
"location": {"type": "string", "description": "城市名，如 北京、上海"},
"days": {"type": "integer", "description": "预报天数，默认 3，最大 30", "default": 3},
"hours": {"type": "integer", "description": "逐时预报小时数，默认 0（不查询），最大 168", "default": 0}
"location": {"type": "string", "description": "城市名，如 北京、上海"}
},
"required": ["location"]
}),
timeout_secs: 15,
timeout_secs: 10,
}
}
fn extract_params(params: &serde_json::Value) -> (&str, u32, u32) {
pub async fn now_execute(params: serde_json::Value) -> SkillResult {
let location = params
.get("location")
.and_then(|v| v.as_str())
.unwrap_or("合肥");
let days = params
.get("days")
.and_then(|v| v.as_u64())
.unwrap_or(3)
.min(30) as u32;
let hours = params
.get("hours")
.and_then(|v| v.as_u64())
.unwrap_or(0)
.min(168) as u32;
(location, days.max(1), hours)
.unwrap_or("");
if location.is_empty() {
return SkillResult::error("请提供 location（城市名）");
}
#[allow(dead_code)]
fn extract_days_from_prompt(prompt: &str) -> u32 {
// 匹配模式: "7天" "3日" "未来7" "7d" "7D"
for pat in &["天", "日", "d", "D"] {
if let Some(pos) = prompt.find(pat) {
// 向前找数字
let before = &prompt[..pos];
if let Some(num_str) = before
.chars()
.rev()
.take_while(|c| c.is_ascii_digit())
.collect::<String>()
.chars()
.rev()
.collect::<String>()
.into()
&& let Ok(n) = num_str.parse::<u32>() {
return n.min(30).max(1);
}
}
}
3 // 默认
}
#[allow(dead_code)]
fn extract_city_from_prompt(prompt: &str) -> &str {
// 去掉常见前缀，取第一个2-3字中文词。GeoAPI 支持模糊搜索，
// 即使多取一两个字也能正确匹配到城市
let chars: Vec<char> = prompt.chars().collect();
let byte_pos = |idx: usize| -> usize { chars[..idx].iter().map(|c| c.len_utf8()).sum() };
// 跳过前缀
let mut i = 0;
for prefix in &[
"帮我查一下",
"帮我看看",
"帮我查",
"帮我",
"查一下",
"查询",
"看看",
] {
if prompt.starts_with(prefix) {
i = prefix.chars().count();
while i < chars.len() && chars[i] as u32 <= 0x4e00 {
i += 1;
}
break;
}
}
// 跳过非中文
while i < chars.len() && chars[i] as u32 <= 0x4e00 {
i += 1;
}
// 取至多4个连续中文
let start = i;
while i < chars.len() && chars[i] as u32 > 0x4e00 && i - start < 4 {
i += 1;
}
if i - start >= 2 {
&prompt[byte_pos(start)..byte_pos(i)]
} else {
"北京"
}
}
pub async fn execute(params: serde_json::Value) -> SkillResult {
let (location, days, hours) = extract_params(&params);
tracing::info!(
"🌤️ 查询天气: location={}, days={}, hours={}",
location,
days,
hours
);
let client = match QWeatherClient::new().await {
Ok(c) => c,
Err(e) => return SkillResult::error(format!("初始化天气客户端失败: {}", e)),
};
// 先查城市 ID
let city_id = match client.lookup_city(location).await {
Ok(id) => id,
Err(e) => return SkillResult::error(e),
};
// 并发查询
let now_fut = client.weather_now(&city_id);
let daily_fut = client.daily_forecast(&city_id, days);
match client.weather_now(&city_id).await {
Ok(data) => {
let output = serde_json::json!({"location": location, "now": data});
SkillResult::ok_with_data(serde_json::to_string(&output).unwrap_or_default(), output)
}
Err(e) => SkillResult::error(e),
}
}
// hours == 0 时跳过逐时预报，节省 API 调用
let (now_result, daily_result, hourly_result) = if hours > 0 {
let hourly_fut = client.hourly_forecast(&city_id, hours);
let (n, d, h) = tokio::join!(now_fut, daily_fut, hourly_fut);
(n, d, h)
} else {
let (n, d) = tokio::join!(now_fut, daily_fut);
(n, d, Err("hours=0, skipped".into()))
/// 2. weather_forecast — 每日预报
pub fn forecast_spec() -> SkillSpec {
SkillSpec {
name: "weather_forecast".into(),
description: "查询指定城市未来几天的天气预报（最高/最低温度、白天/夜间天气状况、风向风力、湿度、紫外线指数）。参数 location 为城市名，days 为预报天数（默认3，最大30）".into(),
risk_level: RiskLevel::Low,
parameters: serde_json::json!({
"type": "object",
"properties": {
"location": {"type": "string", "description": "城市名，如 北京、上海"},
"days": {"type": "integer", "description": "预报天数，默认3，最大30"}
},
"required": ["location"]
}),
timeout_secs: 10,
}
}
pub async fn forecast_execute(params: serde_json::Value) -> SkillResult {
let location = params
.get("location")
.and_then(|v| v.as_str())
.unwrap_or("");
if location.is_empty() {
return SkillResult::error("请提供 location（城市名）");
}
let days = params
.get("days")
.and_then(|v| v.as_u64())
.unwrap_or(3)
.min(30)
.max(1) as u32;
let client = match QWeatherClient::new().await {
Ok(c) => c,
Err(e) => return SkillResult::error(format!("初始化天气客户端失败: {}", e)),
};
let city_id = match client.lookup_city(location).await {
Ok(id) => id,
Err(e) => return SkillResult::error(e),
};
let now = match now_result {
Ok(n) => n,
Err(e) => {
tracing::warn!("实时天气查询失败: {}", e);
serde_json::Value::Null
match client.daily_forecast(&city_id, days).await {
Ok(data) => {
let output = serde_json::json!({"location": location, "forecast": data});
SkillResult::ok_with_data(serde_json::to_string(&output).unwrap_or_default(), output)
}
Err(e) => SkillResult::error(e),
}
}
/// 3. weather_hourly — 逐时预报
pub fn hourly_spec() -> SkillSpec {
SkillSpec {
name: "weather_hourly".into(),
description: "查询指定城市未来逐小时的天气预报（温度、天气状况、风向风力、湿度、降水概率、降水量）。参数 location 为城市名，hours 为预报小时数（默认24，最大168）".into(),
risk_level: RiskLevel::Low,
parameters: serde_json::json!({
"type": "object",
"properties": {
"location": {"type": "string", "description": "城市名，如 北京、上海"},
"hours": {"type": "integer", "description": "预报小时数，默认24，最大168"}
},
"required": ["location"]
}),
timeout_secs: 10,
}
}
pub async fn hourly_execute(params: serde_json::Value) -> SkillResult {
let location = params
.get("location")
.and_then(|v| v.as_str())
.unwrap_or("");
if location.is_empty() {
return SkillResult::error("请提供 location（城市名）");
}
let hours = params
.get("hours")
.and_then(|v| v.as_u64())
.unwrap_or(24)
.min(168)
.max(1) as u32;
let client = match QWeatherClient::new().await {
Ok(c) => c,
Err(e) => return SkillResult::error(format!("初始化天气客户端失败: {}", e)),
};
let city_id = match client.lookup_city(location).await {
Ok(id) => id,
Err(e) => return SkillResult::error(e),
};
let daily = match daily_result {
Ok(d) => d,
Err(e) => {
tracing::warn!("每日预报查询失败: {}", e);
serde_json::Value::Null
match client.hourly_forecast(&city_id, hours).await {
Ok(data) => {
let output = serde_json::json!({"location": location, "hourly": data});
SkillResult::ok_with_data(serde_json::to_string(&output).unwrap_or_default(), output)
}
};
let hourly = match hourly_result {
Ok(h) => h,
Err(e) => {
tracing::warn!("逐时预报查询失败: {}", e);
serde_json::Value::Null
Err(e) => SkillResult::error(e),
}
};
let output = serde_json::json!({
"location": location,
"now": now,
"forecast": daily,
"hourly": hourly,
});
// 过滤 null 字段
let mut output = output;
if let Some(obj) = output.as_object_mut() {
obj.retain(|_, v| !v.is_null());
// 过滤 forecast/hourly 中每个元素内的空值
for key in &["forecast", "hourly"] {
if let Some(arr) = obj.get_mut(*key).and_then(|v| v.as_array_mut()) {
for item in arr.iter_mut() {
if let Some(m) = item.as_object_mut() {
m.retain(|_, v| !v.is_null() && v.as_str().is_none_or(|s| !s.is_empty()));
}
}
}
}
}
let text = serde_json::to_string(&output).unwrap_or_default();
SkillResult::ok_with_data(text, output)
}
#[cfg(test)]
@@ -640,48 +590,31 @@ mod tests {
}
}
#[test]
fn test_extract_params() {
let params = serde_json::json!({"location": "上海", "days": 5});
let (loc, days, hours) = extract_params(&params);
assert_eq!(loc, "上海");
assert_eq!(days, 5);
assert_eq!(hours, 0);
}
#[test]
fn test_extract_city_from_prompt() {
// 取前缀后的连续中文给 GeoAPI 模糊搜索（多取不影响结果）
assert_eq!(
extract_city_from_prompt("查询合肥未来7天的天气预报"),
"合肥未来"
);
assert_eq!(extract_city_from_prompt("北京今天天气"), "北京今天");
assert_eq!(extract_city_from_prompt("上海"), "上海");
assert_eq!(extract_city_from_prompt("帮我查一下广州"), "广州");
}
#[test]
fn test_extract_days() {
assert_eq!(extract_days_from_prompt("查询合肥未来7天的天气预报"), 7);
assert_eq!(extract_days_from_prompt("3天预报"), 3);
assert_eq!(extract_days_from_prompt("北京天气"), 3);
}
/// 端到端测试：需要有效的 QWeather 环境变量
#[tokio::test]
async fn test_execute_weather_integration() {
// 跳过无环境变量的情况
async fn test_now_execute_integration() {
if std::env::var("QWEATHER_JWT_KEY_ID").is_err() {
eprintln!("跳过集成测试: 未设置 QWeather 环境变量");
return;
}
let result = super::execute(serde_json::json!({"location": "合肥", "days": 3})).await;
let result = super::now_execute(serde_json::json!({"location": "合肥"})).await;
assert!(result.success, "查询失败: {}", result.output);
let data = result.data.expect("应有 data 字段");
assert_eq!(data["location"], "合肥");
assert!(data["now"].is_object(), "应有实时天气");
}
#[tokio::test]
async fn test_forecast_execute_integration() {
if std::env::var("QWEATHER_JWT_KEY_ID").is_err() {
eprintln!("跳过集成测试: 未设置 QWeather 环境变量");
return;
}
let result =
super::forecast_execute(serde_json::json!({"location": "合肥", "days": 3})).await;
assert!(result.success, "查询失败: {}", result.output);
let data = result.data.expect("应有 data 字段");
assert_eq!(data["location"], "合肥");
assert!(data["forecast"].is_array(), "应有预报");
}
}
src/tools/mod.rs
+131 -34
View File
@@ -17,8 +17,8 @@
//!
//! ### 工具执行流程
//! ```text
//! LLM 请求 call_capability("query_weather", {"location":"北京"})
//! → daemon 解析出 target_name="query_weather"，提取参数
//! LLM 请求 call_capability("weather_now", {"location":"北京"})
//! → daemon 解析出 target_name="weather_now"，提取参数
//! → 高风险？→ 走 ApprovalManager 审批流
//! → 低风险？→ ToolRegistry::execute()
//! → 结果返回 LLM 继续对话
@@ -113,23 +113,59 @@ fn build_registry() -> ToolRegistry {
}),
)));
// ── ApiTool: 备忘录（文件 I/O） ──
let memos_spec = self::builtins::memos::spec();
// ── ApiTool: 备忘录 — 列出（文件 I/O） ──
registry.register(Box::new(ApiTool::new(
memos_spec,
self::builtins::memos::list_spec(),
std::sync::Arc::new(|params| {
Box::pin(async move {
let args = serde_json::to_string(&params).unwrap_or_default();
self::builtins::memos::execute(&args).await
})
Box::pin(async { self::builtins::memos::list_execute(params).await })
}),
)));
// ── ApiTool: 天气（HTTP API） ──
// ── ApiTool: 备忘录 — 添加（文件 I/O） ──
registry.register(Box::new(ApiTool::new(
self::builtins::weather::spec(),
self::builtins::memos::add_spec(),
std::sync::Arc::new(|params| {
Box::pin(async { self::builtins::weather::execute(params).await })
Box::pin(async { self::builtins::memos::add_execute(params).await })
}),
)));
// ── ApiTool: 备忘录 — 查看（文件 I/O） ──
registry.register(Box::new(ApiTool::new(
self::builtins::memos::get_spec(),
std::sync::Arc::new(|params| {
Box::pin(async { self::builtins::memos::get_execute(params).await })
}),
)));
// ── ApiTool: 备忘录 — 删除（文件 I/O） ──
registry.register(Box::new(ApiTool::new(
self::builtins::memos::delete_spec(),
std::sync::Arc::new(|params| {
Box::pin(async { self::builtins::memos::delete_execute(params).await })
}),
)));
// ── ApiTool: 天气 — 实时天气（HTTP API） ──
registry.register(Box::new(ApiTool::new(
self::builtins::weather::now_spec(),
std::sync::Arc::new(|params| {
Box::pin(async { self::builtins::weather::now_execute(params).await })
}),
)));
// ── ApiTool: 天气 — 每日预报（HTTP API） ──
registry.register(Box::new(ApiTool::new(
self::builtins::weather::forecast_spec(),
std::sync::Arc::new(|params| {
Box::pin(async { self::builtins::weather::forecast_execute(params).await })
}),
)));
// ── ApiTool: 天气 — 逐时预报（HTTP API） ──
registry.register(Box::new(ApiTool::new(
self::builtins::weather::hourly_spec(),
std::sync::Arc::new(|params| {
Box::pin(async { self::builtins::weather::hourly_execute(params).await })
}),
)));
@@ -149,30 +185,91 @@ fn build_registry() -> ToolRegistry {
}),
)));
// ── ApiTool: 高德地图系列（HTTP API） ──
for amap_spec in self::builtins::amap::specs() {
let name = amap_spec.name.clone();
// ── ApiTool: 高德地图 — POI搜索 ──
registry.register(Box::new(ApiTool::new(
amap_spec,
std::sync::Arc::new(move |params| {
let n = name.clone();
Box::pin(async move {
self::builtins::amap::execute(&n, params)
.await
.unwrap_or_else(|| SkillResult::error("执行失败"))
})
self::builtins::amap::poi_search_spec(),
std::sync::Arc::new(|params| {
Box::pin(async { self::builtins::amap::poi_search_execute(params).await })
}),
)));
}
// ── ApiTool: 定时任务管理（PostgreSQL） ──
// ── ApiTool: 高德地图 — 地理编码 ──
registry.register(Box::new(ApiTool::new(
self::builtins::scheduled_task::spec(),
self::builtins::amap::geocode_spec(),
std::sync::Arc::new(|params| {
Box::pin(async move {
let args = serde_json::to_string(&params).unwrap_or_default();
self::builtins::scheduled_task::execute(&args).await
})
Box::pin(async { self::builtins::amap::geocode_execute(params).await })
}),
)));
// ── ApiTool: 高德地图 — 逆地理编码 ──
registry.register(Box::new(ApiTool::new(
self::builtins::amap::reverse_geocode_spec(),
std::sync::Arc::new(|params| {
Box::pin(async { self::builtins::amap::reverse_geocode_execute(params).await })
}),
)));
// ── ApiTool: 高德地图 — 路径规划 ──
registry.register(Box::new(ApiTool::new(
self::builtins::amap::route_plan_spec(),
std::sync::Arc::new(|params| {
Box::pin(async { self::builtins::amap::route_plan_execute(params).await })
}),
)));
// ── ApiTool: 高德地图 — 旅游规划 ──
registry.register(Box::new(ApiTool::new(
self::builtins::amap::travel_plan_spec(),
std::sync::Arc::new(|params| {
Box::pin(async { self::builtins::amap::travel_plan_execute(params).await })
}),
)));
// ── ApiTool: 高德地图 — 地图链接 ──
registry.register(Box::new(ApiTool::new(
self::builtins::amap::map_link_spec(),
std::sync::Arc::new(|params| {
Box::pin(async { self::builtins::amap::map_link_execute(params).await })
}),
)));
// ── ApiTool: 定时任务 — 列出（PostgreSQL） ──
registry.register(Box::new(ApiTool::new(
self::builtins::scheduled_task::list_spec(),
std::sync::Arc::new(|params| {
Box::pin(async { self::builtins::scheduled_task::list_execute(params).await })
}),
)));
// ── ApiTool: 定时任务 — 添加（PostgreSQL） ──
registry.register(Box::new(ApiTool::new(
self::builtins::scheduled_task::add_spec(),
std::sync::Arc::new(|params| {
Box::pin(async { self::builtins::scheduled_task::add_execute(params).await })
}),
)));
// ── ApiTool: 定时任务 — 更新（PostgreSQL） ──
registry.register(Box::new(ApiTool::new(
self::builtins::scheduled_task::update_spec(),
std::sync::Arc::new(|params| {
Box::pin(async { self::builtins::scheduled_task::update_execute(params).await })
}),
)));
// ── ApiTool: 定时任务 — 删除（PostgreSQL） ──
registry.register(Box::new(ApiTool::new(
self::builtins::scheduled_task::delete_spec(),
std::sync::Arc::new(|params| {
Box::pin(async { self::builtins::scheduled_task::delete_execute(params).await })
}),
)));
// ── ApiTool: 定时任务 — 切换状态（PostgreSQL） ──
registry.register(Box::new(ApiTool::new(
self::builtins::scheduled_task::toggle_spec(),
std::sync::Arc::new(|params| {
Box::pin(async { self::builtins::scheduled_task::toggle_execute(params).await })
}),
)));
@@ -222,22 +319,22 @@ pub fn build_capability_guide(specs: &[SkillSpec]) -> String {
lines.push(String::new());
lines.push("## 🌤️ 天气工具 — 天气相关查询：".to_string());
for spec in specs.iter().filter(|s| s.name == "query_weather") {
for spec in specs.iter().filter(|s| s.name.starts_with("weather_")) {
format_spec(&mut lines, spec);
}
lines.push(String::new());
lines.push(" ▸ 任何天气相关问题（实时天气、预报、逐时预报）→ 使用 query_weather".to_string());
lines.push(" ▸ 实时天气 → weather_now | 每日预报 → weather_forecast | 逐时预报 → weather_hourly".to_string());
lines.push(String::new());
lines.push("## 🔍 其他工具：".to_string());
for spec in specs.iter().filter(|s| !s.name.starts_with("amap_") && s.name != "query_weather") {
for spec in specs.iter().filter(|s| !s.name.starts_with("amap_") && !s.name.starts_with("weather_")) {
format_spec(&mut lines, spec);
}
lines.push(String::new());
lines.push("## ⚡ 工具选择优先级：".to_string());
lines.push(" 1. 旅游/路线/地点/地址查询 → amap_* 系列".to_string());
lines.push(" 2. 天气查询 → query_weather".to_string());
lines.push(" 2. 天气查询 → weather_now / weather_forecast / weather_hourly".to_string());
lines.push(" 3. 联网搜索/网页抓取 → web_search / fetch_page".to_string());
lines.push(" 4. 以上都不匹配的需求 → web_search 搜索最新信息".to_string());
lines.push(String::new());
src/tools/types.rs
+4 -17
View File
@@ -87,23 +87,6 @@ impl SkillResult {
}
}
pub fn rejected(skill_name: &str) -> Self {
Self {
success: false,
output: format!("用户拒绝了你的调用：{}", skill_name),
data: None,
}
}
#[allow(dead_code)]
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,
@@ -124,6 +107,7 @@ pub type WechatSender = Arc<
dyn Fn(&str, &str) -> Pin<Box<dyn Future<Output = Result<(), String>> + Send>> + Send + Sync,
>;
#[allow(dead_code)]
pub struct ExecutionContext {
pub user_id: String,
pub approval_manager: Option<Arc<crate::tools::approval::ApprovalManager>>,
@@ -131,6 +115,7 @@ pub struct ExecutionContext {
}
impl ExecutionContext {
#[allow(dead_code)]
pub fn new(user_id: impl Into<String>) -> Self {
Self {
user_id: user_id.into(),
@@ -139,11 +124,13 @@ impl ExecutionContext {
}
}
#[allow(dead_code)]
pub fn with_approval(mut self, mgr: Arc<crate::tools::approval::ApprovalManager>) -> Self {
self.approval_manager = Some(mgr);
self
}
#[allow(dead_code)]
pub fn with_wechat(mut self, send: WechatSender) -> Self {
self.send_wechat = Some(send);
self
src/wechat/client.rs
+2
View File
@@ -284,6 +284,7 @@ impl WeChatClient {
}
/// 注销监听器
#[allow(dead_code)]
pub async fn notify_stop(&self) -> Result<(), String> {
let body = serde_json::json!({ "base_info": BaseInfo::new() });
@@ -398,6 +399,7 @@ impl WeChatClient {
}
/// 获取当前 updates_buf（用于持久化）
#[allow(dead_code)]
pub async fn get_updates_buf(&self) -> String {
self.updates_buf.lock().await.clone()
}
src/worker.rs
-300
View File
@@ -1,300 +0,0 @@
//! ## Worker —— 无状态执行进程（已废弃）
//!
//! ### 旧架构说明
//! 在旧架构中，每条消息会 spawn 一个独立的 Worker 进程，
//! Worker 通过 UDS 连接到 daemon，处理完一条消息后退出。
//!
//! 优势：Worker 代码热更新（编译后下一条消息自动用新版本）
//! 劣势：进程启动开销、IPC 序列化开销、状态管理复杂
//!
//! ### 新架构
//! 现在改为 daemon 内的 LLM/Tool/Send 三个消费者 tokio task。
//! 此模块保留供编译参考，不再被 daemon 使用。
//! 1. 连接 daemon 的 Unix Domain Socket
//! 2. 读取 TaskFrame (用户消息 + 上下文 + 环境变量)
//! 3. 注入环境变量，构建 Conversation，执行 LLM 对话 + 工具调用
//! 4. 逐帧输出结果 (usage, reply, bye)
//! 5. 高风险工具遇审批 → 发送 RequestApproval → 退出，由 daemon 重新 spawn
use crate::ipc::{OutputFrame, TaskFrame, recv_task, send_frame};
use crate::llm::{Conversation, ConversationConfig, DEFAULT_SYSTEM_PROMPT};
use crate::tools::types::SkillResult;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::net::UnixStream;
use tracing::info;
/// Worker 与工具执行器之间的共享状态
struct WorkerShared {
/// 待发送的审批请求 (tool, reason)
pending_approval: tokio::sync::Mutex<Option<(String, String)>>,
/// 本次对话中已由用户批准的工具名（避免重复审批）
approved_tool: tokio::sync::Mutex<Option<String>>,
/// 本地记忆缓存
memories: tokio::sync::Mutex<HashMap<String, Vec<String>>>,
/// 本次写入的新记忆（退出前发送 DbWrite 帧同步到 daemon）
new_memories: tokio::sync::Mutex<Vec<String>>,
/// 预载的摘要
summaries: Vec<String>,
/// 用户 ID
user_id: String,
}
/// Worker 入口: 连接 daemon，处理一条消息后退出
pub async fn run(sock_path: &str) -> Result<(), String> {
let mut stream = UnixStream::connect(sock_path)
.await
.map_err(|e| format!("连接 daemon 失败 ({}): {}", sock_path, e))?;
info!("Worker 已连接 daemon");
// 1. 读 task
let task: TaskFrame = recv_task(&mut stream).await?;
info!(
"Worker 收到任务: user={} msg={:.60}",
task.user_id, task.msg.text
);
// 2. 注入环境变量
for (k, v) in &task.env {
unsafe {
std::env::set_var(k, v);
}
}
// 3. 构建 ConversationConfig
let system_prompt = if task.system_prompt.is_empty() {
DEFAULT_SYSTEM_PROMPT.to_string()
} else {
task.system_prompt.clone()
};
let model = if task.model.is_empty() {
std::env::var("DEEPSEEK_MODEL").unwrap_or_else(|_| "deepseek-v4-flash".to_string())
} else {
task.model.clone()
};
let cfg = ConversationConfig {
system_prompt,
model,
tools: Some(task.tools.clone()),
..Default::default()
};
// 4. 构建 Conversation 并加载上下文
let mut conv =
Conversation::new(cfg).map_err(|e| format!("初始化 Conversation 失败: {}", e))?;
info!("Worker LLM: {}/{}", conv.provider_name(), conv.model());
{
let session = conv.session();
let mut s = session.lock().await;
s.load_from_history(&task.history);
s.current_user_id = task.user_id.clone();
}
// 5. 共享状态
let approved_tool = task.approved_tool.clone().filter(|t| !t.is_empty());
let shared = Arc::new(WorkerShared {
pending_approval: tokio::sync::Mutex::new(None),
approved_tool: tokio::sync::Mutex::new(approved_tool),
memories: tokio::sync::Mutex::new(HashMap::from([(
task.user_id.clone(),
task.memories.clone(),
)])),
new_memories: tokio::sync::Mutex::new(Vec::new()),
summaries: task.summaries.clone(),
user_id: task.user_id.clone(),
});
// 6. 构建工具执行器
let executor = build_worker_executor(conv.session(), shared.clone());
conv.set_tool_executor(executor);
// 7. 执行 LLM 对话（带审批检测）
use crate::llm::ChatResult;
let ChatResult { reply, used_tools: _used_tools, usage, .. } = conv
.chat_with_tools(task.msg.text.clone())
.await
.map_err(|e| format!("LLM 对话失败: {}", e))?;
// 8. 检查是否有审批请求（由工具执行器设置）
let approval_req = shared.pending_approval.lock().await.take();
if let Some((tool, reason)) = approval_req {
info!("Worker 发送审批请求: tool={}", tool);
let _ = send_frame(&mut stream, &OutputFrame::RequestApproval { tool, reason }).await;
let _ = send_frame(&mut stream, &OutputFrame::Bye).await;
return Ok(());
}
// 9. 发送结果帧
// 先同步新写入的记忆到 daemon
let new_mems = shared.new_memories.lock().await;
for content in new_mems.iter() {
let _ = send_frame(
&mut stream,
&OutputFrame::DbWrite {
table: "user_memories".into(),
row: serde_json::json!({"content": content}),
},
)
.await;
}
drop(new_mems);
if let Some(u) = &usage {
let _ = send_frame(
&mut stream,
&OutputFrame::UsageReport {
model: conv.model().to_string(),
provider: conv.provider_name().to_string(),
prompt_tokens: u.prompt_tokens,
completion_tokens: u.completion_tokens,
total_tokens: u.total_tokens,
cache_hit_tokens: u.prompt_cache_hit_tokens,
cache_miss_tokens: u.prompt_cache_miss_tokens,
user_id: task.user_id.clone(),
},
)
.await;
}
let _ = send_frame(
&mut stream,
&OutputFrame::Reply {
text: reply.clone(),
},
)
.await;
let _ = send_frame(&mut stream, &OutputFrame::Bye).await;
info!(
"Worker 完成: user={} reply_len={}",
task.user_id,
reply.len()
);
Ok(())
}
/// 构建 Worker 侧的工具执行器（无 DB 连接）
fn build_worker_executor(
_session: Arc<tokio::sync::Mutex<crate::context::types::ChatSession>>,
shared: Arc<WorkerShared>,
) -> crate::llm::conversation::ToolExecutor {
Arc::new(move |name: &str, args_json: &str| {
let shared = shared.clone();
let n = name.to_string();
let args = args_json.to_string();
Box::pin(async move {
match n.as_str() {
"read_memories" => {
let cache = shared.memories.lock().await;
let mems = cache.get(&shared.user_id);
match mems {
Some(m) if !m.is_empty() => {
let lines: Vec<String> = m
.iter()
.enumerate()
.map(|(i, v)| format!("{}. {}", i + 1, v))
.collect();
Ok(lines.join("\n"))
}
_ => Ok("暂无长期记忆".to_string()),
}
}
"write_memory" => {
let params: serde_json::Value = serde_json::from_str(&args).unwrap_or_default();
let content = params.get("content").and_then(|v| v.as_str()).unwrap_or("");
if content.is_empty() {
return Ok("请提供 content 参数".to_string());
}
shared
.memories
.lock()
.await
.entry(shared.user_id.clone())
.or_default()
.push(content.to_string());
shared.new_memories.lock().await.push(content.to_string());
Ok(format!("已添加长期记忆: {}", content))
}
"read_summaries" => {
if shared.summaries.is_empty() {
Ok("暂无历史摘要".to_string())
} else {
Ok(shared
.summaries
.iter()
.enumerate()
.map(|(i, s)| format!("{}. {}", i + 1, s))
.collect::<Vec<_>>()
.join("\n---\n"))
}
}
"query_capabilities" => {
let specs = crate::tools::specs();
Ok(crate::tools::build_capability_guide(&specs))
}
// call_capability 或直接调用
_ => {
let (target_name, target_args) = if n == "call_capability" {
let cp: serde_json::Value = serde_json::from_str(&args).unwrap_or_default();
let name = cp
.get("name")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string();
let unpacked = crate::tools::unpack_call_params(&cp);
(name, serde_json::to_string(&unpacked).unwrap_or_default())
} else {
(n.clone(), args.clone())
};
if target_name.is_empty() {
return Ok(format!("未知工具: {}", n));
}
// 高风险工具 → 请求审批（除非已被用户批准）
if crate::tools::is_high_risk(&target_name) {
let is_approved = shared.approved_tool.lock().await.as_deref() == Some(&target_name);
if !is_approved {
let reason = format!("LLM 请求执行高风险工具: {}", target_name);
*shared.pending_approval.lock().await = Some((target_name.clone(), reason));
return Ok(SkillResult::rejected(&target_name).output);
}
// 批准过 → 执行一次后清除（同一对话内同一工具不会再被拦截）
*shared.approved_tool.lock().await = None;
}
// 执行内置工具
if let Some(result) =
crate::tools::execute(&target_name, &target_args)
.await
{
if target_name == "manage_memos" {
let cp: serde_json::Value =
serde_json::from_str(&target_args).unwrap_or_default();
if cp.get("action").and_then(|v| v.as_str()) == Some("add")
&& let Some(content) = cp.get("content").and_then(|v| v.as_str()) {
shared
.memories
.lock()
.await
.entry(shared.user_id.clone())
.or_default()
.push(content.to_string());
shared.new_memories.lock().await.push(content.to_string());
}
}
return Ok(result.output);
}
Ok(format!("未知工具: {}", target_name))
}
}
})
})
}