ClawHDF5 Coding Agent
45a38ba260
blas_cosine_batch and accelerate_cosine_batch_vecs re-flattened the entire Vec<Vec<f32>> corpus into a fresh Vec<f32> on every single query before running the batch matmul — an O(N·dim) copy paid per query when fast-math/accelerate/openblas is enabled, even though a flat fast-path (blas_cosine_batch_flat / accelerate_cosine_batch) already existed for pre-flattened input. Add MemoryCache::embeddings_flat, a contiguous [N × embedding_dim] buffer maintained incrementally in push/update/compact (O(1) amortized append, O(dim) in-place overwrite, O(n) rebuild only on compact/bulk load). schema.rs's direct-push load path calls the new rebuild_flat() explicitly. flat_embeddings() now just clones the already-maintained buffer instead of rebuilding it. Thread the flat buffer through strategy::search_with_metrics as a new vectors_flat parameter, used only by the Blas/Accelerate arms (now calling the *_flat variants); other strategies are unaffected. No current caller wires search_with_metrics into the production query path yet (only its own tests exercise it) — this fixes the identified per-query re-flatten and makes the flat buffer available for whenever that wiring lands. INT-16
329 lines
10 KiB
Rust
329 lines
10 KiB
Rust
//! In-memory cache for memory entries, sessions, and knowledge graph.
|
||
|
||
use crate::vector_search;
|
||
|
||
/// In-memory cache for the /memory group data.
|
||
#[derive(Debug, Clone)]
|
||
pub struct MemoryCache {
|
||
pub chunks: Vec<String>,
|
||
pub embeddings: Vec<Vec<f32>>,
|
||
/// `embeddings` flattened into one contiguous `[N × embedding_dim]`
|
||
/// buffer, maintained incrementally alongside `embeddings` (push/update/
|
||
/// compact) so BLAS/Accelerate batch search can read it directly instead
|
||
/// of re-flattening the whole corpus on every query.
|
||
pub embeddings_flat: Vec<f32>,
|
||
pub source_channels: Vec<String>,
|
||
pub timestamps: Vec<f64>,
|
||
pub session_ids: Vec<String>,
|
||
pub tags: Vec<String>,
|
||
pub tombstones: Vec<u8>,
|
||
pub embedding_dim: usize,
|
||
/// Pre-computed L2 norms for each embedding.
|
||
pub norms: Vec<f32>,
|
||
/// Hebbian activation weights (default 1.0 per entry).
|
||
pub activation_weights: Vec<f32>,
|
||
}
|
||
|
||
impl MemoryCache {
|
||
pub fn new(embedding_dim: usize) -> Self {
|
||
Self {
|
||
chunks: Vec::new(),
|
||
embeddings: Vec::new(),
|
||
embeddings_flat: Vec::new(),
|
||
source_channels: Vec::new(),
|
||
timestamps: Vec::new(),
|
||
session_ids: Vec::new(),
|
||
tags: Vec::new(),
|
||
tombstones: Vec::new(),
|
||
embedding_dim,
|
||
norms: Vec::new(),
|
||
activation_weights: Vec::new(),
|
||
}
|
||
}
|
||
|
||
/// Rebuild `embeddings_flat` from `embeddings` from scratch. Callers that
|
||
/// populate `embeddings` directly (bulk loads) must call this afterward.
|
||
pub fn rebuild_flat(&mut self) {
|
||
self.embeddings_flat.clear();
|
||
self.embeddings_flat
|
||
.reserve(self.embeddings.len() * self.embedding_dim);
|
||
for emb in &self.embeddings {
|
||
self.embeddings_flat.extend_from_slice(emb);
|
||
}
|
||
}
|
||
|
||
/// Total number of entries (including tombstoned).
|
||
pub fn len(&self) -> usize {
|
||
self.chunks.len()
|
||
}
|
||
|
||
pub fn is_empty(&self) -> bool {
|
||
self.chunks.is_empty()
|
||
}
|
||
|
||
/// Number of active (non-tombstoned) entries.
|
||
pub fn count_active(&self) -> usize {
|
||
self.tombstones.iter().filter(|&&t| t == 0).count()
|
||
}
|
||
|
||
/// Push a new entry, returns its index.
|
||
pub fn push(
|
||
&mut self,
|
||
chunk: String,
|
||
embedding: Vec<f32>,
|
||
source_channel: String,
|
||
timestamp: f64,
|
||
session_id: String,
|
||
tags: String,
|
||
) -> usize {
|
||
let idx = self.chunks.len();
|
||
let norm = vector_search::compute_norm(&embedding);
|
||
self.chunks.push(chunk);
|
||
self.embeddings_flat.extend_from_slice(&embedding);
|
||
self.embeddings.push(embedding);
|
||
self.source_channels.push(source_channel);
|
||
self.timestamps.push(timestamp);
|
||
self.session_ids.push(session_id);
|
||
self.tags.push(tags);
|
||
self.tombstones.push(0);
|
||
self.norms.push(norm);
|
||
self.activation_weights.push(1.0);
|
||
idx
|
||
}
|
||
|
||
/// Find an active (non-tombstoned) entry by tags (used as key for dedup).
|
||
/// Returns the index of the first matching active entry, or None.
|
||
pub fn find_by_tags(&self, tags: &str) -> Option<usize> {
|
||
if tags.is_empty() {
|
||
return None;
|
||
}
|
||
for (i, t) in self.tags.iter().enumerate() {
|
||
if self.tombstones.get(i).copied().unwrap_or(1) == 0 && t == tags {
|
||
return Some(i);
|
||
}
|
||
}
|
||
None
|
||
}
|
||
|
||
/// Update an existing entry in-place (for upsert dedup).
|
||
pub fn update(
|
||
&mut self,
|
||
idx: usize,
|
||
chunk: String,
|
||
embedding: Vec<f32>,
|
||
source_channel: String,
|
||
timestamp: f64,
|
||
session_id: String,
|
||
) {
|
||
if idx < self.chunks.len() {
|
||
let norm = vector_search::compute_norm(&embedding);
|
||
self.chunks[idx] = chunk;
|
||
let dim = self.embedding_dim;
|
||
let flat_start = idx * dim;
|
||
let matches_dim =
|
||
embedding.len() == dim && flat_start + dim <= self.embeddings_flat.len();
|
||
self.embeddings[idx] = embedding;
|
||
if matches_dim {
|
||
self.embeddings_flat[flat_start..flat_start + dim]
|
||
.copy_from_slice(&self.embeddings[idx]);
|
||
} else {
|
||
// Embedding length doesn't match embedding_dim (shouldn't
|
||
// happen in practice) — fall back to a full rebuild rather
|
||
// than leave embeddings_flat misaligned with embeddings.
|
||
self.rebuild_flat();
|
||
}
|
||
self.source_channels[idx] = source_channel;
|
||
self.timestamps[idx] = timestamp;
|
||
self.session_ids[idx] = session_id;
|
||
self.norms[idx] = norm;
|
||
self.activation_weights[idx] = 1.0; // reset activation on update
|
||
}
|
||
}
|
||
|
||
/// Mark an entry as deleted (tombstoned).
|
||
pub fn mark_deleted(&mut self, id: usize) -> bool {
|
||
if id < self.tombstones.len() && self.tombstones[id] == 0 {
|
||
self.tombstones[id] = 1;
|
||
true
|
||
} else {
|
||
false
|
||
}
|
||
}
|
||
|
||
/// Fraction of entries that are tombstoned.
|
||
pub fn tombstone_fraction(&self) -> f32 {
|
||
if self.chunks.is_empty() {
|
||
return 0.0;
|
||
}
|
||
let tombstoned = self.tombstones.iter().filter(|&&t| t == 1).count();
|
||
tombstoned as f32 / self.chunks.len() as f32
|
||
}
|
||
|
||
/// Remove all tombstoned entries, returns number removed.
|
||
/// Also returns a mapping from old indices to new indices (None if removed).
|
||
/// Recomputes norms for remaining entries.
|
||
pub fn compact(&mut self) -> (usize, Vec<Option<usize>>) {
|
||
let old_len = self.chunks.len();
|
||
let mut index_map = vec![None; old_len];
|
||
let mut new_idx = 0usize;
|
||
|
||
let mut new_chunks = Vec::new();
|
||
let mut new_embeddings = Vec::new();
|
||
let mut new_source_channels = Vec::new();
|
||
let mut new_timestamps = Vec::new();
|
||
let mut new_session_ids = Vec::new();
|
||
let mut new_tags = Vec::new();
|
||
let mut new_tombstones = Vec::new();
|
||
let mut new_norms = Vec::new();
|
||
let mut new_activation_weights = Vec::new();
|
||
|
||
for (i, slot) in index_map.iter_mut().enumerate() {
|
||
if self.tombstones[i] == 0 {
|
||
*slot = Some(new_idx);
|
||
new_idx += 1;
|
||
let norm = vector_search::compute_norm(&self.embeddings[i]);
|
||
new_chunks.push(self.chunks[i].clone());
|
||
new_embeddings.push(self.embeddings[i].clone());
|
||
new_source_channels.push(self.source_channels[i].clone());
|
||
new_timestamps.push(self.timestamps[i]);
|
||
new_session_ids.push(self.session_ids[i].clone());
|
||
new_tags.push(self.tags[i].clone());
|
||
new_tombstones.push(0u8);
|
||
new_norms.push(norm);
|
||
new_activation_weights.push(self.activation_weights[i]);
|
||
}
|
||
}
|
||
|
||
let removed = old_len - new_chunks.len();
|
||
self.chunks = new_chunks;
|
||
self.embeddings = new_embeddings;
|
||
self.source_channels = new_source_channels;
|
||
self.timestamps = new_timestamps;
|
||
self.session_ids = new_session_ids;
|
||
self.tags = new_tags;
|
||
self.tombstones = new_tombstones;
|
||
self.norms = new_norms;
|
||
self.activation_weights = new_activation_weights;
|
||
self.rebuild_flat();
|
||
|
||
(removed, index_map)
|
||
}
|
||
|
||
/// Flatten all embeddings into a single Vec<f32> for HDF5 storage.
|
||
/// `embeddings_flat` is already maintained incrementally, so this just
|
||
/// clones it — kept as a method for callers that want an owned copy.
|
||
pub fn flat_embeddings(&self) -> Vec<f32> {
|
||
self.embeddings_flat.clone()
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
/// `embeddings_flat` must always equal a from-scratch flatten of `embeddings`.
|
||
fn assert_flat_in_sync(cache: &MemoryCache) {
|
||
let expected: Vec<f32> = cache.embeddings.iter().flatten().copied().collect();
|
||
assert_eq!(cache.embeddings_flat, expected);
|
||
}
|
||
|
||
#[test]
|
||
fn push_keeps_flat_buffer_in_sync() {
|
||
let mut cache = MemoryCache::new(3);
|
||
cache.push(
|
||
"a".into(),
|
||
vec![1.0, 2.0, 3.0],
|
||
"chan".into(),
|
||
0.0,
|
||
"s1".into(),
|
||
String::new(),
|
||
);
|
||
cache.push(
|
||
"b".into(),
|
||
vec![4.0, 5.0, 6.0],
|
||
"chan".into(),
|
||
1.0,
|
||
"s1".into(),
|
||
String::new(),
|
||
);
|
||
assert_flat_in_sync(&cache);
|
||
assert_eq!(cache.embeddings_flat, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
|
||
}
|
||
|
||
#[test]
|
||
fn update_keeps_flat_buffer_in_sync() {
|
||
let mut cache = MemoryCache::new(3);
|
||
cache.push(
|
||
"a".into(),
|
||
vec![1.0, 2.0, 3.0],
|
||
"chan".into(),
|
||
0.0,
|
||
"s1".into(),
|
||
String::new(),
|
||
);
|
||
cache.push(
|
||
"b".into(),
|
||
vec![4.0, 5.0, 6.0],
|
||
"chan".into(),
|
||
1.0,
|
||
"s1".into(),
|
||
String::new(),
|
||
);
|
||
cache.update(
|
||
0,
|
||
"a2".into(),
|
||
vec![7.0, 8.0, 9.0],
|
||
"chan".into(),
|
||
2.0,
|
||
"s1".into(),
|
||
);
|
||
assert_flat_in_sync(&cache);
|
||
assert_eq!(
|
||
cache.embeddings_flat,
|
||
vec![7.0, 8.0, 9.0, 4.0, 5.0, 6.0],
|
||
"update must overwrite the correct flat slice, not just append"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn compact_keeps_flat_buffer_in_sync() {
|
||
let mut cache = MemoryCache::new(2);
|
||
cache.push(
|
||
"a".into(),
|
||
vec![1.0, 1.0],
|
||
"chan".into(),
|
||
0.0,
|
||
"s1".into(),
|
||
String::new(),
|
||
);
|
||
cache.push(
|
||
"b".into(),
|
||
vec![2.0, 2.0],
|
||
"chan".into(),
|
||
1.0,
|
||
"s1".into(),
|
||
String::new(),
|
||
);
|
||
cache.push(
|
||
"c".into(),
|
||
vec![3.0, 3.0],
|
||
"chan".into(),
|
||
2.0,
|
||
"s1".into(),
|
||
String::new(),
|
||
);
|
||
cache.mark_deleted(1);
|
||
cache.compact();
|
||
assert_flat_in_sync(&cache);
|
||
assert_eq!(cache.embeddings_flat, vec![1.0, 1.0, 3.0, 3.0]);
|
||
}
|
||
|
||
#[test]
|
||
fn rebuild_flat_matches_manual_flatten() {
|
||
let mut cache = MemoryCache::new(2);
|
||
cache.embeddings = vec![vec![1.0, 2.0], vec![3.0, 4.0]];
|
||
cache.rebuild_flat();
|
||
assert_eq!(cache.embeddings_flat, vec![1.0, 2.0, 3.0, 4.0]);
|
||
}
|
||
}
|