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//! ## Graph store: query, export, validation (since 0.1.4-alpha)
36//!
37//! - **[`SqliteGraphStore`]**: SQLite backend with **`PRAGMA foreign_keys = ON`**, `FOREIGN KEY` constraints
38//!   on `ainl_graph_edges`, one-time migration for legacy DBs (see [CHANGELOG.md](../CHANGELOG.md)).
39//! - **[`GraphQuery`]**: `store.query(agent_id)` — agent-scoped SQL helpers (episodes, lineage, tags, …).
40//! - **Snapshots**: [`AgentGraphSnapshot`], [`SnapshotEdge`], [`SNAPSHOT_SCHEMA_VERSION`];
41//!   [`SqliteGraphStore::export_graph`] / [`SqliteGraphStore::import_graph`] (strict vs repair via
42//!   `allow_dangling_edges`).
43//! - **Validation**: [`GraphValidationReport`], [`DanglingEdgeDetail`]; [`SqliteGraphStore::validate_graph`]
44//!   for agent-scoped semantics beyond raw FK enforcement.
45//! - **[`GraphMemory`]** forwards the above where hosts should not reach past the facade (see impl block).
46//!
47//! ## Node Types
48//!
49//! - **Episode**: What happened during an agent turn (tool calls, delegations)
50//! - **Semantic**: Facts learned with confidence scores
51//! - **Procedural**: Reusable compiled workflow patterns
52//! - **Persona**: Agent traits learned over time
53//! - **Runtime state** (`RuntimeStateNode`, `node_type = runtime_state`): Optional persisted session
54//!   counters and persona snapshot JSON for **ainl-runtime** (see [`GraphMemory::read_runtime_state`] /
55//!   [`GraphMemory::write_runtime_state`]).
56
57pub 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
82/// High-level graph memory API - the main entry point for AINL memory.
83///
84/// Wraps a GraphStore implementation with a simplified 5-method API.
85pub struct GraphMemory {
86    store: SqliteGraphStore,
87}
88
89impl GraphMemory {
90    /// Create a new graph memory at the given database path.
91    ///
92    /// This will create the database file if it doesn't exist, and
93    /// ensure the AINL graph schema is initialized.
94    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    /// Create from an existing SQLite connection (for integration with existing memory pools)
100    pub fn from_connection(conn: rusqlite::Connection) -> Result<Self, String> {
101        let store = SqliteGraphStore::from_connection(conn)?;
102        Ok(Self { store })
103    }
104
105    /// Wrap an already-open [`SqliteGraphStore`] (for hosts that manage connections externally).
106    pub fn from_sqlite_store(store: SqliteGraphStore) -> Self {
107        Self { store }
108    }
109
110    /// Write an episode node (what happened during an agent turn).
111    ///
112    /// # Arguments
113    /// * `tool_calls` - List of tools executed during this turn
114    /// * `delegation_to` - Agent ID this turn delegated to (if any)
115    /// * `trace_event` - Optional orchestration trace event (serialized JSON)
116    ///
117    /// # Returns
118    /// The ID of the created episode node
119    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    /// Write a semantic fact (learned information with confidence).
137    ///
138    /// # Arguments
139    /// * `fact` - The fact in natural language
140    /// * `confidence` - Confidence score (0.0-1.0)
141    /// * `source_turn_id` - Turn ID that generated this fact
142    ///
143    /// # Returns
144    /// The ID of the created semantic node
145    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    /// Store a procedural pattern (compiled workflow).
158    ///
159    /// # Arguments
160    /// * `pattern_name` - Name/identifier for the pattern
161    /// * `compiled_graph` - Binary representation of the compiled graph
162    ///
163    /// # Returns
164    /// The ID of the created procedural node
165    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    /// Store a procedural pattern derived from a live tool sequence (heuristic extraction).
177    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    /// Write a graph edge between nodes (e.g. episode timeline `follows`).
194    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    /// Recall recent episodes (within the last N seconds).
199    ///
200    /// # Arguments
201    /// * `seconds_ago` - Only return episodes from the last N seconds
202    ///
203    /// # Returns
204    /// Vector of episode nodes, most recent first
205    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    /// Recall nodes of a specific kind written in the last `seconds_ago` seconds.
211    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    /// Write a persona trait node.
222    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    /// Get direct access to the underlying store for advanced queries
235    pub fn store(&self) -> &dyn GraphStore {
236        &self.store
237    }
238
239    /// SQLite backing store (for components such as `ainl-graph-extractor` that require concrete SQL access).
240    pub fn sqlite_store(&self) -> &SqliteGraphStore {
241        &self.store
242    }
243
244    /// [`SqliteGraphStore::validate_graph`] for the same backing database (checkpoint / boot gate).
245    pub fn validate_graph(&self, agent_id: &str) -> Result<GraphValidationReport, String> {
246        self.store.validate_graph(agent_id)
247    }
248
249    /// [`SqliteGraphStore::export_graph`].
250    pub fn export_graph(&self, agent_id: &str) -> Result<AgentGraphSnapshot, String> {
251        self.store.export_graph(agent_id)
252    }
253
254    /// [`SqliteGraphStore::import_graph`] — use `allow_dangling_edges: false` for normal loads; `true` only for repair.
255    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    /// [`SqliteGraphStore::agent_subgraph_edges`].
264    pub fn agent_subgraph_edges(&self, agent_id: &str) -> Result<Vec<SnapshotEdge>, String> {
265        self.store.agent_subgraph_edges(agent_id)
266    }
267
268    /// [`SqliteGraphStore::write_node_with_edges`] (transactional; fails if embedded edge targets are missing).
269    pub fn write_node_with_edges(&mut self, node: &AinlMemoryNode) -> Result<(), String> {
270        self.store.write_node_with_edges(node)
271    }
272
273    /// [`SqliteGraphStore::insert_graph_edge_checked`].
274    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    /// Read persisted [`RuntimeStateNode`] for `agent_id` (most recent row).
284    pub fn read_runtime_state(&self, agent_id: &str) -> Result<Option<RuntimeStateNode>, String> {
285        self.store.read_runtime_state(agent_id)
286    }
287
288    /// Upsert persisted [`RuntimeStateNode`] for the given agent (stable node id per `agent_id`).
289    pub fn write_runtime_state(&self, state: &RuntimeStateNode) -> Result<(), String> {
290        self.store.write_runtime_state(state)
291    }
292
293    /// Write a fully constructed node (additive API for callers that set extended metadata).
294    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        // Write an episode
312        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        // Write a fact
323        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        // Recall recent episodes
334        let recent = memory.recall_recent(60).expect("Failed to recall");
335        assert_eq!(recent.len(), 1);
336
337        // Verify the episode content
338        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        // Query it back
361        let patterns = find_patterns(memory.store(), "research").expect("Query failed");
362        assert_eq!(patterns.len(), 1);
363    }
364}