diff --git a/BENCHMARKS.md b/BENCHMARKS.md index a672bbe..1e4b98d 100644 --- a/BENCHMARKS.md +++ b/BENCHMARKS.md @@ -149,6 +149,49 @@ vectors, and recall is measured against brute-force ground truth rather than against the f32 index, whose own approximation errors a re-scored search is entitled to get right. +### Search options: source filters, re-ranking, confidence + +Measured 2026-09-24 on tank (AMD Ryzen 7 7800X3D). `HDF5Memory::search` +with `SearchOptions`, clustered 384-dim data, k = 10, Hebbian boosting off. +Filters keep 50%, 10% or 1% of the store at random, or two whole clusters +chosen *away* from each query's own — the case an ANN index handles worst, +because nothing it finds near the query is allowed. Recall is vector-only +against an exact scan of the allowed records; latency is full hybrid search +(vector + BM25). 200 queries, medians of three runs (recall was identical in +every run). + +```bash +cargo run --release -p clawhdf5-bench --bin search_harness -- --options-study --full +``` + +| N | options | filtered recall@10 | p50 ms | p99 ms | +|---:|---|---:|---:|---:| +| 10 000 | no filter | 1.0000 | 0.488 | 0.518 | +| 10 000 | random 50% | 1.0000 | 0.520 | 0.563 | +| 10 000 | random 10% | 1.0000 | 0.342 | 0.367 | +| 10 000 | random 1% | 1.0000 | 0.220 | 0.239 | +| 10 000 | 2 clusters away from the query | 1.0000 | 0.252 | 0.278 | +| 10 000 | re-rank | — | 0.562 | 0.674 | +| 10 000 | re-rank + confidence | — | 0.554 | 0.575 | +| 100 000 | no filter | 0.9995 | 4.600 | 5.245 | +| 100 000 | random 50% | 1.0000 | 4.770 | 5.614 | +| 100 000 | random 10% | 1.0000 | 3.244 | 4.008 | +| 100 000 | random 1% | 1.0000 | 2.288 | 2.943 | +| 100 000 | 2 clusters away from the query | 1.0000 | 2.298 | 3.048 | +| 100 000 | re-rank | — | 4.747 | 5.411 | +| 100 000 | re-rank + confidence | — | 4.748 | 5.526 | + +A filtered search finds the exact filtered top 10 and is never slower than an +unfiltered one. The filter applies before ranking — over-fetching the index in +proportion to what the filter removes, and scanning the allowed records +exactly whenever that costs fewer distance evaluations than the index would +(roughly `pool × M`). The first version compared the over-fetch to the store +size instead, and that measured badly at 100K: the 1% filter took 5.9 ms at +recall 0.9965 and the away-from-query filter 12.3 ms at 0.976, both through an +index asked for ~16 000 candidates, where scanning the few hundred or thousand +allowed records is exact and cheap. Re-ranking a 3k candidate pool and +confidence rejection add about 3%. + ### float16 embedding storage (`MemoryConfig::float16`) Measured 2026-09-23 on tank (AMD Ryzen 7 7800X3D). The same clustered @@ -183,6 +226,30 @@ The in-memory cache holds the half-rounded values, so the store searches the same before and after a reopen; RAM use is unchanged (the cache is still `f32`). What `float16` saves is disk, and the I/O that goes with it. +**On real embeddings.** The table above is synthetic clustered data. The full +LongMemEval haystack (`longmemeval_s`, 500 questions, ~494 turns each) with +real all-MiniLM-L6-v2 embeddings, run once with `f32` stores and once with +`--float16`, measured 2026-09-24 on tank (embeddings on an RTX 5060 Ti): + +```bash +cargo run --release -p clawhdf5-bench --bin longmemeval_bench --features embeddings-cuda -- \ + benchmarks/longmemeval/longmemeval_s.json --embeddings weights/all-minilm-l6-v2 [--float16] +``` + +| mode | turn Hit@1 | turn Hit@5 | turn Hit@10 | turn MRR | session Hit@5 | session MRR | +|---|---:|---:|---:|---:|---:|---:| +| Hybrid 0.4 / 0.6, f32 | 51.6% | 81.4% | 87.8% | 0.6430 | 96.8% | 0.9347 | +| Hybrid 0.4 / 0.6, float16 | 51.6% | 81.4% | 87.8% | 0.6430 | 96.8% | 0.9347 | +| Vector only, f32 | 36.0% | 71.8% | 81.6% | 0.5031 | 94.2% | 0.8901 | +| Vector only, float16 | 36.0% | 71.8% | 81.6% | 0.5031 | 94.2% | 0.8901 | + +All eight modes the harness runs (BM25, vector, hybrid, RRF, stemmed, and +both re-rank variants) were identical at every Hit@k and MRR, turn and session +level, except RRF's session MRR (0.9253 vs 0.9254) and one or two flips in +which of two gold sessions ranks first, out of ~320. Those flips show the +half-precision path was in effect; they do not change a single hit. The f32 +run reproduces the published hybrid numbers exactly. + ### Opening a store (`read_from_disk`) `HDF5Memory::open` memory-mapped the file, copied the whole mapping into a diff --git a/CHANGELOG.md b/CHANGELOG.md index 7d2a4f1..2996da6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,12 +11,18 @@ each checkpoint, so it becomes readable at its next checkpoint on this version; other files with `f32` or empty datasets need rewriting. Details in `docs/known-issues.md`. -- **`MemoryConfig::float16` now does what it says.** It was persisted and - otherwise ignored; embeddings were always stored as `f32`. A store created - with it on now writes half-precision embeddings (48% smaller files) and - rounds embeddings to half precision as they are saved. A store that already - had `float16 = true` rounds its embeddings when next opened and writes them - as `float16` at its next checkpoint. Off by default. +- **New stores store embeddings as half precision by default.** + `MemoryConfig::float16` was persisted and otherwise ignored; it now writes + `float16` embeddings (48% smaller files at 100K) and rounds each embedding + to half precision as it is saved — and it defaults to `true` for new + stores. On the full LongMemEval haystack with real MiniLM embeddings every + retrieval metric matched `f32`. **Existing stores are unaffected**: every + agent store has recorded `float16 = false`, and keeps it (a v2.5.0 fixture + guards this). A store that already had `float16 = true` rounds its + embeddings when next opened and writes them as `float16` at its next + checkpoint. Opt out with `float16 = false` or `create --f32`; the CLI's + `--float16` is still accepted and now a no-op. Values beyond ±65504 are + refused, so keep `f32` for unnormalised vectors. - **Breaking:** `MemoryError` gained `InvalidEntry`, returned when a `float16` store is given an embedding value beyond ±65504. Exhaustive matches need the new arm. @@ -40,6 +46,28 @@ `quantized_index = false`, or pass `create --f32-index` to the CLI, to opt out. The CLI's `--quantized-index` is still accepted but is now a no-op. +### Search +- `clawhdf5-agent`: **`HDF5Memory::search` with `SearchOptions`** — source + filtering, re-ranking and confidence rejection in the store's own search + path. Re-ranking and confidence rejection used to be reachable only + through the OpenClaw backend, which now calls `search` with both on. + - `with_sources([..])` restricts a search to records from those source + channels. It applies before ranking, so a filtered search still returns up + to `k` results, normalised over what it can return. Measured at 100K: the + exact filtered top 10 for filters keeping 50%, 10% and 1% of the store and + for records far from the query, and never slower than an unfiltered search + (2.3 ms for a 1% filter vs 4.6 ms unfiltered). See `BENCHMARKS.md`, + "Search options". + - `with_rerank(ReRankConfig)` re-ranks a pool of `max(3k, 10)` candidates + (`rerank_pool` to change it) by relevance, recency, source authority and + activation; `with_confidence(ConfidenceConfig)` drops low-confidence + results; `at_time(now)` pins the clock for recency. About 3% on latency. + - `hybrid_search` and `hybrid_search_with` are unchanged (tested bit for + bit against `search` with default options). +- `clawhdf5-agent`: the OpenClaw backend's search now boosts the Hebbian + activation of the `k` results it returns, not of the whole `3k` candidate + pool it re-ranks. + ### Interop - `clawhdf5-format`: **every `f32` dataset was unreadable by h5py and libhdf5.** The float datatype encoder hard-coded the sign bit's position to @@ -67,7 +95,9 @@ precision.** At 100K x 384 the file goes from 154.0 to 80.8 MiB (−48%), a checkpoint from 752 to 512 ms and open from 300 to 252 ms, with the same vector recall@10 against an exact scan (0.999 vs 0.994) and the same - `hybrid_search` latency; at 10K open is 3 ms slower. The cache rounds each + `hybrid_search` latency; at 10K open is 3 ms slower. On the full + LongMemEval haystack with real MiniLM embeddings every retrieval metric is + identical to `f32` (`longmemeval_bench --float16`). The cache rounds each embedding as it is saved, so memory and file agree bit for bit and a store returns the same results before and after a reopen (tested). Out-of-range values are refused with `MemoryError::InvalidEntry` rather than stored as @@ -99,6 +129,11 @@ now an error. ### Defaults +- `clawhdf5-agent`: `MemoryConfig::float16` defaults to `true` for new stores, + measured rather than assumed: identical LongMemEval retrieval on real + embeddings, 48% smaller files and faster checkpoints and opens at 100K. + `clawhdf5-cli create --f32` opts out; like `--f32-index`, it only ever + switches the default off. - `clawhdf5-agent`: `MemoryConfig::quantized_index` defaults to `true` for new stores. The reason it had been off — that int8 search was slower on ARM — did not survive measurement (see Corrections). Stores that predate the diff --git a/CLAUDE.md b/CLAUDE.md index 5423c38..21b47d8 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -87,16 +87,26 @@ Cargo workspace with 16 crates under `crates/` (plus `libaec-sys`, an internal F `export` do). An unreadable WAL (torn header, bad magic) is quarantined to `.h5.wal.corrupt-` rather than blocking `open()`; a WAL with an unknown *newer* version still fails and is left untouched. -- `MemoryConfig::float16` (off by default, persisted; CLI `create --float16`) - writes `/memory/embeddings` as IEEE half precision (48% smaller file at - 100K, same recall). `MemoryCache::half_precision` rounds each embedding as - it enters the cache (push, update, WAL replay, and on load of a store still +- `MemoryConfig::float16` (**on by default** for new stores, persisted; + existing stores keep their recorded `false` — guarded by the v2.5.0 + fixture in `tests/float16_store.rs`; CLI opt-out is `create --f32`) writes + `/memory/embeddings` as IEEE half precision (48% smaller file at 100K; + LongMemEval with real MiniLM embeddings identical to f32). + `MemoryCache::half_precision` rounds each embedding as it enters the cache (push, update, WAL replay, and on load of a store still `f32` on disk), so memory and file agree bit for bit; the conversions live in `clawhdf5_format::float16` and must stay the single implementation. Values beyond ±65504 are `MemoryError::InvalidEntry`. Interop: every file must open in h5py — `f32` datasets and empty datasets did not until 2026-09-23 (see `docs/known-issues.md`); the agent's `h5py_interop` test guards a whole store. +- `HDF5Memory::search(query_emb, text, &SearchOptions)` is the full search + path: optional source-channel filter (applied before ranking; exact scan of + the allowed records whenever cheaper than `pool × M` index distance + evaluations, and as the fallback when the pool comes back short), fusion, + activation scaling, optional re-ranking and confidence rejection. + `hybrid_search`/`hybrid_search_with` are thin wrappers; the OpenClaw + backend is `search` with re-rank + confidence on. Measure changes with + `search_harness --options-study`. - `MemoryConfig::compression` is off by default; when on, embeddings are deflate-compressed, or Zstd with the agent's `zstd` feature (links libzstd). - `Dataset::verify_provenance()` (clawhdf5 facade, `provenance` feature, on by diff --git a/README.md b/README.md index f661c4b..a96553b 100644 --- a/README.md +++ b/README.md @@ -82,14 +82,18 @@ breaking change, are in [CHANGELOG.md](CHANGELOG.md). 100K × 384 store to 1.74× the raw vectors. At equal recall it is also faster than `f32`: 1.63× QPS on AVX2, 1.18× on a Raspberry Pi 5 (NEON `SDOT`). -**Interop (unreleased)** +**Interop and search (unreleased)** - **Files we write now open in h5py and libhdf5.** Every `f32` dataset — including every agent store's embeddings — and every empty dataset was refused by libhdf5. Both were write-side bugs in every release; agent stores fix themselves at their next checkpoint. See [docs/known-issues.md](docs/known-issues.md). - `MemoryConfig::float16` now stores half-precision embeddings (it was - ignored): 48% smaller files at the same recall. + ignored), and is on by default for new stores: 48% smaller files, and + identical LongMemEval retrieval on real embeddings. +- `HDF5Memory::search` with `SearchOptions`: filter by source channel (exact + filtered top-k, never slower than unfiltered), and opt-in re-ranking and + confidence rejection, which used to be OpenClaw-only. **Tooling** - CI now runs the h5py/netCDF4 interop suites for real (they had been skipping @@ -253,8 +257,9 @@ retrieval recall reported as QA accuracy typically overstates by 20–30 points. | 10K | 17.0 MB | 1.7 KB | 2.7 MB (6.2x) | | 100K | 169.8 MB | 1.7 KB | 26.9 MB (6.2x) | -With `MemoryConfig::float16` the embeddings take half the space: an agent -store of 100K × 384 records is 80.8 MiB instead of 154.0. +These figures are `f32` embeddings. New agent stores default to +`MemoryConfig::float16`, which halves them: 100K × 384 records take 80.8 MiB +instead of 154.0. **In memory** — a store reopened from disk, 384-dim `f32`, measured with a counting allocator ([BENCHMARKS.md § Memory footprint](BENCHMARKS.md#memory-footprint)): @@ -293,12 +298,14 @@ ClawhDF5's agent memory engine draws on 15+ recent papers on agentic memory syst └────────┬────────┘ │ ┌─────────────────▼──────────────────┐ - │ HDF5Memory::hybrid_search │ + │ HDF5Memory::search │ + │ optional source-channel filter │ │ HNSW vector + BM25 keyword │ │ weighted fusion (0.4 / 0.6) │ │ × √(Hebbian activation) │ └─────────────────┬──────────────────┘ - │ OpenClaw backend adds: + │ opt-in (SearchOptions); + │ the OpenClaw backend turns both on ┌─────────────────▼──────────────────┐ │ Multi-factor re-ranking │ │ relevance · recency · authority · │ @@ -334,8 +341,8 @@ directly; the store persists the records, sessions and graph they work over. | **`knowledge`** | Entity/relation graph with BFS traversal, spreading activation, fuzzy (Levenshtein) entity resolution | | **`consolidation`** | Three-tier memory (Working → Episodic → Semantic) with importance scoring, novelty, and time-decay | | **`hybrid`** | Vector + BM25 fusion. Default is a min-max-normalised weighted sum, vector 0.4 / keyword 0.6 (`hybrid::DEFAULT_FUSION`, tuned on LongMemEval); RRF is available via `Fusion::Rrf` / `hybrid_search_with`. The vector stage uses the HNSW index by default (`hnsw` feature); disable with `--no-default-features --features float16` for an exact linear scan | -| **`reranker`** | Multi-factor re-ranking: retrieval relevance (leads, weight 1.0), temporal recency, source authority, activation weight. Used by the OpenClaw backend | -| **`confidence`** | Low-confidence rejection — suppresses spurious recalls when nothing matches (OpenClaw backend) | +| **`reranker`** | Multi-factor re-ranking: retrieval relevance (leads, weight 1.0), temporal recency, source authority, activation weight. Opt-in via `SearchOptions::with_rerank`; on in the OpenClaw backend | +| **`confidence`** | Low-confidence rejection — suppresses spurious recalls when nothing matches. Opt-in via `SearchOptions::with_confidence`; on in the OpenClaw backend | | **`temporal`** | Sorted timestamp index, session DAG, entity timeline, temporal query hints | | **`multimodal`** | Cross-modal search across text/image/audio/video embeddings | | **`provenance`** | Source attribution and an unkeyed FNV-1a content hash per record, held in memory for the session, for detecting accidental corruption (not tamper-proof) | @@ -401,6 +408,31 @@ for result in results { } ``` +### Search Options + +```rust +use clawhdf5_agent::SearchOptions; +use clawhdf5_agent::confidence::ConfidenceConfig; +use clawhdf5_agent::reranker::ReRankConfig; + +// Only memories from these source channels; still a full page of k results. +let work = memory.search( + &query_embedding, + "deadline", + &SearchOptions::new(5).with_sources(["slack", "email"]), +); + +// Re-rank by relevance, recency, source authority and activation, then drop +// low-confidence results — the pipeline the OpenClaw backend runs. +let careful = memory.search( + &query_embedding, + "user preferences", + &SearchOptions::new(5) + .with_rerank(ReRankConfig::default()) + .with_confidence(ConfidenceConfig::default()), +); +``` + ### Knowledge Graph ```rust @@ -580,12 +612,14 @@ setting existed keep their `f32` index; opt out for new stores with `quantized_index = false` or `clawhdf5-cli create --f32-index`. See [BENCHMARKS.md § Quantising the index copy](BENCHMARKS.md#quantising-the-index-copy-quantized_index). -`MemoryConfig::float16` (off by default; CLI `create --float16`) stores the +`MemoryConfig::float16` (**on by default** for new stores) stores the embeddings on disk as IEEE half precision (numpy `float16`): at 100K × 384 the -file drops from 154 to 81 MiB, checkpoints and opens get faster, and vector -recall and search latency do not change. Embeddings are rounded as they are -saved, so the store searches the same before and after a reopen; values must -lie within ±65504. See +file drops from 154 to 81 MiB, checkpoints and opens get faster, and on the +full LongMemEval haystack with real MiniLM embeddings every retrieval metric +matches `f32`. Embeddings are rounded as they are saved, so the store searches +the same before and after a reopen; values must lie within ±65504. Existing +stores keep their setting. Opt out with `float16 = false` or +`clawhdf5-cli create --f32` — e.g. for unnormalised vectors. See [BENCHMARKS.md § float16 embedding storage](BENCHMARKS.md#float16-embedding-storage-memoryconfigfloat16). ### `clawhdf5-format` diff --git a/crates/clawhdf5-agent/src/hybrid.rs b/crates/clawhdf5-agent/src/hybrid.rs index 21d94d5..42f1fc5 100644 --- a/crates/clawhdf5-agent/src/hybrid.rs +++ b/crates/clawhdf5-agent/src/hybrid.rs @@ -65,31 +65,34 @@ pub fn hybrid_search_fused( ) -> Vec<(usize, f32)> { // Get raw scores from both systems. Request all results so normalization // covers the full distribution. - // Use parallel search when rayon feature is enabled and vector count > 10K. - let vec_scores = { - #[cfg(feature = "parallel")] - { - if vectors.count() > 10_000 { - vector_search::parallel_cosine_batch( - query_embedding, - vectors, - tombstones, - vectors.count(), - ) - } else { - vector_search::cosine_similarity_batch(query_embedding, vectors, tombstones) - } - } - #[cfg(not(feature = "parallel"))] - { - vector_search::cosine_similarity_batch(query_embedding, vectors, tombstones) - } - }; + let vec_scores = exact_vector_scores(query_embedding, vectors, tombstones); let kw_scores = bm25_index.scores(query_text); fuse(vec_scores, kw_scores, fusion, k) } +/// Cosine similarity of `query_embedding` to every vector whose `skip` byte is +/// 0 (a tombstone, or any other exclusion mask). Parallel above 10K vectors +/// when the `parallel` feature is on. +pub fn exact_vector_scores( + query_embedding: &[f32], + vectors: &(impl crate::vector_search::VectorSet + Sync + ?Sized), + skip: &[u8], +) -> Vec<(usize, f32)> { + #[cfg(feature = "parallel")] + { + if vectors.count() > 10_000 { + return vector_search::parallel_cosine_batch( + query_embedding, + vectors, + skip, + vectors.count(), + ); + } + } + vector_search::cosine_similarity_batch(query_embedding, vectors, skip) +} + /// 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 diff --git a/crates/clawhdf5-agent/src/lib.rs b/crates/clawhdf5-agent/src/lib.rs index 4607449..b74507f 100644 --- a/crates/clawhdf5-agent/src/lib.rs +++ b/crates/clawhdf5-agent/src/lib.rs @@ -72,6 +72,7 @@ use ephemeral::{EphemeralConfig, EphemeralStore}; pub use ephemeral::{EphemeralEntry, EphemeralStats}; use knowledge::KnowledgeCache; use memory_strategy::{Exchange, MemoryStrategy, StrategyOutput}; +pub use search::SearchOptions; use session::SessionCache; // --- Error type --- @@ -135,6 +136,13 @@ pub struct MemoryConfig { /// so search results are the same before and after a reopen. Values must /// lie within ±65504; a save outside that is `MemoryError::InvalidEntry`. /// Fixed when the store is created (persisted in `/meta`). + /// + /// **On by default for new stores**: on the full LongMemEval haystack with + /// real MiniLM embeddings every retrieval metric matched `f32`, and at + /// 100K records the file is 48% smaller (`BENCHMARKS.md`). Existing + /// stores keep the setting they were created with. Set it to `false` for + /// full-precision embeddings, e.g. for unnormalised vectors that may + /// exceed the half-precision range. pub float16: bool, pub compression: bool, pub compression_level: u32, @@ -190,7 +198,7 @@ impl MemoryConfig { embedding_dim, chunk_size: 512, overlap: 50, - float16: false, + float16: true, compression: false, compression_level: 0, compact_threshold: 0.3, diff --git a/crates/clawhdf5-agent/src/openclaw.rs b/crates/clawhdf5-agent/src/openclaw.rs index 40e58de..f90f8f6 100644 --- a/crates/clawhdf5-agent/src/openclaw.rs +++ b/crates/clawhdf5-agent/src/openclaw.rs @@ -13,9 +13,8 @@ use std::path::{Path, PathBuf}; use std::time::{SystemTime, UNIX_EPOCH}; use crate::{ - AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry, - confidence::{ConfidenceConfig, ScoredResult, reject_low_confidence}, - reranker::{ReRankConfig, RerankInput, rerank}, + AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry, SearchOptions, + confidence::ConfidenceConfig, reranker::ReRankConfig, }; // ───────────────────────────────────────────────────────────────────────────── @@ -524,71 +523,27 @@ impl ClawhdfBackend { impl MemoryBackend for ClawhdfBackend { /// Search using hybrid vector + BM25 retrieval, then re-rank and - /// confidence-filter. + /// confidence-filter — [`HDF5Memory::search`] with both stages on. fn search( &mut self, query_text: &str, query_embedding: &[f32], k: usize, ) -> Vec { - // 1. Hybrid retrieval (vector + BM25, fused by score). - let candidates = k.saturating_mul(3).max(10); - let raw = self.memory.hybrid_search_with( - query_embedding, - query_text, - crate::hybrid::DEFAULT_FUSION, - candidates, - ); - - if raw.is_empty() { - return Vec::new(); - } - - let now = Self::now_secs(); - - // 2. Re-rank using temporal recency, source authority, Hebbian weight. - let rerank_inputs: Vec = raw - .iter() - .map(|r| RerankInput { - index: r.index, - timestamp: r.timestamp, - source_channel: r.source_channel.clone(), - raw_activation: r.activation, - relevance: r.score, - }) - .collect(); - - let reranked = rerank(&rerank_inputs, &self.rerank_config, now); - - // 3. Confidence rejection. - let scored: Vec = reranked - .iter() - .map(|r| ScoredResult { - index: r.index, - score: r.combined_score, - }) - .collect(); - - let confident = reject_low_confidence(&scored, &self.confidence_config); - - // 4. Map back to MemorySearchResult; preserve raw text via index lookup. - let raw_by_idx: HashMap = - raw.iter().map(|r| (r.index, r)).collect(); - - confident + let options = SearchOptions::new(k) + .with_rerank(self.rerank_config) + .with_confidence(self.confidence_config.clone()) + .at_time(Self::now_secs()); + self.memory + .search(query_embedding, query_text, &options) .into_iter() - .take(k) - .filter_map(|sr| { - let r = raw_by_idx.get(&sr.index)?; - let path = r.source_channel.clone(); - Some(MemorySearchResult { - text: r.chunk.clone(), - score: sr.score, - path: path.clone(), - line_range: None, - timestamp: Some(r.timestamp), - source: path, - }) + .map(|r| MemorySearchResult { + text: r.chunk, + score: r.score, + path: r.source_channel.clone(), + line_range: None, + timestamp: Some(r.timestamp), + source: r.source_channel, }) .collect() } diff --git a/crates/clawhdf5-agent/src/search.rs b/crates/clawhdf5-agent/src/search.rs index 51398c0..3687c51 100644 --- a/crates/clawhdf5-agent/src/search.rs +++ b/crates/clawhdf5-agent/src/search.rs @@ -2,18 +2,107 @@ use std::path::Path; +use std::collections::HashSet; + use crate::bm25; +use crate::confidence::{ConfidenceConfig, ScoredResult, reject_low_confidence}; use crate::hybrid; +use crate::reranker::{ReRankConfig, RerankInput, rerank}; use crate::{HDF5Memory, MAX_ACTIVATION_WEIGHT, MemoryError, Result, SearchResult}; +/// Options for [`HDF5Memory::search`]. +/// +/// [`SearchOptions::new`] is plain hybrid search with the tuned default +/// fusion — the same as `hybrid_search_with(.., hybrid::DEFAULT_FUSION, k)`. +/// Every stage beyond that is opt-in. +#[derive(Debug, Clone)] +pub struct SearchOptions { + /// Number of results to return. + pub k: usize, + /// How the vector and keyword stages are combined. + pub fusion: hybrid::Fusion, + /// Only consider records whose `source_channel` is one of these. The + /// filter applies *before* ranking, so a filtered search still returns up + /// to `k` results and scores are normalised over the records it can + /// return. `None` searches everything; an empty list matches nothing. + pub source_channels: Option>, + /// Re-rank a candidate pool by retrieval relevance, recency, source + /// authority and activation — the pipeline the OpenClaw backend runs. + pub rerank: Option, + /// Candidates retrieved for re-ranking; 0 means `max(3k, 10)`. + pub rerank_pool: usize, + /// Drop low-confidence results (after re-ranking, when that is on). + pub confidence: Option, + /// The time recency is measured from, in seconds since the epoch. + /// `None` uses the system clock. + pub now: Option, +} + +impl SearchOptions { + pub fn new(k: usize) -> Self { + Self { + k, + fusion: hybrid::DEFAULT_FUSION, + source_channels: None, + rerank: None, + rerank_pool: 0, + confidence: None, + now: None, + } + } + + pub fn with_fusion(mut self, fusion: hybrid::Fusion) -> Self { + self.fusion = fusion; + self + } + + /// Search only records from these source channels. + pub fn with_sources>(mut self, channels: impl IntoIterator) -> Self { + self.source_channels = Some(channels.into_iter().map(Into::into).collect()); + self + } + + pub fn with_rerank(mut self, config: ReRankConfig) -> Self { + self.rerank = Some(config); + self + } + + pub fn with_confidence(mut self, config: ConfidenceConfig) -> Self { + self.confidence = Some(config); + self + } + + /// Measure recency from `now` (seconds since the epoch) instead of the + /// system clock — for reproducible results and tests. + pub fn at_time(mut self, now: f64) -> Self { + self.now = Some(now); + self + } +} + +impl Default for SearchOptions { + fn default() -> Self { + Self::new(10) + } +} + impl HDF5Memory { - /// Vector + keyword scoring stage of [`HDF5Memory::hybrid_search`]. + /// Vector + keyword scoring stage of [`HDF5Memory::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`]. + /// + /// `exclude`, when given, marks records that must not be returned (1 = + /// excluded; it covers tombstones too). The index is over-fetched in + /// proportion to how much the mask removes. Surfacing `pool` candidates + /// costs the index roughly `pool × M` distance evaluations, while an exact + /// scan of the allowed records costs one each — so whenever that scan is + /// the cheaper of the two it is used instead, and it is also the fallback + /// if the pool comes back with too few allowed hits (the allowed records + /// sit away from the query). A filtered search never comes back short. #[cfg(feature = "hnsw")] fn vector_keyword_search( &mut self, @@ -22,16 +111,30 @@ impl HDF5Memory { bm25: &bm25::BM25Index, fusion: hybrid::Fusion, k: usize, + exclude: Option<&[u8]>, ) -> Vec<(usize, f32)> { self.ensure_hnsw_fresh(); + let n = self.cache.len(); + // Over-fetch so the merge sees a useful vector pool. `ef` is + // configurable, but the pool the fusion stage sees is not tied to it: + // a caller lowering `ef` for speed should not silently narrow what + // fusion has to work with. + let mut pool = (k * 8).max(64); + let mut allowed = n; + if let Some(ex) = exclude { + allowed = ex.iter().filter(|&&e| e == 0).count(); + if allowed == 0 { + return Vec::new(); + } + // Expect `pool` allowed hits if the filter is independent of the + // query's neighbourhood. + pool = pool.saturating_mul(n).div_ceil(allowed); + if allowed <= pool.saturating_mul(self.hnsw_m()) { + return self.exact_masked_search(query_embedding, query_text, bm25, fusion, k, ex); + } + } 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). - // `ef` is configurable, but the pool the fusion stage sees is - // not tied to it: a caller lowering `ef` for speed should not - // silently narrow what fusion has to work with. - let pool = (k * 8).max(64); let ef = self.hnsw_ef_search(k).max(pool); let candidates = index.search(query_embedding, pool, ef); // A quantised index returns approximate distances, and no @@ -42,6 +145,7 @@ impl HDF5Memory { let exact = index.storage() == clawhdf5_ann::Storage::Int8; let vec_scores: Vec<(usize, f32)> = candidates .into_iter() + .filter(|(id, _)| exclude.is_none_or(|ex| ex[*id] == 0)) .map(|(id, dist)| { let score = if exact { crate::vector_search::cosine_similarity( @@ -56,10 +160,54 @@ impl HDF5Memory { .collect(); // Fusion normalises over every keyword match, so it needs all // the scores — but not ranked. - let kw_scores = bm25.scores(query_text); + let mut kw_scores = bm25.scores(query_text); + if let Some(ex) = exclude { + if vec_scores.len() < k.min(allowed) { + // The allowed records are not where the index looked. + return self.exact_masked_search( + query_embedding, + query_text, + bm25, + fusion, + k, + ex, + ); + } + kw_scores.retain(|(id, _)| ex[*id] == 0); + } hybrid::fuse(vec_scores, kw_scores, fusion, k) } - _ => hybrid::hybrid_search_fused( + _ => match exclude { + Some(ex) => { + self.exact_masked_search(query_embedding, query_text, bm25, fusion, k, ex) + } + None => hybrid::hybrid_search_fused( + query_embedding, + query_text, + &self.cache.embeddings, + &self.cache.chunks, + &self.cache.tombstones, + bm25, + fusion, + k, + ), + }, + } + } + + #[cfg(not(feature = "hnsw"))] + fn vector_keyword_search( + &mut self, + query_embedding: &[f32], + query_text: &str, + bm25: &bm25::BM25Index, + fusion: hybrid::Fusion, + k: usize, + exclude: Option<&[u8]>, + ) -> Vec<(usize, f32)> { + match exclude { + Some(ex) => self.exact_masked_search(query_embedding, query_text, bm25, fusion, k, ex), + None => hybrid::hybrid_search_fused( query_embedding, query_text, &self.cache.embeddings, @@ -72,25 +220,33 @@ impl HDF5Memory { } } - #[cfg(not(feature = "hnsw"))] - fn vector_keyword_search( - &mut self, + /// Exact hybrid search over the records `exclude` leaves (0 = allowed). + fn exact_masked_search( + &self, query_embedding: &[f32], query_text: &str, bm25: &bm25::BM25Index, fusion: hybrid::Fusion, k: usize, + exclude: &[u8], ) -> Vec<(usize, f32)> { - hybrid::hybrid_search_fused( - query_embedding, - query_text, - &self.cache.embeddings, - &self.cache.chunks, - &self.cache.tombstones, - bm25, - fusion, - k, - ) + let vec_scores = + hybrid::exact_vector_scores(query_embedding, &self.cache.embeddings, exclude); + let mut kw_scores = bm25.scores(query_text); + kw_scores.retain(|(id, _)| exclude.get(*id) == Some(&0)); + hybrid::fuse(vec_scores, kw_scores, fusion, k) + } + + /// The exclusion mask for a source-channel filter: 1 for a tombstoned + /// record or one from a channel not in `channels`. + fn source_mask(&self, channels: &[String]) -> Vec { + let allowed: HashSet<&str> = channels.iter().map(String::as_str).collect(); + self.cache + .source_channels + .iter() + .zip(&self.cache.tombstones) + .map(|(ch, &t)| u8::from(t != 0 || !allowed.contains(ch.as_str()))) + .collect() } /// Perform hybrid search combining cosine vector similarity and BM25 keyword search. @@ -125,12 +281,53 @@ impl HDF5Memory { fusion: hybrid::Fusion, k: usize, ) -> Vec { + self.search( + query_embedding, + query_text, + &SearchOptions::new(k).with_fusion(fusion), + ) + } + + /// Hybrid search with optional source filtering, re-ranking and + /// confidence rejection — see [`SearchOptions`]. + /// + /// Stages, in order: vector + keyword retrieval over the records the + /// source filter allows; fusion; scaling by Hebbian activation; re-ranking + /// (if on) of a `rerank_pool` of candidates; confidence rejection (if on); + /// the top `k`. The records returned with a positive score get their + /// Hebbian boost. + pub fn search( + &mut self, + query_embedding: &[f32], + query_text: &str, + options: &SearchOptions, + ) -> Vec { + let k = options.k; + let fetch = match options.rerank { + Some(_) if options.rerank_pool > 0 => options.rerank_pool.max(k), + Some(_) => k.saturating_mul(3).max(10), + None => k, + }; + let exclude = options + .source_channels + .as_deref() + .map(|channels| self.source_mask(channels)); + // 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, &bm25, fusion, k); + let scored = self.vector_keyword_search( + query_embedding, + query_text, + &bm25, + options.fusion, + fetch, + exclude.as_deref(), + ); + self.bm25 = Some(bm25); + let mut results: Vec = scored .into_iter() .map(|(idx, score)| { @@ -154,6 +351,25 @@ impl HDF5Memory { .then(a.index.cmp(&b.index)) }); + if let Some(config) = &options.rerank { + results = Self::rerank_results(results, config, options.now); + } + if let Some(config) = &options.confidence { + let scored: Vec = results + .iter() + .map(|r| ScoredResult { + index: r.index, + score: r.score, + }) + .collect(); + let keep: HashSet = reject_low_confidence(&scored, config) + .into_iter() + .map(|r| r.index) + .collect(); + results.retain(|r| keep.contains(&r.index)); + } + results.truncate(k); + // Only reinforce records that actually matched. When fewer than `k` // records are relevant, the rest of the list is zero-score filler; // boosting it would teach the store that arbitrary records are @@ -164,11 +380,45 @@ impl HDF5Memory { .map(|r| r.index) .collect(); self.apply_hebbian_boost(&hit_indices); - self.bm25 = Some(bm25); results } + /// Reorder by the re-ranker's combined score, which also becomes each + /// result's `score`. + fn rerank_results( + results: Vec, + config: &ReRankConfig, + now: Option, + ) -> Vec { + let now = now.unwrap_or_else(|| { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_secs_f64()) + .unwrap_or(0.0) + }); + let inputs: Vec = results + .iter() + .map(|r| RerankInput { + index: r.index, + timestamp: r.timestamp, + source_channel: r.source_channel.clone(), + raw_activation: r.activation, + relevance: r.score, + }) + .collect(); + let mut by_index: std::collections::HashMap = + results.into_iter().map(|r| (r.index, r)).collect(); + rerank(&inputs, config, now) + .into_iter() + .filter_map(|rr| { + let mut r = by_index.remove(&rr.index)?; + r.score = rr.combined_score; + Some(r) + }) + .collect() + } + /// 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 diff --git a/crates/clawhdf5-agent/tests/float16_store.rs b/crates/clawhdf5-agent/tests/float16_store.rs index 5c83107..0c20d62 100644 --- a/crates/clawhdf5-agent/tests/float16_store.rs +++ b/crates/clawhdf5-agent/tests/float16_store.rs @@ -206,3 +206,55 @@ fn wal_replay_rounds_like_a_live_save() { assert_eq!(recovered.count(), 30); assert_eq!(search_bits(&mut recovered, 77), live); } + +#[test] +fn new_stores_default_to_float16() { + let dir = TempDir::new().unwrap(); + let path = dir.path().join("default.h5"); + let mut m = HDF5Memory::create(MemoryConfig::new(path.clone(), "agent", DIM)).unwrap(); + assert!(m.config().float16); + m.save_batch((0..10).map(entry).collect()).unwrap(); + drop(m); + assert_eq!(embeddings_dtype_and_values(&path).0, "Other(\"float16\")"); + assert!(HDF5Memory::open(&path).unwrap().config().float16); +} + +#[test] +fn an_existing_f32_store_stays_f32() { + // Written by the v2.5.0 CLI, with `float16 = 0` in /meta (every agent + // store has recorded it). Flipping the default for new stores must not + // reach back and round an existing store's embeddings. + let dir = TempDir::new().unwrap(); + let path = dir.path().join("legacy.h5"); + std::fs::copy( + concat!( + env!("CARGO_MANIFEST_DIR"), + "/tests/fixtures/store_v2_5_0.h5" + ), + &path, + ) + .unwrap(); + let before = embeddings_dtype_and_values(&path); + assert_eq!(before.0, "F32"); + + let mut m = HDF5Memory::open(&path).unwrap(); + assert!(!m.config().float16, "an old store must reopen as f32"); + let dim = m.config().embedding_dim; + let odd: Vec = (0..dim).map(|i| 0.1 + i as f32 * 1e-4).collect(); + m.save_batch(vec![MemoryEntry { + chunk: "added after the upgrade".into(), + embedding: odd.clone(), + source_channel: "test".into(), + timestamp: 1.0, + session_id: "s".into(), + tags: String::new(), + }]) + .unwrap(); + drop(m); + + // Checkpointed: still f32, the old rows untouched and the new one exact. + let (dtype, values) = embeddings_dtype_and_values(&path); + assert_eq!(dtype, "F32"); + assert_eq!(&values[..before.1.len()], before.1.as_slice()); + assert_eq!(&values[before.1.len()..], odd.as_slice()); +} diff --git a/crates/clawhdf5-agent/tests/search_options.rs b/crates/clawhdf5-agent/tests/search_options.rs new file mode 100644 index 0000000..14fba9b --- /dev/null +++ b/crates/clawhdf5-agent/tests/search_options.rs @@ -0,0 +1,344 @@ +//! `HDF5Memory::search` with `SearchOptions`: source filtering, re-ranking and +//! confidence rejection in the store's own search path. + +use std::collections::HashSet; + +use clawhdf5_agent::confidence::ConfidenceConfig; +use clawhdf5_agent::reranker::ReRankConfig; +use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry, SearchOptions, hybrid}; +use tempfile::TempDir; + +const DIM: usize = 32; +const N: usize = 3000; +const CLUSTERS: usize = 20; + +struct Rng(u64); +impl Rng { + fn next(&mut self) -> u64 { + self.0 = self.0.wrapping_add(0x9E37_79B9_7F4A_7C15); + let mut z = self.0; + z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9); + z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB); + z ^ (z >> 31) + } + fn unit(&mut self) -> f32 { + (self.next() >> 40) as f32 / (1u64 << 24) as f32 - 0.5 + } +} + +fn normalize(v: &mut [f32]) { + let n = v.iter().map(|x| x * x).sum::().sqrt(); + v.iter_mut().for_each(|x| *x /= n); +} + +struct Data { + vectors: Vec>, + cluster: Vec, + centres: Vec>, +} + +fn data() -> Data { + let mut rng = Rng(42); + let centres: Vec> = (0..CLUSTERS) + .map(|_| { + let mut c: Vec = (0..DIM).map(|_| rng.unit()).collect(); + normalize(&mut c); + c + }) + .collect(); + let mut vectors = Vec::new(); + let mut cluster = Vec::new(); + for i in 0..N { + let c = i % CLUSTERS; + let mut v: Vec = centres[c].iter().map(|x| x + rng.unit() * 0.3).collect(); + normalize(&mut v); + vectors.push(v); + cluster.push(c); + } + Data { + vectors, + cluster, + centres, + } +} + +/// Channel of record `i` for a filter keeping `percent`% of the store at +/// random (independent of the vectors). +fn random_channel(i: usize, rng_seed: u64, percent: u64) -> String { + let mut r = Rng(rng_seed ^ (i as u64 * 7919)); + if r.next() % 100 < percent { + "keep".into() + } else { + "other".into() + } +} + +fn build(data: &Data, channel: impl Fn(usize) -> String) -> (TempDir, HDF5Memory) { + let dir = TempDir::new().unwrap(); + let mut cfg = MemoryConfig::new(dir.path().join("s.h5"), "agent", DIM); + cfg.hebbian_boost = 0.0; // every query sees the same store + let mut m = HDF5Memory::create(cfg).unwrap(); + let entries = data + .vectors + .iter() + .enumerate() + .map(|(i, v)| MemoryEntry { + chunk: format!("record {i} cluster {}", data.cluster[i]), + embedding: v.clone(), + source_channel: channel(i), + timestamp: i as f64, + session_id: "s".into(), + tags: format!("t{i}"), + }) + .collect(); + m.save_batch(entries).unwrap(); + (dir, m) +} + +/// Exact top-k by cosine among the records `allowed` keeps. +fn exact_top(data: &Data, q: &[f32], k: usize, allowed: impl Fn(usize) -> bool) -> Vec { + let mut s: Vec<(usize, f32)> = (0..N) + .filter(|&i| allowed(i)) + .map(|i| (i, data.vectors[i].iter().zip(q).map(|(a, b)| a * b).sum())) + .collect(); + s.sort_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0))); + s.into_iter().take(k).map(|(i, _)| i).collect() +} + +fn query(data: &Data, i: usize) -> Vec { + let mut rng = Rng(1000 + i as u64); + let mut q: Vec = data.centres[i % CLUSTERS] + .iter() + .map(|x| x + rng.unit() * 0.3) + .collect(); + normalize(&mut q); + q +} + +fn vector_only(k: usize) -> SearchOptions { + SearchOptions::new(k).with_fusion(hybrid::Fusion::Weighted { + vector: 1.0, + keyword: 0.0, + }) +} + +#[test] +fn source_filter_returns_only_allowed_records_and_a_full_page() { + let d = data(); + // At N = 3000 and k = 10 the index serves a filter only when that is + // cheaper than scanning the allowed records: pool = 80 * N / allowed + // candidates at ~M = 16 distances each, against `allowed` distances. So + // 90% goes through the index, 50% and 1% to the exact scan. + for percent in [90, 50, 1] { + let (_dir, mut m) = build(&d, |i| random_channel(i, 5, percent)); + let allowed = |i: usize| random_channel(i, 5, percent) == "keep"; + let mut hits = 0; + for qi in 0..40 { + let q = query(&d, qi); + let got = m.search(&q, "", &vector_only(10).with_sources(["keep"])); + assert_eq!(got.len(), 10, "{percent}%: short page"); + assert!(got.iter().all(|r| r.source_channel == "keep")); + let want: HashSet = exact_top(&d, &q, 10, allowed).into_iter().collect(); + hits += got.iter().filter(|r| want.contains(&r.index)).count(); + } + let recall = hits as f64 / 400.0; + let floor = if percent == 90 { 0.95 } else { 1.0 }; + assert!(recall >= floor, "{percent}%: recall@10 {recall}"); + } +} + +#[test] +fn filter_away_from_the_query_falls_back_to_an_exact_scan() { + // Channel = cluster, and the filter keeps two clusters (10% of the + // store) that are not the query's: the index's neighbourhood of the + // query holds none of them. The search must still return the exact + // top 10 among the allowed records, not a short or empty page. + let d = data(); + let (_dir, mut m) = build(&d, |i| format!("c{}", d.cluster[i])); + for qi in 0..20 { + let q = query(&d, qi); + let a = format!("c{}", (qi + 7) % CLUSTERS); + let b = format!("c{}", (qi + 13) % CLUSTERS); + let got: Vec = m + .search( + &q, + "", + &vector_only(10).with_sources([a.clone(), b.clone()]), + ) + .iter() + .map(|r| r.index) + .collect(); + let want = exact_top(&d, &q, 10, |i| { + let c = format!("c{}", d.cluster[i]); + c == a || c == b + }); + assert_eq!(got, want, "query {qi}"); + } +} + +#[test] +fn filter_edge_cases() { + let d = data(); + let (_dir, mut m) = build(&d, |i| random_channel(i, 9, 50)); + let q = query(&d, 0); + assert!( + m.search( + &q, + "cluster", + &SearchOptions::new(10).with_sources(Vec::::new()) + ) + .is_empty() + ); + assert!( + m.search( + &q, + "cluster", + &SearchOptions::new(10).with_sources(["nope"]) + ) + .is_empty() + ); + // Keyword matches from other channels are filtered too. + let got = m.search( + &q, + "record cluster", + &SearchOptions::new(50).with_sources(["keep"]), + ); + assert_eq!(got.len(), 50); + assert!(got.iter().all(|r| r.source_channel == "keep")); + // Deleted records never come back, filtered or not. + let first = got[0].index; + m.delete(first).unwrap(); + let again = m.search( + &q, + "record cluster", + &SearchOptions::new(50).with_sources(["keep"]), + ); + assert!(again.iter().all(|r| r.index != first)); +} + +#[test] +fn plain_options_equal_hybrid_search_with() { + // Two identical stores, so neither query sees the other's boosts. + let d = data(); + let (_a, mut a) = build(&d, |i| random_channel(i, 3, 50)); + let (_b, mut b) = build(&d, |i| random_channel(i, 3, 50)); + for qi in 0..10 { + let q = query(&d, qi); + let x: Vec<(usize, u32)> = a + .search(&q, "record cluster 3", &SearchOptions::new(10)) + .iter() + .map(|r| (r.index, r.score.to_bits())) + .collect(); + let y: Vec<(usize, u32)> = b + .hybrid_search_with(&q, "record cluster 3", hybrid::DEFAULT_FUSION, 10) + .iter() + .map(|r| (r.index, r.score.to_bits())) + .collect(); + assert_eq!(x, y); + } +} + +fn small_store(entries: &[(&str, &str, f64)]) -> (TempDir, HDF5Memory) { + let dir = TempDir::new().unwrap(); + let mut m = HDF5Memory::create(MemoryConfig::new(dir.path().join("r.h5"), "a", 4)).unwrap(); + m.save_batch( + entries + .iter() + .map(|(chunk, channel, ts)| MemoryEntry { + chunk: chunk.to_string(), + embedding: vec![1.0, 0.0, 0.0, 0.0], + source_channel: channel.to_string(), + timestamp: *ts, + session_id: "s".into(), + tags: String::new(), + }) + .collect(), + ) + .unwrap(); + (dir, m) +} + +#[test] +fn rerank_breaks_relevance_ties_by_recency() { + // Identical text and vectors, so retrieval ties; re-ranking must put the + // newer record first and report the combined score. + let now = 1_000_000.0; + let (_d, mut m) = small_store(&[ + ("user prefers dark mode", "chat", now - 30.0 * 86_400.0), + ("user prefers dark mode", "chat", now - 60.0), + ]); + let q = [1.0, 0.0, 0.0, 0.0]; + let plain = m.search(&q, "dark mode", &SearchOptions::new(2)); + assert_eq!(plain[0].index, 0, "ties break by index without re-ranking"); + let reranked = m.search( + &q, + "dark mode", + &SearchOptions::new(2) + .with_rerank(ReRankConfig::default()) + .at_time(now), + ); + assert_eq!(reranked[0].index, 1); + assert!(reranked[0].score > reranked[1].score); + assert_ne!(reranked[0].score.to_bits(), plain[0].score.to_bits()); +} + +#[test] +fn confidence_rejects_when_nothing_is_good_enough() { + let (_d, mut m) = small_store(&[("alpha", "chat", 0.0), ("beta", "chat", 0.0)]); + let q = [1.0, 0.0, 0.0, 0.0]; + let strict = ConfidenceConfig { + min_score: 10.0, + ..ConfidenceConfig::default() + }; + assert!( + m.search(&q, "alpha", &SearchOptions::new(2).with_confidence(strict)) + .is_empty() + ); + let lenient = ConfidenceConfig { + min_score: 0.0, + min_gap: f32::INFINITY, + max_results: 1, + }; + assert_eq!( + m.search(&q, "alpha", &SearchOptions::new(2).with_confidence(lenient)) + .len(), + 1 + ); +} + +#[test] +fn only_returned_results_are_reinforced() { + // With re-ranking, a pool of max(3k, 10) candidates is retrieved; only + // the k returned should gain activation. + let d = data(); + let dir = TempDir::new().unwrap(); + let path = dir.path().join("h.h5"); + let mut m = HDF5Memory::create(MemoryConfig::new(path, "a", DIM)).unwrap(); + m.save_batch( + (0..200) + .map(|i| MemoryEntry { + chunk: format!("record {i}"), + embedding: d.vectors[i].clone(), + source_channel: "chat".into(), + timestamp: i as f64, + session_id: "s".into(), + tags: String::new(), + }) + .collect(), + ) + .unwrap(); + let q = query(&d, 0); + let got = m.search( + &q, + "record", + &SearchOptions::new(3).with_rerank(ReRankConfig::default()), + ); + assert_eq!(got.len(), 3); + let returned: HashSet = got.iter().map(|r| r.index).collect(); + // A second plain search reports each record's current activation. + let all = m.search(&q, "record", &SearchOptions::new(200)); + for r in &all { + let boosted = r.activation > 1.0; + assert_eq!(boosted, returned.contains(&r.index), "record {}", r.index); + } +} diff --git a/crates/clawhdf5-bench/src/bin/longmemeval_bench.rs b/crates/clawhdf5-bench/src/bin/longmemeval_bench.rs index 7d4858c..7943f53 100644 --- a/crates/clawhdf5-bench/src/bin/longmemeval_bench.rs +++ b/crates/clawhdf5-bench/src/bin/longmemeval_bench.rs @@ -64,6 +64,11 @@ use tempfile::TempDir; const EMBEDDING_DIM: usize = 384; +/// `--float16`: build every per-question store with `MemoryConfig::float16`, +/// so embeddings are rounded to half precision as they are saved — exactly +/// what such a store searches over. +static FLOAT16: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false); + /// A mode's fusion, as one short string for the reports. fn describe(mode: Mode) -> String { let fusion = match mode.fusion { @@ -431,6 +436,7 @@ fn evaluate_question( let mut config = MemoryConfig::new(dir.path().join("lme.h5"), "lme-bench", EMBEDDING_DIM); config.wal_enabled = false; config.compact_threshold = 0.0; + config.float16 = FLOAT16.load(std::sync::atomic::Ordering::Relaxed); let mut memory = HDF5Memory::create(config).expect("failed to create HDF5Memory"); memory.set_token_filter(mode.tokens); @@ -940,6 +946,10 @@ fn main() { limit = Some(v.parse().expect("--limit must be a positive integer")); } "--sweep" => sweep = true, + "--float16" => { + FLOAT16.store(true, std::sync::atomic::Ordering::Relaxed); + eprintln!("Stores use MemoryConfig::float16 (half-precision embeddings)"); + } "--rerank-sweep" => { // Re-ranking needs the vector stage to have candidates worth // reordering, so this is an embeddings-only comparison. @@ -971,6 +981,9 @@ fn main() { --rerank-sweep\n\ compare re-ranking off, metadata-only (the old\n\ behaviour) and blended at several half-lives.\n\ + --float16\n\ + build each store with MemoryConfig::float16, to\n\ + compare retrieval on half-precision embeddings.\n\ --sweep instead of the three named modes, sweep vector_weight\n\ from 0.0 to 1.0 in 0.1 steps. The 0.7/0.3 default was\n\ never searched; this is what searches it." diff --git a/crates/clawhdf5-bench/src/bin/search_harness.rs b/crates/clawhdf5-bench/src/bin/search_harness.rs index 1df37a9..920a44a 100644 --- a/crates/clawhdf5-bench/src/bin/search_harness.rs +++ b/crates/clawhdf5-bench/src/bin/search_harness.rs @@ -20,6 +20,7 @@ //! cargo run --release -p clawhdf5-bench --bin search_harness -- --json out.json //! cargo run --release -p clawhdf5-bench --bin search_harness -- --ann-only --uniform //! cargo run --release -p clawhdf5-bench --bin search_harness -- --float16-study --full +//! cargo run --release -p clawhdf5-bench --bin search_harness -- --options-study --full //! ``` use std::time::{Duration, Instant}; @@ -487,6 +488,177 @@ fn bench_end_to_end(n: usize, json: &mut Vec) { })); } +// --------------------------------------------------------------------------- +// Search options study: source filters, re-ranking, confidence rejection +// --------------------------------------------------------------------------- + +/// `--options-study`: what `HDF5Memory::search`'s options cost and whether a +/// filtered search finds the right records. Filters keep 50%, 10% or 1% of +/// the store at random, or two whole clusters away from the query (the case +/// the index cannot serve, which falls back to an exact scan). Recall is +/// vector-only against an exact scan of the allowed records; latency is full +/// hybrid search. Hebbian boosting is off. +fn options_study(n: usize) { + use clawhdf5_agent::SearchOptions; + use clawhdf5_agent::confidence::ConfidenceConfig; + use clawhdf5_agent::hybrid::Fusion; + use clawhdf5_agent::reranker::ReRankConfig; + + let data = make_dataset(n, 0x0B7 ^ n as u64); + let n_clusters = data.cluster_of.iter().max().map_or(1, |m| m + 1); + let mut rng = Rng(5); + let bucket_of: Vec = (0..n).map(|_| rng.below(100)).collect(); + let bucket = &bucket_of; + let query_texts: Vec = data + .query_cluster + .iter() + .enumerate() + .map(|(i, c)| text_for(*c, i, &mut rng)) + .collect(); + let exact_top = |q: &[f32], allowed: &dyn Fn(usize) -> bool| -> Vec { + let mut s: Vec<(usize, f32)> = (0..n) + .filter(|&i| allowed(i)) + .map(|i| (i, data.vectors[i].iter().zip(q).map(|(a, b)| a * b).sum())) + .collect(); + s.sort_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0))); + s.into_iter().take(K).map(|(i, _)| i).collect() + }; + + // Two stores: channel = random bucket, and channel = cluster. + let dir = tempfile::tempdir().unwrap(); + let mut stores = Vec::new(); + for by_cluster in [false, true] { + let mut rng = Rng(3); + let entries: Vec = data + .vectors + .iter() + .enumerate() + .map(|(i, v)| MemoryEntry { + chunk: text_for(data.cluster_of[i], i, &mut rng), + embedding: v.clone(), + source_channel: if by_cluster { + format!("c{}", data.cluster_of[i]) + } else { + format!("b{}", bucket[i]) + }, + timestamp: i as f64, + session_id: format!("s{}", i % 50), + tags: format!("t{i}"), + }) + .collect(); + let mut config = MemoryConfig::new( + dir.path().join(format!("opt_{by_cluster}.h5")), + "bench", + DIM, + ); + config.hebbian_boost = 0.0; + let mut mem = HDF5Memory::create(config).unwrap(); + mem.save_batch(entries).unwrap(); + std::hint::black_box(mem.search(&data.queries[0], "", &SearchOptions::new(K))); + stores.push(mem); + } + + let vector_only = SearchOptions::new(K).with_fusion(Fusion::Weighted { + vector: 1.0, + keyword: 0.0, + }); + // (label, store, channels for query i, allowed(i, record)) + type Case<'a> = ( + String, + usize, + Box Option> + 'a>, + Box bool + 'a>, + ); + let mut cases: Vec = vec![( + "no filter".into(), + 0, + Box::new(|_| None), + Box::new(|_, _| true), + )]; + for pct in [50usize, 10, 1] { + cases.push(( + format!("random {pct}%"), + 0, + Box::new(move |_| Some((0..pct).map(|b| format!("b{b}")).collect())), + Box::new(move |_, i| bucket[i] < pct), + )); + } + let d = &data; + let away = move |qi: usize| { + let qc = d.query_cluster[qi]; + [ + (qc + n_clusters / 3) % n_clusters, + (qc + 2 * n_clusters / 3) % n_clusters, + ] + }; + cases.push(( + "2 clusters away from the query".into(), + 1, + Box::new(move |qi| Some(away(qi).iter().map(|c| format!("c{c}")).collect())), + Box::new(move |qi, i| away(qi).contains(&d.cluster_of[i])), + )); + + for (label, store, channels, allowed) in &cases { + let mem = &mut stores[*store]; + let mut hits = 0; + let mut kept = 0; + for (qi, q) in data.queries.iter().enumerate() { + let mut opts = vector_only.clone(); + opts.source_channels = channels(qi); + let got = mem.search(q, "", &opts); + let want = exact_top(q, &|i| allowed(qi, i)); + kept += want.len(); + hits += got.iter().filter(|r| want.contains(&r.index)).count(); + } + let latency = summarize( + (0..N_QUERIES) + .map(|qi| { + let mut opts = SearchOptions::new(K); + opts.source_channels = channels(qi); + let t = Instant::now(); + std::hint::black_box(mem.search(&data.queries[qi], &query_texts[qi], &opts)); + t.elapsed() + }) + .collect(), + ); + println!( + "| {n} | {label} | {:.4} | {:.3} | {:.3} |", + hits as f64 / kept.max(1) as f64, + millis(latency.p50), + millis(latency.p99), + ); + } + + let mem = &mut stores[0]; + for (label, opts) in [ + ( + "re-rank", + SearchOptions::new(K).with_rerank(ReRankConfig::default()), + ), + ( + "re-rank + confidence", + SearchOptions::new(K) + .with_rerank(ReRankConfig::default()) + .with_confidence(ConfidenceConfig::default()), + ), + ] { + let latency = summarize( + (0..N_QUERIES) + .map(|qi| { + let t = Instant::now(); + std::hint::black_box(mem.search(&data.queries[qi], &query_texts[qi], &opts)); + t.elapsed() + }) + .collect(), + ); + println!( + "| {n} | {label} | — | {:.3} | {:.3} |", + millis(latency.p50), + millis(latency.p99) + ); + } +} + // --------------------------------------------------------------------------- // float16 study: what does half-precision embedding storage cost? // --------------------------------------------------------------------------- @@ -788,6 +960,19 @@ fn main() { } return; } + if args.iter().any(|a| a == "--options-study") { + println!("## Search options ({DIM}-dim, k = {K}, Hebbian boost off)\n"); + println!("| N | options | filtered recall@10 | p50 ms | p99 ms |"); + println!("|---:|---|---:|---:|---:|"); + for &n in if full { + &[10_000, 100_000][..] + } else { + &[10_000][..] + } { + options_study(n); + } + return; + } if args.iter().any(|a| a == "--f16-first") { F16_FIRST.store(true, std::sync::atomic::Ordering::Relaxed); } diff --git a/crates/clawhdf5-cli/src/main.rs b/crates/clawhdf5-cli/src/main.rs index c3f15e3..2fc07c0 100644 --- a/crates/clawhdf5-cli/src/main.rs +++ b/crates/clawhdf5-cli/src/main.rs @@ -36,10 +36,13 @@ enum Commands { /// Accepted for compatibility; int8 is now the default #[arg(long, hide = true, conflicts_with = "f32_index")] quantized_index: bool, - /// Store embeddings on disk as IEEE half precision (float16): half - /// the bytes, about three significant digits; values must lie within - /// ±65504 + /// Store embeddings as full-precision f32 instead of the default + /// half precision (float16: half the bytes, about three significant + /// digits, values within ±65504) #[arg(long)] + f32: bool, + /// Accepted for compatibility; float16 is now the default + #[arg(long, hide = true, conflicts_with = "f32")] float16: bool, }, /// Save a memory entry (reads JSON from stdin or --json) @@ -107,11 +110,16 @@ fn run(cli: Cli) -> Result<(), Box> { wal, f32_index, quantized_index: _, - float16, + f32, + float16: _, } => { let mut config = MemoryConfig::new(cli.path.clone(), &agent_id, dim); config.wal_enabled = wal; - config.float16 = float16; + // As with --f32-index: only ever switch the library default off. + if f32 { + config.float16 = false; + } + let config_float16 = config.float16; // Only ever switch *off* the library default: assigning the flag // outright would force every CLI-created store back to f32 unless // the caller knew to ask for int8. @@ -127,7 +135,7 @@ fn run(cli: Cli) -> Result<(), Box> { "embedding_dim": dim, "wal_enabled": wal, "quantized_index": config_quantized, - "float16": float16, + "float16": config_float16, "count": mem.count(), }); println!("{}", serde_json::to_string_pretty(&j)?);