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

mentedb_index/
manager.rs
