From acf74981b21bf691f2a2f72cafb74c08a3e40344 Mon Sep 17 00:00:00 2001 From: xiaoxixi Date: Fri, 21 Aug 2026 17:29:48 +0800 Subject: [PATCH] feat(memory): relevance-gated auto recall with hard timeout Replace the unconditional top-5 keyword recall with a deterministic, layered gate so only relevant Knowledge entries reach the prompt. - New src/memory/recall.rs: tokenization with stopword filtering and a bounded term list; ranking combines lexical relevance, importance and recency; double gate (min_relevance + min_score) drops weak or stale matches; hard tokio::time::timeout wraps the SQL search so a slow FTS5 query never delays a turn. - MemoryConfig gains recall_min_relevance, recall_min_score, recall_recency_half_life_days, recall_timeout_ms (default 1000ms); recall_limit is now actually wired instead of being a dead field. - MemoryManager::recall_for_context exposes the gated path; the memory_recall tool keeps the raw search. - Storage::search_memories / search_memories_by_time share a single jieba-based tokenizer via the new module; search_memories_by_terms accepts pre-tokenized input for the gated path. - Docs and example configs (README, about-picobot references, both config.example.json templates) updated to reflect the new behavior. --- README.md | 8 +- .../about-picobot/assets/config.example.json | 4 + .../about-picobot/references/architecture.md | 2 +- .../skills/about-picobot/references/config.md | 8 +- .../skills/about-picobot/references/faq.md | 2 +- resources/templates/config.example.json | 4 + src/config/mod.rs | 31 +++ src/gateway/mod.rs | 15 +- src/memory/mod.rs | 118 ++++++++- src/memory/recall.rs | 242 ++++++++++++++++++ src/session/turn_input.rs | 21 +- src/storage/memory.rs | 119 ++++----- 12 files changed, 486 insertions(+), 88 deletions(-) create mode 100644 src/memory/recall.rs diff --git a/README.md b/README.md index 927d53a..21468c0 100644 --- a/README.md +++ b/README.md @@ -329,7 +329,7 @@ PicoBot 有两类记忆： | Knowledge | 偏好、事实、项目规则、长期可复用信息 | 长期保留，手动删除 | | Timeline | 长对话压缩后的历史摘要 | 默认保留 90 天 | -每轮处理用户消息时，MemoryManager 会按用户输入召回 Knowledge，并作为运行时上下文附加到本轮用户消息。当前召回上限固定为 5；`memory.recall_limit` 已支持解析但尚未接入 worker。长会话使用一个活动 checkpoint：累计摘要加 `first_retained_seq` 之后的原始消息尾部构成模型上下文，原始消息、工具调用结果、ID 和 seq 均不会被压缩改写。旧工具结果会保留在原始历史中，但 checkpoint 边界推进后不再永久占用 Provider 上下文。成功的语义摘要还会 best-effort 保存为 Timeline，供 `timeline_recall` 检索；Timeline 不参与会话恢复正确性。Scheduler 默认创建一个每日维护任务，按 `memory.timeline_retention_days` 清理过期 Timeline；结果通过 `complete_scheduled_run` 结构化提交，Knowledge 不会被自动删除。 +每轮处理用户消息时，MemoryManager 会按用户输入召回 Knowledge，并作为运行时上下文附加到本轮用户消息。自动召回是确定性的关键词检索（jieba 分词 + FTS5），按「词项相关度 + 重要度 + 时效」加权并通过相关性/综合分双门槛过滤，条数受 `memory.recall_limit` 约束，搜索受 `memory.recall_timeout_ms` 硬超时保护，超时或无关时本轮不注入。长会话使用一个活动 checkpoint：累计摘要加 `first_retained_seq` 之后的原始消息尾部构成模型上下文，原始消息、工具调用结果、ID 和 seq 均不会被压缩改写。旧工具结果会保留在原始历史中，但 checkpoint 边界推进后不再永久占用 Provider 上下文。成功的语义摘要还会 best-effort 保存为 Timeline，供 `timeline_recall` 检索；Timeline 不参与会话恢复正确性。Scheduler 默认创建一个每日维护任务，按 `memory.timeline_retention_days` 清理过期 Timeline；结果通过 `complete_scheduled_run` 结构化提交，Knowledge 不会被自动删除。 模型的 `models..token_limit` 给出上下文窗口上限，未配置时默认为 128,000；Agent 的 `agents..token_limit` 是可选的收紧上限，两者都有配置时有效窗口取二者最小值，因此 Agent 不能扩大模型窗口。自动压缩使用保留量阈值 `context_tokens > context_window - effective_reserve`，默认 reserve 为 16,384 tokens，并尽量原样保留最近 20,000 tokens。小窗口会自动把 reserve 限制为窗口的一半、把近期保留量限制为有效阈值的一半。摘要请求不使用固定 32K 输入上限，而是按有效窗口扣除摘要输出、提示词和安全余量；超大历史只在摘要请求副本中按“已有 checkpoint + 最新消息优先”生成有界 head/tail 转录，SQLite 原文不变。手动 `/compact` 跳过自动阈值；换成小模型后若发送前预检已发现硬超限，或首次请求返回真实 context overflow，语义摘要不可用时才使用明确标记的确定性降级裁剪，正式请求最多重试一次。若 overflow 发生在工具已经执行之后，AgentLoop 只在当前内存转录上裁掉旧完整 Turn 并重试当前模型步骤一次，不会从数据库历史重跑工具。 @@ -403,7 +403,11 @@ Skill 是包含 `SKILL.md` 的目录。加载优先级从高到低： | `context_compaction.enabled` | `true` | | `context_compaction.reserve_tokens` | `16384` | | `context_compaction.keep_recent_tokens` | `20000` | -| `memory.recall_limit` | `5`（当前运行时固定为 5） | +| `memory.recall_limit` | `5` | +| `memory.recall_min_relevance` | `0.25` | +| `memory.recall_min_score` | `0.25` | +| `memory.recall_recency_half_life_days` | `30` | +| `memory.recall_timeout_ms` | `1000` | | `memory.timeline_retention_days` | `90` | | `mcp.tool_timeout_secs` | `180` | | `mcp.servers[].tool_settings` | `{}`；可按工具名声明 `read_only` / `exclusive`，并发状态自动推导 | diff --git a/resources/skills/about-picobot/assets/config.example.json b/resources/skills/about-picobot/assets/config.example.json index f734c3b..d4cc377 100644 --- a/resources/skills/about-picobot/assets/config.example.json +++ b/resources/skills/about-picobot/assets/config.example.json @@ -87,6 +87,10 @@ "consolidation_provider": null, "consolidation_model": null, "recall_limit": 5, + "recall_min_relevance": 0.25, + "recall_min_score": 0.25, + "recall_recency_half_life_days": 30, + "recall_timeout_ms": 1000, "idle_consolidation_minutes": 10, "timeline_retention_days": 90, "max_failures_before_degrade": 3 diff --git a/resources/skills/about-picobot/references/architecture.md b/resources/skills/about-picobot/references/architecture.md index 186f586..fa92fdc 100644 --- a/resources/skills/about-picobot/references/architecture.md +++ b/resources/skills/about-picobot/references/architecture.md @@ -228,7 +228,7 @@ WebUI/TUI 的 Active Turn 使用 `send_message(files=...)` 向自身 session 投 | 语义 checkpoint 提交后 | 摘要 best-effort 存储为 Timeline 记忆 | | 会话恢复 | 从 checkpoint 与原始 seq 确定性重建，不读取 Timeline | -`memory.recall_limit`、`idle_consolidation_minutes`、`timeline_retention_days` 和 `max_failures_before_degrade` 当前会被配置解析；其中每轮 Knowledge 召回在 worker 中仍固定为 5，其余自动维护策略尚未接入运行循环。不要把“配置可解析”误认为“行为已生效”。 +`memory.recall_limit` 及新增的 `recall_min_relevance`、`recall_min_score`、`recall_recency_half_life_days`、`recall_timeout_ms` 已生效：每轮自动召回用 jieba 分词 + FTS5 检索候选，按「词项相关度 0.5 + 重要度 0.3 + 时效 0.2」加权并通过相关性/综合分双门槛过滤，搜索受硬超时保护。`idle_consolidation_minutes`、`timeline_retention_days` 和 `max_failures_before_degrade` 仍只是配置解析：自动 idle consolidation、Timeline 清理和失败降级循环尚未接入运行循环。不要把“配置可解析”误认为“行为已生效”。 --- diff --git a/resources/skills/about-picobot/references/config.md b/resources/skills/about-picobot/references/config.md index 0263e6c..aefa13d 100644 --- a/resources/skills/about-picobot/references/config.md +++ b/resources/skills/about-picobot/references/config.md @@ -111,12 +111,16 @@ Gateway WebUI 的“配置”页可以编辑实际加载的配置文件。读取 |------|------|------|------| | `consolidation_provider` | string | 主 Agent provider | 当前记录在 MemoryManager 中供后续归并使用；压缩摘要仍使用 Session provider | | `consolidation_model` | string | 主 Agent model | 当前记录在 MemoryManager 中供后续归并使用；压缩摘要仍使用 Session model | -| `recall_limit` | int | 5 | 预期的每轮知识召回上限；当前 worker 固定使用 5 | +| `recall_limit` | int | 5 | 每轮自动注入上下文的知识记忆条数上限 | +| `recall_min_relevance` | float | 0.25 | 自动召回的相关性硬门槛（命中词项占比低于此值则丢弃） | +| `recall_min_score` | float | 0.25 | 自动召回的综合分门槛（相关度+重要度+时效加权） | +| `recall_recency_half_life_days` | int | 30 | 自动召回时效衰减的半衰期（天） | +| `recall_timeout_ms` | int | 1000 | 自动召回搜索的硬超时；超时本轮不注入记忆 | | `idle_consolidation_minutes` | int | 10 | 预留的空闲归并阈值；当前无对应循环 | | `timeline_retention_days` | int | 90 | 默认日常维护巡检删除超过该期限的 Timeline；Knowledge 不受影响 | | `max_failures_before_degrade` | int | 3 | 预留的归并失败阈值；当前无失败降级循环 | -注意：当前 worker 的 Knowledge 召回数量仍固定为 5；idle consolidation 和失败降级循环尚未接入。Timeline 清理由默认启用的 `picobot-routine-maintenance` Scheduled Run 执行；该任务使用 `never` 策略，结构化结果只进入运行审计和 Health。 +自动召回每轮用当前用户输入做关键词检索（jieba 分词 + FTS5），按「词项相关度 0.5 + 重要度 0.3 + 时效 0.2」加权，通过相关性/综合分双门槛后才注入；搜索有硬超时保证不拖慢 Turn。Timeline 不自动召回，需显式 `timeline_recall`。idle consolidation 和失败降级循环尚未接入。Timeline 清理由默认启用的 `picobot-routine-maintenance` Scheduled Run 执行；该任务使用 `never` 策略，结构化结果只进入运行审计和 Health。 ## channels.feishu 字段 diff --git a/resources/skills/about-picobot/references/faq.md b/resources/skills/about-picobot/references/faq.md index 2e589d0..e325560 100644 --- a/resources/skills/about-picobot/references/faq.md +++ b/resources/skills/about-picobot/references/faq.md @@ -64,7 +64,7 @@ LLM 调用记录存储在 `llm_calls` 表中。可通过 SQLite 客户端直接 ## Q: 为什么修改了某些 memory 配置却没有看到行为变化？ -当前 `recall_limit`、`idle_consolidation_minutes`、`timeline_retention_days` 和 `max_failures_before_degrade` 都能被配置解析，但每轮 Knowledge 召回仍固定为 5，自动 idle consolidation、Timeline 清理和失败降级循环尚未接入。以当前代码行为为准。 +当前 `recall_limit` 以及新增的 `recall_min_relevance`、`recall_min_score`、`recall_recency_half_life_days`、`recall_timeout_ms` 均已生效，控制每轮自动知识召回的相关性门槛、加权与超时。`idle_consolidation_minutes`、`timeline_retention_days` 和 `max_failures_before_degrade` 仍只是配置解析：自动 idle consolidation、Timeline 清理（由维护任务执行）和失败降级循环尚未接入。以当前代码行为为准。 ## Q: Gateway 为什么无法立即退出？ diff --git a/resources/templates/config.example.json b/resources/templates/config.example.json index 0566a4f..6d64162 100644 --- a/resources/templates/config.example.json +++ b/resources/templates/config.example.json @@ -111,6 +111,10 @@ "consolidation_provider": null, "consolidation_model": null, "recall_limit": 5, + "recall_min_relevance": 0.25, + "recall_min_score": 0.25, + "recall_recency_half_life_days": 30, + "recall_timeout_ms": 1000, "idle_consolidation_minutes": 10, "timeline_retention_days": 90, "max_failures_before_degrade": 3 diff --git a/src/config/mod.rs b/src/config/mod.rs index 9c1bb88..5609fee 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -555,6 +555,21 @@ pub struct MemoryConfig { /// Max knowledge entries injected into system prompt per turn. #[serde(default = "default_recall_limit")] pub recall_limit: usize, + /// Minimum lexical relevance (fraction of matched query terms) for a + /// memory to be injected during automatic per-turn recall. + #[serde(default = "default_recall_min_relevance")] + pub recall_min_relevance: f64, + /// Minimum combined score (relevance + importance + recency) for a memory + /// to be injected during automatic per-turn recall. + #[serde(default = "default_recall_min_score")] + pub recall_min_score: f64, + /// Half-life (days) of the recency decay factor used in automatic recall. + #[serde(default = "default_recall_recency_half_life_days")] + pub recall_recency_half_life_days: u64, + /// Hard timeout (ms) for the automatic recall search; on expiry the turn + /// proceeds without injected memories. + #[serde(default = "default_recall_timeout_ms")] + pub recall_timeout_ms: u64, /// Idle minutes before triggering consolidation (for async channels). #[serde(default = "default_idle_consolidation_minutes")] pub idle_consolidation_minutes: u64, @@ -572,6 +587,10 @@ impl Default for MemoryConfig { consolidation_provider: None, consolidation_model: None, recall_limit: 5, + recall_min_relevance: 0.25, + recall_min_score: 0.25, + recall_recency_half_life_days: 30, + recall_timeout_ms: 1000, idle_consolidation_minutes: 10, timeline_retention_days: 90, max_failures_before_degrade: 3, @@ -789,6 +808,18 @@ impl Default for BrowserConfig { fn default_recall_limit() -> usize { 5 } +fn default_recall_min_relevance() -> f64 { + 0.25 +} +fn default_recall_min_score() -> f64 { + 0.25 +} +fn default_recall_recency_half_life_days() -> u64 { + 30 +} +fn default_recall_timeout_ms() -> u64 { + 1000 +} fn default_idle_consolidation_minutes() -> u64 { 10 } diff --git a/src/gateway/mod.rs b/src/gateway/mod.rs index 50f4d69..aac9183 100644 --- a/src/gateway/mod.rs +++ b/src/gateway/mod.rs @@ -156,11 +156,16 @@ impl GatewayState { let consolidation_model = config .memory .resolve_consolidation_model(&provider_config.model_id); - let memory_manager = Arc::new(MemoryManager::new( - storage.clone(), - consolidation_provider, - consolidation_model, - )); + let memory_manager = Arc::new( + MemoryManager::new( + storage.clone(), + consolidation_provider, + consolidation_model, + ) + .with_recall(crate::memory::recall::RecallConfig::from_memory_config( + &config.memory, + )), + ); tracing::info!( consolidation_provider = %memory_manager.consolidation_provider, consolidation_model = %memory_manager.consolidation_model, diff --git a/src/memory/mod.rs b/src/memory/mod.rs index c57282f..40433a7 100644 --- a/src/memory/mod.rs +++ b/src/memory/mod.rs @@ -1,3 +1,4 @@ +pub mod recall; pub mod types; use std::sync::Arc; @@ -6,6 +7,8 @@ use uuid::Uuid; use crate::storage::Storage; pub use types::{ConsolidationFact, ConsolidationResult, MemoryCategory, MemoryEntry}; +use recall::RecallConfig; + /// MemoryManager provides high-level memory operations. /// Wraps the Storage SQLite layer with semantic methods. #[derive(Clone)] @@ -13,6 +16,7 @@ pub struct MemoryManager { storage: Arc, pub consolidation_provider: String, pub consolidation_model: String, + recall: RecallConfig, } impl MemoryManager { @@ -25,9 +29,57 @@ impl MemoryManager { storage, consolidation_provider, consolidation_model, + recall: RecallConfig::default(), } } + /// Override the automatic-recall knobs (derived from `config.memory`). + pub fn with_recall(mut self, config: RecallConfig) -> Self { + self.recall = config; + self + } + + /// Automatic per-turn recall for context injection. + /// + /// Returns up to `recall.limit` Knowledge entries that clear both the + /// relevance and combined-score gates. The search is bounded by a hard + /// timeout: on expiry, error, or an empty normalized query it returns an + /// empty vector so the turn is never delayed by memory lookup. + pub async fn recall_for_context(&self, query: &str) -> Vec { + let terms = recall::normalize_query(query); + if terms.is_empty() { + return Vec::new(); + } + let candidate_limit = (self.recall.limit * recall::CANDIDATE_FACTOR) + .min(recall::MAX_CANDIDATES) + .max(self.recall.limit); + let candidates = tokio::time::timeout( + self.recall.timeout, + self.storage.search_memories_by_terms( + &terms, + Some(&MemoryCategory::Knowledge), + None, + candidate_limit, + ), + ) + .await; + let candidates = match candidates { + Ok(Ok(entries)) => entries, + Ok(Err(error)) => { + tracing::debug!(error = %error, "memory recall failed; skipping injection"); + return Vec::new(); + } + Err(_) => { + tracing::debug!("memory recall timed out; skipping injection"); + return Vec::new(); + } + }; + let now_ms = chrono::Utc::now().timestamp_millis(); + let mut ranked = recall::rank_and_gate(candidates, &terms, &self.recall, now_ms); + ranked.truncate(self.recall.limit); + ranked + } + /// Store or update a memory entry. Generates timestamp and UUID. pub async fn store( &self, @@ -88,7 +140,10 @@ impl MemoryManager { /// Check if the memory system has any entries (for testing/health check). pub async fn is_empty(&self) -> Result { - self.recall("*", 1, None, None).await.map(|r| r.is_empty()) + self.storage + .list_memories(None, None, 1) + .await + .map(|entries| entries.is_empty()) } } @@ -257,4 +312,65 @@ mod tests { assert_eq!(scoped[0].key, "tl_a"); assert_eq!(scoped[0].session_id.as_deref(), Some("chan:chat:dialog_a")); } + + #[tokio::test] + async fn recall_for_context_returns_empty_for_stopword_only_query() { + let (mm, _dir) = setup_memory_manager().await; + mm.store( + "user_pref", + "user prefers python", + MemoryCategory::Knowledge, + None, + Some(0.9), + ) + .await + .unwrap(); + // A query made entirely of stopwords must not surface unrelated memories. + assert!(mm.recall_for_context("你好吗请吧").await.is_empty()); + } + + #[tokio::test] + async fn recall_for_context_recalls_relevant_and_gates_irrelevant() { + let (mm, _dir) = setup_memory_manager().await; + mm.store( + "user_pref", + "user prefers python for scripting", + MemoryCategory::Knowledge, + None, + Some(0.9), + ) + .await + .unwrap(); + mm.store( + "favorite_color", + "user favorite color is blue", + MemoryCategory::Knowledge, + None, + Some(0.9), + ) + .await + .unwrap(); + + let hits = mm.recall_for_context("帮我写一个 python 脚本").await; + assert_eq!(hits.len(), 1); + assert_eq!(hits[0].key, "user_pref"); + } + + #[tokio::test] + async fn recall_for_context_respects_limit() { + let (mm, _dir) = setup_memory_manager().await; + for i in 0..8 { + mm.store( + &format!("pref_{i}"), + format!("user preference about python number {i}").as_str(), + MemoryCategory::Knowledge, + None, + Some(0.9), + ) + .await + .unwrap(); + } + let hits = mm.recall_for_context("python preference").await; + assert!(hits.len() <= 5); + } } diff --git a/src/memory/recall.rs b/src/memory/recall.rs new file mode 100644 index 0000000..90431b1 --- /dev/null +++ b/src/memory/recall.rs @@ -0,0 +1,242 @@ +//! Deterministic, relevance-gated memory recall for per-turn context injection. +//! +//! This is the *automatic* recall path (as opposed to the user-facing +//! `memory_recall` tool, which performs an unfiltered search). It tokenizes +//! and normalizes the query, retrieves a small candidate set through FTS5, +//! then ranks and gates candidates by lexical relevance, importance, and +//! recency. Everything here is bounded and synchronous so a turn is never +//! delayed by a slow search. + +use std::collections::HashSet; +use std::sync::OnceLock; +use std::time::Duration; + +use jieba_rs::Jieba; + +use super::MemoryEntry; + +/// Weights for the combined score (sum to 1.0 for a clean 0..1 range). +const WEIGHT_RELEVANCE: f64 = 0.5; +const WEIGHT_IMPORTANCE: f64 = 0.3; +const WEIGHT_RECENCY: f64 = 0.2; + +/// Candidate set is `limit * CANDIDATE_FACTOR`, capped by `MAX_CANDIDATES`. +pub(crate) const CANDIDATE_FACTOR: usize = 4; +pub(crate) const MAX_CANDIDATES: usize = 50; + +/// Hard bounds on the query to keep tokenization and the FTS5 MATCH cheap. +const MAX_QUERY_CHARS: usize = 512; +const MAX_TERMS: usize = 12; + +const STOPWORDS: &[&str] = &[ + "的", "了", "是", "在", "我", "你", "他", "她", "它", "们", "吗", "呢", "啊", "吧", "这", "那", + "什么", "怎么", "为什么", "帮我", "一下", "请", "要", "想", "能", "可以", "这个", "那个", "今天", + "现在", "以及", "还是", "因为", "所以", "如果", "但是", "就", "都", "也", "还", "很", "被", "把", + "和", "与", "或", "有", "没有", "一个", "一种", "一些", "哪些", "如何", + "a", "an", "the", "to", "of", "for", "and", "or", "what", "how", "please", "can", "do", "is", + "are", "my", "you", "i", "we", "this", "that", "me", "it", "on", "in", "at", "with", "about", +]; + +/// Runtime knobs for automatic recall, derived from `config.memory`. +#[derive(Debug, Clone)] +pub struct RecallConfig { + pub limit: usize, + pub min_relevance: f64, + pub min_score: f64, + pub recency_half_life_days: f64, + pub timeout: Duration, +} + +impl Default for RecallConfig { + fn default() -> Self { + Self { + limit: 5, + min_relevance: 0.25, + min_score: 0.25, + recency_half_life_days: 30.0, + timeout: Duration::from_millis(1000), + } + } +} + +impl RecallConfig { + pub fn from_memory_config(memory: &crate::config::MemoryConfig) -> Self { + Self { + limit: memory.recall_limit, + min_relevance: memory.recall_min_relevance, + min_score: memory.recall_min_score, + recency_half_life_days: memory.recall_recency_half_life_days as f64, + timeout: Duration::from_millis(memory.recall_timeout_ms), + } + } +} + +fn jieba() -> &'static Jieba { + static INSTANCE: OnceLock = OnceLock::new(); + INSTANCE.get_or_init(Jieba::new) +} + +/// Tokenize with jieba and drop tokens that are too short to be meaningful. +pub(crate) fn tokenize(query: &str) -> Vec { + jieba() + .cut(query, true) + .into_iter() + .map(|token| token.word) + .filter(|word| word.len() > 1 || word.bytes().any(|b| b > 127)) + .map(str::to_string) + .collect() +} + +/// Normalize a user query into deduplicated, stopword-free search terms. +pub(crate) fn normalize_query(query: &str) -> Vec { + let bounded: String = if query.chars().count() > MAX_QUERY_CHARS { + query.chars().take(MAX_QUERY_CHARS).collect() + } else { + query.to_string() + }; + let mut seen = HashSet::new(); + let mut terms: Vec = Vec::new(); + for term in tokenize(&bounded) { + let lower = term.to_lowercase(); + if STOPWORDS.contains(&lower.as_str()) { + continue; + } + if seen.insert(lower) { + terms.push(term); + } + } + terms.truncate(MAX_TERMS); + terms +} + +/// Fraction of distinct query terms present in the entry's key or content. +fn relevance(entry: &MemoryEntry, terms: &[String]) -> f64 { + let key = entry.key.to_lowercase(); + let content = entry.content.to_lowercase(); + let matched = terms + .iter() + .filter(|term| { + let term = term.to_lowercase(); + key.contains(&term) || content.contains(&term) + }) + .count(); + matched as f64 / terms.len() as f64 +} + +/// Exponential recency decay in [0,1]; 1.0 for just-updated, 0.0 for very old. +fn recency(updated_at: &str, half_life_days: f64, now_ms: i64) -> f64 { + let Ok(dt) = chrono::DateTime::parse_from_rfc3339(updated_at) else { + return 0.0; + }; + let half_life = half_life_days.max(1.0); + let age_ms = (now_ms - dt.timestamp_millis()).max(0); + let age_days = age_ms as f64 / 86_400_000.0; + (-age_days / half_life).exp() +} + +/// Rank candidates by combined score and drop those below the gates. +pub(crate) fn rank_and_gate( + candidates: Vec, + terms: &[String], + config: &RecallConfig, + now_ms: i64, +) -> Vec { + let mut scored: Vec<(f64, f64, MemoryEntry)> = Vec::with_capacity(candidates.len()); + for entry in candidates { + let rel = relevance(&entry, terms); + if rel < config.min_relevance { + continue; + } + let importance = entry.importance.clamp(0.0, 1.0); + let rec = recency(&entry.updated_at, config.recency_half_life_days, now_ms); + let score = WEIGHT_RELEVANCE * rel + WEIGHT_IMPORTANCE * importance + WEIGHT_RECENCY * rec; + if score < config.min_score { + continue; + } + scored.push((score, importance, entry)); + } + scored.sort_by(|a, b| { + b.0.partial_cmp(&a.0) + .unwrap_or(std::cmp::Ordering::Equal) + .then_with(|| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)) + }); + scored.into_iter().map(|(_, _, entry)| entry).collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::memory::MemoryCategory; + + fn entry(key: &str, content: &str, importance: f64, updated_at: &str) -> MemoryEntry { + MemoryEntry { + id: format!("id-{key}"), + key: key.to_string(), + content: content.to_string(), + category: MemoryCategory::Knowledge, + importance, + session_id: None, + created_at: updated_at.to_string(), + updated_at: updated_at.to_string(), + } + } + + #[test] + fn normalize_query_removes_stopwords_and_deduplicates() { + let terms = normalize_query("请帮我写一个 Python 脚本，用 python 处理"); + assert!(!terms.contains(&"请".to_string())); + assert!(!terms.contains(&"一个".to_string())); + // "python" appears twice but is deduplicated (case-insensitive). + assert_eq!(terms.iter().filter(|t| t.to_lowercase() == "python").count(), 1); + } + + #[test] + fn normalize_query_returns_empty_for_all_stopwords() { + assert!(normalize_query("呢吗啊吧").is_empty()); + } + + #[test] + fn relevance_counts_distinct_term_matches() { + let e = entry("py", "the user prefers python", 0.5, "2024-01-01T00:00:00Z"); + let terms = vec!["python".to_string(), "rust".to_string()]; + assert!((relevance(&e, &terms) - 0.5).abs() < 1e-9); + } + + #[test] + fn gate_drops_low_relevance() { + let e = entry("unrelated", "something else entirely", 1.0, "2024-01-01T00:00:00Z"); + let terms = vec!["python".to_string(), "script".to_string()]; + let cfg = RecallConfig::default(); + let out = rank_and_gate(vec![e], &terms, &cfg, 1_704_067_200_000); + assert!(out.is_empty()); + } + + #[test] + fn gate_keeps_high_relevance_and_sorts_by_score() { + let strong = entry("strong", "python script", 0.5, "2024-01-01T00:00:00Z"); + let weak = entry("weak", "a python note", 0.9, "2024-01-01T00:00:00Z"); + let terms = vec!["python".to_string(), "script".to_string()]; + let cfg = RecallConfig::default(); + let out = rank_and_gate(vec![weak, strong], &terms, &cfg, 1_704_067_200_000); + assert_eq!(out.len(), 2); + assert_eq!(out[0].key, "strong"); + } + + #[test] + fn importance_breaks_relevance_ties() { + let a = entry("a", "python script", 0.2, "2024-01-01T00:00:00Z"); + let b = entry("b", "python script", 0.9, "2024-01-01T00:00:00Z"); + let terms = vec!["python".to_string(), "script".to_string()]; + let cfg = RecallConfig::default(); + let out = rank_and_gate(vec![a, b], &terms, &cfg, 1_704_067_200_000); + assert_eq!(out[0].key, "b"); + } + + #[test] + fn recency_decays_with_age() { + let recent = recency("2024-01-10T00:00:00Z", 30.0, 1_704_067_200_000); + let old = recency("2023-01-10T00:00:00Z", 30.0, 1_704_067_200_000); + assert!(recent > old); + assert!((0.0..=1.0).contains(&recent)); + } +} diff --git a/src/session/turn_input.rs b/src/session/turn_input.rs index 6a758e3..9a85da4 100644 --- a/src/session/turn_input.rs +++ b/src/session/turn_input.rs @@ -2,7 +2,7 @@ use std::sync::Arc; use crate::agent::system_prompt::build_runtime_context; use crate::bus::ChatMessage; -use crate::memory::{MemoryCategory, MemoryManager}; +use crate::memory::MemoryManager; use crate::work::WorkManager; /// Immutable context used to assemble provider input for both the initial call @@ -42,23 +42,20 @@ pub(super) async fn prepare_turn_runtime( query: &str, system_prompt: String, ) -> TurnRuntimeContext { - let memory_future = memory_manager.recall(query, 5, Some(MemoryCategory::Knowledge), None); + let memory_future = memory_manager.recall_for_context(query); let work_future = work_manager.active_plan(session_id); - let (memory_result, work_result) = tokio::join!(memory_future, work_future); + let (memory_entries, work_result) = tokio::join!(memory_future, work_future); - let memory_context = match memory_result { - Ok(entries) if !entries.is_empty() => Some( - entries + let memory_context = if memory_entries.is_empty() { + None + } else { + Some( + memory_entries .iter() .map(|entry| format!("- {}: {}", entry.key, entry.content)) .collect::>() .join("\n"), - ), - Err(error) => { - tracing::warn!(error = %error, "Failed to fetch memory context"); - None - } - _ => None, + ) }; let work_context = match work_result { Ok(Some(plan)) => Some(plan.compact_context()), diff --git a/src/storage/memory.rs b/src/storage/memory.rs index 50566f0..382d887 100644 --- a/src/storage/memory.rs +++ b/src/storage/memory.rs @@ -1,17 +1,9 @@ use sqlx::Row; -use std::sync::OnceLock; - -use jieba_rs::Jieba; use crate::memory::{MemoryCategory, MemoryEntry}; use super::StorageError; -fn jieba() -> &'static Jieba { - static INSTANCE: OnceLock = OnceLock::new(); - INSTANCE.get_or_init(Jieba::new) -} - impl super::Storage { /// List recent memories without requiring a full-text query. pub async fn list_memories( @@ -99,12 +91,27 @@ impl super::Storage { session_id: Option<&str>, limit: usize, ) -> Result, StorageError> { - // Build FTS5 query: segment with jieba, wrap each term in quotes, join with OR - let fts_query = jieba() - .cut(query, true) - .into_iter() - .map(|token| token.word) - .filter(|word| word.len() > 1 || word.bytes().any(|b| b > 127)) + let terms = crate::memory::recall::tokenize(query); + self.search_memories_by_terms(&terms, category, session_id, limit) + .await + } + + /// Search memories using pre-tokenized terms (FTS5 with LIKE fallback). + /// An empty term list returns no results without issuing a query. + pub async fn search_memories_by_terms( + &self, + terms: &[String], + category: Option<&MemoryCategory>, + session_id: Option<&str>, + limit: usize, + ) -> Result, StorageError> { + if terms.is_empty() { + return Ok(Vec::new()); + } + + // Build FTS5 query: wrap each term in quotes, join with OR. + let fts_query = terms + .iter() .map(|word| format!("\"{}\"", word.replace('"', ""))) .collect::>() .join(" OR "); @@ -136,51 +143,41 @@ impl super::Storage { // Fallback to term-based LIKE query if FTS5 returned nothing if entries.is_empty() { - let terms: Vec = jieba() - .cut(query, true) - .into_iter() - .map(|token| token.word) - .filter(|word| word.len() > 1 || word.bytes().any(|b| b > 127)) - .map(|word| word.replace(['%', '_'], "")) - .collect(); + let like_clauses = terms + .iter() + .map(|_| "(key LIKE ? OR content LIKE ?)") + .collect::>() + .join(" OR "); - if !terms.is_empty() { - let like_clauses = terms - .iter() - .map(|_| "(key LIKE ? OR content LIKE ?)") - .collect::>() - .join(" OR "); + let sql = format!( + r#" + SELECT id, key, content, category, importance, + session_id, created_at, updated_at + FROM memories + WHERE ({}) + AND (? IS NULL OR category = ?) + AND (? IS NULL OR session_id = ?) + ORDER BY importance DESC, updated_at DESC + LIMIT ? + "#, + like_clauses + ); - let sql = format!( - r#" - SELECT id, key, content, category, importance, - session_id, created_at, updated_at - FROM memories - WHERE ({}) - AND (? IS NULL OR category = ?) - AND (? IS NULL OR session_id = ?) - ORDER BY importance DESC, updated_at DESC - LIMIT ? - "#, - like_clauses - ); - - // The only interpolated fragment is a generated sequence of bind placeholders. - let mut query_builder = sqlx::query(sqlx::AssertSqlSafe(sql)); - for term in &terms { - let pattern = format!("%{}%", term); - query_builder = query_builder.bind(pattern.clone()).bind(pattern); - } - query_builder = query_builder - .bind(category_filter) - .bind(category_filter) - .bind(session_id) - .bind(session_id) - .bind(limit as i64); - - let rows = query_builder.fetch_all(self.pool()).await?; - entries = parse_memory_rows(&rows)?; + // The only interpolated fragment is a generated sequence of bind placeholders. + let mut query_builder = sqlx::query(sqlx::AssertSqlSafe(sql)); + for term in terms { + let pattern = format!("%{}%", term.replace(['%', '_'], "")); + query_builder = query_builder.bind(pattern.clone()).bind(pattern); } + query_builder = query_builder + .bind(category_filter) + .bind(category_filter) + .bind(session_id) + .bind(session_id) + .bind(limit as i64); + + let rows = query_builder.fetch_all(self.pool()).await?; + entries = parse_memory_rows(&rows)?; } Ok(entries) @@ -205,13 +202,7 @@ impl super::Storage { .to_rfc3339(); let rows = if let Some(q) = query { - let terms: Vec = jieba() - .cut(q, true) - .into_iter() - .map(|token| token.word) - .filter(|word| word.len() > 1 || word.bytes().any(|b| b > 127)) - .map(|word| word.replace(['%', '_'], "")) - .collect(); + let terms: Vec = crate::memory::recall::tokenize(q); if terms.is_empty() { return Ok(Vec::new()); @@ -241,7 +232,7 @@ impl super::Storage { // The only interpolated fragment is a generated sequence of bind placeholders. let mut query_builder = sqlx::query(sqlx::AssertSqlSafe(sql)); for term in &terms { - let pattern = format!("%{}%", term); + let pattern = format!("%{}%", term.replace(['%', '_'], "")); query_builder = query_builder.bind(pattern.clone()).bind(pattern); } query_builder = query_builder