llm_kernel/graph/
traversal.rs

1//! Graph traversal: 1-hop neighbors and BFS via recursive CTEs.
2
3use std::collections::HashSet;
4
5use rusqlite::{Connection, params};
6
7use super::store::{read_edges, read_nodes};
8use super::types::{Graph, GraphEdge, GraphNodeSummary};
9
10/// Maximum edges in a graph snapshot (prevents unbounded memory).
11const MAX_GRAPH_EDGES: usize = 2000;
12
13/// Maximum seed IDs per neighbor query (keeps SQLite bind variables under limit).
14const MAX_SEED_IDS: usize = 100;
15
16/// Get 1-hop neighbors from seed IDs. Returns `(neighbor_id, total_weight)` sorted by weight DESC.
17///
18/// Follows edges in both directions (source→target and target→source).
19/// Seed nodes are excluded from results.
20pub fn graph_neighbors(conn: &Connection, seed_ids: &[String]) -> Vec<(String, f64)> {
21    if seed_ids.is_empty() {
22        return vec![];
23    }
24    let seed_ids = if seed_ids.len() > MAX_SEED_IDS {
25        &seed_ids[..MAX_SEED_IDS]
26    } else {
27        seed_ids
28    };
29
30    let ph: String = seed_ids.iter().map(|_| "?").collect::<Vec<_>>().join(",");
31    let sql = format!(
32        "SELECT target AS nb, SUM(weight) AS w FROM edges WHERE source IN ({ph}) GROUP BY target \
33         UNION ALL \
34         SELECT source AS nb, SUM(weight) AS w FROM edges WHERE target IN ({ph}) GROUP BY source"
35    );
36
37    let mut stmt = match conn.prepare(&sql) {
38        Ok(s) => s,
39        Err(_) => return vec![],
40    };
41
42    let rows: Vec<(String, f64)> = stmt
43        .query_map(
44            rusqlite::params_from_iter(seed_ids.iter().chain(seed_ids.iter())),
45            |row| Ok((row.get::<_, String>(0)?, row.get::<_, f64>(1)?)),
46        )
47        .map(|rows| rows.flatten().collect())
48        .unwrap_or_default();
49
50    let seed_set: HashSet<&str> = seed_ids.iter().map(String::as_str).collect();
51    let mut weights: std::collections::HashMap<String, f64> = std::collections::HashMap::new();
52    for (nid, w) in rows {
53        if !seed_set.contains(nid.as_str()) {
54            *weights.entry(nid).or_default() += w;
55        }
56    }
57
58    let mut result: Vec<(String, f64)> = weights.into_iter().collect();
59    result.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
60    result
61}
62
63/// BFS traversal from `start_id` via SQL recursive CTE.
64///
65/// Returns all reachable node IDs (excluding start), capped at 500.
66/// Follows edges in both directions.
67pub fn related_nodes(conn: &Connection, start_id: &str, depth: usize) -> Vec<String> {
68    let sql = "
69        WITH RECURSIVE bfs(node_id, lvl) AS (
70            SELECT target, 1 FROM edges WHERE source = ?1
71            UNION SELECT source, 1 FROM edges WHERE target = ?1
72            UNION SELECT e.target, bfs.lvl + 1 FROM edges e
73                JOIN bfs ON e.source = bfs.node_id
74                WHERE e.target != ?1 AND bfs.lvl < ?2
75            UNION SELECT e.source, bfs.lvl + 1 FROM edges e
76                JOIN bfs ON e.target = bfs.node_id
77                WHERE e.source != ?1 AND bfs.lvl < ?2
78        )
79        SELECT DISTINCT node_id FROM bfs
80        LIMIT 500
81    ";
82
83    conn.prepare(sql)
84        .and_then(|mut stmt| {
85            stmt.query_map(params![start_id, depth as i64], |row| {
86                row.get::<_, String>(0)
87            })
88            .map(|rows| rows.flatten().collect())
89        })
90        .unwrap_or_default()
91}
92
93/// Build a graph snapshot (node summaries + edges) from the database.
94pub fn build_graph(conn: &Connection) -> crate::error::Result<Graph> {
95    let ids = super::store::list_node_ids(conn)?;
96    let id_refs: Vec<&str> = ids.iter().map(String::as_str).collect();
97    let nodes: Vec<GraphNodeSummary> = read_nodes(conn, &id_refs)?
98        .into_iter()
99        .map(|node| GraphNodeSummary {
100            id: node.id,
101            title: node.title,
102            node_type: node.node_type,
103            tags: node.tags,
104            importance: node.importance,
105        })
106        .collect();
107    let edges: Vec<GraphEdge> = read_edges(conn, MAX_GRAPH_EDGES)?;
108    Ok(Graph { nodes, edges })
109}
110
111#[cfg(test)]
112mod tests {
113    use super::*;
114    use crate::graph::schema::init_graph_schema;
115    use crate::graph::store::append_edge;
116    use rusqlite::Connection;
117
118    fn mem_db() -> Connection {
119        let conn = Connection::open_in_memory().unwrap();
120        init_graph_schema(&conn).unwrap();
121        conn
122    }
123
124    fn insert_edge(conn: &Connection, id: &str, src: &str, tgt: &str) {
125        let e = GraphEdge {
126            id: id.to_string(),
127            source: src.to_string(),
128            target: tgt.to_string(),
129            relation: "related".to_string(),
130            weight: 1.0,
131            ts: "2026-01-01T00:00:00Z".to_string(),
132        };
133        append_edge(conn, &e).unwrap();
134    }
135
136    #[test]
137    fn neighbors_returns_direct_connections() {
138        let conn = mem_db();
139        insert_edge(&conn, "e1", "A", "B");
140        insert_edge(&conn, "e2", "A", "C");
141        insert_edge(&conn, "e3", "D", "A");
142
143        let mut result = graph_neighbors(&conn, &["A".to_string()]);
144        result.sort_by(|a, b| a.0.cmp(&b.0));
145        let ids: Vec<&str> = result.iter().map(|r| r.0.as_str()).collect();
146        assert!(ids.contains(&"B"));
147        assert!(ids.contains(&"C"));
148        assert!(ids.contains(&"D"));
149        assert!(!ids.contains(&"A"));
150    }
151
152    #[test]
153    fn neighbors_excludes_seeds() {
154        let conn = mem_db();
155        insert_edge(&conn, "e1", "A", "B");
156        insert_edge(&conn, "e2", "B", "C");
157
158        let result = graph_neighbors(&conn, &["A".to_string(), "B".to_string()]);
159        let ids: Vec<&str> = result.iter().map(|r| r.0.as_str()).collect();
160        assert!(ids.contains(&"C"));
161        assert!(!ids.contains(&"A"));
162        assert!(!ids.contains(&"B"));
163    }
164
165    #[test]
166    fn neighbors_empty_seeds() {
167        let conn = mem_db();
168        assert!(graph_neighbors(&conn, &[]).is_empty());
169    }
170
171    #[test]
172    fn related_nodes_recursive_bfs() {
173        let conn = mem_db();
174        insert_edge(&conn, "e1", "A", "B");
175        insert_edge(&conn, "e2", "B", "C");
176
177        let result = related_nodes(&conn, "A", 2);
178        assert!(result.contains(&"B".to_string()));
179        assert!(result.contains(&"C".to_string()));
180        assert!(!result.contains(&"A".to_string()));
181    }
182
183    #[test]
184    fn related_nodes_handles_cycles() {
185        let conn = mem_db();
186        insert_edge(&conn, "e1", "A", "B");
187        insert_edge(&conn, "e2", "B", "C");
188        insert_edge(&conn, "e3", "C", "A");
189
190        let result = related_nodes(&conn, "A", 3);
191        let unique: HashSet<_> = result.iter().collect();
192        assert_eq!(result.len(), unique.len(), "no duplicates in cycle");
193        assert!(!result.contains(&"A".to_string()));
194    }
195
196    #[test]
197    fn neighbor_weight_accumulation() {
198        let conn = mem_db();
199        insert_edge(&conn, "e1", "A", "C");
200        insert_edge(&conn, "e2", "B", "C");
201
202        let result = graph_neighbors(&conn, &["A".to_string(), "B".to_string()]);
203        let c_weight = result.iter().find(|(id, _)| id == "C").map(|(_, w)| *w);
204        assert_eq!(c_weight, Some(2.0));
205    }
206}
llm_kernel/graph/traversal.rs

llm_kernel/graph/
traversal.rs
