ainl_memory/
lib.rs

1//! AINL graph-memory substrate - Spike implementation
2//!
3//! Proof-of-concept: Agent memory as execution graph nodes.
4//! Every delegation, tool call, and agent turn becomes a typed graph node.
5//! This spike proves the concept before extraction to `ainl-memory` crate.
6
7use serde::{Deserialize, Serialize};
8use uuid::Uuid;
9
10/// Core AINL node types - the vocabulary of agent memory
11#[derive(Serialize, Deserialize, Debug, Clone)]
12#[serde(tag = "type", rename_all = "snake_case")]
13pub enum AinlNodeType {
14    /// Episodic memory: what happened during an agent turn
15    Episode {
16        turn_id: Uuid,
17        agent_id: String,
18        tool_calls: Vec<String>,
19        delegation_to: Option<String>,
20        trace_id: Option<String>,
21        depth: u32,
22    },
23    /// Semantic memory: facts learned, with confidence
24    Semantic {
25        fact: String,
26        confidence: f32,
27        source_turn: Uuid,
28    },
29    /// Procedural memory: reusable compiled workflow patterns
30    Procedural {
31        pattern_name: String,
32        compiled_graph: Vec<u8>,
33    },
34    /// Persona memory: traits learned over time
35    Persona {
36        trait_name: String,
37        strength: f32,
38        learned_from: Vec<Uuid>,
39    },
40}
41
42/// A node in the AINL memory graph
43#[derive(Serialize, Deserialize, Debug, Clone)]
44pub struct AinlMemoryNode {
45    pub id: Uuid,
46    pub node_type: AinlNodeType,
47    pub timestamp: i64,
48    pub edges: Vec<AinlEdge>,
49}
50
51/// Typed edge connecting memory nodes
52#[derive(Serialize, Deserialize, Debug, Clone)]
53pub struct AinlEdge {
54    pub target_id: Uuid,
55    pub label: String,
56}
57
58impl AinlMemoryNode {
59    /// Create a new episode node for an orchestration delegation
60    pub fn new_delegation_episode(
61        agent_id: String,
62        delegated_to: String,
63        trace_id: String,
64        depth: u32,
65    ) -> Self {
66        let turn_id = Uuid::new_v4();
67        Self {
68            id: Uuid::new_v4(),
69            node_type: AinlNodeType::Episode {
70                turn_id,
71                agent_id,
72                tool_calls: vec!["agent_delegate".to_string()],
73                delegation_to: Some(delegated_to),
74                trace_id: Some(trace_id),
75                depth,
76            },
77            timestamp: chrono::Utc::now().timestamp(),
78            edges: Vec::new(),
79        }
80    }
81
82    /// Create a new episode node for a tool use
83    pub fn new_tool_episode(agent_id: String, tool_name: String) -> Self {
84        let turn_id = Uuid::new_v4();
85        Self {
86            id: Uuid::new_v4(),
87            node_type: AinlNodeType::Episode {
88                turn_id,
89                agent_id,
90                tool_calls: vec![tool_name],
91                delegation_to: None,
92                trace_id: None,
93                depth: 0,
94            },
95            timestamp: chrono::Utc::now().timestamp(),
96            edges: Vec::new(),
97        }
98    }
99
100    /// Add an edge to another node
101    pub fn add_edge(&mut self, target_id: Uuid, label: impl Into<String>) {
102        self.edges.push(AinlEdge {
103            target_id,
104            label: label.into(),
105        });
106    }
107}
108
109/// Graph memory storage - trait for swappable backends
110/// (Note: Send + Sync removed for spike - will use Arc<Mutex<>> in production)
111pub trait GraphStore {
112    /// Write a node to storage
113    fn write_node(&self, node: &AinlMemoryNode) -> Result<(), String>;
114
115    /// Read a node by ID
116    fn read_node(&self, id: Uuid) -> Result<Option<AinlMemoryNode>, String>;
117
118    /// Query nodes by type
119    fn query_by_type(&self, type_name: &str) -> Result<Vec<AinlMemoryNode>, String>;
120
121    /// Query recent episodes for an agent
122    fn query_recent_episodes(
123        &self,
124        agent_id: &str,
125        limit: usize,
126    ) -> Result<Vec<AinlMemoryNode>, String>;
127
128    /// Walk the graph from a starting node
129    fn walk_edges(&self, from_id: Uuid, label: &str) -> Result<Vec<AinlMemoryNode>, String>;
130}
131
132/// SQLite implementation of GraphStore (spike - uses existing connection)
133pub struct SqliteGraphStore {
134    db: rusqlite::Connection,
135}
136
137impl SqliteGraphStore {
138    /// Create tables if they don't exist
139    pub fn ensure_schema(db: &rusqlite::Connection) -> Result<(), rusqlite::Error> {
140        db.execute(
141            "CREATE TABLE IF NOT EXISTS ainl_graph_nodes (
142                id TEXT PRIMARY KEY,
143                node_type TEXT NOT NULL,
144                payload TEXT NOT NULL,
145                timestamp INTEGER NOT NULL
146            )",
147            [],
148        )?;
149
150        db.execute(
151            "CREATE INDEX IF NOT EXISTS idx_ainl_nodes_timestamp
152             ON ainl_graph_nodes(timestamp DESC)",
153            [],
154        )?;
155
156        db.execute(
157            "CREATE INDEX IF NOT EXISTS idx_ainl_nodes_type
158             ON ainl_graph_nodes(node_type)",
159            [],
160        )?;
161
162        db.execute(
163            "CREATE TABLE IF NOT EXISTS ainl_graph_edges (
164                from_id TEXT NOT NULL,
165                to_id TEXT NOT NULL,
166                label TEXT NOT NULL,
167                PRIMARY KEY (from_id, to_id, label)
168            )",
169            [],
170        )?;
171
172        db.execute(
173            "CREATE INDEX IF NOT EXISTS idx_ainl_edges_from
174             ON ainl_graph_edges(from_id, label)",
175            [],
176        )?;
177
178        Ok(())
179    }
180
181    /// Open/create a graph store at the given path
182    pub fn open(path: &std::path::Path) -> Result<Self, String> {
183        let db = rusqlite::Connection::open(path).map_err(|e| e.to_string())?;
184        Self::ensure_schema(&db).map_err(|e| e.to_string())?;
185        Ok(Self { db })
186    }
187}
188
189impl GraphStore for SqliteGraphStore {
190    fn write_node(&self, node: &AinlMemoryNode) -> Result<(), String> {
191        let payload = serde_json::to_string(node).map_err(|e| e.to_string())?;
192        let type_name = match &node.node_type {
193            AinlNodeType::Episode { .. } => "episode",
194            AinlNodeType::Semantic { .. } => "semantic",
195            AinlNodeType::Procedural { .. } => "procedural",
196            AinlNodeType::Persona { .. } => "persona",
197        };
198
199        self.db
200            .execute(
201                "INSERT OR REPLACE INTO ainl_graph_nodes (id, node_type, payload, timestamp)
202                 VALUES (?1, ?2, ?3, ?4)",
203                rusqlite::params![
204                    node.id.to_string(),
205                    type_name,
206                    payload,
207                    node.timestamp,
208                ],
209            )
210            .map_err(|e| e.to_string())?;
211
212        // Write edges
213        for edge in &node.edges {
214            self.db
215                .execute(
216                    "INSERT OR REPLACE INTO ainl_graph_edges (from_id, to_id, label)
217                     VALUES (?1, ?2, ?3)",
218                    rusqlite::params![
219                        node.id.to_string(),
220                        edge.target_id.to_string(),
221                        edge.label,
222                    ],
223                )
224                .map_err(|e| e.to_string())?;
225        }
226
227        Ok(())
228    }
229
230    fn read_node(&self, id: Uuid) -> Result<Option<AinlMemoryNode>, String> {
231        use rusqlite::OptionalExtension;
232
233        let payload: Option<String> = self
234            .db
235            .query_row(
236                "SELECT payload FROM ainl_graph_nodes WHERE id = ?1",
237                [id.to_string()],
238                |row| row.get::<_, String>(0),
239            )
240            .optional()
241            .map_err(|e: rusqlite::Error| e.to_string())?;
242
243        match payload {
244            Some(p) => {
245                let node: AinlMemoryNode =
246                    serde_json::from_str(&p).map_err(|e| e.to_string())?;
247                Ok(Some(node))
248            }
249            None => Ok(None),
250        }
251    }
252
253    fn query_by_type(&self, type_name: &str) -> Result<Vec<AinlMemoryNode>, String> {
254        let mut stmt = self
255            .db
256            .prepare("SELECT payload FROM ainl_graph_nodes WHERE node_type = ?1 ORDER BY timestamp DESC")
257            .map_err(|e| e.to_string())?;
258
259        let rows = stmt
260            .query_map([type_name], |row| {
261                let payload: String = row.get(0)?;
262                Ok(payload)
263            })
264            .map_err(|e| e.to_string())?;
265
266        let mut nodes = Vec::new();
267        for row in rows {
268            let payload = row.map_err(|e| e.to_string())?;
269            let node: AinlMemoryNode =
270                serde_json::from_str(&payload).map_err(|e| e.to_string())?;
271            nodes.push(node);
272        }
273
274        Ok(nodes)
275    }
276
277    fn query_recent_episodes(
278        &self,
279        agent_id: &str,
280        limit: usize,
281    ) -> Result<Vec<AinlMemoryNode>, String> {
282        let mut stmt = self
283            .db
284            .prepare(
285                "SELECT payload FROM ainl_graph_nodes
286                 WHERE node_type = 'episode'
287                 ORDER BY timestamp DESC
288                 LIMIT ?1",
289            )
290            .map_err(|e| e.to_string())?;
291
292        let rows = stmt
293            .query_map([limit], |row| {
294                let payload: String = row.get(0)?;
295                Ok(payload)
296            })
297            .map_err(|e| e.to_string())?;
298
299        let mut nodes = Vec::new();
300        for row in rows {
301            let payload = row.map_err(|e| e.to_string())?;
302            let node: AinlMemoryNode =
303                serde_json::from_str(&payload).map_err(|e| e.to_string())?;
304
305            // Filter by agent_id in the Episode variant
306            if let AinlNodeType::Episode {
307                agent_id: node_agent,
308                ..
309            } = &node.node_type
310            {
311                if node_agent == agent_id {
312                    nodes.push(node);
313                }
314            }
315        }
316
317        Ok(nodes)
318    }
319
320    fn walk_edges(&self, from_id: Uuid, label: &str) -> Result<Vec<AinlMemoryNode>, String> {
321        let mut stmt = self
322            .db
323            .prepare(
324                "SELECT to_id FROM ainl_graph_edges
325                 WHERE from_id = ?1 AND label = ?2",
326            )
327            .map_err(|e| e.to_string())?;
328
329        let target_ids: Vec<String> = stmt
330            .query_map([from_id.to_string(), label.to_string()], |row| row.get(0))
331            .map_err(|e| e.to_string())?
332            .collect::<Result<Vec<_>, _>>()
333            .map_err(|e| e.to_string())?;
334
335        let mut nodes = Vec::new();
336        for target_id in target_ids {
337            let id = Uuid::parse_str(&target_id).map_err(|e| e.to_string())?;
338            if let Some(node) = self.read_node(id)? {
339                nodes.push(node);
340            }
341        }
342
343        Ok(nodes)
344    }
345}
346
347#[cfg(test)]
348mod tests {
349    use super::*;
350
351    #[test]
352    fn test_create_delegation_episode() {
353        let node = AinlMemoryNode::new_delegation_episode(
354            "agent-123".to_string(),
355            "agent-456".to_string(),
356            "trace-xyz".to_string(),
357            1,
358        );
359
360        assert!(matches!(node.node_type, AinlNodeType::Episode { .. }));
361        if let AinlNodeType::Episode {
362            delegation_to,
363            depth,
364            ..
365        } = node.node_type
366        {
367            assert_eq!(delegation_to, Some("agent-456".to_string()));
368            assert_eq!(depth, 1);
369        }
370    }
371
372    #[test]
373    fn test_sqlite_store_write_read() {
374        let temp_dir = std::env::temp_dir();
375        let db_path = temp_dir.join("ainl_test.db");
376        let _ = std::fs::remove_file(&db_path); // Clean up from previous run
377
378        let store = SqliteGraphStore::open(&db_path).expect("Failed to open store");
379
380        let node = AinlMemoryNode::new_delegation_episode(
381            "agent-123".to_string(),
382            "agent-456".to_string(),
383            "trace-xyz".to_string(),
384            1,
385        );
386
387        store.write_node(&node).expect("Failed to write node");
388
389        let retrieved = store
390            .read_node(node.id)
391            .expect("Failed to read node")
392            .expect("Node not found");
393
394        assert_eq!(retrieved.id, node.id);
395        assert_eq!(retrieved.timestamp, node.timestamp);
396    }
397
398    #[test]
399    fn test_query_recent_episodes() {
400        let temp_dir = std::env::temp_dir();
401        let db_path = temp_dir.join("ainl_test_query.db");
402        let _ = std::fs::remove_file(&db_path);
403
404        let store = SqliteGraphStore::open(&db_path).expect("Failed to open store");
405
406        // Write 3 episodes for the same agent
407        for i in 0..3 {
408            let node = AinlMemoryNode::new_delegation_episode(
409                "agent-123".to_string(),
410                format!("agent-{}", i),
411                format!("trace-{}", i),
412                i as u32,
413            );
414            store.write_node(&node).expect("Failed to write node");
415            std::thread::sleep(std::time::Duration::from_millis(10)); // Ensure different timestamps
416        }
417
418        let episodes = store
419            .query_recent_episodes("agent-123", 10)
420            .expect("Failed to query");
421
422        assert_eq!(episodes.len(), 3);
423    }
424}
ainl_memory/lib.rs

ainl_memory/
lib.rs
