mod cli; mod context; mod db; mod llm; mod logger; mod scheduler; mod state; mod tools; mod wechat; use clap::Parser; use cli::{Cli, Commands}; use context::MemoryStore; use db::Database; use llm::{Conversation, ConversationConfig, DEFAULT_SYSTEM_PROMPT, ToolExecutor, Usage}; use std::sync::Arc; use tools::approval::{ApprovalDecision, ApprovalManager}; use tools::types::ExecutionContext; use tracing::{error, info}; use 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 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, } } // ─── 命令实现 ─── async fn cmd_login( timeout_secs: u64, database: &Option>, 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>, file_state: &state::StateManager, memory_store: &Arc, ) { // 从数据库或文件加载认证 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 = 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(|_| ()) .map_err(|e| e) }) }) }; 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(build_capability_list(&specs)); } // 确定目标工具名和参数(call_capability 提取 name + 透传完整参数) 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(); (name, serde_json::to_string(&cp).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>, file_state: &state::StateManager, conversation: Option>, approval_manager: &ApprovalManager, memory_store: &Arc, ) { 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 { if 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 { if 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 { if 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 { if 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>, ) -> Result { 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>, ) -> Result { let (reply, _used_tools, usage) = conv.chat_with_tools(user_text).await?; if let Some(u) = &usage { info!( "📊 Token: {} 总 / {} 提示 / {} 生成 | 缓存命中 {} | 缓存未命中 {}", u.total_tokens, u.prompt_tokens, u.completion_tokens, u.prompt_cache_hit_tokens, u.prompt_cache_miss_tokens ); store_usage(db, "", conv.model(), conv.provider_name(), u).await; } Ok(reply) } async fn store_usage( db: &Option>, user_id: &str, model: &str, provider: &str, u: &Usage, ) { if let Some(db) = db { if 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>, file_state: &state::StateManager) { let auth: Option = 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>, since: Option, until: Option, model: Option, ) { 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>, file_state: &state::StateManager, ) { let auth: Option = 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 { if 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 { 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 { if let Err(e) = send(&ctx.user_id, &msg).await { return Err(e); } } match tokio::time::timeout(std::time::Duration::from_secs(300), rx).await { Ok(Ok(ApprovalDecision::Approved)) => Ok(true), _ => Ok(false), } } fn build_capability_list(specs: &[tools::types::SkillSpec]) -> String { let mut lines = vec!["📋 可用工具:".to_string()]; for spec in specs { let risk = if spec.risk_level == tools::types::RiskLevel::High { " ⚠️需确认" } else { "" }; lines.push(format!( "\n🔹 {} {} — {}", spec.name, risk, spec.description )); if !spec .parameters .get("properties") .and_then(|v| v.as_object()) .map_or(true, |o| o.is_empty()) { lines.push(format!( " 参数: {}", serde_json::to_string(&spec.parameters).unwrap_or_default() )); } } lines.join("\n") }