Search options (source filters, re-ranking, confidence); float16 default #5

Open
osobh wants to merge 3 commits from feat/search-options into main
13 changed files with 1100 additions and 136 deletions
+67
View File
@@ -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
+42 -7
View File
@@ -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
+14 -4
View File
@@ -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
`<store>.h5.wal.corrupt-<ts>` 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
+47 -13
View File
@@ -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`
+23 -20
View File
@@ -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
+9 -1
View File
@@ -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,
+16 -61
View File
@@ -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<MemorySearchResult> {
// 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<RerankInput> = 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<ScoredResult> = 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<usize, &crate::SearchResult> =
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()
}
+274 -24
View File
@@ -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<Vec<String>>,
/// Re-rank a candidate pool by retrieval relevance, recency, source
/// authority and activation — the pipeline the OpenClaw backend runs.
pub rerank: Option<ReRankConfig>,
/// 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<ConfidenceConfig>,
/// The time recency is measured from, in seconds since the epoch.
/// `None` uses the system clock.
pub now: Option<f64>,
}
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<S: Into<String>>(mut self, channels: impl IntoIterator<Item = S>) -> 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<u8> {
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<SearchResult> {
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<SearchResult> {
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<SearchResult> = 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<ScoredResult> = results
.iter()
.map(|r| ScoredResult {
index: r.index,
score: r.score,
})
.collect();
let keep: HashSet<usize> = 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<SearchResult>,
config: &ReRankConfig,
now: Option<f64>,
) -> Vec<SearchResult> {
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<RerankInput> = 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<usize, SearchResult> =
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
@@ -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<f32> = (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());
}
@@ -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::<f32>().sqrt();
v.iter_mut().for_each(|x| *x /= n);
}
struct Data {
vectors: Vec<Vec<f32>>,
cluster: Vec<usize>,
centres: Vec<Vec<f32>>,
}
fn data() -> Data {
let mut rng = Rng(42);
let centres: Vec<Vec<f32>> = (0..CLUSTERS)
.map(|_| {
let mut c: Vec<f32> = (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<f32> = 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<usize> {
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<f32> {
let mut rng = Rng(1000 + i as u64);
let mut q: Vec<f32> = 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<usize> = 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<usize> = 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::<String>::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<usize> = 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);
}
}
@@ -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."
@@ -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<serde_json::Value>) {
}));
}
// ---------------------------------------------------------------------------
// 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<usize> = (0..n).map(|_| rng.below(100)).collect();
let bucket = &bucket_of;
let query_texts: Vec<String> = 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<usize> {
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<MemoryEntry> = 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<dyn Fn(usize) -> Option<Vec<String>> + 'a>,
Box<dyn Fn(usize, usize) -> bool + 'a>,
);
let mut cases: Vec<Case> = 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);
}
+14 -6
View File
@@ -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<dyn std::error::Error>> {
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<dyn std::error::Error>> {
"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)?);