1use std::path::{Path, PathBuf};
46
47use mentedb_cognitive::EntityResolver;
48use mentedb_cognitive::interference::{InterferenceDetector, InterferencePair};
49use mentedb_cognitive::llm::EntityMergeGroup;
50use mentedb_cognitive::pain::{PainRegistry, PainSignal};
51use mentedb_cognitive::phantom::{PhantomConfig, PhantomMemory, PhantomTracker};
52use mentedb_cognitive::speculative::{CacheEntry, CacheStats, SpeculativeCache};
53use mentedb_cognitive::stream::{CognitionStream, StreamAlert, StreamConfig};
54use mentedb_cognitive::trajectory::{TrajectoryNode, TrajectoryTracker};
55use mentedb_cognitive::write_inference::{
56 InferredAction, WriteInferenceConfig, WriteInferenceEngine,
57};
58use mentedb_consolidation::archival::{ArchivalConfig, ArchivalDecision, ArchivalPipeline};
59use mentedb_consolidation::compression::{CompressedMemory, MemoryCompressor};
60use mentedb_consolidation::consolidation::{ConsolidationCandidate, ConsolidationEngine};
61use mentedb_consolidation::decay::{DecayConfig, DecayEngine};
62use mentedb_context::{AssemblyConfig, ContextAssembler, ContextWindow, ScoredMemory};
63use mentedb_core::edge::EdgeType;
64use mentedb_core::error::MenteResult;
65use mentedb_core::memory::MemoryType;
66use mentedb_core::types::{AgentId, MemoryId, Timestamp};
67use mentedb_core::{MemoryEdge, MemoryNode, MenteError};
68use mentedb_embedding::provider::EmbeddingProvider;
69use mentedb_graph::GraphManager;
70use mentedb_index::IndexManager;
71use mentedb_query::{Mql, QueryPlan};
72use mentedb_storage::StorageEngine;
73use parking_lot::RwLock;
74use tracing::{debug, info, warn};
75
76pub const VERSION: &str = env!("CARGO_PKG_VERSION");
80
81pub use mentedb_cognitive as cognitive;
83pub use mentedb_consolidation as consolidation;
85pub use mentedb_context as context;
87pub use mentedb_core as core;
89pub use mentedb_graph as graph;
91pub use mentedb_index as index;
93pub use mentedb_query as query;
95pub use mentedb_storage as storage;
97
98pub mod process_turn;
100
101pub mod injection;
103
104#[cfg(feature = "enrichment")]
106pub mod enrichment;
107
108pub mod prelude {
110 pub use mentedb_core::edge::EdgeType;
111 pub use mentedb_core::error::MenteResult;
112 pub use mentedb_core::memory::MemoryType;
113 pub use mentedb_core::types::*;
114 pub use mentedb_core::{MemoryEdge, MemoryNode, MemoryTier, MenteError};
115
116 pub use crate::MenteDb;
117}
118
119use mentedb_storage::PageId;
120use std::collections::HashMap;
122
123#[derive(Debug, Clone)]
129pub struct EnrichmentConfig {
130 pub enabled: bool,
132 pub trigger_interval: u64,
134 pub min_confidence: f32,
136 pub max_enrichment_confidence: f32,
138 pub enable_user_model: bool,
140 pub entity_merge_threshold: f32,
142 pub entity_separate_threshold: f32,
144}
145
146impl Default for EnrichmentConfig {
147 fn default() -> Self {
148 Self {
149 enabled: false,
150 trigger_interval: 50,
151 min_confidence: 0.6,
152 max_enrichment_confidence: 0.7,
153 enable_user_model: false,
154 entity_merge_threshold: 0.7,
155 entity_separate_threshold: 0.4,
156 }
157 }
158}
159
160#[derive(Debug, Clone, Default)]
162pub struct EnrichmentResult {
163 pub memories_stored: usize,
165 pub entities_processed: usize,
167 pub edges_created: usize,
169 pub duplicates_skipped: usize,
171 pub contradictions_found: usize,
173 pub completed_at_turn: u64,
175 pub entities_linked: usize,
177 pub entities_ambiguous: usize,
179}
180
181#[derive(Debug, Clone, Default)]
183pub struct EntityLinkResult {
184 pub linked: usize,
186 pub ambiguous: usize,
188 pub edges_created: usize,
190}
191
192#[derive(Debug, Clone)]
197pub struct EntityLinkResolution {
198 pub canonical: String,
200 pub aliases: Vec<String>,
202 pub confidence: f32,
204}
205
206#[derive(Debug, Clone)]
208pub struct EntitySeparation {
209 pub name_a: String,
210 pub name_b: String,
211}
212
213#[derive(Debug, Clone)]
215pub struct CognitiveConfig {
216 pub write_inference: bool,
218 pub decay_on_recall: bool,
220 pub pain_tracking: bool,
222 pub interference_detection: bool,
224 pub phantom_tracking: bool,
226 pub speculative_cache: bool,
228 pub archival_evaluation: bool,
230 pub inference_config: WriteInferenceConfig,
232 pub decay_config: DecayConfig,
234 pub phantom_config: PhantomConfig,
236 pub archival_config: ArchivalConfig,
238 pub stream_config: StreamConfig,
240 pub enrichment_config: EnrichmentConfig,
242 pub interference_threshold: f32,
244 pub trajectory_max_turns: usize,
246 pub speculative_cache_size: usize,
248 pub pain_max_warnings: usize,
250 pub injection_config: injection::InjectionConfig,
252}
253
254impl Default for CognitiveConfig {
255 fn default() -> Self {
256 Self {
257 write_inference: true,
258 decay_on_recall: true,
259 pain_tracking: true,
260 interference_detection: true,
261 phantom_tracking: true,
262 speculative_cache: true,
263 archival_evaluation: true,
264 inference_config: WriteInferenceConfig::default(),
265 decay_config: DecayConfig::default(),
266 phantom_config: PhantomConfig::default(),
267 archival_config: ArchivalConfig::default(),
268 stream_config: StreamConfig::default(),
269 enrichment_config: EnrichmentConfig::default(),
270 interference_threshold: 0.8,
271 trajectory_max_turns: 100,
272 speculative_cache_size: 10,
273 pain_max_warnings: 5,
274 injection_config: injection::InjectionConfig::default(),
275 }
276 }
277}
278
279pub struct MenteDb {
288 storage: StorageEngine,
289 index: IndexManager,
290 graph: GraphManager,
291 page_map: RwLock<HashMap<MemoryId, PageId>>,
293 embedding_dim: usize,
295 path: PathBuf,
297 embedder: Option<Box<dyn EmbeddingProvider>>,
299 cognitive_config: CognitiveConfig,
301 write_inference: WriteInferenceEngine,
303 decay: DecayEngine,
305 consolidation: ConsolidationEngine,
307 pain: RwLock<PainRegistry>,
309 trajectory: RwLock<TrajectoryTracker>,
311 stream: CognitionStream,
313 phantom: RwLock<PhantomTracker>,
315 speculative: RwLock<SpeculativeCache>,
317 interference: InterferenceDetector,
319 entity_resolver: RwLock<EntityResolver>,
321 compressor: MemoryCompressor,
323 archival: ArchivalPipeline,
325 last_enrichment_turn: RwLock<u64>,
327 enrichment_pending: RwLock<bool>,
329}
330
331pub(crate) fn agent_visible(owner: AgentId, scope: Option<AgentId>) -> bool {
335 match scope {
336 None => true,
337 Some(a) => owner == a || owner.is_nil(),
338 }
339}
340
341impl MenteDb {
342 pub fn open(path: &Path) -> MenteResult<Self> {
344 Self::open_with_config(path, CognitiveConfig::default())
345 }
346
347 pub fn open_with_config(path: &Path, cognitive_config: CognitiveConfig) -> MenteResult<Self> {
349 info!("Opening MenteDB at {}", path.display());
350 let storage = StorageEngine::open(path)?;
351
352 let index_dir = path.join("indexes");
353 let graph_dir = path.join("graph");
354
355 let index = if index_dir.join("hnsw.bin").exists() || index_dir.join("hnsw.json").exists() {
356 debug!("Loading indexes from {}", index_dir.display());
357 IndexManager::load(&index_dir)?
358 } else {
359 IndexManager::default()
360 };
361
362 let graph = GraphManager::open(&graph_dir)?;
365
366 let entries = storage.scan_all_memories();
368 let mut page_map = HashMap::new();
369 for (memory_id, page_id) in &entries {
370 page_map.insert(*memory_id, *page_id);
371 }
372 if !page_map.is_empty() {
373 info!(memories = page_map.len(), "rebuilt page map from storage");
374 }
375
376 for memory_id in page_map.keys() {
380 if !graph.read_graph().contains_node(*memory_id) {
381 graph.add_memory(*memory_id);
382 }
383 }
384
385 let write_inference =
386 WriteInferenceEngine::with_config(cognitive_config.inference_config.clone());
387 let decay = DecayEngine::new(cognitive_config.decay_config.clone());
388 let consolidation = ConsolidationEngine::new();
389 let pain = RwLock::new(PainRegistry::new(cognitive_config.pain_max_warnings));
390 let trajectory = RwLock::new(TrajectoryTracker::new(
391 cognitive_config.trajectory_max_turns,
392 ));
393 let stream = CognitionStream::with_config(cognitive_config.stream_config.clone());
394 let phantom = RwLock::new(PhantomTracker::new(cognitive_config.phantom_config.clone()));
395 let speculative = RwLock::new(SpeculativeCache::new(
396 cognitive_config.speculative_cache_size,
397 0.5,
398 0.4,
399 ));
400 let interference = InterferenceDetector::new(cognitive_config.interference_threshold);
401 let entity_resolver = RwLock::new(EntityResolver::new());
402 let compressor = MemoryCompressor::new();
403 let archival = ArchivalPipeline::new(cognitive_config.archival_config.clone());
404
405 let cognitive_dir = path.join("cognitive");
407 if cognitive_dir.exists() {
408 let _ = trajectory
409 .write()
410 .transitions
411 .load(&cognitive_dir.join("transitions.json"));
412 let _ = speculative
413 .write()
414 .load(&cognitive_dir.join("speculative.json"));
415 let _ = entity_resolver
416 .write()
417 .load(&cognitive_dir.join("entities.json"));
418 }
419
420 Ok(Self {
421 storage,
422 index,
423 graph,
424 page_map: RwLock::new(page_map),
425 embedding_dim: 0,
426 path: path.to_path_buf(),
427 embedder: None,
428 cognitive_config,
429 write_inference,
430 decay,
431 consolidation,
432 pain,
433 trajectory,
434 stream,
435 phantom,
436 speculative,
437 interference,
438 entity_resolver,
439 compressor,
440 archival,
441 last_enrichment_turn: RwLock::new(0),
442 enrichment_pending: RwLock::new(false),
443 })
444 }
445
446 pub fn open_with_embedder(
448 path: &Path,
449 embedder: Box<dyn EmbeddingProvider>,
450 ) -> MenteResult<Self> {
451 let mut db = Self::open(path)?;
452 db.embedding_dim = embedder.dimensions();
453 db.embedder = Some(embedder);
454 Ok(db)
455 }
456
457 pub fn open_with_embedder_and_config(
459 path: &Path,
460 embedder: Box<dyn EmbeddingProvider>,
461 cognitive_config: CognitiveConfig,
462 ) -> MenteResult<Self> {
463 let mut db = Self::open_with_config(path, cognitive_config)?;
464 db.embedding_dim = embedder.dimensions();
465 db.embedder = Some(embedder);
466 Ok(db)
467 }
468
469 pub fn set_embedder(&mut self, embedder: Box<dyn EmbeddingProvider>) {
471 self.embedding_dim = embedder.dimensions();
472 self.embedder = Some(embedder);
473 }
474
475 pub fn embed_text(&self, text: &str) -> MenteResult<Option<Vec<f32>>> {
478 match &self.embedder {
479 Some(e) => Ok(Some(e.embed(text)?)),
480 None => Ok(None),
481 }
482 }
483
484 pub fn store(&self, node: MemoryNode) -> MenteResult<()> {
496 let id = node.id;
497 debug!("Storing memory {}", id);
498
499 if self.embedding_dim > 0
501 && !node.embedding.is_empty()
502 && node.embedding.len() != self.embedding_dim
503 {
504 return Err(MenteError::EmbeddingDimensionMismatch {
505 got: node.embedding.len(),
506 expected: self.embedding_dim,
507 });
508 }
509
510 let page_id = self.storage.store_memory(&node)?;
511 self.page_map.write().insert(id, page_id);
512 self.index.index_memory(&node);
513 self.graph.add_memory(id);
514
515 if self.cognitive_config.write_inference {
517 self.run_write_inference(&node);
518 }
519
520 Ok(())
521 }
522
523 pub fn store_batch(&self, nodes: Vec<MemoryNode>) -> MenteResult<Vec<MemoryId>> {
529 for node in &nodes {
531 if self.embedding_dim > 0
532 && !node.embedding.is_empty()
533 && node.embedding.len() != self.embedding_dim
534 {
535 return Err(MenteError::EmbeddingDimensionMismatch {
536 got: node.embedding.len(),
537 expected: self.embedding_dim,
538 });
539 }
540 }
541
542 let page_ids = self.storage.store_memory_batch(&nodes)?;
543
544 let mut ids = Vec::with_capacity(nodes.len());
545 let mut page_map = self.page_map.write();
546 for (node, page_id) in nodes.iter().zip(page_ids.iter()) {
547 page_map.insert(node.id, *page_id);
548 self.index.index_memory(node);
549 self.graph.add_memory(node.id);
550 ids.push(node.id);
551 }
552 drop(page_map);
553
554 if self.cognitive_config.write_inference {
557 for node in &nodes {
558 self.run_write_inference(node);
559 }
560 }
561
562 Ok(ids)
563 }
564
565 pub fn recall(&self, query: &str) -> MenteResult<ContextWindow> {
571 debug!("Recalling with query: {}", query);
572 let plan = Mql::parse(query)?;
573
574 let scored = self.execute_plan(&plan)?;
575 let config = AssemblyConfig::default();
576 let window = ContextAssembler::assemble(scored, vec![], &config);
577 Ok(window)
578 }
579
580 pub fn recall_similar(&self, embedding: &[f32], k: usize) -> MenteResult<Vec<(MemoryId, f32)>> {
586 self.recall_similar_filtered(embedding, k, None, None)
587 }
588
589 pub fn recall_similar_filtered(
591 &self,
592 embedding: &[f32],
593 k: usize,
594 tags: Option<&[&str]>,
595 time_range: Option<(Timestamp, Timestamp)>,
596 ) -> MenteResult<Vec<(MemoryId, f32)>> {
597 let now = std::time::SystemTime::now()
598 .duration_since(std::time::UNIX_EPOCH)
599 .unwrap_or_default()
600 .as_micros() as u64;
601 self.recall_similar_filtered_at(embedding, k, now, tags, time_range)
602 }
603
604 pub fn recall_similar_at(
610 &self,
611 embedding: &[f32],
612 k: usize,
613 at: Timestamp,
614 ) -> MenteResult<Vec<(MemoryId, f32)>> {
615 self.recall_similar_filtered_at(embedding, k, at, None, None)
616 }
617
618 pub fn recall_similar_filtered_at(
624 &self,
625 embedding: &[f32],
626 k: usize,
627 at: Timestamp,
628 tags: Option<&[&str]>,
629 time_range: Option<(Timestamp, Timestamp)>,
630 ) -> MenteResult<Vec<(MemoryId, f32)>> {
631 self.recall_hybrid_at(embedding, None, k, at, tags, time_range)
632 }
633
634 pub fn recall_hybrid_at(
640 &self,
641 embedding: &[f32],
642 query_text: Option<&str>,
643 k: usize,
644 at: Timestamp,
645 tags: Option<&[&str]>,
646 time_range: Option<(Timestamp, Timestamp)>,
647 ) -> MenteResult<Vec<(MemoryId, f32)>> {
648 self.recall_hybrid_at_mode(embedding, query_text, k, at, tags, false, time_range)
649 }
650
651 #[allow(clippy::too_many_arguments)]
653 pub fn recall_hybrid_at_mode(
654 &self,
655 embedding: &[f32],
656 query_text: Option<&str>,
657 k: usize,
658 at: Timestamp,
659 tags: Option<&[&str]>,
660 tags_or: bool,
661 time_range: Option<(Timestamp, Timestamp)>,
662 ) -> MenteResult<Vec<(MemoryId, f32)>> {
663 self.recall_hybrid_scoped_at_mode(
664 embedding, query_text, k, at, tags, tags_or, time_range, None,
665 )
666 }
667
668 #[allow(clippy::too_many_arguments)]
673 pub fn recall_hybrid_scoped_at_mode(
674 &self,
675 embedding: &[f32],
676 query_text: Option<&str>,
677 k: usize,
678 at: Timestamp,
679 tags: Option<&[&str]>,
680 tags_or: bool,
681 time_range: Option<(Timestamp, Timestamp)>,
682 agent: Option<AgentId>,
683 ) -> MenteResult<Vec<(MemoryId, f32)>> {
684 debug!(
685 "Recall hybrid, k={}, at={}, bm25={}, tags_or={}",
686 k,
687 at,
688 query_text.is_some(),
689 tags_or
690 );
691 let results = self.index.hybrid_search_with_query_mode(
693 embedding,
694 query_text,
695 tags,
696 tags_or,
697 time_range,
698 k * 3,
699 );
700 let graph = self.graph.graph();
701 let pm = self.page_map.read();
702 let filtered: Vec<(MemoryId, f32)> = results
703 .into_iter()
704 .filter(|(id, _)| {
705 let incoming = graph.incoming(*id);
706 let has_active_supersede = incoming.iter().any(|(_, e)| {
707 (e.edge_type == EdgeType::Supersedes || e.edge_type == EdgeType::Contradicts)
708 && e.is_valid_at(at)
709 });
710 !has_active_supersede
711 })
712 .filter(|(id, _)| {
713 if let Some(&page_id) = pm.get(id)
714 && let Ok(node) = self.storage.load_memory(page_id)
715 {
716 node.is_valid_at(at) && agent_visible(node.agent_id, agent)
717 } else {
718 true
719 }
720 })
721 .take(k)
722 .collect();
723 Ok(filtered)
724 }
725
726 pub fn recall_similar_multi(
733 &self,
734 embeddings: &[Vec<f32>],
735 k: usize,
736 tags: Option<&[&str]>,
737 time_range: Option<(Timestamp, Timestamp)>,
738 ) -> MenteResult<Vec<(MemoryId, f32)>> {
739 self.recall_hybrid_multi(embeddings, None, k, tags, time_range)
740 }
741
742 pub fn recall_hybrid_multi(
747 &self,
748 embeddings: &[Vec<f32>],
749 query_texts: Option<&[String]>,
750 k: usize,
751 tags: Option<&[&str]>,
752 time_range: Option<(Timestamp, Timestamp)>,
753 ) -> MenteResult<Vec<(MemoryId, f32)>> {
754 self.recall_hybrid_multi_mode(embeddings, query_texts, k, tags, false, time_range)
755 }
756
757 pub fn recall_hybrid_multi_mode(
759 &self,
760 embeddings: &[Vec<f32>],
761 query_texts: Option<&[String]>,
762 k: usize,
763 tags: Option<&[&str]>,
764 tags_or: bool,
765 time_range: Option<(Timestamp, Timestamp)>,
766 ) -> MenteResult<Vec<(MemoryId, f32)>> {
767 use std::collections::HashMap;
768
769 let rrf_k: f32 = 60.0;
770 let mut rrf_scores: HashMap<MemoryId, f32> = HashMap::new();
771
772 let now = std::time::SystemTime::now()
773 .duration_since(std::time::UNIX_EPOCH)
774 .unwrap_or_default()
775 .as_micros() as u64;
776
777 for (i, emb) in embeddings.iter().enumerate() {
778 let qt = query_texts.and_then(|texts| texts.get(i).map(|s| s.as_str()));
779 let results = self.recall_hybrid_at_mode(emb, qt, k, now, tags, tags_or, time_range)?;
780 for (rank, (id, _score)) in results.iter().enumerate() {
781 *rrf_scores.entry(*id).or_insert(0.0) += 1.0 / (rrf_k + rank as f32);
782 }
783 }
784
785 let mut merged: Vec<(MemoryId, f32)> = rrf_scores.into_iter().collect();
786 merged.sort_unstable_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
787 merged.truncate(k);
788 Ok(merged)
789 }
790
791 pub fn invalidate_memory(&self, id: MemoryId, at: Timestamp) -> MenteResult<()> {
796 debug!("Invalidating memory {} at {}", id, at);
797 let page_id = self
798 .page_map
799 .read()
800 .get(&id)
801 .copied()
802 .ok_or(MenteError::MemoryNotFound(id))?;
803 let mut node = self.storage.load_memory(page_id)?;
804 node.invalidate(at);
805 self.storage.update_memory(page_id, &node)?;
806 Ok(())
807 }
808
809 pub fn relate(&self, edge: MemoryEdge) -> MenteResult<()> {
811 debug!("Relating {} -> {}", edge.source, edge.target);
812 self.graph.add_relationship(&edge)?;
813 Ok(())
814 }
815
816 pub fn get_memory(&self, id: MemoryId) -> MenteResult<MemoryNode> {
818 let page_id = self
819 .page_map
820 .read()
821 .get(&id)
822 .copied()
823 .ok_or(MenteError::MemoryNotFound(id))?;
824 self.storage.load_memory(page_id)
825 }
826
827 pub fn memory_ids(&self) -> Vec<MemoryId> {
829 self.page_map.read().keys().copied().collect()
830 }
831
832 pub fn memory_count(&self) -> usize {
834 self.page_map.read().len()
835 }
836
837 pub fn forget(&self, id: MemoryId) -> MenteResult<()> {
839 debug!("Forgetting memory {}", id);
840
841 if let Some(&page_id) = self.page_map.read().get(&id) {
842 if let Ok(node) = self.storage.load_memory(page_id) {
843 self.index.remove_memory(id, &node);
844 }
845 self.storage.delete_memory(page_id)?;
848 }
849
850 self.graph.remove_memory(id);
851 self.page_map.write().remove(&id);
852 Ok(())
853 }
854
855 pub fn graph(&self) -> &GraphManager {
857 &self.graph
858 }
859
860 #[deprecated(note = "GraphManager now uses interior mutability; use graph() instead")]
862 pub fn graph_mut(&mut self) -> &mut GraphManager {
863 &mut self.graph
864 }
865
866 pub fn cognitive_config(&self) -> &CognitiveConfig {
868 &self.cognitive_config
869 }
870
871 fn run_write_inference(&self, new_memory: &MemoryNode) {
881 let candidates = if !new_memory.embedding.is_empty() {
884 let now = std::time::SystemTime::now()
885 .duration_since(std::time::UNIX_EPOCH)
886 .unwrap_or_default()
887 .as_micros() as u64;
888 self.recall_hybrid_at(&new_memory.embedding, None, 20, now, None, None)
889 .unwrap_or_default()
890 } else {
891 vec![]
892 };
893
894 if candidates.is_empty() {
895 return;
896 }
897
898 let pm = self.page_map.read();
900 let existing: Vec<MemoryNode> = candidates
901 .iter()
902 .filter(|(id, _)| *id != new_memory.id)
903 .filter_map(|(id, _)| {
904 pm.get(id)
905 .and_then(|&pid| self.storage.load_memory(pid).ok())
906 })
907 .collect();
908 drop(pm);
909
910 if existing.is_empty() {
911 return;
912 }
913
914 let actions = self
915 .write_inference
916 .infer_on_write(new_memory, &existing, &[]);
917
918 let action_count = actions.len();
919 for action in actions {
920 if let Err(e) = self.apply_inferred_action(action) {
921 warn!("Failed to apply inferred action: {}", e);
922 }
923 }
924 if action_count > 0 {
925 debug!(
926 "Write inference for {} produced {} actions",
927 new_memory.id, action_count
928 );
929 }
930 }
931
932 fn apply_inferred_action(&self, action: InferredAction) -> MenteResult<()> {
934 match action {
935 InferredAction::CreateEdge {
936 source,
937 target,
938 edge_type,
939 weight,
940 } => {
941 let now = std::time::SystemTime::now()
942 .duration_since(std::time::UNIX_EPOCH)
943 .unwrap_or_default()
944 .as_micros() as u64;
945 let edge = MemoryEdge {
946 source,
947 target,
948 edge_type,
949 weight,
950 created_at: now,
951 valid_from: None,
952 valid_until: None,
953 label: None,
954 };
955 debug!(
956 "Auto-creating {:?} edge {} -> {}",
957 edge_type, source, target
958 );
959 self.graph.add_relationship(&edge)?;
960 }
961 InferredAction::InvalidateMemory {
962 memory,
963 superseded_by,
964 valid_until,
965 } => {
966 debug!(
967 "Invalidating memory {} (superseded by {})",
968 memory, superseded_by
969 );
970 self.invalidate_memory(memory, valid_until)?;
971 let now = std::time::SystemTime::now()
973 .duration_since(std::time::UNIX_EPOCH)
974 .unwrap_or_default()
975 .as_micros() as u64;
976 let edge = MemoryEdge {
977 source: superseded_by,
978 target: memory,
979 edge_type: EdgeType::Supersedes,
980 weight: 1.0,
981 created_at: now,
982 valid_from: None,
983 valid_until: None,
984 label: None,
985 };
986 self.graph.add_relationship(&edge)?;
987 }
988 InferredAction::MarkObsolete {
989 memory,
990 superseded_by,
991 } => {
992 debug!(
993 "Marking {} obsolete (superseded by {})",
994 memory, superseded_by
995 );
996 let now = std::time::SystemTime::now()
997 .duration_since(std::time::UNIX_EPOCH)
998 .unwrap_or_default()
999 .as_micros() as u64;
1000 self.invalidate_memory(memory, now)?;
1001 let edge = MemoryEdge {
1002 source: superseded_by,
1003 target: memory,
1004 edge_type: EdgeType::Supersedes,
1005 weight: 1.0,
1006 created_at: now,
1007 valid_from: None,
1008 valid_until: None,
1009 label: None,
1010 };
1011 self.graph.add_relationship(&edge)?;
1012 }
1013 InferredAction::FlagContradiction {
1014 existing,
1015 new,
1016 reason,
1017 } => {
1018 debug!(
1019 "Contradiction detected: {} vs {} — {}",
1020 existing, new, reason
1021 );
1022 let now = std::time::SystemTime::now()
1023 .duration_since(std::time::UNIX_EPOCH)
1024 .unwrap_or_default()
1025 .as_micros() as u64;
1026 let edge = MemoryEdge {
1027 source: new,
1028 target: existing,
1029 edge_type: EdgeType::Contradicts,
1030 weight: 1.0,
1031 created_at: now,
1032 valid_from: None,
1033 valid_until: None,
1034 label: Some(reason),
1035 };
1036 self.graph.add_relationship(&edge)?;
1037 }
1038 InferredAction::UpdateConfidence {
1039 memory,
1040 new_confidence,
1041 } => {
1042 debug!("Updating confidence for {} to {}", memory, new_confidence);
1043 if let Ok(mut node) = self.get_memory(memory) {
1044 node.confidence = new_confidence;
1045 if let Some(&pid) = self.page_map.read().get(&memory) {
1046 self.storage.update_memory(pid, &node)?;
1047 }
1048 }
1049 }
1050 InferredAction::PropagateBeliefChange { root, delta } => {
1051 debug!("Propagating belief change from {} (delta={})", root, delta);
1052 if let Ok(node) = self.get_memory(root) {
1053 let new_confidence = (node.confidence + delta).clamp(0.0, 1.0);
1054 let affected = self.graph.propagate_belief_change(root, new_confidence);
1055 for (affected_id, new_conf) in affected {
1056 if let Ok(mut affected_node) = self.get_memory(affected_id) {
1057 affected_node.confidence = new_conf;
1058 if let Some(&pid) = self.page_map.read().get(&affected_id)
1059 && let Err(e) = self.storage.update_memory(pid, &affected_node)
1060 {
1061 warn!("Failed to persist belief update for {affected_id}: {e}");
1062 }
1063 }
1064 }
1065 }
1066 }
1067 InferredAction::UpdateContent {
1068 memory,
1069 new_content,
1070 reason,
1071 } => {
1072 debug!("Updating content of {}: {}", memory, reason);
1073 if let Ok(mut node) = self.get_memory(memory) {
1074 node.content = new_content;
1075 if let Some(&pid) = self.page_map.read().get(&memory) {
1076 self.storage.update_memory(pid, &node)?;
1077 }
1078 self.index.remove_memory(memory, &node);
1079 self.index.index_memory(&node);
1080 }
1081 }
1082 }
1083 Ok(())
1084 }
1085
1086 pub fn apply_decay(&self, memories: &mut [MemoryNode]) {
1095 let now = std::time::SystemTime::now()
1096 .duration_since(std::time::UNIX_EPOCH)
1097 .unwrap_or_default()
1098 .as_micros() as u64;
1099 self.decay.apply_decay_batch(memories, now);
1100 }
1101
1102 pub fn compute_decayed_salience(&self, memory: &MemoryNode) -> f32 {
1104 let now = std::time::SystemTime::now()
1105 .duration_since(std::time::UNIX_EPOCH)
1106 .unwrap_or_default()
1107 .as_micros() as u64;
1108 self.decay.compute_decay(
1109 memory.salience,
1110 memory.created_at,
1111 memory.accessed_at,
1112 memory.access_count,
1113 now,
1114 )
1115 }
1116
1117 pub fn apply_decay_global(&self) -> MenteResult<usize> {
1122 let now = std::time::SystemTime::now()
1123 .duration_since(std::time::UNIX_EPOCH)
1124 .unwrap_or_default()
1125 .as_micros() as u64;
1126 let page_ids: Vec<PageId> = self.page_map.read().values().copied().collect();
1127
1128 let mut updated = 0;
1129 for pid in &page_ids {
1130 if let Ok(mut node) = self.storage.load_memory(*pid) {
1131 let new_salience = self.decay.compute_decay(
1132 node.salience,
1133 node.created_at,
1134 node.accessed_at,
1135 node.access_count,
1136 now,
1137 );
1138 if (new_salience - node.salience).abs() > 0.001 {
1139 node.salience = new_salience;
1140 self.storage.update_memory(*pid, &node)?;
1141 updated += 1;
1142 }
1143 }
1144 }
1145 if updated > 0 {
1146 info!("Decay pass updated {} memories", updated);
1147 }
1148 Ok(updated)
1149 }
1150
1151 pub fn find_consolidation_candidates(
1160 &self,
1161 min_cluster_size: usize,
1162 similarity_threshold: f32,
1163 ) -> MenteResult<Vec<ConsolidationCandidate>> {
1164 let now = std::time::SystemTime::now()
1165 .duration_since(std::time::UNIX_EPOCH)
1166 .unwrap_or_default()
1167 .as_micros() as u64;
1168
1169 let pm = self.page_map.read();
1171 let eligible: Vec<MemoryNode> = pm
1172 .values()
1173 .filter_map(|pid| self.storage.load_memory(*pid).ok())
1174 .filter(|node| ConsolidationEngine::should_consolidate(node, now))
1175 .collect();
1176 drop(pm);
1177
1178 if eligible.is_empty() {
1179 return Ok(vec![]);
1180 }
1181
1182 Ok(self
1183 .consolidation
1184 .find_candidates(&eligible, min_cluster_size, similarity_threshold))
1185 }
1186
1187 pub fn consolidate_cluster(&self, memory_ids: &[MemoryId]) -> MenteResult<MemoryId> {
1192 let pm = self.page_map.read();
1193 let cluster: Vec<MemoryNode> = memory_ids
1194 .iter()
1195 .filter_map(|id| {
1196 pm.get(id)
1197 .and_then(|&pid| self.storage.load_memory(pid).ok())
1198 })
1199 .collect();
1200 drop(pm);
1201
1202 if cluster.len() < 2 {
1203 return Err(MenteError::Query(
1204 "consolidation requires at least 2 memories".into(),
1205 ));
1206 }
1207
1208 let result = self.consolidation.consolidate(&cluster);
1209
1210 let agent_id = cluster[0].agent_id;
1212 let mut consolidated = MemoryNode::new(
1213 agent_id,
1214 result.new_type,
1215 result.summary,
1216 result.combined_embedding,
1217 );
1218 consolidated.confidence = result.combined_confidence;
1219
1220 let consolidated_id = consolidated.id;
1221 self.store(consolidated)?;
1222
1223 let now = std::time::SystemTime::now()
1225 .duration_since(std::time::UNIX_EPOCH)
1226 .unwrap_or_default()
1227 .as_micros() as u64;
1228 for source_id in &result.source_memories {
1229 let _ = self.invalidate_memory(*source_id, now);
1230 let edge = MemoryEdge {
1231 source: consolidated_id,
1232 target: *source_id,
1233 edge_type: EdgeType::Derived,
1234 weight: 1.0,
1235 created_at: now,
1236 valid_from: None,
1237 valid_until: None,
1238 label: None,
1239 };
1240 let _ = self.graph.add_relationship(&edge);
1241 }
1242
1243 info!(
1244 "Consolidated {} memories into {}",
1245 result.source_memories.len(),
1246 consolidated_id
1247 );
1248 Ok(consolidated_id)
1249 }
1250
1251 pub fn close(&self) -> MenteResult<()> {
1253 info!("Closing MenteDB");
1254 self.flush()?;
1255 self.storage.close()?;
1256 Ok(())
1257 }
1258
1259 #[doc(hidden)]
1263 pub fn simulate_crash(self) {
1264 self.storage.release_process_lock();
1265 std::mem::forget(self);
1266 }
1267
1268 pub fn rebuild_indexes(&self) -> MenteResult<usize> {
1273 info!("Rebuilding indexes from storage...");
1274 let ids: Vec<MemoryId> = self.page_map.read().keys().copied().collect();
1275 let total = ids.len();
1276 let mut indexed = 0usize;
1277 for id in ids {
1278 if let Ok(node) = self.get_memory(id) {
1279 self.index.index_memory(&node);
1280 indexed += 1;
1281 }
1282 }
1283 self.index.save(&self.path.join("indexes"))?;
1284 info!(indexed, total, "index rebuild complete");
1285 Ok(indexed)
1286 }
1287
1288 pub fn flush(&self) -> MenteResult<()> {
1293 debug!("Flushing MenteDB to disk");
1294 self.index.save(&self.path.join("indexes"))?;
1295 self.graph.save(&self.path.join("graph"))?;
1296 self.storage.checkpoint()?;
1297
1298 let cognitive_dir = self.path.join("cognitive");
1300 if std::fs::create_dir_all(&cognitive_dir).is_ok() {
1301 let _ = self
1302 .trajectory
1303 .read()
1304 .transitions
1305 .save(&cognitive_dir.join("transitions.json"), 1);
1306 let _ = self
1307 .speculative
1308 .read()
1309 .save(&cognitive_dir.join("speculative.json"), 0);
1310 let _ = self
1311 .entity_resolver
1312 .read()
1313 .save(&cognitive_dir.join("entities.json"));
1314 }
1315 Ok(())
1316 }
1317
1318 fn execute_plan(&self, plan: &QueryPlan) -> MenteResult<Vec<ScoredMemory>> {
1320 match plan {
1321 QueryPlan::VectorSearch { query, k, .. } => {
1322 let hits = self.index.hybrid_search(query, None, None, *k);
1323 self.load_scored_memories(&hits)
1324 }
1325 QueryPlan::TagScan { tags, limit, .. } => {
1326 let tag_refs: Vec<&str> = tags.iter().map(|s| s.as_str()).collect();
1327 let k = limit.unwrap_or(10);
1328 let hits = self.index.hybrid_search(&[], Some(&tag_refs), None, k);
1330 self.load_scored_memories(&hits)
1331 }
1332 QueryPlan::TemporalScan { start, end, .. } => {
1333 let hits = self
1334 .index
1335 .hybrid_search(&[], None, Some((*start, *end)), 100);
1336 self.load_scored_memories(&hits)
1337 }
1338 QueryPlan::GraphTraversal { start, depth, .. } => {
1339 let (ids, _edges) = self.graph.get_context_subgraph(*start, *depth);
1340 let pm = self.page_map.read();
1341 let scored: Vec<ScoredMemory> = ids
1342 .iter()
1343 .filter_map(|id| {
1344 pm.get(id).and_then(|&pid| {
1345 self.storage.load_memory(pid).ok().map(|node| ScoredMemory {
1346 memory: node,
1347 score: 1.0,
1348 })
1349 })
1350 })
1351 .collect();
1352 Ok(scored)
1353 }
1354 QueryPlan::PointLookup { id } => {
1355 let page_id = self
1356 .page_map
1357 .read()
1358 .get(id)
1359 .copied()
1360 .ok_or(MenteError::MemoryNotFound(*id))?;
1361 let node = self.storage.load_memory(page_id)?;
1362 Ok(vec![ScoredMemory {
1363 memory: node,
1364 score: 1.0,
1365 }])
1366 }
1367 _ => Ok(vec![]),
1368 }
1369 }
1370
1371 fn load_scored_memories(&self, hits: &[(MemoryId, f32)]) -> MenteResult<Vec<ScoredMemory>> {
1376 let pm = self.page_map.read();
1377 let now = if self.cognitive_config.decay_on_recall {
1378 std::time::SystemTime::now()
1379 .duration_since(std::time::UNIX_EPOCH)
1380 .unwrap_or_default()
1381 .as_micros() as u64
1382 } else {
1383 0
1384 };
1385
1386 let mut scored = Vec::with_capacity(hits.len());
1387 for &(id, score) in hits {
1388 if let Some(&page_id) = pm.get(&id)
1389 && let Ok(node) = self.storage.load_memory(page_id)
1390 {
1391 let final_score = if self.cognitive_config.decay_on_recall {
1392 let decayed_salience = self.decay.compute_decay(
1393 node.salience,
1394 node.created_at,
1395 node.accessed_at,
1396 node.access_count,
1397 now,
1398 );
1399 score * 0.7 + decayed_salience * 0.3
1403 } else {
1404 score
1405 };
1406 scored.push(ScoredMemory {
1407 memory: node,
1408 score: final_score,
1409 });
1410 }
1411 }
1412 if self.cognitive_config.decay_on_recall {
1414 scored.sort_unstable_by(|a, b| {
1415 b.score
1416 .partial_cmp(&a.score)
1417 .unwrap_or(std::cmp::Ordering::Equal)
1418 });
1419 }
1420 Ok(scored)
1421 }
1422
1423 pub fn record_pain(&self, signal: PainSignal) {
1432 if self.cognitive_config.pain_tracking {
1433 self.pain.write().record_pain(signal);
1434 }
1435 }
1436
1437 pub fn get_pain_warnings(&self, context_keywords: &[String]) -> Vec<PainSignal> {
1442 if !self.cognitive_config.pain_tracking {
1443 return vec![];
1444 }
1445 let registry = self.pain.read();
1446 registry
1447 .get_pain_for_context(context_keywords)
1448 .into_iter()
1449 .cloned()
1450 .collect()
1451 }
1452
1453 pub fn format_pain_warnings(&self, signals: &[&PainSignal]) -> String {
1455 self.pain.read().format_pain_warnings(signals)
1456 }
1457
1458 pub fn decay_pain(&self) {
1460 let now = std::time::SystemTime::now()
1461 .duration_since(std::time::UNIX_EPOCH)
1462 .unwrap_or_default()
1463 .as_micros() as u64;
1464 self.pain.write().decay_all(now);
1465 }
1466
1467 pub fn all_pain_signals(&self) -> Vec<PainSignal> {
1469 self.pain.read().all_signals().to_vec()
1470 }
1471
1472 pub fn record_trajectory_turn(&self, turn: TrajectoryNode) {
1481 self.trajectory.write().record_turn(turn);
1482 }
1483
1484 pub fn get_resume_context(&self) -> Option<String> {
1488 self.trajectory.read().get_resume_context()
1489 }
1490
1491 pub fn predict_next_topics(&self) -> Vec<String> {
1495 self.trajectory.read().predict_next_topics()
1496 }
1497
1498 pub fn get_trajectory(&self) -> Vec<TrajectoryNode> {
1500 self.trajectory.read().get_trajectory().to_vec()
1501 }
1502
1503 pub fn reinforce_transition(&self, hit_topic: &str) {
1505 self.trajectory.write().reinforce_transition(hit_topic);
1506 }
1507
1508 pub fn feed_stream_token(&self, token: &str) {
1517 self.stream.feed_token(token);
1518 }
1519
1520 pub fn check_stream_alerts(&self, known_facts: &[(MemoryId, String)]) -> Vec<StreamAlert> {
1525 self.stream.check_alerts(known_facts)
1526 }
1527
1528 pub fn drain_stream_buffer(&self) -> String {
1530 self.stream.drain_buffer()
1531 }
1532
1533 pub fn detect_phantoms(
1542 &self,
1543 content: &str,
1544 known_entities: &[String],
1545 turn_id: u64,
1546 ) -> Vec<PhantomMemory> {
1547 if !self.cognitive_config.phantom_tracking {
1548 return vec![];
1549 }
1550 self.phantom
1551 .write()
1552 .detect_gaps(content, known_entities, turn_id)
1553 }
1554
1555 pub fn resolve_phantom(&self, phantom_id: MemoryId) {
1557 self.phantom.write().resolve(phantom_id.into());
1558 }
1559
1560 pub fn get_active_phantoms(&self) -> Vec<PhantomMemory> {
1562 self.phantom
1563 .read()
1564 .get_active_phantoms()
1565 .into_iter()
1566 .cloned()
1567 .collect()
1568 }
1569
1570 pub fn format_phantom_warnings(&self) -> String {
1572 self.phantom.read().format_phantom_warnings()
1573 }
1574
1575 pub fn register_entity(&self, entity: &str) {
1577 self.phantom.write().register_entity(entity);
1578 }
1579
1580 pub fn register_entities(&self, entities: &[&str]) {
1582 self.phantom.write().register_entities(entities);
1583 }
1584
1585 pub fn try_speculative_hit(
1594 &self,
1595 query: &str,
1596 query_embedding: Option<&[f32]>,
1597 ) -> Option<CacheEntry> {
1598 if !self.cognitive_config.speculative_cache {
1599 return None;
1600 }
1601 self.speculative.write().try_hit(query, query_embedding)
1602 }
1603
1604 pub fn pre_assemble_speculative<F>(&self, predictions: Vec<String>, builder: F)
1609 where
1610 F: Fn(&str) -> Option<(String, Vec<MemoryId>, Option<Vec<f32>>)>,
1611 {
1612 if self.cognitive_config.speculative_cache {
1613 self.speculative.write().pre_assemble(predictions, builder);
1614 }
1615 }
1616
1617 pub fn evict_stale_speculative(&self, max_age_us: u64) {
1619 let now = std::time::SystemTime::now()
1620 .duration_since(std::time::UNIX_EPOCH)
1621 .unwrap_or_default()
1622 .as_micros() as u64;
1623 self.speculative.write().evict_stale(max_age_us, now);
1624 }
1625
1626 pub fn speculative_cache_stats(&self) -> CacheStats {
1628 self.speculative.read().stats()
1629 }
1630
1631 pub fn detect_interference(&self, memories: &[MemoryNode]) -> Vec<InterferencePair> {
1641 if !self.cognitive_config.interference_detection {
1642 return vec![];
1643 }
1644 self.interference.detect_interference(memories)
1645 }
1646
1647 pub fn generate_disambiguation(&self, a: &MemoryNode, b: &MemoryNode) -> String {
1649 self.interference.generate_disambiguation(a, b)
1650 }
1651
1652 pub fn arrange_with_separation(
1654 memories: Vec<MemoryId>,
1655 pairs: &[InterferencePair],
1656 ) -> Vec<MemoryId> {
1657 InterferenceDetector::arrange_with_separation(memories, pairs)
1658 }
1659
1660 pub fn resolve_entity(&self, name: &str) -> mentedb_cognitive::ResolvedEntity {
1668 self.entity_resolver.read().resolve(name)
1669 }
1670
1671 pub fn add_entity_alias(&self, alias: &str, canonical: &str, confidence: f32) {
1673 self.entity_resolver
1674 .write()
1675 .add_alias(alias, canonical, confidence);
1676 }
1677
1678 pub fn get_canonical_entity(&self, name: &str) -> Option<String> {
1680 self.entity_resolver.read().get_canonical(name).cloned()
1681 }
1682
1683 pub fn known_entities(&self) -> Vec<String> {
1685 self.entity_resolver.read().known_entities()
1686 }
1687
1688 pub fn compress_memory(&self, memory: &MemoryNode) -> CompressedMemory {
1697 self.compressor.compress(memory)
1698 }
1699
1700 pub fn compress_memories(&self, memories: &[MemoryNode]) -> Vec<CompressedMemory> {
1702 self.compressor.compress_batch(memories)
1703 }
1704
1705 pub fn estimate_tokens(text: &str) -> usize {
1707 MemoryCompressor::estimate_tokens(text)
1708 }
1709
1710 pub fn evaluate_archival(&self, memory: &MemoryNode) -> ArchivalDecision {
1718 if !self.cognitive_config.archival_evaluation {
1719 return ArchivalDecision::Keep;
1720 }
1721 let now = std::time::SystemTime::now()
1722 .duration_since(std::time::UNIX_EPOCH)
1723 .unwrap_or_default()
1724 .as_micros() as u64;
1725 self.archival.evaluate(memory, now)
1726 }
1727
1728 pub fn evaluate_archival_batch(
1730 &self,
1731 memories: &[MemoryNode],
1732 ) -> Vec<(MemoryId, ArchivalDecision)> {
1733 if !self.cognitive_config.archival_evaluation {
1734 return memories
1735 .iter()
1736 .map(|m| (m.id, ArchivalDecision::Keep))
1737 .collect();
1738 }
1739 let now = std::time::SystemTime::now()
1740 .duration_since(std::time::UNIX_EPOCH)
1741 .unwrap_or_default()
1742 .as_micros() as u64;
1743 self.archival.evaluate_batch(memories, now)
1744 }
1745
1746 pub fn evaluate_archival_global(&self) -> MenteResult<Vec<(MemoryId, ArchivalDecision)>> {
1751 let now = std::time::SystemTime::now()
1752 .duration_since(std::time::UNIX_EPOCH)
1753 .unwrap_or_default()
1754 .as_micros() as u64;
1755 let pm = self.page_map.read();
1756 let memories: Vec<MemoryNode> = pm
1757 .values()
1758 .filter_map(|pid| self.storage.load_memory(*pid).ok())
1759 .collect();
1760 drop(pm);
1761 Ok(self.archival.evaluate_batch(&memories, now))
1762 }
1763
1764 pub fn needs_enrichment(&self) -> bool {
1770 if !self.cognitive_config.enrichment_config.enabled {
1771 return false;
1772 }
1773 *self.enrichment_pending.read()
1774 }
1775
1776 pub fn last_enrichment_turn(&self) -> u64 {
1778 *self.last_enrichment_turn.read()
1779 }
1780
1781 pub fn request_enrichment(&self) {
1783 *self.enrichment_pending.write() = true;
1784 }
1785
1786 pub fn enrichment_candidates(&self) -> Vec<MemoryNode> {
1791 let last_turn = *self.last_enrichment_turn.read();
1792 let page_ids: Vec<PageId> = self.page_map.read().values().copied().collect();
1793 let mut candidates: Vec<MemoryNode> = page_ids
1794 .iter()
1795 .filter_map(|pid| self.storage.load_memory(*pid).ok())
1796 .filter(|m| {
1797 m.memory_type == mentedb_core::memory::MemoryType::Episodic
1798 && !m.tags.contains(&"source:enrichment".to_string())
1799 && m.created_at > last_turn
1800 })
1801 .collect();
1802 candidates.sort_by_key(|m| m.created_at);
1803 candidates
1804 }
1805
1806 pub fn store_enrichment_memories(
1815 &self,
1816 memories: Vec<MemoryNode>,
1817 source_ids: &[MemoryId],
1818 ) -> MenteResult<(usize, usize)> {
1819 let max_conf = self
1820 .cognitive_config
1821 .enrichment_config
1822 .max_enrichment_confidence;
1823 let mut stored = 0usize;
1824 let mut edges = 0usize;
1825
1826 for mut mem in memories {
1827 if !mem.tags.contains(&"source:enrichment".to_string()) {
1829 mem.tags.push("source:enrichment".to_string());
1830 }
1831 if mem.confidence > max_conf {
1833 mem.confidence = max_conf;
1834 }
1835
1836 let mem_id = mem.id;
1837 self.store(mem)?;
1838 stored += 1;
1839
1840 let now = std::time::SystemTime::now()
1842 .duration_since(std::time::UNIX_EPOCH)
1843 .unwrap_or_default()
1844 .as_micros() as u64;
1845 for src_id in source_ids {
1846 let edge = MemoryEdge {
1847 source: mem_id,
1848 target: *src_id,
1849 edge_type: EdgeType::Derived,
1850 weight: 0.8,
1851 created_at: now,
1852 valid_from: None,
1853 valid_until: None,
1854 label: Some("enrichment".to_string()),
1855 };
1856 if self.relate(edge).is_ok() {
1857 edges += 1;
1858 }
1859 }
1860 }
1861
1862 debug!(stored, edges, "enrichment memories stored");
1863 Ok((stored, edges))
1864 }
1865
1866 pub fn mark_enrichment_complete(&self, turn_id: u64) {
1868 *self.last_enrichment_turn.write() = turn_id;
1869 *self.enrichment_pending.write() = false;
1870 debug!(turn_id, "enrichment cycle complete");
1871 }
1872
1873 pub fn enrichment_config(&self) -> &EnrichmentConfig {
1875 &self.cognitive_config.enrichment_config
1876 }
1877
1878 pub fn all_entity_names(&self) -> Vec<String> {
1883 let page_ids: Vec<PageId> = self.page_map.read().values().copied().collect();
1884 let mut names = std::collections::HashSet::new();
1885 for pid in &page_ids {
1886 if let Ok(mem) = self.storage.load_memory(*pid) {
1887 for tag in &mem.tags {
1888 if let Some(name) = tag.strip_prefix("entity:") {
1889 names.insert(name.to_lowercase().trim().to_string());
1890 }
1891 }
1892 }
1893 }
1894 let mut sorted: Vec<String> = names.into_iter().collect();
1895 sorted.sort();
1896 sorted
1897 }
1898
1899 pub fn unresolved_entity_names(&self) -> Vec<String> {
1904 let all_names = self.all_entity_names();
1905 self.entity_resolver.read().unresolved_names(&all_names)
1906 }
1907
1908 pub fn entity_names_with_context(&self) -> Vec<(String, Option<String>)> {
1914 let page_ids: Vec<PageId> = self.page_map.read().values().copied().collect();
1915 let mut entity_contexts: HashMap<String, String> = HashMap::new();
1916
1917 for pid in &page_ids {
1918 if let Ok(mem) = self.storage.load_memory(*pid) {
1919 for tag in &mem.tags {
1920 if let Some(name) = tag.strip_prefix("entity:") {
1921 let normalized = name.to_lowercase().trim().to_string();
1922 entity_contexts
1924 .entry(normalized)
1925 .and_modify(|existing| {
1926 if existing.len() < 300 {
1928 existing.push_str(" | ");
1929 let remaining = 500usize.saturating_sub(existing.len());
1930 existing
1931 .push_str(&mem.content[..mem.content.len().min(remaining)]);
1932 }
1933 })
1934 .or_insert_with(|| {
1935 mem.content[..mem.content.len().min(300)].to_string()
1936 });
1937 break;
1938 }
1939 }
1940 }
1941 }
1942
1943 entity_contexts
1944 .into_iter()
1945 .map(|(name, ctx)| (name, Some(ctx)))
1946 .collect()
1947 }
1948
1949 pub fn apply_entity_link_resolutions(
1956 &self,
1957 merge_groups: &[EntityLinkResolution],
1958 separations: &[EntitySeparation],
1959 ) -> MenteResult<EntityLinkResult> {
1960 let mut result = EntityLinkResult::default();
1961 let now = std::time::SystemTime::now()
1962 .duration_since(std::time::UNIX_EPOCH)
1963 .unwrap_or_default()
1964 .as_micros() as u64;
1965
1966 let entity_memory_map = self.build_entity_memory_map();
1968
1969 let mut resolver = self.entity_resolver.write();
1970
1971 for group in merge_groups {
1972 let mut aliases: Vec<String> = group.aliases.clone();
1974 aliases.retain(|a| a.to_lowercase() != group.canonical.to_lowercase());
1975 resolver.learn_group(&EntityMergeGroup {
1976 canonical: group.canonical.clone(),
1977 aliases,
1978 confidence: group.confidence,
1979 });
1980
1981 let mut group_memory_ids: Vec<MemoryId> = Vec::new();
1983
1984 let canonical_norm = group.canonical.to_lowercase();
1986 if let Some(ids) = entity_memory_map.get(&canonical_norm) {
1987 group_memory_ids.extend(ids);
1988 }
1989
1990 for alias in &group.aliases {
1992 let alias_norm = alias.to_lowercase();
1993 if let Some(ids) = entity_memory_map.get(&alias_norm) {
1994 group_memory_ids.extend(ids);
1995 }
1996 }
1997
1998 group_memory_ids.sort();
1999 group_memory_ids.dedup();
2000
2001 let label = format!("entity_link:{}", canonical_norm);
2003 for i in 0..group_memory_ids.len() {
2004 for j in (i + 1)..group_memory_ids.len() {
2005 let a_id = group_memory_ids[i];
2006 let b_id = group_memory_ids[j];
2007
2008 let graph = self.graph.read_graph();
2010 let already_linked = graph.outgoing(a_id).iter().any(|(tid, e)| {
2011 *tid == b_id
2012 && e.edge_type == EdgeType::Related
2013 && e.label
2014 .as_ref()
2015 .is_some_and(|l| l.starts_with("entity_link:"))
2016 });
2017 drop(graph);
2018
2019 if already_linked {
2020 continue;
2021 }
2022
2023 let edge = MemoryEdge {
2024 source: a_id,
2025 target: b_id,
2026 edge_type: EdgeType::Related,
2027 weight: group.confidence,
2028 created_at: now,
2029 valid_from: None,
2030 valid_until: None,
2031 label: Some(label.clone()),
2032 };
2033 if self.relate(edge).is_ok() {
2034 result.edges_created += 1;
2035 }
2036 result.linked += 1;
2037 }
2038 }
2039
2040 debug!(
2041 canonical = group.canonical,
2042 aliases = ?group.aliases,
2043 memories = group_memory_ids.len(),
2044 "entity resolution: merged group"
2045 );
2046 }
2047
2048 for sep in separations {
2050 resolver.mark_different(&sep.name_a, &sep.name_b);
2051 debug!(
2052 a = sep.name_a,
2053 b = sep.name_b,
2054 "entity resolution: confirmed different"
2055 );
2056 }
2057
2058 let cognitive_dir = self.path.join("cognitive");
2060 if cognitive_dir.exists() || std::fs::create_dir_all(&cognitive_dir).is_ok() {
2061 let _ = resolver.save(&cognitive_dir.join("entities.json"));
2062 }
2063
2064 debug!(
2065 linked = result.linked,
2066 edges = result.edges_created,
2067 groups = merge_groups.len(),
2068 separations = separations.len(),
2069 "entity link resolutions applied"
2070 );
2071 Ok(result)
2072 }
2073
2074 pub fn link_entities(&self) -> MenteResult<EntityLinkResult> {
2080 let entity_memory_map = self.build_entity_memory_map();
2081 let resolver = self.entity_resolver.read();
2082
2083 let mut canonical_groups: HashMap<String, Vec<String>> = HashMap::new();
2085 for entity_name in entity_memory_map.keys() {
2086 let resolved = resolver.resolve(entity_name);
2087 if resolved.source != mentedb_cognitive::ResolutionSource::Identity {
2088 canonical_groups
2089 .entry(resolved.canonical.clone())
2090 .or_default()
2091 .push(entity_name.clone());
2092 }
2093 }
2094
2095 drop(resolver);
2096
2097 let mut result = EntityLinkResult::default();
2098 let now = std::time::SystemTime::now()
2099 .duration_since(std::time::UNIX_EPOCH)
2100 .unwrap_or_default()
2101 .as_micros() as u64;
2102
2103 for (canonical, names) in &canonical_groups {
2104 let mut group_memory_ids: Vec<MemoryId> = Vec::new();
2106 for name in names {
2107 if let Some(ids) = entity_memory_map.get(name) {
2108 group_memory_ids.extend(ids);
2109 }
2110 }
2111 if let Some(ids) = entity_memory_map.get(canonical) {
2113 group_memory_ids.extend(ids);
2114 }
2115 group_memory_ids.sort();
2116 group_memory_ids.dedup();
2117
2118 if group_memory_ids.len() < 2 {
2119 continue;
2120 }
2121
2122 let label = format!("entity_link:{}", canonical);
2123 for i in 0..group_memory_ids.len() {
2124 for j in (i + 1)..group_memory_ids.len() {
2125 let a_id = group_memory_ids[i];
2126 let b_id = group_memory_ids[j];
2127
2128 let graph = self.graph.read_graph();
2129 let already_linked = graph.outgoing(a_id).iter().any(|(tid, e)| {
2130 *tid == b_id
2131 && e.edge_type == EdgeType::Related
2132 && e.label
2133 .as_ref()
2134 .is_some_and(|l| l.starts_with("entity_link:"))
2135 });
2136 drop(graph);
2137
2138 if already_linked {
2139 continue;
2140 }
2141
2142 let edge = MemoryEdge {
2143 source: a_id,
2144 target: b_id,
2145 edge_type: EdgeType::Related,
2146 weight: 1.0,
2147 created_at: now,
2148 valid_from: None,
2149 valid_until: None,
2150 label: Some(label.clone()),
2151 };
2152 if self.relate(edge).is_ok() {
2153 result.edges_created += 1;
2154 }
2155 result.linked += 1;
2156 }
2157 }
2158 }
2159
2160 debug!(
2161 linked = result.linked,
2162 edges = result.edges_created,
2163 groups = canonical_groups.len(),
2164 "sync entity linking complete"
2165 );
2166 Ok(result)
2167 }
2168
2169 fn build_entity_memory_map(&self) -> HashMap<String, Vec<MemoryId>> {
2171 let page_ids: Vec<PageId> = self.page_map.read().values().copied().collect();
2172 let mut map: HashMap<String, Vec<MemoryId>> = HashMap::new();
2173 for pid in &page_ids {
2174 if let Ok(mem) = self.storage.load_memory(*pid) {
2175 for tag in &mem.tags {
2176 if let Some(name) = tag.strip_prefix("entity:") {
2177 let normalized = name.to_lowercase().trim().to_string();
2178 map.entry(normalized).or_default().push(mem.id);
2180 break;
2181 }
2182 }
2183 }
2184 }
2185 map
2186 }
2187
2188 pub fn entity_memories(&self) -> Vec<MemoryNode> {
2190 let page_ids: Vec<PageId> = self.page_map.read().values().copied().collect();
2191 page_ids
2192 .iter()
2193 .filter_map(|pid| self.storage.load_memory(*pid).ok())
2194 .filter(|m| m.tags.iter().any(|t| t.starts_with("entity:")))
2195 .collect()
2196 }
2197
2198 pub fn entity_communities(&self) -> HashMap<String, Vec<(String, String)>> {
2203 let page_ids: Vec<PageId> = self.page_map.read().values().copied().collect();
2204 let mut categories: HashMap<String, Vec<(String, String)>> = HashMap::new();
2205
2206 for pid in &page_ids {
2207 if let Ok(mem) = self.storage.load_memory(*pid) {
2208 if mem.tags.iter().any(|t| t == "community_summary") {
2210 continue;
2211 }
2212
2213 let entity_name = mem
2214 .tags
2215 .iter()
2216 .find_map(|t| t.strip_prefix("entity:"))
2217 .map(|n| n.to_string());
2218
2219 if let Some(name) = entity_name {
2220 let entity_type = mem
2221 .tags
2222 .iter()
2223 .find_map(|t| t.strip_prefix("entity_type:"))
2224 .unwrap_or("general")
2225 .to_lowercase();
2226
2227 let context: String = mem.content.chars().take(200).collect();
2228 categories
2229 .entry(entity_type)
2230 .or_default()
2231 .push((name, context));
2232 }
2233 }
2234 }
2235
2236 categories.retain(|_, members| members.len() >= 2);
2238 categories
2239 }
2240
2241 pub fn store_community_summary(
2246 &self,
2247 category: &str,
2248 summary: &str,
2249 member_names: &[String],
2250 ) -> MenteResult<MemoryId> {
2251 if category.is_empty() {
2252 return Err(MenteError::Storage(
2253 "community category cannot be empty".into(),
2254 ));
2255 }
2256
2257 let now = std::time::SystemTime::now()
2258 .duration_since(std::time::UNIX_EPOCH)
2259 .unwrap_or_default()
2260 .as_micros() as u64;
2261
2262 let community_tag = format!("community:{}", category);
2264 let page_ids: Vec<PageId> = self.page_map.read().values().copied().collect();
2265 let mut existing_id = None;
2266 for pid in &page_ids {
2267 if let Ok(mem) = self.storage.load_memory(*pid)
2268 && mem.tags.iter().any(|t| t == &community_tag)
2269 {
2270 let mut updated = mem.clone();
2272 updated.content = summary.to_string();
2273 if let Some(ref embedder) = self.embedder {
2274 updated.embedding = embedder
2275 .embed(summary)
2276 .unwrap_or_else(|_| updated.embedding.clone());
2277 }
2278 self.storage.update_memory(*pid, &updated)?;
2279 existing_id = Some(updated.id);
2280 break;
2281 }
2282 }
2283
2284 let node_id = if let Some(id) = existing_id {
2285 id
2286 } else {
2287 let embedding = self
2289 .embedder
2290 .as_ref()
2291 .and_then(|e| e.embed(summary).ok())
2292 .unwrap_or_default();
2293
2294 let mut node = MemoryNode::new(
2295 mentedb_core::types::AgentId::new(),
2296 MemoryType::Semantic,
2297 summary.to_string(),
2298 embedding,
2299 );
2300 node.tags = vec![
2301 "community_summary".to_string(),
2302 community_tag,
2303 "source:enrichment".to_string(),
2304 ];
2305 node.confidence = 0.7;
2306 let id = node.id;
2307 self.store(node)?;
2308 id
2309 };
2310
2311 let entity_map = self.build_entity_memory_map();
2314 for name in member_names {
2315 let normalized = name.to_lowercase();
2316 if let Some(member_ids) = entity_map.get(&normalized) {
2317 for member_id in member_ids {
2318 self.relate(MemoryEdge {
2319 source: node_id,
2320 target: *member_id,
2321 edge_type: EdgeType::Derived,
2322 weight: 0.8,
2323 created_at: now,
2324 valid_from: None,
2325 valid_until: None,
2326 label: Some(format!("community_member:{}", category)),
2327 })?;
2328 }
2329 }
2330 }
2331
2332 Ok(node_id)
2333 }
2334
2335 pub fn community_summaries(&self) -> Vec<MemoryNode> {
2337 let page_ids: Vec<PageId> = self.page_map.read().values().copied().collect();
2338 page_ids
2339 .iter()
2340 .filter_map(|pid| self.storage.load_memory(*pid).ok())
2341 .filter(|m| m.tags.iter().any(|t| t == "community_summary"))
2342 .collect()
2343 }
2344
2345 pub fn profile_facts(&self) -> Vec<String> {
2349 let page_ids: Vec<PageId> = self.page_map.read().values().copied().collect();
2350 let mut facts = Vec::new();
2351
2352 for pid in &page_ids {
2353 if let Ok(mem) = self.storage.load_memory(*pid) {
2354 if mem.confidence < 0.5 {
2356 continue;
2357 }
2358 match mem.memory_type {
2359 MemoryType::Semantic | MemoryType::Procedural => {
2360 if mem
2362 .tags
2363 .iter()
2364 .any(|t| t == "community_summary" || t.starts_with("entity:"))
2365 {
2366 continue;
2367 }
2368 facts.push(mem.content.chars().take(300).collect());
2369 }
2370 _ => {}
2371 }
2372 }
2373 }
2374
2375 facts.truncate(100);
2377 facts
2378 }
2379
2380 pub fn store_user_profile(&self, profile: &str) -> MenteResult<MemoryId> {
2385 let page_ids: Vec<PageId> = self.page_map.read().values().copied().collect();
2387 for pid in &page_ids {
2388 if let Ok(mem) = self.storage.load_memory(*pid)
2389 && mem.tags.iter().any(|t| t == "user_profile")
2390 {
2391 let mut updated = mem.clone();
2393 updated.content = profile.to_string();
2394 if let Some(ref embedder) = self.embedder {
2395 updated.embedding = embedder
2396 .embed(profile)
2397 .unwrap_or_else(|_| updated.embedding.clone());
2398 }
2399 self.storage.update_memory(*pid, &updated)?;
2400 return Ok(updated.id);
2401 }
2402 }
2403
2404 let embedding = self
2406 .embedder
2407 .as_ref()
2408 .and_then(|e| e.embed(profile).ok())
2409 .unwrap_or_default();
2410
2411 let mut node = MemoryNode::new(
2412 mentedb_core::types::AgentId::new(),
2413 MemoryType::Semantic,
2414 profile.to_string(),
2415 embedding,
2416 );
2417 node.tags = vec![
2418 "user_profile".to_string(),
2419 "scope:always".to_string(),
2420 "source:enrichment".to_string(),
2421 ];
2422 node.confidence = 0.8;
2423 let node_id = node.id;
2424 self.store(node)?;
2425
2426 Ok(node_id)
2427 }
2428
2429 pub fn user_profile(&self) -> Option<MemoryNode> {
2431 let page_ids: Vec<PageId> = self.page_map.read().values().copied().collect();
2432 for pid in &page_ids {
2433 if let Ok(mem) = self.storage.load_memory(*pid)
2434 && mem.tags.iter().any(|t| t == "user_profile")
2435 {
2436 return Some(mem);
2437 }
2438 }
2439 None
2440 }
2441}