feat(agent): expose the HNSW parameters in MemoryConfig
Graph degree and the build- and query-time candidate list sizes were constants, so a deployment had no way to trade recall against memory or query speed. They are now `MemoryConfig::hnsw_m`, `hnsw_ef_construction` and `hnsw_ef_search`, persisted with the store and defaulting to exactly the previous behaviour (16, 64, and a query list that scales with `k`). Two things the straightforward version would have got wrong: `clawhdf5-ann` asserts a graph degree of at least 2, so a configured 0 — from a file, or from a caller reading 0 as "use the default" — aborted the process inside the index builder. The store clamps instead, and a test covers it: removing the clamp makes that test panic rather than fail. `ef_search` and the candidate pool handed to score fusion were the same number. Tying the pool to the new setting would mean lowering `ef` for speed also narrows what fusion sees, quietly degrading hybrid results through a knob that looks like it only costs time. They are now independent. Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
This commit is contained in:
@@ -119,6 +119,9 @@ mod tests {
|
||||
wal_enabled: false,
|
||||
wal_max_entries: 500,
|
||||
quantized_index: false,
|
||||
hnsw_m: 16,
|
||||
hnsw_ef_construction: 64,
|
||||
hnsw_ef_search: 0,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -65,12 +65,6 @@ use cache::MemoryCache;
|
||||
use clawhdf5_ann::{DistanceMetric, HnswIndex, Storage};
|
||||
use ephemeral::{EphemeralConfig, EphemeralStore};
|
||||
|
||||
/// HNSW construction parameters used for the agent's vector index. Cosine is the
|
||||
/// agent's similarity metric, so the index is built with cosine distance.
|
||||
#[cfg(feature = "hnsw")]
|
||||
const HNSW_M: usize = 16;
|
||||
#[cfg(feature = "hnsw")]
|
||||
const HNSW_EF_CONSTRUCTION: usize = 64;
|
||||
// EphemeralEntry and EphemeralStats are part of the crate public API via
|
||||
// the `ephemeral` module; they are not needed directly in lib.rs internals.
|
||||
#[allow(unused_imports)]
|
||||
@@ -150,6 +144,23 @@ pub struct MemoryConfig {
|
||||
///
|
||||
/// Has no effect without the `hnsw` feature.
|
||||
pub quantized_index: bool,
|
||||
/// HNSW graph degree. Higher means a denser graph: better recall, more
|
||||
/// memory and slower builds. Clamped to at least 2 when the index is
|
||||
/// built, since a graph with fewer connections is not one.
|
||||
///
|
||||
/// Has no effect without the `hnsw` feature.
|
||||
pub hnsw_m: usize,
|
||||
/// Candidate list size while building the HNSW graph. Higher means a
|
||||
/// better graph and a slower build; it does not affect query cost.
|
||||
///
|
||||
/// Has no effect without the `hnsw` feature.
|
||||
pub hnsw_ef_construction: usize,
|
||||
/// Candidate list size for a query, trading throughput for recall. `0`
|
||||
/// keeps the default, which scales with the requested `k`
|
||||
/// (`max(k * 8, 64)`) so that fusion still sees a useful pool.
|
||||
///
|
||||
/// Has no effect without the `hnsw` feature.
|
||||
pub hnsw_ef_search: usize,
|
||||
}
|
||||
|
||||
impl MemoryConfig {
|
||||
@@ -172,6 +183,9 @@ impl MemoryConfig {
|
||||
wal_enabled: true,
|
||||
wal_max_entries: 500,
|
||||
quantized_index: false,
|
||||
hnsw_m: 16,
|
||||
hnsw_ef_construction: 64,
|
||||
hnsw_ef_search: 0,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -826,6 +840,32 @@ impl HDF5Memory {
|
||||
// the index length drifts from the cache length (covering any mutation path
|
||||
// that doesn't call a hook, e.g. consolidation pushes).
|
||||
|
||||
/// Graph degree for the index, never below the 2 the builder requires:
|
||||
/// a config value of 0 or 1 would otherwise panic inside `clawhdf5-ann`.
|
||||
#[cfg(feature = "hnsw")]
|
||||
fn hnsw_m(&self) -> usize {
|
||||
self.config.hnsw_m.max(2)
|
||||
}
|
||||
|
||||
/// Build-time candidate list size, never below the graph degree — a
|
||||
/// smaller one cannot fill a node's connections.
|
||||
#[cfg(feature = "hnsw")]
|
||||
fn hnsw_ef_construction(&self) -> usize {
|
||||
self.config.hnsw_ef_construction.max(self.hnsw_m())
|
||||
}
|
||||
|
||||
/// Query-time candidate list size for a `k`-result search. `0` means the
|
||||
/// default, which scales with `k`.
|
||||
#[cfg(feature = "hnsw")]
|
||||
pub(crate) fn hnsw_ef_search(&self, k: usize) -> usize {
|
||||
let default = (k * 8).max(64);
|
||||
if self.config.hnsw_ef_search == 0 {
|
||||
default
|
||||
} else {
|
||||
self.config.hnsw_ef_search.max(k)
|
||||
}
|
||||
}
|
||||
|
||||
/// How the index should store its copy of the vectors, per the config.
|
||||
#[cfg(feature = "hnsw")]
|
||||
fn index_storage(&self) -> Storage {
|
||||
@@ -856,8 +896,8 @@ impl HDF5Memory {
|
||||
let rows: Vec<Vec<f32>> = self.cache.embeddings.iter().map(<[f32]>::to_vec).collect();
|
||||
let mut index = HnswIndex::build_with(
|
||||
&rows,
|
||||
HNSW_M,
|
||||
HNSW_EF_CONSTRUCTION,
|
||||
self.hnsw_m(),
|
||||
self.hnsw_ef_construction(),
|
||||
DistanceMetric::Cosine,
|
||||
self.index_storage(),
|
||||
);
|
||||
|
||||
@@ -108,6 +108,15 @@ pub fn build_hdf5_file_with_meta(
|
||||
"quantized_index",
|
||||
AttrValue::I64(config.quantized_index.into()),
|
||||
);
|
||||
meta.set_attr("hnsw_m", AttrValue::I64(config.hnsw_m as i64));
|
||||
meta.set_attr(
|
||||
"hnsw_ef_construction",
|
||||
AttrValue::I64(config.hnsw_ef_construction as i64),
|
||||
);
|
||||
meta.set_attr(
|
||||
"hnsw_ef_search",
|
||||
AttrValue::I64(config.hnsw_ef_search as i64),
|
||||
);
|
||||
meta.set_attr(
|
||||
"edgehdf5_version",
|
||||
AttrValue::String(ZEROCLAW_VERSION.into()),
|
||||
@@ -489,6 +498,15 @@ pub fn validate_and_load(
|
||||
.and_then(|v| usize::try_from(v).ok())
|
||||
.unwrap_or(500),
|
||||
quantized_index: optional_bool_attr(&attrs, "quantized_index", false),
|
||||
hnsw_m: optional_i64_attr(&attrs, "hnsw_m")
|
||||
.and_then(|v| usize::try_from(v).ok())
|
||||
.unwrap_or(16),
|
||||
hnsw_ef_construction: optional_i64_attr(&attrs, "hnsw_ef_construction")
|
||||
.and_then(|v| usize::try_from(v).ok())
|
||||
.unwrap_or(64),
|
||||
hnsw_ef_search: optional_i64_attr(&attrs, "hnsw_ef_search")
|
||||
.and_then(|v| usize::try_from(v).ok())
|
||||
.unwrap_or(0),
|
||||
};
|
||||
|
||||
// Load /memory group
|
||||
|
||||
@@ -28,8 +28,12 @@ impl HDF5Memory {
|
||||
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 candidates = index.search(query_embedding, pool, pool);
|
||||
let ef = self.hnsw_ef_search(k).max(pool);
|
||||
let candidates = index.search(query_embedding, pool, ef);
|
||||
// A quantised index returns approximate distances, and no
|
||||
// amount of `ef` fixes that — the loss is in the distances,
|
||||
// not the graph. Re-score the pool against the cache's exact
|
||||
|
||||
@@ -234,3 +234,53 @@ fn quantized_index_setting_survives_a_reopen() {
|
||||
let reopened = HDF5Memory::open(&path).unwrap();
|
||||
assert!(reopened.config().quantized_index);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hnsw_parameters_are_configurable_and_persisted() {
|
||||
// The graph degree and both candidate-list sizes used to be constants, so
|
||||
// a deployment could not trade recall against memory or speed at all.
|
||||
let dir = TempDir::new().unwrap();
|
||||
let path = dir.path().join("mem.h5");
|
||||
let mut config = MemoryConfig::new(path.clone(), "agent", 16);
|
||||
config.hnsw_m = 8;
|
||||
config.hnsw_ef_construction = 32;
|
||||
config.hnsw_ef_search = 128;
|
||||
let mut mem = HDF5Memory::create(config).unwrap();
|
||||
|
||||
let mut seed = 99;
|
||||
let vectors: Vec<Vec<f32>> = (0..300).map(|_| make_vector(&mut seed, 16)).collect();
|
||||
for (i, v) in vectors.iter().enumerate() {
|
||||
mem.save(entry(&format!("c{i}"), v.clone(), "t")).unwrap();
|
||||
}
|
||||
// Still correct with a smaller graph: an exact match must rank first.
|
||||
let top = mem.hybrid_search(&vectors[42], "", 1.0, 0.0, 1);
|
||||
assert_eq!(top[0].index, 42);
|
||||
|
||||
mem.flush_wal().unwrap();
|
||||
drop(mem);
|
||||
let reopened = HDF5Memory::open(&path).unwrap();
|
||||
assert_eq!(reopened.config().hnsw_m, 8);
|
||||
assert_eq!(reopened.config().hnsw_ef_construction, 32);
|
||||
assert_eq!(reopened.config().hnsw_ef_search, 128);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn degenerate_hnsw_parameters_do_not_panic() {
|
||||
// `clawhdf5-ann` asserts m >= 2, so a zero from a config file — or from a
|
||||
// caller who assumed 0 meant "default" — would abort the process inside
|
||||
// the index builder. The store clamps instead.
|
||||
let dir = TempDir::new().unwrap();
|
||||
let mut config = MemoryConfig::new(dir.path().join("mem.h5"), "agent", 8);
|
||||
config.hnsw_m = 0;
|
||||
config.hnsw_ef_construction = 0;
|
||||
config.hnsw_ef_search = 1;
|
||||
let mut mem = HDF5Memory::create(config).unwrap();
|
||||
|
||||
let mut seed = 5;
|
||||
let vectors: Vec<Vec<f32>> = (0..50).map(|_| make_vector(&mut seed, 8)).collect();
|
||||
for (i, v) in vectors.iter().enumerate() {
|
||||
mem.save(entry(&format!("c{i}"), v.clone(), "t")).unwrap();
|
||||
}
|
||||
let results = mem.hybrid_search(&vectors[7], "", 1.0, 0.0, 5);
|
||||
assert_eq!(results[0].index, 7, "exact match should still rank first");
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user