engram_core/
indexes.rs

1use std::collections::HashMap;
2use std::io::{Read, Write};
3
4use engram_embeddings::ThreeFieldEmbedding;
5use engram_hnsw::{HnswError, HnswGraph, HnswParams};
6
7use crate::error::CoreError;
8
9pub struct IndexSet {
10    context_index: HnswGraph,
11    action_index: HnswGraph,
12    result_index: HnswGraph,
13    id_map: HashMap<u64, String>,
14}
15
16impl IndexSet {
17    pub fn new(
18        build_params: impl Fn() -> Result<HnswParams, CoreError>,
19    ) -> Result<Self, CoreError> {
20        let context_index = HnswGraph::new(build_params()?);
21        let action_index = HnswGraph::new(build_params()?);
22        let result_index = HnswGraph::new(build_params()?);
23        Ok(Self {
24            context_index,
25            action_index,
26            result_index,
27            id_map: HashMap::new(),
28        })
29    }
30
31    pub fn insert(
32        &mut self,
33        id: u64,
34        memory_id: &str,
35        embedding: &ThreeFieldEmbedding,
36        rng_value: f64,
37    ) -> Result<(), HnswError> {
38        self.context_index
39            .insert(id, embedding.context.clone(), rng_value)?;
40        self.action_index
41            .insert(id, embedding.action.clone(), rng_value)?;
42        self.result_index
43            .insert(id, embedding.result.clone(), rng_value)?;
44        self.id_map.insert(id, memory_id.to_string());
45        Ok(())
46    }
47
48    pub fn resolve_node_id(&self, node_id: u64) -> Option<&str> {
49        self.id_map.get(&node_id).map(String::as_str)
50    }
51
52    pub fn rebuild_id_map(&mut self, entries: impl Iterator<Item = (u64, String)>) {
53        self.id_map.clear();
54        for (hash, memory_id) in entries {
55            self.id_map.insert(hash, memory_id);
56        }
57    }
58
59    pub fn search(&self, query: &[f32], top_k: usize) -> Result<Vec<(u64, f32)>, HnswError> {
60        let mut scores: HashMap<u64, f32> = HashMap::new();
61        merge_scores(&mut scores, &self.context_index.search(query, top_k)?);
62        merge_scores(&mut scores, &self.action_index.search(query, top_k)?);
63        merge_scores(&mut scores, &self.result_index.search(query, top_k)?);
64        let mut ranked: Vec<(u64, f32)> = scores.into_iter().collect();
65        ranked.sort_by(|a, b| b.1.total_cmp(&a.1));
66        ranked.truncate(top_k);
67        Ok(ranked)
68    }
69
70    pub fn delete(&mut self, id: u64) -> Result<(), HnswError> {
71        self.context_index.delete(id)?;
72        self.action_index.delete(id)?;
73        self.result_index.delete(id)?;
74        self.id_map.remove(&id);
75        Ok(())
76    }
77
78    pub fn contains(&self, id: u64) -> bool {
79        self.context_index.contains(id)
80    }
81
82    pub fn len(&self) -> usize {
83        self.context_index.len()
84    }
85
86    pub fn is_empty(&self) -> bool {
87        self.context_index.is_empty()
88    }
89
90    pub fn serialize<W: Write>(&self, writer: &mut W) -> std::io::Result<()> {
91        self.context_index.serialize(writer)?;
92        self.action_index.serialize(writer)?;
93        self.result_index.serialize(writer)?;
94        Ok(())
95    }
96
97    pub fn deserialize<R: Read>(reader: &mut R) -> Result<Self, HnswError> {
98        let context_index = HnswGraph::deserialize(reader)?;
99        let action_index = HnswGraph::deserialize(reader)?;
100        let result_index = HnswGraph::deserialize(reader)?;
101        Ok(Self {
102            context_index,
103            action_index,
104            result_index,
105            id_map: HashMap::new(),
106        })
107    }
108}
109
110fn merge_scores(scores: &mut HashMap<u64, f32>, results: &[(u64, f32)]) {
111    for &(id, similarity) in results {
112        scores
113            .entry(id)
114            .and_modify(|existing| *existing = existing.max(similarity))
115            .or_insert(similarity);
116    }
117}
118
119pub mod instrumentation {
120    use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
121
122    static READER_TRACKING_ENABLED: AtomicBool = AtomicBool::new(false);
123    static CURRENT_READERS: AtomicUsize = AtomicUsize::new(0);
124    static MAX_READERS: AtomicUsize = AtomicUsize::new(0);
125
126    pub fn enable_reader_tracking() {
127        READER_TRACKING_ENABLED.store(true, Ordering::Relaxed);
128    }
129
130    pub fn disable_reader_tracking() {
131        READER_TRACKING_ENABLED.store(false, Ordering::Relaxed);
132    }
133
134    pub fn reset_reader_counters() {
135        CURRENT_READERS.store(0, Ordering::Relaxed);
136        MAX_READERS.store(0, Ordering::Relaxed);
137    }
138
139    pub fn concurrent_readers_max() -> usize {
140        MAX_READERS.load(Ordering::Relaxed)
141    }
142
143    pub struct ReaderTracker;
144
145    impl ReaderTracker {
146        pub fn new() -> Self {
147            if READER_TRACKING_ENABLED.load(Ordering::Relaxed) {
148                let current = CURRENT_READERS.fetch_add(1, Ordering::AcqRel) + 1;
149                MAX_READERS.fetch_max(current, Ordering::AcqRel);
150            }
151            Self
152        }
153    }
154
155    impl Drop for ReaderTracker {
156        fn drop(&mut self) {
157            if READER_TRACKING_ENABLED.load(Ordering::Relaxed) {
158                CURRENT_READERS.fetch_sub(1, Ordering::AcqRel);
159            }
160        }
161    }
162
163    impl Default for ReaderTracker {
164        fn default() -> Self {
165            Self::new()
166        }
167    }
168}
engram_core/indexes.rs

engram_core/
indexes.rs
