use sqlx::PgPool; use std::collections::HashMap; use std::sync::Arc; use tokio::sync::Mutex; /// ## 长期记忆管理器(内存 + PostgreSQL 双写) /// /// 为每个用户维护独立的长期记忆列表,支持读写操作。 /// /// ### 存储策略 /// * **内存缓存** — `HashMap>`,按 user_id 隔离 /// * **数据库持久化** — 可选的 `PgPool`,写入时同时写到 PostgreSQL /// /// ### 数据流 /// ```text /// 启动时 load(user_id) → 从数据库加载到缓存 /// 对话中 read_for() → 从缓存读取,返回 "1. ...\n2. ..." 格式 /// 对话中 write_for() → 写入缓存 + 写入数据库 /// ``` pub struct MemoryStore { pool: Option>, cache: Arc>>>, } impl MemoryStore { pub fn new(pool: Option>) -> Self { Self { pool, cache: Arc::new(Mutex::new(HashMap::new())), } } 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", ) .bind(uid) .fetch_all(pool.as_ref()) .await { let mems: Vec = rows.into_iter().map(|(c,)| c).collect(); self.cache.lock().await.insert(uid.to_string(), mems); } } pub async fn read_for(&self, user_id: &str) -> String { let uid = if user_id.is_empty() { "default" } else { user_id }; 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)) .collect::>() .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)") .bind(uid) .bind(content) .execute(pool.as_ref()) .await; } "已添加长期记忆".to_string() } } /// ## `read_summaries` 工具 —— 读取历史摘要(内存 + 数据库) /// /// 返回当前 session 中的摘要列表 + 数据库中存档的摘要。 /// 去重:数据库中的摘要如果与 session 中的摘要前 30 个字符相同,则跳过。 pub async fn read_summaries( session: &Arc>, ) -> String { let s = session.lock().await; 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 { 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::())) { lines.push(format!("[数据库存档] {}", text)); } } } } if lines.is_empty() { "暂无历史摘要".to_string() } else { lines.join("\n---\n") } }