mirror of
https://github.com/tinyhumansai/openhuman.git
synced 2026-07-27 21:08:00 +00:00
feat(memory-conversations): trigram inverted index for cross-thread search (#2756)
This commit is contained in:
mysma-9403
Generated
+1
@@ -5161,6 +5161,7 @@ dependencies = [
|
||||
"tracing-appender",
|
||||
"tracing-log",
|
||||
"tracing-subscriber",
|
||||
"unicode-normalization",
|
||||
"unicode-segmentation",
|
||||
"unicode-width",
|
||||
"url",
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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<Hit>
|
||||
//! ```
|
||||
//!
|
||||
//! ## 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<str>`** (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<str>`** (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<u32>`** for ergonomic ordered
|
||||
//! iteration. The Phase 1 intersection is performed against a
|
||||
//! single-allocation `Vec<u32>` 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<str>` 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<str>,
|
||||
message_id: String,
|
||||
role: Arc<str>,
|
||||
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<str>` to shave 8 bytes per entry vs `String`.
|
||||
postings: HashMap<Box<str>, BTreeSet<u32>>,
|
||||
/// Tombstoned: `docs[i] == None` means the message was deleted. We
|
||||
/// keep the slot so existing doc-ids in posting lists stay valid.
|
||||
docs: Vec<Option<DocEntry>>,
|
||||
/// 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<str>` 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<String, Arc<str>>,
|
||||
role_pool: HashMap<String, Arc<str>>,
|
||||
}
|
||||
|
||||
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<u32> = 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<CrossThreadHit> {
|
||||
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<String> = 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<u32>> = 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::<Vec<u32>>(),
|
||||
};
|
||||
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<u32, usize> = 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<CrossThreadHit> = 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<Vec<u32>> {
|
||||
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<u32> = 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<str> {
|
||||
if let Some(existing) = self.thread_id_pool.get(thread_id) {
|
||||
return Arc::clone(existing);
|
||||
}
|
||||
let arc: Arc<str> = 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<str> {
|
||||
if let Some(existing) = self.role_pool.get(role) {
|
||||
return Arc::clone(existing);
|
||||
}
|
||||
let arc: Arc<str> = 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<CrossThreadHit> {
|
||||
let mut hits: Vec<CrossThreadHit> = 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<u32>, other: &BTreeSet<u32>) {
|
||||
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<str> 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<u32> = [2, 4, 5, 9].into_iter().collect();
|
||||
let mut acc: Vec<u32> = 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<u32> = BTreeSet::new();
|
||||
let mut acc: Vec<u32> = vec![1, 2, 3];
|
||||
intersect_sorted_with_btreeset(&mut acc, &other);
|
||||
assert!(acc.is_empty());
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
@@ -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<Mutex<()>> = 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<Mutex<HashMap<PathBuf, InvertedIndex>>> =
|
||||
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<Vec<CrossThreadHit>, 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<R>(&self, f: impl FnOnce(&mut InvertedIndex) -> R) -> Result<R, String> {
|
||||
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<CrossThreadHit> = 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::<ConversationMessage>(&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,
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user