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
95impl Default for GraphManager {
96    fn default() -> Self {
97        Self::new()
98    }
99}
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104    use mentedb_core::edge::EdgeType;
105
106    fn make_edge(src: MemoryId, tgt: MemoryId, etype: EdgeType) -> MemoryEdge {
107        MemoryEdge {
108            source: src,
109            target: tgt,
110            edge_type: etype,
111            weight: 0.8,
112            created_at: 1000,
113            valid_from: None,
114            valid_until: None,
115        }
116    }
117
118    #[test]
119    fn test_add_memory_and_relationship() {
120        let mut mgr = GraphManager::new();
121        let a = MemoryId::new();
122        let b = MemoryId::new();
123        mgr.add_memory(a);
124        mgr.add_memory(b);
125        assert!(
126            mgr.add_relationship(&make_edge(a, b, EdgeType::Caused))
127                .is_ok()
128        );
129    }
130
131    #[test]
132    fn test_relationship_missing_node() {
133        let mut mgr = GraphManager::new();
134        let a = MemoryId::new();
135        let b = MemoryId::new();
136        mgr.add_memory(a);
137        // b not added
138        assert!(
139            mgr.add_relationship(&make_edge(a, b, EdgeType::Caused))
140                .is_err()
141        );
142    }
143
144    #[test]
145    fn test_context_subgraph() {
146        let mut mgr = GraphManager::new();
147        let a = MemoryId::new();
148        let b = MemoryId::new();
149        let c = MemoryId::new();
150        mgr.add_memory(a);
151        mgr.add_memory(b);
152        mgr.add_memory(c);
153        mgr.add_relationship(&make_edge(a, b, EdgeType::Caused))
154            .unwrap();
155        mgr.add_relationship(&make_edge(b, c, EdgeType::Related))
156            .unwrap();
157
158        let (nodes, edges) = mgr.get_context_subgraph(a, 2);
159        assert_eq!(nodes.len(), 3);
160        assert_eq!(edges.len(), 2);
161    }
162
163    #[test]
164    fn test_compact() {
165        let mut mgr = GraphManager::new();
166        let a = MemoryId::new();
167        let b = MemoryId::new();
168        mgr.add_memory(a);
169        mgr.add_memory(b);
170        mgr.add_relationship(&make_edge(a, b, EdgeType::Caused))
171            .unwrap();
172        mgr.compact();
173
174        let out = mgr.graph().outgoing(a);
175        assert_eq!(out.len(), 1);
176    }
177
178    #[test]
179    fn test_belief_propagation() {
180        let mut mgr = GraphManager::new();
181        let a = MemoryId::new();
182        let b = MemoryId::new();
183        mgr.add_memory(a);
184        mgr.add_memory(b);
185        mgr.add_relationship(&MemoryEdge {
186            source: a,
187            target: b,
188            edge_type: EdgeType::Caused,
189            weight: 1.0,
190            created_at: 1000,
191            valid_from: None,
192            valid_until: None,
193        })
194        .unwrap();
195
196        let results = mgr.propagate_belief_change(a, 0.5);
197        assert!(results.len() >= 2);
198    }
199}
mentedb_graph/manager.rs

mentedb_graph/
manager.rs
