Compare commits
6
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
dce5559ff2 | ||
|
|
1b3bbb054a | ||
|
|
a8fb758489 | ||
|
|
5c8323cb1e | ||
|
|
dbaf3f505d | ||
|
|
c470244a6f |
+681
-139
File diff suppressed because it is too large
Load Diff
+93
-7
@@ -3,6 +3,17 @@
|
||||
## Unreleased
|
||||
|
||||
### Upgrade Notes
|
||||
- **`clawhdf5-migrate` now writes a real agent store.** Its output used to be
|
||||
a layout of its own (`/chunks`, `/sessions`, `/entities`, `/relations`, no
|
||||
`/meta`) that `HDF5Memory::open` rejected, so a migrated file could not be
|
||||
used as agent memory. Files it wrote before this release are not agent
|
||||
stores; re-run the migration. Also: embeddings default to `float16` like
|
||||
any new store (`--f32` opts out; `--float16` is a hidden no-op); a row with
|
||||
the wrong embedding length is an error instead of being truncated or
|
||||
padded; `--incremental` now matches rows by content against an existing
|
||||
store and follows the source's deleted flags; a source with no memory rows
|
||||
needs `--embedding-dim`. The per-dataset SHA-256 provenance attributes of
|
||||
the old layout are gone (the agent schema has no place for them).
|
||||
- **Files written by clawhdf5 now open in h5py and libhdf5.** Every `f32`
|
||||
dataset we wrote — including every agent store's embeddings — was refused
|
||||
with "sign bit position out of bounds", and every empty dataset with
|
||||
@@ -11,12 +22,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 +57,60 @@
|
||||
`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.
|
||||
|
||||
### Migration
|
||||
- `clawhdf5-migrate`: writes through the agent's own API (`HDF5Memory::create`
|
||||
/ `open`, `save_batch`, the session cache and knowledge graph), so there is
|
||||
no second copy of the schema. Sessions and entities/relations carry over;
|
||||
deleted rows become deleted records (or are left out with
|
||||
`--skip-deleted`). Every source row is checked before the output is created,
|
||||
so a source that cannot be migrated leaves an existing store untouched.
|
||||
Validation reads the result back with `HDF5Memory::open_read_only`, compares
|
||||
every field (embeddings bit for bit — `round_to_f16` of the source for a
|
||||
`float16` store) and checks that a migrated record is found by search. The
|
||||
`half`-based conversion is gone; `clawhdf5_format::float16` is the only one.
|
||||
42 tests, including h5py opening a migrated store; an adversarial review's
|
||||
two blocker and four major findings are fixed with regression tests.
|
||||
- `clawhdf5-agent`: `HDF5Memory::sessions()` / `sessions_mut()`,
|
||||
`HDF5Memory::delete_batch(&[usize])` (one save, all-or-nothing, never
|
||||
auto-compacts), `SessionCache::add_at`, and `SessionCache` / `SessionEntry`
|
||||
re-exported from the crate root.
|
||||
|
||||
### 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.
|
||||
|
||||
### Benchmarks
|
||||
- Every undated or pre-September section of `BENCHMARKS.md` re-run on one
|
||||
machine on one day (tank, 2026-09-24, commit 5c8323c), with the command for
|
||||
each and every number traced back to the raw output by a separate check.
|
||||
Where a figure moved, the section says so. Two apparent regressions were
|
||||
isolated rather than published: knowledge-graph traversal (a real bug,
|
||||
fixed above) and the write path, which measures the same at v2.3.0 on this
|
||||
machine — the old 18 µs / 6.17 ms figures came from an undated run on other
|
||||
hardware; `float16` adds ~2 µs per save and the int8 index nothing.
|
||||
- New `multimodal_bench`: cross-modal search at 1K and 10K records, which the
|
||||
README claimed but nothing measured.
|
||||
- `footprint_bench` reports whether it built `float16` or `f32` stores and
|
||||
takes `--f32`; it had kept printing "f32" after the default changed.
|
||||
|
||||
### 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 +138,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 +172,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
|
||||
@@ -111,6 +189,14 @@
|
||||
knew to ask; it now only ever switches the default off.
|
||||
|
||||
### Performance
|
||||
- `clawhdf5-agent`: **knowledge-graph traversal was 6.5x slower than it
|
||||
should be.** `bfs_neighbors` and `spreading_activation` built an adjacency
|
||||
index over the whole graph on every call (1efd82c), so a 2-hop BFS over 1K
|
||||
entities took 155 µs. The index is now cached on `KnowledgeCache` and
|
||||
checked against a fingerprint of the graph on each use — one pass over
|
||||
entity ids and relation endpoints, no allocation — so any change, including
|
||||
direct edits of its public `Vec`s, still rebuilds it (tested). BFS over 1K
|
||||
entities: 155.1 -> 23.1 µs; spreading activation over 100: 22.8 -> 10.1 µs.
|
||||
- `clawhdf5-format`, `clawhdf5-filters`: both deflate paths hand the codec the
|
||||
whole chunk in one call, into a buffer allocated once, instead of streaming
|
||||
it through a 32 KiB buffer: about 5% on chunked writes and 10% on zlib-ng's
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -6,7 +6,7 @@
|
||||
[](https://www.rust-lang.org)
|
||||
[](#building)
|
||||
[](BENCHMARKS.md#longmemeval-results)
|
||||
[](BENCHMARKS.md#memory-footprint-1)
|
||||
[](BENCHMARKS.md#memory-footprint-1)
|
||||
|
||||
ClawHDF5 is a pure-Rust HDF5 implementation combined with a research-grade agent memory engine. It gives AI agents persistent, searchable, integrity-checked memory — all stored in a single portable 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
|
||||
@@ -107,8 +111,8 @@ Every AI agent needs memory. Today that means scattered Markdown files, SQLite d
|
||||
| Keyword search | Separate FTS engine | Integrated BM25 |
|
||||
| Knowledge graph | Neo4j or none | In-file graph with spreading activation |
|
||||
| Memory consolidation | Manual pruning | Hippocampal-inspired automatic tiers |
|
||||
| Temporal queries | Custom code | Native temporal index (716ns) |
|
||||
| Multi-modal | Multiple stores | Unified cross-modal search |
|
||||
| Temporal queries | Custom code | Native temporal index (622 ns range query over 10K) |
|
||||
| Multi-modal | Multiple stores | Unified cross-modal search (exact scan: 842 µs over 1K records) |
|
||||
| Integrity | Hope for the best | Chained-CRC WAL, checksummed chunk indexes, write-anomaly alerts, opt-in SHA-256 dataset provenance |
|
||||
| Portability | Config + DB + files | **One `.h5` file. Copy it anywhere.** |
|
||||
|
||||
@@ -116,7 +120,7 @@ Every AI agent needs memory. Today that means scattered Markdown files, SQLite d
|
||||
|
||||
## Performance
|
||||
|
||||
Vector search and agent-memory operations below are benchmarked on Intel i7-12650H (10C/16T), 384-dim embeddings, Criterion.rs. The HDF5 Core I/O table immediately below is from a separate, independently reproduced run (see its own hardware note).
|
||||
The brute-force/IVF vector search, agent-memory, on-disk footprint and consolidation figures below were measured 2026-09-24 on tank (AMD Ryzen 7 7800X3D, 8C/16T), commit 5c8323c, 384-dim embeddings; the commands are in [BENCHMARKS.md](BENCHMARKS.md). Exceptions are marked where they appear: the HDF5 Core I/O table immediately below is from a separate, independently reproduced run (see its own hardware note), and the HNSW `f32`/`i8` table and the in-memory `i8` column were not re-measured on 2026-09-24.
|
||||
|
||||
### HDF5 Core I/O (vs libhdf5 1.14.6)
|
||||
|
||||
@@ -150,33 +154,45 @@ and [§ Quantising the index copy](BENCHMARKS.md#quantising-the-index-copy-quant
|
||||
| `f32` index | 0.9945 | 13 399 | 3.2 s |
|
||||
| `i8` index + exact re-score (**default for new stores**) | 0.9940 | **21 848** | **1.8 s** |
|
||||
|
||||
Before the v2.4.0 neighbour-selection fix, recall@10 at 100K was 0.31.
|
||||
Before the v2.4.0 neighbour-selection fix, recall@10 at 100K was 0.31. These
|
||||
two rows are a paired comparison (medians of alternating runs, same binary).
|
||||
A single `f32` run on 2026-09-24 measured recall 0.9945, 19 001 QPS and a
|
||||
2.7 s build; the int8 row was not re-run, so the pair has not been re-checked
|
||||
([§ Quantising the index copy](BENCHMARKS.md#quantising-the-index-copy-quantized_index)).
|
||||
|
||||
**Brute-force and IVF paths** (Criterion, i7-12650H):
|
||||
**Brute-force and IVF paths** (Criterion, tank, 2026-09-24):
|
||||
|
||||
| Scale | Flat | IVF (nprobe=10) | IVF-PQ | vs MemX¹ |
|
||||
| Scale | Flat | IVF (nprobe=10) | IVF-PQ | MemX¹ (claimed, end-to-end) |
|
||||
|-------|------|-----------------|--------|----------|
|
||||
| 1K | **54 µs** | — | — | — |
|
||||
| 10K | 753 µs | **27 µs** | — | — |
|
||||
| 100K | 11.4 ms | 1.32 ms | **1.19 ms** | ~8–76× (see caveat) |
|
||||
| 1K | **47.4 µs** | — | — | — |
|
||||
| 10K | 500.5 µs | **24.8 µs** | — | — |
|
||||
| 100K | 6.58 ms | 592 µs | **869 µs** | <90 ms |
|
||||
|
||||
> Reproduced on the same second machine (Ryzen 7 7800X3D) with a corrected,
|
||||
> apples-to-apples SIMD/scalar/parallel comparison methodology — see
|
||||
> [BENCHMARKS.md § Independent Validation: tank — LongMemEval & Vector
|
||||
> Search](BENCHMARKS.md#independent-validation-tank--longmemeval--vector-search-ryzen-7-7800x3d-2026-08-05).
|
||||
> These replace figures from the original i7-12650H run (flat 54 µs / 753 µs /
|
||||
> 11.4 ms); a 2026-08-05 run on tank had already matched the new ones — see
|
||||
> [BENCHMARKS.md § Vector Search Latency](BENCHMARKS.md#vector-search-latency).
|
||||
|
||||
### Agent Memory Operations
|
||||
|
||||
| Operation | Latency | Scale |
|
||||
|-----------|---------|-------|
|
||||
| Hybrid search (`HDF5Memory::hybrid_search`, p50) | **70 µs** / 0.49 ms / 4.65 ms | 1K / 10K / 100K records |
|
||||
| BM25 keyword search | **67 µs** | 1K records |
|
||||
| Knowledge graph BFS | **24 µs** | 1K entities |
|
||||
| Spreading activation | **17 µs** | 100 entities |
|
||||
| Temporal range query | **716 ns** | 10K timestamps |
|
||||
| Consolidation cycle | **164 µs** | 1K records |
|
||||
| Memory write (WAL) | **18 µs** | per record (group-commit append; HDF5 batched at flush) |
|
||||
| Importance gate | **61 ns** | per record |
|
||||
| Hybrid search (`HDF5Memory::hybrid_search`, p50) | **0.07 ms** / 0.49 ms / 4.69 ms | 1K / 10K / 100K records |
|
||||
| BM25 keyword search | **20.4 µs** | 1K records |
|
||||
| Knowledge graph BFS | **23.1 µs** | 1K entities |
|
||||
| Spreading activation | **10.1 µs** | 100 entities |
|
||||
| Temporal range query | **622 ns** | 10K timestamps |
|
||||
| Consolidation cycle | **115.2 µs** | 1K records |
|
||||
| Cross-modal search (exact scan, 2 embeddings per record) | **842.0 µs** / 8.44 ms | 1K / 10K records |
|
||||
| Memory write (WAL) | **26.1 µs** | per record (group-commit append; HDF5 batched at flush) |
|
||||
| Importance gate | **57.6 ns** | per record (trivial skip) |
|
||||
|
||||
The old 18 µs WAL write was undated, from another machine: v2.3.0 measures
|
||||
24.3 µs on the same hardware as this table, the same as an `f32` store today.
|
||||
`float16` stores (the new default) add ~2 µs for rounding; the int8 index adds
|
||||
nothing. See [BENCHMARKS.md § Write Path](BENCHMARKS.md#write-path).
|
||||
Knowledge-graph traversal was briefly 6.5x slower (155 µs) until this re-run
|
||||
found and fixed an adjacency index rebuilt on every traversal; see
|
||||
[§ Knowledge Graph](BENCHMARKS.md#knowledge-graph).
|
||||
|
||||
### Chunked Write Throughput (codec comparison)
|
||||
|
||||
@@ -191,7 +207,7 @@ by default (AoS→SoA byte transpose, +157–204% throughput for float data):
|
||||
|
||||
Use `.with_zstd(3)` or `.with_deflate(6)` for write-heavy workloads — both now perform at ~720–750 MiB/s on large matrices. Use `.with_pcodec()` for write-once/read-many workloads where compression ratio matters more than encode speed. Disable auto-shuffle with `.without_shuffle()` for byte arrays that don't benefit from AoS→SoA transposition.
|
||||
|
||||
> ¹ MemX ([arxiv:2603.16171](https://arxiv.org/abs/2603.16171), March 2026): Rust + libSQL, claims <90ms at 100K records. **Not like-for-like:** MemX's figure is *end-to-end* (embeddings + FTS5 + four-factor re-ranking); ours is a *single component* (raw vector search). The ratio overstates the real advantage by an unquantified margin — order-of-magnitude indication only. See [BENCHMARKS.md](BENCHMARKS.md#comparison-to-memx-arxiv260316171).
|
||||
> ¹ MemX ([arxiv:2603.16171](https://arxiv.org/abs/2603.16171), March 2026): Rust + libSQL, claims <90ms at 100K records. **Not like-for-like:** MemX's figure is *end-to-end* (embeddings + FTS5 + four-factor re-ranking); ours is a *single component* (raw vector search), so the two columns are not comparable and no ratio is given. See [BENCHMARKS.md](BENCHMARKS.md#comparison-to-memx-arxiv260316171).
|
||||
|
||||
### LongMemEval Retrieval Recall
|
||||
|
||||
@@ -244,17 +260,26 @@ retrieval recall reported as QA accuracy typically overstates by 20–30 points.
|
||||
|
||||
### Memory Footprint
|
||||
|
||||
**On disk** — 384-dim embeddings, 200-char text
|
||||
**On disk** — 384-dim `float16` embeddings (the default for new stores),
|
||||
200-char text, `footprint_bench`
|
||||
([BENCHMARKS.md § Memory Footprint](BENCHMARKS.md#memory-footprint-1)):
|
||||
|
||||
| Records | File Size | Bytes/Record | Gzip-6 compressed |
|
||||
|---------|-----------|--------------|-------------------|
|
||||
| 1K | 1.7 MB | 1.8 KB | 277 KB (6.1x) |
|
||||
| 10K | 17.0 MB | 1.7 KB | 2.7 MB (6.2x) |
|
||||
| 100K | 169.8 MB | 1.7 KB | 26.9 MB (6.2x) |
|
||||
| 1K | 810.4 KB | 829 B | 56.4 KB |
|
||||
| 10K | 7.8 MB | 820 B | 471.3 KB |
|
||||
| 100K | 76.7 MB | 803 B | 4.5 MB |
|
||||
|
||||
With `MemoryConfig::float16` the embeddings take half the space: an agent
|
||||
store of 100K × 384 records is 80.8 MiB instead of 154.0.
|
||||
The benchmark's synthetic embeddings and text are far more repetitive than
|
||||
real data (only 40 distinct texts), so no column here is an expectation for
|
||||
real data. The compressed column is an upper bound, and the Bytes/Record
|
||||
column is optimistic too: it is not an uncompressed figure, because the store
|
||||
always deflates its text (any string dataset of 4 KiB or more) whatever
|
||||
`MemoryConfig::compression` says. The `float16` embeddings alone are 768 B per
|
||||
record, so 200 characters of real text would take a record above 820 B.
|
||||
This table used to show `f32` stores (1.7 KB per record, 169.8 MB at 100K);
|
||||
those were not re-measured. The float16 study compares the two on the same
|
||||
data: 100K × 384 records take 80.8 MiB as `float16` and 154.0 MiB as `f32`.
|
||||
|
||||
**In memory** — a store reopened from disk, 384-dim `f32`, measured with a
|
||||
counting allocator ([BENCHMARKS.md § Memory footprint](BENCHMARKS.md#memory-footprint)):
|
||||
@@ -266,7 +291,8 @@ counting allocator ([BENCHMARKS.md § Memory footprint](BENCHMARKS.md#memory-foo
|
||||
| 100K | 146 MiB | 399 MiB (2.72x) | **256 MiB (1.74x)** |
|
||||
|
||||
Down from 505 MiB (3.44x) at 100K before v2.6.0, when the cache held every
|
||||
embedding twice.
|
||||
embedding twice. The `f32` column was re-measured on 2026-09-24 and reproduced
|
||||
exactly; the `i8` column was not re-run.
|
||||
|
||||
### Consolidation Efficiency
|
||||
|
||||
@@ -277,7 +303,10 @@ embedding twice.
|
||||
|--------|--------|-------|-------|
|
||||
| Records in store | 1,000 | 100 | −90% |
|
||||
| Hit@1 recall (signal records) | 100% | 100% | no loss |
|
||||
| Search latency | 2.75 ms | 0.31 ms | **8.8x faster** |
|
||||
| Search latency (avg) | 2.22 ms | 0.24 ms | **9.3x faster** |
|
||||
|
||||
The consolidation cycle that does this took 0.13 ms; at 10K records a cycle
|
||||
takes 2.16 ms.
|
||||
|
||||
**Full benchmark details: [BENCHMARKS.md](BENCHMARKS.md)**
|
||||
|
||||
@@ -293,12 +322,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 +365,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 +432,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 +636,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`
|
||||
@@ -720,9 +778,37 @@ Replace in `Cargo.toml` and source:
|
||||
|
||||
```bash
|
||||
cargo install --path crates/clawhdf5-migrate
|
||||
clawhdf5-migrate --sqlite old.db --hdf5 memory.h5 --agent-id my-agent --embedding-dim 384
|
||||
clawhdf5-migrate --sqlite old.db --hdf5 memory.h5 --agent-id my-agent --embedder minilm
|
||||
```
|
||||
|
||||
The output is an ordinary `clawhdf5-agent` store, written through the agent's
|
||||
own API: open it with `HDF5Memory::open` (or `clawhdf5-cli --path memory.h5 …`)
|
||||
and search it straight away. What carries over from the ZeroClaw tables:
|
||||
|
||||
| SQLite | Agent store |
|
||||
|--------|-------------|
|
||||
| `memory_chunks` | memory records (text, embedding, source channel, timestamp, session id, tags); rows with `deleted = 1` become deleted records, or are left out with `--skip-deleted` |
|
||||
| `sessions` | sessions (id, start/end index, channel, summary, timestamp) |
|
||||
| `entities`, `relations` | knowledge graph entities and relations; entities get new ids and relations are re-pointed at them |
|
||||
|
||||
The chunk `id` column has no counterpart in the agent store, so records are
|
||||
written in `id` order and numbered from 0. Embeddings are stored as float16
|
||||
like any new store; `--f32` keeps full precision (and is required for values
|
||||
beyond ±65504). The embedding dimension is detected from the first row unless
|
||||
`--embedding-dim` is given, and every row must have it: a row of another length
|
||||
is an error, never truncated or padded. A source with no memory records (only
|
||||
sessions or the graph) needs `--embedding-dim`, since a store's dimension is
|
||||
fixed when it is created. Every row is checked before the output is created,
|
||||
so a source that cannot be migrated leaves an existing store at `--hdf5` as it
|
||||
was. `--incremental` adds to an existing store only the rows it does not
|
||||
already hold; the source must have the store's dimension, and records already
|
||||
in the store take the source's deleted flag (a row deleted in SQLite since the
|
||||
last run is deleted in the store; one un-deleted there is written again, as
|
||||
the agent has no un-delete). The tool reads the result back with
|
||||
`HDF5Memory::open_read_only`, compares it with the source (every row with
|
||||
`--validate-full`) and checks that a migrated record is found by search;
|
||||
`--dry-run` only counts the rows.
|
||||
|
||||
---
|
||||
|
||||
## Roadmap
|
||||
|
||||
@@ -45,6 +45,10 @@ harness = false
|
||||
name = "memory_bench"
|
||||
harness = false
|
||||
|
||||
[[bench]]
|
||||
name = "multimodal_bench"
|
||||
harness = false
|
||||
|
||||
[features]
|
||||
default = ["float16", "hnsw", "parallel"]
|
||||
float16 = ["half"]
|
||||
|
||||
@@ -0,0 +1,107 @@
|
||||
//! Multi-modal memory search benchmarks (`clawhdf5_agent::multimodal`).
|
||||
//!
|
||||
//! Covers `MultiModalStore::search_cross_modal` (every embedding of every
|
||||
//! record, whatever its modality) and, for comparison,
|
||||
//! `MultiModalStore::search_by_modality` restricted to one modality.
|
||||
//!
|
||||
//! Corpus: N records (1K and 10K), each carrying two 384-dim embeddings —
|
||||
//! a text embedding of its caption plus one embedding of its primary modality,
|
||||
//! cycling Image / Audio / Video — so a cross-modal query scores 2N vectors.
|
||||
//! All data comes from a fixed-seed LCG, so every run sees the same corpus.
|
||||
//!
|
||||
//! Run: `cargo bench -p clawhdf5-agent --bench multimodal_bench`
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use clawhdf5_agent::multimodal::{
|
||||
MediaRef, ModalEmbedding, Modality, MultiModalRecord, MultiModalStore,
|
||||
};
|
||||
use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Simple deterministic PRNG (LCG), same as the other agent benches
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
struct Rng(u32);
|
||||
|
||||
impl Rng {
|
||||
fn new(seed: u32) -> Self {
|
||||
Self(seed)
|
||||
}
|
||||
fn next_u32(&mut self) -> u32 {
|
||||
self.0 = self.0.wrapping_mul(1103515245).wrapping_add(12345);
|
||||
self.0 >> 16
|
||||
}
|
||||
fn next_f32(&mut self) -> f32 {
|
||||
self.next_u32() as f32 / 65536.0 - 0.5
|
||||
}
|
||||
}
|
||||
|
||||
fn make_vec(rng: &mut Rng, dim: usize) -> Vec<f32> {
|
||||
(0..dim).map(|_| rng.next_f32()).collect()
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Corpus
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
const DIM: usize = 384;
|
||||
const K: usize = 10;
|
||||
|
||||
const MEDIA: [(Modality, &str, &str); 3] = [
|
||||
(Modality::Image, "image/png", "clip-vit-base"),
|
||||
(Modality::Audio, "audio/wav", "clap-base"),
|
||||
(Modality::Video, "video/mp4", "xclip-base"),
|
||||
];
|
||||
|
||||
fn build_store(n: usize, seed: u32) -> MultiModalStore {
|
||||
let mut rng = Rng::new(seed);
|
||||
let mut store = MultiModalStore::new();
|
||||
for i in 0..n {
|
||||
let (modality, mime, model) = &MEDIA[i % MEDIA.len()];
|
||||
let embeddings = vec![
|
||||
ModalEmbedding::new(Modality::Text, make_vec(&mut rng, DIM), "minilm-l6"),
|
||||
ModalEmbedding::new(modality.clone(), make_vec(&mut rng, DIM), *model),
|
||||
];
|
||||
store.add_record(MultiModalRecord {
|
||||
id: 0,
|
||||
primary_modality: modality.clone(),
|
||||
text_content: Some(format!("{modality} memory {i}")),
|
||||
media_ref: Some(MediaRef::path(format!("/media/{i}"), *mime)),
|
||||
embeddings,
|
||||
observation: None,
|
||||
timestamp: 1_700_000_000.0 + i as f64,
|
||||
metadata: HashMap::new(),
|
||||
});
|
||||
}
|
||||
store
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Benchmarks
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
fn multimodal_search_benches(c: &mut Criterion) {
|
||||
let query = make_vec(&mut Rng::new(99), DIM);
|
||||
|
||||
let mut group = c.benchmark_group("multimodal_search");
|
||||
group.sample_size(50);
|
||||
|
||||
for (label, n) in [("1k", 1_000usize), ("10k", 10_000)] {
|
||||
let store = build_store(n, 42);
|
||||
assert_eq!(store.count(), n);
|
||||
|
||||
group.bench_with_input(BenchmarkId::new("cross_modal", label), &n, |b, _| {
|
||||
b.iter(|| store.search_cross_modal(&query, K));
|
||||
});
|
||||
|
||||
group.bench_with_input(BenchmarkId::new("by_modality_image", label), &n, |b, _| {
|
||||
b.iter(|| store.search_by_modality(&Modality::Image, &query, K));
|
||||
});
|
||||
}
|
||||
|
||||
group.finish();
|
||||
}
|
||||
|
||||
criterion_group!(multimodal_benches, multimodal_search_benches);
|
||||
criterion_main!(multimodal_benches);
|
||||
@@ -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
|
||||
|
||||
@@ -163,12 +163,13 @@ fn levenshtein(a: &str, b: &str) -> usize {
|
||||
/// entities-slice-index map, and an entity-id -> relation-indices map (edges
|
||||
/// touching that entity as either source or target).
|
||||
///
|
||||
/// Built fresh per traversal call rather than cached on `KnowledgeCache`:
|
||||
/// entities/relations are plain `pub` `Vec`s that get pushed to directly
|
||||
/// (e.g. `schema.rs`'s load path bypasses `add_entity`/`add_relation`), so a
|
||||
/// persistent index would need extra bookkeeping to avoid drifting stale. A
|
||||
/// one-off O(V+E) build per call is still a large win over the O(V·E) (BFS)
|
||||
/// / O(steps·active·E) (spreading activation) scans it replaces.
|
||||
/// Cached on `KnowledgeCache` and checked against a fingerprint of the graph
|
||||
/// on every use ([`graph_fingerprint`]). entities/relations are plain `pub`
|
||||
/// `Vec`s that get changed directly (e.g. `schema.rs`'s load path bypasses
|
||||
/// `add_entity`/`add_relation`), so the cache cannot rely on being told about
|
||||
/// changes; the fingerprint notices any of them. Rebuilding it on every
|
||||
/// traversal instead made a 2-hop BFS over 1K entities 6.5x slower than the
|
||||
/// scan it replaced (24 -> 155 µs; `BENCHMARKS.md`, "Knowledge Graph").
|
||||
struct AdjacencyIndex {
|
||||
entity_index: HashMap<u64, usize>,
|
||||
by_entity: HashMap<u64, Vec<usize>>,
|
||||
@@ -204,6 +205,45 @@ impl AdjacencyIndex {
|
||||
}
|
||||
}
|
||||
|
||||
/// A hash of everything [`AdjacencyIndex`] depends on — each entity's id and
|
||||
/// position, each relation's endpoints and position. One linear pass, no
|
||||
/// allocation: far cheaper than building the index, which hashes the same
|
||||
/// values into two maps.
|
||||
fn graph_fingerprint(entities: &[Entity], relations: &[Relation]) -> u64 {
|
||||
// splitmix64-style mixing; order matters, so positions are covered.
|
||||
fn mix(h: u64, v: u64) -> u64 {
|
||||
let mut z = (h ^ v).wrapping_add(0x9E37_79B9_7F4A_7C15);
|
||||
z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
|
||||
z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
|
||||
z ^ (z >> 31)
|
||||
}
|
||||
let mut h = mix(entities.len() as u64, relations.len() as u64);
|
||||
for e in entities {
|
||||
h = mix(h, e.id);
|
||||
}
|
||||
for r in relations {
|
||||
h = mix(mix(h, r.src), r.tgt);
|
||||
}
|
||||
h
|
||||
}
|
||||
|
||||
/// The cached [`AdjacencyIndex`] and the fingerprint it was built for.
|
||||
/// Cloning a `KnowledgeCache` starts the clone with an empty cache.
|
||||
#[derive(Default)]
|
||||
struct AdjacencyCache(std::sync::Mutex<Option<(u64, std::sync::Arc<AdjacencyIndex>)>>);
|
||||
|
||||
impl Clone for AdjacencyCache {
|
||||
fn clone(&self) -> Self {
|
||||
Self::default()
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for AdjacencyCache {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.write_str("AdjacencyCache")
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// KnowledgeCache
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -216,6 +256,7 @@ pub struct KnowledgeCache {
|
||||
pub alias_strings: Vec<String>,
|
||||
pub alias_entity_ids: Vec<i64>,
|
||||
next_entity_id: u64,
|
||||
adjacency: AdjacencyCache,
|
||||
}
|
||||
|
||||
impl KnowledgeCache {
|
||||
@@ -226,6 +267,7 @@ impl KnowledgeCache {
|
||||
alias_strings: Vec::new(),
|
||||
alias_entity_ids: Vec::new(),
|
||||
next_entity_id: 0,
|
||||
adjacency: AdjacencyCache::default(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -236,9 +278,29 @@ impl KnowledgeCache {
|
||||
alias_strings: Vec::new(),
|
||||
alias_entity_ids: Vec::new(),
|
||||
next_entity_id: next_id,
|
||||
adjacency: AdjacencyCache::default(),
|
||||
}
|
||||
}
|
||||
|
||||
/// The adjacency index for the graph as it is now: the cached one if the
|
||||
/// graph's fingerprint still matches, otherwise rebuilt and cached.
|
||||
fn adjacency_index(&self) -> std::sync::Arc<AdjacencyIndex> {
|
||||
let fp = graph_fingerprint(&self.entities, &self.relations);
|
||||
let mut slot = self
|
||||
.adjacency
|
||||
.0
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
if let Some((cached_fp, idx)) = slot.as_ref()
|
||||
&& *cached_fp == fp
|
||||
{
|
||||
return idx.clone();
|
||||
}
|
||||
let idx = std::sync::Arc::new(AdjacencyIndex::build(&self.entities, &self.relations));
|
||||
*slot = Some((fp, idx.clone()));
|
||||
idx
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// Entity management
|
||||
// -----------------------------------------------------------------------
|
||||
@@ -397,7 +459,7 @@ impl KnowledgeCache {
|
||||
/// together with their discovered depth. The seed entity itself is NOT
|
||||
/// included. Traversal follows both outgoing and incoming relation edges.
|
||||
pub fn bfs_neighbors(&self, entity_id: u64, max_depth: usize) -> Vec<(Entity, usize)> {
|
||||
let idx = AdjacencyIndex::build(&self.entities, &self.relations);
|
||||
let idx = self.adjacency_index();
|
||||
let mut visited: HashSet<u64> = HashSet::new();
|
||||
let mut queue: VecDeque<(u64, usize)> = VecDeque::new();
|
||||
let mut results: Vec<(Entity, usize)> = Vec::new();
|
||||
@@ -502,7 +564,7 @@ impl KnowledgeCache {
|
||||
min_activation: f32,
|
||||
max_steps: usize,
|
||||
) -> Vec<(u64, f32)> {
|
||||
let idx = AdjacencyIndex::build(&self.entities, &self.relations);
|
||||
let idx = self.adjacency_index();
|
||||
let mut activation: HashMap<u64, f32> = HashMap::new();
|
||||
|
||||
// Initialise seeds with activation 1.0.
|
||||
@@ -631,6 +693,51 @@ impl Default for KnowledgeCache {
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn cached_adjacency_sees_direct_changes_to_the_graph() {
|
||||
// The index is cached across traversals, but entities/relations are
|
||||
// pub Vecs anyone can edit; every kind of edit must be seen.
|
||||
let mut kg = KnowledgeCache::new();
|
||||
let a = kg.add_entity("a", "t", -1);
|
||||
let b = kg.add_entity("b", "t", -1);
|
||||
let c = kg.add_entity("c", "t", -1);
|
||||
kg.add_relation(a, b, "r", 1.0);
|
||||
let ids = |kg: &KnowledgeCache| -> Vec<u64> {
|
||||
let mut v: Vec<u64> = kg.bfs_neighbors(a, 3).iter().map(|(e, _)| e.id).collect();
|
||||
v.sort();
|
||||
v
|
||||
};
|
||||
assert_eq!(ids(&kg), vec![b]);
|
||||
assert_eq!(ids(&kg), vec![b], "cached index reused");
|
||||
|
||||
// Pushed directly, bypassing add_relation.
|
||||
kg.relations.push(Relation {
|
||||
src: b,
|
||||
tgt: c,
|
||||
..Relation::default()
|
||||
});
|
||||
assert_eq!(ids(&kg), vec![b, c]);
|
||||
|
||||
// Rewired in place: same lengths, different edge.
|
||||
kg.relations[1].tgt = a;
|
||||
assert_eq!(ids(&kg), vec![b]);
|
||||
|
||||
// Removed and replaced: same lengths again.
|
||||
kg.relations.pop();
|
||||
kg.relations.push(Relation {
|
||||
src: a,
|
||||
tgt: c,
|
||||
..Relation::default()
|
||||
});
|
||||
assert_eq!(ids(&kg), vec![b, c]);
|
||||
let act: Vec<u64> = kg
|
||||
.spreading_activation(&[a], 0.5, 0.0, 2)
|
||||
.iter()
|
||||
.map(|(id, _)| *id)
|
||||
.collect();
|
||||
assert!(act.contains(&c));
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// Original tests — must remain passing
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
@@ -72,7 +72,8 @@ use ephemeral::{EphemeralConfig, EphemeralStore};
|
||||
pub use ephemeral::{EphemeralEntry, EphemeralStats};
|
||||
use knowledge::KnowledgeCache;
|
||||
use memory_strategy::{Exchange, MemoryStrategy, StrategyOutput};
|
||||
use session::SessionCache;
|
||||
pub use search::SearchOptions;
|
||||
pub use session::{SessionCache, SessionEntry};
|
||||
|
||||
// --- 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,
|
||||
@@ -1020,6 +1028,18 @@ impl HDF5Memory {
|
||||
&self.config
|
||||
}
|
||||
|
||||
/// The sessions recorded in this store.
|
||||
pub fn sessions(&self) -> &SessionCache {
|
||||
&self.sessions
|
||||
}
|
||||
|
||||
/// Mutable access to the sessions, e.g. to add many at once. Changes
|
||||
/// reach the disk at the next checkpoint (any flushing call, such as
|
||||
/// [`HDF5Memory::flush_wal`] or `save_batch`), not immediately.
|
||||
pub fn sessions_mut(&mut self) -> &mut SessionCache {
|
||||
&mut self.sessions
|
||||
}
|
||||
|
||||
/// Get a reference to the knowledge cache.
|
||||
pub fn knowledge(&self) -> &KnowledgeCache {
|
||||
&self.knowledge
|
||||
@@ -1393,6 +1413,35 @@ impl HDF5Memory {
|
||||
}
|
||||
|
||||
impl HDF5Memory {
|
||||
/// Delete many records with a single checkpoint, where
|
||||
/// [`AgentMemory::delete`] checkpoints once per record.
|
||||
///
|
||||
/// All or nothing: if any id is out of range or already deleted (or
|
||||
/// repeated), nothing is deleted and `MemoryError::NotFound` is returned.
|
||||
/// Unlike `delete`, this never auto-compacts, so the records stay in the
|
||||
/// store as tombstones (their indices unchanged) until [`AgentMemory::compact`]
|
||||
/// is called — importers use it to carry over records that were already
|
||||
/// deleted in the source.
|
||||
pub fn delete_batch(&mut self, ids: &[usize]) -> Result<()> {
|
||||
let mut seen = std::collections::HashSet::with_capacity(ids.len());
|
||||
for &id in ids {
|
||||
if self.cache.tombstones.get(id).copied() != Some(0) || !seen.insert(id) {
|
||||
return Err(MemoryError::NotFound(format!(
|
||||
"entry {id} not found or already deleted"
|
||||
)));
|
||||
}
|
||||
}
|
||||
if ids.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
for &id in ids {
|
||||
self.cache.mark_deleted(id);
|
||||
self.hnsw_on_delete(id);
|
||||
self.bm25_on_delete(id);
|
||||
}
|
||||
self.flush()
|
||||
}
|
||||
|
||||
pub fn tick_session(&mut self) -> Result<()> {
|
||||
let d = self.config.decay_factor;
|
||||
for w in self.cache.activation_weights.iter_mut() {
|
||||
@@ -1591,6 +1640,79 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn delete_batch_tombstones_without_compacting() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let path = dir.path().join("test.h5");
|
||||
let mut mem = HDF5Memory::create(make_config(&dir)).unwrap();
|
||||
mem.save_batch(
|
||||
(0..4)
|
||||
.map(|i| make_entry(&format!("record {i}"), &[i as f32, 1.0, 0.0, 0.0]))
|
||||
.collect(),
|
||||
)
|
||||
.unwrap();
|
||||
// 3 of 4 is far past compact_threshold (0.3): delete() would compact.
|
||||
mem.delete_batch(&[0, 1, 3]).unwrap();
|
||||
assert_eq!(mem.count(), 4);
|
||||
assert_eq!(mem.count_active(), 1);
|
||||
drop(mem);
|
||||
|
||||
let mut mem = HDF5Memory::open(&path).unwrap();
|
||||
assert_eq!(mem.cache.tombstones, vec![1, 1, 0, 1]);
|
||||
let hits = mem.hybrid_search(&[0.0, 1.0, 0.0, 0.0], "record", 0.5, 0.5, 10);
|
||||
assert!(
|
||||
hits.iter().all(|r| r.index == 2),
|
||||
"tombstoned record returned"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn delete_batch_is_all_or_nothing() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let mut mem = HDF5Memory::create(make_config(&dir)).unwrap();
|
||||
mem.save_batch(vec![
|
||||
make_entry("a", &[1.0, 0.0, 0.0, 0.0]),
|
||||
make_entry("b", &[0.0, 1.0, 0.0, 0.0]),
|
||||
])
|
||||
.unwrap();
|
||||
for bad in [&[0, 5][..], &[1, 1][..]] {
|
||||
assert!(matches!(
|
||||
mem.delete_batch(bad),
|
||||
Err(MemoryError::NotFound(_))
|
||||
));
|
||||
assert_eq!(mem.count_active(), 2, "{bad:?} deleted something");
|
||||
}
|
||||
mem.delete_batch(&[]).unwrap();
|
||||
assert_eq!(mem.count_active(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sessions_mut_add_at_keeps_timestamp_across_reopen() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let path = dir.path().join("test.h5");
|
||||
let mut mem = HDF5Memory::create(make_config(&dir)).unwrap();
|
||||
mem.sessions_mut()
|
||||
.add_at("s-old", 2, 7, "discord", "old summary", 1.7e15);
|
||||
mem.flush_wal().unwrap();
|
||||
drop(mem);
|
||||
|
||||
let mem = HDF5Memory::open_read_only(&path).unwrap();
|
||||
let s = mem.sessions();
|
||||
assert_eq!(s.len(), 1);
|
||||
let e = &s.entries[0];
|
||||
assert_eq!(
|
||||
(
|
||||
e.id.as_str(),
|
||||
e.start_idx,
|
||||
e.end_idx,
|
||||
e.channel.as_str(),
|
||||
e.ts
|
||||
),
|
||||
("s-old", 2, 7, "discord", 1.7e15)
|
||||
);
|
||||
assert_eq!(s.summaries[0], "old summary");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn create_new_file() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -33,7 +33,7 @@ impl SessionCache {
|
||||
self.entries.is_empty()
|
||||
}
|
||||
|
||||
/// Add a new session with its summary.
|
||||
/// Add a new session with its summary, timestamped now.
|
||||
pub fn add(
|
||||
&mut self,
|
||||
id: &str,
|
||||
@@ -47,6 +47,21 @@ impl SessionCache {
|
||||
.unwrap_or_default()
|
||||
.as_secs_f64()
|
||||
* 1_000_000.0; // microseconds
|
||||
self.add_at(id, start_idx, end_idx, channel, summary, ts);
|
||||
}
|
||||
|
||||
/// Add a session with an explicit timestamp (Unix **microseconds**, the
|
||||
/// unit [`SessionEntry::ts`] uses) — for importers carrying sessions over
|
||||
/// from another store, whose original time should be kept.
|
||||
pub fn add_at(
|
||||
&mut self,
|
||||
id: &str,
|
||||
start_idx: usize,
|
||||
end_idx: usize,
|
||||
channel: &str,
|
||||
summary: &str,
|
||||
ts: f64,
|
||||
) {
|
||||
self.entries.push(SessionEntry {
|
||||
id: id.to_string(),
|
||||
start_idx: start_idx as u64,
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -11,12 +11,14 @@
|
||||
//!
|
||||
//! Configuration matrix:
|
||||
//! - Text lengths: short (50 chars), medium (200 chars), long (1000 chars)
|
||||
//! - Embedding: 384-dim f32 (1536 bytes raw per record)
|
||||
//! - Embedding: 384-dim, stored as float16 (the default for new stores) or
|
||||
//! f32 with `--f32`; "raw" bytes are counted as f32 input either way
|
||||
//! - WAL: enabled and disabled
|
||||
//!
|
||||
//! # Usage
|
||||
//! ```
|
||||
//! cargo run --release --bin footprint_bench
|
||||
//! cargo run --release --bin footprint_bench # float16 stores
|
||||
//! cargo run --release --bin footprint_bench -- --f32 # f32 stores
|
||||
//! ```
|
||||
|
||||
use std::time::Instant;
|
||||
@@ -24,6 +26,9 @@ use std::time::Instant;
|
||||
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};
|
||||
use tempfile::TempDir;
|
||||
|
||||
/// `--f32`: build f32 stores instead of the library's float16 default.
|
||||
static F32: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
|
||||
|
||||
const EMBEDDING_DIM: usize = 384;
|
||||
|
||||
// Raw bytes per record: 384 f32 embeddings + median text + overhead
|
||||
@@ -152,6 +157,9 @@ fn measure_footprint(
|
||||
config.compression = compression;
|
||||
config.compression_level = if compression { 6 } else { 0 };
|
||||
config.compact_threshold = 0.0;
|
||||
if F32.load(std::sync::atomic::Ordering::Relaxed) {
|
||||
config.float16 = false;
|
||||
}
|
||||
|
||||
let mut memory = HDF5Memory::create(config).expect("HDF5Memory::create failed");
|
||||
|
||||
@@ -241,11 +249,19 @@ fn fmt_n(n: usize) -> String {
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
fn main() {
|
||||
if std::env::args().skip(1).any(|a| a == "--f32") {
|
||||
F32.store(true, std::sync::atomic::Ordering::Relaxed);
|
||||
}
|
||||
let stored = if F32.load(std::sync::atomic::Ordering::Relaxed) {
|
||||
"f32 (1,536 bytes per record)"
|
||||
} else {
|
||||
"float16 (768 bytes per record; the default for new stores)"
|
||||
};
|
||||
println!("=================================================================");
|
||||
println!(" ClawhDF5 Memory Footprint Benchmark");
|
||||
println!("=================================================================");
|
||||
println!();
|
||||
println!("Embedding: 384-dim f32 = 1,536 bytes raw per record");
|
||||
println!("Embedding: 384-dim, stored as {stored}; raw input counted as f32");
|
||||
println!("Text lengths: short=50 chars, medium=200 chars, long=1000 chars");
|
||||
println!();
|
||||
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
@@ -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)?);
|
||||
|
||||
@@ -3,7 +3,7 @@ name = "clawhdf5-migrate"
|
||||
version = "2.7.0"
|
||||
edition = "2024"
|
||||
rust-version.workspace = true
|
||||
description = "CLI to migrate SQLite agent memory databases to HDF5 format"
|
||||
description = "CLI to migrate SQLite agent memory databases to clawhdf5-agent stores"
|
||||
license = "MIT"
|
||||
repository = "https://git.redclaw.dev/quantumclaw/clawhdf5"
|
||||
readme = "README.md"
|
||||
@@ -17,10 +17,8 @@ path = "src/main.rs"
|
||||
[dependencies]
|
||||
clawhdf5-agent = { path = "../clawhdf5-agent", version = "2.7.0" }
|
||||
clawhdf5-format = { path = "../clawhdf5-format", version = "2.7.0" }
|
||||
clawhdf5 = { path = "../clawhdf5", version = "2.7.0" }
|
||||
rusqlite = { version = "0.31", features = ["bundled"] }
|
||||
clap = { version = "4", features = ["derive"] }
|
||||
half = { workspace = true }
|
||||
|
||||
[dev-dependencies]
|
||||
tempfile = { workspace = true }
|
||||
|
||||
@@ -3,9 +3,12 @@
|
||||
[](https://crates.io/crates/clawhdf5-migrate)
|
||||
[](https://docs.rs/clawhdf5-migrate)
|
||||
|
||||
CLI tool to migrate SQLite agent memory databases to HDF5 format.
|
||||
CLI tool to migrate SQLite agent memory databases (the ZeroClaw layout) to a
|
||||
[clawhdf5-agent](https://crates.io/crates/clawhdf5-agent) store.
|
||||
|
||||
Converts existing SQLite-based agent memory stores (embeddings, text chunks, metadata) into the HDF5 format used by [clawhdf5-agent](https://crates.io/crates/clawhdf5-agent).
|
||||
The output is written through `clawhdf5-agent`'s own API, so it opens with
|
||||
`HDF5Memory::open` and is searchable immediately: memory records, sessions and
|
||||
the knowledge graph (entities and relations) are carried over.
|
||||
|
||||
## Installation
|
||||
|
||||
@@ -16,9 +19,19 @@ cargo install clawhdf5-migrate
|
||||
## Usage
|
||||
|
||||
```bash
|
||||
clawhdf5-migrate --input agent.db --output agent.h5
|
||||
clawhdf5-migrate --sqlite agent.db --hdf5 agent.h5 --agent-id my-agent
|
||||
```
|
||||
|
||||
Embeddings are stored as float16 (the library default for new stores); pass
|
||||
`--f32` for full precision. Every embedding must have the same dimension
|
||||
(the first row's, or `--embedding-dim`, which a source with no memory records
|
||||
requires); rows are never truncated, and the whole source is checked before an
|
||||
existing output store is replaced. `--incremental` adds only new rows to an
|
||||
existing store of the same dimension and carries over changes to rows'
|
||||
deleted flags, `--skip-deleted` leaves out tombstoned rows, and `--dry-run`
|
||||
only counts.
|
||||
See `clawhdf5-migrate --help` for every option.
|
||||
|
||||
## License
|
||||
|
||||
MIT
|
||||
|
||||
@@ -1,163 +0,0 @@
|
||||
//! Read a migration HDF5 file back into the in-memory data model.
|
||||
//!
|
||||
//! Used to verify migrated content (real validation) and to merge new rows into
|
||||
//! an existing output (incremental migration). Mirrors the layout produced by
|
||||
//! [`crate::hdf5_writer`].
|
||||
|
||||
use clawhdf5::reader::{File, Group};
|
||||
use clawhdf5_format::type_builders::AttrValue;
|
||||
|
||||
use crate::sqlite_reader::{Entity, MemoryChunk, Relation, Session, SqliteData};
|
||||
|
||||
type BoxErr = Box<dyn std::error::Error>;
|
||||
|
||||
fn read_strings(group: &Group<'_>, name: &str) -> Result<Vec<String>, BoxErr> {
|
||||
Ok(group.dataset(name)?.read_string()?)
|
||||
}
|
||||
|
||||
fn read_i64s(group: &Group<'_>, name: &str) -> Result<Vec<i64>, BoxErr> {
|
||||
Ok(group.dataset(name)?.read_i64()?)
|
||||
}
|
||||
|
||||
fn read_f64s(group: &Group<'_>, name: &str) -> Result<Vec<f64>, BoxErr> {
|
||||
Ok(group.dataset(name)?.read_f64()?)
|
||||
}
|
||||
|
||||
/// Read the embeddings dataset as a flat `Vec<f32>` of `n * dim` values,
|
||||
/// handling both f32 and (lossy) f16 storage.
|
||||
fn read_embeddings_flat(group: &Group<'_>) -> Result<Vec<f32>, BoxErr> {
|
||||
Ok(group.dataset("embeddings")?.read_f32()?)
|
||||
}
|
||||
|
||||
/// Read a migration HDF5 file into a [`SqliteData`].
|
||||
pub fn read_hdf5(path: &str) -> Result<SqliteData, BoxErr> {
|
||||
let file = File::open(path)?;
|
||||
|
||||
let embedding_dim = match file.root().attrs()?.get("embedding_dim") {
|
||||
Some(AttrValue::I64(d)) => *d as usize,
|
||||
_ => 0,
|
||||
};
|
||||
|
||||
let chunks = read_chunks(&file, embedding_dim)?;
|
||||
let sessions = read_sessions(&file)?;
|
||||
let entities = read_entities(&file)?;
|
||||
let relations = read_relations(&file)?;
|
||||
|
||||
Ok(SqliteData {
|
||||
chunks,
|
||||
sessions,
|
||||
entities,
|
||||
relations,
|
||||
embedding_dim,
|
||||
// Not a SQLite read — the caller (incremental migration) carries
|
||||
// forward the current run's actual `source_path` from the fresh
|
||||
// SQLite read instead of using this placeholder.
|
||||
source_path: String::new(),
|
||||
})
|
||||
}
|
||||
|
||||
fn read_chunks(file: &File, dim: usize) -> Result<Vec<MemoryChunk>, BoxErr> {
|
||||
let g = file.group("chunks")?;
|
||||
let count = group_count(&g)?;
|
||||
if count == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let ids = read_i64s(&g, "id")?;
|
||||
let texts = read_strings(&g, "text")?;
|
||||
let channels = read_strings(&g, "source_channel")?;
|
||||
let timestamps = read_f64s(&g, "timestamp")?;
|
||||
let session_ids = read_strings(&g, "session_id")?;
|
||||
let tags = read_strings(&g, "tags")?;
|
||||
let deleted = g.dataset("deleted")?.read_i32()?;
|
||||
let emb_flat = read_embeddings_flat(&g)?;
|
||||
let dim = dim.max(1);
|
||||
|
||||
let mut chunks = Vec::with_capacity(ids.len());
|
||||
for (i, &id) in ids.iter().enumerate() {
|
||||
let embedding = emb_flat
|
||||
.get(i * dim..(i + 1) * dim)
|
||||
.map(|s| s.to_vec())
|
||||
.unwrap_or_default();
|
||||
chunks.push(MemoryChunk {
|
||||
id,
|
||||
chunk: texts.get(i).cloned().unwrap_or_default(),
|
||||
embedding,
|
||||
source_channel: channels.get(i).cloned().unwrap_or_default(),
|
||||
timestamp: timestamps.get(i).copied().unwrap_or(0.0),
|
||||
session_id: session_ids.get(i).cloned().unwrap_or_default(),
|
||||
tags: tags.get(i).cloned().unwrap_or_default(),
|
||||
deleted: deleted.get(i).copied().unwrap_or(0),
|
||||
});
|
||||
}
|
||||
Ok(chunks)
|
||||
}
|
||||
|
||||
fn read_sessions(file: &File) -> Result<Vec<Session>, BoxErr> {
|
||||
let g = file.group("sessions")?;
|
||||
if group_count(&g)? == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let ids = read_strings(&g, "id")?;
|
||||
let starts = read_i64s(&g, "start_idx")?;
|
||||
let ends = read_i64s(&g, "end_idx")?;
|
||||
let channels = read_strings(&g, "channel")?;
|
||||
let timestamps = read_f64s(&g, "timestamp")?;
|
||||
let summaries = read_strings(&g, "summary")?;
|
||||
Ok((0..ids.len())
|
||||
.map(|i| Session {
|
||||
id: ids[i].clone(),
|
||||
start_idx: starts.get(i).copied().unwrap_or(0),
|
||||
end_idx: ends.get(i).copied().unwrap_or(0),
|
||||
channel: channels.get(i).cloned().unwrap_or_default(),
|
||||
timestamp: timestamps.get(i).copied().unwrap_or(0.0),
|
||||
summary: summaries.get(i).cloned().unwrap_or_default(),
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
fn read_entities(file: &File) -> Result<Vec<Entity>, BoxErr> {
|
||||
let g = file.group("entities")?;
|
||||
if group_count(&g)? == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let ids = read_i64s(&g, "id")?;
|
||||
let names = read_strings(&g, "name")?;
|
||||
let types = read_strings(&g, "type")?;
|
||||
let emb_idxs = read_i64s(&g, "embedding_idx")?;
|
||||
Ok((0..ids.len())
|
||||
.map(|i| Entity {
|
||||
id: ids[i],
|
||||
name: names.get(i).cloned().unwrap_or_default(),
|
||||
entity_type: types.get(i).cloned().unwrap_or_default(),
|
||||
embedding_idx: emb_idxs.get(i).copied().unwrap_or(-1),
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
fn read_relations(file: &File) -> Result<Vec<Relation>, BoxErr> {
|
||||
let g = file.group("relations")?;
|
||||
if group_count(&g)? == 0 {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let srcs = read_i64s(&g, "src")?;
|
||||
let tgts = read_i64s(&g, "tgt")?;
|
||||
let rels = read_strings(&g, "relation")?;
|
||||
let weights = read_f64s(&g, "weight")?;
|
||||
let timestamps = read_f64s(&g, "timestamp")?;
|
||||
Ok((0..srcs.len())
|
||||
.map(|i| Relation {
|
||||
src: srcs[i],
|
||||
tgt: tgts.get(i).copied().unwrap_or(0),
|
||||
relation: rels.get(i).cloned().unwrap_or_default(),
|
||||
weight: weights.get(i).copied().unwrap_or(1.0),
|
||||
timestamp: timestamps.get(i).copied().unwrap_or(0.0),
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
fn group_count(group: &Group<'_>) -> Result<u64, BoxErr> {
|
||||
match group.attrs()?.get("count") {
|
||||
Some(AttrValue::I64(n)) => Ok(*n as u64),
|
||||
_ => Ok(0),
|
||||
}
|
||||
}
|
||||
@@ -1,366 +0,0 @@
|
||||
use clawhdf5::writer::FileBuilder;
|
||||
use clawhdf5_format::datatype::{CharacterSet, Datatype, StringPadding};
|
||||
use clawhdf5_format::type_builders::AttrValue;
|
||||
|
||||
use crate::sqlite_reader::SqliteData;
|
||||
|
||||
/// Options controlling HDF5 output.
|
||||
pub struct WriteOptions {
|
||||
pub agent_id: String,
|
||||
pub embedder: String,
|
||||
pub compression: bool,
|
||||
pub compression_level: u32,
|
||||
pub float16: bool,
|
||||
}
|
||||
|
||||
/// Write SQLite data to an HDF5 file.
|
||||
pub fn write_hdf5(
|
||||
path: &str,
|
||||
data: &SqliteData,
|
||||
opts: &WriteOptions,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let mut builder = FileBuilder::new();
|
||||
let timestamp = iso8601_now();
|
||||
|
||||
// Root-level metadata attributes
|
||||
builder.set_attr("agent_id", AttrValue::String(opts.agent_id.clone()));
|
||||
builder.set_attr("embedder", AttrValue::String(opts.embedder.clone()));
|
||||
builder.set_attr("embedding_dim", AttrValue::I64(data.embedding_dim as i64));
|
||||
builder.set_attr("source", AttrValue::String("sqlite-migration".into()));
|
||||
builder.set_attr("version", AttrValue::I64(1));
|
||||
// Lineage: which SQLite database this output was migrated from and when,
|
||||
// plus the migrator tool version — so a chain of `--incremental` runs
|
||||
// still has an audit trail instead of every run overwriting the same
|
||||
// static attributes (see research/03_provenance.md, INT-03).
|
||||
builder.set_attr("source_path", AttrValue::String(data.source_path.clone()));
|
||||
builder.set_attr("migrated_at", AttrValue::String(timestamp.clone()));
|
||||
builder.set_attr(
|
||||
"migrator_version",
|
||||
AttrValue::String(env!("CARGO_PKG_VERSION").to_owned()),
|
||||
);
|
||||
|
||||
write_chunks_group(&mut builder, data, opts, ×tamp);
|
||||
write_sessions_group(&mut builder, data);
|
||||
write_entities_group(&mut builder, data);
|
||||
write_relations_group(&mut builder, data);
|
||||
|
||||
builder.write(path)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Current UTC time formatted as an ISO-8601 / RFC-3339 timestamp
|
||||
/// (`YYYY-MM-DDTHH:MM:SSZ`), with no external date/time dependency.
|
||||
fn iso8601_now() -> String {
|
||||
let secs = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs();
|
||||
let days = (secs / 86_400) as i64;
|
||||
let time_of_day = secs % 86_400;
|
||||
let (h, m, s) = (
|
||||
time_of_day / 3600,
|
||||
(time_of_day % 3600) / 60,
|
||||
time_of_day % 60,
|
||||
);
|
||||
let (y, mo, d) = civil_from_days(days);
|
||||
format!("{y:04}-{mo:02}-{d:02}T{h:02}:{m:02}:{s:02}Z")
|
||||
}
|
||||
|
||||
/// Days-since-epoch to (year, month, day), Howard Hinnant's `civil_from_days`
|
||||
/// algorithm (proleptic Gregorian calendar, valid for the full `i64` range).
|
||||
fn civil_from_days(z: i64) -> (i64, u32, u32) {
|
||||
let z = z + 719_468;
|
||||
let era = if z >= 0 { z } else { z - 146_096 } / 146_097;
|
||||
let doe = (z - era * 146_097) as u64; // [0, 146096]
|
||||
let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365; // [0, 399]
|
||||
let y = yoe as i64 + era * 400;
|
||||
let doy = doe - (365 * yoe + yoe / 4 - yoe / 100); // [0, 365]
|
||||
let mp = (5 * doy + 2) / 153; // [0, 11]
|
||||
let d = (doy - (153 * mp + 2) / 5 + 1) as u32; // [1, 31]
|
||||
let m = (if mp < 10 { mp + 3 } else { mp - 9 }) as u32; // [1, 12]
|
||||
let y = if m <= 2 { y + 1 } else { y };
|
||||
(y, m, d)
|
||||
}
|
||||
|
||||
/// Build a fixed-length string Datatype from the max byte length of the items.
|
||||
fn string_dtype(max_len: usize) -> Datatype {
|
||||
Datatype::String {
|
||||
size: max_len.max(1) as u32,
|
||||
padding: StringPadding::NullPad,
|
||||
charset: CharacterSet::Utf8,
|
||||
}
|
||||
}
|
||||
|
||||
/// Pack a slice of strings into null-padded raw bytes of uniform width.
|
||||
fn pack_strings(strings: &[String]) -> (Vec<u8>, usize) {
|
||||
let max_len = strings.iter().map(|s| s.len()).max().unwrap_or(0).max(1);
|
||||
let mut buf = vec![0u8; strings.len() * max_len];
|
||||
for (i, s) in strings.iter().enumerate() {
|
||||
let start = i * max_len;
|
||||
let bytes = s.as_bytes();
|
||||
let copy_len = bytes.len().min(max_len);
|
||||
buf[start..start + copy_len].copy_from_slice(&bytes[..copy_len]);
|
||||
}
|
||||
(buf, max_len)
|
||||
}
|
||||
|
||||
fn apply_compression(ds: &mut clawhdf5_format::type_builders::DatasetBuilder, opts: &WriteOptions) {
|
||||
if opts.compression {
|
||||
ds.with_deflate(opts.compression_level);
|
||||
ds.with_shuffle();
|
||||
}
|
||||
}
|
||||
|
||||
fn write_chunks_group(
|
||||
builder: &mut FileBuilder,
|
||||
data: &SqliteData,
|
||||
opts: &WriteOptions,
|
||||
timestamp: &str,
|
||||
) {
|
||||
let mut group = builder.create_group("chunks");
|
||||
let n = data.chunks.len() as u64;
|
||||
|
||||
if n == 0 {
|
||||
group.set_attr("count", AttrValue::I64(0));
|
||||
builder.add_group(group.finish());
|
||||
return;
|
||||
}
|
||||
|
||||
group.set_attr("count", AttrValue::I64(n as i64));
|
||||
|
||||
// Source attribution attached directly to the content-bearing datasets
|
||||
// (SHA-256 of the raw bytes + creator/timestamp/source), so the chunk
|
||||
// text and embeddings each carry their own verifiable provenance
|
||||
// (see clawhdf5_format::provenance / `Dataset::verify_provenance`).
|
||||
let source_opt = if data.source_path.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(data.source_path.as_str())
|
||||
};
|
||||
|
||||
// ids
|
||||
let ids: Vec<i64> = data.chunks.iter().map(|c| c.id).collect();
|
||||
group.create_dataset("id").with_i64_data(&ids);
|
||||
|
||||
// text
|
||||
let texts: Vec<String> = data.chunks.iter().map(|c| c.chunk.clone()).collect();
|
||||
let (text_raw, text_len) = pack_strings(&texts);
|
||||
group
|
||||
.create_dataset("text")
|
||||
.with_compound_data(string_dtype(text_len), text_raw, n)
|
||||
.with_provenance("clawhdf5-migrate", timestamp, source_opt);
|
||||
|
||||
// embeddings - flatten to [N, dim]
|
||||
let dim = data.embedding_dim;
|
||||
if opts.float16 {
|
||||
let f16_data: Vec<u16> = data
|
||||
.chunks
|
||||
.iter()
|
||||
.flat_map(|c| {
|
||||
c.embedding
|
||||
.iter()
|
||||
.map(|&v| half::f16::from_f32(v).to_bits())
|
||||
})
|
||||
.collect();
|
||||
let raw: Vec<u8> = f16_data.iter().flat_map(|v| v.to_le_bytes()).collect();
|
||||
let f16_dtype = Datatype::FloatingPoint {
|
||||
size: 2,
|
||||
byte_order: clawhdf5_format::datatype::DatatypeByteOrder::LittleEndian,
|
||||
bit_offset: 0,
|
||||
bit_precision: 16,
|
||||
exponent_location: 10,
|
||||
exponent_size: 5,
|
||||
mantissa_location: 0,
|
||||
mantissa_size: 10,
|
||||
exponent_bias: 15,
|
||||
};
|
||||
let ds = group
|
||||
.create_dataset("embeddings")
|
||||
.with_compound_data(f16_dtype, raw, n)
|
||||
.with_shape(&[n, dim as u64])
|
||||
.with_provenance("clawhdf5-migrate", timestamp, source_opt);
|
||||
apply_compression(ds, opts);
|
||||
} else {
|
||||
let flat: Vec<f32> = data
|
||||
.chunks
|
||||
.iter()
|
||||
.flat_map(|c| c.embedding.iter().copied())
|
||||
.collect();
|
||||
let ds = group
|
||||
.create_dataset("embeddings")
|
||||
.with_f32_data(&flat)
|
||||
.with_shape(&[n, dim as u64])
|
||||
.with_provenance("clawhdf5-migrate", timestamp, source_opt);
|
||||
apply_compression(ds, opts);
|
||||
}
|
||||
|
||||
// source_channel
|
||||
let channels: Vec<String> = data
|
||||
.chunks
|
||||
.iter()
|
||||
.map(|c| c.source_channel.clone())
|
||||
.collect();
|
||||
let (ch_raw, ch_len) = pack_strings(&channels);
|
||||
group
|
||||
.create_dataset("source_channel")
|
||||
.with_compound_data(string_dtype(ch_len), ch_raw, n);
|
||||
|
||||
// timestamp
|
||||
let timestamps: Vec<f64> = data.chunks.iter().map(|c| c.timestamp).collect();
|
||||
group.create_dataset("timestamp").with_f64_data(×tamps);
|
||||
|
||||
// session_id
|
||||
let sess_ids: Vec<String> = data.chunks.iter().map(|c| c.session_id.clone()).collect();
|
||||
let (sid_raw, sid_len) = pack_strings(&sess_ids);
|
||||
group
|
||||
.create_dataset("session_id")
|
||||
.with_compound_data(string_dtype(sid_len), sid_raw, n);
|
||||
|
||||
// tags
|
||||
let tags: Vec<String> = data.chunks.iter().map(|c| c.tags.clone()).collect();
|
||||
let (tag_raw, tag_len) = pack_strings(&tags);
|
||||
group
|
||||
.create_dataset("tags")
|
||||
.with_compound_data(string_dtype(tag_len), tag_raw, n);
|
||||
|
||||
// deleted
|
||||
let deleted: Vec<i32> = data.chunks.iter().map(|c| c.deleted).collect();
|
||||
group.create_dataset("deleted").with_i32_data(&deleted);
|
||||
|
||||
builder.add_group(group.finish());
|
||||
}
|
||||
|
||||
fn write_sessions_group(builder: &mut FileBuilder, data: &SqliteData) {
|
||||
let mut group = builder.create_group("sessions");
|
||||
let n = data.sessions.len() as u64;
|
||||
group.set_attr("count", AttrValue::I64(n as i64));
|
||||
|
||||
if n == 0 {
|
||||
builder.add_group(group.finish());
|
||||
return;
|
||||
}
|
||||
|
||||
let ids: Vec<String> = data.sessions.iter().map(|s| s.id.clone()).collect();
|
||||
let (id_raw, id_len) = pack_strings(&ids);
|
||||
group
|
||||
.create_dataset("id")
|
||||
.with_compound_data(string_dtype(id_len), id_raw, n);
|
||||
|
||||
let start_idxs: Vec<i64> = data.sessions.iter().map(|s| s.start_idx).collect();
|
||||
group.create_dataset("start_idx").with_i64_data(&start_idxs);
|
||||
|
||||
let end_idxs: Vec<i64> = data.sessions.iter().map(|s| s.end_idx).collect();
|
||||
group.create_dataset("end_idx").with_i64_data(&end_idxs);
|
||||
|
||||
let channels: Vec<String> = data.sessions.iter().map(|s| s.channel.clone()).collect();
|
||||
let (ch_raw, ch_len) = pack_strings(&channels);
|
||||
group
|
||||
.create_dataset("channel")
|
||||
.with_compound_data(string_dtype(ch_len), ch_raw, n);
|
||||
|
||||
let timestamps: Vec<f64> = data.sessions.iter().map(|s| s.timestamp).collect();
|
||||
group.create_dataset("timestamp").with_f64_data(×tamps);
|
||||
|
||||
let summaries: Vec<String> = data.sessions.iter().map(|s| s.summary.clone()).collect();
|
||||
let (sum_raw, sum_len) = pack_strings(&summaries);
|
||||
group
|
||||
.create_dataset("summary")
|
||||
.with_compound_data(string_dtype(sum_len), sum_raw, n);
|
||||
|
||||
builder.add_group(group.finish());
|
||||
}
|
||||
|
||||
fn write_entities_group(builder: &mut FileBuilder, data: &SqliteData) {
|
||||
let mut group = builder.create_group("entities");
|
||||
let n = data.entities.len() as u64;
|
||||
group.set_attr("count", AttrValue::I64(n as i64));
|
||||
|
||||
if n == 0 {
|
||||
builder.add_group(group.finish());
|
||||
return;
|
||||
}
|
||||
|
||||
let ids: Vec<i64> = data.entities.iter().map(|e| e.id).collect();
|
||||
group.create_dataset("id").with_i64_data(&ids);
|
||||
|
||||
let names: Vec<String> = data.entities.iter().map(|e| e.name.clone()).collect();
|
||||
let (name_raw, name_len) = pack_strings(&names);
|
||||
group
|
||||
.create_dataset("name")
|
||||
.with_compound_data(string_dtype(name_len), name_raw, n);
|
||||
|
||||
let types: Vec<String> = data
|
||||
.entities
|
||||
.iter()
|
||||
.map(|e| e.entity_type.clone())
|
||||
.collect();
|
||||
let (type_raw, type_len) = pack_strings(&types);
|
||||
group
|
||||
.create_dataset("type")
|
||||
.with_compound_data(string_dtype(type_len), type_raw, n);
|
||||
|
||||
let emb_idxs: Vec<i64> = data.entities.iter().map(|e| e.embedding_idx).collect();
|
||||
group
|
||||
.create_dataset("embedding_idx")
|
||||
.with_i64_data(&emb_idxs);
|
||||
|
||||
builder.add_group(group.finish());
|
||||
}
|
||||
|
||||
fn write_relations_group(builder: &mut FileBuilder, data: &SqliteData) {
|
||||
let mut group = builder.create_group("relations");
|
||||
let n = data.relations.len() as u64;
|
||||
group.set_attr("count", AttrValue::I64(n as i64));
|
||||
|
||||
if n == 0 {
|
||||
builder.add_group(group.finish());
|
||||
return;
|
||||
}
|
||||
|
||||
let srcs: Vec<i64> = data.relations.iter().map(|r| r.src).collect();
|
||||
group.create_dataset("src").with_i64_data(&srcs);
|
||||
|
||||
let tgts: Vec<i64> = data.relations.iter().map(|r| r.tgt).collect();
|
||||
group.create_dataset("tgt").with_i64_data(&tgts);
|
||||
|
||||
let rels: Vec<String> = data.relations.iter().map(|r| r.relation.clone()).collect();
|
||||
let (rel_raw, rel_len) = pack_strings(&rels);
|
||||
group
|
||||
.create_dataset("relation")
|
||||
.with_compound_data(string_dtype(rel_len), rel_raw, n);
|
||||
|
||||
let weights: Vec<f64> = data.relations.iter().map(|r| r.weight).collect();
|
||||
group.create_dataset("weight").with_f64_data(&weights);
|
||||
|
||||
let timestamps: Vec<f64> = data.relations.iter().map(|r| r.timestamp).collect();
|
||||
group.create_dataset("timestamp").with_f64_data(×tamps);
|
||||
|
||||
builder.add_group(group.finish());
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod time_tests {
|
||||
use super::civil_from_days;
|
||||
|
||||
#[test]
|
||||
fn epoch_day_zero_is_1970_01_01() {
|
||||
assert_eq!(civil_from_days(0), (1970, 1, 1));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn known_dates_roundtrip() {
|
||||
// 2026-08-16 is 20,681 days after 1970-01-01.
|
||||
assert_eq!(civil_from_days(20_681), (2026, 8, 16));
|
||||
// 2000-02-29 (leap day itself) and 2000-03-01 (the day after).
|
||||
assert_eq!(civil_from_days(11_016), (2000, 2, 29));
|
||||
assert_eq!(civil_from_days(11_017), (2000, 3, 1));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn iso8601_now_has_expected_shape() {
|
||||
let ts = super::iso8601_now();
|
||||
assert_eq!(ts.len(), "2026-08-16T00:00:00Z".len());
|
||||
assert!(ts.starts_with("20")); // sanity: 21st-century year
|
||||
assert!(ts.ends_with('Z'));
|
||||
}
|
||||
}
|
||||
+1041
-421
File diff suppressed because it is too large
Load Diff
@@ -50,12 +50,8 @@ pub struct SqliteData {
|
||||
pub sessions: Vec<Session>,
|
||||
pub entities: Vec<Entity>,
|
||||
pub relations: Vec<Relation>,
|
||||
/// `--embedding-dim`, or the first row's; 0 when neither exists.
|
||||
pub embedding_dim: usize,
|
||||
/// Filesystem path of the SQLite database this data was read from, for
|
||||
/// provenance attribution on the HDF5 output. Empty when the data did
|
||||
/// not come directly from a SQLite read (e.g. re-read of a prior HDF5
|
||||
/// migration output for an incremental merge).
|
||||
pub source_path: String,
|
||||
}
|
||||
|
||||
/// A table name plus the ordered column names the reader maps by position.
|
||||
@@ -167,11 +163,13 @@ pub fn read_counts(
|
||||
})
|
||||
}
|
||||
|
||||
/// Auto-detect embedding dimension from the first chunk's BLOB size.
|
||||
/// Auto-detect embedding dimension from the BLOB size of the first chunk (in
|
||||
/// id order, deleted or not).
|
||||
fn detect_embedding_dim(conn: &Connection, config: &SchemaConfig) -> SqlResult<Option<usize>> {
|
||||
let emb_col = config.chunks.columns.get(2).copied().unwrap_or("embedding");
|
||||
let id_col = config.chunks.columns.first().copied().unwrap_or("id");
|
||||
let mut stmt = conn.prepare(&format!(
|
||||
"SELECT {emb_col} FROM {} LIMIT 1",
|
||||
"SELECT {emb_col} FROM {} ORDER BY {id_col} LIMIT 1",
|
||||
config.chunks.table
|
||||
))?;
|
||||
let mut rows = stmt.query([])?;
|
||||
@@ -195,24 +193,16 @@ fn blob_to_f32(blob: &[u8]) -> Vec<f32> {
|
||||
/// Read all data from a ZeroClaw SQLite database.
|
||||
///
|
||||
/// If `skip_deleted` is true, rows with `deleted=1` are excluded from chunks.
|
||||
/// If `embedding_dim` is `None`, auto-detect from the first row.
|
||||
/// If `embedding_dim` is `None`, auto-detect from the first row (0 when there
|
||||
/// are no rows). Embeddings are returned at their full stored length whatever
|
||||
/// the dimension: checking that every row matches it is the writer's job
|
||||
/// (`store_writer::write_store`), so a mismatch is an error, not silent
|
||||
/// truncation.
|
||||
pub fn read_sqlite(
|
||||
path: &str,
|
||||
skip_deleted: bool,
|
||||
embedding_dim: Option<usize>,
|
||||
config: &SchemaConfig,
|
||||
) -> Result<SqliteData, Box<dyn std::error::Error>> {
|
||||
read_sqlite_filtered(path, skip_deleted, embedding_dim, config, 0)
|
||||
}
|
||||
|
||||
/// Like [`read_sqlite`] but only reads chunks whose id is greater than
|
||||
/// `min_chunk_id` (0 = all). Used for incremental migration.
|
||||
pub fn read_sqlite_filtered(
|
||||
path: &str,
|
||||
skip_deleted: bool,
|
||||
embedding_dim: Option<usize>,
|
||||
config: &SchemaConfig,
|
||||
min_chunk_id: i64,
|
||||
) -> Result<SqliteData, Box<dyn std::error::Error>> {
|
||||
let conn = Connection::open(path)?;
|
||||
|
||||
@@ -221,7 +211,7 @@ pub fn read_sqlite_filtered(
|
||||
None => detect_embedding_dim(&conn, config)?.unwrap_or(0),
|
||||
};
|
||||
|
||||
let chunks = read_chunks(&conn, skip_deleted, dim, config, min_chunk_id)?;
|
||||
let chunks = read_chunks(&conn, skip_deleted, config)?;
|
||||
let sessions = read_sessions(&conn, config)?;
|
||||
let entities = read_entities(&conn, config)?;
|
||||
let relations = read_relations(&conn, config)?;
|
||||
@@ -232,42 +222,43 @@ pub fn read_sqlite_filtered(
|
||||
entities,
|
||||
relations,
|
||||
embedding_dim: dim,
|
||||
source_path: path.to_owned(),
|
||||
})
|
||||
}
|
||||
|
||||
fn read_chunks(
|
||||
conn: &Connection,
|
||||
skip_deleted: bool,
|
||||
expected_dim: usize,
|
||||
config: &SchemaConfig,
|
||||
min_chunk_id: i64,
|
||||
) -> SqlResult<Vec<MemoryChunk>> {
|
||||
let id_col = config.chunks.columns.first().copied().unwrap_or("id");
|
||||
let deleted_col = config.chunks.columns.get(7).copied().unwrap_or("deleted");
|
||||
let mut conds = Vec::new();
|
||||
let mut where_clause = String::new();
|
||||
if skip_deleted {
|
||||
conds.push(format!("{deleted_col} = 0"));
|
||||
where_clause = format!(" WHERE {deleted_col} = 0");
|
||||
}
|
||||
if min_chunk_id > 0 {
|
||||
conds.push(format!("{id_col} > {min_chunk_id}"));
|
||||
}
|
||||
let where_clause = if conds.is_empty() {
|
||||
String::new()
|
||||
} else {
|
||||
format!(" WHERE {}", conds.join(" AND "))
|
||||
};
|
||||
// In id order, so the store's records follow the source's order.
|
||||
where_clause.push_str(&format!(" ORDER BY {id_col}"));
|
||||
let sql = config.chunks.select(&where_clause);
|
||||
|
||||
let mut stmt = conn.prepare(&sql)?;
|
||||
let rows = stmt.query_map([], |row| {
|
||||
let blob: Vec<u8> = row.get(2)?;
|
||||
let mut embedding = blob_to_f32(&blob);
|
||||
|
||||
// Validate/truncate to expected dimension
|
||||
if expected_dim > 0 {
|
||||
embedding.truncate(expected_dim);
|
||||
if !blob.len().is_multiple_of(4) {
|
||||
let id: i64 = row.get(0)?;
|
||||
return Err(rusqlite::Error::FromSqlConversionFailure(
|
||||
2,
|
||||
rusqlite::types::Type::Blob,
|
||||
format!(
|
||||
"chunk id {id}: embedding BLOB is {} bytes, not a whole number of \
|
||||
little-endian f32 values",
|
||||
blob.len()
|
||||
)
|
||||
.into(),
|
||||
));
|
||||
}
|
||||
// Read at full length: rows of the wrong dimension are rejected by
|
||||
// the writer, never truncated to fit.
|
||||
let embedding = blob_to_f32(&blob);
|
||||
|
||||
Ok(MemoryChunk {
|
||||
id: row.get(0)?,
|
||||
|
||||
@@ -0,0 +1,407 @@
|
||||
//! Write migrated SQLite data into a clawhdf5-agent store.
|
||||
//!
|
||||
//! Everything goes through `clawhdf5-agent`'s own API — `HDF5Memory::create`
|
||||
//! (or `open` for `--incremental`), `save_batch`, `delete_batch`, the session
|
||||
//! cache and the knowledge graph — so the result is an ordinary agent store
|
||||
//! that `HDF5Memory::open` accepts, not a second hand-built copy of its schema.
|
||||
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::path::Path;
|
||||
|
||||
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};
|
||||
use clawhdf5_format::float16::round_to_f16;
|
||||
|
||||
use crate::sqlite_reader::{MemoryChunk, SqliteData};
|
||||
|
||||
type BoxErr = Box<dyn std::error::Error>;
|
||||
|
||||
/// SQLite timestamps are Unix seconds; the agent's session and relation
|
||||
/// timestamps are Unix microseconds (memory records stay in seconds).
|
||||
pub const US_PER_SEC: f64 = 1_000_000.0;
|
||||
|
||||
/// Options controlling the output store.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct WriteOptions {
|
||||
pub agent_id: String,
|
||||
pub embedder: String,
|
||||
pub compression: bool,
|
||||
pub compression_level: u32,
|
||||
/// Store full-precision `f32` embeddings instead of the library default
|
||||
/// (half precision). Only applies to a newly created store: an existing
|
||||
/// store keeps the precision it was created with.
|
||||
pub f32: bool,
|
||||
/// Add to the store at the output path if there is one, instead of
|
||||
/// replacing it.
|
||||
pub incremental: bool,
|
||||
/// Leave out deleted source rows that are not in the store. (A deleted
|
||||
/// row that matches an active store record still tombstones it, so pass
|
||||
/// deleted rows in `data` for an incremental run.)
|
||||
pub skip_deleted: bool,
|
||||
}
|
||||
|
||||
/// What the migration wrote, and where each source row went, so validation
|
||||
/// can compare the store with the source row by row.
|
||||
#[derive(Debug, Default)]
|
||||
pub struct Migration {
|
||||
/// Whether the output store existed and was added to (`--incremental`).
|
||||
pub appended_to_existing: bool,
|
||||
/// The store's embedding precision.
|
||||
pub float16: bool,
|
||||
pub embedding_dim: usize,
|
||||
/// Records in the store after the migration (including tombstones).
|
||||
pub store_count: usize,
|
||||
/// `(store index, source chunk index)` of every record written.
|
||||
pub records: Vec<(usize, usize)>,
|
||||
/// Source chunks already in the store (incremental), not written again.
|
||||
pub chunks_present: usize,
|
||||
/// `(store index, source chunk index)` of records that were active in
|
||||
/// the store but whose source row is now deleted (incremental): they were
|
||||
/// tombstoned by this run.
|
||||
pub deleted_in_store: Vec<(usize, usize)>,
|
||||
/// Source rows that were deleted in the store but are active in the
|
||||
/// source (incremental): the agent has no un-delete, so each was written
|
||||
/// again as a new record (counted in `records` too).
|
||||
pub restored: usize,
|
||||
/// Deleted source rows left out because of `skip_deleted`.
|
||||
pub deleted_skipped: usize,
|
||||
/// `(store session index, source session index)` of each session written.
|
||||
pub sessions: Vec<(usize, usize)>,
|
||||
pub sessions_present: usize,
|
||||
/// `(store entity id, source entity index)` of each entity written.
|
||||
pub entities: Vec<(u64, usize)>,
|
||||
pub entities_present: usize,
|
||||
/// SQLite entity id -> store entity id, for every source entity.
|
||||
pub entity_ids: HashMap<i64, u64>,
|
||||
/// `(store relation index, source relation index)` of each relation written.
|
||||
pub relations: Vec<(usize, usize)>,
|
||||
pub relations_present: usize,
|
||||
/// Source relations naming an entity id that is not in the entities
|
||||
/// table; the knowledge graph cannot hold them, so they are skipped.
|
||||
pub dangling_relations: Vec<usize>,
|
||||
/// Messages of the write-anomaly alerts the agent raised while importing
|
||||
/// (informational; they never block a save — a bulk import typically
|
||||
/// trips the write-rate check).
|
||||
pub anomaly_alerts: Vec<String>,
|
||||
}
|
||||
|
||||
/// Identity of a memory record for incremental de-duplication: every field
|
||||
/// the agent stores except the embedding (whose stored form depends on the
|
||||
/// store's precision).
|
||||
type RecordKey = (String, String, String, String, u64);
|
||||
|
||||
fn record_key(
|
||||
chunk: &str,
|
||||
source_channel: &str,
|
||||
session_id: &str,
|
||||
tags: &str,
|
||||
ts: f64,
|
||||
) -> RecordKey {
|
||||
(
|
||||
chunk.to_owned(),
|
||||
source_channel.to_owned(),
|
||||
session_id.to_owned(),
|
||||
tags.to_owned(),
|
||||
ts.to_bits(),
|
||||
)
|
||||
}
|
||||
|
||||
/// Reject rows the agent would otherwise store differently from the source,
|
||||
/// or not at all: an embedding of a different length from the store's
|
||||
/// dimension (the agent pads/truncates silently), an empty embedding, or, in
|
||||
/// a float16 store, a value beyond the half-precision range.
|
||||
///
|
||||
/// Every source row is checked, including ones that end up not being written
|
||||
/// (already in the store, or deleted and skipped): the source must be
|
||||
/// consistent as a whole, and the check runs before the store is touched.
|
||||
fn check_chunks(chunks: &[MemoryChunk], dim: usize, float16: bool) -> Result<(), BoxErr> {
|
||||
for c in chunks {
|
||||
if c.embedding.is_empty() {
|
||||
return Err(format!(
|
||||
"chunk id {}: the embedding is empty; an agent store needs an embedding \
|
||||
for every record",
|
||||
c.id
|
||||
)
|
||||
.into());
|
||||
}
|
||||
if c.embedding.len() != dim {
|
||||
return Err(format!(
|
||||
"chunk id {}: embedding has {} values, expected {dim}; every row must have \
|
||||
the store's dimension (detected from the first row unless --embedding-dim \
|
||||
is given), and rows are never truncated or padded to fit",
|
||||
c.id,
|
||||
c.embedding.len()
|
||||
)
|
||||
.into());
|
||||
}
|
||||
if float16
|
||||
&& let Some((k, v)) = c
|
||||
.embedding
|
||||
.iter()
|
||||
.enumerate()
|
||||
.find(|&(_, &v)| v.is_finite() && round_to_f16(v).is_infinite())
|
||||
{
|
||||
return Err(format!(
|
||||
"chunk id {}: embedding[{k}] = {v} is outside the half-precision range \
|
||||
(±65504) of a float16 store; migrate with --f32",
|
||||
c.id
|
||||
)
|
||||
.into());
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Migrate `data` into the agent store at `path`.
|
||||
///
|
||||
/// Without `opts.incremental` (or when nothing exists at `path`) a new store
|
||||
/// is created, replacing any file there — but only once every source row has
|
||||
/// passed [`check_chunks`], so a source that cannot be migrated leaves an
|
||||
/// existing store untouched. With it, the existing store is opened and only
|
||||
/// source rows it does not already hold are added: memory records are
|
||||
/// matched on their content, sessions on their id, entities on name and
|
||||
/// type, relations on (source, target, relation). A matched record then
|
||||
/// takes the source row's deleted flag: see [`Migration::deleted_in_store`]
|
||||
/// and [`Migration::restored`].
|
||||
pub fn write_store(
|
||||
path: &Path,
|
||||
data: &SqliteData,
|
||||
opts: &WriteOptions,
|
||||
) -> Result<Migration, BoxErr> {
|
||||
let existing = opts.incremental && path.exists();
|
||||
let mut mem = if existing {
|
||||
// `open` does not modify the store beyond what the agent itself does
|
||||
// on open; the checks below run before anything is written.
|
||||
let mem = HDF5Memory::open(path)?;
|
||||
let dim = mem.config().embedding_dim;
|
||||
// `data.embedding_dim` is 0 only for a source with no records and no
|
||||
// --embedding-dim, which has no dimension to disagree with.
|
||||
if data.embedding_dim != 0 && dim != data.embedding_dim {
|
||||
let hint = if dim == 0 {
|
||||
" (a store created from a source with no memory records; re-create it \
|
||||
with --embedding-dim)"
|
||||
} else {
|
||||
""
|
||||
};
|
||||
return Err(format!(
|
||||
"the store at {} has embedding_dim {dim}{hint}, the source {}; \
|
||||
embeddings of a different dimension cannot be added to it",
|
||||
path.display(),
|
||||
data.embedding_dim
|
||||
)
|
||||
.into());
|
||||
}
|
||||
check_chunks(&data.chunks, dim, mem.config().float16)?;
|
||||
mem
|
||||
} else {
|
||||
// (With records, a dimension of 0 means an empty first embedding,
|
||||
// which `check_chunks` reports more precisely.)
|
||||
if data.embedding_dim == 0 && data.chunks.is_empty() {
|
||||
return Err(
|
||||
"the source has no memory records to detect the embedding dimension \
|
||||
from; pass --embedding-dim (the dimension of the agent's embedder), \
|
||||
or the store could never hold a record"
|
||||
.into(),
|
||||
);
|
||||
}
|
||||
let mut config = MemoryConfig::new(path.to_path_buf(), &opts.agent_id, data.embedding_dim);
|
||||
config.embedder = opts.embedder.clone();
|
||||
config.compression = opts.compression;
|
||||
config.compression_level = opts.compression_level;
|
||||
// Only ever switch the library default off (as `clawhdf5-cli create`).
|
||||
if opts.f32 {
|
||||
config.float16 = false;
|
||||
}
|
||||
// Before `create`, which replaces whatever is at `path`.
|
||||
check_chunks(&data.chunks, config.embedding_dim, config.float16)?;
|
||||
HDF5Memory::create(config)?
|
||||
};
|
||||
let float16 = mem.config().float16;
|
||||
let dim = mem.config().embedding_dim;
|
||||
|
||||
let mut m = Migration {
|
||||
appended_to_existing: existing,
|
||||
float16,
|
||||
embedding_dim: dim,
|
||||
..Migration::default()
|
||||
};
|
||||
|
||||
// ---- Memory records --------------------------------------------------
|
||||
// Store indices of every record the store already holds, by content, so
|
||||
// a source row that appears twice is only treated as present as often
|
||||
// as the store has it.
|
||||
let mut present: HashMap<RecordKey, Vec<usize>> = HashMap::new();
|
||||
if existing {
|
||||
let c = &mem.cache;
|
||||
for i in 0..c.len() {
|
||||
let key = record_key(
|
||||
&c.chunks[i],
|
||||
&c.source_channels[i],
|
||||
&c.session_ids[i],
|
||||
&c.tags[i],
|
||||
c.timestamps[i],
|
||||
);
|
||||
present.entry(key).or_default().push(i);
|
||||
}
|
||||
}
|
||||
let key_of = |c: &MemoryChunk| {
|
||||
record_key(
|
||||
&c.chunk,
|
||||
&c.source_channel,
|
||||
&c.session_id,
|
||||
&c.tags,
|
||||
c.timestamp,
|
||||
)
|
||||
};
|
||||
let tombstoned = |idx: usize| mem.cache.tombstones[idx] != 0;
|
||||
// Pass 1: a store record in the same deleted state as the source row.
|
||||
let mut unmatched: Vec<usize> = Vec::new();
|
||||
for (i, c) in data.chunks.iter().enumerate() {
|
||||
let src_deleted = c.deleted != 0;
|
||||
let hit = present.get_mut(&key_of(c)).and_then(|idxs| {
|
||||
let at = idxs.iter().position(|&x| tombstoned(x) == src_deleted)?;
|
||||
Some(idxs.remove(at))
|
||||
});
|
||||
match hit {
|
||||
Some(_) => m.chunks_present += 1,
|
||||
None => unmatched.push(i),
|
||||
}
|
||||
}
|
||||
// Pass 2: a store record whose deleted state differs — the source row
|
||||
// was deleted or restored since the last migration. The source wins.
|
||||
let mut new_chunks: Vec<usize> = Vec::with_capacity(unmatched.len());
|
||||
let mut delete_in_store: Vec<usize> = Vec::new();
|
||||
for i in unmatched {
|
||||
let c = &data.chunks[i];
|
||||
let hit = present
|
||||
.get_mut(&key_of(c))
|
||||
.and_then(|idxs| (!idxs.is_empty()).then(|| idxs.remove(0)));
|
||||
match hit {
|
||||
// Active in the store, deleted in the source: tombstone it.
|
||||
Some(idx) if c.deleted != 0 => {
|
||||
m.deleted_in_store.push((idx, i));
|
||||
delete_in_store.push(idx);
|
||||
}
|
||||
// Deleted in the store, active in the source. The agent has no
|
||||
// un-delete, so the row is written again as a new active record
|
||||
// (the tombstone stays until the store is compacted).
|
||||
Some(_) => {
|
||||
m.restored += 1;
|
||||
new_chunks.push(i);
|
||||
}
|
||||
None if c.deleted != 0 && opts.skip_deleted => m.deleted_skipped += 1,
|
||||
None => new_chunks.push(i),
|
||||
}
|
||||
}
|
||||
new_chunks.sort_unstable();
|
||||
let to_write: Vec<&MemoryChunk> = new_chunks.iter().map(|&i| &data.chunks[i]).collect();
|
||||
|
||||
// ---- Sessions (in the cache; persisted by the save_batch checkpoint) ---
|
||||
let known_sessions: HashSet<String> = mem
|
||||
.sessions()
|
||||
.entries
|
||||
.iter()
|
||||
.map(|e| e.id.clone())
|
||||
.collect();
|
||||
for (i, s) in data.sessions.iter().enumerate() {
|
||||
if known_sessions.contains(&s.id) {
|
||||
m.sessions_present += 1;
|
||||
continue;
|
||||
}
|
||||
let sessions = mem.sessions_mut();
|
||||
let at = sessions.len();
|
||||
sessions.add_at(
|
||||
&s.id,
|
||||
s.start_idx.max(0) as usize,
|
||||
s.end_idx.max(0) as usize,
|
||||
&s.channel,
|
||||
&s.summary,
|
||||
s.timestamp * US_PER_SEC,
|
||||
);
|
||||
m.sessions.push((at, i));
|
||||
}
|
||||
|
||||
// ---- Knowledge graph -------------------------------------------------
|
||||
let kg = mem.knowledge_mut();
|
||||
// Matched only against what the store held before this run: the source
|
||||
// itself is copied as it is, duplicates included.
|
||||
let by_name_type: HashMap<(String, String), u64> = kg
|
||||
.entities
|
||||
.iter()
|
||||
.map(|e| ((e.name.clone(), e.entity_type.clone()), e.id))
|
||||
.collect();
|
||||
for (i, e) in data.entities.iter().enumerate() {
|
||||
let key = (e.name.clone(), e.entity_type.clone());
|
||||
let id = match by_name_type.get(&key) {
|
||||
Some(&id) => {
|
||||
m.entities_present += 1;
|
||||
id
|
||||
}
|
||||
None => {
|
||||
let id = kg.add_entity(&e.name, &e.entity_type, e.embedding_idx);
|
||||
m.entities.push((id, i));
|
||||
id
|
||||
}
|
||||
};
|
||||
m.entity_ids.insert(e.id, id);
|
||||
}
|
||||
let known_relations: HashSet<(u64, u64, String)> = kg
|
||||
.relations
|
||||
.iter()
|
||||
.map(|r| (r.src, r.tgt, r.relation.clone()))
|
||||
.collect();
|
||||
for (i, r) in data.relations.iter().enumerate() {
|
||||
let (Some(&src), Some(&tgt)) = (m.entity_ids.get(&r.src), m.entity_ids.get(&r.tgt)) else {
|
||||
m.dangling_relations.push(i);
|
||||
continue;
|
||||
};
|
||||
if known_relations.contains(&(src, tgt, r.relation.clone())) {
|
||||
m.relations_present += 1;
|
||||
continue;
|
||||
}
|
||||
let at = kg.relations.len();
|
||||
kg.add_relation(src, tgt, &r.relation, r.weight as f32);
|
||||
kg.relations[at].ts = r.timestamp * US_PER_SEC;
|
||||
m.relations.push((at, i));
|
||||
}
|
||||
|
||||
// ---- Write: one checkpoint for records, sessions and graph -----------
|
||||
let entries: Vec<MemoryEntry> = to_write
|
||||
.iter()
|
||||
.map(|c| MemoryEntry {
|
||||
chunk: c.chunk.clone(),
|
||||
embedding: c.embedding.clone(),
|
||||
source_channel: c.source_channel.clone(),
|
||||
timestamp: c.timestamp,
|
||||
session_id: c.session_id.clone(),
|
||||
tags: c.tags.clone(),
|
||||
})
|
||||
.collect();
|
||||
let indices = mem.save_batch(entries)?;
|
||||
m.records = indices
|
||||
.iter()
|
||||
.copied()
|
||||
.zip(new_chunks.iter().copied())
|
||||
.collect();
|
||||
|
||||
// Rows deleted in the source stay deleted: tombstones, as the agent's own
|
||||
// `delete` leaves them (not compacted away).
|
||||
// Records matched in the store whose source row has since been deleted
|
||||
// are tombstoned too.
|
||||
let tombstones: Vec<usize> = m
|
||||
.records
|
||||
.iter()
|
||||
.filter(|&&(_, src)| data.chunks[src].deleted != 0)
|
||||
.map(|&(idx, _)| idx)
|
||||
.chain(delete_in_store)
|
||||
.collect();
|
||||
mem.delete_batch(&tombstones)?;
|
||||
|
||||
m.anomaly_alerts = mem
|
||||
.take_anomaly_alerts()
|
||||
.into_iter()
|
||||
.map(|a| a.message)
|
||||
.collect();
|
||||
m.store_count = mem.count();
|
||||
drop(mem); // release the single-writer lock before anyone re-opens it
|
||||
Ok(m)
|
||||
}
|
||||
@@ -1,192 +1,266 @@
|
||||
use clawhdf5::reader::File as Hdf5File;
|
||||
use clawhdf5_format::provenance::VerifyResult;
|
||||
//! Validate a migration by reading the store back the way an agent would:
|
||||
//! through `HDF5Memory::open_read_only`, comparing what it loads with the
|
||||
//! SQLite source, and running a search for a migrated record.
|
||||
|
||||
use std::path::Path;
|
||||
|
||||
use clawhdf5_agent::{AgentMemory, HDF5Memory, SearchOptions};
|
||||
use clawhdf5_format::float16::round_to_f16;
|
||||
|
||||
use crate::hdf5_reader::read_hdf5;
|
||||
use crate::sqlite_reader::SqliteData;
|
||||
use crate::store_writer::{Migration, US_PER_SEC};
|
||||
|
||||
type BoxErr = Box<dyn std::error::Error>;
|
||||
|
||||
/// Summary of a migration validation.
|
||||
#[derive(Debug)]
|
||||
pub struct ValidationSummary {
|
||||
pub chunks: u64,
|
||||
pub sessions: u64,
|
||||
pub entities: u64,
|
||||
pub relations: u64,
|
||||
pub embedding_dim: u64,
|
||||
/// Number of rows whose full content was compared against the source.
|
||||
/// Records in the store (including tombstones).
|
||||
pub count: usize,
|
||||
/// Records in the store that are not deleted.
|
||||
pub active: usize,
|
||||
pub sessions: usize,
|
||||
pub entities: usize,
|
||||
pub relations: usize,
|
||||
pub embedding_dim: usize,
|
||||
pub float16: bool,
|
||||
/// Rows whose full content was compared against the source.
|
||||
pub rows_checked: u64,
|
||||
/// Whether the `chunks/text` and `chunks/embeddings` SHINES provenance
|
||||
/// hashes (written via [`crate::hdf5_writer`]) were both present and
|
||||
/// matched their recomputed SHA-256 on read-back. `false` when either
|
||||
/// dataset has no provenance metadata (e.g. an older output file) or
|
||||
/// there are zero chunks to check.
|
||||
pub provenance_verified: bool,
|
||||
/// Whether a search for a migrated record found it (`false` when there
|
||||
/// was no active migrated record with an embedding to search for).
|
||||
pub search_checked: bool,
|
||||
}
|
||||
|
||||
/// Validate a migrated HDF5 file against the source data.
|
||||
/// Validate the store at `path` against the source rows `migration` wrote.
|
||||
///
|
||||
/// Reads the written file back and compares actual content — chunk text,
|
||||
/// embeddings, and every session/entity/relation field — to the source, not
|
||||
/// just the row counts. When `full` is false a representative sample of chunk
|
||||
/// rows is content-checked (counts and all other groups are always checked in
|
||||
/// full); when `full` is true every chunk row is compared too. `float16` widens
|
||||
/// the embedding tolerance to allow for half-precision quantization.
|
||||
pub fn validate_hdf5(
|
||||
path: &str,
|
||||
/// Counts and the session / entity / relation rows are always checked in
|
||||
/// full. Memory records are content-checked on a representative sample, or
|
||||
/// all of them with `full`. Embeddings must match exactly: the source values
|
||||
/// themselves in an `f32` store, their [`round_to_f16`] in a `float16` one.
|
||||
pub fn validate_store(
|
||||
path: &Path,
|
||||
source: &SqliteData,
|
||||
migration: &Migration,
|
||||
full: bool,
|
||||
float16: bool,
|
||||
) -> Result<ValidationSummary, BoxErr> {
|
||||
let got = read_hdf5(path)?;
|
||||
let provenance_verified = verify_chunk_provenance(path)?;
|
||||
let mut mem = HDF5Memory::open_read_only(path)?;
|
||||
let float16 = mem.config().float16;
|
||||
let dim = mem.config().embedding_dim;
|
||||
|
||||
// ---- Counts ----
|
||||
check_count("chunk", got.chunks.len(), source.chunks.len())?;
|
||||
check_count("session", got.sessions.len(), source.sessions.len())?;
|
||||
check_count("entity", got.entities.len(), source.entities.len())?;
|
||||
check_count("relation", got.relations.len(), source.relations.len())?;
|
||||
if got.embedding_dim != source.embedding_dim {
|
||||
check_count("record", mem.count(), migration.store_count)?;
|
||||
if float16 != migration.float16 {
|
||||
return Err(format!(
|
||||
"embedding_dim mismatch: HDF5 has {}, source has {}",
|
||||
got.embedding_dim, source.embedding_dim
|
||||
"float16 mismatch: store {float16}, expected {}",
|
||||
migration.float16
|
||||
)
|
||||
.into());
|
||||
}
|
||||
if dim != migration.embedding_dim {
|
||||
return Err(format!(
|
||||
"embedding_dim mismatch: store has {dim}, expected {}",
|
||||
migration.embedding_dim
|
||||
)
|
||||
.into());
|
||||
}
|
||||
if !migration.appended_to_existing {
|
||||
check_count("record", mem.count(), migration.records.len())?;
|
||||
check_count("session", mem.sessions().len(), migration.sessions.len())?;
|
||||
check_count(
|
||||
"entity",
|
||||
mem.knowledge().entities.len(),
|
||||
migration.entities.len(),
|
||||
)?;
|
||||
check_count(
|
||||
"relation",
|
||||
mem.knowledge().relations.len(),
|
||||
migration.relations.len(),
|
||||
)?;
|
||||
}
|
||||
|
||||
// ---- Chunk content (sampled or full) ----
|
||||
let (emb_abs, emb_rel) = if float16 { (1e-2, 1e-2) } else { (1e-4, 0.0) };
|
||||
// ---- Memory records (sampled or full) ----
|
||||
let mut rows_checked = 0u64;
|
||||
for i in sample_indices(source.chunks.len(), full) {
|
||||
let (s, g) = (&source.chunks[i], &got.chunks[i]);
|
||||
if s.id != g.id {
|
||||
return Err(field_err("chunk", i, "id", s.id, g.id));
|
||||
let expected_value = |v: f32| if float16 { round_to_f16(v) } else { v };
|
||||
for k in sample_indices(migration.records.len(), full) {
|
||||
let (idx, src) = migration.records[k];
|
||||
let s = &source.chunks[src];
|
||||
let c = &mem.cache;
|
||||
if idx >= c.len() {
|
||||
return Err(
|
||||
format!("record {idx} (chunk id {}) is missing from the store", s.id).into(),
|
||||
);
|
||||
}
|
||||
if s.chunk != g.chunk {
|
||||
let id = s.id;
|
||||
if c.chunks[idx] != s.chunk {
|
||||
return Err(format!(
|
||||
"chunk[{i}].text mismatch: source {:?}, HDF5 {:?}",
|
||||
"record {idx} (chunk id {id}) text mismatch: source {:?}, store {:?}",
|
||||
truncate(&s.chunk),
|
||||
truncate(&g.chunk)
|
||||
truncate(&c.chunks[idx])
|
||||
)
|
||||
.into());
|
||||
}
|
||||
if s.session_id != g.session_id || s.source_channel != g.source_channel || s.tags != g.tags
|
||||
if c.source_channels[idx] != s.source_channel
|
||||
|| c.session_ids[idx] != s.session_id
|
||||
|| c.tags[idx] != s.tags
|
||||
{
|
||||
return Err(format!("chunk[{i}] string field mismatch").into());
|
||||
return Err(format!("record {idx} (chunk id {id}) string field mismatch").into());
|
||||
}
|
||||
if s.deleted != g.deleted {
|
||||
return Err(field_err("chunk", i, "deleted", s.deleted, g.deleted));
|
||||
}
|
||||
if s.embedding.len() != g.embedding.len() {
|
||||
if c.timestamps[idx].to_bits() != s.timestamp.to_bits() {
|
||||
return Err(format!(
|
||||
"chunk[{i}] embedding length mismatch: {} vs {}",
|
||||
s.embedding.len(),
|
||||
g.embedding.len()
|
||||
"record {idx} (chunk id {id}) timestamp mismatch: source {}, store {}",
|
||||
s.timestamp, c.timestamps[idx]
|
||||
)
|
||||
.into());
|
||||
}
|
||||
for (k, (&a, &b)) in s.embedding.iter().zip(g.embedding.iter()).enumerate() {
|
||||
if (a - b).abs() > emb_abs + emb_rel * a.abs() {
|
||||
return Err(
|
||||
format!("chunk[{i}].embedding[{k}] mismatch: source {a}, HDF5 {b}").into(),
|
||||
);
|
||||
let deleted = c.tombstones[idx] != 0;
|
||||
if deleted != (s.deleted != 0) {
|
||||
return Err(format!(
|
||||
"record {idx} (chunk id {id}) deleted mismatch: source {}, store {deleted}",
|
||||
s.deleted != 0
|
||||
)
|
||||
.into());
|
||||
}
|
||||
let got = c.embeddings.get(idx).unwrap_or(&[]);
|
||||
if got.len() != s.embedding.len() {
|
||||
return Err(format!(
|
||||
"record {idx} (chunk id {id}) embedding length mismatch: source {}, store {}",
|
||||
s.embedding.len(),
|
||||
got.len()
|
||||
)
|
||||
.into());
|
||||
}
|
||||
for (j, (&a, &b)) in s.embedding.iter().zip(got).enumerate() {
|
||||
let want = expected_value(a);
|
||||
if want.to_bits() != b.to_bits() && !(want.is_nan() && b.is_nan()) {
|
||||
return Err(format!(
|
||||
"record {idx} (chunk id {id}) embedding[{j}] mismatch: source {a}, \
|
||||
expected {want}, store {b}"
|
||||
)
|
||||
.into());
|
||||
}
|
||||
}
|
||||
rows_checked += 1;
|
||||
}
|
||||
|
||||
// ---- Other groups (always full — they are small) ----
|
||||
for (i, (s, g)) in source.sessions.iter().zip(got.sessions.iter()).enumerate() {
|
||||
if s.id != g.id
|
||||
|| s.start_idx != g.start_idx
|
||||
|| s.end_idx != g.end_idx
|
||||
|| s.channel != g.channel
|
||||
|| s.summary != g.summary
|
||||
{
|
||||
return Err(format!("session[{i}] mismatch").into());
|
||||
// ---- Records tombstoned because their source row was deleted ----
|
||||
for &(idx, src) in &migration.deleted_in_store {
|
||||
let s = &source.chunks[src];
|
||||
let c = &mem.cache;
|
||||
if idx >= c.len() || c.chunks[idx] != s.chunk || c.timestamps[idx] != s.timestamp {
|
||||
return Err(format!("record {idx} (chunk id {}) mismatch or missing", s.id).into());
|
||||
}
|
||||
rows_checked += 1;
|
||||
}
|
||||
for (i, (s, g)) in source.entities.iter().zip(got.entities.iter()).enumerate() {
|
||||
if s.id != g.id
|
||||
|| s.name != g.name
|
||||
|| s.entity_type != g.entity_type
|
||||
|| s.embedding_idx != g.embedding_idx
|
||||
{
|
||||
return Err(format!("entity[{i}] mismatch").into());
|
||||
}
|
||||
rows_checked += 1;
|
||||
}
|
||||
for (i, (s, g)) in source
|
||||
.relations
|
||||
.iter()
|
||||
.zip(got.relations.iter())
|
||||
.enumerate()
|
||||
{
|
||||
if s.src != g.src || s.tgt != g.tgt || s.relation != g.relation {
|
||||
return Err(format!("relation[{i}] mismatch").into());
|
||||
if c.tombstones[idx] == 0 {
|
||||
return Err(format!(
|
||||
"record {idx} (chunk id {}) is deleted in the source but active in the store",
|
||||
s.id
|
||||
)
|
||||
.into());
|
||||
}
|
||||
rows_checked += 1;
|
||||
}
|
||||
|
||||
// ---- Sessions ----
|
||||
let sessions = mem.sessions();
|
||||
for &(at, src) in &migration.sessions {
|
||||
let s = &source.sessions[src];
|
||||
let (Some(e), Some(summary)) = (sessions.entries.get(at), sessions.summaries.get(at))
|
||||
else {
|
||||
return Err(format!("session {:?} is missing from the store", s.id).into());
|
||||
};
|
||||
if e.id != s.id
|
||||
|| e.start_idx != s.start_idx.max(0) as u64
|
||||
|| e.end_idx != s.end_idx.max(0) as u64
|
||||
|| e.channel != s.channel
|
||||
|| *summary != s.summary
|
||||
|| e.ts != s.timestamp * US_PER_SEC
|
||||
{
|
||||
return Err(format!("session {:?} mismatch", s.id).into());
|
||||
}
|
||||
rows_checked += 1;
|
||||
}
|
||||
|
||||
// ---- Knowledge graph ----
|
||||
let kg = mem.knowledge();
|
||||
for &(id, src) in &migration.entities {
|
||||
let s = &source.entities[src];
|
||||
let Some(e) = kg.get_entity(id) else {
|
||||
return Err(format!(
|
||||
"entity {:?} (id {}) is missing from the store",
|
||||
s.name, s.id
|
||||
)
|
||||
.into());
|
||||
};
|
||||
if e.name != s.name || e.entity_type != s.entity_type || e.embedding_idx != s.embedding_idx
|
||||
{
|
||||
return Err(format!("entity {:?} (id {}) mismatch", s.name, s.id).into());
|
||||
}
|
||||
rows_checked += 1;
|
||||
}
|
||||
for &(at, src) in &migration.relations {
|
||||
let s = &source.relations[src];
|
||||
let r = kg.relations.get(at);
|
||||
let ok = r.is_some_and(|r| {
|
||||
Some(&r.src) == migration.entity_ids.get(&s.src)
|
||||
&& Some(&r.tgt) == migration.entity_ids.get(&s.tgt)
|
||||
&& r.relation == s.relation
|
||||
&& r.weight == s.weight as f32
|
||||
&& r.ts == s.timestamp * US_PER_SEC
|
||||
});
|
||||
if !ok {
|
||||
return Err(format!(
|
||||
"relation {} -[{}]-> {} mismatch or missing",
|
||||
s.src, s.relation, s.tgt
|
||||
)
|
||||
.into());
|
||||
}
|
||||
rows_checked += 1;
|
||||
}
|
||||
|
||||
// ---- A migrated record must be findable by search ----
|
||||
let probe = migration
|
||||
.records
|
||||
.iter()
|
||||
.copied()
|
||||
.find(|&(idx, _)| dim > 0 && mem.cache.tombstones[idx] == 0);
|
||||
let search_checked = match probe {
|
||||
None => false,
|
||||
Some((idx, _)) => {
|
||||
let query = mem.cache.embeddings[idx].to_vec();
|
||||
let text = mem.cache.chunks[idx].clone();
|
||||
let hits = mem.search(&query, &text, &SearchOptions::new(10));
|
||||
// A record with the same text is as good a hit: the source may
|
||||
// hold duplicates, and they tie.
|
||||
if !hits.iter().any(|h| h.index == idx || h.chunk == text) {
|
||||
return Err(format!(
|
||||
"search for migrated record {idx} ({:?}) did not return it",
|
||||
truncate(&text)
|
||||
)
|
||||
.into());
|
||||
}
|
||||
true
|
||||
}
|
||||
};
|
||||
|
||||
Ok(ValidationSummary {
|
||||
chunks: got.chunks.len() as u64,
|
||||
sessions: got.sessions.len() as u64,
|
||||
entities: got.entities.len() as u64,
|
||||
relations: got.relations.len() as u64,
|
||||
embedding_dim: got.embedding_dim as u64,
|
||||
count: mem.count(),
|
||||
active: mem.count_active(),
|
||||
sessions: mem.sessions().len(),
|
||||
entities: mem.knowledge().entities.len(),
|
||||
relations: mem.knowledge().relations.len(),
|
||||
embedding_dim: dim,
|
||||
float16,
|
||||
rows_checked,
|
||||
provenance_verified,
|
||||
search_checked,
|
||||
})
|
||||
}
|
||||
|
||||
fn check_count(kind: &str, got: usize, expected: usize) -> Result<(), BoxErr> {
|
||||
if got != expected {
|
||||
return Err(format!("{kind} count mismatch: HDF5 has {got}, source has {expected}").into());
|
||||
return Err(format!("{kind} count mismatch: store has {got}, expected {expected}").into());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Re-verify the SHA-256 provenance hash of `chunks/text` and
|
||||
/// `chunks/embeddings` against their actual stored bytes, catching
|
||||
/// post-write corruption that a plain content comparison against the
|
||||
/// in-memory source wouldn't (the source is compared against what
|
||||
/// `read_hdf5` decoded, not against the raw bytes on disk).
|
||||
///
|
||||
/// Returns `Ok(true)` only if both datasets exist and both hashes match.
|
||||
/// Returns `Ok(false)` (not an error) if a dataset has no provenance
|
||||
/// attributes at all (e.g. a file written before this check existed) or
|
||||
/// there are zero chunks. Returns an error only on an actual hash mismatch —
|
||||
/// that indicates real corruption.
|
||||
fn verify_chunk_provenance(path: &str) -> Result<bool, BoxErr> {
|
||||
let file = Hdf5File::open(path)?;
|
||||
let Ok(chunks) = file.group("chunks") else {
|
||||
return Ok(false);
|
||||
};
|
||||
let mut all_present = true;
|
||||
for name in ["text", "embeddings"] {
|
||||
let Ok(ds) = chunks.dataset(name) else {
|
||||
all_present = false;
|
||||
continue;
|
||||
};
|
||||
match ds.verify_provenance()? {
|
||||
VerifyResult::Ok => {}
|
||||
VerifyResult::NoHash => all_present = false,
|
||||
VerifyResult::Mismatch { stored, computed } => {
|
||||
return Err(format!(
|
||||
"provenance hash mismatch on chunks/{name}: stored {stored}, recomputed {computed} — data may be corrupted"
|
||||
)
|
||||
.into());
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(all_present)
|
||||
}
|
||||
|
||||
fn field_err<T: std::fmt::Display>(kind: &str, i: usize, field: &str, s: T, g: T) -> BoxErr {
|
||||
format!("{kind}[{i}].{field} mismatch: source {s}, HDF5 {g}").into()
|
||||
}
|
||||
|
||||
fn truncate(s: &str) -> String {
|
||||
if s.len() <= 40 {
|
||||
s.to_string()
|
||||
@@ -196,7 +270,7 @@ fn truncate(s: &str) -> String {
|
||||
}
|
||||
}
|
||||
|
||||
/// Indices of chunk rows to content-check. Full = all; otherwise a spread of
|
||||
/// Indices of records to content-check. Full = all; otherwise a spread of
|
||||
/// representative rows (first/last and evenly-spaced interior samples).
|
||||
fn sample_indices(n: usize, full: bool) -> Vec<usize> {
|
||||
if n == 0 {
|
||||
|
||||
Reference in New Issue
Block a user