1use 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
14pub struct GraphManager {
16 graph: CsrGraph,
17}
18
19impl GraphManager {
20 pub fn new() -> Self {
22 Self {
23 graph: CsrGraph::new(),
24 }
25 }
26
27 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 pub fn load(dir: &Path) -> MenteResult<Self> {
35 let graph = CsrGraph::load(&dir.join("graph.json"))?;
36 Ok(Self { graph })
37 }
38
39 pub fn graph(&self) -> &CsrGraph {
41 &self.graph
42 }
43
44 pub fn add_memory(&mut self, id: MemoryId) {
46 self.graph.add_node(id);
47 }
48
49 pub fn remove_memory(&mut self, id: MemoryId) {
51 self.graph.remove_node(id);
52 }
53
54 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 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 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 pub fn find_all_contradictions(&self, id: MemoryId) -> Vec<MemoryId> {
86 find_contradictions(&self.graph, id)
87 }
88
89 pub fn compact(&mut self) {
91 self.graph.compact();
92 }
93
94 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 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}