Merge HNSW agent integration + Python 3.14 fix

Lands the changes from PR #1, whose server-side squash merge landed as an
empty commit (Gitea merge working-tree error). See branch
feat/hnsw-agent-integration-py314 for the full history.

Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>

crates/clawhdf5-agent/Cargo.toml

@@ -14,6 +14,7 @@ clawhdf5-format = { path = "../clawhdf5-format", version = "2.0.0", features = [
 clawhdf5 = { path = "../clawhdf5", version = "2.0.0" }
 clawhdf5-io = { path = "../clawhdf5-io", version = "2.0.0", features = ["mmap"] }
 clawhdf5-accel = { path = "../clawhdf5-accel", version = "2.0.0" }
+clawhdf5-ann = { path = "../clawhdf5-ann", version = "2.0.0", optional = true }
 clawhdf5-gpu = { path = "../clawhdf5-gpu", version = "2.0.0", optional = true, default-features = false }
 serde = { version = "1", features = ["derive"] }
 byteorder = "1"
@@ -44,9 +45,15 @@ name = "memory_bench"
 harness = false
 
 [features]
-default = ["float16"]
+default = ["float16", "hnsw"]
 float16 = ["half"]
 parallel = ["rayon"]
+# HNSW approximate-nearest-neighbour acceleration for the vector stage of
+# hybrid_search. On by default; the index is rebuilt from the cache on demand
+# and stays self-consistent with the persisted memory store. Disable with
+# `--no-default-features` (plus re-enabling other defaults) to force the exact
+# linear cosine scan.
+hnsw = ["clawhdf5-ann"]
 agent = []
 gpu = ["clawhdf5-gpu/gpu-wgpu"]
 fast-math = ["matrixmultiply"]

crates/clawhdf5-agent/src/hybrid.rs

@@ -60,6 +60,22 @@ pub fn hybrid_search(
     };
     let kw_scores = bm25_index.search(query_text, vectors.len());
 
+    merge_vector_keyword(vec_scores, kw_scores, vector_weight, keyword_weight, k)
+}
+
+/// Merge pre-computed vector-similarity and keyword scores into a single ranking.
+///
+/// Both score sets are independently min-max normalized to [0, 1] and combined
+/// with the given weights. This is the shared core of [`hybrid_search`]; it is
+/// also used by the optional HNSW path, which supplies vector scores from an
+/// approximate-nearest-neighbour index instead of a full linear scan.
+pub fn merge_vector_keyword(
+    vec_scores: Vec<(usize, f32)>,
+    kw_scores: Vec<(usize, f32)>,
+    vector_weight: f32,
+    keyword_weight: f32,
+    k: usize,
+) -> Vec<(usize, f32)> {
     // Normalize each set to [0, 1].
     let vec_normalized = normalize_scores(&vec_scores);
     let kw_normalized = normalize_scores(&kw_scores);

crates/clawhdf5-agent/src/lib.rs

@@ -60,7 +60,16 @@ pub fn cosine_similarity_prenorm(
 use std::path::{Path, PathBuf};
 
 use cache::MemoryCache;
+#[cfg(feature = "hnsw")]
+use clawhdf5_ann::{DistanceMetric, HnswIndex};
 use ephemeral::{EphemeralConfig, EphemeralStore};
+
+/// HNSW construction parameters used for the agent's vector index. Cosine is the
+/// agent's similarity metric, so the index is built with cosine distance.
+#[cfg(feature = "hnsw")]
+const HNSW_M: usize = 16;
+#[cfg(feature = "hnsw")]
+const HNSW_EF_CONSTRUCTION: usize = 64;
 // EphemeralEntry and EphemeralStats are part of the crate public API via
 // the `ephemeral` module; they are not needed directly in lib.rs internals.
 #[allow(unused_imports)]
@@ -202,6 +211,22 @@ pub struct HDF5Memory {
     wal: Option<wal::WalFile>,
     strategy: Option<Box<dyn MemoryStrategy>>,
     pub ephemeral: Option<EphemeralStore>,
+    /// Optional HNSW index accelerating the vector stage of `hybrid_search`.
+    /// `None` when the store isn't indexable (no/zero-dim/mixed-dim embeddings);
+    /// rebuilt from the cache whenever it drifts out of sync (see
+    /// [`HDF5Memory::ensure_hnsw_fresh`]).
+    #[cfg(feature = "hnsw")]
+    hnsw: Option<HnswIndex>,
+    /// Set when an in-place update/compaction may have invalidated `hnsw`,
+    /// forcing a rebuild before the next search.
+    #[cfg(feature = "hnsw")]
+    hnsw_dirty: bool,
+    /// Cache length the current `hnsw` value reflects. A mismatch with the live
+    /// cache length triggers a rebuild — this both picks up unhooked cache
+    /// pushes and avoids re-attempting to build an unindexable store every
+    /// search.
+    #[cfg(feature = "hnsw")]
+    hnsw_synced_len: usize,
 }
 
 impl std::fmt::Debug for HDF5Memory {
@@ -235,6 +260,12 @@ impl HDF5Memory {
             wal,
             strategy: None,
             ephemeral: None,
+            #[cfg(feature = "hnsw")]
+            hnsw: None,
+            #[cfg(feature = "hnsw")]
+            hnsw_dirty: false,
+            #[cfg(feature = "hnsw")]
+            hnsw_synced_len: 0,
         })
     }
 
@@ -262,6 +293,14 @@ impl HDF5Memory {
             wal,
             strategy: None,
             ephemeral: None,
+            // Existing data is loaded from disk + WAL replay; mark the index
+            // dirty so it is (re)built from the cache on the first search.
+            #[cfg(feature = "hnsw")]
+            hnsw: None,
+            #[cfg(feature = "hnsw")]
+            hnsw_dirty: true,
+            #[cfg(feature = "hnsw")]
+            hnsw_synced_len: 0,
         })
     }
 
