agentdb/memory/
graph.rs

1use crate::error::{AgentDbError, Result};
2use crate::schema::now_ms;
3use rusqlite::params;
4use rusqlite::Connection;
5use serde_json::Value;
6use std::sync::{Arc, Mutex};
7
8/// A typed node in the memory graph.
9#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
10pub struct Node {
11    /// Unique identifier for this node.
12    pub id: String,
13    /// Semantic type label (e.g. `"session"`, `"thought"`, `"tool"`).
14    pub kind: String,
15    /// Arbitrary JSON payload attached to the node.
16    pub data: Option<Value>,
17    /// Unix-millisecond timestamp when the node was first created.
18    pub created_at: i64,
19    /// Unix-millisecond timestamp of the most recent update.
20    pub updated_at: i64,
21}
22
23/// A directed, weighted edge between two nodes in the memory graph.
24#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
25pub struct Edge {
26    /// ID of the source node.
27    pub src: String,
28    /// ID of the destination node.
29    pub dst: String,
30    /// Semantic relationship label (e.g. `"recalled"`, `"leads_to"`).
31    pub relation: String,
32    /// Importance or strength of the relationship, conventionally in `[0.0, 1.0]`.
33    pub weight: f64,
34    /// Unix-millisecond timestamp when the edge was created or last updated.
35    pub created_at: i64,
36}
37
38/// Options controlling how the memory graph is traversed.
39#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
40pub struct TraversalOptions {
41    /// If set, only follow edges whose relation label matches this string.
42    pub relation: Option<String>,
43    /// Maximum number of hops from the anchor node.
44    pub max_depth: usize,
45    /// Discard edges whose weight is below this threshold.
46    pub min_weight: Option<f64>,
47}
48
49/// A single node returned by a graph traversal, with path metadata.
50#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
51pub struct TraversalResult {
52    /// The reached node.
53    pub node: Node,
54    /// Number of hops from the anchor node.
55    pub depth: usize,
56    /// Weight of the edge that connected this node to its parent in the traversal.
57    pub weight: f64,
58}
59
60/// In-process memory graph backed by a SQLite WAL database.
61pub struct MemoryGraph {
62    conn: Arc<Mutex<Connection>>,
63}
64
65impl MemoryGraph {
66    pub(crate) fn new(conn: Arc<Mutex<Connection>>) -> Self {
67        Self { conn }
68    }
69
70    /// Add or update a node. If a node with this `id` already exists it is overwritten
71    /// in place (kind, data, and updated_at are refreshed; created_at is preserved).
72    pub fn add_node(&self, id: &str, kind: &str, data: Option<Value>) -> Result<()> {
73        let conn = self.conn.lock().unwrap();
74        let data_str = data.as_ref().map(|d| d.to_string());
75        let now = now_ms();
76        conn.execute(
77            "INSERT INTO _adb_nodes (id, kind, data, created_at, updated_at)
78             VALUES (?1, ?2, ?3, ?4, ?5)
79             ON CONFLICT(id) DO UPDATE SET
80               kind       = excluded.kind,
81               data       = excluded.data,
82               updated_at = excluded.updated_at",
83            params![id, kind, data_str, now, now],
84        )?;
85        Ok(())
86    }
87
88    /// Retrieve a node by its ID.
89    ///
90    /// Returns [`AgentDbError::NodeNotFound`] if no node with that ID exists.
91    pub fn get_node(&self, id: &str) -> Result<Node> {
92        let conn = self.conn.lock().unwrap();
93        conn.query_row(
94            "SELECT id, kind, data, created_at, updated_at FROM _adb_nodes WHERE id = ?1",
95            params![id],
96            |row| {
97                let data_str: Option<String> = row.get(2)?;
98                Ok(Node {
99                    id: row.get(0)?,
100                    kind: row.get(1)?,
101                    data: data_str
102                        .as_deref()
103                        .and_then(|s| serde_json::from_str(s).ok()),
104                    created_at: row.get(3)?,
105                    updated_at: row.get(4)?,
106                })
107            },
108        )
109        .map_err(|_| AgentDbError::NodeNotFound(id.to_string()))
110    }
111
112    /// Remove a node by its ID. Associated edges are also deleted via ON DELETE CASCADE.
113    pub fn delete_node(&self, id: &str) -> Result<()> {
114        let conn = self.conn.lock().unwrap();
115        conn.execute("DELETE FROM _adb_nodes WHERE id = ?1", params![id])?;
116        Ok(())
117    }
118
119    /// Add or update a directed edge `src → dst` with the given `relation` label and `weight`.
120    ///
121    /// Both `src` and `dst` must already exist; returns [`AgentDbError::NodeNotFound`] otherwise.
122    /// If an edge with the same `(src, dst, relation)` triple already exists its weight is updated.
123    pub fn add_edge(&self, src: &str, dst: &str, relation: &str, weight: f64) -> Result<()> {
124        self.get_node(src)?;
125        self.get_node(dst)?;
126        let conn = self.conn.lock().unwrap();
127        conn.execute(
128            "INSERT INTO _adb_edges (src, dst, relation, weight, created_at)
129             VALUES (?1, ?2, ?3, ?4, ?5)
130             ON CONFLICT(src, dst, relation) DO UPDATE SET
131               weight     = excluded.weight,
132               created_at = excluded.created_at",
133            params![src, dst, relation, weight, now_ms()],
134        )?;
135        Ok(())
136    }
137
138    /// Remove the edge `src → dst` with the given `relation`.
139    ///
140    /// Returns [`AgentDbError::EdgeNotFound`] if no such edge exists.
141    pub fn delete_edge(&self, src: &str, dst: &str, relation: &str) -> Result<()> {
142        let conn = self.conn.lock().unwrap();
143        let changed = conn.execute(
144            "DELETE FROM _adb_edges WHERE src = ?1 AND dst = ?2 AND relation = ?3",
145            params![src, dst, relation],
146        )?;
147        if changed == 0 {
148            return Err(AgentDbError::EdgeNotFound {
149                src: src.to_string(),
150                dst: dst.to_string(),
151            });
152        }
153        Ok(())
154    }
155
156    /// Traverse the graph from `node_id` and return all reachable nodes within the
157    /// constraints defined by `opts`.
158    ///
159    /// Uses a recursive Common Table Expression (CTE) for efficient SQLite-side
160    /// traversal. Results are ordered by ascending depth, then descending edge weight.
161    pub fn neighbors(&self, node_id: &str, opts: TraversalOptions) -> Result<Vec<TraversalResult>> {
162        let conn = self.conn.lock().unwrap();
163        let max_depth = opts.max_depth.max(1) as i64;
164        let min_weight = opts.min_weight.unwrap_or(0.0);
165
166        let results = if let Some(ref relation) = opts.relation {
167            let sql = "
168                WITH RECURSIVE traverse(node_id, depth, weight, visited) AS (
169                    SELECT dst, 1, weight, ',' || ?1 || ',' || dst || ','
170                    FROM _adb_edges
171                    WHERE src = ?1 AND relation = ?2 AND weight >= ?4
172                    UNION ALL
173                    SELECT e.dst, t.depth + 1, e.weight,
174                           t.visited || e.dst || ','
175                    FROM _adb_edges e
176                    JOIN traverse t ON e.src = t.node_id
177                    WHERE t.depth < ?3
178                      AND e.relation = ?2
179                      AND e.weight >= ?4
180                      AND INSTR(t.visited, ',' || e.dst || ',') = 0
181                )
182                SELECT n.id, n.kind, n.data, n.created_at, n.updated_at,
183                       MIN(t.depth) AS depth, MAX(t.weight) AS weight
184                FROM traverse t
185                JOIN _adb_nodes n ON n.id = t.node_id
186                GROUP BY n.id
187                ORDER BY depth ASC, weight DESC
188            ";
189            let mut stmt = conn.prepare(sql)?;
190            let rows =
191                stmt.query_map(params![node_id, relation, max_depth, min_weight], parse_row)?;
192            rows.map(|r| r.map_err(AgentDbError::Sqlite))
193                .collect::<Result<Vec<_>>>()?
194        } else {
195            let sql = "
196                WITH RECURSIVE traverse(node_id, depth, weight, visited) AS (
197                    SELECT dst, 1, weight, ',' || ?1 || ',' || dst || ','
198                    FROM _adb_edges
199                    WHERE src = ?1 AND weight >= ?3
200                    UNION ALL
201                    SELECT e.dst, t.depth + 1, e.weight,
202                           t.visited || e.dst || ','
203                    FROM _adb_edges e
204                    JOIN traverse t ON e.src = t.node_id
205                    WHERE t.depth < ?2
206                      AND e.weight >= ?3
207                      AND INSTR(t.visited, ',' || e.dst || ',') = 0
208                )
209                SELECT n.id, n.kind, n.data, n.created_at, n.updated_at,
210                       MIN(t.depth) AS depth, MAX(t.weight) AS weight
211                FROM traverse t
212                JOIN _adb_nodes n ON n.id = t.node_id
213                GROUP BY n.id
214                ORDER BY depth ASC, weight DESC
215            ";
216            let mut stmt = conn.prepare(sql)?;
217            let rows = stmt.query_map(params![node_id, max_depth, min_weight], parse_row)?;
218            rows.map(|r| r.map_err(AgentDbError::Sqlite))
219                .collect::<Result<Vec<_>>>()?
220        };
221
222        Ok(results)
223    }
224
225    /// Return all nodes whose `kind` field matches the given string.
226    pub fn nodes_by_kind(&self, kind: &str) -> Result<Vec<Node>> {
227        let conn = self.conn.lock().unwrap();
228        let mut stmt = conn.prepare(
229            "SELECT id, kind, data, created_at, updated_at FROM _adb_nodes WHERE kind = ?1",
230        )?;
231        let rows = stmt.query_map(params![kind], |row| {
232            let data_str: Option<String> = row.get(2)?;
233            Ok(Node {
234                id: row.get(0)?,
235                kind: row.get(1)?,
236                data: data_str
237                    .as_deref()
238                    .and_then(|s| serde_json::from_str(s).ok()),
239                created_at: row.get(3)?,
240                updated_at: row.get(4)?,
241            })
242        })?;
243        rows.map(|r| r.map_err(AgentDbError::Sqlite)).collect()
244    }
245
246    /// Return the total node and edge counts as `(nodes, edges)`.
247    pub fn stats(&self) -> Result<(i64, i64)> {
248        let conn = self.conn.lock().unwrap();
249        let nodes: i64 = conn.query_row("SELECT COUNT(*) FROM _adb_nodes", [], |r| r.get(0))?;
250        let edges: i64 = conn.query_row("SELECT COUNT(*) FROM _adb_edges", [], |r| r.get(0))?;
251        Ok((nodes, edges))
252    }
253}
254
255fn parse_row(row: &rusqlite::Row) -> rusqlite::Result<TraversalResult> {
256    let data_str: Option<String> = row.get(2)?;
257    Ok(TraversalResult {
258        node: Node {
259            id: row.get(0)?,
260            kind: row.get(1)?,
261            data: data_str
262                .as_deref()
263                .and_then(|s| serde_json::from_str(s).ok()),
264            created_at: row.get(3)?,
265            updated_at: row.get(4)?,
266        },
267        depth: row.get::<_, i64>(5)? as usize,
268        weight: row.get(6)?,
269    })
270}
agentdb/memory/graph.rs

agentdb/memory/
graph.rs
