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

llm_kernel/graph/
async_graph.rs