@@ -284,6 +323,108 @@ impl HDF5Memory {
         Ok(())
     }
 
+    // ---- HNSW index maintenance --------------------------------------------
+    //
+    // The index mirrors the cache: HNSW node id == cache index, kept aligned by
+    // appending to both in lock-step and mirroring deletes. The incremental
+    // hooks below are an optimization for the hot path; correctness is
+    // guaranteed by `ensure_hnsw_fresh`, which rebuilds from the cache whenever
+    // the index length drifts from the cache length (covering any mutation path
+    // that doesn't call a hook, e.g. consolidation pushes).
+
+    /// Build an HNSW index over the entire cache, re-applying tombstones as
+    /// soft-deletions so node ids stay aligned with cache indices.
+    ///
+    /// Returns `None` for stores that aren't usefully indexable: no embeddings,
+    /// a zero embedding dimension, or embeddings of mixed dimension (in which
+    /// case `hybrid_search` keeps using the linear scan).
+    #[cfg(feature = "hnsw")]
+    fn build_hnsw_from_cache(&self) -> Option<HnswIndex> {
+        let dim = self.cache.embedding_dim;
+        if dim == 0 || self.cache.embeddings.is_empty() {
+            return None;
+        }
+        if self.cache.embeddings.iter().any(|e| e.len() != dim) {
+            return None;
+        }
+        let mut index = HnswIndex::build_with_metric(
+            &self.cache.embeddings,
+            HNSW_M,
+            HNSW_EF_CONSTRUCTION,
+            DistanceMetric::Cosine,
+        );
+        for (i, &t) in self.cache.tombstones.iter().enumerate() {
+            if t != 0 {
+                index.mark_deleted(i);
+            }
+        }
+        Some(index)
+    }
+
+    /// Ensure the HNSW index reflects the current cache. Rebuilds when marked
+    /// dirty or when the cache length no longer matches what the index reflects.
+    #[cfg(feature = "hnsw")]
+    fn ensure_hnsw_fresh(&mut self) {
+        let n = self.cache.embeddings.len();
+        if self.hnsw_dirty || self.hnsw_synced_len != n {
+            self.hnsw = self.build_hnsw_from_cache();
+            self.hnsw_synced_len = n;
+            self.hnsw_dirty = false;
+        }
+    }
+
+    /// Incrementally index the embedding just pushed at `idx`. Falls back to a
+    /// rebuild (via the dirty flag) for the first vector, dimension mismatches,
+    /// or id drift.
+    #[cfg(feature = "hnsw")]
+    fn hnsw_on_insert(&mut self, idx: usize) {
+        if self.hnsw_dirty {
+            return; // a rebuild is already pending; it will pick this up
+        }
+        let emb_len = self.cache.embeddings[idx].len();
+        match self.hnsw.as_mut() {
+            Some(index) if emb_len == index.dimension() => {
+                let id = index.insert(self.cache.embeddings[idx].clone());
+                if id == idx {
+                    self.hnsw_synced_len = self.cache.embeddings.len();
+                } else {
+                    self.hnsw_dirty = true;
+                }
+            }
+            // Dimension mismatch, first-ever vector, or no index yet: defer to a
+            // rebuild, which decides indexability uniformly.
+            _ => self.hnsw_dirty = true,
+        }
+    }
+
+    #[cfg(not(feature = "hnsw"))]
+    #[inline]
+    fn hnsw_on_insert(&mut self, _idx: usize) {}
+
+    /// Mirror a cache deletion into the index.
+    #[cfg(feature = "hnsw")]
+    fn hnsw_on_delete(&mut self, id: usize) {
+        if let Some(index) = self.hnsw.as_mut() {
+            index.mark_deleted(id);
+        }
+    }
+
+    #[cfg(not(feature = "hnsw"))]
+    #[inline]
+    fn hnsw_on_delete(&mut self, _id: usize) {}
+
+    /// Mark the index for rebuild after a mutation that may have changed
+    /// existing embeddings or renumbered ids (in-place update, compaction).
+    #[cfg(feature = "hnsw")]
+    #[inline]
+    fn hnsw_mark_dirty(&mut self) {
+        self.hnsw_dirty = true;
+    }
+
+    #[cfg(not(feature = "hnsw"))]
+    #[inline]
+    fn hnsw_mark_dirty(&mut self) {}
+
     /// Get a reference to the config.
     pub fn config(&self) -> &MemoryConfig {
         &self.config
@@ -374,6 +515,8 @@ impl HDF5Memory {
                 entry.timestamp,
                 entry.session_id,
             );
+            // In-place embedding change: the index node is stale, force rebuild.
+            self.hnsw_mark_dirty();
             let needs_flush = self
                 .wal
                 .as_ref()
@@ -414,6 +557,7 @@ impl AgentMemory for HDF5Memory {
             entry.session_id,
             entry.tags,
         );
+        self.hnsw_on_insert(idx);
         let needs_flush = self
             .wal
             .as_ref()
@@ -440,6 +584,8 @@ impl AgentMemory for HDF5Memory {
             );
             indices.push(idx);
         }
