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    // v2.6: refresh the entity index before deriving edges. A brain that was
76    // migrated with an unreadable corpus, or one whose extractor rules have
77    // since changed, would otherwise build its graph from a stale index —
78    // and this command is exactly what a user runs to fix that.
79    if !dry_run {
80        let _ = crate::graph::reproject_all_entities(&conn);
81    }
82
83    let rule = crate::graph::build_relates_to_edges(&conn, max_fan_out)?;
84    let rule_edges = rule.len();
85    let enrich_edges = extra_edges.len();
86
87    // Merge rule + enrichment, de-duplicating on (src, dst, type). Self-loops are
88    // dropped by add_memory_edges; we also drop them here for an accurate summary.
89    let mut seen: BTreeSet<(String, String, String)> = BTreeSet::new();
90    let mut by_type: BTreeMap<String, usize> = BTreeMap::new();
91    let mut merged: Vec<(String, String, String)> = Vec::new();
92    let push = |src: String,
93                dst: String,
94                ty: String,
95                seen: &mut BTreeSet<(String, String, String)>,
96                by_type: &mut BTreeMap<String, usize>,
97                merged: &mut Vec<(String, String, String)>| {
98        if src == dst {
99            return;
100        }
101        let key = (src.clone(), dst.clone(), ty.clone());
102        if seen.insert(key) {
103            *by_type.entry(ty.clone()).or_insert(0) += 1;
104            merged.push((src, dst, ty));
105        }
106    };
107    for e in &rule {
108        push(
109            e.src_id.clone(),
110            e.dst_id.clone(),
111            e.edge_type.clone(),
112            &mut seen,
113            &mut by_type,
114            &mut merged,
115        );
116    }
117    for (src, dst, ty) in extra_edges {
118        push(
119            src.clone(),
120            dst.clone(),
121            ty.clone(),
122            &mut seen,
123            &mut by_type,
124            &mut merged,
125        );
126    }
127
128    let written = if dry_run {
129        0
130    } else {
131        projector::add_memory_edges(&conn, &merged)?
132    };
133
134    Ok(GraphBuildSummary {
135        active_memories,
136        rule_edges,
137        enrich_edges,
138        written,
139        by_type,
140        dry_run,
141    })
142}