1pub mod node;
58pub mod query;
59pub mod snapshot;
60pub mod store;
61
62pub use node::{
63 AinlEdge, AinlMemoryNode, AinlNodeKind, AinlNodeType, EpisodicNode, MemoryCategory,
64 PersonaLayer, PersonaNode, PersonaSource, ProceduralNode, ProcedureType, RuntimeStateNode,
65 SemanticNode, Sentiment, StrengthEvent,
66};
67pub use query::{
68 count_by_topic_cluster, find_high_confidence_facts, find_patterns, find_strong_traits,
69 recall_by_procedure_type, recall_by_topic_cluster, recall_contradictions,
70 recall_delta_by_relevance, recall_episodes_by_conversation, recall_episodes_with_signal,
71 recall_flagged_episodes, recall_low_success_procedures, recall_recent, recall_strength_history,
72 walk_from, GraphQuery,
73};
74pub use snapshot::{
75 AgentGraphSnapshot, DanglingEdgeDetail, GraphValidationReport, SnapshotEdge,
76 SNAPSHOT_SCHEMA_VERSION,
77};
78pub use store::{GraphStore, SqliteGraphStore};
79
80use uuid::Uuid;
81
82pub struct GraphMemory {
86 store: SqliteGraphStore,
87}
88
89impl GraphMemory {
90 pub fn new(db_path: &std::path::Path) -> Result<Self, String> {
95 let store = SqliteGraphStore::open(db_path)?;
96 Ok(Self { store })
97 }
98
99 pub fn from_connection(conn: rusqlite::Connection) -> Result<Self, String> {
101 let store = SqliteGraphStore::from_connection(conn)?;
102 Ok(Self { store })
103 }
104
105 pub fn from_sqlite_store(store: SqliteGraphStore) -> Self {
107 Self { store }
108 }
109
110 pub fn write_episode(
120 &self,
121 tool_calls: Vec<String>,
122 delegation_to: Option<String>,
123 trace_event: Option<serde_json::Value>,
124 ) -> Result<Uuid, String> {
125 let turn_id = Uuid::new_v4();
126 let timestamp = chrono::Utc::now().timestamp();
127
128 let node =
129 AinlMemoryNode::new_episode(turn_id, timestamp, tool_calls, delegation_to, trace_event);
130
131 let node_id = node.id;
132 self.store.write_node(&node)?;
133 Ok(node_id)
134 }
135
136 pub fn write_fact(
146 &self,
147 fact: String,
148 confidence: f32,
149 source_turn_id: Uuid,
150 ) -> Result<Uuid, String> {
151 let node = AinlMemoryNode::new_fact(fact, confidence, source_turn_id);
152 let node_id = node.id;
153 self.store.write_node(&node)?;
154 Ok(node_id)
155 }
156
157 pub fn store_pattern(
166 &self,
167 pattern_name: String,
168 compiled_graph: Vec<u8>,
169 ) -> Result<Uuid, String> {
170 let node = AinlMemoryNode::new_pattern(pattern_name, compiled_graph);
171 let node_id = node.id;
172 self.store.write_node(&node)?;
173 Ok(node_id)
174 }
175
176 pub fn write_procedural(
178 &self,
179 pattern_name: &str,
180 tool_sequence: Vec<String>,
181 confidence: f32,
182 ) -> Result<Uuid, String> {
183 let node = AinlMemoryNode::new_procedural_tools(
184 pattern_name.to_string(),
185 tool_sequence,
186 confidence,
187 );
188 let node_id = node.id;
189 self.store.write_node(&node)?;
190 Ok(node_id)
191 }
192
193 pub fn write_edge(&self, source: Uuid, target: Uuid, rel: &str) -> Result<(), String> {
195 self.store.insert_graph_edge(source, target, rel)
196 }
197
198 pub fn recall_recent(&self, seconds_ago: i64) -> Result<Vec<AinlMemoryNode>, String> {
206 let since = chrono::Utc::now().timestamp() - seconds_ago;
207 self.store.query_episodes_since(since, 100)
208 }
209
210 pub fn recall_by_type(
212 &self,
213 kind: AinlNodeKind,
214 seconds_ago: i64,
215 ) -> Result<Vec<AinlMemoryNode>, String> {
216 let since = chrono::Utc::now().timestamp() - seconds_ago;
217 self.store
218 .query_nodes_by_type_since(kind.as_str(), since, 500)
219 }
220
221 pub fn write_persona(
223 &self,
224 trait_name: &str,
225 strength: f32,
226 learned_from: Vec<Uuid>,
227 ) -> Result<Uuid, String> {
228 let node = AinlMemoryNode::new_persona(trait_name.to_string(), strength, learned_from);
229 let node_id = node.id;
230 self.store.write_node(&node)?;
231 Ok(node_id)
232 }
233
234 pub fn store(&self) -> &dyn GraphStore {
236 &self.store
237 }
238
239 pub fn sqlite_store(&self) -> &SqliteGraphStore {
241 &self.store
242 }
243
244 pub fn validate_graph(&self, agent_id: &str) -> Result<GraphValidationReport, String> {
246 self.store.validate_graph(agent_id)
247 }
248
249 pub fn export_graph(&self, agent_id: &str) -> Result<AgentGraphSnapshot, String> {
251 self.store.export_graph(agent_id)
252 }
253
254 pub fn import_graph(
256 &mut self,
257 snapshot: &AgentGraphSnapshot,
258 allow_dangling_edges: bool,
259 ) -> Result<(), String> {
260 self.store.import_graph(snapshot, allow_dangling_edges)
261 }
262
263 pub fn agent_subgraph_edges(&self, agent_id: &str) -> Result<Vec<SnapshotEdge>, String> {
265 self.store.agent_subgraph_edges(agent_id)
266 }
267
268 pub fn write_node_with_edges(&mut self, node: &AinlMemoryNode) -> Result<(), String> {
270 self.store.write_node_with_edges(node)
271 }
272
273 pub fn insert_graph_edge_checked(
275 &self,
276 from_id: Uuid,
277 to_id: Uuid,
278 label: &str,
279 ) -> Result<(), String> {
280 self.store.insert_graph_edge_checked(from_id, to_id, label)
281 }
282
283 pub fn read_runtime_state(&self, agent_id: &str) -> Result<Option<RuntimeStateNode>, String> {
285 self.store.read_runtime_state(agent_id)
286 }
287
288 pub fn write_runtime_state(&self, state: &RuntimeStateNode) -> Result<(), String> {
290 self.store.write_runtime_state(state)
291 }
292
293 pub fn write_node(&self, node: &AinlMemoryNode) -> Result<(), String> {
295 self.store.write_node(node)
296 }
297}
298
299#[cfg(test)]
300mod tests {
301 use super::*;
302
303 #[test]
304 fn test_graph_memory_api() {
305 let temp_dir = std::env::temp_dir();
306 let db_path = temp_dir.join("ainl_lib_test.db");
307 let _ = std::fs::remove_file(&db_path);
308
309 let memory = GraphMemory::new(&db_path).expect("Failed to create memory");
310
311 let episode_id = memory
313 .write_episode(
314 vec!["file_read".to_string(), "agent_delegate".to_string()],
315 Some("agent-B".to_string()),
316 None,
317 )
318 .expect("Failed to write episode");
319
320 assert_ne!(episode_id, Uuid::nil());
321
322 let fact_id = memory
324 .write_fact(
325 "User prefers concise responses".to_string(),
326 0.85,
327 episode_id,
328 )
329 .expect("Failed to write fact");
330
331 assert_ne!(fact_id, Uuid::nil());
332
333 let recent = memory.recall_recent(60).expect("Failed to recall");
335 assert_eq!(recent.len(), 1);
336
337 if let AinlNodeType::Episode { episodic } = &recent[0].node_type {
339 assert_eq!(episodic.delegation_to, Some("agent-B".to_string()));
340 assert_eq!(episodic.tool_calls.len(), 2);
341 } else {
342 panic!("Wrong node type");
343 }
344 }
345
346 #[test]
347 fn test_store_pattern() {
348 let temp_dir = std::env::temp_dir();
349 let db_path = temp_dir.join("ainl_lib_test_pattern.db");
350 let _ = std::fs::remove_file(&db_path);
351
352 let memory = GraphMemory::new(&db_path).expect("Failed to create memory");
353
354 let pattern_id = memory
355 .store_pattern("research_workflow".to_string(), vec![1, 2, 3, 4])
356 .expect("Failed to store pattern");
357
358 assert_ne!(pattern_id, Uuid::nil());
359
360 let patterns = find_patterns(memory.store(), "research").expect("Query failed");
362 assert_eq!(patterns.len(), 1);
363 }
364}