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::{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    /// Delete an edge by ID. Returns `true` if a row was deleted.
119    pub async fn delete_edge(&self, id: impl Into<String>) -> Result<bool> {
120        let id = id.into();
121        self.with_conn(move |c| crate::graph::store::delete_edge(c, &id))
122            .await
123            .map_err(|e| KernelError::Store(e.to_string()))?
124    }
125
126    /// Run smart recall with composite scoring.
127    pub async fn smart_recall(
128        &self,
129        project: Option<String>,
130        hint: Option<String>,
131        limit: usize,
132    ) -> Result<Vec<ScoredNode>> {
133        self.with_conn(move |c| {
134            crate::graph::recall::smart_recall(c, project.as_deref(), hint.as_deref(), limit)
135        })
136        .await
137        .map_err(|e| KernelError::Store(e.to_string()))?
138    }
139
140    /// Full-text search over node titles and bodies.
141    pub async fn search_nodes(
142        &self,
143        query: impl Into<String>,
144        limit: usize,
145    ) -> Result<Vec<GraphNode>> {
146        let query = query.into();
147        self.with_conn(move |c| crate::graph::search::search_nodes(c, &query, limit))
148            .await
149            .map_err(|e| KernelError::Store(e.to_string()))?
150    }
151
152    /// Compute graph statistics (node/edge counts, avg importance).
153    pub async fn stats(&self) -> Result<GraphStats> {
154        self.with_conn(crate::graph::lifecycle::compute_stats)
155            .await
156            .map_err(|e| KernelError::Store(e.to_string()))?
157    }
158}
159
160#[cfg(test)]
161mod tests {
162    use super::*;
163    use crate::graph::schema::init_graph_schema;
164
165    fn mem_graph() -> AsyncGraph {
166        let conn = Connection::open_in_memory().unwrap();
167        init_graph_schema(&conn).unwrap();
168        AsyncGraph::new(conn)
169    }
170
171    fn node(id: &str) -> GraphNode {
172        GraphNode {
173            id: id.into(),
174            node_type: "concept".into(),
175            title: format!("Node {id}"),
176            body: "body".into(),
177            tags: vec![],
178            projects: vec![],
179            agents: vec![],
180            created: "2026-01-01T00:00:00Z".into(),
181            updated: "2026-01-01T00:00:00Z".into(),
182            importance: 0.5,
183            access_count: 0,
184            accessed_at: String::new(),
185        }
186    }
187
188    #[tokio::test]
189    async fn upsert_and_read() {
190        let g = mem_graph();
191        g.upsert_node(node("n1")).await.unwrap();
192        let loaded = g.read_node("n1").await.unwrap().unwrap();
193        assert_eq!(loaded.id, "n1");
194    }
195
196    #[tokio::test]
197    async fn read_missing_returns_none() {
198        let g = mem_graph();
199        assert!(g.read_node("ghost").await.unwrap().is_none());
200    }
201
202    #[tokio::test]
203    async fn delete_node() {
204        let g = mem_graph();
205        g.upsert_node(node("n1")).await.unwrap();
206        assert!(g.delete_node("n1").await.unwrap());
207        assert!(g.read_node("n1").await.unwrap().is_none());
208    }
209
210    #[tokio::test]
211    async fn append_and_delete_edge() {
212        let g = mem_graph();
213        g.upsert_node(node("a")).await.unwrap();
214        g.upsert_node(node("b")).await.unwrap();
215        g.append_edge(GraphEdge {
216            id: "e1".into(),
217            source: "a".into(),
218            target: "b".into(),
219            relation: "related".into(),
220            weight: 1.0,
221            ts: "2026-01-01T00:00:00Z".into(),
222        })
223        .await
224        .unwrap();
225        assert!(g.delete_edge("e1").await.unwrap());
226    }
227
228    #[tokio::test]
229    async fn smart_recall_async() {
230        let g = mem_graph();
231        g.upsert_node(node("x")).await.unwrap();
232        let results = g.smart_recall(None, None, 10).await.unwrap();
233        assert_eq!(results.len(), 1);
234    }
235
236    #[tokio::test]
237    async fn stats_returns_counts() {
238        let g = mem_graph();
239        g.upsert_node(node("a")).await.unwrap();
240        g.upsert_node(node("b")).await.unwrap();
241        let s = g.stats().await.unwrap();
242        assert_eq!(s.total_nodes, 2);
243        assert_eq!(s.total_edges, 0);
244    }
245
246    #[tokio::test]
247    async fn clone_shares_connection() {
248        let g = mem_graph();
249        let g2 = g.clone();
250        g.upsert_node(node("n1")).await.unwrap();
251        // g2 sees the same DB
252        assert!(g2.read_node("n1").await.unwrap().is_some());
253    }
254}
llm_kernel/graph/async_graph.rs

llm_kernel/graph/
async_graph.rs
