llm_kernel/graph/
async_graph.rs

1//! Async wrappers for the knowledge graph.
2//!
3//! All operations run on the Tokio blocking thread pool via
4//! [`tokio::task::spawn_blocking`], keeping the async executor free.
5//!
6//! # Usage
7//!
8//! ```no_run
9//! use std::sync::{Arc, Mutex};
10//! use rusqlite::Connection;
11//! use llm_kernel::graph::{init_graph_schema, GraphNode};
12//! use llm_kernel::graph::async_graph::AsyncGraph;
13//!
14//! # async fn example() {
15//! let conn = Connection::open_in_memory().unwrap();
16//! init_graph_schema(&conn).unwrap();
17//! let graph = AsyncGraph::new(conn);
18//!
19//! graph.upsert_node(GraphNode {
20//!     id: "n1".into(),
21//!     node_type: "concept".into(),
22//!     title: "Example".into(),
23//!     body: "body text".into(),
24//!     tags: vec![],
25//!     projects: vec![],
26//!     agents: vec![],
27//!     created: "2026-01-01T00:00:00Z".into(),
28//!     updated: "2026-01-01T00:00:00Z".into(),
29//!     importance: 0.5,
30//!     access_count: 0,
31//!     accessed_at: String::new(),
32//! }).await.unwrap();
33//! # }
34//! ```
35
36use std::sync::{Arc, Mutex};
37
38use rusqlite::Connection;
39use tokio::task;
40
41use crate::error::{KernelError, Result};
42use crate::graph::types::{EdgeDirection, GraphEdge, GraphNode, GraphStats, ScoredNode};
43
44/// Async handle to a knowledge graph backed by a `rusqlite::Connection`.
45///
46/// Internally wraps the connection in `Arc<Mutex<_>>` so it can be shared
47/// across async tasks. All blocking SQL calls are offloaded to the Tokio
48/// blocking-thread pool.
49#[derive(Clone)]
50pub struct AsyncGraph {
51    conn: Arc<Mutex<Connection>>,
52}
53
54impl AsyncGraph {
55    /// Wrap an already-initialised connection.
56    pub fn new(conn: Connection) -> Self {
57        Self {
58            conn: Arc::new(Mutex::new(conn)),
59        }
60    }
61
62    /// Open (or create) a database at `path` and initialise the graph schema.
63    pub async fn open(path: impl Into<String>) -> Result<Self> {
64        let path = path.into();
65        task::spawn_blocking(move || {
66            let conn = Connection::open(&path).map_err(|e| KernelError::Store(e.to_string()))?;
67            crate::graph::schema::init_graph_schema(&conn)?;
68            Ok(Self::new(conn))
69        })
70        .await
71        .map_err(|e| KernelError::Store(e.to_string()))?
72    }
73
74    fn with_conn<F, T>(&self, f: F) -> task::JoinHandle<Result<T>>
75    where
76        F: FnOnce(&Connection) -> Result<T> + Send + 'static,
77        T: Send + 'static,
78    {
79        let conn = Arc::clone(&self.conn);
80        task::spawn_blocking(move || {
81            let guard = conn
82                .lock()
83                .map_err(|_| KernelError::Store("mutex poisoned".into()))?;
84            f(&guard)
85        })
86    }
87
88    /// Insert or replace a node.
89    pub async fn upsert_node(&self, node: GraphNode) -> Result<()> {
90        self.with_conn(move |c| crate::graph::store::upsert_node(c, &node))
91            .await
92            .map_err(|e| KernelError::Store(e.to_string()))?
93    }
94
95    /// Read a node by ID. Returns `None` if not found.
96    pub async fn read_node(&self, id: impl Into<String>) -> Result<Option<GraphNode>> {
97        let id = id.into();
98        self.with_conn(move |c| crate::graph::store::read_node(c, &id))
99            .await
100            .map_err(|e| KernelError::Store(e.to_string()))?
101    }
102
103    /// Delete a node by ID. Returns `true` if a row was deleted.
104    pub async fn delete_node(&self, id: impl Into<String>) -> Result<bool> {
105        let id = id.into();
106        self.with_conn(move |c| crate::graph::store::delete_node(c, &id))
107            .await
108            .map_err(|e| KernelError::Store(e.to_string()))?
109    }
110
111    /// Append an edge (duplicates by ID are ignored).
112    pub async fn append_edge(&self, edge: GraphEdge) -> Result<()> {
113        self.with_conn(move |c| crate::graph::store::append_edge(c, &edge))
114            .await
115            .map_err(|e| KernelError::Store(e.to_string()))?
116    }
117
118    /// Append many edges in one transaction (duplicates by ID *or* by the
119    /// `(source, target, relation)` unique index are ignored).
120    pub async fn append_edges(&self, edges: Vec<GraphEdge>) -> Result<()> {
121        self.with_conn(move |c| crate::graph::store::append_edges(c, &edges))
122            .await
123            .map_err(|e| KernelError::Store(e.to_string()))?
124    }
125
126    /// Read edges touching `node_id`, filtered by direction and optional relation.
127    pub async fn edges_for_node_dir(
128        &self,
129        node_id: impl Into<String>,
130        dir: EdgeDirection,
131        relation: Option<String>,
132    ) -> Result<Vec<GraphEdge>> {
133        let node_id = node_id.into();
134        self.with_conn(move |c| {
135            crate::graph::store::edges_for_node_dir(c, &node_id, dir, relation.as_deref())
136        })
137        .await
138        .map_err(|e| KernelError::Store(e.to_string()))?
139    }
140
141    /// 1-hop neighbors of `seed_ids` (weighted sum), filtered by direction and
142    /// an optional relation. Seed nodes are excluded.
143    pub async fn neighbors_weighted(
144        &self,
145        seed_ids: Vec<String>,
146        dir: EdgeDirection,
147        relation: Option<String>,
148    ) -> Result<Vec<(String, f64)>> {
149        self.with_conn(move |c| {
150            Ok(crate::graph::traversal::neighbors_weighted(
151                c,
152                &seed_ids,
153                dir,
154                relation.as_deref(),
155            ))
156        })
157        .await
158        .map_err(|e| KernelError::Store(e.to_string()))?
159    }
160
161    /// Delete an edge by ID. Returns `true` if a row was deleted.
162    pub async fn delete_edge(&self, id: impl Into<String>) -> Result<bool> {
163        let id = id.into();
164        self.with_conn(move |c| crate::graph::store::delete_edge(c, &id))
165            .await
166            .map_err(|e| KernelError::Store(e.to_string()))?
167    }
168
169    /// Run smart recall with composite scoring.
170    pub async fn smart_recall(
171        &self,
172        project: Option<String>,
173        hint: Option<String>,
174        limit: usize,
175    ) -> Result<Vec<ScoredNode>> {
176        self.with_conn(move |c| {
177            crate::graph::recall::smart_recall(c, project.as_deref(), hint.as_deref(), limit)
178        })
179        .await
180        .map_err(|e| KernelError::Store(e.to_string()))?
181    }
182
183    /// Full-text search over node titles and bodies.
184    pub async fn search_nodes(
185        &self,
186        query: impl Into<String>,
187        limit: usize,
188    ) -> Result<Vec<GraphNode>> {
189        let query = query.into();
190        self.with_conn(move |c| crate::graph::search::search_nodes(c, &query, limit))
191            .await
192            .map_err(|e| KernelError::Store(e.to_string()))?
193    }
194
195    /// Compute graph statistics (node/edge counts, avg importance).
196    pub async fn stats(&self) -> Result<GraphStats> {
197        self.with_conn(crate::graph::lifecycle::compute_stats)
198            .await
199            .map_err(|e| KernelError::Store(e.to_string()))?
200    }
201}
202
203#[cfg(test)]
204mod tests {
205    use super::*;
206    use crate::graph::schema::init_graph_schema;
207
208    fn mem_graph() -> AsyncGraph {
209        let conn = Connection::open_in_memory().unwrap();
210        init_graph_schema(&conn).unwrap();
211        AsyncGraph::new(conn)
212    }
213
214    fn node(id: &str) -> GraphNode {
215        GraphNode {
216            id: id.into(),
217            node_type: "concept".into(),
218            title: format!("Node {id}"),
219            body: "body".into(),
220            tags: vec![],
221            projects: vec![],
222            agents: vec![],
223            created: "2026-01-01T00:00:00Z".into(),
224            updated: "2026-01-01T00:00:00Z".into(),
225            importance: 0.5,
226            access_count: 0,
227            accessed_at: String::new(),
228        }
229    }
230
231    #[tokio::test]
232    async fn upsert_and_read() {
233        let g = mem_graph();
234        g.upsert_node(node("n1")).await.unwrap();
235        let loaded = g.read_node("n1").await.unwrap().unwrap();
236        assert_eq!(loaded.id, "n1");
237    }
238
239    #[tokio::test]
240    async fn read_missing_returns_none() {
241        let g = mem_graph();
242        assert!(g.read_node("ghost").await.unwrap().is_none());
243    }
244
245    #[tokio::test]
246    async fn delete_node() {
247        let g = mem_graph();
248        g.upsert_node(node("n1")).await.unwrap();
249        assert!(g.delete_node("n1").await.unwrap());
250        assert!(g.read_node("n1").await.unwrap().is_none());
251    }
252
253    #[tokio::test]
254    async fn append_and_delete_edge() {
255        let g = mem_graph();
256        g.upsert_node(node("a")).await.unwrap();
257        g.upsert_node(node("b")).await.unwrap();
258        g.append_edge(GraphEdge {
259            id: "e1".into(),
260            source: "a".into(),
261            target: "b".into(),
262            relation: "related".into(),
263            weight: 1.0,
264            ts: "2026-01-01T00:00:00Z".into(),
265        })
266        .await
267        .unwrap();
268        assert!(g.delete_edge("e1").await.unwrap());
269    }
270
271    #[tokio::test]
272    async fn smart_recall_async() {
273        let g = mem_graph();
274        g.upsert_node(node("x")).await.unwrap();
275        let results = g.smart_recall(None, None, 10).await.unwrap();
276        assert_eq!(results.len(), 1);
277    }
278
279    #[tokio::test]
280    async fn stats_returns_counts() {
281        let g = mem_graph();
282        g.upsert_node(node("a")).await.unwrap();
283        g.upsert_node(node("b")).await.unwrap();
284        let s = g.stats().await.unwrap();
285        assert_eq!(s.total_nodes, 2);
286        assert_eq!(s.total_edges, 0);
287    }
288
289    #[tokio::test]
290    async fn clone_shares_connection() {
291        let g = mem_graph();
292        let g2 = g.clone();
293        g.upsert_node(node("n1")).await.unwrap();
294        // g2 sees the same DB
295        assert!(g2.read_node("n1").await.unwrap().is_some());
296    }
297}
llm_kernel/graph/async_graph.rs

llm_kernel/graph/
async_graph.rs
