Skip to main content

kimetsu_brain/
graph_build.rs

1//! Knowledge-graph edge building over active memories.
2//! Split out of `project.rs` (v2.5.1); re-exported by [`crate::project`].
3
4use std::path::Path;
5
6use kimetsu_core::KimetsuResult;
7use rusqlite::{Connection, OpenFlags};
8
9use crate::projector;
10use crate::schema;
11
12// ── #2 knowledge graph: build relation edges ─────────────────────────────────
13
14/// Summary of a `kimetsu brain graph build` run.
15#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
16pub struct GraphBuildSummary {
17    /// Active (non-invalidated, non-superseded) memories scanned.
18    pub active_memories: usize,
19    /// Rule-derived `relates_to` edges proposed.
20    pub rule_edges: usize,
21    /// Enrichment (LLM typed) edges proposed by the caller.
22    pub enrich_edges: usize,
23    /// Edges actually written (0 when `dry_run`).
24    pub written: usize,
25    /// Proposed edge counts grouped by edge_type (rule + enrichment, pre-write).
26    pub by_type: std::collections::BTreeMap<String, usize>,
27    /// True when no edges were persisted (preview only).
28    pub dry_run: bool,
29}
30
31/// Read every active memory as `(id, text)` for graph enrichment. Read-only.
32/// Exposed so the CLI (which owns the cheap-model provider) can compute typed
33/// enrichment edges before calling [`build_graph`].
34pub fn active_memory_texts(start: &Path) -> KimetsuResult<Vec<(String, String)>> {
35    let paths = kimetsu_core::paths::ProjectPaths::discover(start)?;
36    let conn = Connection::open_with_flags(&paths.brain_db, OpenFlags::SQLITE_OPEN_READ_ONLY)?;
37    let mut stmt = conn.prepare(
38        "SELECT memory_id, text
39         FROM memories
40         WHERE invalidated_at IS NULL AND superseded_by IS NULL
41         ORDER BY memory_id",
42    )?;
43    let rows = stmt
44        .query_map([], |r| Ok((r.get::<_, String>(0)?, r.get::<_, String>(1)?)))?
45        .collect::<Result<Vec<_>, _>>()?;
46    Ok(rows)
47}
48
49/// #2: build the knowledge-graph edges for the workspace brain.
50///
51/// Combines the deterministic rule layer ([`crate::graph::build_relates_to_edges`])
52/// with any caller-supplied `extra_edges` (LLM enrichment computed in the CLI),
53/// de-duplicates, and — unless `dry_run` — persists them as rebuild-safe
54/// `memory.edge` events via [`projector::add_memory_edges`]. Returns a summary.
55///
56/// `max_fan_out` caps rule edges per source memory (0 = the module default).
57pub fn build_graph(
58    start: &Path,
59    extra_edges: &[(String, String, String)],
60    max_fan_out: usize,
61    dry_run: bool,
62) -> KimetsuResult<GraphBuildSummary> {
63    use std::collections::{BTreeMap, BTreeSet};
64
65    let paths = kimetsu_core::paths::ProjectPaths::discover(start)?;
66    let conn = Connection::open(&paths.brain_db)?;
67    schema::initialize(&conn)?;
68
69    let active_memories = conn.query_row(
70        "SELECT COUNT(*) FROM memories WHERE invalidated_at IS NULL AND superseded_by IS NULL",
71        [],
72        |r| r.get::<_, i64>(0),
73    )? as usize;
74
75    let rule = crate::graph::build_relates_to_edges(&conn, max_fan_out)?;
76    let rule_edges = rule.len();
77    let enrich_edges = extra_edges.len();
78
79    // Merge rule + enrichment, de-duplicating on (src, dst, type). Self-loops are
80    // dropped by add_memory_edges; we also drop them here for an accurate summary.
81    let mut seen: BTreeSet<(String, String, String)> = BTreeSet::new();
82    let mut by_type: BTreeMap<String, usize> = BTreeMap::new();
83    let mut merged: Vec<(String, String, String)> = Vec::new();
84    let push = |src: String,
85                dst: String,
86                ty: String,
87                seen: &mut BTreeSet<(String, String, String)>,
88                by_type: &mut BTreeMap<String, usize>,
89                merged: &mut Vec<(String, String, String)>| {
90        if src == dst {
91            return;
92        }
93        let key = (src.clone(), dst.clone(), ty.clone());
94        if seen.insert(key) {
95            *by_type.entry(ty.clone()).or_insert(0) += 1;
96            merged.push((src, dst, ty));
97        }
98    };
99    for e in &rule {
100        push(
101            e.src_id.clone(),
102            e.dst_id.clone(),
103            e.edge_type.clone(),
104            &mut seen,
105            &mut by_type,
106            &mut merged,
107        );
108    }
109    for (src, dst, ty) in extra_edges {
110        push(
111            src.clone(),
112            dst.clone(),
113            ty.clone(),
114            &mut seen,
115            &mut by_type,
116            &mut merged,
117        );
118    }
119
120    let written = if dry_run {
121        0
122    } else {
123        projector::add_memory_edges(&conn, &merged)?
124    };
125
126    Ok(GraphBuildSummary {
127        active_memories,
128        rule_edges,
129        enrich_edges,
130        written,
131        by_type,
132        dry_run,
133    })
134}