kimetsu_brain/
graph_build.rs1use std::path::Path;
5
6use kimetsu_core::KimetsuResult;
7use rusqlite::{Connection, OpenFlags};
8
9use crate::projector;
10use crate::schema;
11
12#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
16pub struct GraphBuildSummary {
17 pub active_memories: usize,
19 pub rule_edges: usize,
21 pub enrich_edges: usize,
23 pub written: usize,
25 pub by_type: std::collections::BTreeMap<String, usize>,
27 pub dry_run: bool,
29}
30
31pub 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
49pub 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 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 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}