1091 lines
39 KiB
Rust
1091 lines
39 KiB
Rust
//! ## iAs —— 微信 AI 智能助手（入口）
|
||
//!
|
||
//! ### 系统架构概览
|
||
//!
|
||
//! ```text
|
||
//! CLI (main.rs → clap)
|
||
//! ├── ias login → 扫码登录，保存 token 到 DB 或文件
|
||
//! ├── ias listen --llm → Daemon 模式（主运行模式）
|
||
//! │ └── daemon.rs: 长轮询 → MessageQueue → 3 个 tokio 消费者
|
||
//! │ ├── LLM Consumer → Conversation.chat_with_tools() → DeepSeek API
|
||
//! │ ├── Tool Consumer → BuiltinRegistry.execute() / ApprovalManager
|
||
//! │ └── Send Consumer → WeChatClient.send_message()
|
||
//! ├── ias service → 等同 listen --llm 但启用文件日志
|
||
//! ├── ias daemon → 守护进程入口（等同 listen --llm）
|
||
//! ├── ias tool <tool> → 直接调用内置工具（无需登录）
|
||
//! └── ias usage → Token 消耗统计查询
|
||
//! ```
|
||
//!
|
||
//! ### 数据流
|
||
//! ```text
|
||
//! 微信消息 → 存 chat_records → 加载历史 + 记忆 → 构建 LLM 上下文
|
||
//! → Conversation.chat_with_tools() → LLM 回复或调工具
|
||
//! → 高风险工具走审批 → 工具结果回喂 LLM → 最终回复发送
|
||
//! ```
|
||
//!
|
||
//! ### 关键设计决策
|
||
//! 1. **单进程架构** — 旧版 daemon+worker 分离模式已废弃，统一用 tokio task
|
||
//! 2. **公平队列** — MessageQueue 按用户轮转，防止一个用户刷屏饿死其他人
|
||
//! 3. **可选数据库** — PostgreSQL 不可用时回退到文件存储
|
||
//! 4. **两层元工具** — LLM 通过 query_capabilities / call_capability 间接调用工具
|
||
|
||
mod channel;
|
||
mod cli;
|
||
mod context;
|
||
mod daemon;
|
||
mod db;
|
||
mod ipc;
|
||
mod llm;
|
||
mod logger;
|
||
mod queue;
|
||
mod scheduler;
|
||
mod state;
|
||
mod tools;
|
||
mod wechat;
|
||
mod worker;
|
||
|
||
use clap::Parser;
|
||
use cli::{AmapAction, Cli, Commands, MemosAction, 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 wechat::client::WeChatClient;
|
||
|
||
#[tokio::main]
|
||
async fn main() {
|
||
let cli = Cli::parse();
|
||
let is_service = matches!(cli.command, Commands::Service);
|
||
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);
|
||
|
||
// .env 加载：先当前目录，再 ~/.ias/.env
|
||
if std::path::Path::new(".env").exists() {
|
||
dotenvy::dotenv().ok();
|
||
} else if let Ok(home) = std::env::var("HOME") {
|
||
let fallback = std::path::PathBuf::from(&home).join(".ias").join(".env");
|
||
if fallback.exists() {
|
||
dotenvy::from_path(&fallback).ok();
|
||
tracing::info!("使用全局环境配置: {}", fallback.display());
|
||
}
|
||
}
|
||
|
||
let file_state = state::StateManager::new(".data/weixin-ilink");
|
||
|
||
// 工具命令无需登录/数据库，提前处理
|
||
let tool_cmd = match &cli.command {
|
||
Commands::Tool(cmd) => Some(cmd),
|
||
_ => None,
|
||
};
|
||
if let Some(cmd) = tool_cmd {
|
||
cmd_tool(cmd.clone()).await;
|
||
return;
|
||
}
|
||
|
||
let database = match Database::connect().await {
|
||
Ok(db) => {
|
||
info!("✅ 数据库连接成功");
|
||
Some(Arc::new(db))
|
||
}
|
||
Err(e) => {
|
||
info!("数据库不可用，使用文件存储: {}", e);
|
||
None
|
||
}
|
||
};
|
||
|
||
// 长期记忆存储
|
||
let memory_store = Arc::new(MemoryStore::new(
|
||
database.as_ref().map(|db| Arc::new(db.pool().clone())),
|
||
));
|
||
memory_store.load("").await;
|
||
|
||
match cli.command {
|
||
Commands::Login { timeout } => {
|
||
cmd_login(timeout, &database, &file_state).await;
|
||
}
|
||
Commands::Listen {
|
||
llm: enable_llm,
|
||
echo,
|
||
} => {
|
||
cmd_listen(enable_llm, echo, &database, &file_state, &memory_store).await;
|
||
}
|
||
Commands::Send {
|
||
to,
|
||
text,
|
||
context_token,
|
||
} => {
|
||
cmd_send(&to, &text, context_token.as_deref(), &database, &file_state).await;
|
||
}
|
||
Commands::Whoami => {
|
||
cmd_whoami(&database, &file_state).await;
|
||
}
|
||
Commands::Usage {
|
||
since,
|
||
until,
|
||
model,
|
||
} => {
|
||
cmd_usage(&database, since, until, model).await;
|
||
}
|
||
Commands::Service => cmd_listen(true, false, &database, &file_state, &memory_store).await,
|
||
Commands::Daemon { sock } => {
|
||
cmd_daemon(sock, &database, &file_state, &memory_store).await;
|
||
}
|
||
Commands::Worker { sock } => {
|
||
warn!("Worker 模式已废弃，请使用 daemon 模式");
|
||
#[allow(deprecated)]
|
||
if let Err(e) = worker::run(&sock).await {
|
||
error!("Worker 失败: {}", e);
|
||
}
|
||
}
|
||
Commands::Tool(..) => unreachable!("Tool 命令已提前处理"),
|
||
}
|
||
}
|
||
|
||
// ─── 命令实现 ───
|
||
|
||
async fn cmd_daemon(
|
||
sock_path: String,
|
||
database: &Option<Arc<Database>>,
|
||
file_state: &state::StateManager,
|
||
memory_store: &Arc<MemoryStore>,
|
||
) {
|
||
daemon::run(sock_path, database, file_state, memory_store).await;
|
||
}
|
||
|
||
async fn cmd_login(
|
||
timeout_secs: u64,
|
||
database: &Option<Arc<Database>>,
|
||
file_state: &state::StateManager,
|
||
) {
|
||
let client = WeChatClient::new(None);
|
||
|
||
match client
|
||
.login(|url| println!("二维码链接: {}", url), timeout_secs)
|
||
.await
|
||
{
|
||
Ok(result) => {
|
||
let auth = state::AuthState {
|
||
token: result.token.clone(),
|
||
account_id: result.account_id.clone(),
|
||
base_url: result.base_url.clone(),
|
||
};
|
||
|
||
// 优先存数据库
|
||
if let Some(db) = database {
|
||
if let Err(e) = db::models::save_auth(db.pool(), &auth).await {
|
||
error!("保存 auth 到数据库失败: {}", e);
|
||
file_state.save_auth(&auth);
|
||
} else {
|
||
info!("认证信息已保存到数据库");
|
||
// 清理文件状态（迁移完成）
|
||
file_state.clear_auth();
|
||
}
|
||
} else {
|
||
file_state.save_auth(&auth);
|
||
}
|
||
|
||
println!("\n✅ 登录成功！");
|
||
println!(" 账号: {}", result.account_id);
|
||
println!(
|
||
" Token: {}...",
|
||
&result.token[..std::cmp::min(16, result.token.len())]
|
||
);
|
||
println!(" API: {}", result.base_url);
|
||
}
|
||
Err(e) => error!("登录失败: {}", e),
|
||
}
|
||
}
|
||
|
||
async fn cmd_listen(
|
||
enable_llm: bool,
|
||
echo: bool,
|
||
database: &Option<Arc<Database>>,
|
||
file_state: &state::StateManager,
|
||
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::builtin::BuiltinRegistry::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::builtin::BuiltinRegistry::is_builtin(&target_name) {
|
||
if tools::builtin::BuiltinRegistry::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::builtin::BuiltinRegistry::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 {
|
||
db::models::load_auth(db.pool()).await
|
||
} else {
|
||
file_state.load_auth()
|
||
};
|
||
|
||
match auth {
|
||
Some(a) => {
|
||
println!("账号: {}", a.account_id);
|
||
println!("Token: {}...", &a.token[..std::cmp::min(16, a.token.len())]);
|
||
println!("API: {}", a.base_url);
|
||
if database.is_some() {
|
||
println!("存储: PostgreSQL");
|
||
} else {
|
||
println!("存储: 文件 ({})", file_state.state_dir().display());
|
||
}
|
||
}
|
||
None => println!("未登录，请先执行 login 命令"),
|
||
}
|
||
}
|
||
|
||
async fn cmd_usage(
|
||
database: &Option<Arc<Database>>,
|
||
since: Option<String>,
|
||
until: Option<String>,
|
||
model: Option<String>,
|
||
) {
|
||
let db = match database {
|
||
Some(d) => d,
|
||
None => {
|
||
println!("数据库未配置，无法查询用量");
|
||
return;
|
||
}
|
||
};
|
||
|
||
let since_dt = since.and_then(|s| {
|
||
chrono::DateTime::parse_from_rfc3339(&s)
|
||
.ok()
|
||
.map(|d| d.with_timezone(&chrono::Utc))
|
||
});
|
||
let until_dt = until.and_then(|s| {
|
||
chrono::DateTime::parse_from_rfc3339(&s)
|
||
.ok()
|
||
.map(|d| d.with_timezone(&chrono::Utc))
|
||
});
|
||
let model_ref = model.as_deref();
|
||
|
||
match db::models::query_llm_usage_stats(db.pool(), since_dt, until_dt, model_ref).await {
|
||
Ok(stats) => {
|
||
let hit_rate = if stats.total_cache_hit + stats.total_cache_miss > 0 {
|
||
stats.total_cache_hit as f64
|
||
/ (stats.total_cache_hit + stats.total_cache_miss) as f64
|
||
* 100.0
|
||
} else {
|
||
0.0
|
||
};
|
||
|
||
println!("📊 LLM Token 使用统计");
|
||
println!("{:=<40}", "");
|
||
println!("调用次数: {}", stats.total_calls);
|
||
println!(
|
||
"Prompt Tokens: {} ({:.1}K)",
|
||
stats.total_prompt_tokens,
|
||
stats.total_prompt_tokens as f64 / 1000.0
|
||
);
|
||
println!(
|
||
"生成 Tokens: {} ({:.1}K)",
|
||
stats.total_completion_tokens,
|
||
stats.total_completion_tokens as f64 / 1000.0
|
||
);
|
||
println!(
|
||
"总 Tokens: {} ({:.1}K)",
|
||
stats.total_tokens,
|
||
stats.total_tokens as f64 / 1000.0
|
||
);
|
||
println!(
|
||
"缓存命中: {} ({:.1}%)",
|
||
stats.total_cache_hit, hit_rate
|
||
);
|
||
println!(
|
||
"缓存未命中: {} ({:.1}%)",
|
||
stats.total_cache_miss,
|
||
100.0 - hit_rate
|
||
);
|
||
}
|
||
Err(e) => println!("查询失败: {}", e),
|
||
}
|
||
}
|
||
|
||
async fn cmd_send(
|
||
to: &str,
|
||
text: &str,
|
||
context_token: Option<&str>,
|
||
database: &Option<Arc<Database>>,
|
||
file_state: &state::StateManager,
|
||
) {
|
||
let auth: Option<state::AuthState> = if let Some(db) = database {
|
||
db::models::load_auth(db.pool()).await
|
||
} else {
|
||
file_state.load_auth()
|
||
};
|
||
|
||
let auth = match auth {
|
||
Some(a) => a,
|
||
None => {
|
||
error!("未登录，请先执行 login 命令");
|
||
return;
|
||
}
|
||
};
|
||
|
||
let base_url = auth.base_url.clone();
|
||
let account_id = auth.account_id.clone();
|
||
let mut client = WeChatClient::new(Some(auth.base_url));
|
||
client
|
||
.set_auth(&auth.token, &auth.account_id, &base_url)
|
||
.await;
|
||
|
||
match client.send_text(to, text, context_token).await {
|
||
Ok(msg_id) => {
|
||
println!("✅ 消息已发送, msg_id={}", msg_id);
|
||
|
||
// 入库：发送的消息
|
||
if let Some(db) = database
|
||
&& let Err(e) = db::models::insert_chat_record(
|
||
db.pool(),
|
||
"outbound",
|
||
to,
|
||
&account_id,
|
||
text,
|
||
"manual",
|
||
context_token,
|
||
&msg_id,
|
||
)
|
||
.await
|
||
{
|
||
error!("保存聊天记录失败: {}", e);
|
||
}
|
||
}
|
||
Err(e) => error!("发送失败: {}", e),
|
||
}
|
||
}
|
||
|
||
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 命令 ───
|
||
|
||
async fn cmd_tool(cmd: ToolCommand) {
|
||
match cmd {
|
||
ToolCommand::Datetime => {
|
||
let result = tools::builtins::datetime::execute();
|
||
println!("{}", result.output);
|
||
}
|
||
ToolCommand::Memos(action) => {
|
||
let args = match action {
|
||
MemosAction::List => serde_json::json!({"action": "list"}),
|
||
MemosAction::Add { content } => {
|
||
serde_json::json!({"action": "add", "content": content})
|
||
}
|
||
MemosAction::Delete { id } => {
|
||
serde_json::json!({"action": "delete", "id": id})
|
||
}
|
||
};
|
||
let args_str = serde_json::to_string(&args).unwrap_or_default();
|
||
let result = tools::builtins::memos::execute(&args_str).await;
|
||
println!("{}", result.output);
|
||
}
|
||
ToolCommand::Weather { location, days } => {
|
||
let params = serde_json::json!({"location": location, "days": days});
|
||
let result = tools::builtins::weather::execute(params).await;
|
||
println!("{}", result.output);
|
||
}
|
||
ToolCommand::Search {
|
||
query,
|
||
max_results,
|
||
no_answer,
|
||
topic,
|
||
time_range,
|
||
} => {
|
||
let mut params = serde_json::json!({
|
||
"query": query,
|
||
"max_results": max_results,
|
||
});
|
||
if no_answer {
|
||
params["include_answer"] = serde_json::json!(false);
|
||
}
|
||
if let Some(t) = topic {
|
||
params["topic"] = serde_json::json!(t);
|
||
}
|
||
if let Some(tr) = time_range {
|
||
params["time_range"] = serde_json::json!(tr);
|
||
}
|
||
let result = tools::builtins::web_search::execute(params).await;
|
||
println!("{}", result.output);
|
||
}
|
||
ToolCommand::Fetch { url } => {
|
||
let params = serde_json::json!({"url": url});
|
||
let result = tools::builtins::fetch_page::execute(params).await;
|
||
println!("{}", result.output);
|
||
}
|
||
ToolCommand::Amap(action) => {
|
||
cmd_amap(action).await;
|
||
}
|
||
}
|
||
}
|
||
|
||
async fn cmd_amap(action: AmapAction) {
|
||
match action {
|
||
AmapAction::PoiSearch {
|
||
keywords,
|
||
city,
|
||
location,
|
||
radius,
|
||
page,
|
||
offset,
|
||
} => {
|
||
let mut params = serde_json::json!({
|
||
"keywords": keywords,
|
||
"page": page,
|
||
"offset": offset,
|
||
});
|
||
if let Some(c) = city {
|
||
params["city"] = serde_json::json!(c);
|
||
}
|
||
if let Some(l) = location {
|
||
params["location"] = serde_json::json!(l);
|
||
}
|
||
if let Some(r) = radius {
|
||
params["radius"] = serde_json::json!(r);
|
||
}
|
||
let result = tools::builtins::amap::execute("amap_poi_search", params).await;
|
||
if let Some(r) = result {
|
||
println!(
|
||
"{}\n{}",
|
||
r.output,
|
||
serde_json::to_string(&r.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 {
|
||
println!(
|
||
"{}\n{}",
|
||
r.output,
|
||
serde_json::to_string(&r.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 {
|
||
println!(
|
||
"{}\n{}",
|
||
r.output,
|
||
serde_json::to_string(&r.data).unwrap_or_default()
|
||
);
|
||
}
|
||
}
|
||
AmapAction::RoutePlan {
|
||
r#type,
|
||
origin,
|
||
destination,
|
||
city,
|
||
waypoints,
|
||
strategy,
|
||
} => {
|
||
let mut params = serde_json::json!({
|
||
"type": r#type,
|
||
"origin": origin,
|
||
"destination": destination,
|
||
});
|
||
if let Some(c) = city {
|
||
params["city"] = serde_json::json!(c);
|
||
}
|
||
if let Some(w) = waypoints {
|
||
params["waypoints"] = serde_json::json!(w);
|
||
}
|
||
if let Some(s) = strategy {
|
||
params["strategy"] = serde_json::json!(s);
|
||
}
|
||
let result = tools::builtins::amap::execute("amap_route_plan", params).await;
|
||
if let Some(r) = result {
|
||
println!(
|
||
"{}\n{}",
|
||
r.output,
|
||
serde_json::to_string(&r.data).unwrap_or_default()
|
||
);
|
||
}
|
||
}
|
||
AmapAction::TravelPlan {
|
||
city,
|
||
interests,
|
||
route_type,
|
||
} => {
|
||
let interests_arr: Vec<serde_json::Value> = interests
|
||
.split(',')
|
||
.map(|s| serde_json::json!(s.trim()))
|
||
.collect();
|
||
let params = serde_json::json!({
|
||
"city": city,
|
||
"interests": interests_arr,
|
||
"route_type": route_type,
|
||
});
|
||
let result = tools::builtins::amap::execute("amap_travel_plan", params).await;
|
||
if let Some(r) = result {
|
||
println!(
|
||
"{}\n{}",
|
||
r.output,
|
||
serde_json::to_string(&r.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 {
|
||
println!(
|
||
"{}\n{}",
|
||
r.output,
|
||
serde_json::to_string(&r.data).unwrap_or_default()
|
||
);
|
||
}
|
||
}
|
||
}
|
||
}
|