Skip to main content

ainl_memory/
lib.rs

1//! AINL Memory - Graph-based agent memory substrate
2//!
3//! **Graph-as-memory for AI agents. Execution IS the memory.**
4//!
5//! AINL Memory implements agent memory as an execution graph. Every agent turn,
6//! tool call, and delegation becomes a typed graph node. No separate retrieval
7//! layer—the graph itself is the memory.
8//!
9//! # Quick Start
10//!
11//! ```no_run
12//! use ainl_memory::GraphMemory;
13//! use std::path::Path;
14//!
15//! let memory = GraphMemory::new(Path::new("memory.db")).unwrap();
16//!
17//! // Record an episode
18//! memory.write_episode(
19//!     vec!["file_read".to_string(), "agent_delegate".to_string()],
20//!     Some("agent-B".to_string()),
21//!     None,
22//! ).unwrap();
23//!
24//! // Recall recent episodes
25//! let recent = memory.recall_recent(100).unwrap();
26//! ```
27//!
28//! # Architecture
29//!
30//! AINL Memory is designed as infrastructure that any agent framework can adopt:
31//! - Zero dependencies on specific agent runtimes
32//! - Simple trait-based API via `GraphStore`
33//! - Bring your own storage backend
34//!
35//! ## Node Types
36//!
37//! - **Episode**: What happened during an agent turn (tool calls, delegations)
38//! - **Semantic**: Facts learned with confidence scores
39//! - **Procedural**: Reusable compiled workflow patterns
40//! - **Persona**: Agent traits learned over time
41
42pub 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
62/// High-level graph memory API - the main entry point for AINL memory.
63///
64/// Wraps a GraphStore implementation with a simplified 5-method API.
65pub struct GraphMemory {
66    store: SqliteGraphStore,
67}
68
69impl GraphMemory {
70    /// Create a new graph memory at the given database path.
71    ///
72    /// This will create the database file if it doesn't exist, and
73    /// ensure the AINL graph schema is initialized.
74    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    /// Create from an existing SQLite connection (for integration with existing memory pools)
80    pub fn from_connection(conn: rusqlite::Connection) -> Result<Self, String> {
81        let store = SqliteGraphStore::from_connection(conn)?;
82        Ok(Self { store })
83    }
84
85    /// Write an episode node (what happened during an agent turn).
86    ///
87    /// # Arguments
88    /// * `tool_calls` - List of tools executed during this turn
89    /// * `delegation_to` - Agent ID this turn delegated to (if any)
90    /// * `trace_event` - Optional orchestration trace event (serialized JSON)
91    ///
92    /// # Returns
93    /// The ID of the created episode node
94    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    /// Write a semantic fact (learned information with confidence).
112    ///
113    /// # Arguments
114    /// * `fact` - The fact in natural language
115    /// * `confidence` - Confidence score (0.0-1.0)
116    /// * `source_turn_id` - Turn ID that generated this fact
117    ///
118    /// # Returns
119    /// The ID of the created semantic node
120    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    /// Store a procedural pattern (compiled workflow).
133    ///
134    /// # Arguments
135    /// * `pattern_name` - Name/identifier for the pattern
136    /// * `compiled_graph` - Binary representation of the compiled graph
137    ///
138    /// # Returns
139    /// The ID of the created procedural node
140    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    /// Store a procedural pattern derived from a live tool sequence (heuristic extraction).
152    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    /// Write a graph edge between nodes (e.g. episode timeline `follows`).
169    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    /// Recall recent episodes (within the last N seconds).
174    ///
175    /// # Arguments
176    /// * `seconds_ago` - Only return episodes from the last N seconds
177    ///
178    /// # Returns
179    /// Vector of episode nodes, most recent first
180    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    /// Recall nodes of a specific kind written in the last `seconds_ago` seconds.
186    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    /// Write a persona trait node.
197    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    /// Get direct access to the underlying store for advanced queries
210    pub fn store(&self) -> &dyn GraphStore {
211        &self.store
212    }
213
214    /// Write a fully constructed node (additive API for callers that set extended metadata).
215    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        // Write an episode
233        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        // Write a fact
244        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        // Recall recent episodes
255        let recent = memory.recall_recent(60).expect("Failed to recall");
256        assert_eq!(recent.len(), 1);
257
258        // Verify the episode content
259        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        // Query it back
282        let patterns = find_patterns(memory.store(), "research").expect("Query failed");
283        assert_eq!(patterns.len(), 1);
284    }
285}