mentedb_index/
manager.rs

1//! Composite index manager that owns and coordinates all index types.
2
3use std::collections::HashSet;
4use std::path::Path;
5
6use mentedb_core::MemoryNode;
7use mentedb_core::error::MenteResult;
8use mentedb_core::types::{MemoryId, Timestamp};
9
10use crate::bitmap::BitmapIndex;
11use crate::hnsw::{HnswConfig, HnswIndex};
12use crate::salience::SalienceIndex;
13use crate::temporal::TemporalIndex;
14
15/// Configuration for the composite index manager.
16#[derive(Default)]
17pub struct IndexManagerConfig {
18    /// HNSW configuration parameters.
19    pub hnsw: HnswConfig,
20}
21
22/// Owns all index types and provides unified indexing and hybrid search.
23pub struct IndexManager {
24    /// Vector similarity index.
25    pub hnsw: HnswIndex,
26    /// Tag and attribute bitmap index.
27    pub bitmap: BitmapIndex,
28    /// Timestamp range index.
29    pub temporal: TemporalIndex,
30    /// Importance score index.
31    pub salience: SalienceIndex,
32}
33
34impl IndexManager {
35    /// Create a new index manager with the given configuration.
36    pub fn new(config: IndexManagerConfig) -> Self {
37        Self {
38            hnsw: HnswIndex::new(config.hnsw),
39            bitmap: BitmapIndex::new(),
40            temporal: TemporalIndex::new(),
41            salience: SalienceIndex::new(),
42        }
43    }
44
45    /// Save all indexes to the given directory.
46    pub fn save(&self, dir: &Path) -> MenteResult<()> {
47        std::fs::create_dir_all(dir)?;
48        self.hnsw.save(&dir.join("hnsw.json"))?;
49        self.bitmap.save(&dir.join("bitmap.json"))?;
50        self.temporal.save(&dir.join("temporal.json"))?;
51        self.salience.save(&dir.join("salience.json"))?;
52        Ok(())
53    }
54
55    /// Load all indexes from the given directory.
56    pub fn load(dir: &Path) -> MenteResult<Self> {
57        let hnsw = HnswIndex::load(&dir.join("hnsw.json"), HnswConfig::default().ef_search)?;
58        let bitmap = BitmapIndex::load(&dir.join("bitmap.json"))?;
59        let temporal = TemporalIndex::load(&dir.join("temporal.json"))?;
60        let salience = SalienceIndex::load(&dir.join("salience.json"))?;
61        Ok(Self {
62            hnsw,
63            bitmap,
64            temporal,
65            salience,
66        })
67    }
68
69    /// Index a memory node across all indexes.
70    pub fn index_memory(&self, node: &MemoryNode) {
71        // Vector index
72        if !node.embedding.is_empty() {
73            let _ = self.hnsw.insert(node.id, &node.embedding);
74        }
75
76        // Tag bitmap index
77        for tag in &node.tags {
78            self.bitmap.add_tag(node.id, tag);
79        }
80
81        // Temporal index
82        self.temporal.insert(node.id, node.created_at);
83
84        // Salience index
85        self.salience.insert(node.id, node.salience);
86    }
87
88    /// Remove a memory from all indexes.
89    pub fn remove_memory(&self, id: MemoryId, node: &MemoryNode) {
90        let _ = self.hnsw.remove(id);
91        self.bitmap.remove_all(id);
92        self.temporal.remove(id, node.created_at);
93        self.salience.remove(id, node.salience);
94    }
95
96    /// Hybrid search combining vector similarity, tag filtering, time range, and salience.
97    ///
98    /// Strategy:
99    /// 1. Vector search to get top k*4 candidates
100    /// 2. Filter by tags (if provided) and time range (if provided)
101    /// 3. Re-rank by combined score: vector_sim * 0.6 + salience * 0.3 + recency * 0.1
102    /// 4. Return top k results
103    pub fn hybrid_search(
104        &self,
105        query_embedding: &[f32],
106        tags: Option<&[&str]>,
107        time_range: Option<(Timestamp, Timestamp)>,
108        k: usize,
109    ) -> Vec<(MemoryId, f32)> {
110        if k == 0 {
111            return Vec::new();
112        }
113
114        // Step 1: vector search for top k*4 candidates
115        let vector_candidates = self.hnsw.search(query_embedding, k * 4);
116
117        if vector_candidates.is_empty() {
118            return Vec::new();
119        }
120
121        // Build set of tag-filtered ids (if tags are specified)
122        let tag_filter: Option<HashSet<MemoryId>> = tags.map(|t| {
123            if t.is_empty() {
124                HashSet::new()
125            } else {
126                self.bitmap.query_tags_and(t).into_iter().collect()
127            }
128        });
129
130        // Build set of time-range-filtered ids (if time range is specified)
131        let time_filter: Option<HashSet<MemoryId>> =
132            time_range.map(|(start, end)| self.temporal.range(start, end).into_iter().collect());
133
134        // Find the max distance for normalization
135        let max_dist = vector_candidates
136            .iter()
137            .map(|(_, d)| *d)
138            .fold(f32::NEG_INFINITY, f32::max)
139            .max(f32::EPSILON);
140
141        // Find the latest timestamp among candidates for recency normalization
142        let max_ts = vector_candidates
143            .iter()
144            .filter_map(|(id, _)| self.temporal.get_timestamp(*id))
145            .max()
146            .unwrap_or(1) as f64;
147
148        // Step 2 & 3: filter and re-rank
149        let mut scored: Vec<(MemoryId, f32)> = vector_candidates
150            .into_iter()
151            .filter(|(id, _)| {
152                if let Some(ref tf) = tag_filter
153                    && !tf.contains(id)
154                {
155                    return false;
156                }
157                if let Some(ref trf) = time_filter
158                    && !trf.contains(id)
159                {
160                    return false;
161                }
162                true
163            })
164            .map(|(id, dist)| {
165                // Vector similarity: 1.0 - normalized distance
166                let vector_sim = 1.0 - (dist / max_dist);
167
168                // Salience score (default 0.5 if not found)
169                let salience = self.salience.get_salience(id).unwrap_or(0.5);
170
171                // Recency: normalize timestamp to [0, 1]
172                let ts = self.temporal.get_timestamp(id).unwrap_or(0) as f64;
173                let recency = if max_ts > 0.0 {
174                    (ts / max_ts) as f32
175                } else {
176                    0.0
177                };
178
179                let combined = vector_sim * 0.6 + salience * 0.3 + recency * 0.1;
180                (id, combined)
181            })
182            .collect();
183
184        // Sort descending by combined score
185        scored.sort_unstable_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
186        scored.truncate(k);
187        scored
188    }
189}
190
191impl Default for IndexManager {
192    fn default() -> Self {
193        Self::new(IndexManagerConfig::default())
194    }
195}
196
197#[cfg(test)]
198mod tests {
199    use super::*;
200    use mentedb_core::memory::MemoryType;
201    use mentedb_core::types::AgentId;
202
203    fn make_node(
204        embedding: Vec<f32>,
205        tags: Vec<String>,
206        salience: f32,
207        created_at: u64,
208    ) -> MemoryNode {
209        let mut node = MemoryNode::new(
210            AgentId::new(),
211            MemoryType::Episodic,
212            "test".into(),
213            embedding,
214        );
215        node.tags = tags;
216        node.salience = salience;
217        node.created_at = created_at;
218        node
219    }
220
221    #[test]
222    fn test_index_and_search() {
223        let mgr = IndexManager::default();
224        let node = make_node(vec![1.0, 0.0, 0.0, 0.0], vec!["test".into()], 0.8, 1000);
225        mgr.index_memory(&node);
226
227        let results = mgr.hybrid_search(&[1.0, 0.0, 0.0, 0.0], None, None, 1);
228        assert_eq!(results.len(), 1);
229        assert_eq!(results[0].0, node.id);
230    }
231
232    #[test]
233    fn test_tag_filter() {
234        let mgr = IndexManager::default();
235        let a = make_node(vec![1.0, 0.0, 0.0, 0.0], vec!["alpha".into()], 0.8, 1000);
236        let b = make_node(vec![0.9, 0.1, 0.0, 0.0], vec!["beta".into()], 0.8, 1000);
237        mgr.index_memory(&a);
238        mgr.index_memory(&b);
239
240        let results = mgr.hybrid_search(&[1.0, 0.0, 0.0, 0.0], Some(&["alpha"]), None, 10);
241        assert_eq!(results.len(), 1);
242        assert_eq!(results[0].0, a.id);
243    }
244
245    #[test]
246    fn test_time_filter() {
247        let mgr = IndexManager::default();
248        let a = make_node(vec![1.0, 0.0, 0.0, 0.0], vec![], 0.8, 100);
249        let b = make_node(vec![0.9, 0.1, 0.0, 0.0], vec![], 0.8, 500);
250        mgr.index_memory(&a);
251        mgr.index_memory(&b);
252
253        let results = mgr.hybrid_search(&[1.0, 0.0, 0.0, 0.0], None, Some((400, 600)), 10);
254        assert_eq!(results.len(), 1);
255        assert_eq!(results[0].0, b.id);
256    }
257
258    #[test]
259    fn test_remove_memory() {
260        let mgr = IndexManager::default();
261        let node = make_node(vec![1.0, 0.0, 0.0, 0.0], vec!["t".into()], 0.5, 100);
262        let id = node.id;
263        mgr.index_memory(&node);
264        mgr.remove_memory(id, &node);
265
266        let results = mgr.hybrid_search(&[1.0, 0.0, 0.0, 0.0], None, None, 10);
267        assert!(results.is_empty());
268    }
269
270    #[test]
271    fn test_empty_search() {
272        let mgr = IndexManager::default();
273        let results = mgr.hybrid_search(&[1.0, 0.0], None, None, 5);
274        assert!(results.is_empty());
275    }
276}
mentedb_index/manager.rs

mentedb_index/
manager.rs