+        // Batch inserts rebuild the index once rather than node-by-node.
+        self.hnsw_mark_dirty();
         self.flush()?;
         Ok(indices)
     }
@@ -450,6 +596,7 @@ impl AgentMemory for HDF5Memory {
                 "entry {id} not found or already deleted"
             )));
         }
+        self.hnsw_on_delete(id);
         self.flush()?;
 
         // Auto-compact if threshold exceeded
@@ -465,6 +612,8 @@ impl AgentMemory for HDF5Memory {
     fn compact(&mut self) -> Result<usize> {
         let (removed, _index_map) = self.cache.compact();
         if removed > 0 {
+            // Compaction renumbers cache indices; rebuild the index to match.
+            self.hnsw_mark_dirty();
             self.flush()?;
         }
         Ok(removed)

crates/clawhdf5-agent/src/search.rs

@@ -7,6 +7,76 @@ use crate::hybrid;
 use crate::{HDF5Memory, MemoryError, Result, SearchResult};
 
 impl HDF5Memory {
+    /// Vector + keyword scoring stage of [`HDF5Memory::hybrid_search`].
+    ///
+    /// Without the `hnsw` feature this is a full linear cosine scan (the exact
+    /// previous behaviour, also used as the correctness oracle in tests). With
+    /// `hnsw` enabled and an index available, the vector candidates come from an
+    /// approximate-nearest-neighbour search over an over-fetched pool, then merge
+    /// with BM25 via the shared [`hybrid::merge_vector_keyword`].
+    #[cfg(feature = "hnsw")]
+    fn vector_keyword_search(
+        &mut self,
+        query_embedding: &[f32],
+        query_text: &str,
+        bm25: &bm25::BM25Index,
+        vector_weight: f32,
+        keyword_weight: f32,
+        k: usize,
+    ) -> Vec<(usize, f32)> {
+        self.ensure_hnsw_fresh();
+        match self.hnsw.as_ref() {
+            Some(index)
+                if !index.is_empty() && index.dimension() == query_embedding.len() =>
+            {
+                // Over-fetch so the merge sees a useful vector pool; cosine
+                // distance from the index converts back to similarity (1 - d).
+                let pool = (k * 8).max(64);
+                let vec_scores: Vec<(usize, f32)> = index
+                    .search(query_embedding, pool, pool)
+                    .into_iter()
+                    .map(|(id, dist)| (id, 1.0 - dist))
+                    .collect();
+                let kw_scores = bm25.search(query_text, self.cache.len());
+                hybrid::merge_vector_keyword(vec_scores, kw_scores, vector_weight, keyword_weight, k)
+            }
+            _ => hybrid::hybrid_search(
+                query_embedding,
+                query_text,
+                &self.cache.embeddings,
+                &self.cache.chunks,
+                &self.cache.tombstones,
+                bm25,
+                vector_weight,
+                keyword_weight,
+                k,
+            ),
+        }
+    }
+
+    #[cfg(not(feature = "hnsw"))]
+    fn vector_keyword_search(
+        &mut self,
+        query_embedding: &[f32],
+        query_text: &str,
+        bm25: &bm25::BM25Index,
+        vector_weight: f32,
+        keyword_weight: f32,
+        k: usize,
+    ) -> Vec<(usize, f32)> {
+        hybrid::hybrid_search(
+            query_embedding,
+            query_text,
+            &self.cache.embeddings,
+            &self.cache.chunks,
+            &self.cache.tombstones,
+            bm25,
+            vector_weight,
+            keyword_weight,
+            k,
+        )
+    }
+
     /// Perform hybrid search combining cosine vector similarity and BM25 keyword search.
     pub fn hybrid_search(
         &mut self,
@@ -17,12 +87,9 @@ impl HDF5Memory {
         k: usize,
     ) -> Vec<SearchResult> {
         let bm25 = bm25::BM25Index::build(&self.cache.chunks, &self.cache.tombstones);
-        let scored = hybrid::hybrid_search(
+        let scored = self.vector_keyword_search(
             query_embedding,
             query_text,
-            &self.cache.embeddings,
-            &self.cache.chunks,
-            &self.cache.tombstones,
             &bm25,
             vector_weight,
             keyword_weight,

crates/clawhdf5-agent/tests/hnsw_integration.rs

@@ -0,0 +1,163 @@
+//! Integration tests for the optional HNSW-accelerated vector search path.
+//!
+//! These only run when the crate is built with `--features hnsw`. They drive the
+//! real `HDF5Memory` API (save / save_batch / delete / hybrid_search) and check
+//! the approximate results against a brute-force cosine oracle, plus confirm that
+//! deletions are honoured end-to-end.
+#![cfg(feature = "hnsw")]
+
+use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};
+use tempfile::TempDir;
+
+/// Deterministic splitmix64 so tests are reproducible without an RNG crate.
+fn splitmix64(state: &mut u64) -> u64 {
+    *state = state.wrapping_add(0x9e37_79b9_7f4a_7c15);
+    let mut z = *state;
+    z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
+    z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
+    z ^ (z >> 31)
+}
+
+fn make_vector(seed: &mut u64, dim: usize) -> Vec<f32> {
+    (0..dim)
+        .map(|_| (splitmix64(seed) >> 40) as f32 / 16_777_216.0 - 0.5)
+        .collect()
+}
+
+fn cosine(a: &[f32], b: &[f32]) -> f32 {
+    let dot: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum();
+    let na: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
+    let nb: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
+    if na == 0.0 || nb == 0.0 {
+        0.0
+    } else {
+        dot / (na * nb)
+    }
+}
+
+fn entry(chunk: &str, embedding: Vec<f32>, tags: &str) -> MemoryEntry {
+    MemoryEntry {
+        chunk: chunk.to_string(),
+        embedding,
+        source_channel: "test".to_string(),
+        timestamp: 0.0,
+        session_id: "s".to_string(),
+        tags: tags.to_string(),
+    }
+}
+
+fn new_memory(dir: &TempDir, dim: usize) -> HDF5Memory {
+    let config = MemoryConfig::new(dir.path().join("mem.h5"), "agent", dim);
+    HDF5Memory::create(config).unwrap()
+}
+
+#[test]
+fn hnsw_matches_bruteforce_oracle() {
+    let dir = TempDir::new().unwrap();
+    let dim = 16;
+    let n = 250;
+    let mut mem = new_memory(&dir, dim);
+
+    let mut seed = 0xC0FF_EE12_3456_789A;
+    let vectors: Vec<Vec<f32>> = (0..n).map(|_| make_vector(&mut seed, dim)).collect();
+    for (i, v) in vectors.iter().enumerate() {
+        mem.save(entry(&format!("chunk {i}"), v.clone(), &format!("k{i}")))
+            .unwrap();
+    }
+
+    // Vector-only query: keyword weight 0 isolates the HNSW vector stage.
+    let query = make_vector(&mut seed, dim);
+    let k = 10;
+    let results = mem.hybrid_search(&query, "", 1.0, 0.0, k);
+    assert_eq!(results.len(), k, "should return k results");
+
+    // Brute-force cosine top-k oracle.
+    let mut oracle: Vec<(usize, f32)> = vectors
+        .iter()
+        .enumerate()
+        .map(|(i, v)| (i, cosine(&query, v)))
+        .collect();
+    oracle.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
+    let oracle_ids: std::collections::HashSet<usize> =
+        oracle.iter().take(k).map(|(i, _)| *i).collect();
+    let hnsw_ids: std::collections::HashSet<usize> =
+        results.iter().map(|r| r.index).collect();
+
+    let overlap = oracle_ids.intersection(&hnsw_ids).count();
+    assert!(
+        overlap >= 7,
+        "HNSW recall too low vs brute force: {overlap}/{k} (hnsw={hnsw_ids:?})"
+    );
+}
+
+#[test]
+fn deleted_entry_excluded_from_search() {
+    let dir = TempDir::new().unwrap();
+    let dim = 8;
+    let mut mem = new_memory(&dir, dim);
+
+    let mut seed = 42;
+    let vectors: Vec<Vec<f32>> = (0..60).map(|_| make_vector(&mut seed, dim)).collect();
+    for (i, v) in vectors.iter().enumerate() {
+        mem.save(entry(&format!("c{i}"), v.clone(), &format!("t{i}")))
+            .unwrap();
+    }
+
+    // Query exactly equal to vector 5 — it must be the top hit.
+    let query = vectors[5].clone();
+    let top = mem.hybrid_search(&query, "", 1.0, 0.0, 1);
+    assert_eq!(top[0].index, 5, "exact match should rank first");
+
+    mem.delete(5).unwrap();
+
+    let after = mem.hybrid_search(&query, "", 1.0, 0.0, 5);
+    assert!(
+        after.iter().all(|r| r.index != 5),
+        "deleted entry must not appear in results"
+    );
+}
+
+#[test]
+fn incremental_inserts_after_search_are_found() {
+    let dir = TempDir::new().unwrap();
+    let dim = 8;
+    let mut mem = new_memory(&dir, dim);
+
+    let mut seed = 7;
+    // First batch, then a search to force the index to build.
+    for i in 0..40 {
+        let v = make_vector(&mut seed, dim);
+        mem.save(entry(&format!("a{i}"), v, &format!("a{i}"))).unwrap();
+    }
+    let _ = mem.hybrid_search(&make_vector(&mut seed, dim), "", 1.0, 0.0, 5);
+
+    // Now insert a distinctive vector incrementally and confirm we can find it.
+    let needle = vec![10.0f32; dim];
+    let idx = mem
+        .save(entry("needle", needle.clone(), "needle"))
+        .unwrap();
+    let hits = mem.hybrid_search(&needle, "", 1.0, 0.0, 1);
+    assert_eq!(hits[0].index, idx, "incrementally inserted vector must be found");
+}
+
+#[test]
+fn save_batch_then_search_is_consistent() {
+    let dir = TempDir::new().unwrap();
+    let dim = 8;
+    let mut mem = new_memory(&dir, dim);
+
+    let mut seed = 99;
+    let vectors: Vec<Vec<f32>> = (0..50).map(|_| make_vector(&mut seed, dim)).collect();
+    let entries: Vec<MemoryEntry> = vectors
+        .iter()
+        .enumerate()
+        .map(|(i, v)| entry(&format!("b{i}"), v.clone(), &format!("b{i}")))
+        .collect();
+    mem.save_batch(entries).unwrap();
+
+    // Exact-match queries should resolve to themselves after a batch insert.
+    for probe in [0usize, 17, 49] {
+        let hits = mem.hybrid_search(&vectors[probe], "", 1.0, 0.0, 1);
+        assert_eq!(hits[0].index, probe, "batch-inserted vector {probe} not found");
+    }
+}