perf(agent): persistent incremental BM25 index; no store rewrite per query
hybrid_search rebuilt the BM25 index from scratch (re-tokenising every record) and rewrote the whole .h5 file on every single query, so a query cost O(store size) in both CPU and disk I/O. Steady-state p50 per the search harness: 5.5 -> 0.24 ms (1K), 49 -> 2.1 ms (10K), 884 -> 23 ms (100K). - BM25Index is incremental: add_document / remove_document keep it exactly equivalent to a fresh build over the same live documents (property test: 60 random op sequences compared against BM25Index::build after every step). IDF moves to query time since it depends on the live document count. Top-k uses a bounded heap, ties break by doc id (results were HashMap-ordered), and the "WAND" code that computed a bound and then discarded it is removed. - HDF5Memory keeps one index for its lifetime, built lazily. Appends are picked up by ensure_bm25_fresh whatever path added them; delete and in-place update report themselves; compaction drops the index. A test drives every mutation and compares against a fresh build. - A query no longer calls flush(). Activation boosts are marked dirty and persisted by the next checkpoint, including a best-effort one on drop so a search-only session keeps them (approved behaviour change). Activation weights are capped at 16.0; they previously grew without bound. The archived mission branch's BM25 cache was reviewed and not used: it was invalidated by every write, so interleaved save/search still rebuilt per query, and it changed the default fusion weights. Co-Authored-By: Claude Fable 5.1 <[email protected]>
This commit is contained in:
osobh
co-authored by
Claude Fable 5.1
Claude Fable 5.1
parent
61424d1418
commit
2bfbb7fb4b
+217
-113
@@ -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<String, Vec<(usize, u32)>>,
|
||||
/// Cached IDF scores per token.
|
||||
idf_cache: HashMap<String, f32>,
|
||||
/// Number of tokens in each document (0 for tombstoned docs).
|
||||
doc_lengths: Vec<u32>,
|
||||
/// 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<QueryTerm<'_>> = 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<usize, f32> = HashMap::new();
|
||||
|
||||
// Compute maximum possible contribution per term for WAND
|
||||
let max_tf_score: Vec<f32> = 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<Reverse<HeapScore>> = 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<Reverse<(HeapScore, Reverse<usize>)>> =
|
||||
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::<Vec<_>>()
|
||||
.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<String> = Vec::new();
|
||||
let mut tombstones: Vec<u8> = 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<usize> = (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<String> = (0..6).map(|_| "same text".to_string()).collect();
|
||||
let index = BM25Index::build(&docs, &[0; 6]);
|
||||
let ids: Vec<usize> = index.search("same", 3).into_iter().map(|r| r.0).collect();
|
||||
assert_eq!(ids, [0, 1, 2]);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -205,6 +205,12 @@ pub trait AgentMemory {
|
||||
fn get_session_summary(&self, session_id: &str) -> Result<Option<String>>;
|
||||
}
|
||||
|
||||
/// 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<anomaly::AnomalyAlert>,
|
||||
/// 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<bm25::BM25Index>,
|
||||
/// 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();
|
||||
|
||||
@@ -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<SearchResult> {
|
||||
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.
|
||||
|
||||
Reference in New Issue
Block a user