1pub mod anchored_summary;
61pub mod node;
62pub mod pattern_promotion;
63pub mod query;
64pub mod snapshot;
65pub mod store;
66pub mod trajectory_table;
67mod trajectory_persist;
68
69pub use anchored_summary::{anchored_summary_id, ANCHORED_SUMMARY_TAG};
70
71pub use trajectory_persist::{
72 persist_trajectory_coarse_tools, persist_trajectory_for_episode, trajectory_env_enabled,
73};
74
75pub use node::{
76 AinlEdge, AinlMemoryNode, AinlNodeKind, AinlNodeType, EpisodicNode, FailureNode, MemoryCategory,
77 PersonaLayer, PersonaNode, PersonaSource, ProceduralNode, ProcedureType, RuntimeStateNode,
78 SemanticNode, Sentiment, StrengthEvent, TrajectoryNode,
79};
80pub use query::{
81 count_by_topic_cluster, find_high_confidence_facts, find_patterns, find_strong_traits,
82 recall_by_procedure_type, recall_by_topic_cluster, recall_contradictions,
83 recall_delta_by_relevance, recall_episodes_by_conversation, recall_episodes_with_signal,
84 recall_flagged_episodes, recall_low_success_procedures, recall_recent, recall_strength_history,
85 recall_task_scoped_episodes, walk_from, GraphQuery,
86};
87pub use snapshot::{
88 AgentGraphSnapshot, DanglingEdgeDetail, GraphValidationReport, SnapshotEdge,
89 SNAPSHOT_SCHEMA_VERSION,
90};
91pub use store::{GraphStore, GraphValidationError, SnapshotImportError, SqliteGraphStore};
92pub use trajectory_table::TrajectoryDetailRecord;
93
94use uuid::Uuid;
95
96pub struct GraphMemory {
100 store: SqliteGraphStore,
101}
102
103impl GraphMemory {
104 pub fn new(db_path: &std::path::Path) -> Result<Self, String> {
109 let store = SqliteGraphStore::open(db_path)?;
110 Ok(Self { store })
111 }
112
113 pub fn from_connection(conn: rusqlite::Connection) -> Result<Self, String> {
115 let store = SqliteGraphStore::from_connection(conn)?;
116 Ok(Self { store })
117 }
118
119 pub fn from_sqlite_store(store: SqliteGraphStore) -> Self {
121 Self { store }
122 }
123
124 pub fn write_episode(
134 &self,
135 tool_calls: Vec<String>,
136 delegation_to: Option<String>,
137 trace_event: Option<serde_json::Value>,
138 ) -> Result<Uuid, String> {
139 let turn_id = Uuid::new_v4();
140 let timestamp = chrono::Utc::now().timestamp();
141
142 let node =
143 AinlMemoryNode::new_episode(turn_id, timestamp, tool_calls, delegation_to, trace_event);
144
145 let node_id = node.id;
146 self.store.write_node(&node)?;
147 Ok(node_id)
148 }
149
150 pub fn write_fact(
160 &self,
161 fact: String,
162 confidence: f32,
163 source_turn_id: Uuid,
164 ) -> Result<Uuid, String> {
165 let node = AinlMemoryNode::new_fact(fact, confidence, source_turn_id);
166 let node_id = node.id;
167 self.store.write_node(&node)?;
168 Ok(node_id)
169 }
170
171 pub fn store_pattern(
180 &self,
181 pattern_name: String,
182 compiled_graph: Vec<u8>,
183 ) -> Result<Uuid, String> {
184 let node = AinlMemoryNode::new_pattern(pattern_name, compiled_graph);
185 let node_id = node.id;
186 self.store.write_node(&node)?;
187 Ok(node_id)
188 }
189
190 pub fn write_procedural(
196 &self,
197 pattern_name: &str,
198 tool_sequence: Vec<String>,
199 confidence: f32,
200 ) -> Result<Uuid, String> {
201 let mut node = AinlMemoryNode::new_procedural_tools(
202 pattern_name.to_string(),
203 tool_sequence,
204 confidence,
205 );
206 if let AinlNodeType::Procedural { ref mut procedural } = node.node_type {
207 procedural.pattern_observation_count =
208 procedural
209 .pattern_observation_count
210 .max(crate::pattern_promotion::DEFAULT_MIN_OBSERVATIONS);
211 let floor = crate::pattern_promotion::DEFAULT_FITNESS_FLOOR;
212 if let Some(f) = procedural.fitness {
213 procedural.fitness = Some(f.max(floor));
214 } else {
215 procedural.fitness = Some(floor);
216 }
217 procedural.prompt_eligible = true;
218 }
219 let node_id = node.id;
220 self.store.write_node(&node)?;
221 Ok(node_id)
222 }
223
224 pub fn write_edge(&self, source: Uuid, target: Uuid, rel: &str) -> Result<(), String> {
226 self.store.insert_graph_edge(source, target, rel)
227 }
228
229 pub fn recall_recent(&self, seconds_ago: i64) -> Result<Vec<AinlMemoryNode>, String> {
237 let since = chrono::Utc::now().timestamp() - seconds_ago;
238 self.store.query_episodes_since(since, 100)
239 }
240
241 pub fn recall_by_type(
243 &self,
244 kind: AinlNodeKind,
245 seconds_ago: i64,
246 ) -> Result<Vec<AinlMemoryNode>, String> {
247 let since = chrono::Utc::now().timestamp() - seconds_ago;
248 self.store
249 .query_nodes_by_type_since(kind.as_str(), since, 500)
250 }
251
252 pub fn find_procedural_by_tool_sequence(
256 &self,
257 agent_id: &str,
258 tool_sequence: &[String],
259 ) -> Result<Option<AinlMemoryNode>, String> {
260 let norm: Vec<String> = tool_sequence.iter().map(|s| s.trim().to_string()).collect();
261 if norm.is_empty() {
262 return Ok(None);
263 }
264 let nodes = self.recall_by_type(AinlNodeKind::Procedural, 60 * 60 * 24 * 365 * 5)?;
265 for n in nodes {
266 if n.agent_id != agent_id {
267 continue;
268 }
269 let AinlNodeType::Procedural { ref procedural } = n.node_type else {
270 continue;
271 };
272 if procedural.tool_sequence.len() != norm.len() {
273 continue;
274 }
275 let same = procedural
276 .tool_sequence
277 .iter()
278 .zip(norm.iter())
279 .all(|(a, b)| a.trim() == b.trim());
280 if same {
281 return Ok(Some(n));
283 }
284 }
285 Ok(None)
286 }
287
288 pub fn write_persona(
290 &self,
291 trait_name: &str,
292 strength: f32,
293 learned_from: Vec<Uuid>,
294 ) -> Result<Uuid, String> {
295 let node = AinlMemoryNode::new_persona(trait_name.to_string(), strength, learned_from);
296 let node_id = node.id;
297 self.store.write_node(&node)?;
298 Ok(node_id)
299 }
300
301 pub fn store(&self) -> &dyn GraphStore {
303 &self.store
304 }
305
306 pub fn sqlite_store(&self) -> &SqliteGraphStore {
308 &self.store
309 }
310
311 pub fn validate_graph(&self, agent_id: &str) -> Result<GraphValidationReport, String> {
313 self.store.validate_graph(agent_id)
314 }
315
316 pub fn export_graph(&self, agent_id: &str) -> Result<AgentGraphSnapshot, String> {
318 self.store.export_graph(agent_id)
319 }
320
321 pub fn import_graph(
323 &mut self,
324 snapshot: &AgentGraphSnapshot,
325 allow_dangling_edges: bool,
326 ) -> Result<(), String> {
327 self.store.import_graph(snapshot, allow_dangling_edges)
328 }
329
330 pub fn agent_subgraph_edges(&self, agent_id: &str) -> Result<Vec<SnapshotEdge>, String> {
332 self.store.agent_subgraph_edges(agent_id)
333 }
334
335 pub fn write_node_with_edges(&mut self, node: &AinlMemoryNode) -> Result<(), String> {
337 self.store.write_node_with_edges(node)
338 }
339
340 pub fn insert_graph_edge_checked(
342 &self,
343 from_id: Uuid,
344 to_id: Uuid,
345 label: &str,
346 ) -> Result<(), String> {
347 self.store.insert_graph_edge_checked(from_id, to_id, label)
348 }
349
350 pub fn read_runtime_state(&self, agent_id: &str) -> Result<Option<RuntimeStateNode>, String> {
352 self.store.read_runtime_state(agent_id)
353 }
354
355 pub fn write_runtime_state(&self, state: &RuntimeStateNode) -> Result<(), String> {
357 self.store.write_runtime_state(state)
358 }
359
360 pub fn write_node(&self, node: &AinlMemoryNode) -> Result<(), String> {
362 self.store.write_node(node)
363 }
364
365 pub fn insert_trajectory_detail(&self, row: &TrajectoryDetailRecord) -> Result<(), String> {
367 self.store.insert_trajectory_detail(row)
368 }
369
370 pub fn list_trajectories_for_agent(
372 &self,
373 agent_id: &str,
374 limit: usize,
375 since_timestamp: Option<i64>,
376 ) -> Result<Vec<TrajectoryDetailRecord>, String> {
377 self.store
378 .list_trajectories_for_agent(agent_id, limit, since_timestamp)
379 }
380
381 pub fn count_trajectory_details_before(
384 &self,
385 agent_id: &str,
386 before_recorded_at: i64,
387 ) -> Result<usize, String> {
388 self.store
389 .count_trajectory_details_before(agent_id, before_recorded_at)
390 }
391
392 pub fn prune_trajectory_details_before(
397 &self,
398 agent_id: &str,
399 before_recorded_at: i64,
400 ) -> Result<usize, String> {
401 self.store
402 .delete_trajectory_details_before(agent_id, before_recorded_at)
403 }
404
405 pub fn search_failures_for_agent(
407 &self,
408 agent_id: &str,
409 query: &str,
410 limit: usize,
411 ) -> Result<Vec<AinlMemoryNode>, String> {
412 self.store
413 .search_failures_fts_for_agent(agent_id, query, limit)
414 }
415
416 pub fn search_all_nodes_fts(
418 &self,
419 agent_id: &str,
420 query: &str,
421 project_id: Option<&str>,
422 limit: usize,
423 ) -> Result<Vec<AinlMemoryNode>, String> {
424 self.store
425 .search_all_nodes_fts_for_agent(agent_id, query, project_id, limit)
426 }
427}
428
429#[cfg(test)]
430mod tests {
431 use super::*;
432
433 #[test]
434 fn test_graph_memory_api() {
435 let temp_dir = std::env::temp_dir();
436 let db_path = temp_dir.join("ainl_lib_test.db");
437 let _ = std::fs::remove_file(&db_path);
438
439 let memory = GraphMemory::new(&db_path).expect("Failed to create memory");
440
441 let episode_id = memory
443 .write_episode(
444 vec!["file_read".to_string(), "agent_delegate".to_string()],
445 Some("agent-B".to_string()),
446 None,
447 )
448 .expect("Failed to write episode");
449
450 assert_ne!(episode_id, Uuid::nil());
451
452 let fact_id = memory
454 .write_fact(
455 "User prefers concise responses".to_string(),
456 0.85,
457 episode_id,
458 )
459 .expect("Failed to write fact");
460
461 assert_ne!(fact_id, Uuid::nil());
462
463 let recent = memory.recall_recent(60).expect("Failed to recall");
465 assert_eq!(recent.len(), 1);
466
467 if let AinlNodeType::Episode { episodic } = &recent[0].node_type {
469 assert_eq!(episodic.delegation_to, Some("agent-B".to_string()));
470 assert_eq!(episodic.tool_calls.len(), 2);
471 } else {
472 panic!("Wrong node type");
473 }
474 }
475
476 #[test]
477 fn test_store_pattern() {
478 let temp_dir = std::env::temp_dir();
479 let db_path = temp_dir.join("ainl_lib_test_pattern.db");
480 let _ = std::fs::remove_file(&db_path);
481
482 let memory = GraphMemory::new(&db_path).expect("Failed to create memory");
483
484 let pattern_id = memory
485 .store_pattern("research_workflow".to_string(), vec![1, 2, 3, 4])
486 .expect("Failed to store pattern");
487
488 assert_ne!(pattern_id, Uuid::nil());
489
490 let patterns = find_patterns(memory.store(), "research").expect("Query failed");
492 assert_eq!(patterns.len(), 1);
493 }
494
495 #[test]
497 fn failure_write_and_fts_search_roundtrip() {
498 let dir = tempfile::tempdir().expect("tempdir");
499 let db_path = dir.path().join("ainl_failure_fts_smoke.db");
500 let memory = GraphMemory::new(&db_path).expect("graph memory");
501 let agent_id = "agent-smoke-fts";
502
503 let mut node = AinlMemoryNode::new_loop_guard_failure(
504 "block",
505 Some("shell_exec"),
506 "repeated identical tool invocation blocked by loop guard",
507 Some("session-xyz"),
508 );
509 node.agent_id = agent_id.to_string();
510 let nid = node.id;
511 memory.write_node(&node).expect("write failure node");
512
513 let hits = memory
514 .search_failures_for_agent(agent_id, "loop", 10)
515 .expect("search loop");
516 assert_eq!(hits.len(), 1, "expected one FTS hit for token 'loop'");
517 assert_eq!(hits[0].id, nid);
518 assert!(
519 matches!(&hits[0].node_type, AinlNodeType::Failure { .. }),
520 "expected Failure node type"
521 );
522
523 let hits2 = memory
524 .search_failures_for_agent(agent_id, "shell_exec", 10)
525 .expect("search tool name");
526 assert_eq!(hits2.len(), 1);
527 assert_eq!(hits2[0].id, nid);
528
529 let empty = memory
530 .search_failures_for_agent(agent_id, " ", 10)
531 .expect("whitespace-only query");
532 assert!(empty.is_empty());
533
534 let wrong_agent = memory
535 .search_failures_for_agent("other-agent", "loop", 10)
536 .expect("wrong agent id");
537 assert!(wrong_agent.is_empty());
538 }
539
540 #[test]
542 fn all_nodes_fts_write_and_search_roundtrip() {
543 let dir = tempfile::tempdir().expect("tempdir");
544 let db_path = dir.path().join("ainl_all_nodes_fts.db");
545 let memory = GraphMemory::new(&db_path).expect("graph memory");
546 let agent_id = "agent-fts-all";
547 let mut node = AinlMemoryNode::new_fact("unique-fts-violet-cat-42".into(), 0.8, Uuid::new_v4());
548 node.agent_id = agent_id.to_string();
549 let nid = node.id;
550 memory.write_node(&node).expect("write fact");
551
552 let hits = memory
553 .search_all_nodes_fts(agent_id, "violet", None, 10)
554 .expect("search");
555 assert_eq!(hits.len(), 1, "expected one all-nodes FTS hit");
556 assert_eq!(hits[0].id, nid);
557 }
558
559 #[test]
560 fn tool_execution_failure_write_and_fts_search_roundtrip() {
561 let dir = tempfile::tempdir().expect("tempdir");
562 let db_path = dir.path().join("ainl_tool_failure_fts.db");
563 let memory = GraphMemory::new(&db_path).expect("graph memory");
564 let agent_id = "agent-tool-ft";
565
566 let mut node = AinlMemoryNode::new_tool_execution_failure(
567 "file_read",
568 "ENOENT: no such file or directory",
569 Some("sess-tool-1"),
570 );
571 node.agent_id = agent_id.to_string();
572 let nid = node.id;
573 memory.write_node(&node).expect("write tool failure node");
574
575 let hits = memory
576 .search_failures_for_agent(agent_id, "ENOENT", 10)
577 .expect("search ENOENT");
578 assert_eq!(hits.len(), 1);
579 assert_eq!(hits[0].id, nid);
580
581 let src_hits = memory
582 .search_failures_for_agent(agent_id, "tool_runner", 10)
583 .expect("search source");
584 assert_eq!(src_hits.len(), 1);
585 assert_eq!(src_hits[0].id, nid);
586 }
587
588 #[test]
589 fn agent_loop_precheck_failure_write_and_fts_search_roundtrip() {
590 let dir = tempfile::tempdir().expect("tempdir");
591 let db_path = dir.path().join("ainl_precheck_failure_fts.db");
592 let memory = GraphMemory::new(&db_path).expect("graph memory");
593 let agent_id = "agent-precheck-ft";
594
595 let mut node = AinlMemoryNode::new_agent_loop_precheck_failure(
596 "param_validation",
597 "file_write",
598 "missing required field: path",
599 Some("sess-pv-1"),
600 );
601 node.agent_id = agent_id.to_string();
602 let nid = node.id;
603 memory.write_node(&node).expect("write precheck failure");
604
605 let hits = memory
606 .search_failures_for_agent(agent_id, "param_validation", 10)
607 .expect("search kind");
608 assert_eq!(hits.len(), 1);
609 assert_eq!(hits[0].id, nid);
610
611 let hits2 = memory
612 .search_failures_for_agent(agent_id, "agent_loop", 10)
613 .expect("search agent_loop prefix");
614 assert_eq!(hits2.len(), 1);
615 }
616
617 #[test]
618 fn ainl_runtime_graph_validation_failure_write_and_fts_search_roundtrip() {
619 let dir = tempfile::tempdir().expect("tempdir");
620 let db_path = dir.path().join("ainl_graph_validation_failure_fts.db");
621 let memory = GraphMemory::new(&db_path).expect("graph memory");
622 let agent_id = "agent-graph-val-ft";
623
624 let mut node = AinlMemoryNode::new_ainl_runtime_graph_validation_failure(
625 "graph validation failed before turn: dangling edges …",
626 Some("sess-gv-1"),
627 );
628 node.agent_id = agent_id.to_string();
629 let nid = node.id;
630 memory.write_node(&node).expect("write graph validation failure");
631
632 let hits = memory
633 .search_failures_for_agent(agent_id, "graph_validation", 10)
634 .expect("search source label");
635 assert_eq!(hits.len(), 1);
636 assert_eq!(hits[0].id, nid);
637
638 let hits2 = memory
639 .search_failures_for_agent(agent_id, "dangling", 10)
640 .expect("search message body");
641 assert_eq!(hits2.len(), 1);
642 }
643
644 #[test]
645 fn trajectory_detail_prune_before_drops_only_old_rows() {
646 use ainl_contracts::{TrajectoryOutcome, TrajectoryStep};
647
648 let dir = tempfile::tempdir().expect("tempdir");
649 let db_path = dir.path().join("ainl_traj_prune.db");
650 let memory = GraphMemory::new(&db_path).expect("graph memory");
651 let agent = "agent-traj-prune";
652 let ep_old = memory
653 .write_episode(vec![], None, None)
654 .expect("episode for old traj");
655 let ep_new = memory
656 .write_episode(vec![], None, None)
657 .expect("episode for new traj");
658 let mk_step = |sid: &str| TrajectoryStep {
659 step_id: sid.to_string(),
660 timestamp_ms: 0,
661 adapter: "a".into(),
662 operation: "o".into(),
663 inputs_preview: None,
664 outputs_preview: None,
665 duration_ms: 1,
666 success: true,
667 error: None,
668 vitals: None,
669 freshness_at_step: None,
670 frame_vars: None,
671 tool_telemetry: None,
672 };
673 let r_old = TrajectoryDetailRecord {
674 id: Uuid::new_v4(),
675 episode_id: ep_old,
676 graph_trajectory_node_id: None,
677 agent_id: agent.to_string(),
678 session_id: "s-old".into(),
679 project_id: None,
680 recorded_at: 100,
681 outcome: TrajectoryOutcome::Success,
682 ainl_source_hash: None,
683 duration_ms: 1,
684 steps: vec![mk_step("1")],
685 frame_vars: None,
686 fitness_delta: None,
687 };
688 let r_new = TrajectoryDetailRecord {
689 id: Uuid::new_v4(),
690 episode_id: ep_new,
691 graph_trajectory_node_id: None,
692 agent_id: agent.to_string(),
693 session_id: "s-new".into(),
694 project_id: None,
695 recorded_at: 200,
696 outcome: TrajectoryOutcome::Success,
697 ainl_source_hash: None,
698 duration_ms: 1,
699 steps: vec![mk_step("2")],
700 frame_vars: None,
701 fitness_delta: None,
702 };
703 memory.insert_trajectory_detail(&r_old).expect("insert old");
704 memory.insert_trajectory_detail(&r_new).expect("insert new");
705 let before = memory
706 .list_trajectories_for_agent(agent, 10, None)
707 .expect("list");
708 assert_eq!(before.len(), 2);
709 let removed = memory
710 .prune_trajectory_details_before(agent, 200)
711 .expect("prune");
712 assert_eq!(removed, 1);
713 let after = memory
714 .list_trajectories_for_agent(agent, 10, None)
715 .expect("list after");
716 assert_eq!(after.len(), 1);
717 assert_eq!(after[0].recorded_at, 200);
718 }
719}