1pub mod node;
43pub mod query;
44pub mod store;
45
46pub use node::{
47 AinlEdge, AinlMemoryNode, AinlNodeKind, AinlNodeType, EpisodicNode, MemoryCategory,
48 PersonaLayer, PersonaNode, PersonaSource, ProcedureType, ProceduralNode, SemanticNode,
49 Sentiment, StrengthEvent,
50};
51pub use query::{
52 count_by_topic_cluster, find_high_confidence_facts, find_patterns, find_strong_traits,
53 recall_by_procedure_type, recall_by_topic_cluster, recall_contradictions,
54 recall_delta_by_relevance, recall_episodes_by_conversation, recall_episodes_with_signal,
55 recall_flagged_episodes, recall_low_success_procedures, recall_recent, recall_strength_history,
56 walk_from,
57};
58pub use store::{GraphStore, SqliteGraphStore};
59
60use uuid::Uuid;
61
62pub struct GraphMemory {
66 store: SqliteGraphStore,
67}
68
69impl GraphMemory {
70 pub fn new(db_path: &std::path::Path) -> Result<Self, String> {
75 let store = SqliteGraphStore::open(db_path)?;
76 Ok(Self { store })
77 }
78
79 pub fn from_connection(conn: rusqlite::Connection) -> Result<Self, String> {
81 let store = SqliteGraphStore::from_connection(conn)?;
82 Ok(Self { store })
83 }
84
85 pub fn write_episode(
95 &self,
96 tool_calls: Vec<String>,
97 delegation_to: Option<String>,
98 trace_event: Option<serde_json::Value>,
99 ) -> Result<Uuid, String> {
100 let turn_id = Uuid::new_v4();
101 let timestamp = chrono::Utc::now().timestamp();
102
103 let node =
104 AinlMemoryNode::new_episode(turn_id, timestamp, tool_calls, delegation_to, trace_event);
105
106 let node_id = node.id;
107 self.store.write_node(&node)?;
108 Ok(node_id)
109 }
110
111 pub fn write_fact(
121 &self,
122 fact: String,
123 confidence: f32,
124 source_turn_id: Uuid,
125 ) -> Result<Uuid, String> {
126 let node = AinlMemoryNode::new_fact(fact, confidence, source_turn_id);
127 let node_id = node.id;
128 self.store.write_node(&node)?;
129 Ok(node_id)
130 }
131
132 pub fn store_pattern(
141 &self,
142 pattern_name: String,
143 compiled_graph: Vec<u8>,
144 ) -> Result<Uuid, String> {
145 let node = AinlMemoryNode::new_pattern(pattern_name, compiled_graph);
146 let node_id = node.id;
147 self.store.write_node(&node)?;
148 Ok(node_id)
149 }
150
151 pub fn write_procedural(
153 &self,
154 pattern_name: &str,
155 tool_sequence: Vec<String>,
156 confidence: f32,
157 ) -> Result<Uuid, String> {
158 let node = AinlMemoryNode::new_procedural_tools(
159 pattern_name.to_string(),
160 tool_sequence,
161 confidence,
162 );
163 let node_id = node.id;
164 self.store.write_node(&node)?;
165 Ok(node_id)
166 }
167
168 pub fn write_edge(&self, source: Uuid, target: Uuid, rel: &str) -> Result<(), String> {
170 self.store.insert_graph_edge(source, target, rel)
171 }
172
173 pub fn recall_recent(&self, seconds_ago: i64) -> Result<Vec<AinlMemoryNode>, String> {
181 let since = chrono::Utc::now().timestamp() - seconds_ago;
182 self.store.query_episodes_since(since, 100)
183 }
184
185 pub fn recall_by_type(
187 &self,
188 kind: AinlNodeKind,
189 seconds_ago: i64,
190 ) -> Result<Vec<AinlMemoryNode>, String> {
191 let since = chrono::Utc::now().timestamp() - seconds_ago;
192 self.store
193 .query_nodes_by_type_since(kind.as_str(), since, 500)
194 }
195
196 pub fn write_persona(
198 &self,
199 trait_name: &str,
200 strength: f32,
201 learned_from: Vec<Uuid>,
202 ) -> Result<Uuid, String> {
203 let node = AinlMemoryNode::new_persona(trait_name.to_string(), strength, learned_from);
204 let node_id = node.id;
205 self.store.write_node(&node)?;
206 Ok(node_id)
207 }
208
209 pub fn store(&self) -> &dyn GraphStore {
211 &self.store
212 }
213
214 pub fn write_node(&self, node: &AinlMemoryNode) -> Result<(), String> {
216 self.store.write_node(node)
217 }
218}
219
220#[cfg(test)]
221mod tests {
222 use super::*;
223
224 #[test]
225 fn test_graph_memory_api() {
226 let temp_dir = std::env::temp_dir();
227 let db_path = temp_dir.join("ainl_lib_test.db");
228 let _ = std::fs::remove_file(&db_path);
229
230 let memory = GraphMemory::new(&db_path).expect("Failed to create memory");
231
232 let episode_id = memory
234 .write_episode(
235 vec!["file_read".to_string(), "agent_delegate".to_string()],
236 Some("agent-B".to_string()),
237 None,
238 )
239 .expect("Failed to write episode");
240
241 assert_ne!(episode_id, Uuid::nil());
242
243 let fact_id = memory
245 .write_fact(
246 "User prefers concise responses".to_string(),
247 0.85,
248 episode_id,
249 )
250 .expect("Failed to write fact");
251
252 assert_ne!(fact_id, Uuid::nil());
253
254 let recent = memory.recall_recent(60).expect("Failed to recall");
256 assert_eq!(recent.len(), 1);
257
258 if let AinlNodeType::Episode { episodic } = &recent[0].node_type {
260 assert_eq!(episodic.delegation_to, Some("agent-B".to_string()));
261 assert_eq!(episodic.tool_calls.len(), 2);
262 } else {
263 panic!("Wrong node type");
264 }
265 }
266
267 #[test]
268 fn test_store_pattern() {
269 let temp_dir = std::env::temp_dir();
270 let db_path = temp_dir.join("ainl_lib_test_pattern.db");
271 let _ = std::fs::remove_file(&db_path);
272
273 let memory = GraphMemory::new(&db_path).expect("Failed to create memory");
274
275 let pattern_id = memory
276 .store_pattern("research_workflow".to_string(), vec![1, 2, 3, 4])
277 .expect("Failed to store pattern");
278
279 assert_ne!(pattern_id, Uuid::nil());
280
281 let patterns = find_patterns(memory.store(), "research").expect("Query failed");
283 assert_eq!(patterns.len(), 1);
284 }
285}