Skip to main content

edgehdf5_memory/
cache.rs

1//! In-memory cache for memory entries, sessions, and knowledge graph.
2
3use crate::vector_search;
4
5/// In-memory cache for the /memory group data.
6#[derive(Debug, Clone)]
7pub struct MemoryCache {
8    pub chunks: Vec<String>,
9    pub embeddings: Vec<Vec<f32>>,
10    pub source_channels: Vec<String>,
11    pub timestamps: Vec<f64>,
12    pub session_ids: Vec<String>,
13    pub tags: Vec<String>,
14    pub tombstones: Vec<u8>,
15    pub embedding_dim: usize,
16    /// Pre-computed L2 norms for each embedding.
17    pub norms: Vec<f32>,
18    /// Hebbian activation weights (default 1.0 per entry).
19    pub activation_weights: Vec<f32>,
20}
21
22impl MemoryCache {
23    pub fn new(embedding_dim: usize) -> Self {
24        Self {
25            chunks: Vec::new(),
26            embeddings: Vec::new(),
27            source_channels: Vec::new(),
28            timestamps: Vec::new(),
29            session_ids: Vec::new(),
30            tags: Vec::new(),
31            tombstones: Vec::new(),
32            embedding_dim,
33            norms: Vec::new(),
34            activation_weights: Vec::new(),
35        }
36    }
37
38    /// Total number of entries (including tombstoned).
39    pub fn len(&self) -> usize {
40        self.chunks.len()
41    }
42
43    pub fn is_empty(&self) -> bool {
44        self.chunks.is_empty()
45    }
46
47    /// Number of active (non-tombstoned) entries.
48    pub fn count_active(&self) -> usize {
49        self.tombstones.iter().filter(|&&t| t == 0).count()
50    }
51
52    /// Push a new entry, returns its index.
53    pub fn push(
54        &mut self,
55        chunk: String,
56        embedding: Vec<f32>,
57        source_channel: String,
58        timestamp: f64,
59        session_id: String,
60        tags: String,
61    ) -> usize {
62        let idx = self.chunks.len();
63        let norm = vector_search::compute_norm(&embedding);
64        self.chunks.push(chunk);
65        self.embeddings.push(embedding);
66        self.source_channels.push(source_channel);
67        self.timestamps.push(timestamp);
68        self.session_ids.push(session_id);
69        self.tags.push(tags);
70        self.tombstones.push(0);
71        self.norms.push(norm);
72        self.activation_weights.push(1.0);
73        idx
74    }
75
76    /// Mark an entry as deleted (tombstoned).
77    pub fn mark_deleted(&mut self, id: usize) -> bool {
78        if id < self.tombstones.len() && self.tombstones[id] == 0 {
79            self.tombstones[id] = 1;
80            true
81        } else {
82            false
83        }
84    }
85
86    /// Fraction of entries that are tombstoned.
87    pub fn tombstone_fraction(&self) -> f32 {
88        if self.chunks.is_empty() {
89            return 0.0;
90        }
91        let tombstoned = self.tombstones.iter().filter(|&&t| t == 1).count();
92        tombstoned as f32 / self.chunks.len() as f32
93    }
94
95    /// Remove all tombstoned entries, returns number removed.
96    /// Also returns a mapping from old indices to new indices (None if removed).
97    /// Recomputes norms for remaining entries.
98    pub fn compact(&mut self) -> (usize, Vec<Option<usize>>) {
99        let old_len = self.chunks.len();
100        let mut index_map = vec![None; old_len];
101        let mut new_idx = 0usize;
102
103        let mut new_chunks = Vec::new();
104        let mut new_embeddings = Vec::new();
105        let mut new_source_channels = Vec::new();
106        let mut new_timestamps = Vec::new();
107        let mut new_session_ids = Vec::new();
108        let mut new_tags = Vec::new();
109        let mut new_tombstones = Vec::new();
110        let mut new_norms = Vec::new();
111        let mut new_activation_weights = Vec::new();
112
113        for (i, slot) in index_map.iter_mut().enumerate() {
114            if self.tombstones[i] == 0 {
115                *slot = Some(new_idx);
116                new_idx += 1;
117                let norm = vector_search::compute_norm(&self.embeddings[i]);
118                new_chunks.push(self.chunks[i].clone());
119                new_embeddings.push(self.embeddings[i].clone());
120                new_source_channels.push(self.source_channels[i].clone());
121                new_timestamps.push(self.timestamps[i]);
122                new_session_ids.push(self.session_ids[i].clone());
123                new_tags.push(self.tags[i].clone());
124                new_tombstones.push(0u8);
125                new_norms.push(norm);
126                new_activation_weights.push(self.activation_weights[i]);
127            }
128        }
129
130        let removed = old_len - new_chunks.len();
131        self.chunks = new_chunks;
132        self.embeddings = new_embeddings;
133        self.source_channels = new_source_channels;
134        self.timestamps = new_timestamps;
135        self.session_ids = new_session_ids;
136        self.tags = new_tags;
137        self.tombstones = new_tombstones;
138        self.norms = new_norms;
139        self.activation_weights = new_activation_weights;
140
141        (removed, index_map)
142    }
143
144    /// Flatten all embeddings into a single Vec<f32> for HDF5 storage.
145    pub fn flat_embeddings(&self) -> Vec<f32> {
146        let mut flat = Vec::with_capacity(self.embeddings.len() * self.embedding_dim);
147        for emb in &self.embeddings {
148            flat.extend_from_slice(emb);
149        }
150        flat
151    }
152}