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 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
67/// High-level graph memory API - the main entry point for AINL memory.
68///
69/// Wraps a GraphStore implementation with a simplified 5-method API.
70pub struct GraphMemory {
71    store: SqliteGraphStore,
72}
73
74impl GraphMemory {
75    /// Create a new graph memory at the given database path.
76    ///
77    /// This will create the database file if it doesn't exist, and
78    /// ensure the AINL graph schema is initialized.
79    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    /// Create from an existing SQLite connection (for integration with existing memory pools)
85    pub fn from_connection(conn: rusqlite::Connection) -> Result<Self, String> {
86        let store = SqliteGraphStore::from_connection(conn)?;
87        Ok(Self { store })
88    }
89
90    /// Wrap an already-open [`SqliteGraphStore`] (for hosts that manage connections externally).
91    pub fn from_sqlite_store(store: SqliteGraphStore) -> Self {
92        Self { store }
93    }
94
95    /// Write an episode node (what happened during an agent turn).
96    ///
97    /// # Arguments
98    /// * `tool_calls` - List of tools executed during this turn
99    /// * `delegation_to` - Agent ID this turn delegated to (if any)
100    /// * `trace_event` - Optional orchestration trace event (serialized JSON)
101    ///
102    /// # Returns
103    /// The ID of the created episode node
104    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    /// Write a semantic fact (learned information with confidence).
122    ///
123    /// # Arguments
124    /// * `fact` - The fact in natural language
125    /// * `confidence` - Confidence score (0.0-1.0)
126    /// * `source_turn_id` - Turn ID that generated this fact
127    ///
128    /// # Returns
129    /// The ID of the created semantic node
130    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    /// Store a procedural pattern (compiled workflow).
143    ///
144    /// # Arguments
145    /// * `pattern_name` - Name/identifier for the pattern
146    /// * `compiled_graph` - Binary representation of the compiled graph
147    ///
148    /// # Returns
149    /// The ID of the created procedural node
150    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    /// Store a procedural pattern derived from a live tool sequence (heuristic extraction).
162    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    /// Write a graph edge between nodes (e.g. episode timeline `follows`).
179    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    /// Recall recent episodes (within the last N seconds).
184    ///
185    /// # Arguments
186    /// * `seconds_ago` - Only return episodes from the last N seconds
187    ///
188    /// # Returns
189    /// Vector of episode nodes, most recent first
190    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    /// Recall nodes of a specific kind written in the last `seconds_ago` seconds.
196    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    /// Write a persona trait node.
207    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    /// Get direct access to the underlying store for advanced queries
220    pub fn store(&self) -> &dyn GraphStore {
221        &self.store
222    }
223
224    /// SQLite backing store (for components such as `ainl-graph-extractor` that require concrete SQL access).
225    pub fn sqlite_store(&self) -> &SqliteGraphStore {
226        &self.store
227    }
228
229    /// [`SqliteGraphStore::validate_graph`] for the same backing database (checkpoint / boot gate).
230    pub fn validate_graph(&self, agent_id: &str) -> Result<GraphValidationReport, String> {
231        self.store.validate_graph(agent_id)
232    }
233
234    /// [`SqliteGraphStore::export_graph`].
235    pub fn export_graph(&self, agent_id: &str) -> Result<AgentGraphSnapshot, String> {
236        self.store.export_graph(agent_id)
237    }
238
239    /// [`SqliteGraphStore::import_graph`] — use `allow_dangling_edges: false` for normal loads; `true` only for repair.
240    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    /// [`SqliteGraphStore::agent_subgraph_edges`].
249    pub fn agent_subgraph_edges(&self, agent_id: &str) -> Result<Vec<SnapshotEdge>, String> {
250        self.store.agent_subgraph_edges(agent_id)
251    }
252
253    /// [`SqliteGraphStore::write_node_with_edges`] (transactional; fails if embedded edge targets are missing).
254    pub fn write_node_with_edges(&mut self, node: &AinlMemoryNode) -> Result<(), String> {
255        self.store.write_node_with_edges(node)
256    }
257
258    /// [`SqliteGraphStore::insert_graph_edge_checked`].
259    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    /// Write a fully constructed node (additive API for callers that set extended metadata).
269    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        // Write an episode
287        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        // Write a fact
298        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        // Recall recent episodes
309        let recent = memory.recall_recent(60).expect("Failed to recall");
310        assert_eq!(recent.len(), 1);
311
312        // Verify the episode content
313        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        // Query it back
336        let patterns = find_patterns(memory.store(), "research").expect("Query failed");
337        assert_eq!(patterns.len(), 1);
338    }
339}