Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
73bb068264 |
+139
-681
File diff suppressed because it is too large
Load Diff
+7
-93
@@ -3,17 +3,6 @@
|
|||||||
## Unreleased
|
## Unreleased
|
||||||
|
|
||||||
### Upgrade Notes
|
### 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`
|
- **Files written by clawhdf5 now open in h5py and libhdf5.** Every `f32`
|
||||||
dataset we wrote — including every agent store's embeddings — was refused
|
dataset we wrote — including every agent store's embeddings — was refused
|
||||||
with "sign bit position out of bounds", and every empty dataset with
|
with "sign bit position out of bounds", and every empty dataset with
|
||||||
@@ -22,18 +11,12 @@
|
|||||||
each checkpoint, so it becomes readable at its next checkpoint on this
|
each checkpoint, so it becomes readable at its next checkpoint on this
|
||||||
version; other files with `f32` or empty datasets need rewriting. Details in
|
version; other files with `f32` or empty datasets need rewriting. Details in
|
||||||
`docs/known-issues.md`.
|
`docs/known-issues.md`.
|
||||||
- **New stores store embeddings as half precision by default.**
|
- **`MemoryConfig::float16` now does what it says.** It was persisted and
|
||||||
`MemoryConfig::float16` was persisted and otherwise ignored; it now writes
|
otherwise ignored; embeddings were always stored as `f32`. A store created
|
||||||
`float16` embeddings (48% smaller files at 100K) and rounds each embedding
|
with it on now writes half-precision embeddings (48% smaller files) and
|
||||||
to half precision as it is saved — and it defaults to `true` for new
|
rounds embeddings to half precision as they are saved. A store that already
|
||||||
stores. On the full LongMemEval haystack with real MiniLM embeddings every
|
had `float16 = true` rounds its embeddings when next opened and writes them
|
||||||
retrieval metric matched `f32`. **Existing stores are unaffected**: every
|
as `float16` at its next checkpoint. Off by default.
|
||||||
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
|
- **Breaking:** `MemoryError` gained `InvalidEntry`, returned when a
|
||||||
`float16` store is given an embedding value beyond ±65504. Exhaustive
|
`float16` store is given an embedding value beyond ±65504. Exhaustive
|
||||||
matches need the new arm.
|
matches need the new arm.
|
||||||
@@ -57,60 +40,6 @@
|
|||||||
`quantized_index = false`, or pass `create --f32-index` to the CLI, to opt
|
`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.
|
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
|
### Interop
|
||||||
- `clawhdf5-format`: **every `f32` dataset was unreadable by h5py and
|
- `clawhdf5-format`: **every `f32` dataset was unreadable by h5py and
|
||||||
libhdf5.** The float datatype encoder hard-coded the sign bit's position to
|
libhdf5.** The float datatype encoder hard-coded the sign bit's position to
|
||||||
@@ -138,9 +67,7 @@
|
|||||||
precision.** At 100K x 384 the file goes from 154.0 to 80.8 MiB (−48%), a
|
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
|
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
|
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. On the full
|
`hybrid_search` latency; at 10K open is 3 ms slower. The cache rounds each
|
||||||
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
|
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
|
returns the same results before and after a reopen (tested). Out-of-range
|
||||||
values are refused with `MemoryError::InvalidEntry` rather than stored as
|
values are refused with `MemoryError::InvalidEntry` rather than stored as
|
||||||
@@ -172,11 +99,6 @@
|
|||||||
now an error.
|
now an error.
|
||||||
|
|
||||||
### Defaults
|
### 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
|
- `clawhdf5-agent`: `MemoryConfig::quantized_index` defaults to `true` for new
|
||||||
stores. The reason it had been off — that int8 search was slower on ARM —
|
stores. The reason it had been off — that int8 search was slower on ARM —
|
||||||
did not survive measurement (see Corrections). Stores that predate the
|
did not survive measurement (see Corrections). Stores that predate the
|
||||||
@@ -189,14 +111,6 @@
|
|||||||
knew to ask; it now only ever switches the default off.
|
knew to ask; it now only ever switches the default off.
|
||||||
|
|
||||||
### Performance
|
### 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
|
- `clawhdf5-format`, `clawhdf5-filters`: both deflate paths hand the codec the
|
||||||
whole chunk in one call, into a buffer allocated once, instead of streaming
|
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
|
it through a 32 KiB buffer: about 5% on chunked writes and 10% on zlib-ng's
|
||||||
|
|||||||
@@ -87,26 +87,16 @@ 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
|
`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
|
`<store>.h5.wal.corrupt-<ts>` rather than blocking `open()`; a WAL with an
|
||||||
unknown *newer* version still fails and is left untouched.
|
unknown *newer* version still fails and is left untouched.
|
||||||
- `MemoryConfig::float16` (**on by default** for new stores, persisted;
|
- `MemoryConfig::float16` (off by default, persisted; CLI `create --float16`)
|
||||||
existing stores keep their recorded `false` — guarded by the v2.5.0
|
writes `/memory/embeddings` as IEEE half precision (48% smaller file at
|
||||||
fixture in `tests/float16_store.rs`; CLI opt-out is `create --f32`) writes
|
100K, same recall). `MemoryCache::half_precision` rounds each embedding as
|
||||||
`/memory/embeddings` as IEEE half precision (48% smaller file at 100K;
|
it enters the cache (push, update, WAL replay, and on load of a store still
|
||||||
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
|
`f32` on disk), so memory and file agree bit for bit; the conversions live
|
||||||
in `clawhdf5_format::float16` and must stay the single implementation.
|
in `clawhdf5_format::float16` and must stay the single implementation.
|
||||||
Values beyond ±65504 are `MemoryError::InvalidEntry`. Interop: every file
|
Values beyond ±65504 are `MemoryError::InvalidEntry`. Interop: every file
|
||||||
must open in h5py — `f32` datasets and empty datasets did not until
|
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
|
2026-09-23 (see `docs/known-issues.md`); the agent's `h5py_interop` test
|
||||||
guards a whole store.
|
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
|
- `MemoryConfig::compression` is off by default; when on, embeddings are
|
||||||
deflate-compressed, or Zstd with the agent's `zstd` feature (links libzstd).
|
deflate-compressed, or Zstd with the agent's `zstd` feature (links libzstd).
|
||||||
- `Dataset::verify_provenance()` (clawhdf5 facade, `provenance` feature, on by
|
- `Dataset::verify_provenance()` (clawhdf5 facade, `provenance` feature, on by
|
||||||
|
|||||||
@@ -6,7 +6,7 @@
|
|||||||
[](https://www.rust-lang.org)
|
[](https://www.rust-lang.org)
|
||||||
[](#building)
|
[](#building)
|
||||||
[](BENCHMARKS.md#longmemeval-results)
|
[](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.
|
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,18 +82,14 @@ breaking change, are in [CHANGELOG.md](CHANGELOG.md).
|
|||||||
100K × 384 store to 1.74× the raw vectors. At equal recall it is also faster
|
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`).
|
than `f32`: 1.63× QPS on AVX2, 1.18× on a Raspberry Pi 5 (NEON `SDOT`).
|
||||||
|
|
||||||
**Interop and search (unreleased)**
|
**Interop (unreleased)**
|
||||||
- **Files we write now open in h5py and libhdf5.** Every `f32` dataset —
|
- **Files we write now open in h5py and libhdf5.** Every `f32` dataset —
|
||||||
including every agent store's embeddings — and every empty dataset was
|
including every agent store's embeddings — and every empty dataset was
|
||||||
refused by libhdf5. Both were write-side bugs in every release; agent stores
|
refused by libhdf5. Both were write-side bugs in every release; agent stores
|
||||||
fix themselves at their next checkpoint. See
|
fix themselves at their next checkpoint. See
|
||||||
[docs/known-issues.md](docs/known-issues.md).
|
[docs/known-issues.md](docs/known-issues.md).
|
||||||
- `MemoryConfig::float16` now stores half-precision embeddings (it was
|
- `MemoryConfig::float16` now stores half-precision embeddings (it was
|
||||||
ignored), and is on by default for new stores: 48% smaller files, and
|
ignored): 48% smaller files at the same recall.
|
||||||
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**
|
**Tooling**
|
||||||
- CI now runs the h5py/netCDF4 interop suites for real (they had been skipping
|
- CI now runs the h5py/netCDF4 interop suites for real (they had been skipping
|
||||||
@@ -111,8 +107,8 @@ Every AI agent needs memory. Today that means scattered Markdown files, SQLite d
|
|||||||
| Keyword search | Separate FTS engine | Integrated BM25 |
|
| Keyword search | Separate FTS engine | Integrated BM25 |
|
||||||
| Knowledge graph | Neo4j or none | In-file graph with spreading activation |
|
| Knowledge graph | Neo4j or none | In-file graph with spreading activation |
|
||||||
| Memory consolidation | Manual pruning | Hippocampal-inspired automatic tiers |
|
| Memory consolidation | Manual pruning | Hippocampal-inspired automatic tiers |
|
||||||
| Temporal queries | Custom code | Native temporal index (622 ns range query over 10K) |
|
| Temporal queries | Custom code | Native temporal index (716ns) |
|
||||||
| Multi-modal | Multiple stores | Unified cross-modal search (exact scan: 842 µs over 1K records) |
|
| Multi-modal | Multiple stores | Unified cross-modal search |
|
||||||
| Integrity | Hope for the best | Chained-CRC WAL, checksummed chunk indexes, write-anomaly alerts, opt-in SHA-256 dataset provenance |
|
| 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.** |
|
| Portability | Config + DB + files | **One `.h5` file. Copy it anywhere.** |
|
||||||
|
|
||||||
@@ -120,7 +116,7 @@ Every AI agent needs memory. Today that means scattered Markdown files, SQLite d
|
|||||||
|
|
||||||
## Performance
|
## Performance
|
||||||
|
|
||||||
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.
|
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).
|
||||||
|
|
||||||
### HDF5 Core I/O (vs libhdf5 1.14.6)
|
### HDF5 Core I/O (vs libhdf5 1.14.6)
|
||||||
|
|
||||||
@@ -154,45 +150,33 @@ and [§ Quantising the index copy](BENCHMARKS.md#quantising-the-index-copy-quant
|
|||||||
| `f32` index | 0.9945 | 13 399 | 3.2 s |
|
| `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** |
|
| `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. These
|
Before the v2.4.0 neighbour-selection fix, recall@10 at 100K was 0.31.
|
||||||
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, tank, 2026-09-24):
|
**Brute-force and IVF paths** (Criterion, i7-12650H):
|
||||||
|
|
||||||
| Scale | Flat | IVF (nprobe=10) | IVF-PQ | MemX¹ (claimed, end-to-end) |
|
| Scale | Flat | IVF (nprobe=10) | IVF-PQ | vs MemX¹ |
|
||||||
|-------|------|-----------------|--------|----------|
|
|-------|------|-----------------|--------|----------|
|
||||||
| 1K | **47.4 µs** | — | — | — |
|
| 1K | **54 µs** | — | — | — |
|
||||||
| 10K | 500.5 µs | **24.8 µs** | — | — |
|
| 10K | 753 µs | **27 µs** | — | — |
|
||||||
| 100K | 6.58 ms | 592 µs | **869 µs** | <90 ms |
|
| 100K | 11.4 ms | 1.32 ms | **1.19 ms** | ~8–76× (see caveat) |
|
||||||
|
|
||||||
> These replace figures from the original i7-12650H run (flat 54 µs / 753 µs /
|
> Reproduced on the same second machine (Ryzen 7 7800X3D) with a corrected,
|
||||||
> 11.4 ms); a 2026-08-05 run on tank had already matched the new ones — see
|
> apples-to-apples SIMD/scalar/parallel comparison methodology — see
|
||||||
> [BENCHMARKS.md § Vector Search Latency](BENCHMARKS.md#vector-search-latency).
|
> [BENCHMARKS.md § Independent Validation: tank — LongMemEval & Vector
|
||||||
|
> Search](BENCHMARKS.md#independent-validation-tank--longmemeval--vector-search-ryzen-7-7800x3d-2026-08-05).
|
||||||
|
|
||||||
### Agent Memory Operations
|
### Agent Memory Operations
|
||||||
|
|
||||||
| Operation | Latency | Scale |
|
| Operation | Latency | Scale |
|
||||||
|-----------|---------|-------|
|
|-----------|---------|-------|
|
||||||
| Hybrid search (`HDF5Memory::hybrid_search`, p50) | **0.07 ms** / 0.49 ms / 4.69 ms | 1K / 10K / 100K records |
|
| Hybrid search (`HDF5Memory::hybrid_search`, p50) | **70 µs** / 0.49 ms / 4.65 ms | 1K / 10K / 100K records |
|
||||||
| BM25 keyword search | **20.4 µs** | 1K records |
|
| BM25 keyword search | **67 µs** | 1K records |
|
||||||
| Knowledge graph BFS | **23.1 µs** | 1K entities |
|
| Knowledge graph BFS | **24 µs** | 1K entities |
|
||||||
| Spreading activation | **10.1 µs** | 100 entities |
|
| Spreading activation | **17 µs** | 100 entities |
|
||||||
| Temporal range query | **622 ns** | 10K timestamps |
|
| Temporal range query | **716 ns** | 10K timestamps |
|
||||||
| Consolidation cycle | **115.2 µs** | 1K records |
|
| Consolidation cycle | **164 µs** | 1K records |
|
||||||
| Cross-modal search (exact scan, 2 embeddings per record) | **842.0 µs** / 8.44 ms | 1K / 10K records |
|
| Memory write (WAL) | **18 µs** | per record (group-commit append; HDF5 batched at flush) |
|
||||||
| Memory write (WAL) | **26.1 µs** | per record (group-commit append; HDF5 batched at flush) |
|
| Importance gate | **61 ns** | per record |
|
||||||
| 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)
|
### Chunked Write Throughput (codec comparison)
|
||||||
|
|
||||||
@@ -207,7 +191,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.
|
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), so the two columns are not comparable and no ratio is given. 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). The ratio overstates the real advantage by an unquantified margin — order-of-magnitude indication only. See [BENCHMARKS.md](BENCHMARKS.md#comparison-to-memx-arxiv260316171).
|
||||||
|
|
||||||
### LongMemEval Retrieval Recall
|
### LongMemEval Retrieval Recall
|
||||||
|
|
||||||
@@ -260,26 +244,17 @@ retrieval recall reported as QA accuracy typically overstates by 20–30 points.
|
|||||||
|
|
||||||
### Memory Footprint
|
### Memory Footprint
|
||||||
|
|
||||||
**On disk** — 384-dim `float16` embeddings (the default for new stores),
|
**On disk** — 384-dim embeddings, 200-char text
|
||||||
200-char text, `footprint_bench`
|
|
||||||
([BENCHMARKS.md § Memory Footprint](BENCHMARKS.md#memory-footprint-1)):
|
([BENCHMARKS.md § Memory Footprint](BENCHMARKS.md#memory-footprint-1)):
|
||||||
|
|
||||||
| Records | File Size | Bytes/Record | Gzip-6 compressed |
|
| Records | File Size | Bytes/Record | Gzip-6 compressed |
|
||||||
|---------|-----------|--------------|-------------------|
|
|---------|-----------|--------------|-------------------|
|
||||||
| 1K | 810.4 KB | 829 B | 56.4 KB |
|
| 1K | 1.7 MB | 1.8 KB | 277 KB (6.1x) |
|
||||||
| 10K | 7.8 MB | 820 B | 471.3 KB |
|
| 10K | 17.0 MB | 1.7 KB | 2.7 MB (6.2x) |
|
||||||
| 100K | 76.7 MB | 803 B | 4.5 MB |
|
| 100K | 169.8 MB | 1.7 KB | 26.9 MB (6.2x) |
|
||||||
|
|
||||||
The benchmark's synthetic embeddings and text are far more repetitive than
|
With `MemoryConfig::float16` the embeddings take half the space: an agent
|
||||||
real data (only 40 distinct texts), so no column here is an expectation for
|
store of 100K × 384 records is 80.8 MiB instead of 154.0.
|
||||||
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
|
**In memory** — a store reopened from disk, 384-dim `f32`, measured with a
|
||||||
counting allocator ([BENCHMARKS.md § Memory footprint](BENCHMARKS.md#memory-footprint)):
|
counting allocator ([BENCHMARKS.md § Memory footprint](BENCHMARKS.md#memory-footprint)):
|
||||||
@@ -291,8 +266,7 @@ counting allocator ([BENCHMARKS.md § Memory footprint](BENCHMARKS.md#memory-foo
|
|||||||
| 100K | 146 MiB | 399 MiB (2.72x) | **256 MiB (1.74x)** |
|
| 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
|
Down from 505 MiB (3.44x) at 100K before v2.6.0, when the cache held every
|
||||||
embedding twice. The `f32` column was re-measured on 2026-09-24 and reproduced
|
embedding twice.
|
||||||
exactly; the `i8` column was not re-run.
|
|
||||||
|
|
||||||
### Consolidation Efficiency
|
### Consolidation Efficiency
|
||||||
|
|
||||||
@@ -303,10 +277,7 @@ exactly; the `i8` column was not re-run.
|
|||||||
|--------|--------|-------|-------|
|
|--------|--------|-------|-------|
|
||||||
| Records in store | 1,000 | 100 | −90% |
|
| Records in store | 1,000 | 100 | −90% |
|
||||||
| Hit@1 recall (signal records) | 100% | 100% | no loss |
|
| Hit@1 recall (signal records) | 100% | 100% | no loss |
|
||||||
| Search latency (avg) | 2.22 ms | 0.24 ms | **9.3x faster** |
|
| Search latency | 2.75 ms | 0.31 ms | **8.8x 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)**
|
**Full benchmark details: [BENCHMARKS.md](BENCHMARKS.md)**
|
||||||
|
|
||||||
@@ -322,14 +293,12 @@ ClawhDF5's agent memory engine draws on 15+ recent papers on agentic memory syst
|
|||||||
└────────┬────────┘
|
└────────┬────────┘
|
||||||
│
|
│
|
||||||
┌─────────────────▼──────────────────┐
|
┌─────────────────▼──────────────────┐
|
||||||
│ HDF5Memory::search │
|
│ HDF5Memory::hybrid_search │
|
||||||
│ optional source-channel filter │
|
|
||||||
│ HNSW vector + BM25 keyword │
|
│ HNSW vector + BM25 keyword │
|
||||||
│ weighted fusion (0.4 / 0.6) │
|
│ weighted fusion (0.4 / 0.6) │
|
||||||
│ × √(Hebbian activation) │
|
│ × √(Hebbian activation) │
|
||||||
└─────────────────┬──────────────────┘
|
└─────────────────┬──────────────────┘
|
||||||
│ opt-in (SearchOptions);
|
│ OpenClaw backend adds:
|
||||||
│ the OpenClaw backend turns both on
|
|
||||||
┌─────────────────▼──────────────────┐
|
┌─────────────────▼──────────────────┐
|
||||||
│ Multi-factor re-ranking │
|
│ Multi-factor re-ranking │
|
||||||
│ relevance · recency · authority · │
|
│ relevance · recency · authority · │
|
||||||
@@ -365,8 +334,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 |
|
| **`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 |
|
| **`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 |
|
| **`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. Opt-in via `SearchOptions::with_rerank`; on in the OpenClaw backend |
|
| **`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. Opt-in via `SearchOptions::with_confidence`; on in the OpenClaw backend |
|
| **`confidence`** | Low-confidence rejection — suppresses spurious recalls when nothing matches (OpenClaw backend) |
|
||||||
| **`temporal`** | Sorted timestamp index, session DAG, entity timeline, temporal query hints |
|
| **`temporal`** | Sorted timestamp index, session DAG, entity timeline, temporal query hints |
|
||||||
| **`multimodal`** | Cross-modal search across text/image/audio/video embeddings |
|
| **`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) |
|
| **`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) |
|
||||||
@@ -432,31 +401,6 @@ 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
|
### Knowledge Graph
|
||||||
|
|
||||||
```rust
|
```rust
|
||||||
@@ -636,14 +580,12 @@ setting existed keep their `f32` index; opt out for new stores with
|
|||||||
`quantized_index = false` or `clawhdf5-cli create --f32-index`. See
|
`quantized_index = false` or `clawhdf5-cli create --f32-index`. See
|
||||||
[BENCHMARKS.md § Quantising the index copy](BENCHMARKS.md#quantising-the-index-copy-quantized_index).
|
[BENCHMARKS.md § Quantising the index copy](BENCHMARKS.md#quantising-the-index-copy-quantized_index).
|
||||||
|
|
||||||
`MemoryConfig::float16` (**on by default** for new stores) stores the
|
`MemoryConfig::float16` (off by default; CLI `create --float16`) stores the
|
||||||
embeddings on disk as IEEE half precision (numpy `float16`): at 100K × 384 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 on the
|
file drops from 154 to 81 MiB, checkpoints and opens get faster, and vector
|
||||||
full LongMemEval haystack with real MiniLM embeddings every retrieval metric
|
recall and search latency do not change. Embeddings are rounded as they are
|
||||||
matches `f32`. Embeddings are rounded as they are saved, so the store searches
|
saved, so the store searches the same before and after a reopen; values must
|
||||||
the same before and after a reopen; values must lie within ±65504. Existing
|
lie within ±65504. See
|
||||||
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).
|
[BENCHMARKS.md § float16 embedding storage](BENCHMARKS.md#float16-embedding-storage-memoryconfigfloat16).
|
||||||
|
|
||||||
### `clawhdf5-format`
|
### `clawhdf5-format`
|
||||||
@@ -778,37 +720,9 @@ Replace in `Cargo.toml` and source:
|
|||||||
|
|
||||||
```bash
|
```bash
|
||||||
cargo install --path crates/clawhdf5-migrate
|
cargo install --path crates/clawhdf5-migrate
|
||||||
clawhdf5-migrate --sqlite old.db --hdf5 memory.h5 --agent-id my-agent --embedder minilm
|
clawhdf5-migrate --sqlite old.db --hdf5 memory.h5 --agent-id my-agent --embedding-dim 384
|
||||||
```
|
```
|
||||||
|
|
||||||
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
|
## Roadmap
|
||||||
|
|||||||
@@ -45,10 +45,6 @@ harness = false
|
|||||||
name = "memory_bench"
|
name = "memory_bench"
|
||||||
harness = false
|
harness = false
|
||||||
|
|
||||||
[[bench]]
|
|
||||||
name = "multimodal_bench"
|
|
||||||
harness = false
|
|
||||||
|
|
||||||
[features]
|
[features]
|
||||||
default = ["float16", "hnsw", "parallel"]
|
default = ["float16", "hnsw", "parallel"]
|
||||||
float16 = ["half"]
|
float16 = ["half"]
|
||||||
|
|||||||
@@ -1,107 +0,0 @@
|
|||||||
//! 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,32 +65,29 @@ pub fn hybrid_search_fused(
|
|||||||
) -> Vec<(usize, f32)> {
|
) -> Vec<(usize, f32)> {
|
||||||
// Get raw scores from both systems. Request all results so normalization
|
// Get raw scores from both systems. Request all results so normalization
|
||||||
// covers the full distribution.
|
// covers the full distribution.
|
||||||
let vec_scores = exact_vector_scores(query_embedding, vectors, tombstones);
|
// Use parallel search when rayon feature is enabled and vector count > 10K.
|
||||||
let kw_scores = bm25_index.scores(query_text);
|
let vec_scores = {
|
||||||
|
|
||||||
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")]
|
#[cfg(feature = "parallel")]
|
||||||
{
|
{
|
||||||
if vectors.count() > 10_000 {
|
if vectors.count() > 10_000 {
|
||||||
return vector_search::parallel_cosine_batch(
|
vector_search::parallel_cosine_batch(
|
||||||
query_embedding,
|
query_embedding,
|
||||||
vectors,
|
vectors,
|
||||||
skip,
|
tombstones,
|
||||||
vectors.count(),
|
vectors.count(),
|
||||||
);
|
)
|
||||||
|
} else {
|
||||||
|
vector_search::cosine_similarity_batch(query_embedding, vectors, tombstones)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
vector_search::cosine_similarity_batch(query_embedding, vectors, skip)
|
#[cfg(not(feature = "parallel"))]
|
||||||
|
{
|
||||||
|
vector_search::cosine_similarity_batch(query_embedding, vectors, tombstones)
|
||||||
|
}
|
||||||
|
};
|
||||||
|
let kw_scores = bm25_index.scores(query_text);
|
||||||
|
|
||||||
|
fuse(vec_scores, kw_scores, fusion, k)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Merge pre-computed vector-similarity and keyword scores into a single ranking.
|
/// Merge pre-computed vector-similarity and keyword scores into a single ranking.
|
||||||
|
|||||||
@@ -163,13 +163,12 @@ fn levenshtein(a: &str, b: &str) -> usize {
|
|||||||
/// entities-slice-index map, and an entity-id -> relation-indices map (edges
|
/// entities-slice-index map, and an entity-id -> relation-indices map (edges
|
||||||
/// touching that entity as either source or target).
|
/// touching that entity as either source or target).
|
||||||
///
|
///
|
||||||
/// Cached on `KnowledgeCache` and checked against a fingerprint of the graph
|
/// Built fresh per traversal call rather than cached on `KnowledgeCache`:
|
||||||
/// on every use ([`graph_fingerprint`]). entities/relations are plain `pub`
|
/// entities/relations are plain `pub` `Vec`s that get pushed to directly
|
||||||
/// `Vec`s that get changed directly (e.g. `schema.rs`'s load path bypasses
|
/// (e.g. `schema.rs`'s load path bypasses `add_entity`/`add_relation`), so a
|
||||||
/// `add_entity`/`add_relation`), so the cache cannot rely on being told about
|
/// persistent index would need extra bookkeeping to avoid drifting stale. A
|
||||||
/// changes; the fingerprint notices any of them. Rebuilding it on every
|
/// one-off O(V+E) build per call is still a large win over the O(V·E) (BFS)
|
||||||
/// traversal instead made a 2-hop BFS over 1K entities 6.5x slower than the
|
/// / O(steps·active·E) (spreading activation) scans it replaces.
|
||||||
/// scan it replaced (24 -> 155 µs; `BENCHMARKS.md`, "Knowledge Graph").
|
|
||||||
struct AdjacencyIndex {
|
struct AdjacencyIndex {
|
||||||
entity_index: HashMap<u64, usize>,
|
entity_index: HashMap<u64, usize>,
|
||||||
by_entity: HashMap<u64, Vec<usize>>,
|
by_entity: HashMap<u64, Vec<usize>>,
|
||||||
@@ -205,45 +204,6 @@ 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
|
// KnowledgeCache
|
||||||
// ---------------------------------------------------------------------------
|
// ---------------------------------------------------------------------------
|
||||||
@@ -256,7 +216,6 @@ pub struct KnowledgeCache {
|
|||||||
pub alias_strings: Vec<String>,
|
pub alias_strings: Vec<String>,
|
||||||
pub alias_entity_ids: Vec<i64>,
|
pub alias_entity_ids: Vec<i64>,
|
||||||
next_entity_id: u64,
|
next_entity_id: u64,
|
||||||
adjacency: AdjacencyCache,
|
|
||||||
}
|
}
|
||||||
|
|
||||||
impl KnowledgeCache {
|
impl KnowledgeCache {
|
||||||
@@ -267,7 +226,6 @@ impl KnowledgeCache {
|
|||||||
alias_strings: Vec::new(),
|
alias_strings: Vec::new(),
|
||||||
alias_entity_ids: Vec::new(),
|
alias_entity_ids: Vec::new(),
|
||||||
next_entity_id: 0,
|
next_entity_id: 0,
|
||||||
adjacency: AdjacencyCache::default(),
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -278,29 +236,9 @@ impl KnowledgeCache {
|
|||||||
alias_strings: Vec::new(),
|
alias_strings: Vec::new(),
|
||||||
alias_entity_ids: Vec::new(),
|
alias_entity_ids: Vec::new(),
|
||||||
next_entity_id: next_id,
|
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
|
// Entity management
|
||||||
// -----------------------------------------------------------------------
|
// -----------------------------------------------------------------------
|
||||||
@@ -459,7 +397,7 @@ impl KnowledgeCache {
|
|||||||
/// together with their discovered depth. The seed entity itself is NOT
|
/// together with their discovered depth. The seed entity itself is NOT
|
||||||
/// included. Traversal follows both outgoing and incoming relation edges.
|
/// included. Traversal follows both outgoing and incoming relation edges.
|
||||||
pub fn bfs_neighbors(&self, entity_id: u64, max_depth: usize) -> Vec<(Entity, usize)> {
|
pub fn bfs_neighbors(&self, entity_id: u64, max_depth: usize) -> Vec<(Entity, usize)> {
|
||||||
let idx = self.adjacency_index();
|
let idx = AdjacencyIndex::build(&self.entities, &self.relations);
|
||||||
let mut visited: HashSet<u64> = HashSet::new();
|
let mut visited: HashSet<u64> = HashSet::new();
|
||||||
let mut queue: VecDeque<(u64, usize)> = VecDeque::new();
|
let mut queue: VecDeque<(u64, usize)> = VecDeque::new();
|
||||||
let mut results: Vec<(Entity, usize)> = Vec::new();
|
let mut results: Vec<(Entity, usize)> = Vec::new();
|
||||||
@@ -564,7 +502,7 @@ impl KnowledgeCache {
|
|||||||
min_activation: f32,
|
min_activation: f32,
|
||||||
max_steps: usize,
|
max_steps: usize,
|
||||||
) -> Vec<(u64, f32)> {
|
) -> Vec<(u64, f32)> {
|
||||||
let idx = self.adjacency_index();
|
let idx = AdjacencyIndex::build(&self.entities, &self.relations);
|
||||||
let mut activation: HashMap<u64, f32> = HashMap::new();
|
let mut activation: HashMap<u64, f32> = HashMap::new();
|
||||||
|
|
||||||
// Initialise seeds with activation 1.0.
|
// Initialise seeds with activation 1.0.
|
||||||
@@ -693,51 +631,6 @@ impl Default for KnowledgeCache {
|
|||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
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
|
// Original tests — must remain passing
|
||||||
// -----------------------------------------------------------------------
|
// -----------------------------------------------------------------------
|
||||||
|
|||||||
@@ -72,8 +72,7 @@ use ephemeral::{EphemeralConfig, EphemeralStore};
|
|||||||
pub use ephemeral::{EphemeralEntry, EphemeralStats};
|
pub use ephemeral::{EphemeralEntry, EphemeralStats};
|
||||||
use knowledge::KnowledgeCache;
|
use knowledge::KnowledgeCache;
|
||||||
use memory_strategy::{Exchange, MemoryStrategy, StrategyOutput};
|
use memory_strategy::{Exchange, MemoryStrategy, StrategyOutput};
|
||||||
pub use search::SearchOptions;
|
use session::SessionCache;
|
||||||
pub use session::{SessionCache, SessionEntry};
|
|
||||||
|
|
||||||
// --- Error type ---
|
// --- Error type ---
|
||||||
|
|
||||||
@@ -136,13 +135,6 @@ pub struct MemoryConfig {
|
|||||||
/// so search results are the same before and after a reopen. Values must
|
/// so search results are the same before and after a reopen. Values must
|
||||||
/// lie within ±65504; a save outside that is `MemoryError::InvalidEntry`.
|
/// lie within ±65504; a save outside that is `MemoryError::InvalidEntry`.
|
||||||
/// Fixed when the store is created (persisted in `/meta`).
|
/// 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 float16: bool,
|
||||||
pub compression: bool,
|
pub compression: bool,
|
||||||
pub compression_level: u32,
|
pub compression_level: u32,
|
||||||
@@ -198,7 +190,7 @@ impl MemoryConfig {
|
|||||||
embedding_dim,
|
embedding_dim,
|
||||||
chunk_size: 512,
|
chunk_size: 512,
|
||||||
overlap: 50,
|
overlap: 50,
|
||||||
float16: true,
|
float16: false,
|
||||||
compression: false,
|
compression: false,
|
||||||
compression_level: 0,
|
compression_level: 0,
|
||||||
compact_threshold: 0.3,
|
compact_threshold: 0.3,
|
||||||
@@ -1028,18 +1020,6 @@ impl HDF5Memory {
|
|||||||
&self.config
|
&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.
|
/// Get a reference to the knowledge cache.
|
||||||
pub fn knowledge(&self) -> &KnowledgeCache {
|
pub fn knowledge(&self) -> &KnowledgeCache {
|
||||||
&self.knowledge
|
&self.knowledge
|
||||||
@@ -1413,35 +1393,6 @@ impl HDF5Memory {
|
|||||||
}
|
}
|
||||||
|
|
||||||
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<()> {
|
pub fn tick_session(&mut self) -> Result<()> {
|
||||||
let d = self.config.decay_factor;
|
let d = self.config.decay_factor;
|
||||||
for w in self.cache.activation_weights.iter_mut() {
|
for w in self.cache.activation_weights.iter_mut() {
|
||||||
@@ -1640,79 +1591,6 @@ 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]
|
#[test]
|
||||||
fn create_new_file() {
|
fn create_new_file() {
|
||||||
let dir = TempDir::new().unwrap();
|
let dir = TempDir::new().unwrap();
|
||||||
|
|||||||
@@ -13,8 +13,9 @@ use std::path::{Path, PathBuf};
|
|||||||
use std::time::{SystemTime, UNIX_EPOCH};
|
use std::time::{SystemTime, UNIX_EPOCH};
|
||||||
|
|
||||||
use crate::{
|
use crate::{
|
||||||
AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry, SearchOptions,
|
AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry,
|
||||||
confidence::ConfidenceConfig, reranker::ReRankConfig,
|
confidence::{ConfidenceConfig, ScoredResult, reject_low_confidence},
|
||||||
|
reranker::{ReRankConfig, RerankInput, rerank},
|
||||||
};
|
};
|
||||||
|
|
||||||
// ─────────────────────────────────────────────────────────────────────────────
|
// ─────────────────────────────────────────────────────────────────────────────
|
||||||
@@ -523,27 +524,71 @@ impl ClawhdfBackend {
|
|||||||
|
|
||||||
impl MemoryBackend for ClawhdfBackend {
|
impl MemoryBackend for ClawhdfBackend {
|
||||||
/// Search using hybrid vector + BM25 retrieval, then re-rank and
|
/// Search using hybrid vector + BM25 retrieval, then re-rank and
|
||||||
/// confidence-filter — [`HDF5Memory::search`] with both stages on.
|
/// confidence-filter.
|
||||||
fn search(
|
fn search(
|
||||||
&mut self,
|
&mut self,
|
||||||
query_text: &str,
|
query_text: &str,
|
||||||
query_embedding: &[f32],
|
query_embedding: &[f32],
|
||||||
k: usize,
|
k: usize,
|
||||||
) -> Vec<MemorySearchResult> {
|
) -> Vec<MemorySearchResult> {
|
||||||
let options = SearchOptions::new(k)
|
// 1. Hybrid retrieval (vector + BM25, fused by score).
|
||||||
.with_rerank(self.rerank_config)
|
let candidates = k.saturating_mul(3).max(10);
|
||||||
.with_confidence(self.confidence_config.clone())
|
let raw = self.memory.hybrid_search_with(
|
||||||
.at_time(Self::now_secs());
|
query_embedding,
|
||||||
self.memory
|
query_text,
|
||||||
.search(query_embedding, query_text, &options)
|
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
|
||||||
.into_iter()
|
.into_iter()
|
||||||
.map(|r| MemorySearchResult {
|
.take(k)
|
||||||
text: r.chunk,
|
.filter_map(|sr| {
|
||||||
score: r.score,
|
let r = raw_by_idx.get(&sr.index)?;
|
||||||
path: r.source_channel.clone(),
|
let path = r.source_channel.clone();
|
||||||
|
Some(MemorySearchResult {
|
||||||
|
text: r.chunk.clone(),
|
||||||
|
score: sr.score,
|
||||||
|
path: path.clone(),
|
||||||
line_range: None,
|
line_range: None,
|
||||||
timestamp: Some(r.timestamp),
|
timestamp: Some(r.timestamp),
|
||||||
source: r.source_channel,
|
source: path,
|
||||||
|
})
|
||||||
})
|
})
|
||||||
.collect()
|
.collect()
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -2,107 +2,18 @@
|
|||||||
|
|
||||||
use std::path::Path;
|
use std::path::Path;
|
||||||
|
|
||||||
use std::collections::HashSet;
|
|
||||||
|
|
||||||
use crate::bm25;
|
use crate::bm25;
|
||||||
use crate::confidence::{ConfidenceConfig, ScoredResult, reject_low_confidence};
|
|
||||||
use crate::hybrid;
|
use crate::hybrid;
|
||||||
use crate::reranker::{ReRankConfig, RerankInput, rerank};
|
|
||||||
use crate::{HDF5Memory, MAX_ACTIVATION_WEIGHT, MemoryError, Result, SearchResult};
|
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 {
|
impl HDF5Memory {
|
||||||
/// Vector + keyword scoring stage of [`HDF5Memory::search`].
|
/// Vector + keyword scoring stage of [`HDF5Memory::hybrid_search`].
|
||||||
///
|
///
|
||||||
/// Without the `hnsw` feature this is a full linear cosine scan (the exact
|
/// Without the `hnsw` feature this is a full linear cosine scan (the exact
|
||||||
/// previous behaviour, also used as the correctness oracle in tests). With
|
/// previous behaviour, also used as the correctness oracle in tests). With
|
||||||
/// `hnsw` enabled and an index available, the vector candidates come from an
|
/// `hnsw` enabled and an index available, the vector candidates come from an
|
||||||
/// approximate-nearest-neighbour search over an over-fetched pool, then merge
|
/// approximate-nearest-neighbour search over an over-fetched pool, then merge
|
||||||
/// with BM25 via the shared [`hybrid::merge_vector_keyword`].
|
/// 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")]
|
#[cfg(feature = "hnsw")]
|
||||||
fn vector_keyword_search(
|
fn vector_keyword_search(
|
||||||
&mut self,
|
&mut self,
|
||||||
@@ -111,30 +22,16 @@ impl HDF5Memory {
|
|||||||
bm25: &bm25::BM25Index,
|
bm25: &bm25::BM25Index,
|
||||||
fusion: hybrid::Fusion,
|
fusion: hybrid::Fusion,
|
||||||
k: usize,
|
k: usize,
|
||||||
exclude: Option<&[u8]>,
|
|
||||||
) -> Vec<(usize, f32)> {
|
) -> Vec<(usize, f32)> {
|
||||||
self.ensure_hnsw_fresh();
|
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() {
|
match self.hnsw.as_ref() {
|
||||||
Some(index) if !index.is_empty() && index.dimension() == query_embedding.len() => {
|
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 ef = self.hnsw_ef_search(k).max(pool);
|
||||||
let candidates = index.search(query_embedding, pool, ef);
|
let candidates = index.search(query_embedding, pool, ef);
|
||||||
// A quantised index returns approximate distances, and no
|
// A quantised index returns approximate distances, and no
|
||||||
@@ -145,7 +42,6 @@ impl HDF5Memory {
|
|||||||
let exact = index.storage() == clawhdf5_ann::Storage::Int8;
|
let exact = index.storage() == clawhdf5_ann::Storage::Int8;
|
||||||
let vec_scores: Vec<(usize, f32)> = candidates
|
let vec_scores: Vec<(usize, f32)> = candidates
|
||||||
.into_iter()
|
.into_iter()
|
||||||
.filter(|(id, _)| exclude.is_none_or(|ex| ex[*id] == 0))
|
|
||||||
.map(|(id, dist)| {
|
.map(|(id, dist)| {
|
||||||
let score = if exact {
|
let score = if exact {
|
||||||
crate::vector_search::cosine_similarity(
|
crate::vector_search::cosine_similarity(
|
||||||
@@ -160,28 +56,10 @@ impl HDF5Memory {
|
|||||||
.collect();
|
.collect();
|
||||||
// Fusion normalises over every keyword match, so it needs all
|
// Fusion normalises over every keyword match, so it needs all
|
||||||
// the scores — but not ranked.
|
// the scores — but not ranked.
|
||||||
let mut kw_scores = bm25.scores(query_text);
|
let 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::fuse(vec_scores, kw_scores, fusion, k)
|
||||||
}
|
}
|
||||||
_ => match exclude {
|
_ => hybrid::hybrid_search_fused(
|
||||||
Some(ex) => {
|
|
||||||
self.exact_masked_search(query_embedding, query_text, bm25, fusion, k, ex)
|
|
||||||
}
|
|
||||||
None => hybrid::hybrid_search_fused(
|
|
||||||
query_embedding,
|
query_embedding,
|
||||||
query_text,
|
query_text,
|
||||||
&self.cache.embeddings,
|
&self.cache.embeddings,
|
||||||
@@ -191,7 +69,6 @@ impl HDF5Memory {
|
|||||||
fusion,
|
fusion,
|
||||||
k,
|
k,
|
||||||
),
|
),
|
||||||
},
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -203,11 +80,8 @@ impl HDF5Memory {
|
|||||||
bm25: &bm25::BM25Index,
|
bm25: &bm25::BM25Index,
|
||||||
fusion: hybrid::Fusion,
|
fusion: hybrid::Fusion,
|
||||||
k: usize,
|
k: usize,
|
||||||
exclude: Option<&[u8]>,
|
|
||||||
) -> Vec<(usize, f32)> {
|
) -> Vec<(usize, f32)> {
|
||||||
match exclude {
|
hybrid::hybrid_search_fused(
|
||||||
Some(ex) => self.exact_masked_search(query_embedding, query_text, bm25, fusion, k, ex),
|
|
||||||
None => hybrid::hybrid_search_fused(
|
|
||||||
query_embedding,
|
query_embedding,
|
||||||
query_text,
|
query_text,
|
||||||
&self.cache.embeddings,
|
&self.cache.embeddings,
|
||||||
@@ -216,37 +90,7 @@ impl HDF5Memory {
|
|||||||
bm25,
|
bm25,
|
||||||
fusion,
|
fusion,
|
||||||
k,
|
k,
|
||||||
),
|
)
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// 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)> {
|
|
||||||
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.
|
/// Perform hybrid search combining cosine vector similarity and BM25 keyword search.
|
||||||
@@ -281,53 +125,12 @@ impl HDF5Memory {
|
|||||||
fusion: hybrid::Fusion,
|
fusion: hybrid::Fusion,
|
||||||
k: usize,
|
k: usize,
|
||||||
) -> Vec<SearchResult> {
|
) -> 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
|
// 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
|
// incrementally. Take it out for the duration of the call so the
|
||||||
// vector stage can borrow `self` mutably, then put it back.
|
// vector stage can borrow `self` mutably, then put it back.
|
||||||
self.ensure_bm25_fresh();
|
self.ensure_bm25_fresh();
|
||||||
let bm25 = self.bm25.take().expect("ensure_bm25_fresh leaves an index");
|
let bm25 = self.bm25.take().expect("ensure_bm25_fresh leaves an index");
|
||||||
let scored = self.vector_keyword_search(
|
let scored = self.vector_keyword_search(query_embedding, query_text, &bm25, fusion, k);
|
||||||
query_embedding,
|
|
||||||
query_text,
|
|
||||||
&bm25,
|
|
||||||
options.fusion,
|
|
||||||
fetch,
|
|
||||||
exclude.as_deref(),
|
|
||||||
);
|
|
||||||
self.bm25 = Some(bm25);
|
|
||||||
|
|
||||||
let mut results: Vec<SearchResult> = scored
|
let mut results: Vec<SearchResult> = scored
|
||||||
.into_iter()
|
.into_iter()
|
||||||
.map(|(idx, score)| {
|
.map(|(idx, score)| {
|
||||||
@@ -351,25 +154,6 @@ impl HDF5Memory {
|
|||||||
.then(a.index.cmp(&b.index))
|
.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`
|
// Only reinforce records that actually matched. When fewer than `k`
|
||||||
// records are relevant, the rest of the list is zero-score filler;
|
// records are relevant, the rest of the list is zero-score filler;
|
||||||
// boosting it would teach the store that arbitrary records are
|
// boosting it would teach the store that arbitrary records are
|
||||||
@@ -380,45 +164,11 @@ impl HDF5Memory {
|
|||||||
.map(|r| r.index)
|
.map(|r| r.index)
|
||||||
.collect();
|
.collect();
|
||||||
self.apply_hebbian_boost(&hit_indices);
|
self.apply_hebbian_boost(&hit_indices);
|
||||||
|
self.bm25 = Some(bm25);
|
||||||
|
|
||||||
results
|
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
|
/// Reinforce the records a query returned. The new weights are persisted by
|
||||||
/// the next checkpoint (any write that flushes, `flush_wal`, or drop) — not
|
/// the next checkpoint (any write that flushes, `flush_wal`, or drop) — not
|
||||||
/// by rewriting the whole store inside the query, which is what made
|
/// by rewriting the whole store inside the query, which is what made
|
||||||
|
|||||||
@@ -33,7 +33,7 @@ impl SessionCache {
|
|||||||
self.entries.is_empty()
|
self.entries.is_empty()
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Add a new session with its summary, timestamped now.
|
/// Add a new session with its summary.
|
||||||
pub fn add(
|
pub fn add(
|
||||||
&mut self,
|
&mut self,
|
||||||
id: &str,
|
id: &str,
|
||||||
@@ -47,21 +47,6 @@ impl SessionCache {
|
|||||||
.unwrap_or_default()
|
.unwrap_or_default()
|
||||||
.as_secs_f64()
|
.as_secs_f64()
|
||||||
* 1_000_000.0; // microseconds
|
* 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 {
|
self.entries.push(SessionEntry {
|
||||||
id: id.to_string(),
|
id: id.to_string(),
|
||||||
start_idx: start_idx as u64,
|
start_idx: start_idx as u64,
|
||||||
|
|||||||
@@ -206,55 +206,3 @@ fn wal_replay_rounds_like_a_live_save() {
|
|||||||
assert_eq!(recovered.count(), 30);
|
assert_eq!(recovered.count(), 30);
|
||||||
assert_eq!(search_bits(&mut recovered, 77), live);
|
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());
|
|
||||||
}
|
|
||||||
|
|||||||
@@ -1,344 +0,0 @@
|
|||||||
//! `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,14 +11,12 @@
|
|||||||
//!
|
//!
|
||||||
//! Configuration matrix:
|
//! Configuration matrix:
|
||||||
//! - Text lengths: short (50 chars), medium (200 chars), long (1000 chars)
|
//! - Text lengths: short (50 chars), medium (200 chars), long (1000 chars)
|
||||||
//! - Embedding: 384-dim, stored as float16 (the default for new stores) or
|
//! - Embedding: 384-dim f32 (1536 bytes raw per record)
|
||||||
//! f32 with `--f32`; "raw" bytes are counted as f32 input either way
|
|
||||||
//! - WAL: enabled and disabled
|
//! - WAL: enabled and disabled
|
||||||
//!
|
//!
|
||||||
//! # Usage
|
//! # Usage
|
||||||
//! ```
|
//! ```
|
||||||
//! cargo run --release --bin footprint_bench # float16 stores
|
//! cargo run --release --bin footprint_bench
|
||||||
//! cargo run --release --bin footprint_bench -- --f32 # f32 stores
|
|
||||||
//! ```
|
//! ```
|
||||||
|
|
||||||
use std::time::Instant;
|
use std::time::Instant;
|
||||||
@@ -26,9 +24,6 @@ use std::time::Instant;
|
|||||||
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};
|
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};
|
||||||
use tempfile::TempDir;
|
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;
|
const EMBEDDING_DIM: usize = 384;
|
||||||
|
|
||||||
// Raw bytes per record: 384 f32 embeddings + median text + overhead
|
// Raw bytes per record: 384 f32 embeddings + median text + overhead
|
||||||
@@ -157,9 +152,6 @@ fn measure_footprint(
|
|||||||
config.compression = compression;
|
config.compression = compression;
|
||||||
config.compression_level = if compression { 6 } else { 0 };
|
config.compression_level = if compression { 6 } else { 0 };
|
||||||
config.compact_threshold = 0.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");
|
let mut memory = HDF5Memory::create(config).expect("HDF5Memory::create failed");
|
||||||
|
|
||||||
@@ -249,19 +241,11 @@ fn fmt_n(n: usize) -> String {
|
|||||||
// ---------------------------------------------------------------------------
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
fn main() {
|
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!("=================================================================");
|
||||||
println!(" ClawhDF5 Memory Footprint Benchmark");
|
println!(" ClawhDF5 Memory Footprint Benchmark");
|
||||||
println!("=================================================================");
|
println!("=================================================================");
|
||||||
println!();
|
println!();
|
||||||
println!("Embedding: 384-dim, stored as {stored}; raw input counted as f32");
|
println!("Embedding: 384-dim f32 = 1,536 bytes raw per record");
|
||||||
println!("Text lengths: short=50 chars, medium=200 chars, long=1000 chars");
|
println!("Text lengths: short=50 chars, medium=200 chars, long=1000 chars");
|
||||||
println!();
|
println!();
|
||||||
|
|
||||||
|
|||||||
@@ -64,11 +64,6 @@ use tempfile::TempDir;
|
|||||||
|
|
||||||
const EMBEDDING_DIM: usize = 384;
|
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.
|
/// A mode's fusion, as one short string for the reports.
|
||||||
fn describe(mode: Mode) -> String {
|
fn describe(mode: Mode) -> String {
|
||||||
let fusion = match mode.fusion {
|
let fusion = match mode.fusion {
|
||||||
@@ -436,7 +431,6 @@ fn evaluate_question(
|
|||||||
let mut config = MemoryConfig::new(dir.path().join("lme.h5"), "lme-bench", EMBEDDING_DIM);
|
let mut config = MemoryConfig::new(dir.path().join("lme.h5"), "lme-bench", EMBEDDING_DIM);
|
||||||
config.wal_enabled = false;
|
config.wal_enabled = false;
|
||||||
config.compact_threshold = 0.0;
|
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");
|
let mut memory = HDF5Memory::create(config).expect("failed to create HDF5Memory");
|
||||||
memory.set_token_filter(mode.tokens);
|
memory.set_token_filter(mode.tokens);
|
||||||
@@ -946,10 +940,6 @@ fn main() {
|
|||||||
limit = Some(v.parse().expect("--limit must be a positive integer"));
|
limit = Some(v.parse().expect("--limit must be a positive integer"));
|
||||||
}
|
}
|
||||||
"--sweep" => sweep = true,
|
"--sweep" => sweep = true,
|
||||||
"--float16" => {
|
|
||||||
FLOAT16.store(true, std::sync::atomic::Ordering::Relaxed);
|
|
||||||
eprintln!("Stores use MemoryConfig::float16 (half-precision embeddings)");
|
|
||||||
}
|
|
||||||
"--rerank-sweep" => {
|
"--rerank-sweep" => {
|
||||||
// Re-ranking needs the vector stage to have candidates worth
|
// Re-ranking needs the vector stage to have candidates worth
|
||||||
// reordering, so this is an embeddings-only comparison.
|
// reordering, so this is an embeddings-only comparison.
|
||||||
@@ -981,9 +971,6 @@ fn main() {
|
|||||||
--rerank-sweep\n\
|
--rerank-sweep\n\
|
||||||
compare re-ranking off, metadata-only (the old\n\
|
compare re-ranking off, metadata-only (the old\n\
|
||||||
behaviour) and blended at several half-lives.\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\
|
--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\
|
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."
|
never searched; this is what searches it."
|
||||||
|
|||||||
@@ -20,7 +20,6 @@
|
|||||||
//! cargo run --release -p clawhdf5-bench --bin search_harness -- --json out.json
|
//! 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 -- --ann-only --uniform
|
||||||
//! cargo run --release -p clawhdf5-bench --bin search_harness -- --float16-study --full
|
//! 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};
|
use std::time::{Duration, Instant};
|
||||||
@@ -488,177 +487,6 @@ 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?
|
// float16 study: what does half-precision embedding storage cost?
|
||||||
// ---------------------------------------------------------------------------
|
// ---------------------------------------------------------------------------
|
||||||
@@ -960,19 +788,6 @@ fn main() {
|
|||||||
}
|
}
|
||||||
return;
|
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") {
|
if args.iter().any(|a| a == "--f16-first") {
|
||||||
F16_FIRST.store(true, std::sync::atomic::Ordering::Relaxed);
|
F16_FIRST.store(true, std::sync::atomic::Ordering::Relaxed);
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -36,13 +36,10 @@ enum Commands {
|
|||||||
/// Accepted for compatibility; int8 is now the default
|
/// Accepted for compatibility; int8 is now the default
|
||||||
#[arg(long, hide = true, conflicts_with = "f32_index")]
|
#[arg(long, hide = true, conflicts_with = "f32_index")]
|
||||||
quantized_index: bool,
|
quantized_index: bool,
|
||||||
/// Store embeddings as full-precision f32 instead of the default
|
/// Store embeddings on disk as IEEE half precision (float16): half
|
||||||
/// half precision (float16: half the bytes, about three significant
|
/// the bytes, about three significant digits; values must lie within
|
||||||
/// digits, values within ±65504)
|
/// ±65504
|
||||||
#[arg(long)]
|
#[arg(long)]
|
||||||
f32: bool,
|
|
||||||
/// Accepted for compatibility; float16 is now the default
|
|
||||||
#[arg(long, hide = true, conflicts_with = "f32")]
|
|
||||||
float16: bool,
|
float16: bool,
|
||||||
},
|
},
|
||||||
/// Save a memory entry (reads JSON from stdin or --json)
|
/// Save a memory entry (reads JSON from stdin or --json)
|
||||||
@@ -110,16 +107,11 @@ fn run(cli: Cli) -> Result<(), Box<dyn std::error::Error>> {
|
|||||||
wal,
|
wal,
|
||||||
f32_index,
|
f32_index,
|
||||||
quantized_index: _,
|
quantized_index: _,
|
||||||
f32,
|
float16,
|
||||||
float16: _,
|
|
||||||
} => {
|
} => {
|
||||||
let mut config = MemoryConfig::new(cli.path.clone(), &agent_id, dim);
|
let mut config = MemoryConfig::new(cli.path.clone(), &agent_id, dim);
|
||||||
config.wal_enabled = wal;
|
config.wal_enabled = wal;
|
||||||
// As with --f32-index: only ever switch the library default off.
|
config.float16 = float16;
|
||||||
if f32 {
|
|
||||||
config.float16 = false;
|
|
||||||
}
|
|
||||||
let config_float16 = config.float16;
|
|
||||||
// Only ever switch *off* the library default: assigning the flag
|
// Only ever switch *off* the library default: assigning the flag
|
||||||
// outright would force every CLI-created store back to f32 unless
|
// outright would force every CLI-created store back to f32 unless
|
||||||
// the caller knew to ask for int8.
|
// the caller knew to ask for int8.
|
||||||
@@ -135,7 +127,7 @@ fn run(cli: Cli) -> Result<(), Box<dyn std::error::Error>> {
|
|||||||
"embedding_dim": dim,
|
"embedding_dim": dim,
|
||||||
"wal_enabled": wal,
|
"wal_enabled": wal,
|
||||||
"quantized_index": config_quantized,
|
"quantized_index": config_quantized,
|
||||||
"float16": config_float16,
|
"float16": float16,
|
||||||
"count": mem.count(),
|
"count": mem.count(),
|
||||||
});
|
});
|
||||||
println!("{}", serde_json::to_string_pretty(&j)?);
|
println!("{}", serde_json::to_string_pretty(&j)?);
|
||||||
|
|||||||
@@ -3,7 +3,7 @@ name = "clawhdf5-migrate"
|
|||||||
version = "2.7.0"
|
version = "2.7.0"
|
||||||
edition = "2024"
|
edition = "2024"
|
||||||
rust-version.workspace = true
|
rust-version.workspace = true
|
||||||
description = "CLI to migrate SQLite agent memory databases to clawhdf5-agent stores"
|
description = "CLI to migrate SQLite agent memory databases to HDF5 format"
|
||||||
license = "MIT"
|
license = "MIT"
|
||||||
repository = "https://git.redclaw.dev/quantumclaw/clawhdf5"
|
repository = "https://git.redclaw.dev/quantumclaw/clawhdf5"
|
||||||
readme = "README.md"
|
readme = "README.md"
|
||||||
@@ -17,8 +17,10 @@ path = "src/main.rs"
|
|||||||
[dependencies]
|
[dependencies]
|
||||||
clawhdf5-agent = { path = "../clawhdf5-agent", version = "2.7.0" }
|
clawhdf5-agent = { path = "../clawhdf5-agent", version = "2.7.0" }
|
||||||
clawhdf5-format = { path = "../clawhdf5-format", 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"] }
|
rusqlite = { version = "0.31", features = ["bundled"] }
|
||||||
clap = { version = "4", features = ["derive"] }
|
clap = { version = "4", features = ["derive"] }
|
||||||
|
half = { workspace = true }
|
||||||
|
|
||||||
[dev-dependencies]
|
[dev-dependencies]
|
||||||
tempfile = { workspace = true }
|
tempfile = { workspace = true }
|
||||||
|
|||||||
@@ -3,12 +3,9 @@
|
|||||||
[](https://crates.io/crates/clawhdf5-migrate)
|
[](https://crates.io/crates/clawhdf5-migrate)
|
||||||
[](https://docs.rs/clawhdf5-migrate)
|
[](https://docs.rs/clawhdf5-migrate)
|
||||||
|
|
||||||
CLI tool to migrate SQLite agent memory databases (the ZeroClaw layout) to a
|
CLI tool to migrate SQLite agent memory databases to HDF5 format.
|
||||||
[clawhdf5-agent](https://crates.io/crates/clawhdf5-agent) store.
|
|
||||||
|
|
||||||
The output is written through `clawhdf5-agent`'s own API, so it opens with
|
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).
|
||||||
`HDF5Memory::open` and is searchable immediately: memory records, sessions and
|
|
||||||
the knowledge graph (entities and relations) are carried over.
|
|
||||||
|
|
||||||
## Installation
|
## Installation
|
||||||
|
|
||||||
@@ -19,19 +16,9 @@ cargo install clawhdf5-migrate
|
|||||||
## Usage
|
## Usage
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
clawhdf5-migrate --sqlite agent.db --hdf5 agent.h5 --agent-id my-agent
|
clawhdf5-migrate --input agent.db --output agent.h5
|
||||||
```
|
```
|
||||||
|
|
||||||
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
|
## License
|
||||||
|
|
||||||
MIT
|
MIT
|
||||||
|
|||||||
@@ -0,0 +1,163 @@
|
|||||||
|
//! 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),
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,366 @@
|
|||||||
|
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'));
|
||||||
|
}
|
||||||
|
}
|
||||||
+418
-1038
File diff suppressed because it is too large
Load Diff
@@ -50,8 +50,12 @@ pub struct SqliteData {
|
|||||||
pub sessions: Vec<Session>,
|
pub sessions: Vec<Session>,
|
||||||
pub entities: Vec<Entity>,
|
pub entities: Vec<Entity>,
|
||||||
pub relations: Vec<Relation>,
|
pub relations: Vec<Relation>,
|
||||||
/// `--embedding-dim`, or the first row's; 0 when neither exists.
|
|
||||||
pub embedding_dim: usize,
|
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.
|
/// A table name plus the ordered column names the reader maps by position.
|
||||||
@@ -163,13 +167,11 @@ pub fn read_counts(
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Auto-detect embedding dimension from the BLOB size of the first chunk (in
|
/// Auto-detect embedding dimension from the first chunk's BLOB size.
|
||||||
/// id order, deleted or not).
|
|
||||||
fn detect_embedding_dim(conn: &Connection, config: &SchemaConfig) -> SqlResult<Option<usize>> {
|
fn detect_embedding_dim(conn: &Connection, config: &SchemaConfig) -> SqlResult<Option<usize>> {
|
||||||
let emb_col = config.chunks.columns.get(2).copied().unwrap_or("embedding");
|
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!(
|
let mut stmt = conn.prepare(&format!(
|
||||||
"SELECT {emb_col} FROM {} ORDER BY {id_col} LIMIT 1",
|
"SELECT {emb_col} FROM {} LIMIT 1",
|
||||||
config.chunks.table
|
config.chunks.table
|
||||||
))?;
|
))?;
|
||||||
let mut rows = stmt.query([])?;
|
let mut rows = stmt.query([])?;
|
||||||
@@ -193,16 +195,24 @@ fn blob_to_f32(blob: &[u8]) -> Vec<f32> {
|
|||||||
/// Read all data from a ZeroClaw SQLite database.
|
/// Read all data from a ZeroClaw SQLite database.
|
||||||
///
|
///
|
||||||
/// If `skip_deleted` is true, rows with `deleted=1` are excluded from chunks.
|
/// If `skip_deleted` is true, rows with `deleted=1` are excluded from chunks.
|
||||||
/// If `embedding_dim` is `None`, auto-detect from the first row (0 when there
|
/// If `embedding_dim` is `None`, auto-detect from the first row.
|
||||||
/// 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(
|
pub fn read_sqlite(
|
||||||
path: &str,
|
path: &str,
|
||||||
skip_deleted: bool,
|
skip_deleted: bool,
|
||||||
embedding_dim: Option<usize>,
|
embedding_dim: Option<usize>,
|
||||||
config: &SchemaConfig,
|
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>> {
|
) -> Result<SqliteData, Box<dyn std::error::Error>> {
|
||||||
let conn = Connection::open(path)?;
|
let conn = Connection::open(path)?;
|
||||||
|
|
||||||
@@ -211,7 +221,7 @@ pub fn read_sqlite(
|
|||||||
None => detect_embedding_dim(&conn, config)?.unwrap_or(0),
|
None => detect_embedding_dim(&conn, config)?.unwrap_or(0),
|
||||||
};
|
};
|
||||||
|
|
||||||
let chunks = read_chunks(&conn, skip_deleted, config)?;
|
let chunks = read_chunks(&conn, skip_deleted, dim, config, min_chunk_id)?;
|
||||||
let sessions = read_sessions(&conn, config)?;
|
let sessions = read_sessions(&conn, config)?;
|
||||||
let entities = read_entities(&conn, config)?;
|
let entities = read_entities(&conn, config)?;
|
||||||
let relations = read_relations(&conn, config)?;
|
let relations = read_relations(&conn, config)?;
|
||||||
@@ -222,43 +232,42 @@ pub fn read_sqlite(
|
|||||||
entities,
|
entities,
|
||||||
relations,
|
relations,
|
||||||
embedding_dim: dim,
|
embedding_dim: dim,
|
||||||
|
source_path: path.to_owned(),
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
fn read_chunks(
|
fn read_chunks(
|
||||||
conn: &Connection,
|
conn: &Connection,
|
||||||
skip_deleted: bool,
|
skip_deleted: bool,
|
||||||
|
expected_dim: usize,
|
||||||
config: &SchemaConfig,
|
config: &SchemaConfig,
|
||||||
|
min_chunk_id: i64,
|
||||||
) -> SqlResult<Vec<MemoryChunk>> {
|
) -> SqlResult<Vec<MemoryChunk>> {
|
||||||
let id_col = config.chunks.columns.first().copied().unwrap_or("id");
|
let id_col = config.chunks.columns.first().copied().unwrap_or("id");
|
||||||
let deleted_col = config.chunks.columns.get(7).copied().unwrap_or("deleted");
|
let deleted_col = config.chunks.columns.get(7).copied().unwrap_or("deleted");
|
||||||
let mut where_clause = String::new();
|
let mut conds = Vec::new();
|
||||||
if skip_deleted {
|
if skip_deleted {
|
||||||
where_clause = format!(" WHERE {deleted_col} = 0");
|
conds.push(format!("{deleted_col} = 0"));
|
||||||
}
|
}
|
||||||
// In id order, so the store's records follow the source's order.
|
if min_chunk_id > 0 {
|
||||||
where_clause.push_str(&format!(" ORDER BY {id_col}"));
|
conds.push(format!("{id_col} > {min_chunk_id}"));
|
||||||
|
}
|
||||||
|
let where_clause = if conds.is_empty() {
|
||||||
|
String::new()
|
||||||
|
} else {
|
||||||
|
format!(" WHERE {}", conds.join(" AND "))
|
||||||
|
};
|
||||||
let sql = config.chunks.select(&where_clause);
|
let sql = config.chunks.select(&where_clause);
|
||||||
|
|
||||||
let mut stmt = conn.prepare(&sql)?;
|
let mut stmt = conn.prepare(&sql)?;
|
||||||
let rows = stmt.query_map([], |row| {
|
let rows = stmt.query_map([], |row| {
|
||||||
let blob: Vec<u8> = row.get(2)?;
|
let blob: Vec<u8> = row.get(2)?;
|
||||||
if !blob.len().is_multiple_of(4) {
|
let mut embedding = blob_to_f32(&blob);
|
||||||
let id: i64 = row.get(0)?;
|
|
||||||
return Err(rusqlite::Error::FromSqlConversionFailure(
|
// Validate/truncate to expected dimension
|
||||||
2,
|
if expected_dim > 0 {
|
||||||
rusqlite::types::Type::Blob,
|
embedding.truncate(expected_dim);
|
||||||
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 {
|
Ok(MemoryChunk {
|
||||||
id: row.get(0)?,
|
id: row.get(0)?,
|
||||||
|
|||||||
@@ -1,407 +0,0 @@
|
|||||||
//! 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,266 +1,192 @@
|
|||||||
//! Validate a migration by reading the store back the way an agent would:
|
use clawhdf5::reader::File as Hdf5File;
|
||||||
//! through `HDF5Memory::open_read_only`, comparing what it loads with the
|
use clawhdf5_format::provenance::VerifyResult;
|
||||||
//! 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::sqlite_reader::SqliteData;
|
||||||
use crate::store_writer::{Migration, US_PER_SEC};
|
|
||||||
|
|
||||||
type BoxErr = Box<dyn std::error::Error>;
|
type BoxErr = Box<dyn std::error::Error>;
|
||||||
|
|
||||||
/// Summary of a migration validation.
|
/// Summary of a migration validation.
|
||||||
#[derive(Debug)]
|
#[derive(Debug)]
|
||||||
pub struct ValidationSummary {
|
pub struct ValidationSummary {
|
||||||
/// Records in the store (including tombstones).
|
pub chunks: u64,
|
||||||
pub count: usize,
|
pub sessions: u64,
|
||||||
/// Records in the store that are not deleted.
|
pub entities: u64,
|
||||||
pub active: usize,
|
pub relations: u64,
|
||||||
pub sessions: usize,
|
pub embedding_dim: u64,
|
||||||
pub entities: usize,
|
/// Number of rows whose full content was compared against the source.
|
||||||
pub relations: usize,
|
|
||||||
pub embedding_dim: usize,
|
|
||||||
pub float16: bool,
|
|
||||||
/// Rows whose full content was compared against the source.
|
|
||||||
pub rows_checked: u64,
|
pub rows_checked: u64,
|
||||||
/// Whether a search for a migrated record found it (`false` when there
|
/// Whether the `chunks/text` and `chunks/embeddings` SHINES provenance
|
||||||
/// was no active migrated record with an embedding to search for).
|
/// hashes (written via [`crate::hdf5_writer`]) were both present and
|
||||||
pub search_checked: bool,
|
/// 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,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Validate the store at `path` against the source rows `migration` wrote.
|
/// Validate a migrated HDF5 file against the source data.
|
||||||
///
|
///
|
||||||
/// Counts and the session / entity / relation rows are always checked in
|
/// Reads the written file back and compares actual content — chunk text,
|
||||||
/// full. Memory records are content-checked on a representative sample, or
|
/// embeddings, and every session/entity/relation field — to the source, not
|
||||||
/// all of them with `full`. Embeddings must match exactly: the source values
|
/// just the row counts. When `full` is false a representative sample of chunk
|
||||||
/// themselves in an `f32` store, their [`round_to_f16`] in a `float16` one.
|
/// rows is content-checked (counts and all other groups are always checked in
|
||||||
pub fn validate_store(
|
/// full); when `full` is true every chunk row is compared too. `float16` widens
|
||||||
path: &Path,
|
/// the embedding tolerance to allow for half-precision quantization.
|
||||||
|
pub fn validate_hdf5(
|
||||||
|
path: &str,
|
||||||
source: &SqliteData,
|
source: &SqliteData,
|
||||||
migration: &Migration,
|
|
||||||
full: bool,
|
full: bool,
|
||||||
|
float16: bool,
|
||||||
) -> Result<ValidationSummary, BoxErr> {
|
) -> Result<ValidationSummary, BoxErr> {
|
||||||
let mut mem = HDF5Memory::open_read_only(path)?;
|
let got = read_hdf5(path)?;
|
||||||
let float16 = mem.config().float16;
|
let provenance_verified = verify_chunk_provenance(path)?;
|
||||||
let dim = mem.config().embedding_dim;
|
|
||||||
|
|
||||||
// ---- Counts ----
|
// ---- Counts ----
|
||||||
check_count("record", mem.count(), migration.store_count)?;
|
check_count("chunk", got.chunks.len(), source.chunks.len())?;
|
||||||
if float16 != migration.float16 {
|
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 {
|
||||||
return Err(format!(
|
return Err(format!(
|
||||||
"float16 mismatch: store {float16}, expected {}",
|
"embedding_dim mismatch: HDF5 has {}, source has {}",
|
||||||
migration.float16
|
got.embedding_dim, source.embedding_dim
|
||||||
)
|
)
|
||||||
.into());
|
.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(),
|
|
||||||
)?;
|
|
||||||
}
|
|
||||||
|
|
||||||
// ---- Memory records (sampled or full) ----
|
// ---- Chunk content (sampled or full) ----
|
||||||
|
let (emb_abs, emb_rel) = if float16 { (1e-2, 1e-2) } else { (1e-4, 0.0) };
|
||||||
let mut rows_checked = 0u64;
|
let mut rows_checked = 0u64;
|
||||||
let expected_value = |v: f32| if float16 { round_to_f16(v) } else { v };
|
for i in sample_indices(source.chunks.len(), full) {
|
||||||
for k in sample_indices(migration.records.len(), full) {
|
let (s, g) = (&source.chunks[i], &got.chunks[i]);
|
||||||
let (idx, src) = migration.records[k];
|
if s.id != g.id {
|
||||||
let s = &source.chunks[src];
|
return Err(field_err("chunk", i, "id", s.id, g.id));
|
||||||
let c = &mem.cache;
|
}
|
||||||
if idx >= c.len() {
|
if s.chunk != g.chunk {
|
||||||
|
return Err(format!(
|
||||||
|
"chunk[{i}].text mismatch: source {:?}, HDF5 {:?}",
|
||||||
|
truncate(&s.chunk),
|
||||||
|
truncate(&g.chunk)
|
||||||
|
)
|
||||||
|
.into());
|
||||||
|
}
|
||||||
|
if s.session_id != g.session_id || s.source_channel != g.source_channel || s.tags != g.tags
|
||||||
|
{
|
||||||
|
return Err(format!("chunk[{i}] 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() {
|
||||||
|
return Err(format!(
|
||||||
|
"chunk[{i}] embedding length mismatch: {} vs {}",
|
||||||
|
s.embedding.len(),
|
||||||
|
g.embedding.len()
|
||||||
|
)
|
||||||
|
.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(
|
return Err(
|
||||||
format!("record {idx} (chunk id {}) is missing from the store", s.id).into(),
|
format!("chunk[{i}].embedding[{k}] mismatch: source {a}, HDF5 {b}").into(),
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
let id = s.id;
|
|
||||||
if c.chunks[idx] != s.chunk {
|
|
||||||
return Err(format!(
|
|
||||||
"record {idx} (chunk id {id}) text mismatch: source {:?}, store {:?}",
|
|
||||||
truncate(&s.chunk),
|
|
||||||
truncate(&c.chunks[idx])
|
|
||||||
)
|
|
||||||
.into());
|
|
||||||
}
|
}
|
||||||
if c.source_channels[idx] != s.source_channel
|
rows_checked += 1;
|
||||||
|| c.session_ids[idx] != s.session_id
|
}
|
||||||
|| c.tags[idx] != s.tags
|
|
||||||
|
// ---- 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!("record {idx} (chunk id {id}) string field mismatch").into());
|
return Err(format!("session[{i}] mismatch").into());
|
||||||
}
|
|
||||||
if c.timestamps[idx].to_bits() != s.timestamp.to_bits() {
|
|
||||||
return Err(format!(
|
|
||||||
"record {idx} (chunk id {id}) timestamp mismatch: source {}, store {}",
|
|
||||||
s.timestamp, c.timestamps[idx]
|
|
||||||
)
|
|
||||||
.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;
|
rows_checked += 1;
|
||||||
}
|
}
|
||||||
|
for (i, (s, g)) in source.entities.iter().zip(got.entities.iter()).enumerate() {
|
||||||
// ---- Records tombstoned because their source row was deleted ----
|
if s.id != g.id
|
||||||
for &(idx, src) in &migration.deleted_in_store {
|
|| s.name != g.name
|
||||||
let s = &source.chunks[src];
|
|| s.entity_type != g.entity_type
|
||||||
let c = &mem.cache;
|
|| s.embedding_idx != g.embedding_idx
|
||||||
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());
|
|
||||||
}
|
|
||||||
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());
|
return Err(format!("entity[{i}] mismatch").into());
|
||||||
}
|
}
|
||||||
rows_checked += 1;
|
rows_checked += 1;
|
||||||
}
|
}
|
||||||
|
for (i, (s, g)) in source
|
||||||
// ---- Knowledge graph ----
|
.relations
|
||||||
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()
|
.iter()
|
||||||
.copied()
|
.zip(got.relations.iter())
|
||||||
.find(|&(idx, _)| dim > 0 && mem.cache.tombstones[idx] == 0);
|
.enumerate()
|
||||||
let search_checked = match probe {
|
{
|
||||||
None => false,
|
if s.src != g.src || s.tgt != g.tgt || s.relation != g.relation {
|
||||||
Some((idx, _)) => {
|
return Err(format!("relation[{i}] mismatch").into());
|
||||||
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
|
rows_checked += 1;
|
||||||
}
|
}
|
||||||
};
|
|
||||||
|
|
||||||
Ok(ValidationSummary {
|
Ok(ValidationSummary {
|
||||||
count: mem.count(),
|
chunks: got.chunks.len() as u64,
|
||||||
active: mem.count_active(),
|
sessions: got.sessions.len() as u64,
|
||||||
sessions: mem.sessions().len(),
|
entities: got.entities.len() as u64,
|
||||||
entities: mem.knowledge().entities.len(),
|
relations: got.relations.len() as u64,
|
||||||
relations: mem.knowledge().relations.len(),
|
embedding_dim: got.embedding_dim as u64,
|
||||||
embedding_dim: dim,
|
|
||||||
float16,
|
|
||||||
rows_checked,
|
rows_checked,
|
||||||
search_checked,
|
provenance_verified,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
fn check_count(kind: &str, got: usize, expected: usize) -> Result<(), BoxErr> {
|
fn check_count(kind: &str, got: usize, expected: usize) -> Result<(), BoxErr> {
|
||||||
if got != expected {
|
if got != expected {
|
||||||
return Err(format!("{kind} count mismatch: store has {got}, expected {expected}").into());
|
return Err(format!("{kind} count mismatch: HDF5 has {got}, source has {expected}").into());
|
||||||
}
|
}
|
||||||
Ok(())
|
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 {
|
fn truncate(s: &str) -> String {
|
||||||
if s.len() <= 40 {
|
if s.len() <= 40 {
|
||||||
s.to_string()
|
s.to_string()
|
||||||
@@ -270,7 +196,7 @@ fn truncate(s: &str) -> String {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Indices of records to content-check. Full = all; otherwise a spread of
|
/// Indices of chunk rows to content-check. Full = all; otherwise a spread of
|
||||||
/// representative rows (first/last and evenly-spaced interior samples).
|
/// representative rows (first/last and evenly-spaced interior samples).
|
||||||
fn sample_indices(n: usize, full: bool) -> Vec<usize> {
|
fn sample_indices(n: usize, full: bool) -> Vec<usize> {
|
||||||
if n == 0 {
|
if n == 0 {
|
||||||
|
|||||||
Reference in New Issue
Block a user