engram_core/
persistence.rs

1use std::fs;
2use std::io::{BufReader, BufWriter};
3use std::path::Path;
4
5use engram_storage::Database;
6
7use crate::error::CoreError;
8use crate::indexes::IndexSet;
9
10const INDEX_FILENAME: &str = "indexes.hnsw";
11
12const FNV_OFFSET_BASIS: u64 = 0xcbf2_9ce4_8422_2325;
13const FNV_PRIME: u64 = 0x0100_0000_01b3;
14const DETERMINISTIC_RNG_MULTIPLIER: u64 = 0x9e37_79b9_7f4a_7c15;
15
16pub fn load_or_rebuild(
17    index_directory: &str,
18    database: &Database,
19    build_params: impl Fn() -> Result<engram_hnsw::HnswParams, CoreError>,
20) -> Result<IndexSet, CoreError> {
21    let index_path = Path::new(index_directory).join(INDEX_FILENAME);
22    if index_path.exists()
23        && let Ok(mut indexes) = load_from_disk(&index_path)
24    {
25        rebuild_id_map_from_database(&mut indexes, database)?;
26        return Ok(indexes);
27    }
28    rebuild_from_database(database, build_params)
29}
30
31fn load_from_disk(path: &Path) -> Result<IndexSet, CoreError> {
32    let file =
33        fs::File::open(path).map_err(|error| CoreError::IndexCorrupted(error.to_string()))?;
34    let mut reader = BufReader::new(file);
35    IndexSet::deserialize(&mut reader).map_err(CoreError::Hnsw)
36}
37
38fn rebuild_from_database(
39    database: &Database,
40    build_params: impl Fn() -> Result<engram_hnsw::HnswParams, CoreError>,
41) -> Result<IndexSet, CoreError> {
42    let mut indexes = IndexSet::new(build_params)?;
43    let memories = database.get_unindexed_memories(usize::MAX)?;
44    for memory in &memories {
45        let id = parse_memory_id(&memory.id)?;
46        let embedding = match extract_embedding(memory) {
47            Some(emb) => emb,
48            None => continue,
49        };
50        indexes.insert_atomic(id, &memory.id, &embedding, deterministic_rng(id))?;
51    }
52    Ok(indexes)
53}
54
55fn rebuild_id_map_from_database(
56    indexes: &mut IndexSet,
57    database: &Database,
58) -> Result<(), CoreError> {
59    let indexed_ids = database.get_indexed_memory_ids()?;
60    let entries: Vec<(u64, String)> = indexed_ids
61        .into_iter()
62        .filter_map(|memory_id| {
63            let hash = hash_string_to_u64(&memory_id);
64            if indexes.contains(hash) {
65                Some((hash, memory_id))
66            } else {
67                None
68            }
69        })
70        .collect();
71    indexes.rebuild_id_map(entries.into_iter());
72    Ok(())
73}
74
75pub fn save_to_disk(index_directory: &str, indexes: &IndexSet) -> Result<(), CoreError> {
76    let index_path = Path::new(index_directory).join(INDEX_FILENAME);
77    fs::create_dir_all(index_directory)
78        .map_err(|error| CoreError::RebuildFailed(error.to_string()))?;
79    let file = fs::File::create(&index_path)
80        .map_err(|error| CoreError::RebuildFailed(error.to_string()))?;
81    let mut writer = BufWriter::new(file);
82    indexes
83        .serialize(&mut writer)
84        .map_err(|error| CoreError::RebuildFailed(error.to_string()))?;
85    Ok(())
86}
87
88fn parse_memory_id(id: &str) -> Result<u64, CoreError> {
89    let hash = hash_string_to_u64(id);
90    Ok(hash)
91}
92
93pub fn hash_string_to_u64(value: &str) -> u64 {
94    let mut hash: u64 = FNV_OFFSET_BASIS;
95    for byte in value.as_bytes() {
96        hash ^= *byte as u64;
97        hash = hash.wrapping_mul(FNV_PRIME);
98    }
99    hash
100}
101
102pub fn extract_embeddings_from_memory(
103    memory: &engram_storage::Memory,
104) -> Result<engram_embeddings::ThreeFieldEmbedding, CoreError> {
105    extract_embedding(memory).ok_or_else(|| CoreError::IndexCorrupted("missing embeddings".into()))
106}
107
108fn extract_embedding(
109    memory: &engram_storage::Memory,
110) -> Option<engram_embeddings::ThreeFieldEmbedding> {
111    let context = bytes_to_f32_vec(memory.embedding_context.as_deref()?)?;
112    let action = bytes_to_f32_vec(memory.embedding_action.as_deref()?)?;
113    let result = bytes_to_f32_vec(memory.embedding_result.as_deref()?)?;
114    Some(engram_embeddings::ThreeFieldEmbedding {
115        context,
116        action,
117        result,
118    })
119}
120
121fn bytes_to_f32_vec(bytes: &[u8]) -> Option<Vec<f32>> {
122    if !bytes.len().is_multiple_of(4) {
123        return None;
124    }
125    let floats = bytes
126        .chunks_exact(4)
127        .map(|chunk| f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]))
128        .collect();
129    Some(floats)
130}
131
132pub fn deterministic_rng(id: u64) -> f64 {
133    let mixed = id.wrapping_mul(DETERMINISTIC_RNG_MULTIPLIER);
134    (mixed >> 11) as f64 / (1u64 << 53) as f64
135}
136
137pub fn f32_vec_to_bytes(values: &[f32]) -> Vec<u8> {
138    values.iter().flat_map(|v| v.to_le_bytes()).collect()
139}
engram_core/persistence.rs

engram_core/
persistence.rs
