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
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        }
114    }
115
116    #[test]
117    fn test_add_memory_and_relationship() {
118        let mut mgr = GraphManager::new();
119        let a = MemoryId::new();
120        let b = MemoryId::new();
121        mgr.add_memory(a);
122        mgr.add_memory(b);
123        assert!(
124            mgr.add_relationship(&make_edge(a, b, EdgeType::Caused))
125                .is_ok()
126        );
127    }
128
129    #[test]
130    fn test_relationship_missing_node() {
131        let mut mgr = GraphManager::new();
132        let a = MemoryId::new();
133        let b = MemoryId::new();
134        mgr.add_memory(a);
135        // b not added
136        assert!(
137            mgr.add_relationship(&make_edge(a, b, EdgeType::Caused))
138                .is_err()
139        );
140    }
141
142    #[test]
143    fn test_context_subgraph() {
144        let mut mgr = GraphManager::new();
145        let a = MemoryId::new();
146        let b = MemoryId::new();
147        let c = MemoryId::new();
148        mgr.add_memory(a);
149        mgr.add_memory(b);
150        mgr.add_memory(c);
151        mgr.add_relationship(&make_edge(a, b, EdgeType::Caused))
152            .unwrap();
153        mgr.add_relationship(&make_edge(b, c, EdgeType::Related))
154            .unwrap();
155
156        let (nodes, edges) = mgr.get_context_subgraph(a, 2);
157        assert_eq!(nodes.len(), 3);
158        assert_eq!(edges.len(), 2);
159    }
160
161    #[test]
162    fn test_compact() {
163        let mut mgr = GraphManager::new();
164        let a = MemoryId::new();
165        let b = MemoryId::new();
166        mgr.add_memory(a);
167        mgr.add_memory(b);
168        mgr.add_relationship(&make_edge(a, b, EdgeType::Caused))
169            .unwrap();
170        mgr.compact();
171
172        let out = mgr.graph().outgoing(a);
173        assert_eq!(out.len(), 1);
174    }
175
176    #[test]
177    fn test_belief_propagation() {
178        let mut mgr = GraphManager::new();
179        let a = MemoryId::new();
180        let b = MemoryId::new();
181        mgr.add_memory(a);
182        mgr.add_memory(b);
183        mgr.add_relationship(&MemoryEdge {
184            source: a,
185            target: b,
186            edge_type: EdgeType::Caused,
187            weight: 1.0,
188            created_at: 1000,
189        })
190        .unwrap();
191
192        let results = mgr.propagate_belief_change(a, 0.5);
193        assert!(results.len() >= 2);
194    }
195}