From 00002e474069e44d6b819f963c7b0bdfab0eb8cb Mon Sep 17 00:00:00 2001 From: wunianxiao Date: Mon, 1 Jun 2026 17:45:17 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E5=90=8E=E5=8F=B0=E6=9C=8D=E5=8A=A1?= =?UTF-8?q?=E6=A8=A1=E5=BC=8F=20+=20=E6=97=A5=E5=BF=97=E6=BB=9A=E5=8A=A8?= =?UTF-8?q?=20+=20=E4=BF=AE=E5=A4=8D=20llm=5Fusage=20=E8=A1=A8=E5=85=BC?= =?UTF-8?q?=E5=AE=B9=E6=80=A7?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 新增 logger.rs（tracing-appender，日滚日志，.data/logs/） - service 子命令：等同 listen --llm 但日志同时写入文件 - systemd/ias.service 单元文件 - 迁移 0007：ALTER TABLE llm_usage ADD COLUMN IF NOT EXISTS 修复与 iPet 旧表的结构冲突 - ias usage 命令可正常查询 Token 统计 - 日志滚动保留 7 天，自动删除旧文件 --- Cargo.lock | 81 +++++++++++++++++++ Cargo.toml | 1 + .../20260601000006_session_summaries.sql | 12 +++ migrations/20260601000007_fix_llm_usage.sql | 11 +++ src/cli.rs | 3 + src/context/builder.rs | 58 +++++++++---- src/context/mod.rs | 1 - src/context/tools.rs | 48 ++++++----- src/context/types.rs | 55 +++++++++++++ src/db/models.rs | 35 ++++++++ src/logger.rs | 38 +++++++++ src/main.rs | 26 +++--- systemd/ias.service | 17 ++++ 13 files changed, 341 insertions(+), 45 deletions(-) create mode 100644 migrations/20260601000006_session_summaries.sql create mode 100644 migrations/20260601000007_fix_llm_usage.sql create mode 100644 src/logger.rs create mode 100644 systemd/ias.service diff --git a/Cargo.lock b/Cargo.lock index d73eebb..9ff213a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -467,6 +467,15 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "crossbeam-channel" +version = "0.5.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "82b8f8f868b36967f9606790d1903570de9ceaf870a7bf9fbbd3016d636a2cb2" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "crossbeam-deque" version = "0.8.6" @@ -528,6 +537,15 @@ dependencies = [ "zeroize", ] +[[package]] +name = "deranged" +version = "0.5.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7cd812cc2bc1d69d4764bd80df88b4317eaef9e773c75226407d9bc0876b211c" +dependencies = [ + "powerfmt", +] + [[package]] name = "digest" version = "0.10.7" @@ -1127,6 +1145,7 @@ dependencies = [ "sqlx", "tokio", "tracing", + "tracing-appender", "tracing-subscriber", "uuid", ] @@ -1652,6 +1671,12 @@ dependencies = [ "zeroize", ] +[[package]] +name = "num-conv" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "521739c6d2bac4aa25192232afe6841231376b2b26d4d9fae5ecf8ca5772e441" + [[package]] name = "num-derive" version = "0.4.2" @@ -1876,6 +1901,12 @@ dependencies = [ "zerovec", ] +[[package]] +name = "powerfmt" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "439ee305def115ba05938db6eb1644ff94165c5ab5e9420d1c1bcedbba909391" + [[package]] name = "ppv-lite86" version = "0.2.21" @@ -2752,6 +2783,12 @@ version = "2.6.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" +[[package]] +name = "symlink" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7973cce6668464ea31f176d85b13c7ab3bba2cb3b77a2ed26abd7801688010a" + [[package]] name = "syn" version = "2.0.117" @@ -2860,6 +2897,37 @@ dependencies = [ "zune-jpeg", ] +[[package]] +name = "time" +version = "0.3.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "743bd48c283afc0388f9b8827b976905fb217ad9e647fae3a379a9283c4def2c" +dependencies = [ + "deranged", + "itoa", + "num-conv", + "powerfmt", + "serde_core", + "time-core", + "time-macros", +] + +[[package]] +name = "time-core" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7694e1cfe791f8d31026952abf09c69ca6f6fa4e1a1229e18988f06a04a12dca" + +[[package]] +name = "time-macros" +version = "0.2.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e70e4c5a0e0a8a4823ad65dfe1a6930e4f4d756dcd9dd7939022b5e8c501215" +dependencies = [ + "num-conv", + "time-core", +] + [[package]] name = "tinystr" version = "0.8.3" @@ -3014,6 +3082,19 @@ dependencies = [ "tracing-core", ] +[[package]] +name = "tracing-appender" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "050686193eb999b4bb3bc2acfa891a13da00f79734704c4b8b4ef1a10b368a3c" +dependencies = [ + "crossbeam-channel", + "symlink", + "thiserror", + "time", + "tracing-subscriber", +] + [[package]] name = "tracing-attributes" version = "0.1.31" diff --git a/Cargo.toml b/Cargo.toml index 3a12129..f976fc0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -37,3 +37,4 @@ async-trait = "0.1" # 流式处理 futures-util = "0.3" sqlx = { version = "0.8.6", features = ["runtime-tokio-native-tls", "postgres", "uuid", "chrono", "derive", "migrate"] } +tracing-appender = "0.2.5" diff --git a/migrations/20260601000006_session_summaries.sql b/migrations/20260601000006_session_summaries.sql new file mode 100644 index 0000000..13adf1a --- /dev/null +++ b/migrations/20260601000006_session_summaries.sql @@ -0,0 +1,12 @@ +-- 会话摘要持久化 +CREATE TABLE IF NOT EXISTS session_summaries ( + id SERIAL PRIMARY KEY, + user_id TEXT NOT NULL DEFAULT 'default', + summary_text TEXT NOT NULL, + reason TEXT NOT NULL DEFAULT 'overflow', + message_count INTEGER NOT NULL DEFAULT 0, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +CREATE INDEX IF NOT EXISTS idx_session_summaries_created + ON session_summaries (created_at DESC); diff --git a/migrations/20260601000007_fix_llm_usage.sql b/migrations/20260601000007_fix_llm_usage.sql new file mode 100644 index 0000000..76d6857 --- /dev/null +++ b/migrations/20260601000007_fix_llm_usage.sql @@ -0,0 +1,11 @@ +-- 修复 llm_usage 表：从 iPet 旧格式迁移到 iAs 新格式 +-- 旧表有 kind/entry 字段，新表需要 prompt_tokens/completion_tokens 等 + +ALTER TABLE llm_usage ADD COLUMN IF NOT EXISTS prompt_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE llm_usage ADD COLUMN IF NOT EXISTS completion_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE llm_usage ADD COLUMN IF NOT EXISTS total_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE llm_usage ADD COLUMN IF NOT EXISTS cache_hit_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE llm_usage ADD COLUMN IF NOT EXISTS cache_miss_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE llm_usage ADD COLUMN IF NOT EXISTS user_id TEXT NOT NULL DEFAULT ''; +ALTER TABLE llm_usage ADD COLUMN IF NOT EXISTS model TEXT NOT NULL DEFAULT ''; +ALTER TABLE llm_usage ADD COLUMN IF NOT EXISTS provider TEXT NOT NULL DEFAULT 'deepseek'; diff --git a/src/cli.rs b/src/cli.rs index 4eeaa8d..f0b5d82 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -60,4 +60,7 @@ pub enum Commands { #[arg(long)] model: Option, }, + + /// 后台服务模式（等同 listen --llm，日志写入文件） + Service, } diff --git a/src/context/builder.rs b/src/context/builder.rs index d5dfd7f..2b4fee8 100644 --- a/src/context/builder.rs +++ b/src/context/builder.rs @@ -30,11 +30,8 @@ fn is_cjk(c: char) -> bool { /// 构建给 LLM 的消息数组（带 token budget 管理） /// /// 返回消息数组和是否需要摘要的信息 -pub async fn build_context( - session: &Arc>, - system_prompt: &str, -) -> Vec { - let mut s = session.lock().await; +pub async fn build_context(session: &Arc>, system_prompt: &str) -> Vec { + let s = session.lock().await; // ── 1. 空闲超时检查（消息到达前由调用方检查）── // 这里只做构建，超时触发在上层 @@ -74,10 +71,14 @@ pub async fn build_context( included.reverse(); // 保护：确保最近一条用户消息始终在上下文中 - if let Some(last_user) = recent.iter().rev().find(|m| m.role == crate::llm::types::Role::User) { - let already_included = included.iter().any(|m| { - m.role == last_user.role && m.content == last_user.content - }); + if let Some(last_user) = recent + .iter() + .rev() + .find(|m| m.role == crate::llm::types::Role::User) + { + let already_included = included + .iter() + .any(|m| m.role == last_user.role && m.content == last_user.content); if !already_included { let t = estimate_tokens(last_user); if used + tail_used + t <= s.token_budget + 2000 { @@ -112,7 +113,7 @@ pub async fn trigger_overflow_summary( let to_summarize: Vec<_> = s.messages[start..end] .iter() - .take(40) // 最多 40 条 + .take(40) // 最多 40 条 .cloned() .collect(); @@ -124,11 +125,15 @@ pub async fn trigger_overflow_summary( s.summaries.push(super::types::SummaryEntry { checkpoint: prev_checkpoint, - text: summary_text, + text: summary_text.clone(), reason: super::types::SummaryReason::Overflow, created_at: chrono::Utc::now(), }); + // 持久化到数据库 + s.save_summary_to_db(&summary_text, "overflow", to_summarize.len() as i32) + .await; + s.checkpoint = end; // 清理太旧的消息（checkpoint 之前的保留最近 100 条用于调试） @@ -163,11 +168,15 @@ pub async fn trigger_idle_summary( s.summaries.push(super::types::SummaryEntry { checkpoint: cp, - text: summary_text, + text: summary_text.clone(), reason: super::types::SummaryReason::Timeout, created_at: chrono::Utc::now(), }); + let msg_count = cp as i32; + s.save_summary_to_db(&summary_text, "timeout", msg_count) + .await; + s.checkpoint = s.messages.len(); // 清理旧消息（checkpoint 之前的全部清除） @@ -204,9 +213,15 @@ async fn generate_summary(messages: &[Message], summarizer: Option<&Summarizer>) fn build_summary_prompt(messages: &[Message]) -> String { let convo: String = messages .iter() - .filter(|m| m.role == crate::llm::types::Role::User || m.role == crate::llm::types::Role::Assistant) + .filter(|m| { + m.role == crate::llm::types::Role::User || m.role == crate::llm::types::Role::Assistant + }) .map(|m| { - let role = if m.role == crate::llm::types::Role::User { "用户" } else { "助手" }; + let role = if m.role == crate::llm::types::Role::User { + "用户" + } else { + "助手" + }; let content: String = m.content.chars().take(200).collect(); format!("{}: {}", role, content) }) @@ -224,7 +239,9 @@ fn summarize_messages(messages: &[Message]) -> String { // 简单实现：提取用户消息的前 80 个字符作为摘要 let lines: Vec = messages .iter() - .filter(|m| m.role == crate::llm::types::Role::User || m.role == crate::llm::types::Role::Assistant) + .filter(|m| { + m.role == crate::llm::types::Role::User || m.role == crate::llm::types::Role::Assistant + }) .map(|m| { let role = match m.role { crate::llm::types::Role::User => "用户", @@ -239,9 +256,16 @@ fn summarize_messages(messages: &[Message]) -> String { if lines.is_empty() { "暂无历史对话".to_string() } else { - let result = format!("历史对话摘要（{} 条消息）:\n{}", lines.len(), lines.join("\n")); + let result = format!( + "历史对话摘要（{} 条消息）:\n{}", + lines.len(), + lines.join("\n") + ); if result.chars().count() > 2000 { - format!("{}...（已截断）", result.chars().take(2000).collect::()) + format!( + "{}...（已截断）", + result.chars().take(2000).collect::() + ) } else { result } diff --git a/src/context/mod.rs b/src/context/mod.rs index 8bd3b8f..c7178ca 100644 --- a/src/context/mod.rs +++ b/src/context/mod.rs @@ -3,4 +3,3 @@ pub mod tools; pub mod types; pub use tools::MemoryStore; -pub use types::ChatSession; diff --git a/src/context/tools.rs b/src/context/tools.rs index ba953e2..b5183f9 100644 --- a/src/context/tools.rs +++ b/src/context/tools.rs @@ -64,29 +64,41 @@ impl MemoryStore { } } -/// read_summaries 工具：读取历史摘要 +/// read_summaries 工具：读取历史摘要（内存 + 数据库） pub async fn read_summaries( session: &Arc>, ) -> String { let s = session.lock().await; - if s.summaries.is_empty() { + let mut lines: Vec = Vec::new(); + + // 内存中的摘要 + for sum in &s.summaries { + let reason = match sum.reason { + super::types::SummaryReason::Overflow => "上下文压缩", + super::types::SummaryReason::Timeout => "空闲超时", + }; + lines.push(format!( + "[{}] {} (原因: {})", + sum.created_at.format("%Y-%m-%d %H:%M"), + sum.text, + reason, + )); + } + + // 数据库中的历史摘要 + if let Some(pool) = &s.db_pool { + if let Ok(db_summaries) = crate::db::models::load_summaries(pool, 5).await { + for text in db_summaries { + if !lines.iter().any(|l| l.contains(&text.chars().take(30).collect::())) { + lines.push(format!("[数据库存档] {}", text)); + } + } + } + } + + if lines.is_empty() { "暂无历史摘要".to_string() } else { - s.summaries - .iter() - .map(|sum| { - let reason = match sum.reason { - super::types::SummaryReason::Overflow => "上下文压缩", - super::types::SummaryReason::Timeout => "空闲超时", - }; - format!( - "[{}] {} (原因: {})", - sum.created_at.format("%Y-%m-%d %H:%M"), - sum.text, - reason, - ) - }) - .collect::>() - .join("\n---\n") + lines.join("\n---\n") } } diff --git a/src/context/types.rs b/src/context/types.rs index 29bcbde..9a96a2e 100644 --- a/src/context/types.rs +++ b/src/context/types.rs @@ -1,6 +1,7 @@ use crate::llm::types::Message; use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; +use std::sync::Arc; /// 摘要原因 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] @@ -27,6 +28,8 @@ pub struct SummaryEntry { /// 聊天会话状态 #[derive(Debug, Clone)] pub struct ChatSession { + pub user_id: String, + pub db_pool: Option>, /// 摘要总结点——此位置之前的消息已被压缩 pub checkpoint: usize, /// 全部消息（checkpoint 之后的保持完整） @@ -44,6 +47,8 @@ pub struct ChatSession { impl Default for ChatSession { fn default() -> Self { Self { + user_id: "default".to_string(), + db_pool: None, checkpoint: 0, messages: Vec::new(), summaries: Vec::new(), @@ -108,4 +113,54 @@ impl ChatSession { pub fn latest_overflow_summary(&self) -> Option<&SummaryEntry> { self.summaries.iter().rev().find(|s| s.reason == SummaryReason::Overflow) } + + /// 从数据库加载最近的聊天记录到会话中 + pub async fn load_recent_messages(&mut self, limit: i64) { + let pool = match &self.db_pool { + Some(p) => p.clone(), + None => return, + }; + + let rows = 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; + + if let Ok(rows) = rows { + let mut msgs: Vec = 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()); + } + } + } + + /// 保存摘要到数据库 + pub async fn save_summary_to_db(&self, text: &str, reason: &str, message_count: i32) { + if let Some(pool) = &self.db_pool { + let _ = sqlx::query( + "INSERT INTO session_summaries (user_id, summary_text, reason, message_count) VALUES ($1, $2, $3, $4)", + ) + .bind(&self.user_id) + .bind(text) + .bind(reason) + .bind(message_count) + .execute(pool.as_ref()) + .await; + } + } } diff --git a/src/db/models.rs b/src/db/models.rs index 81065ea..b6f365e 100644 --- a/src/db/models.rs +++ b/src/db/models.rs @@ -191,3 +191,38 @@ pub async fn query_llm_usage_stats( Ok(row) } + +// ─── 会话摘要 ─── + +/// 保存一条摘要到数据库 +pub async fn save_summary( + pool: &PgPool, + user_id: &str, + text: &str, + reason: &str, + message_count: i32, +) -> Result<(), String> { + sqlx::query( + "INSERT INTO session_summaries (user_id, summary_text, reason, message_count) VALUES ($1, $2, $3, $4)", + ) + .bind(user_id) + .bind(text) + .bind(reason) + .bind(message_count) + .execute(pool) + .await + .map_err(|e| format!("保存摘要失败: {}", e))?; + Ok(()) +} + +/// 加载最近的摘要（用于 read_summaries 工具） +pub async fn load_summaries(pool: &PgPool, limit: i64) -> Result, String> { + let rows = sqlx::query_as::<_, (String,)>( + "SELECT summary_text FROM session_summaries ORDER BY created_at DESC LIMIT $1", + ) + .bind(limit) + .fetch_all(pool) + .await + .map_err(|e| format!("查询摘要失败: {}", e))?; + Ok(rows.into_iter().map(|(t,)| t).collect()) +} diff --git a/src/logger.rs b/src/logger.rs new file mode 100644 index 0000000..6398855 --- /dev/null +++ b/src/logger.rs @@ -0,0 +1,38 @@ +use std::path::PathBuf; +use tracing_subscriber::layer::SubscriberExt; +use tracing_subscriber::util::SubscriberInitExt; +use tracing_subscriber::EnvFilter; + +pub fn init_logger(log_dir: Option<&str>, with_file: bool) { + let env_filter = EnvFilter::try_from_default_env() + .unwrap_or_else(|_| EnvFilter::new("info")); + + let terminal = tracing_subscriber::fmt::layer() + .with_target(false) + .with_ansi(true); + + if with_file { + if let Some(dir) = log_dir { + let dir_path = PathBuf::from(dir); + std::fs::create_dir_all(&dir_path).ok(); + + let file_appender = tracing_appender::rolling::daily(&dir_path, "ias.log"); + let file_layer = tracing_subscriber::fmt::layer() + .with_target(false) + .with_ansi(false) + .with_writer(file_appender); + + tracing_subscriber::registry() + .with(env_filter) + .with(terminal) + .with(file_layer) + .init(); + return; + } + } + + tracing_subscriber::registry() + .with(env_filter) + .with(terminal) + .init(); +} diff --git a/src/main.rs b/src/main.rs index 67ce15a..9c66c51 100644 --- a/src/main.rs +++ b/src/main.rs @@ -3,6 +3,7 @@ mod config; mod context; mod db; mod llm; +mod logger; mod scheduler; mod state; mod tools; @@ -24,12 +25,9 @@ use wechat::client::WeChatClient; #[tokio::main] async fn main() { - tracing_subscriber::fmt() - .with_env_filter( - EnvFilter::try_from_default_env() - .unwrap_or_else(|_| EnvFilter::new("info")), - ) - .init(); + let cli = Cli::parse(); + let is_service = matches!(cli.command, Commands::Service); + logger::init_logger(Some(".data/logs"), is_service); dotenvy::dotenv().ok(); @@ -54,8 +52,6 @@ async fn main() { // 文件状态管理器（降级方案） let file_state = state::StateManager::new(&config.storage.state_dir); - let cli = Cli::parse(); - match cli.command { Commands::Login { timeout } => { cmd_login(timeout, &database, &file_state).await; @@ -79,6 +75,9 @@ async fn main() { Commands::Usage { since, until, model } => { cmd_usage(&database, since, until, model).await; } + Commands::Service => { + cmd_listen(true, false, &database, &file_state, &config, &memory_store).await + } } } @@ -258,7 +257,16 @@ async fn cmd_listen( // 创建 Conversation let mut conv = match Conversation::new(cfg) { - Ok(c) => { + Ok(mut 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())); + s.load_recent_messages(20).await; + } + } info!("LLM 已初始化: {} / {}", c.provider_name(), c.model()); c } diff --git a/systemd/ias.service b/systemd/ias.service new file mode 100644 index 0000000..fc1019b --- /dev/null +++ b/systemd/ias.service @@ -0,0 +1,17 @@ +[Unit] +Description=iAs - 微信 AI 智能助手 +After=network.target postgresql.service + +[Service] +Type=simple +User=xiao +WorkingDirectory=/home/xiao/project/iAs +EnvironmentFile=/home/xiao/project/iAs/.env +ExecStart=/home/xiao/project/iAs/target/release/iAs service +Restart=on-failure +RestartSec=10 +StandardOutput=null +StandardError=null + +[Install] +WantedBy=multi-user.target