edgehdf5_memory/
lib.rs

1//! ZeroClaw agent memory HDF5 backend.
2//!
3//! Provides persistent memory storage for AI agents using HDF5 files.
4//! All data is cached in-memory for fast access and flushed to disk
5//! on mutations.
6
7pub mod bm25;
8#[cfg(any(feature = "accelerate", feature = "openblas"))]
9pub mod accelerate_search;
10#[cfg(feature = "fast-math")]
11pub mod blas_search;
12pub mod gpu_search;
13pub mod hybrid;
14pub mod ivf;
15pub mod pq;
16pub mod strategy;
17pub mod vector_search;
18
19pub mod agents_md;
20pub mod cache;
21pub mod decision_gate;
22pub mod knowledge;
23pub mod memory_strategy;
24pub mod schema;
25pub mod search;
26pub mod session;
27pub mod storage;
28pub mod wal;
29
30/// Cosine similarity with pre-computed norms using rustyhdf5_accel primitives.
31///
32/// Avoids recomputing the query/vector norms on every comparison.
33#[inline]
34pub fn cosine_similarity_prenorm(
35    query: &[f32],
36    query_norm: f32,
37    vec: &[f32],
38    vec_norm: f32,
39) -> f32 {
40    let denom = query_norm * vec_norm;
41    if denom == 0.0 {
42        return 0.0;
43    }
44    rustyhdf5_accel::dot_product(query, vec) / denom
45}
46
47use std::path::{Path, PathBuf};
48
49use cache::MemoryCache;
50use knowledge::KnowledgeCache;
51use memory_strategy::{Exchange, MemoryStrategy, StrategyOutput};
52use session::SessionCache;
53
54// --- Error type ---
55
56#[derive(Debug)]
57pub enum MemoryError {
58    Io(std::io::Error),
59    Hdf5(String),
60    Schema(String),
61    NotFound(String),
62}
63
64impl std::fmt::Display for MemoryError {
65    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
66        match self {
67            MemoryError::Io(e) => write!(f, "I/O error: {e}"),
68            MemoryError::Hdf5(e) => write!(f, "HDF5 error: {e}"),
69            MemoryError::Schema(e) => write!(f, "schema error: {e}"),
70            MemoryError::NotFound(e) => write!(f, "not found: {e}"),
71        }
72    }
73}
74
75impl std::error::Error for MemoryError {
76    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
77        match self {
78            MemoryError::Io(e) => Some(e),
79            _ => None,
80        }
81    }
82}
83
84impl From<std::io::Error> for MemoryError {
85    fn from(e: std::io::Error) -> Self {
86        MemoryError::Io(e)
87    }
88}
89
90pub type Result<T> = std::result::Result<T, MemoryError>;
91
92// --- Config and data types ---
93
94#[derive(Debug, Clone)]
95pub struct MemoryConfig {
96    pub path: PathBuf,
97    pub agent_id: String,
98    pub embedder: String,
99    pub embedding_dim: usize,
100    pub chunk_size: usize,
101    pub overlap: usize,
102    pub float16: bool,
103    pub compression: bool,
104    pub compression_level: u32,
105    pub compact_threshold: f32,
106    pub hebbian_boost: f32,
107    pub decay_factor: f32,
108    pub created_at: String,
109    pub wal_enabled: bool,
110    pub wal_max_entries: usize,
111}
112
113impl MemoryConfig {
114    pub fn new(path: PathBuf, agent_id: &str, embedding_dim: usize) -> Self {
115        let created_at = now_iso8601();
116        Self {
117            path,
118            agent_id: agent_id.to_string(),
119            embedder: "openai:text-embedding-3-small".to_string(),
120            embedding_dim,
121            chunk_size: 512,
122            overlap: 50,
123            float16: false,
124            compression: false,
125            compression_level: 0,
126            compact_threshold: 0.3,
127            hebbian_boost: 0.15,
128            decay_factor: 0.98,
129            created_at,
130            wal_enabled: true,
131            wal_max_entries: 500,
132        }
133    }
134}
135
136#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
137pub struct MemoryEntry {
138    pub chunk: String,
139    pub embedding: Vec<f32>,
140    pub source_channel: String,
141    pub timestamp: f64,
142    pub session_id: String,
143    pub tags: String,
144}
145
146#[derive(Debug, Clone)]
147pub struct SearchResult {
148    pub score: f32,
149    pub chunk: String,
150    pub index: usize,
151    pub timestamp: f64,
152    pub source_channel: String,
153    pub activation: f32,
154}
155
156// --- Trait ---
157
158pub trait AgentMemory {
159    fn save(&mut self, entry: MemoryEntry) -> Result<usize>;
160    fn save_batch(&mut self, entries: Vec<MemoryEntry>) -> Result<Vec<usize>>;
161    fn delete(&mut self, id: usize) -> Result<()>;
162    fn compact(&mut self) -> Result<usize>;
163    fn count(&self) -> usize;
164    fn count_active(&self) -> usize;
165    fn snapshot(&self, dest: &Path) -> Result<PathBuf>;
166    fn add_session(
167        &mut self,
168        id: &str,
169        start: usize,
170        end: usize,
171        channel: &str,
172        summary: &str,
173    ) -> Result<()>;
174    fn get_session_summary(&self, session_id: &str) -> Result<Option<String>>;
175}
176
177// --- HDF5Memory ---
178
179pub struct HDF5Memory {
180    pub(crate) config: MemoryConfig,
181    pub(crate) cache: MemoryCache,
182    pub(crate) sessions: SessionCache,
183    pub(crate) knowledge: KnowledgeCache,
184    wal: Option<wal::WalFile>,
185    strategy: Option<Box<dyn MemoryStrategy>>,
186}
187
188impl std::fmt::Debug for HDF5Memory {
189    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "HDF5Memory({:?})", self.config.agent_id) }
190}
191
192impl HDF5Memory {
193    /// Create a new HDF5 memory file with the given configuration.
194    pub fn create(config: MemoryConfig) -> Result<Self> {
195        let cache = MemoryCache::new(config.embedding_dim);
196        let sessions = SessionCache::new();
197        let knowledge = KnowledgeCache::new();
198
199        // Write initial empty file
200        storage::write_to_disk(&config.path, &config, &cache, &sessions, &knowledge)?;
201
202        let wal = if config.wal_enabled {
203            let wal_path = config.path.with_extension("h5.wal");
204            Some(wal::WalFile::open(&wal_path)?)
205        } else {
206            None
207        };
208
209        Ok(Self {
210            config,
211            cache,
212            sessions,
213            knowledge,
214            wal,
215            strategy: None,
216        })
217    }
218
219    /// Open an existing HDF5 memory file.
220    pub fn open(path: &Path) -> Result<Self> {
221        let (config, mut cache, sessions, knowledge) = storage::read_from_disk(path)?;
222
223        // Replay WAL if present
224        let wal_path = path.with_extension("h5.wal");
225        let wal = if wal_path.exists() {
226            let entries = wal::WalFile::read_entries(&wal_path)?;
227            wal::replay_into_cache(&entries, &mut cache);
228            Some(wal::WalFile::open(&wal_path)?)
229        } else if config.wal_enabled {
230            Some(wal::WalFile::open(&wal_path)?)
231        } else {
232            None
233        };
234
235        Ok(Self {
236            config,
237            cache,
238            sessions,
239            knowledge,
240            wal,
241            strategy: None,
242        })
243    }
244
245    /// Flush current state to disk.
246    fn flush(&self) -> Result<()> {
247        storage::write_to_disk(
248            &self.config.path,
249            &self.config,
250            &self.cache,
251            &self.sessions,
252            &self.knowledge,
253        )
254    }
255
256    /// Get a reference to the config.
257    pub fn config(&self) -> &MemoryConfig {
258        &self.config
259    }
260
261    /// Get a reference to the knowledge cache.
262    pub fn knowledge(&self) -> &KnowledgeCache {
263        &self.knowledge
264    }
265
266    /// Get a mutable reference to the knowledge cache.
267    pub fn knowledge_mut(&mut self) -> &mut KnowledgeCache {
268        &mut self.knowledge
269    }
270
271    /// Add an entity to the knowledge graph and flush.
272    pub fn add_entity(
273        &mut self,
274        name: &str,
275        entity_type: &str,
276        embedding_idx: i64,
277    ) -> Result<u64> {
278        let id = self.knowledge.add_entity(name, entity_type, embedding_idx);
279        self.flush()?;
280        Ok(id)
281    }
282
283    /// Add an alias for a knowledge graph entity and flush.
284    pub fn add_entity_alias(&mut self, alias: &str, entity_id: i64) -> Result<()> {
285        self.knowledge.add_alias(alias, entity_id);
286        self.flush()
287    }
288
289    /// Add a relation to the knowledge graph and flush.
290    pub fn add_relation(
291        &mut self,
292        src: u64,
293        tgt: u64,
294        relation: &str,
295        weight: f32,
296    ) -> Result<()> {
297        self.knowledge.add_relation(src, tgt, relation, weight);
298        self.flush()?;
299        Ok(())
300    }
301}
302
303impl AgentMemory for HDF5Memory {
304    fn save(&mut self, entry: MemoryEntry) -> Result<usize> {
305        if let Some(ref mut w) = self.wal {
306            let wal_entry = wal::WalEntry {
307                entry_type: wal::WalEntryType::Save,
308                timestamp: entry.timestamp,
309                chunk: entry.chunk.clone(),
310                embedding: entry.embedding.clone(),
311                source_channel: entry.source_channel.clone(),
312                session_id: entry.session_id.clone(),
313                tags: entry.tags.clone(),
314                tombstone_index: None,
315            };
316            w.append_save(&wal_entry)?;
317        }
318        let idx = self.cache.push(
319            entry.chunk,
320            entry.embedding,
321            entry.source_channel,
322            entry.timestamp,
323            entry.session_id,
324            entry.tags,
325        );
326        if self.wal.is_some() {
327            // Check auto-merge threshold
328            if self.wal.as_ref().unwrap().pending_count() as usize > self.config.wal_max_entries {
329                self.flush()?;
330                self.wal.as_mut().unwrap().truncate()?;
331            }
332        } else {
333            self.flush()?;
334        }
335        Ok(idx)
336    }
337
338    fn save_batch(&mut self, entries: Vec<MemoryEntry>) -> Result<Vec<usize>> {
339        let mut indices = Vec::with_capacity(entries.len());
340        for entry in entries {
341            let idx = self.cache.push(
342                entry.chunk,
343                entry.embedding,
344                entry.source_channel,
345                entry.timestamp,
346                entry.session_id,
347                entry.tags,
348            );
349            indices.push(idx);
350        }
351        self.flush()?;
352        Ok(indices)
353    }
354
355    fn delete(&mut self, id: usize) -> Result<()> {
356        if !self.cache.mark_deleted(id) {
357            return Err(MemoryError::NotFound(format!(
358                "entry {id} not found or already deleted"
359            )));
360        }
361        self.flush()?;
362
363        // Auto-compact if threshold exceeded
364        if self.config.compact_threshold > 0.0
365            && self.cache.tombstone_fraction() > self.config.compact_threshold
366        {
367            self.compact()?;
368        }
369
370        Ok(())
371    }
372
373    fn compact(&mut self) -> Result<usize> {
374        let (removed, _index_map) = self.cache.compact();
375        if removed > 0 {
376            self.flush()?;
377        }
378        Ok(removed)
379    }
380
381    fn count(&self) -> usize {
382        self.cache.len()
383    }
384
385    fn count_active(&self) -> usize {
386        self.cache.count_active()
387    }
388
389    fn snapshot(&self, dest: &Path) -> Result<PathBuf> {
390        storage::snapshot_file(&self.config.path, dest)
391    }
392
393    fn add_session(
394        &mut self,
395        id: &str,
396        start: usize,
397        end: usize,
398        channel: &str,
399        summary: &str,
400    ) -> Result<()> {
401        self.sessions.add(id, start, end, channel, summary);
402        self.flush()?;
403        Ok(())
404    }
405
406    fn get_session_summary(&self, session_id: &str) -> Result<Option<String>> {
407        Ok(self.sessions.find_summary(session_id).map(String::from))
408    }
409}
410
411fn now_iso8601() -> String {
412    let d = std::time::SystemTime::now()
413        .duration_since(std::time::UNIX_EPOCH)
414        .unwrap_or_default();
415    let secs = d.as_secs();
416    let time_secs = secs % 86400;
417    let hours = time_secs / 3600;
418    let minutes = (time_secs % 3600) / 60;
419    let seconds = time_secs % 60;
420
421    let mut y = 1970i64;
422    let mut remaining_days = (secs / 86400) as i64;
423    loop {
424        let days_in_year = if is_leap(y) { 366 } else { 365 };
425        if remaining_days < days_in_year {
426            break;
427        }
428        remaining_days -= days_in_year;
429        y += 1;
430    }
431    let month_days = if is_leap(y) {
432        [31, 29, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31]
433    } else {
434        [31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31]
435    };
436    let mut m = 1u32;
437    for &md in &month_days {
438        if remaining_days < md {
439            break;
440        }
441        remaining_days -= md;
442        m += 1;
443    }
444    let day = remaining_days + 1;
445
446    format!("{y:04}-{m:02}-{day:02}T{hours:02}:{minutes:02}:{seconds:02}Z")
447}
448
449fn is_leap(y: i64) -> bool {
450    (y % 4 == 0 && y % 100 != 0) || y % 400 == 0
451}
452
453// --- Tests ---
454
455#[cfg(test)]
456mod tests {
457    use super::*;
458    use tempfile::TempDir;
459
460    fn make_config(dir: &TempDir) -> MemoryConfig {
461        let mut c = MemoryConfig::new(dir.path().join("test.h5"), "agent-test", 4);
462        c.wal_enabled = false;
463        c
464    }
465
466    fn make_entry(chunk: &str, embedding: &[f32]) -> MemoryEntry {
467        MemoryEntry {
468            chunk: chunk.to_string(),
469            embedding: embedding.to_vec(),
470            source_channel: "test".to_string(),
471            timestamp: 1000000.0,
472            session_id: "session-1".to_string(),
473            tags: "tag1,tag2".to_string(),
474        }
475    }
476
477    #[test]
478    fn create_new_file() {
479        let dir = TempDir::new().unwrap();
480        let config = make_config(&dir);
481        let mem = HDF5Memory::create(config).unwrap();
482        assert_eq!(mem.count(), 0);
483        assert_eq!(mem.count_active(), 0);
484        assert!(dir.path().join("test.h5").exists());
485    }
486
487    #[test]
488    fn save_single_entry() {
489        let dir = TempDir::new().unwrap();
490        let config = make_config(&dir);
491        let mut mem = HDF5Memory::create(config).unwrap();
492
493        let idx = mem
494            .save(make_entry("hello world", &[1.0, 2.0, 3.0, 4.0]))
495            .unwrap();
496        assert_eq!(idx, 0);
497        assert_eq!(mem.count(), 1);
498        assert_eq!(mem.count_active(), 1);
499    }
500
501    #[test]
502    fn save_batch() {
503        let dir = TempDir::new().unwrap();
504        let config = make_config(&dir);
505        let mut mem = HDF5Memory::create(config).unwrap();
506
507        let entries = vec![
508            make_entry("chunk 1", &[1.0, 0.0, 0.0, 0.0]),
509            make_entry("chunk 2", &[0.0, 1.0, 0.0, 0.0]),
510            make_entry("chunk 3", &[0.0, 0.0, 1.0, 0.0]),
511        ];
512        let indices = mem.save_batch(entries).unwrap();
513        assert_eq!(indices, vec![0, 1, 2]);
514        assert_eq!(mem.count(), 3);
515    }
516
517    #[test]
518    fn delete_entry() {
519        let dir = TempDir::new().unwrap();
520        let mut config = make_config(&dir);
521        config.compact_threshold = 0.0;
522        let mut mem = HDF5Memory::create(config).unwrap();
523
524        mem.save(make_entry("chunk 1", &[1.0, 0.0, 0.0, 0.0]))
525            .unwrap();
526        mem.save(make_entry("chunk 2", &[0.0, 1.0, 0.0, 0.0]))
527            .unwrap();
528
529        mem.delete(0).unwrap();
530        assert_eq!(mem.count(), 2);
531        assert_eq!(mem.count_active(), 1);
532    }
533
534    #[test]
535    fn compact_removes_tombstoned() {
536        let dir = TempDir::new().unwrap();
537        let mut config = make_config(&dir);
538        config.compact_threshold = 0.0;
539        let mut mem = HDF5Memory::create(config).unwrap();
540
541        mem.save(make_entry("chunk 1", &[1.0, 0.0, 0.0, 0.0]))
542            .unwrap();
543        mem.save(make_entry("chunk 2", &[0.0, 1.0, 0.0, 0.0]))
544            .unwrap();
545        mem.save(make_entry("chunk 3", &[0.0, 0.0, 1.0, 0.0]))
546            .unwrap();
547
548        mem.delete(0).unwrap();
549        mem.delete(2).unwrap();
550
551        let removed = mem.compact().unwrap();
552        assert_eq!(removed, 2);
553        assert_eq!(mem.count(), 1);
554        assert_eq!(mem.count_active(), 1);
555    }
556
557    #[test]
558    fn snapshot_creates_copy() {
559        let dir = TempDir::new().unwrap();
560        let config = make_config(&dir);
561        let mut mem = HDF5Memory::create(config).unwrap();
562
563        mem.save(make_entry("snapshot test", &[1.0, 2.0, 3.0, 4.0]))
564            .unwrap();
565
566        let snap_dir = TempDir::new().unwrap();
567        let snap_path = mem.snapshot(snap_dir.path()).unwrap();
568        assert!(snap_path.exists());
569
570        let snap_mem = HDF5Memory::open(&snap_path).unwrap();
571        assert_eq!(snap_mem.count(), 1);
572    }
573
574    #[test]
575    fn session_tracking() {
576        let dir = TempDir::new().unwrap();
577        let config = make_config(&dir);
578        let mut mem = HDF5Memory::create(config).unwrap();
579
580        mem.add_session("sess-1", 0, 5, "whatsapp", "discussed AI topics")
581            .unwrap();
582        mem.add_session("sess-2", 6, 10, "slack", "code review session")
583            .unwrap();
584
585        let summary = mem.get_session_summary("sess-1").unwrap();
586        assert_eq!(summary.as_deref(), Some("discussed AI topics"));
587
588        let summary2 = mem.get_session_summary("sess-2").unwrap();
589        assert_eq!(summary2.as_deref(), Some("code review session"));
590
591        let missing = mem.get_session_summary("sess-999").unwrap();
592        assert!(missing.is_none());
593    }
594
595    #[test]
596    fn knowledge_add_entity() {
597        let dir = TempDir::new().unwrap();
598        let config = make_config(&dir);
599        let mut mem = HDF5Memory::create(config).unwrap();
600
601        let id1 = mem.add_entity("Rust", "language", -1).unwrap();
602        let id2 = mem.add_entity("HDF5", "format", -1).unwrap();
603
604        assert_eq!(id1, 0);
605        assert_eq!(id2, 1);
606
607        let entity = mem.knowledge().get_entity(0).unwrap();
608        assert_eq!(entity.name, "Rust");
609        assert_eq!(entity.entity_type, "language");
610    }
611
612    #[test]
613    fn knowledge_add_relation() {
614        let dir = TempDir::new().unwrap();
615        let config = make_config(&dir);
616        let mut mem = HDF5Memory::create(config).unwrap();
617
618        let rust_id = mem.add_entity("Rust", "language", -1).unwrap();
619        let hdf5_id = mem.add_entity("HDF5", "format", -1).unwrap();
620        mem.add_relation(rust_id, hdf5_id, "uses", 1.0).unwrap();
621
622        let rels = mem.knowledge().get_relations_from(rust_id);
623        assert_eq!(rels.len(), 1);
624        assert_eq!(rels[0].relation, "uses");
625        assert_eq!(rels[0].tgt, hdf5_id);
626    }
627
628    #[test]
629    fn open_existing() {
630        let dir = TempDir::new().unwrap();
631        let config = make_config(&dir);
632        let path = config.path.clone();
633
634        {
635            let mut mem = HDF5Memory::create(config).unwrap();
636            mem.save(make_entry("persisted chunk", &[1.0, 2.0, 3.0, 4.0]))
637                .unwrap();
638        }
639
640        let mem = HDF5Memory::open(&path).unwrap();
641        assert_eq!(mem.count(), 1);
642        assert_eq!(mem.config().agent_id, "agent-test");
643        assert_eq!(mem.config().embedding_dim, 4);
644    }
645
646    #[test]
647    fn schema_version_mismatch() {
648        let dir = TempDir::new().unwrap();
649        let path = dir.path().join("bad.h5");
650
651        let mut builder = rustyhdf5::FileBuilder::new();
652        let mut meta = builder.create_group("meta");
653        meta.set_attr(
654            "schema_version",
655            rustyhdf5::AttrValue::String("99.0".into()),
656        );
657        meta.set_attr("created_at", rustyhdf5::AttrValue::String("now".into()));
658        meta.set_attr("agent_id", rustyhdf5::AttrValue::String("test".into()));
659        meta.set_attr("embedder", rustyhdf5::AttrValue::String("test".into()));
660        meta.set_attr("embedding_dim", rustyhdf5::AttrValue::I64(4));
661        meta.set_attr("chunk_size", rustyhdf5::AttrValue::I64(512));
662        meta.set_attr("overlap", rustyhdf5::AttrValue::I64(50));
663        meta.create_dataset("_marker").with_u8_data(&[1]);
664        let finished = meta.finish();
665        builder.add_group(finished);
666        builder.write(&path).unwrap();
667
668        let err = HDF5Memory::open(&path).unwrap_err();
669        let msg = err.to_string();
670        assert!(msg.contains("schema version mismatch"), "got: {msg}");
671    }
672
673    #[test]
674    fn round_trip() {
675        let dir = TempDir::new().unwrap();
676        let config = make_config(&dir);
677        let path = config.path.clone();
678
679        {
680            let mut mem = HDF5Memory::create(config).unwrap();
681            mem.save(make_entry("round trip data", &[0.1, 0.2, 0.3, 0.4]))
682                .unwrap();
683            mem.add_session("sess-rt", 0, 0, "api", "round trip session")
684                .unwrap();
685            mem.add_entity("TestEntity", "test", 0).unwrap();
686        }
687
688        let mem = HDF5Memory::open(&path).unwrap();
689        assert_eq!(mem.count(), 1);
690        let summary = mem.get_session_summary("sess-rt").unwrap();
691        assert_eq!(summary.as_deref(), Some("round trip session"));
692        let entity = mem.knowledge().get_entity(0).unwrap();
693        assert_eq!(entity.name, "TestEntity");
694    }
695
696    #[test]
697    fn delete_nonexistent() {
698        let dir = TempDir::new().unwrap();
699        let config = make_config(&dir);
700        let mut mem = HDF5Memory::create(config).unwrap();
701
702        let err = mem.delete(999).unwrap_err();
703        assert!(err.to_string().contains("not found"));
704    }
705
706    #[test]
707    fn double_delete() {
708        let dir = TempDir::new().unwrap();
709        let mut config = make_config(&dir);
710        config.compact_threshold = 0.0;
711        let mut mem = HDF5Memory::create(config).unwrap();
712
713        mem.save(make_entry("double del", &[1.0, 0.0, 0.0, 0.0]))
714            .unwrap();
715        mem.delete(0).unwrap();
716        let err = mem.delete(0).unwrap_err();
717        assert!(err.to_string().contains("not found"));
718    }
719
720    #[test]
721    fn compact_no_tombstones() {
722        let dir = TempDir::new().unwrap();
723        let config = make_config(&dir);
724        let mut mem = HDF5Memory::create(config).unwrap();
725
726        mem.save(make_entry("no compact", &[1.0, 0.0, 0.0, 0.0]))
727            .unwrap();
728        let removed = mem.compact().unwrap();
729        assert_eq!(removed, 0);
730        assert_eq!(mem.count(), 1);
731    }
732
733    #[test]
734    fn empty_file_operations() {
735        let dir = TempDir::new().unwrap();
736        let config = make_config(&dir);
737        let path = config.path.clone();
738        let mem = HDF5Memory::create(config).unwrap();
739        assert_eq!(mem.count(), 0);
740        assert_eq!(mem.count_active(), 0);
741
742        let mem2 = HDF5Memory::open(&path).unwrap();
743        assert_eq!(mem2.count(), 0);
744    }
745
746    #[test]
747    fn multiple_sessions() {
748        let dir = TempDir::new().unwrap();
749        let config = make_config(&dir);
750        let path = config.path.clone();
751
752        {
753            let mut mem = HDF5Memory::create(config).unwrap();
754            for i in 0..5 {
755                mem.add_session(
756                    &format!("sess-{i}"),
757                    i * 10,
758                    (i + 1) * 10,
759                    "api",
760                    &format!("session {i} summary"),
761                )
762                .unwrap();
763            }
764        }
765
766        let mem = HDF5Memory::open(&path).unwrap();
767        for i in 0..5 {
768            let summary = mem
769                .get_session_summary(&format!("sess-{i}"))
770                .unwrap()
771                .unwrap();
772            assert_eq!(summary, format!("session {i} summary"));
773        }
774    }
775
776    #[test]
777    fn knowledge_graph_persistence() {
778        let dir = TempDir::new().unwrap();
779        let config = make_config(&dir);
780        let path = config.path.clone();
781
782        {
783            let mut mem = HDF5Memory::create(config).unwrap();
784            let id1 = mem.add_entity("Alice", "person", -1).unwrap();
785            let id2 = mem.add_entity("Bob", "person", -1).unwrap();
786            mem.add_relation(id1, id2, "knows", 0.9).unwrap();
787        }
788
789        let mem = HDF5Memory::open(&path).unwrap();
790        assert_eq!(mem.knowledge().entities.len(), 2);
791        assert_eq!(mem.knowledge().relations.len(), 1);
792        assert_eq!(mem.knowledge().get_entity(0).unwrap().name, "Alice");
793        assert_eq!(mem.knowledge().get_entity(1).unwrap().name, "Bob");
794
795        let rels = mem.knowledge().get_relations_from(0);
796        assert_eq!(rels.len(), 1);
797        assert_eq!(rels[0].relation, "knows");
798    }
799
800    #[test]
801    fn different_channels() {
802        let dir = TempDir::new().unwrap();
803        let config = make_config(&dir);
804        let path = config.path.clone();
805
806        {
807            let mut mem = HDF5Memory::create(config).unwrap();
808            let e1 = MemoryEntry {
809                chunk: "whatsapp msg".into(),
810                embedding: vec![1.0, 0.0, 0.0, 0.0],
811                source_channel: "whatsapp".into(),
812                timestamp: 100.0,
813                session_id: "s1".into(),
814                tags: "chat".into(),
815            };
816            let e2 = MemoryEntry {
817                chunk: "slack msg".into(),
818                embedding: vec![0.0, 1.0, 0.0, 0.0],
819                source_channel: "slack".into(),
820                timestamp: 200.0,
821                session_id: "s2".into(),
822                tags: "work".into(),
823            };
824            mem.save_batch(vec![e1, e2]).unwrap();
825        }
826
827        let mem = HDF5Memory::open(&path).unwrap();
828        assert_eq!(mem.count(), 2);
829    }
830
831    #[test]
832    fn compact_then_reopen() {
833        let dir = TempDir::new().unwrap();
834        let mut config = make_config(&dir);
835        config.compact_threshold = 0.0;
836        let path = config.path.clone();
837
838        {
839            let mut mem = HDF5Memory::create(config).unwrap();
840            mem.save(make_entry("keep", &[1.0, 0.0, 0.0, 0.0]))
841                .unwrap();
842            mem.save(make_entry("delete me", &[0.0, 1.0, 0.0, 0.0]))
843                .unwrap();
844            mem.save(make_entry("also keep", &[0.0, 0.0, 1.0, 0.0]))
845                .unwrap();
846
847            mem.delete(1).unwrap();
848            mem.compact().unwrap();
849        }
850
851        let mem = HDF5Memory::open(&path).unwrap();
852        assert_eq!(mem.count(), 2);
853        assert_eq!(mem.count_active(), 2);
854    }
855
856    #[test]
857    fn config_preserved() {
858        let dir = TempDir::new().unwrap();
859        let mut config = make_config(&dir);
860        config.embedder = "custom-embedder".into();
861        config.chunk_size = 1024;
862        config.overlap = 100;
863        let path = config.path.clone();
864
865        HDF5Memory::create(config).unwrap();
866
867        let mem = HDF5Memory::open(&path).unwrap();
868        assert_eq!(mem.config().embedder, "custom-embedder");
869        assert_eq!(mem.config().chunk_size, 1024);
870        assert_eq!(mem.config().overlap, 100);
871    }
872
873    #[test]
874    fn large_batch() {
875        let dir = TempDir::new().unwrap();
876        let config = make_config(&dir);
877        let path = config.path.clone();
878
879        {
880            let mut mem = HDF5Memory::create(config).unwrap();
881            let entries: Vec<MemoryEntry> = (0..100)
882                .map(|i| MemoryEntry {
883                    chunk: format!("chunk number {i} with some content"),
884                    embedding: vec![i as f32, 0.0, 0.0, 0.0],
885                    source_channel: "api".into(),
886                    timestamp: i as f64 * 1000.0,
887                    session_id: format!("batch-sess-{}", i / 10),
888                    tags: format!("batch,item-{i}"),
889                })
890                .collect();
891            mem.save_batch(entries).unwrap();
892        }
893
894        let mem = HDF5Memory::open(&path).unwrap();
895        assert_eq!(mem.count(), 100);
896        assert_eq!(mem.count_active(), 100);
897    }
898
899    #[test]
900    fn auto_compact() {
901        let dir = TempDir::new().unwrap();
902        let mut config = make_config(&dir);
903        config.compact_threshold = 0.4;
904        let mut mem = HDF5Memory::create(config).unwrap();
905
906        mem.save(make_entry("a", &[1.0, 0.0, 0.0, 0.0])).unwrap();
907        mem.save(make_entry("b", &[0.0, 1.0, 0.0, 0.0])).unwrap();
908        mem.save(make_entry("c", &[0.0, 0.0, 1.0, 0.0])).unwrap();
909
910        mem.delete(0).unwrap();
911        assert_eq!(mem.count(), 3);
912
913        mem.delete(1).unwrap();
914        assert_eq!(mem.count(), 1);
915        assert_eq!(mem.count_active(), 1);
916    }
917
918    #[test]
919    fn snapshot_to_file() {
920        let dir = TempDir::new().unwrap();
921        let config = make_config(&dir);
922        let mut mem = HDF5Memory::create(config).unwrap();
923        mem.save(make_entry("snap", &[1.0, 2.0, 3.0, 4.0]))
924            .unwrap();
925
926        let snap_path = dir.path().join("my_snapshot.h5");
927        let result = mem.snapshot(&snap_path).unwrap();
928        assert_eq!(result, snap_path);
929        assert!(snap_path.exists());
930    }
931
932    #[test]
933    fn entity_id_continuity() {
934        let dir = TempDir::new().unwrap();
935        let config = make_config(&dir);
936        let path = config.path.clone();
937
938        {
939            let mut mem = HDF5Memory::create(config).unwrap();
940            mem.add_entity("First", "test", -1).unwrap();
941            mem.add_entity("Second", "test", -1).unwrap();
942        }
943
944        let mut mem = HDF5Memory::open(&path).unwrap();
945        let id3 = mem.add_entity("Third", "test", -1).unwrap();
946        assert_eq!(id3, 2);
947    }
948
949    #[test]
950    fn multiple_relations() {
951        let dir = TempDir::new().unwrap();
952        let config = make_config(&dir);
953        let mut mem = HDF5Memory::create(config).unwrap();
954
955        let a = mem.add_entity("A", "node", -1).unwrap();
956        let b = mem.add_entity("B", "node", -1).unwrap();
957        let c = mem.add_entity("C", "node", -1).unwrap();
958
959        mem.add_relation(a, b, "connects", 1.0).unwrap();
960        mem.add_relation(a, c, "connects", 0.5).unwrap();
961        mem.add_relation(b, c, "depends_on", 0.8).unwrap();
962
963        assert_eq!(mem.knowledge().get_relations_from(a).len(), 2);
964        assert_eq!(mem.knowledge().get_relations_from(b).len(), 1);
965        assert_eq!(mem.knowledge().get_relations_to(c).len(), 2);
966    }
967
968    #[test]
969    fn empty_strings() {
970        let dir = TempDir::new().unwrap();
971        let config = make_config(&dir);
972        let path = config.path.clone();
973
974        {
975            let mut mem = HDF5Memory::create(config).unwrap();
976            let entry = MemoryEntry {
977                chunk: "content".into(),
978                embedding: vec![1.0, 0.0, 0.0, 0.0],
979                source_channel: "".into(),
980                timestamp: 0.0,
981                session_id: "".into(),
982                tags: "".into(),
983            };
984            mem.save(entry).unwrap();
985        }
986
987        let mem = HDF5Memory::open(&path).unwrap();
988        assert_eq!(mem.count(), 1);
989    }
990
991    // ---------------------------------------------------------------
992    // Hebbian activation & decay tests
993    // ---------------------------------------------------------------
994
995    #[test]
996    fn test_hebbian_activation_boost() {
997        let dir = TempDir::new().unwrap();
998        let config = make_config(&dir);
999        let mut mem = HDF5Memory::create(config).unwrap();
1000
1001        // Save 10 entries; entry 0 has embedding [1,0,0,0]
1002        for i in 0..10 {
1003            let emb = if i == 0 {
1004                vec![1.0, 0.0, 0.0, 0.0]
1005            } else {
1006                // orthogonal-ish embeddings
1007                vec![0.0, (i as f32).sin(), (i as f32).cos(), 0.0]
1008            };
1009            mem.save(make_entry(&format!("chunk {i}"), &emb)).unwrap();
1010        }
1011
1012        // Search 5 times for a query that matches entry 0 best
1013        let query = vec![1.0, 0.0, 0.0, 0.0];
1014        for _ in 0..5 {
1015            let results = mem.hybrid_search(&query, "chunk", 1.0, 0.0, 3);
1016            assert!(!results.is_empty());
1017        }
1018
1019        // Entry 0 should have a higher activation weight than all others
1020        let w0 = mem.cache.activation_weights[0];
1021        for i in 1..10 {
1022            assert!(
1023                w0 > mem.cache.activation_weights[i],
1024                "entry 0 weight ({w0}) should be > entry {i} weight ({})",
1025                mem.cache.activation_weights[i]
1026            );
1027        }
1028    }
1029
1030    #[test]
1031    fn test_hebbian_decay() {
1032        let dir = TempDir::new().unwrap();
1033        let config = make_config(&dir);
1034        let mut mem = HDF5Memory::create(config).unwrap();
1035
1036        for i in 0..5 {
1037            mem.save(make_entry(&format!("decay {i}"), &[i as f32, 1.0, 0.0, 0.0]))
1038                .unwrap();
1039        }
1040
1041        // Call tick_session 100 times with no searches
1042        for _ in 0..100 {
1043            mem.tick_session().unwrap();
1044        }
1045
1046        // All weights should approach 0 (< 0.2)
1047        for (i, &w) in mem.cache.activation_weights.iter().enumerate() {
1048            assert!(
1049                w < 0.2,
1050                "weight[{i}] = {w}, expected < 0.2 after 100 decay ticks"
1051            );
1052        }
1053    }
1054
1055    #[test]
1056    fn test_hebbian_no_effect_at_default() {
1057        let dir = TempDir::new().unwrap();
1058        let config = make_config(&dir);
1059        let mut mem = HDF5Memory::create(config).unwrap();
1060
1061        mem.save(make_entry("alpha", &[1.0, 0.0, 0.0, 0.0]))
1062            .unwrap();
1063        mem.save(make_entry("beta", &[0.0, 1.0, 0.0, 0.0]))
1064            .unwrap();
1065        mem.save(make_entry("gamma", &[0.5, 0.5, 0.0, 0.0]))
1066            .unwrap();
1067
1068        // All weights should be 1.0 (default)
1069        for &w in &mem.cache.activation_weights {
1070            assert!((w - 1.0).abs() < 1e-6, "default weight should be 1.0, got {w}");
1071        }
1072
1073        // Search: since sqrt(1.0) == 1.0, scores should be pure cosine
1074        let query = vec![1.0, 0.0, 0.0, 0.0];
1075        let results = mem.hybrid_search(&query, "", 1.0, 0.0, 3);
1076        // Entry 0 should be best (perfect match)
1077        assert_eq!(results[0].index, 0);
1078        assert!((results[0].activation - 1.0).abs() < 1e-6);
1079    }
1080
1081    #[test]
1082    fn test_hebbian_persistence() {
1083        let dir = TempDir::new().unwrap();
1084        let config = make_config(&dir);
1085        let path = config.path.clone();
1086
1087        {
1088            let mut mem = HDF5Memory::create(config).unwrap();
1089            mem.save(make_entry("persist me", &[1.0, 0.0, 0.0, 0.0]))
1090                .unwrap();
1091
1092            // Boost via search
1093            let query = vec![1.0, 0.0, 0.0, 0.0];
1094            mem.hybrid_search(&query, "", 1.0, 0.0, 1);
1095            let boosted_weight = mem.cache.activation_weights[0];
1096            assert!(boosted_weight > 1.0, "weight should be boosted after search");
1097        }
1098
1099        // Reopen and check weight persisted
1100        let mem = HDF5Memory::open(&path).unwrap();
1101        assert!(
1102            mem.cache.activation_weights[0] > 1.0,
1103            "persisted weight should be > 1.0, got {}",
1104            mem.cache.activation_weights[0]
1105        );
1106    }
1107
1108    #[test]
1109    fn test_hebbian_compact_preserves_weights() {
1110        let dir = TempDir::new().unwrap();
1111        let mut config = make_config(&dir);
1112        config.compact_threshold = 0.0;
1113        let mut mem = HDF5Memory::create(config).unwrap();
1114
1115        // Save 5 entries
1116        for i in 0..5 {
1117            mem.save(make_entry(&format!("compact {i}"), &[i as f32, 1.0, 0.0, 0.0]))
1118                .unwrap();
1119        }
1120
1121        // Manually set distinct weights
1122        mem.cache.activation_weights = vec![1.0, 2.0, 3.0, 4.0, 5.0];
1123
1124        // Delete entries 1 and 3
1125        mem.delete(1).unwrap();
1126        mem.delete(3).unwrap();
1127        mem.compact().unwrap();
1128
1129        // Remaining: indices 0, 2, 4 -> weights 1.0, 3.0, 5.0
1130        assert_eq!(mem.cache.activation_weights.len(), 3);
1131        assert!((mem.cache.activation_weights[0] - 1.0).abs() < 1e-6);
1132        assert!((mem.cache.activation_weights[1] - 3.0).abs() < 1e-6);
1133        assert!((mem.cache.activation_weights[2] - 5.0).abs() < 1e-6);
1134    }
1135
1136    #[test]
1137    fn test_activation_in_search_result() {
1138        let dir = TempDir::new().unwrap();
1139        let config = make_config(&dir);
1140        let mut mem = HDF5Memory::create(config).unwrap();
1141
1142        mem.save(make_entry("search me", &[1.0, 0.0, 0.0, 0.0]))
1143            .unwrap();
1144        mem.save(make_entry("also me", &[0.0, 1.0, 0.0, 0.0]))
1145            .unwrap();
1146
1147        let query = vec![1.0, 0.0, 0.0, 0.0];
1148        let results = mem.hybrid_search(&query, "", 1.0, 0.0, 2);
1149
1150        // Every result should have a populated activation field
1151        for r in &results {
1152            assert!(r.activation > 0.0, "activation should be > 0, got {}", r.activation);
1153        }
1154    }
1155
1156    #[test]
1157    fn test_add_entity_alias_on_memory() {
1158        let dir = TempDir::new().unwrap();
1159        let config = make_config(&dir);
1160        let mut mem = HDF5Memory::create(config).unwrap();
1161
1162        let id = mem.add_entity("Henry", "person", -1).unwrap();
1163        mem.add_entity_alias("my son", id as i64).unwrap();
1164
1165        let aliases = mem.knowledge().get_aliases(id as i64);
1166        assert_eq!(aliases.len(), 1);
1167        assert_eq!(aliases[0], "my son");
1168    }
1169
1170    #[test]
1171    fn tombstone_fraction() {
1172        let dir = TempDir::new().unwrap();
1173        let mut config = make_config(&dir);
1174        config.compact_threshold = 0.0;
1175        let mut mem = HDF5Memory::create(config).unwrap();
1176
1177        assert_eq!(mem.cache.tombstone_fraction(), 0.0);
1178
1179        mem.save(make_entry("a", &[1.0, 0.0, 0.0, 0.0])).unwrap();
1180        mem.save(make_entry("b", &[0.0, 1.0, 0.0, 0.0])).unwrap();
1181        mem.save(make_entry("c", &[0.0, 0.0, 1.0, 0.0])).unwrap();
1182        mem.save(make_entry("d", &[0.0, 0.0, 0.0, 1.0])).unwrap();
1183
1184        mem.delete(0).unwrap();
1185        assert!((mem.cache.tombstone_fraction() - 0.25).abs() < 0.01);
1186
1187        mem.delete(1).unwrap();
1188        assert!((mem.cache.tombstone_fraction() - 0.50).abs() < 0.01);
1189    }
1190}
1191
1192impl HDF5Memory {
1193pub fn set_strategy(&mut self, s: Box<dyn MemoryStrategy>) { self.strategy = Some(s); }
1194    pub fn record(&mut self, exchange: Exchange) -> Result<StrategyOutput> {
1195        let strat = self.strategy.as_ref().expect("call set_strategy() before record()");
1196        let view = memory_strategy::CacheStoreView::new(&self.cache, &self.knowledge);
1197        let output = strat.evaluate(&exchange, &view);
1198        for e in &output.entries {
1199            self.cache.push(e.chunk.clone(), e.embedding.clone(), e.source_channel.clone(), e.timestamp, e.session_id.clone(), e.tags.clone());
1200        }
1201        for eu in &output.entity_updates {
1202            let id = self.knowledge.add_entity(&eu.name, &eu.entity_type, -1);
1203            for a in &eu.aliases { self.knowledge.add_alias(a, id as i64); }
1204        }
1205        if !output.entries.is_empty() || !output.entity_updates.is_empty() { self.flush()?; }
1206        Ok(output)
1207    }
1208}
1209
1210impl HDF5Memory {
1211    pub fn tick_session(&mut self) -> Result<()> {
1212        let d = self.config.decay_factor;
1213        for w in self.cache.activation_weights.iter_mut() {
1214            *w *= d;
1215        }
1216        self.flush()?;
1217        if let Some(ref mut w) = self.wal {
1218            w.truncate()?;
1219        }
1220        Ok(())
1221    }
1222
1223    /// Number of pending WAL entries (0 if WAL disabled).
1224    pub fn wal_pending_count(&self) -> usize {
1225        self.wal.as_ref().map_or(0, |w| w.pending_count() as usize)
1226    }
1227
1228    /// Explicit WAL merge: flush .h5, truncate WAL.
1229    pub fn flush_wal(&mut self) -> Result<()> {
1230        self.flush()?;
1231        if let Some(ref mut w) = self.wal {
1232            w.truncate()?;
1233        }
1234        Ok(())
1235    }
1236}
edgehdf5_memory/lib.rs

edgehdf5_memory/
lib.rs
