Skip to main content

mentedb_graph/
manager.rs

1//! High-level knowledge graph manager.
2
3use std::path::Path;
4
5use mentedb_core::edge::MemoryEdge;
6use mentedb_core::error::{MenteError, MenteResult};
7use mentedb_core::types::MemoryId;
8
9use crate::belief::propagate_update;
10use crate::contradiction::find_contradictions;
11use crate::csr::CsrGraph;
12use crate::traversal::extract_subgraph;
13
14/// Owns a `CsrGraph` and provides high-level graph operations.
15pub struct GraphManager {
16    graph: CsrGraph,
17}
18
19impl GraphManager {
20    /// Creates a new graph manager with an empty graph.
21    pub fn new() -> Self {
22        Self {
23            graph: CsrGraph::new(),
24        }
25    }
26
27    /// Save the graph to the given directory.
28    pub fn save(&self, dir: &Path) -> MenteResult<()> {
29        std::fs::create_dir_all(dir)?;
30        self.graph.save(&dir.join("graph.json"))
31    }
32
33    /// Load the graph from the given directory.
34    pub fn load(dir: &Path) -> MenteResult<Self> {
35        let graph = CsrGraph::load(&dir.join("graph.json"))?;
36        Ok(Self { graph })
37    }
38
39    /// Access the underlying graph (for traversals, etc.).
40    pub fn graph(&self) -> &CsrGraph {
41        &self.graph
42    }
43
44    /// Register a memory node in the graph.
45    pub fn add_memory(&mut self, id: MemoryId) {
46        self.graph.add_node(id);
47    }
48
49    /// Remove a memory node and all its edges.
50    pub fn remove_memory(&mut self, id: MemoryId) {
51        self.graph.remove_node(id);
52    }
53
54    /// Add a relationship (edge) between two memory nodes.
55    pub fn add_relationship(&mut self, edge: &MemoryEdge) -> MenteResult<()> {
56        if !self.graph.contains_node(edge.source) {
57            return Err(MenteError::MemoryNotFound(edge.source));
58        }
59        if !self.graph.contains_node(edge.target) {
60            return Err(MenteError::MemoryNotFound(edge.target));
61        }
62        self.graph.add_edge(edge);
63        Ok(())
64    }
65
66    /// Extract a context subgraph around a center node.
67    pub fn get_context_subgraph(
68        &self,
69        center: MemoryId,
70        depth: usize,
71    ) -> (Vec<MemoryId>, Vec<MemoryEdge>) {
72        extract_subgraph(&self.graph, center, depth)
73    }
74
75    /// Propagate a confidence change through the graph.
76    pub fn propagate_belief_change(
77        &self,
78        id: MemoryId,
79        new_confidence: f32,
80    ) -> Vec<(MemoryId, f32)> {
81        propagate_update(&self.graph, id, new_confidence)
82    }
83
84    /// Find all nodes that contradict the given node.
85    pub fn find_all_contradictions(&self, id: MemoryId) -> Vec<MemoryId> {
86        find_contradictions(&self.graph, id)
87    }
88
89    /// Merge the delta log into CSR/CSC compressed storage.
90    pub fn compact(&mut self) {
91        self.graph.compact();
92    }
93
94    /// Strengthen an edge weight (Hebbian learning: neurons that fire together wire together).
95    pub fn strengthen_edge(&mut self, source: MemoryId, target: MemoryId, delta: f32) {
96        self.graph.strengthen_edge(source, target, delta);
97    }
98}
99
100impl Default for GraphManager {
101    fn default() -> Self {
102        Self::new()
103    }
104}
105
106#[cfg(test)]
107mod tests {
108    use super::*;
109    use mentedb_core::edge::EdgeType;
110
111    fn make_edge(src: MemoryId, tgt: MemoryId, etype: EdgeType) -> MemoryEdge {
112        MemoryEdge {
113            source: src,
114            target: tgt,
115            edge_type: etype,
116            weight: 0.8,
117            created_at: 1000,
118            valid_from: None,
119            valid_until: None,
120            label: None,
121        }
122    }
123
124    #[test]
125    fn test_add_memory_and_relationship() {
126        let mut mgr = GraphManager::new();
127        let a = MemoryId::new();
128        let b = MemoryId::new();
129        mgr.add_memory(a);
130        mgr.add_memory(b);
131        assert!(
132            mgr.add_relationship(&make_edge(a, b, EdgeType::Caused))
133                .is_ok()
134        );
135    }
136
137    #[test]
138    fn test_relationship_missing_node() {
139        let mut mgr = GraphManager::new();
140        let a = MemoryId::new();
141        let b = MemoryId::new();
142        mgr.add_memory(a);
143        // b not added
144        assert!(
145            mgr.add_relationship(&make_edge(a, b, EdgeType::Caused))
146                .is_err()
147        );
148    }
149
150    #[test]
151    fn test_context_subgraph() {
152        let mut mgr = GraphManager::new();
153        let a = MemoryId::new();
154        let b = MemoryId::new();
155        let c = MemoryId::new();
156        mgr.add_memory(a);
157        mgr.add_memory(b);
158        mgr.add_memory(c);
159        mgr.add_relationship(&make_edge(a, b, EdgeType::Caused))
160            .unwrap();
161        mgr.add_relationship(&make_edge(b, c, EdgeType::Related))
162            .unwrap();
163
164        let (nodes, edges) = mgr.get_context_subgraph(a, 2);
165        assert_eq!(nodes.len(), 3);
166        assert_eq!(edges.len(), 2);
167    }
168
169    #[test]
170    fn test_compact() {
171        let mut mgr = GraphManager::new();
172        let a = MemoryId::new();
173        let b = MemoryId::new();
174        mgr.add_memory(a);
175        mgr.add_memory(b);
176        mgr.add_relationship(&make_edge(a, b, EdgeType::Caused))
177            .unwrap();
178        mgr.compact();
179
180        let out = mgr.graph().outgoing(a);
181        assert_eq!(out.len(), 1);
182    }
183
184    #[test]
185    fn test_belief_propagation() {
186        let mut mgr = GraphManager::new();
187        let a = MemoryId::new();
188        let b = MemoryId::new();
189        mgr.add_memory(a);
190        mgr.add_memory(b);
191        mgr.add_relationship(&MemoryEdge {
192            source: a,
193            target: b,
194            edge_type: EdgeType::Caused,
195            weight: 1.0,
196            created_at: 1000,
197            valid_from: None,
198            valid_until: None,
199            label: None,
200        })
201        .unwrap();
202
203        let results = mgr.propagate_belief_change(a, 0.5);
204        assert!(results.len() >= 2);
205    }
206}