From 77802dd9c8f177f1870323b92960f4c3456eee22 Mon Sep 17 00:00:00 2001 From: mysma-9403 <64923976+mysma-9403@users.noreply.github.com> Date: Thu, 28 May 2026 14:45:41 +0200 Subject: [PATCH] feat(memory-conversations): trigram inverted index for cross-thread search (#2756) --- Cargo.lock | 1 + Cargo.toml | 6 + .../memory_conversations/inverted_index.rs | 609 ++++++++++++++++++ src/openhuman/memory_conversations/mod.rs | 2 + src/openhuman/memory_conversations/store.rs | 135 +++- .../memory_conversations/store_tests.rs | 112 ++++ .../memory_conversations/tokenize.rs | 283 ++++++++ 7 files changed, 1114 insertions(+), 34 deletions(-) create mode 100644 src/openhuman/memory_conversations/inverted_index.rs create mode 100644 src/openhuman/memory_conversations/tokenize.rs diff --git a/Cargo.lock b/Cargo.lock index 869fd932c..7c7855737 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5161,6 +5161,7 @@ dependencies = [ "tracing-appender", "tracing-log", "tracing-subscriber", + "unicode-normalization", "unicode-segmentation", "unicode-width", "url", diff --git a/Cargo.toml b/Cargo.toml index 1ed3417f7..92cc7f9b0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -117,6 +117,12 @@ walkdir = "2" glob = "0.3" unicode-segmentation = "1" unicode-width = "0.2" +# NFKC + combining-mark detection for the cross-thread search inverted +# index (`memory_conversations::tokenize`). NFKC unifies CJK half/full- +# width variants and Arabic presentation forms; `canonical_combining_class` +# lets us strip diacritics across all scripts (Polish ą→a, Arabic harakat, +# Hebrew niqqud, etc.) without per-language tables. +unicode-normalization = "0.1" hostname = "0.4.2" rustls = { version = "0.23", features = ["ring"] } rustls-pki-types = "1.14.0" diff --git a/src/openhuman/memory_conversations/inverted_index.rs b/src/openhuman/memory_conversations/inverted_index.rs new file mode 100644 index 000000000..542de4469 --- /dev/null +++ b/src/openhuman/memory_conversations/inverted_index.rs @@ -0,0 +1,609 @@ +//! In-memory inverted index for `search_cross_thread_messages`. +//! +//! ## Architecture (v1) +//! +//! ```text +//! ┌───────────────┐ +//! query "cat" ───▶ │ Phase 1: │ posting-list intersection +//! │ ngram lookup │ on character n-grams +//! └──────┬────────┘ +//! │ candidate doc ids per term +//! ┌──────▼────────┐ +//! │ Phase 2: │ exact substring verify on +//! │ verify+score │ normalized content, score +//! └──────┬────────┘ by `matched_terms / total_terms` +//! │ +//! ▼ +//! Vec +//! ``` +//! +//! ## Ownership choices for scale +//! +//! Conversation corpora can grow to hundreds of thousands of messages per +//! workspace. The data structures here are picked to keep the resident-set +//! size predictable at that scale: +//! +//! - **`thread_id` and `role` are interned `Arc`** (see +//! `intern_thread_id` / `intern_role`). A workspace with one thread of +//! N messages would otherwise store N copies of the same thread id; a +//! role string only ever takes two distinct values in practice. The +//! interner amortises both to a single heap allocation per distinct +//! value, plus one `Arc` clone per `DocEntry`. +//! - **Posting-map keys are `Box`** (16 bytes) rather than `String` +//! (24 bytes). Saves 8 bytes per distinct ngram in the corpus — at +//! ~17k Latin trigrams plus CJK bigrams that adds up. +//! - **Posting lists are still `BTreeSet`** for ergonomic ordered +//! iteration. The Phase 1 intersection is performed against a +//! single-allocation `Vec` accumulator via a two-pointer +//! sort-merge (no per-iteration `BTreeSet` rebuilds), so the BTreeSet +//! shape only affects insertion and removal, not query latency. +//! Roaring Bitmaps + FST + LSM segments are the long-term destination +//! (Gemini Deep Research write-up); we defer that until corpus sizes +//! justify the complexity. +//! - **Whole index lives in RAM**, rebuilt from JSONL on first access in +//! the process. The JSONL files remain the source of truth. +//! - **Scoring matches the previous linear scan**: +//! `score = matched_terms / total_terms` with a `created_at` tiebreaker. +//! - **Pathological query short-circuit**: if Phase 1 produces a +//! candidate set larger than `LARGE_CANDIDATE_LIMIT` for any term, the +//! index returns recency-ordered hits without running Phase 2. This +//! genuinely caps tail latency — the check fires *before* the +//! substring-verification loop, not after. + +use std::collections::{BTreeSet, HashMap}; +use std::sync::Arc; + +use super::tokenize::{ngrams, normalize}; +use super::types::{ConversationMessage, CrossThreadHit}; + +/// Minimum byte length for a query term to be considered. Matches the +/// historical behaviour of `search_cross_thread_messages` so existing +/// callers (and tests) see no change. Single-byte ASCII tokens like "a" +/// or "is" are filtered out; a single CJK character (3 bytes in UTF-8) +/// passes through. +const MIN_TERM_BYTES: usize = 3; + +/// When Phase 1 returns more than this many candidates we skip Phase 2 +/// verification and fall back to a pure recency-ranked truncation. This +/// is the mitigation for the "user types `e`" pathological case. +const LARGE_CANDIDATE_LIMIT: usize = 10_000; + +/// One indexed message. Carries enough state to (a) reconstruct a +/// `CrossThreadHit` without re-reading JSONL on the hot path and (b) +/// verify Phase 1 candidates by exact substring match on the normalized +/// form. +/// +/// `thread_id` and `role` are `Arc` because they repeat heavily +/// across messages (N messages per thread → N references to the same +/// thread id; only ~2 distinct role values across the entire corpus). +/// `message_id`, `content`, `content_normalized` and `created_at` are +/// per-message unique so they stay as `String`. +#[derive(Debug, Clone)] +struct DocEntry { + thread_id: Arc, + message_id: String, + role: Arc, + content: String, // original, returned verbatim in hits + content_normalized: String, // for Phase 2 substring verification + created_at: String, +} + +/// In-memory trigram/bigram inverted index over conversation messages. +/// +/// Documents are addressed by a dense `u32` doc-id assigned in insertion +/// order. Deletes leave tombstones (`docs[i] = None`) rather than shifting +/// the array, so posting-list integers stay valid without rebuilding. +#[derive(Debug, Default)] +pub(crate) struct InvertedIndex { + /// `ngram -> sorted set of doc-ids`. BTreeSet so per-doc removals + /// are O(log n) and iteration is in sorted order (drives the + /// sort-merge intersect in `candidates_for_term`). Keys are + /// `Box` to shave 8 bytes per entry vs `String`. + postings: HashMap, BTreeSet>, + /// Tombstoned: `docs[i] == None` means the message was deleted. We + /// keep the slot so existing doc-ids in posting lists stay valid. + docs: Vec>, + /// Reverse lookup for incremental removal: `(thread_id, message_id)` + /// → `doc_id`. Letting us drop a single message without re-walking + /// the corpus. + by_message: HashMap<(String, String), u32>, + /// Interner pools. Keep a single `Arc` per distinct thread id + /// and role so every `DocEntry` referencing them can hold a cheap + /// 16-byte `Arc` clone instead of a 24-byte `String`. + thread_id_pool: HashMap>, + role_pool: HashMap>, +} + +impl InvertedIndex { + pub fn new() -> Self { + Self::default() + } + + /// Insert one message. Takes the message by value so the caller's + /// owned strings can be moved into the index without an internal + /// clone of each field. If the (thread, message_id) pair is already + /// in the index this is a no-op (messages are append-only in the + /// store, so duplicate IDs indicate a corrupt JSONL — silently + /// ignore rather than panic). + pub fn insert(&mut self, thread_id: &str, msg: ConversationMessage) { + let ConversationMessage { + id, + content, + sender, + created_at, + message_type: _, + extra_metadata: _, + } = msg; + + let key = (thread_id.to_string(), id.clone()); + if self.by_message.contains_key(&key) { + return; + } + let normalized = normalize(&content); + let doc_id = self.docs.len() as u32; + for ngram in ngrams(&normalized) { + if let Some(posting) = self.postings.get_mut(ngram) { + posting.insert(doc_id); + } else { + let mut set = BTreeSet::new(); + set.insert(doc_id); + self.postings.insert(ngram.into(), set); + } + } + let thread_arc = self.intern_thread_id(thread_id); + let role_arc = self.intern_role(&sender); + self.docs.push(Some(DocEntry { + thread_id: thread_arc, + message_id: id, + role: role_arc, + content, + content_normalized: normalized, + created_at, + })); + self.by_message.insert(key, doc_id); + } + + /// Drop every document belonging to a thread. Used by + /// `delete_thread` and during full purge. + pub fn remove_thread(&mut self, thread_id: &str) { + let to_remove: Vec = self + .by_message + .iter() + .filter(|((t, _), _)| t == thread_id) + .map(|(_, id)| *id) + .collect(); + for doc_id in to_remove { + self.remove_doc(doc_id); + } + self.thread_id_pool.remove(thread_id); + } + + /// Reset the index to its empty state. Cheaper than dropping and + /// re-allocating when a workspace is being rebuilt. + pub fn clear(&mut self) { + self.postings.clear(); + self.docs.clear(); + self.by_message.clear(); + self.thread_id_pool.clear(); + self.role_pool.clear(); + } + + fn remove_doc(&mut self, doc_id: u32) { + let idx = doc_id as usize; + let Some(entry) = self.docs.get_mut(idx).and_then(|slot| slot.take()) else { + return; + }; + self.by_message + .remove(&(entry.thread_id.to_string(), entry.message_id.clone())); + // Remove doc_id from every posting list referencing it. We re- + // tokenize the normalized content rather than tracking the + // per-doc ngram set; tokenization is allocation-free now that + // `ngrams` returns borrowed slices, so this stays cheap. + for ngram in ngrams(&entry.content_normalized) { + if let Some(posting) = self.postings.get_mut(ngram) { + posting.remove(&doc_id); + if posting.is_empty() { + self.postings.remove(ngram); + } + } + } + } + + /// The Phase 1 + Phase 2 query pipeline. Mirrors the contract of + /// `ConversationStore::search_cross_thread_messages` so the store + /// method can be a thin shim. + pub fn search( + &self, + query: &str, + limit: usize, + exclude_thread_id: Option<&str>, + ) -> Vec { + if limit == 0 { + return Vec::new(); + } + let query_lower = normalize(query); + // Filter terms by raw byte length (matches the historical + // 3-byte threshold; single CJK chars are 3 bytes and pass). + let terms: Vec = query_lower + .split_whitespace() + .filter(|t| t.len() >= MIN_TERM_BYTES) + .map(|s| s.to_string()) + .collect(); + if terms.is_empty() { + return Vec::new(); + } + + // Phase 1: collect candidate doc-ids per term. Short-circuit to + // recency-only ordering if any single term's candidate set + // already exceeds the pathological threshold — this is the cap + // on tail latency, and it must fire BEFORE we run the substring + // verification loop. + let mut per_term: Vec> = Vec::with_capacity(terms.len()); + for term in &terms { + let candidates = match self.candidates_for_term(term) { + Some(v) => v, + None => self + .docs + .iter() + .enumerate() + .filter_map(|(i, slot)| slot.as_ref().map(|_| i as u32)) + .collect::>(), + }; + if candidates.len() > LARGE_CANDIDATE_LIMIT { + return self.recency_fallback(exclude_thread_id, limit); + } + per_term.push(candidates); + } + + // Phase 2: verify each candidate by exact substring match. + // Count distinct terms per doc for the score. + let mut hit_counts: HashMap = HashMap::new(); + for (term, candidates) in terms.iter().zip(per_term.into_iter()) { + for doc_id in candidates { + let Some(entry) = self.docs[doc_id as usize].as_ref() else { + continue; + }; + if exclude_thread_id == Some(entry.thread_id.as_ref()) { + continue; + } + if entry.content_normalized.contains(term.as_str()) { + *hit_counts.entry(doc_id).or_insert(0) += 1; + } + } + } + + let total_terms = terms.len() as f64; + let mut hits: Vec = hit_counts + .into_iter() + .map(|(doc_id, matched)| { + let entry = self.docs[doc_id as usize] + .as_ref() + .expect("doc_id from hit_counts must be live"); + CrossThreadHit { + thread_id: entry.thread_id.to_string(), + message_id: entry.message_id.clone(), + role: entry.role.to_string(), + content: entry.content.clone(), + created_at: entry.created_at.clone(), + score: matched as f64 / total_terms, + } + }) + .collect(); + + hits.sort_by(|a, b| { + b.score + .partial_cmp(&a.score) + .unwrap_or(std::cmp::Ordering::Equal) + .then_with(|| b.created_at.cmp(&a.created_at)) + }); + hits.truncate(limit); + hits + } + + /// Build the Phase 1 candidate set for one query term. + /// + /// Returns `Some(vec)` containing the sorted intersection of posting + /// lists for every ngram of `term`. If `term` is too short to + /// produce any ngram (e.g. a single CJK char of length 1) returns + /// `None` so the caller can fall back to a linear scan. + /// + /// The intersect is a two-pointer sort-merge over the already-sorted + /// posting lists: `acc` is rewritten in place once per remaining + /// ngram, allocating zero intermediate sets. + fn candidates_for_term(&self, term: &str) -> Option> { + let term_ngrams = ngrams(term); + if term_ngrams.is_empty() { + return None; + } + let mut iter = term_ngrams.iter(); + let first = iter.next().expect("non-empty by check above"); + let mut acc: Vec = match self.postings.get(*first) { + Some(p) => p.iter().copied().collect(), + None => return Some(Vec::new()), + }; + for ng in iter { + if acc.is_empty() { + return Some(acc); + } + match self.postings.get(*ng) { + Some(p) => intersect_sorted_with_btreeset(&mut acc, p), + None => return Some(Vec::new()), + } + } + Some(acc) + } + + fn intern_thread_id(&mut self, thread_id: &str) -> Arc { + if let Some(existing) = self.thread_id_pool.get(thread_id) { + return Arc::clone(existing); + } + let arc: Arc = Arc::from(thread_id); + self.thread_id_pool + .insert(thread_id.to_string(), Arc::clone(&arc)); + arc + } + + fn intern_role(&mut self, role: &str) -> Arc { + if let Some(existing) = self.role_pool.get(role) { + return Arc::clone(existing); + } + let arc: Arc = Arc::from(role); + self.role_pool.insert(role.to_string(), Arc::clone(&arc)); + arc + } + + fn recency_fallback( + &self, + exclude_thread_id: Option<&str>, + limit: usize, + ) -> Vec { + let mut hits: Vec = self + .docs + .iter() + .filter_map(|slot| slot.as_ref()) + .filter(|entry| exclude_thread_id != Some(entry.thread_id.as_ref())) + .map(|entry| CrossThreadHit { + thread_id: entry.thread_id.to_string(), + message_id: entry.message_id.clone(), + role: entry.role.to_string(), + content: entry.content.clone(), + created_at: entry.created_at.clone(), + // Score 0.0 signals "matched via recency fallback only" — + // documented in the function rustdoc above. Callers + // sorting by `(score desc, created_at desc)` still see + // the newest entries first. + score: 0.0, + }) + .collect(); + hits.sort_by(|a, b| b.created_at.cmp(&a.created_at)); + hits.truncate(limit); + hits + } +} + +/// Two-pointer sort-merge intersect. `acc` and `other` are both sorted +/// ascending; on return `acc` contains only the elements present in +/// both, in the same sorted order. Runs in O(|acc| + |other|) with zero +/// allocations. +fn intersect_sorted_with_btreeset(acc: &mut Vec, other: &BTreeSet) { + let mut other_iter = other.iter().copied().peekable(); + let mut write = 0usize; + for read in 0..acc.len() { + let target = acc[read]; + // Advance `other_iter` past everything strictly less than the + // current `target`. After this loop the next peeked value is + // either equal to `target` (keep) or strictly greater (drop). + while let Some(&o) = other_iter.peek() { + if o < target { + other_iter.next(); + } else { + break; + } + } + if other_iter.peek().copied() == Some(target) { + acc[write] = target; + write += 1; + other_iter.next(); + } + } + acc.truncate(write); +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + fn msg(id: &str, content: &str, created: &str) -> ConversationMessage { + ConversationMessage { + id: id.to_string(), + content: content.to_string(), + message_type: "text".to_string(), + extra_metadata: json!({}), + sender: "user".to_string(), + created_at: created.to_string(), + } + } + + #[test] + fn substring_inside_word_matches() { + // Canonical substring-inside-word case: querying "cat" must find + // "concatenate" — a token-boundary tokenizer would miss this. + let mut idx = InvertedIndex::new(); + idx.insert( + "t1", + msg("m1", "concatenate the strings", "2026-04-10T10:00:00Z"), + ); + + let hits = idx.search("cat", 10, None); + assert_eq!(hits.len(), 1); + assert_eq!(hits[0].message_id, "m1"); + } + + #[test] + fn polish_diacritics_normalized_both_sides() { + let mut idx = InvertedIndex::new(); + idx.insert( + "t1", + msg("m1", "Lecę dziś do Krakowa", "2026-04-10T10:00:00Z"), + ); + // Query with no diacritics finds content with diacritics. + let hits = idx.search("krakow", 10, None); + assert_eq!(hits.len(), 1, "krakow should match Krakowa"); + } + + #[test] + fn japanese_bigram_match() { + let mut idx = InvertedIndex::new(); + idx.insert( + "t1", + msg("m1", "東京タワーが見える", "2026-04-10T10:00:00Z"), + ); + let hits = idx.search("東京", 10, None); + assert_eq!(hits.len(), 1, "two-char CJK query should match"); + } + + #[test] + fn arabic_harakat_stripped() { + let mut idx = InvertedIndex::new(); + // "wrote" with full vocalization vs bare consonants. + idx.insert("t1", msg("m1", "كَتَبَ الطالب", "2026-04-10T10:00:00Z")); + // The bare-consonant form should still find the vocalized one. + let hits = idx.search("كتب", 10, None); + assert_eq!(hits.len(), 1, "harakat stripping should equalize forms"); + } + + #[test] + fn excludes_active_thread() { + let mut idx = InvertedIndex::new(); + idx.insert( + "active", + msg("ma", "postgres deploy", "2026-04-10T10:00:00Z"), + ); + idx.insert( + "other", + msg("mo", "postgres deploy", "2026-04-10T10:01:00Z"), + ); + let hits = idx.search("postgres", 10, Some("active")); + assert_eq!(hits.len(), 1); + assert_eq!(hits[0].thread_id, "other"); + } + + #[test] + fn empty_query_returns_empty() { + let mut idx = InvertedIndex::new(); + idx.insert("t", msg("m", "anything here", "2026-04-10T10:00:00Z")); + assert!(idx.search("", 10, None).is_empty()); + } + + #[test] + fn short_terms_only_returns_empty() { + // Mirrors the legacy `search_cross_thread_messages_skips_short_terms` + // behaviour — terms < 3 bytes are dropped. + let mut idx = InvertedIndex::new(); + idx.insert("t", msg("m", "Postgres", "2026-04-10T10:00:00Z")); + assert!(idx.search("a is on", 10, None).is_empty()); + } + + #[test] + fn limit_zero_returns_empty() { + let mut idx = InvertedIndex::new(); + idx.insert("t", msg("m", "Postgres", "2026-04-10T10:00:00Z")); + assert!(idx.search("postgres", 0, None).is_empty()); + } + + #[test] + fn score_matches_legacy_semantics() { + // 5-term query, 2 substring matches → score = 0.4. + let mut idx = InvertedIndex::new(); + idx.insert( + "t1", + msg( + "m1", + "Remember: my project is called Phoenix and uses Go and PostgreSQL.", + "2026-04-10T10:00:00Z", + ), + ); + let hits = idx.search("What database does my project use", 10, None); + assert_eq!(hits.len(), 1); + // "project" + "use" (substring of "uses") → 2 of 5 terms. + assert!( + (hits[0].score - 0.4).abs() < 1e-9, + "score = {}", + hits[0].score + ); + } + + #[test] + fn remove_thread_drops_all_messages() { + let mut idx = InvertedIndex::new(); + idx.insert("t1", msg("m1", "postgres deploy", "2026-04-10T10:00:00Z")); + idx.insert("t1", msg("m2", "postgres backup", "2026-04-10T10:01:00Z")); + idx.insert("t2", msg("m3", "postgres replica", "2026-04-10T10:02:00Z")); + assert_eq!(idx.search("postgres", 10, None).len(), 3); + idx.remove_thread("t1"); + let hits = idx.search("postgres", 10, None); + assert_eq!(hits.len(), 1); + assert_eq!(hits[0].thread_id, "t2"); + } + + #[test] + fn duplicate_message_id_is_idempotent() { + let mut idx = InvertedIndex::new(); + let m = msg("m1", "postgres deploy", "2026-04-10T10:00:00Z"); + idx.insert("t1", m.clone()); + idx.insert("t1", m); // dup — should be ignored + let hits = idx.search("postgres", 10, None); + assert_eq!(hits.len(), 1, "duplicate insert must not duplicate hits"); + } + + #[test] + fn pathological_query_short_circuits_to_recency() { + // Build a corpus big enough to trip LARGE_CANDIDATE_LIMIT. + let mut idx = InvertedIndex::new(); + let n = LARGE_CANDIDATE_LIMIT + 50; + for i in 0..n { + // "common" appears in every doc, so the trigram "com" hits + // all of them. + let created = format!("2026-04-10T10:{:02}:{:02}Z", i / 60, i % 60); + idx.insert("bulk", msg(&format!("m{i}"), "common payload", &created)); + } + let hits = idx.search("common", 5, None); + assert_eq!(hits.len(), 5, "fallback must still respect limit"); + // Score 0.0 is the recency-fallback marker. + assert!(hits.iter().all(|h| h.score == 0.0)); + } + + #[test] + fn intern_pool_dedupes_repeated_thread_id_and_role() { + // Smoke-test the Arc interning: many messages on one + // thread must share a single Arc backing the thread_id. + let mut idx = InvertedIndex::new(); + for i in 0..5 { + let ts = format!("2026-04-10T10:00:{:02}Z", i); + idx.insert("shared-thread", msg(&format!("m{i}"), "payload", &ts)); + } + assert_eq!(idx.thread_id_pool.len(), 1); + assert_eq!(idx.role_pool.len(), 1); + // After removing the thread, the pool entry is dropped too. + idx.remove_thread("shared-thread"); + assert_eq!(idx.thread_id_pool.len(), 0); + } + + #[test] + fn intersect_sorted_with_btreeset_basic() { + let other: BTreeSet = [2, 4, 5, 9].into_iter().collect(); + let mut acc: Vec = vec![1, 2, 3, 5, 7, 9]; + intersect_sorted_with_btreeset(&mut acc, &other); + assert_eq!(acc, vec![2, 5, 9]); + } + + #[test] + fn intersect_sorted_with_btreeset_empty_other() { + let other: BTreeSet = BTreeSet::new(); + let mut acc: Vec = vec![1, 2, 3]; + intersect_sorted_with_btreeset(&mut acc, &other); + assert!(acc.is_empty()); + } +} diff --git a/src/openhuman/memory_conversations/mod.rs b/src/openhuman/memory_conversations/mod.rs index c5f3c1e8c..43a551008 100644 --- a/src/openhuman/memory_conversations/mod.rs +++ b/src/openhuman/memory_conversations/mod.rs @@ -10,7 +10,9 @@ //! this module as `memory::conversations` during the migration. mod bus; +mod inverted_index; mod store; +mod tokenize; mod types; pub use bus::register_conversation_persistence_subscriber; diff --git a/src/openhuman/memory_conversations/store.rs b/src/openhuman/memory_conversations/store.rs index f2539d6cb..8ec83682a 100644 --- a/src/openhuman/memory_conversations/store.rs +++ b/src/openhuman/memory_conversations/store.rs @@ -5,7 +5,7 @@ //! All on-disk mutations serialise through a single process-wide mutex so //! concurrent RPC handlers don't interleave writes. -use std::collections::BTreeMap; +use std::collections::{BTreeMap, HashMap}; use std::fs::{self, File, OpenOptions}; use std::hash::{Hash, Hasher}; use std::io::{BufRead, BufReader, Write}; @@ -16,6 +16,7 @@ use once_cell::sync::Lazy; use parking_lot::Mutex; use tempfile::NamedTempFile; +use super::inverted_index::InvertedIndex; use super::types::{ ConversationMessage, ConversationMessagePatch, ConversationThread, CreateConversationThread, CrossThreadHit, @@ -26,6 +27,28 @@ const THREADS_FILENAME: &str = "threads.jsonl"; const THREAD_MESSAGES_DIR: &str = "threads"; static CONVERSATION_STORE_LOCK: Lazy> = Lazy::new(|| Mutex::new(())); +/// Per-workspace inverted index cache. Keyed by the workspace's +/// `memory/conversations` root so multiple `ConversationStore` clones +/// pointing at the same workspace share one index. The cache outlives +/// individual store handles (which are cloneable PathBuf wrappers); it +/// is bounded by the number of distinct workspaces a single process +/// touches, which in practice is one. Tests using `TempDir` paths leave +/// behind dead entries when the dir is removed — acceptable for an +/// in-process cache. +/// +/// # Lock ordering (INVARIANT) +/// +/// `CONVERSATION_STORE_LOCK` MUST be acquired before +/// `CONVERSATION_INDEX_CACHE`. Every public store method takes the outer +/// lock at the top of its body and only then reaches for the cache, so +/// in current code the inner mutex is uncontended. Any future caller +/// that violates this order — taking the cache lock first, or taking +/// the cache lock without holding the outer lock and then reaching +/// across into another store API — risks a deadlock. Keep the inner +/// mutex strictly nested inside the outer one. +static CONVERSATION_INDEX_CACHE: Lazy>> = + Lazy::new(|| Mutex::new(HashMap::new())); + fn redact_title_for_log(title: &str) -> String { let mut hasher = std::collections::hash_map::DefaultHasher::new(); title.hash(&mut hasher); @@ -148,6 +171,14 @@ impl ConversationStore { /// chat continuity needs a direct cross-thread reader to surface /// context the user shared in chat A when they ask a dependent /// question in chat B. + /// + /// Backed by an in-memory trigram/CJK-bigram inverted index + /// (`super::inverted_index`). The legacy implementation walked every + /// JSONL file and did `content.to_lowercase().contains(term)` per + /// message, which is O(threads × messages × content_len). The index + /// turns that into O(|posting lists|) for typical queries while + /// preserving the previous scoring contract + /// (`score = matched_terms / total_terms`, recency tiebreak). pub fn search_cross_thread_messages( &self, query: &str, @@ -155,26 +186,47 @@ impl ConversationStore { exclude_thread_id: Option<&str>, ) -> Result, String> { let _guard = CONVERSATION_STORE_LOCK.lock(); - let query_lower = query.to_lowercase(); - let terms: Vec<&str> = query_lower - .split_whitespace() - .filter(|t| t.len() >= 3) - .collect(); - if terms.is_empty() { - return Ok(Vec::new()); + self.with_index(|idx| idx.search(query, limit, exclude_thread_id)) + } + + /// Acquire the cached inverted index for this workspace (building it + /// from JSONL on first access) and run `f` against it. Caller MUST + /// hold `CONVERSATION_STORE_LOCK` for the duration of the closure. + fn with_index(&self, f: impl FnOnce(&mut InvertedIndex) -> R) -> Result { + let key = self.root_dir(); + let mut cache = CONVERSATION_INDEX_CACHE.lock(); + if !cache.contains_key(&key) { + let mut idx = InvertedIndex::new(); + self.populate_index_unlocked(&mut idx)?; + cache.insert(key.clone(), idx); } + let idx = cache.get_mut(&key).expect("inserted above if absent"); + Ok(f(idx)) + } + + /// Walk every per-thread JSONL file in the workspace and insert each + /// message into `idx`. Used on first access to a workspace's index; + /// also called after `purge_threads` to reset the cache from a + /// known-empty state. The JSONL files remain the source of truth, so + /// a rebuild after a process crash is always safe. + fn populate_index_unlocked(&self, idx: &mut InvertedIndex) -> Result<(), String> { + // `list_threads_unlocked` already handles a fresh workspace: + // `ensure_root` creates the directory + threads log if missing, and + // `read_jsonl` returns an empty Vec for an empty file. Anything that + // still bubbles up here is a real filesystem/setup failure and must + // propagate — silently returning Ok would mask it and make search + // appear to return zero results for an undiagnosed reason. let threads = self.list_threads_unlocked()?; - let mut hits: Vec = Vec::new(); for thread in threads { - if exclude_thread_id == Some(thread.id.as_str()) { + let path = self.thread_messages_path(&thread.id); + if !path.exists() { continue; } - let path = self.thread_messages_path(&thread.id); let messages = match read_jsonl::(&path) { Ok(m) => m, Err(err) => { tracing::warn!( - "[conversations] cross-thread scan skipped unreadable file path={} error={}", + "{LOG_PREFIX} index build skipped unreadable file path={} error={}", path.display(), err ); @@ -182,30 +234,16 @@ impl ConversationStore { } }; for msg in messages { - let content_lower = msg.content.to_lowercase(); - let matched = terms.iter().filter(|t| content_lower.contains(*t)).count(); - if matched == 0 { - continue; - } - let score = matched as f64 / terms.len() as f64; - hits.push(CrossThreadHit { - thread_id: thread.id.clone(), - message_id: msg.id, - role: msg.sender, - content: msg.content, - created_at: msg.created_at, - score, - }); + // Move each freshly-deserialized message straight into + // the index; no per-field clones on the rebuild path. + idx.insert(&thread.id, msg); } } - hits.sort_by(|a, b| { - b.score - .partial_cmp(&a.score) - .unwrap_or(std::cmp::Ordering::Equal) - .then_with(|| b.created_at.cmp(&a.created_at)) - }); - hits.truncate(limit); - Ok(hits) + debug!( + "{LOG_PREFIX} inverted index populated workspace={}", + self.root_dir().display() + ); + Ok(()) } /// Append a message to the thread's JSONL file. Errors if the thread is missing. @@ -235,6 +273,20 @@ impl ConversationStore { last_message_at: message.created_at.clone(), }, )?; + // Keep the inverted index in sync. We only update if the index + // has already been materialized for this workspace — otherwise + // the next search will lazily rebuild and pick up this message + // anyway, and we avoid paying the rebuild cost on a write path. + // `insert` takes the message by value so it can move owned + // fields straight into its `DocEntry`; we clone here only when + // the cache is actually warm, paying for one extra owned copy + // (instead of cloning each field inside `insert`). + { + let mut cache = CONVERSATION_INDEX_CACHE.lock(); + if let Some(idx) = cache.get_mut(&self.root_dir()) { + idx.insert(thread_id, message.clone()); + } + } debug!( "{LOG_PREFIX} appended message thread_id={} message_id={} path={}", thread_id, @@ -370,6 +422,14 @@ impl ConversationStore { )); } } + // Drop every indexed message for this thread so future searches + // don't surface stale content. + { + let mut cache = CONVERSATION_INDEX_CACHE.lock(); + if let Some(idx) = cache.get_mut(&self.root_dir()) { + idx.remove_thread(thread_id); + } + } debug!( "{LOG_PREFIX} deleted thread id={} path={}", thread_id, @@ -388,6 +448,13 @@ impl ConversationStore { .map_err(|e| format!("remove conversation dir {}: {e}", root.display()))?; } self.ensure_root()?; + // Drop the cached inverted index — the workspace is now empty, + // and any next search will lazily rebuild from the (now empty) + // JSONL tree. + { + let mut cache = CONVERSATION_INDEX_CACHE.lock(); + cache.remove(&root); + } debug!( "{LOG_PREFIX} purged threads={} messages={} root={}", stats.thread_count, diff --git a/src/openhuman/memory_conversations/store_tests.rs b/src/openhuman/memory_conversations/store_tests.rs index 3975464f7..b5b6e73e2 100644 --- a/src/openhuman/memory_conversations/store_tests.rs +++ b/src/openhuman/memory_conversations/store_tests.rs @@ -725,6 +725,118 @@ fn search_cross_thread_messages_skips_short_terms_and_empty_queries() { .is_empty()); } +#[test] +fn search_cross_thread_messages_finds_polish_substring_without_diacritics() { + let (_temp, store) = make_store(); + store + .ensure_thread(CreateConversationThread { + parent_thread_id: None, + id: "thread-pl".to_string(), + title: "PL".to_string(), + created_at: "2026-04-10T12:00:00Z".to_string(), + labels: None, + }) + .unwrap(); + store + .append_message( + "thread-pl", + ConversationMessage { + id: "m1".to_string(), + content: "Lecę w piątek do Łodzi a potem Krakowa".to_string(), + message_type: "text".to_string(), + extra_metadata: json!({}), + sender: "user".to_string(), + created_at: "2026-04-10T12:01:00Z".to_string(), + }, + ) + .unwrap(); + + // Query without diacritics should still find content with them. + let hits = store + .search_cross_thread_messages("Lodzi", 10, None) + .expect("cross-thread search"); + assert_eq!(hits.len(), 1, "ł-fold should match Łodzi via lodzi"); + + let hits = store + .search_cross_thread_messages("krakow", 10, None) + .expect("cross-thread search"); + assert_eq!(hits.len(), 1, "diacritic strip should match Krakowa"); +} + +#[test] +fn search_cross_thread_messages_finds_japanese_bigram_match() { + let (_temp, store) = make_store(); + store + .ensure_thread(CreateConversationThread { + parent_thread_id: None, + id: "thread-jp".to_string(), + title: "JP".to_string(), + created_at: "2026-04-10T12:00:00Z".to_string(), + labels: None, + }) + .unwrap(); + store + .append_message( + "thread-jp", + ConversationMessage { + id: "m1".to_string(), + content: "明日東京に行きます".to_string(), // "Tomorrow I'm going to Tokyo" + message_type: "text".to_string(), + extra_metadata: json!({}), + sender: "user".to_string(), + created_at: "2026-04-10T12:01:00Z".to_string(), + }, + ) + .unwrap(); + + let hits = store + .search_cross_thread_messages("東京", 10, None) + .expect("cross-thread search"); + assert_eq!(hits.len(), 1, "CJK bigram lookup should find 東京"); + assert_eq!(hits[0].message_id, "m1"); +} + +#[test] +fn search_cross_thread_messages_rebuilds_index_from_jsonl_after_reopen() { + // First store handle writes messages, second handle (simulating + // process restart on the same workspace dir) must lazy-rebuild the + // index from JSONL and still answer search queries. + let temp = TempDir::new().expect("tempdir"); + let workspace = temp.path().to_path_buf(); + { + let store = ConversationStore::new(workspace.clone()); + store + .ensure_thread(CreateConversationThread { + parent_thread_id: None, + id: "thread-x".to_string(), + title: "X".to_string(), + created_at: "2026-04-10T12:00:00Z".to_string(), + labels: None, + }) + .unwrap(); + store + .append_message( + "thread-x", + ConversationMessage { + id: "m1".to_string(), + content: "persisted across reopen — checksum kitten".to_string(), + message_type: "text".to_string(), + extra_metadata: json!({}), + sender: "user".to_string(), + created_at: "2026-04-10T12:01:00Z".to_string(), + }, + ) + .unwrap(); + } + // The cache key is per-workspace path; this TempDir was never seen + // before, so a fresh store handle will trigger a lazy rebuild. + let reopened = ConversationStore::new(workspace); + let hits = reopened + .search_cross_thread_messages("kitten", 10, None) + .expect("cross-thread search"); + assert_eq!(hits.len(), 1, "reopened store must rebuild index from disk"); +} + #[test] fn update_thread_labels_missing_thread_returns_error() { let (_temp, store) = make_store(); diff --git a/src/openhuman/memory_conversations/tokenize.rs b/src/openhuman/memory_conversations/tokenize.rs new file mode 100644 index 000000000..da7ba3523 --- /dev/null +++ b/src/openhuman/memory_conversations/tokenize.rs @@ -0,0 +1,283 @@ +//! Multilingual normalization + character n-gram generation for the +//! cross-thread search inverted index. +//! +//! ## Why character n-grams (not word tokens)? +//! +//! The cross-thread search must find substrings *inside* words — querying +//! `cat` should return messages containing `concatenate` or `Kotlin`. A +//! whitespace/word-boundary tokenizer fundamentally cannot do that, and +//! it also breaks down for CJK scripts that have no whitespace at all. +//! Character n-grams sidestep both problems. +//! +//! ## Why a hybrid trigram + CJK-bigram scheme? +//! +//! Trigrams strike a good balance between recall and dictionary size for +//! alphabetic scripts (~26³ ≈ 17k Latin trigrams). For CJK scripts the +//! alphabet is tens of thousands of characters, so character trigrams +//! explode the dictionary while character bigrams stay tractable. The +//! Gemini Deep Research write-up (see PR description) flagged this as the +//! single most important multilingual mitigation. We therefore generate: +//! +//! - **bigrams** for contiguous runs of CJK characters (Han / Hiragana / +//! Katakana / Hangul), +//! - **trigrams** for everything else. +//! +//! Tokens from both schemes coexist in the same posting map. As long as +//! query-time tokenization runs the *same* code path, lookups stay +//! consistent. +//! +//! ## Normalization pipeline +//! +//! Implemented in `normalize()` as a single iterator chain. The order +//! matters — strip-marks must run on the decomposed form, lowercase must +//! run on stripped code points, and the final NFKC re-compose unifies +//! compatibility variants for byte-stable indexing/querying: +//! +//! 1. **NFKD** — decompose so combining marks become standalone code points +//! (Polish ą → `a` + ̨, Arabic kataba+harakat → base letters + marks). +//! 2. **Strip combining marks** — uses `canonical_combining_class` to +//! drop diacritics across all scripts (Polish ą→a, ć→c; Arabic harakat; +//! Hebrew niqqud; combining tone marks; etc.) without needing +//! per-language tables. +//! 3. **Lowercase** — Unicode-aware case folding for cross-alphabet +//! case insensitivity. +//! 4. **NFKC** — re-compose to canonical form (and unify compatibility +//! characters: half/full-width CJK variants, Arabic presentation forms, +//! ligatures) so byte equality lines up at lookup time. +//! 5. **Non-decomposing fold** (`fold_non_decomposing`) — small per-letter +//! table for decorated letters NFKD leaves untouched (Polish ł, German +//! ß, Norwegian ø, Icelandic þ/ð, Latin æ/œ, Turkish ı, Croatian đ, +//! Maltese ħ, Sami ŋ). + +use unicode_normalization::char::canonical_combining_class; +use unicode_normalization::UnicodeNormalization; + +/// Normalize a piece of text for indexing or querying. Idempotent: running +/// the output through `normalize` again yields the same string. +/// +/// After the standard NFKD + strip-combining-marks pass we additionally +/// fold a small set of letters that *look* decorated but don't decompose +/// canonically (so NFKD alone leaves them unchanged). Polish ł/Ł is the +/// motivating case — a Polish user typing `lacka` reasonably expects to +/// find `łącka`. Same idea for German ß, Norwegian ø, Icelandic þ/ð and +/// Latin æ. +pub fn normalize(text: &str) -> String { + let stripped: String = text + .nfkd() // decompose so combining marks become standalone code points + .filter(|c| canonical_combining_class(*c) == 0) + .flat_map(char::to_lowercase) + .nfkc() // re-compose to canonical form for downstream byte equality + .collect(); + fold_non_decomposing(&stripped) +} + +/// Apply per-letter folds for non-decomposing "decorated" letters that +/// NFKD leaves untouched. Run only after lowercase + NFKD so we don't +/// need uppercase entries. +fn fold_non_decomposing(s: &str) -> String { + let mut out = String::with_capacity(s.len()); + for c in s.chars() { + match c { + 'ł' => out.push('l'), + 'ø' => out.push('o'), + 'ß' => out.push_str("ss"), + 'æ' => out.push_str("ae"), + 'œ' => out.push_str("oe"), + 'þ' => out.push_str("th"), + 'ð' => out.push('d'), + 'đ' => out.push('d'), + 'ħ' => out.push('h'), + 'ı' => out.push('i'), // Turkish dotless i + 'ŋ' => out.push('n'), + other => out.push(other), + } + } + out +} + +/// Returns true for code points that should be tokenized as CJK bigrams. +/// +/// Covers Han ideographs (CJK Unified + Ext A + Compatibility), Japanese +/// kana (Hiragana, Katakana), and Hangul (Jamo + precomposed syllables). +/// CJK punctuation and symbols (U+3000..=U+303F) are intentionally +/// excluded — they should be treated as token delimiters, not content. +pub fn is_cjk(c: char) -> bool { + matches!( + c as u32, + 0x3040..=0x309F // Hiragana + | 0x30A0..=0x30FF // Katakana + | 0x3400..=0x4DBF // CJK Unified Ideographs Extension A + | 0x4E00..=0x9FFF // CJK Unified Ideographs + | 0xF900..=0xFAFF // CJK Compatibility Ideographs + | 0x1100..=0x11FF // Hangul Jamo + | 0xAC00..=0xD7AF // Hangul Syllables + ) +} + +/// Tokenize already-normalized text into character n-grams as borrowed +/// slices into `normalized`. +/// +/// Returning `&str` (instead of owned `String`s) keeps the hot search +/// path allocation-free: query-time ngram extraction only needs to look +/// up posting-list keys, never to insert. On the insert side, the index +/// allocates a fresh key only when an ngram is brand-new to the corpus — +/// see `InvertedIndex::insert`. +/// +/// - CJK runs (≥2 chars) → bigrams. +/// - Non-CJK runs (≥3 chars) → trigrams. +/// - Runs shorter than the relevant n are dropped (they cannot be +/// substring-matched against any document containing them anyway, so +/// the Phase 2 verification will catch them via the linear fallback in +/// `InvertedIndex::search`). +/// +/// Word boundaries inside a run do NOT split the n-gram window — we +/// deliberately want substring matches that span punctuation. +pub fn ngrams(normalized: &str) -> Vec<&str> { + let mut out = Vec::new(); + // Capture (byte_offset, is_cjk) per char. Byte offsets let us slice + // `normalized` directly to return `&str` views; the cjk flag drives + // the script-class run partitioning below. + let chars: Vec<(usize, bool)> = normalized + .char_indices() + .map(|(b, c)| (b, is_cjk(c))) + .collect(); + if chars.is_empty() { + return out; + } + let end_byte = normalized.len(); + + // Walk contiguous runs of "same script class" (CJK vs non-CJK) and + // emit the appropriate n-gram size for each run. + let mut i = 0; + while i < chars.len() { + let cjk = chars[i].1; + let mut j = i + 1; + while j < chars.len() && chars[j].1 == cjk { + j += 1; + } + let n = if cjk { 2 } else { 3 }; + if j - i >= n { + for k in i..=j - n { + let start = chars[k].0; + let end = if k + n < chars.len() { + chars[k + n].0 + } else { + end_byte + }; + out.push(&normalized[start..end]); + } + } + i = j; + } + out +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn normalize_lowercases_ascii() { + assert_eq!(normalize("Hello World"), "hello world"); + } + + #[test] + fn normalize_strips_polish_diacritics() { + // Polish: Kraków, żółć, łąka, plus a Spanish parity case. + // ł/Ł are *not* canonically decomposable (no combining ogonek + // form) so we fold them manually via `fold_non_decomposing`. + assert_eq!(normalize("Kraków"), "krakow"); + assert_eq!(normalize("żółć"), "zolc"); + assert_eq!(normalize("łąka"), "laka"); + assert_eq!(normalize("Mañana"), "manana"); + } + + #[test] + fn normalize_folds_non_decomposing_letters() { + // Letters that NFKD leaves untouched but a user typing without + // the diacritic would still expect to find. + assert_eq!(normalize("Łódź"), "lodz"); + assert_eq!(normalize("Straße"), "strasse"); + assert_eq!(normalize("Bjørn"), "bjorn"); + assert_eq!(normalize("Þórr"), "thorr"); + } + + #[test] + fn normalize_strips_arabic_harakat() { + // Arabic: word "kataba" (he wrote) with harakat marks vs without + let with_marks = "كَتَبَ"; + let without_marks = "كتب"; + assert_eq!(normalize(with_marks), normalize(without_marks)); + } + + #[test] + fn normalize_unifies_cjk_halfwidth_fullwidth() { + // NFKC maps half-width katakana to full-width. + let halfwidth = "ｶﾀｶﾅ"; // half-width + let fullwidth = "カタカナ"; // full-width + assert_eq!(normalize(halfwidth), normalize(fullwidth)); + } + + #[test] + fn normalize_is_idempotent() { + let s = "Café — 東京 — żółć"; + let once = normalize(s); + let twice = normalize(&once); + assert_eq!(once, twice); + } + + #[test] + fn ngrams_emits_trigrams_for_latin() { + let g = ngrams("kitten"); + assert_eq!(g, vec!["kit", "itt", "tte", "ten"]); + } + + #[test] + fn ngrams_emits_bigrams_for_cjk() { + // 日本語 → 日本, 本語 + let g = ngrams("日本語"); + assert_eq!(g, vec!["日本", "本語"]); + } + + #[test] + fn ngrams_mixed_script_splits_at_boundary() { + // "東京tokyo" → CJK run [東京] gives bigram "東京", + // Latin run [tokyo] gives trigrams tok, oky, kyo. + let g = ngrams("東京tokyo"); + assert_eq!(g, vec!["東京", "tok", "oky", "kyo"]); + } + + #[test] + fn ngrams_drops_runs_too_short() { + // "ab東" → Latin run [ab] is only 2 chars → dropped; CJK run [東] + // is only 1 char → dropped. Empty result. + let g = ngrams("ab東"); + assert!(g.is_empty(), "got {:?}", g); + } + + #[test] + fn ngrams_substring_inside_word_is_indexable() { + // After normalize, "Concatenate" → "concatenate" → trigrams include + // "cat". This is the canonical substring-inside-word scenario that + // motivates the character-n-gram scheme over word tokenization. + let normalized = normalize("Concatenate"); + let g = ngrams(&normalized); + assert!(g.contains(&"cat"), "trigrams: {:?}", g); + } + + #[test] + fn ngrams_empty_input_returns_empty() { + assert!(ngrams("").is_empty()); + } + + #[test] + fn is_cjk_classifies_common_scripts() { + assert!(is_cjk('東')); + assert!(is_cjk('あ')); // hiragana + assert!(is_cjk('カ')); // katakana + assert!(is_cjk('한')); // hangul syllable + assert!(!is_cjk('a')); + assert!(!is_cjk('ą')); + assert!(!is_cjk('，')); // CJK punctuation — intentionally NOT cjk + } +}