diff --git a/BENCHMARKS.md b/BENCHMARKS.md index c1fe299..2701517 100644 --- a/BENCHMARKS.md +++ b/BENCHMARKS.md @@ -139,6 +139,26 @@ build: 36472.8 ms (2742 vectors/s) · exact scan: 40 QPS, p50 24644 µs | 128 | 0.9990 | 7633 | 126 | 248 | | 256 | 0.9990 | 2823 | 352 | 510 | +### After: persistent keyword index, no store rewrite per query + +`hybrid_search` used to rebuild the BM25 index from scratch (re-tokenising every +record) and rewrite the whole `.h5` file on **every query**. The index is now +kept for the life of the store and updated incrementally, and activation boosts +are persisted by the next checkpoint instead of inside the query. Steady-state +p50: **5.5 → 0.24 ms** (1K), **49 → 2.1 ms** (10K), **884 → 23 ms** (100K). + +The first query after `open()` is slower than before (it pays for the better — +slower — HNSW build plus the one-off keyword index build); persisting the HNSW +index removes that. + +### End to end: `HDF5Memory::hybrid_search` (k = 10, weights 0.7 / 0.3) + +| N | ingest ms | checkpoint ms | open ms | first query ms | p50 ms | p99 ms | QPS | +|---:|---:|---:|---:|---:|---:|---:|---:| +| 1000 | 11 | 3.8 | 0.9 | 195.9 | 0.24 | 0.27 | 4130.4 | +| 10000 | 104 | 31.1 | 10.9 | 2627.1 | 2.09 | 2.11 | 479.5 | +| 100000 | 1436 | 684.7 | 278.0 | 36308.1 | 22.90 | 25.46 | 43.5 | + ## Vector Search Latency Brute-force cosine similarity over 384-dimensional embeddings (OpenAI text-embedding-3-small size). diff --git a/CHANGELOG.md b/CHANGELOG.md index cc08329..e5f16ec 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,20 @@ 0.98 and responds to `ef`. Builds are slower (~3.5x at 10K). Existing persisted indexes keep their old graph until rebuilt; the agent rebuilds its index from the cache, so stores pick this up automatically. +- `clawhdf5-agent`: **`hybrid_search` is 23-39x faster in steady state** (p50 + 5.5 -> 0.24 ms at 1K records, 49 -> 2.1 ms at 10K, 884 -> 23 ms at 100K). + Every query used to rebuild the BM25 index from scratch and rewrite the whole + `.h5` file. The keyword index now lives for the life of the store and is + updated incrementally (add / remove / in-place update, exactly equivalent to + a fresh build - property-tested), and a query no longer writes the store. + **Behaviour change:** Hebbian activation boosts are persisted by the next + checkpoint (any flushing write, `flush_wal`, or drop) rather than + immediately; a crash in between forgets only the boosts since the last + checkpoint. Activation weights are now capped (16.0) - they grew without + bound. +- `clawhdf5-agent`: BM25 results are deterministic (ties break by record id), + top-k uses a bounded heap, and the "WAND early termination" that computed a + bound and then ignored it is gone. IDF is computed per query. - `clawhdf5-bench`: new `search_harness` binary — HNSW recall@10 / QPS / latency per `ef` against an exact scan, and end-to-end `hybrid_search` timings, on deterministic clustered (or `--uniform`) data. Baseline in `BENCHMARKS.md`. diff --git a/crates/clawhdf5-agent/src/bm25.rs b/crates/clawhdf5-agent/src/bm25.rs index 258df33..1aff2d7 100644 --- a/crates/clawhdf5-agent/src/bm25.rs +++ b/crates/clawhdf5-agent/src/bm25.rs @@ -3,10 +3,15 @@ //! Provides a standard BM25 (Okapi BM25) implementation with an in-memory //! inverted index. Tombstoned documents are excluded from indexing and search. //! -//! Optimizations: -//! - Cached IDF scores (don't recompute per query) -//! - Sorted posting lists by doc_id for cache-friendly access -//! - Block-Max WAND early termination +//! The index is **incremental**: [`BM25Index::add_document`] and +//! [`BM25Index::remove_document`] keep it exactly equivalent to one built from +//! scratch over the same live documents, so a store can maintain one index for +//! its lifetime instead of re-tokenising the whole corpus per query. To make +//! that possible IDF is computed at query time (it depends on the live +//! document count) rather than cached at build time. +//! +//! - Posting lists sorted by doc id +//! - Bounded-heap top-k; results ordered by score, then doc id (deterministic) use std::cmp::Reverse; use std::collections::{BinaryHeap, HashMap}; @@ -41,10 +46,11 @@ const DEFAULT_B: f32 = 0.75; pub struct BM25Index { /// Inverted index: token -> sorted list of (doc_id, term_frequency). inverted: HashMap>, - /// Cached IDF scores per token. - idf_cache: HashMap, /// Number of tokens in each document (0 for tombstoned docs). doc_lengths: Vec, + /// Sum of `doc_lengths` over live documents (keeps `avg_dl` exact under + /// incremental updates). + total_length: u64, /// Average document length across non-tombstoned docs. avg_dl: f32, /// Number of non-tombstoned documents. @@ -60,8 +66,8 @@ impl BM25Index { pub fn build(documents: &[String], tombstones: &[u8]) -> Self { let mut index = Self { inverted: HashMap::new(), - idf_cache: HashMap::new(), doc_lengths: vec![0; documents.len()], + total_length: 0, avg_dl: 0.0, num_docs: 0, k1: DEFAULT_K1, @@ -81,103 +87,135 @@ impl BM25Index { return Vec::new(); } - let tokens = tokenize(query); - if tokens.is_empty() { - return Vec::new(); - } - - // Collect posting lists and cached IDF scores for query tokens - type QueryTerm<'a> = (&'a str, f32, &'a [(usize, u32)]); - let mut query_terms: Vec> = Vec::new(); - for token in &tokens { - if let (Some(postings), Some(&idf)) = ( - self.inverted.get(token.as_str()), - self.idf_cache.get(token.as_str()), - ) { - query_terms.push((token, idf, postings)); - } - } - - if query_terms.is_empty() { - return Vec::new(); - } - - // Accumulate BM25 scores per document using WAND-style scoring + // Term-at-a-time accumulation. IDF is computed here rather than cached + // at build time: it depends on the live document count, which changes + // with every incremental add/remove, and costs one `ln` per query term. let mut scores: HashMap = HashMap::new(); - - // Compute maximum possible contribution per term for WAND - let max_tf_score: Vec = query_terms - .iter() - .map(|(_, idf, _)| { - // Upper bound: max TF contribution when tf is high and dl is short - let max_tf_num = 10.0 * (self.k1 + 1.0); - let max_tf_den = 10.0 + self.k1 * (1.0 - self.b); - idf * max_tf_num / max_tf_den - }) - .collect(); - - let total_max_contribution: f32 = max_tf_score.iter().sum(); - - // Threshold for WAND early termination. `top_k_heap` is a min-heap of - // size k (worst-of-the-top-k at the head) so it can be maintained in - // O(log k) per update instead of re-sorting the whole buffer. - let mut threshold = 0.0f32; - let mut top_k_heap: BinaryHeap> = BinaryHeap::with_capacity(k); - - for (term_idx, (_, idf, postings)) in query_terms.iter().enumerate() { - for &(doc_id, freq) in *postings { + for token in tokenize(query) { + let Some(postings) = self.inverted.get(token.as_str()) else { + continue; + }; + let df = postings.len() as f32; + let idf = ((self.num_docs as f32 - df + 0.5) / (df + 0.5) + 1.0).ln(); + for &(doc_id, freq) in postings { let dl = self.doc_lengths[doc_id] as f32; let freq_f = freq as f32; let tf = (freq_f * (self.k1 + 1.0)) / (freq_f + self.k1 * (1.0 - self.b + self.b * dl / self.avg_dl)); - let contribution = idf * tf; - - let entry = scores.entry(doc_id).or_insert(0.0); - *entry += contribution; - - // WAND check: if this doc's current partial score + remaining - // max terms can't beat threshold, we can skip (but we still - // accumulate since we process term-at-a-time) - if term_idx == query_terms.len() - 1 { - // Last term: check if this doc beats threshold - let final_score = *entry; - if top_k_heap.len() >= k { - if final_score > threshold { - // Replace the current worst-of-top-k. - top_k_heap.pop(); - top_k_heap.push(Reverse(HeapScore(final_score))); - threshold = top_k_heap.peek().map(|Reverse(s)| s.0).unwrap_or(0.0); - } - } else { - top_k_heap.push(Reverse(HeapScore(final_score))); - if top_k_heap.len() == k { - threshold = top_k_heap.peek().map(|Reverse(s)| s.0).unwrap_or(0.0); - } - } - } - } - // After processing each term, check if remaining terms can - // possibly produce results above threshold - let remaining_max: f32 = max_tf_score[term_idx + 1..].iter().sum(); - if remaining_max < threshold && total_max_contribution > 0.0 { - // Early termination: remaining terms can't produce new top-k - // entries on their own. But existing partial scores may still - // be updated, so we continue (WAND is approximate here). - let _ = remaining_max; // hint to compiler + *scores.entry(doc_id).or_insert(0.0) += idf * tf; } } - let mut results: Vec<(usize, f32)> = scores.into_iter().collect(); - results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); - results.truncate(k); + // Top-k with a bounded min-heap: O(matches * log k) instead of sorting + // every match. Ties break towards the lower doc id so results are + // deterministic (the accumulator is a HashMap). + let mut heap: BinaryHeap)>> = + BinaryHeap::with_capacity(k + 1); + for (doc_id, score) in scores { + heap.push(Reverse((HeapScore(score), Reverse(doc_id)))); + if heap.len() > k { + heap.pop(); + } + } + let mut results: Vec<(usize, f32)> = heap + .into_iter() + .map(|Reverse((HeapScore(score), Reverse(doc_id)))| (doc_id, score)) + .collect(); + results.sort_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0))); results } + /// Number of document slots (live or not) the index covers. Ids are + /// positions in the document list it mirrors. + pub fn len(&self) -> usize { + self.doc_lengths.len() + } + + /// `true` when the index covers no document slots. + pub fn is_empty(&self) -> bool { + self.doc_lengths.is_empty() + } + + /// Index `text` as document `doc_id`, which must be the next free id + /// (`self.len()`) or an existing slot that is currently empty (removed or + /// tombstoned). After any sequence of `add_document` / `remove_document` + /// calls the index scores exactly as one freshly built from the same live + /// documents. + pub fn add_document(&mut self, doc_id: usize, text: &str) { + if doc_id >= self.doc_lengths.len() { + self.doc_lengths.resize(doc_id + 1, 0); + } + debug_assert_eq!(self.doc_lengths[doc_id], 0, "slot {doc_id} is occupied"); + + let tokens = tokenize(text); + let mut term_freqs: HashMap<&str, u32> = HashMap::new(); + for token in &tokens { + *term_freqs.entry(token).or_insert(0) += 1; + } + for (token, freq) in term_freqs { + let postings = self.inverted.entry(token.to_string()).or_default(); + // Posting lists stay sorted by doc id; appends are the common case. + match postings.last() { + Some(&(last, _)) if last >= doc_id => { + let at = postings.partition_point(|&(id, _)| id < doc_id); + postings.insert(at, (doc_id, freq)); + } + _ => postings.push((doc_id, freq)), + } + } + self.doc_lengths[doc_id] = tokens.len() as u32; + self.total_length += tokens.len() as u64; + self.num_docs += 1; + self.refresh_avg_dl(); + } + + /// Extend the index to cover `len` document slots, leaving new ones empty. + /// Used for slots that hold no live document (tombstoned records). + pub fn pad_to(&mut self, len: usize) { + if len > self.doc_lengths.len() { + self.doc_lengths.resize(len, 0); + } + } + + /// Remove document `doc_id`, whose indexed text was `text`. The text is + /// needed to find its postings; pass exactly what was added. + pub fn remove_document(&mut self, doc_id: usize, text: &str) { + let tokens = tokenize(text); + let mut seen: std::collections::HashSet<&str> = std::collections::HashSet::new(); + for token in &tokens { + if !seen.insert(token) { + continue; + } + if let Some(postings) = self.inverted.get_mut(token.as_str()) { + if let Ok(at) = postings.binary_search_by_key(&doc_id, |&(id, _)| id) { + postings.remove(at); + } + if postings.is_empty() { + self.inverted.remove(token.as_str()); + } + } + } + if let Some(len) = self.doc_lengths.get_mut(doc_id) { + self.total_length = self.total_length.saturating_sub(u64::from(*len)); + *len = 0; + } + self.num_docs = self.num_docs.saturating_sub(1); + self.refresh_avg_dl(); + } + + fn refresh_avg_dl(&mut self) { + self.avg_dl = if self.num_docs > 0 { + self.total_length as f32 / self.num_docs as f32 + } else { + 0.0 + }; + } + /// Rebuild the index from scratch (e.g., after compaction). pub fn rebuild(&mut self, documents: &[String], tombstones: &[u8]) { self.inverted.clear(); - self.idf_cache.clear(); self.doc_lengths = vec![0; documents.len()]; + self.total_length = 0; self.avg_dl = 0.0; self.num_docs = 0; self.index_documents(documents, tombstones); @@ -214,23 +252,13 @@ impl BM25Index { } self.num_docs = count; - self.avg_dl = if count > 0 { - total_length as f32 / count as f32 - } else { - 0.0 - }; + self.total_length = total_length; + self.refresh_avg_dl(); // Sort posting lists by doc_id for cache-friendly access for postings in self.inverted.values_mut() { postings.sort_by_key(|&(doc_id, _)| doc_id); } - - // Pre-compute and cache IDF scores - for (token, postings) in &self.inverted { - let df = postings.len() as f32; - let idf = ((self.num_docs as f32 - df + 0.5) / (df + 0.5) + 1.0).ln(); - self.idf_cache.insert(token.clone(), idf); - } } } @@ -386,24 +414,21 @@ mod tests { } #[test] - fn cached_idf_consistent_with_computed() { + fn score_matches_the_bm25_formula() { let docs = vec![ "rust programming".to_string(), "rust systems".to_string(), "python scripting".to_string(), ]; - let tombstones = vec![0, 0, 0]; - let index = BM25Index::build(&docs, &tombstones); + let index = BM25Index::build(&docs, &[0, 0, 0]); - // IDF for "rust" (appears in 2 of 3 docs) - let idf_rust = index.idf_cache.get("rust").unwrap(); - let expected_idf = ((3.0f32 - 2.0 + 0.5) / (2.0 + 0.5) + 1.0).ln(); - assert!( - (idf_rust - expected_idf).abs() < 1e-6, - "cached IDF mismatch: {} vs {}", - idf_rust, - expected_idf - ); + // "python": df = 1 of N = 3. Every doc has the average length (2) and + // tf = 1, so the tf factor is exactly 1 and the score is the IDF. + let results = index.search("python", 3); + let expected_idf = ((3.0f32 - 1.0 + 0.5) / (1.0 + 0.5) + 1.0).ln(); + assert_eq!(results.len(), 1); + assert_eq!(results[0].0, 2); + assert!((results[0].1 - expected_idf).abs() < 1e-6, "{results:?}"); } #[test] @@ -467,4 +492,83 @@ mod tests { ); } } + /// Documents drawn from a small vocabulary so terms collide heavily. + fn random_doc(state: &mut u64) -> String { + const VOCAB: &[&str] = &[ + "alpha", "beta", "gamma", "delta", "eps", "zeta", "eta", "x1", + ]; + let mut next = || { + *state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + (*state >> 33) as usize + }; + let len = 1 + next() % 9; + (0..len) + .map(|_| VOCAB[next() % VOCAB.len()]) + .collect::>() + .join(" ") + } + + #[test] + fn incremental_updates_match_a_fresh_build_exactly() { + for seed in 0..60u64 { + let mut state = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1; + let mut docs: Vec = Vec::new(); + let mut tombstones: Vec = Vec::new(); + let mut index = BM25Index::build(&docs, &tombstones); + + for step in 0..80 { + state = state.wrapping_mul(6364136223846793005).wrapping_add(1); + let live: Vec = (0..docs.len()).filter(|&i| tombstones[i] == 0).collect(); + match (state >> 40) % 4 { + 0 if !live.is_empty() => { + // delete + let id = live[(state >> 20) as usize % live.len()]; + index.remove_document(id, &docs[id]); + tombstones[id] = 1; + } + 1 if !live.is_empty() => { + // update in place + let id = live[(state >> 20) as usize % live.len()]; + let new_text = random_doc(&mut state); + index.remove_document(id, &docs[id]); + index.add_document(id, &new_text); + docs[id] = new_text; + } + _ => { + let text = random_doc(&mut state); + index.add_document(docs.len(), &text); + docs.push(text); + tombstones.push(0); + } + } + + let fresh = BM25Index::build(&docs, &tombstones); + for query in ["alpha", "beta gamma", "x1 zeta alpha delta", "missing"] { + let got = index.search(query, 5); + let want = fresh.search(query, 5); + assert_eq!(got.len(), want.len(), "seed {seed} step {step} {query:?}"); + for (g, w) in got.iter().zip(&want) { + assert_eq!( + g.0, w.0, + "seed {seed} step {step} {query:?}: {got:?} vs {want:?}" + ); + assert!( + (g.1 - w.1).abs() < 1e-5, + "seed {seed} step {step} {query:?}" + ); + } + } + } + } + } + + #[test] + fn ties_break_towards_the_lower_doc_id() { + let docs: Vec = (0..6).map(|_| "same text".to_string()).collect(); + let index = BM25Index::build(&docs, &[0; 6]); + let ids: Vec = index.search("same", 3).into_iter().map(|r| r.0).collect(); + assert_eq!(ids, [0, 1, 2]); + } } diff --git a/crates/clawhdf5-agent/src/lib.rs b/crates/clawhdf5-agent/src/lib.rs index f07c755..ad624f0 100644 --- a/crates/clawhdf5-agent/src/lib.rs +++ b/crates/clawhdf5-agent/src/lib.rs @@ -205,6 +205,12 @@ pub trait AgentMemory { fn get_session_summary(&self, session_id: &str) -> Result>; } +/// Ceiling for a record's Hebbian activation weight. Each hit adds +/// `hebbian_boost` and the fused score is scaled by `sqrt(weight)`, so without +/// a cap a frequently returned record's advantage grows without limit and it +/// eventually outranks better matches purely on popularity. +pub(crate) const MAX_ACTIVATION_WEIGHT: f32 = 16.0; + /// Most anomaly alerts kept between `take_anomaly_alerts` calls. const MAX_PENDING_ALERTS: usize = 1024; @@ -247,6 +253,14 @@ pub struct HDF5Memory { /// via [`HDF5Memory::take_anomaly_alerts`]. Saves are never blocked on /// these — surfacing is opt-in for callers that want to act on them. anomaly_alerts: Vec, + /// Keyword index over `cache.chunks`, kept for the life of the store and + /// updated incrementally — it used to be rebuilt from scratch, re-tokenising + /// every record, on every single query. Built lazily on first use; see + /// [`HDF5Memory::ensure_bm25_fresh`] for how it stays in sync. + bm25: Option, + /// Activation weights changed since the last checkpoint (searches boost + /// the records they return). Cleared by `flush`. + activations_dirty: bool, /// Opened with [`HDF5Memory::open_read_only`]: nothing may reach the disk. read_only: bool, /// A WAL that `open()` could not read and moved aside; see @@ -299,6 +313,8 @@ impl HDF5Memory { provenance: provenance::ProvenanceStore::new(), anomaly: anomaly::WriteAnomalyDetector::new(anomaly::AnomalyConfig::default()), anomaly_alerts: Vec::new(), + bm25: None, + activations_dirty: false, read_only: false, quarantined_wal: None, _lock: Some(lock), @@ -434,12 +450,58 @@ impl HDF5Memory { provenance: provenance::ProvenanceStore::new(), anomaly: anomaly::WriteAnomalyDetector::new(anomaly::AnomalyConfig::default()), anomaly_alerts: Vec::new(), + bm25: None, + activations_dirty: false, read_only, quarantined_wal, _lock: lock, }) } + /// Bring the keyword index in line with the cache and return it. + /// + /// Appends need no hook: records the index hasn't seen yet (whatever path + /// added them) are indexed here, in order. Changes that keep the length the + /// same are reported explicitly — [`Self::bm25_on_delete`] and + /// [`Self::bm25_on_update`] — and anything that renumbers records + /// (compaction) drops the index so it is rebuilt. + pub(crate) fn ensure_bm25_fresh(&mut self) -> &bm25::BM25Index { + let n = self.cache.chunks.len(); + let bm25 = match self.bm25.take() { + Some(index) if index.len() <= n => { + let mut index = index; + for id in index.len()..n { + if self.cache.tombstones[id] == 0 { + index.add_document(id, &self.cache.chunks[id]); + } + } + index.pad_to(n); + index + } + _ => bm25::BM25Index::build(&self.cache.chunks, &self.cache.tombstones), + }; + self.bm25.insert(bm25) + } + + /// Record `id` was tombstoned; its text is still in the cache. + fn bm25_on_delete(&mut self, id: usize) { + if let Some(index) = self.bm25.as_mut() + && id < index.len() + { + index.remove_document(id, &self.cache.chunks[id]); + } + } + + /// Record `id`'s text changed from `old_text` to what the cache holds now. + fn bm25_on_update(&mut self, id: usize, old_text: &str) { + if let Some(index) = self.bm25.as_mut() + && id < index.len() + { + index.remove_document(id, old_text); + index.add_document(id, &self.cache.chunks[id]); + } + } + /// Flush current state to disk and truncate the WAL. /// /// Every code path that persists the full cache to the .h5 file must @@ -466,6 +528,7 @@ impl HDF5Memory { if let Some(ref mut w) = self.wal { w.truncate()?; } + self.activations_dirty = false; Ok(()) } @@ -777,6 +840,7 @@ impl HDF5Memory { &entry.session_id, entry.timestamp, ); + let old_text = std::mem::take(&mut self.cache.chunks[existing_idx]); self.cache.update( existing_idx, entry.chunk, @@ -785,6 +849,7 @@ impl HDF5Memory { entry.timestamp, entry.session_id, ); + self.bm25_on_update(existing_idx, &old_text); // In-place embedding change: the index node is stale, force rebuild. self.hnsw_mark_dirty(); let needs_flush = self @@ -876,6 +941,7 @@ impl AgentMemory for HDF5Memory { ))); } self.hnsw_on_delete(id); + self.bm25_on_delete(id); self.flush()?; // Auto-compact if threshold exceeded @@ -893,6 +959,7 @@ impl AgentMemory for HDF5Memory { if removed > 0 { // Record ids are cache indices, which compaction just renumbered. self.provenance.remap(&index_map); + self.bm25 = None; // Compaction renumbers cache indices; rebuild the index to match. self.hnsw_mark_dirty(); self.flush()?; @@ -1172,6 +1239,18 @@ impl HDF5Memory { // --- Tests --- +impl Drop for HDF5Memory { + /// Best-effort checkpoint of activation weights that only searches have + /// touched. Everything else is already durable through the WAL or an + /// earlier checkpoint; without this a search-only session would forget + /// every boost it made. + fn drop(&mut self) { + if self.activations_dirty && !self.read_only { + let _ = self.flush(); + } + } +} + #[cfg(test)] mod tests { use super::*; @@ -1595,6 +1674,83 @@ mod tests { assert_eq!(restored.cache.chunks, ["checkpointed", "wal-only"]); } + #[test] + fn keyword_index_stays_in_sync_through_every_mutation() { + let dir = TempDir::new().unwrap(); + let mut config = make_config(&dir); + config.compact_threshold = 0.0; // compact only when asked + let mut mem = HDF5Memory::create(config).unwrap(); + let check = |mem: &mut HDF5Memory, what: &str| { + let fresh = bm25::BM25Index::build(&mem.cache.chunks, &mem.cache.tombstones); + let n = mem.cache.len(); + for query in ["apple", "banana cherry", "date", "nothing"] { + let kept = mem.ensure_bm25_fresh().search(query, n); + assert_eq!(kept, fresh.search(query, n), "{what}: {query:?}"); + } + }; + let tagged = |chunk: &str, tag: &str| { + let mut e = make_entry(chunk, &[1.0, 0.0, 0.0, 0.0]); + e.tags = tag.into(); + e + }; + + check(&mut mem, "empty"); + mem.save(tagged("apple banana", "a")).unwrap(); + mem.save(tagged("banana cherry cherry", "b")).unwrap(); + check(&mut mem, "after saves"); + mem.save_batch(vec![tagged("date apple", "c"), tagged("cherry", "d")]) + .unwrap(); + check(&mut mem, "after save_batch"); + mem.save_or_update(tagged("date date date", "a")).unwrap(); + check(&mut mem, "after in-place update"); + mem.delete(1).unwrap(); + check(&mut mem, "after delete"); + mem.save(tagged("apple cherry", "e")).unwrap(); + check(&mut mem, "after save following a delete"); + mem.compact().unwrap(); + check(&mut mem, "after compact"); + mem.hybrid_search(&[1.0, 0.0, 0.0, 0.0], "apple", 0.5, 0.5, 3); + check(&mut mem, "after a search"); + } + + #[test] + fn search_does_not_write_the_store_but_boosts_persist_on_drop() { + let dir = TempDir::new().unwrap(); + let config = make_config(&dir); + let path = config.path.clone(); + let mut mem = HDF5Memory::create(config).unwrap(); + mem.save(make_entry("findable", &[1.0, 0.0, 0.0, 0.0])) + .unwrap(); + let before = std::fs::read(&path).unwrap(); + + for _ in 0..3 { + mem.hybrid_search(&[1.0, 0.0, 0.0, 0.0], "findable", 1.0, 0.0, 1); + } + assert_eq!( + std::fs::read(&path).unwrap(), + before, + "a query must not rewrite the store" + ); + let boosted = mem.cache.activation_weights[0]; + assert!(boosted > 1.0); + drop(mem); + + let reopened = HDF5Memory::open(&path).unwrap(); + assert_eq!(reopened.cache.activation_weights[0], boosted); + } + + #[test] + fn activation_weight_is_capped() { + let dir = TempDir::new().unwrap(); + let mut mem = HDF5Memory::create(make_config(&dir)).unwrap(); + mem.save(make_entry("popular", &[1.0, 0.0, 0.0, 0.0])) + .unwrap(); + for _ in 0..500 { + mem.hybrid_search(&[1.0, 0.0, 0.0, 0.0], "popular", 1.0, 0.0, 1); + } + assert_eq!(mem.cache.activation_weights[0], MAX_ACTIVATION_WEIGHT); + } + #[test] fn store_has_a_single_writer() { let dir = TempDir::new().unwrap(); diff --git a/crates/clawhdf5-agent/src/search.rs b/crates/clawhdf5-agent/src/search.rs index b66f2f2..30b679a 100644 --- a/crates/clawhdf5-agent/src/search.rs +++ b/crates/clawhdf5-agent/src/search.rs @@ -4,7 +4,7 @@ use std::path::Path; use crate::bm25; use crate::hybrid; -use crate::{HDF5Memory, MemoryError, Result, SearchResult}; +use crate::{HDF5Memory, MAX_ACTIVATION_WEIGHT, MemoryError, Result, SearchResult}; impl HDF5Memory { /// Vector + keyword scoring stage of [`HDF5Memory::hybrid_search`]. @@ -90,7 +90,11 @@ impl HDF5Memory { keyword_weight: f32, k: usize, ) -> Vec { - let bm25 = bm25::BM25Index::build(&self.cache.chunks, &self.cache.tombstones); + // The keyword index lives for the life of the store and is updated + // incrementally. Take it out for the duration of the call so the + // vector stage can borrow `self` mutably, then put it back. + self.ensure_bm25_fresh(); + let bm25 = self.bm25.take().expect("ensure_bm25_fresh leaves an index"); let scored = self.vector_keyword_search( query_embedding, query_text, @@ -132,15 +136,26 @@ impl HDF5Memory { .map(|r| r.index) .collect(); self.apply_hebbian_boost(&hit_indices); - self.flush().ok(); + self.bm25 = Some(bm25); results } + /// Reinforce the records a query returned. The new weights are persisted by + /// the next checkpoint (any write that flushes, `flush_wal`, or drop) — not + /// by rewriting the whole store inside the query, which is what made + /// `hybrid_search` cost O(store size) in disk I/O. They are a ranking hint, + /// not user data: a crash before the next checkpoint only forgets the + /// boosts since the last one. fn apply_hebbian_boost(&mut self, hit_indices: &[usize]) { - for &idx in hit_indices { - self.cache.activation_weights[idx] += self.config.hebbian_boost; + if hit_indices.is_empty() || self.config.hebbian_boost == 0.0 { + return; } + for &idx in hit_indices { + let w = &mut self.cache.activation_weights[idx]; + *w = (*w + self.config.hebbian_boost).min(MAX_ACTIVATION_WEIGHT); + } + self.activations_dirty = true; } /// Get the chunk text for a memory entry by index.