1pub mod node;
43pub mod query;
44pub mod snapshot;
45pub mod store;
46
47pub use node::{
48 AinlEdge, AinlMemoryNode, AinlNodeKind, AinlNodeType, EpisodicNode, MemoryCategory,
49 PersonaLayer, PersonaNode, PersonaSource, ProcedureType, ProceduralNode, RuntimeStateNode,
50 SemanticNode, Sentiment, StrengthEvent,
51};
52pub use query::{
53 count_by_topic_cluster, find_high_confidence_facts, find_patterns, find_strong_traits,
54 recall_by_procedure_type, recall_by_topic_cluster, recall_contradictions,
55 recall_delta_by_relevance, recall_episodes_by_conversation, recall_episodes_with_signal,
56 recall_flagged_episodes, recall_low_success_procedures, recall_recent, recall_strength_history,
57 walk_from, GraphQuery,
58};
59pub use snapshot::{
60 AgentGraphSnapshot, DanglingEdgeDetail, GraphValidationReport, SnapshotEdge,
61 SNAPSHOT_SCHEMA_VERSION,
62};
63pub use store::{GraphStore, SqliteGraphStore};
64
65use uuid::Uuid;
66
67pub struct GraphMemory {
71 store: SqliteGraphStore,
72}
73
74impl GraphMemory {
75 pub fn new(db_path: &std::path::Path) -> Result<Self, String> {
80 let store = SqliteGraphStore::open(db_path)?;
81 Ok(Self { store })
82 }
83
84 pub fn from_connection(conn: rusqlite::Connection) -> Result<Self, String> {
86 let store = SqliteGraphStore::from_connection(conn)?;
87 Ok(Self { store })
88 }
89
90 pub fn from_sqlite_store(store: SqliteGraphStore) -> Self {
92 Self { store }
93 }
94
95 pub fn write_episode(
105 &self,
106 tool_calls: Vec<String>,
107 delegation_to: Option<String>,
108 trace_event: Option<serde_json::Value>,
109 ) -> Result<Uuid, String> {
110 let turn_id = Uuid::new_v4();
111 let timestamp = chrono::Utc::now().timestamp();
112
113 let node =
114 AinlMemoryNode::new_episode(turn_id, timestamp, tool_calls, delegation_to, trace_event);
115
116 let node_id = node.id;
117 self.store.write_node(&node)?;
118 Ok(node_id)
119 }
120
121 pub fn write_fact(
131 &self,
132 fact: String,
133 confidence: f32,
134 source_turn_id: Uuid,
135 ) -> Result<Uuid, String> {
136 let node = AinlMemoryNode::new_fact(fact, confidence, source_turn_id);
137 let node_id = node.id;
138 self.store.write_node(&node)?;
139 Ok(node_id)
140 }
141
142 pub fn store_pattern(
151 &self,
152 pattern_name: String,
153 compiled_graph: Vec<u8>,
154 ) -> Result<Uuid, String> {
155 let node = AinlMemoryNode::new_pattern(pattern_name, compiled_graph);
156 let node_id = node.id;
157 self.store.write_node(&node)?;
158 Ok(node_id)
159 }
160
161 pub fn write_procedural(
163 &self,
164 pattern_name: &str,
165 tool_sequence: Vec<String>,
166 confidence: f32,
167 ) -> Result<Uuid, String> {
168 let node = AinlMemoryNode::new_procedural_tools(
169 pattern_name.to_string(),
170 tool_sequence,
171 confidence,
172 );
173 let node_id = node.id;
174 self.store.write_node(&node)?;
175 Ok(node_id)
176 }
177
178 pub fn write_edge(&self, source: Uuid, target: Uuid, rel: &str) -> Result<(), String> {
180 self.store.insert_graph_edge(source, target, rel)
181 }
182
183 pub fn recall_recent(&self, seconds_ago: i64) -> Result<Vec<AinlMemoryNode>, String> {
191 let since = chrono::Utc::now().timestamp() - seconds_ago;
192 self.store.query_episodes_since(since, 100)
193 }
194
195 pub fn recall_by_type(
197 &self,
198 kind: AinlNodeKind,
199 seconds_ago: i64,
200 ) -> Result<Vec<AinlMemoryNode>, String> {
201 let since = chrono::Utc::now().timestamp() - seconds_ago;
202 self.store
203 .query_nodes_by_type_since(kind.as_str(), since, 500)
204 }
205
206 pub fn write_persona(
208 &self,
209 trait_name: &str,
210 strength: f32,
211 learned_from: Vec<Uuid>,
212 ) -> Result<Uuid, String> {
213 let node = AinlMemoryNode::new_persona(trait_name.to_string(), strength, learned_from);
214 let node_id = node.id;
215 self.store.write_node(&node)?;
216 Ok(node_id)
217 }
218
219 pub fn store(&self) -> &dyn GraphStore {
221 &self.store
222 }
223
224 pub fn sqlite_store(&self) -> &SqliteGraphStore {
226 &self.store
227 }
228
229 pub fn validate_graph(&self, agent_id: &str) -> Result<GraphValidationReport, String> {
231 self.store.validate_graph(agent_id)
232 }
233
234 pub fn export_graph(&self, agent_id: &str) -> Result<AgentGraphSnapshot, String> {
236 self.store.export_graph(agent_id)
237 }
238
239 pub fn import_graph(
241 &mut self,
242 snapshot: &AgentGraphSnapshot,
243 allow_dangling_edges: bool,
244 ) -> Result<(), String> {
245 self.store.import_graph(snapshot, allow_dangling_edges)
246 }
247
248 pub fn agent_subgraph_edges(&self, agent_id: &str) -> Result<Vec<SnapshotEdge>, String> {
250 self.store.agent_subgraph_edges(agent_id)
251 }
252
253 pub fn write_node_with_edges(&mut self, node: &AinlMemoryNode) -> Result<(), String> {
255 self.store.write_node_with_edges(node)
256 }
257
258 pub fn insert_graph_edge_checked(
260 &self,
261 from_id: Uuid,
262 to_id: Uuid,
263 label: &str,
264 ) -> Result<(), String> {
265 self.store.insert_graph_edge_checked(from_id, to_id, label)
266 }
267
268 pub fn write_node(&self, node: &AinlMemoryNode) -> Result<(), String> {
270 self.store.write_node(node)
271 }
272}
273
274#[cfg(test)]
275mod tests {
276 use super::*;
277
278 #[test]
279 fn test_graph_memory_api() {
280 let temp_dir = std::env::temp_dir();
281 let db_path = temp_dir.join("ainl_lib_test.db");
282 let _ = std::fs::remove_file(&db_path);
283
284 let memory = GraphMemory::new(&db_path).expect("Failed to create memory");
285
286 let episode_id = memory
288 .write_episode(
289 vec!["file_read".to_string(), "agent_delegate".to_string()],
290 Some("agent-B".to_string()),
291 None,
292 )
293 .expect("Failed to write episode");
294
295 assert_ne!(episode_id, Uuid::nil());
296
297 let fact_id = memory
299 .write_fact(
300 "User prefers concise responses".to_string(),
301 0.85,
302 episode_id,
303 )
304 .expect("Failed to write fact");
305
306 assert_ne!(fact_id, Uuid::nil());
307
308 let recent = memory.recall_recent(60).expect("Failed to recall");
310 assert_eq!(recent.len(), 1);
311
312 if let AinlNodeType::Episode { episodic } = &recent[0].node_type {
314 assert_eq!(episodic.delegation_to, Some("agent-B".to_string()));
315 assert_eq!(episodic.tool_calls.len(), 2);
316 } else {
317 panic!("Wrong node type");
318 }
319 }
320
321 #[test]
322 fn test_store_pattern() {
323 let temp_dir = std::env::temp_dir();
324 let db_path = temp_dir.join("ainl_lib_test_pattern.db");
325 let _ = std::fs::remove_file(&db_path);
326
327 let memory = GraphMemory::new(&db_path).expect("Failed to create memory");
328
329 let pattern_id = memory
330 .store_pattern("research_workflow".to_string(), vec![1, 2, 3, 4])
331 .expect("Failed to store pattern");
332
333 assert_ne!(pattern_id, Uuid::nil());
334
335 let patterns = find_patterns(memory.store(), "research").expect("Query failed");
337 assert_eq!(patterns.len(), 1);
338 }
339}