1use chrono::{DateTime, Utc};
2use serde::{Deserialize, Serialize};
3use std::path::PathBuf;
4
5#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
6pub struct KnowledgeNodeRef {
7 pub category: String,
8 pub key: String,
9}
10
11impl KnowledgeNodeRef {
12 pub fn new(category: &str, key: &str) -> Self {
13 Self {
14 category: category.trim().to_string(),
15 key: key.trim().to_string(),
16 }
17 }
18
19 pub fn id(&self) -> String {
20 format!("{}/{}", self.category, self.key)
21 }
22}
23
24#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)]
25#[serde(rename_all = "snake_case")]
26pub enum KnowledgeEdgeKind {
27 DependsOn,
28 RelatedTo,
29 Supports,
30 Contradicts,
31 Supersedes,
32}
33
34impl KnowledgeEdgeKind {
35 pub fn parse(input: &str) -> Option<Self> {
36 match input.trim().to_lowercase().as_str() {
37 "depends_on" | "depends" => Some(Self::DependsOn),
38 "related_to" | "related" => Some(Self::RelatedTo),
39 "supports" | "support" => Some(Self::Supports),
40 "contradicts" | "contradict" => Some(Self::Contradicts),
41 "supersedes" | "supersede" => Some(Self::Supersedes),
42 _ => None,
43 }
44 }
45
46 pub fn as_str(&self) -> &'static str {
47 match self {
48 KnowledgeEdgeKind::DependsOn => "depends_on",
49 KnowledgeEdgeKind::RelatedTo => "related_to",
50 KnowledgeEdgeKind::Supports => "supports",
51 KnowledgeEdgeKind::Contradicts => "contradicts",
52 KnowledgeEdgeKind::Supersedes => "supersedes",
53 }
54 }
55}
56
57fn default_strength() -> f64 {
58 0.5
59}
60fn default_decay_rate() -> f64 {
61 0.02
62}
63
64#[derive(Debug, Clone, Serialize, Deserialize)]
65pub struct KnowledgeEdge {
66 pub from: KnowledgeNodeRef,
67 pub to: KnowledgeNodeRef,
68 pub kind: KnowledgeEdgeKind,
69 pub created_at: DateTime<Utc>,
70 #[serde(default)]
71 pub last_seen: Option<DateTime<Utc>>,
72 #[serde(default)]
73 pub count: u32,
74 pub source_session: String,
75 #[serde(default = "default_strength")]
76 pub strength: f64,
77 #[serde(default = "default_decay_rate")]
78 pub decay_rate: f64,
79}
80
81#[derive(Debug, Clone, Serialize, Deserialize)]
82#[serde(default)]
83pub struct KnowledgeRelationGraph {
84 pub project_hash: String,
85 pub edges: Vec<KnowledgeEdge>,
86 pub updated_at: DateTime<Utc>,
87}
88
89impl Default for KnowledgeRelationGraph {
90 fn default() -> Self {
91 Self {
92 project_hash: String::new(),
93 edges: Vec::new(),
94 updated_at: Utc::now(),
95 }
96 }
97}
98
99impl KnowledgeRelationGraph {
100 pub fn new(project_hash: &str) -> Self {
101 Self {
102 project_hash: project_hash.to_string(),
103 edges: Vec::new(),
104 updated_at: Utc::now(),
105 }
106 }
107
108 pub fn path(project_hash: &str) -> Result<PathBuf, String> {
109 let dir = crate::core::data_dir::lean_ctx_data_dir()?
110 .join("knowledge")
111 .join(project_hash);
112 Ok(dir.join("relations.json"))
113 }
114
115 pub fn load(project_hash: &str) -> Option<Self> {
116 let path = Self::path(project_hash).ok()?;
117 let content = std::fs::read_to_string(&path).ok()?;
118 let mut g = serde_json::from_str::<Self>(&content).ok()?;
119 if g.project_hash.trim().is_empty() {
120 g.project_hash = project_hash.to_string();
121 }
122 Some(g)
123 }
124
125 pub fn load_or_create(project_hash: &str) -> Self {
126 Self::load(project_hash).unwrap_or_else(|| Self::new(project_hash))
127 }
128
129 pub fn save(&mut self) -> Result<(), String> {
130 let path = Self::path(&self.project_hash)?;
131 if let Some(dir) = path.parent() {
132 std::fs::create_dir_all(dir).map_err(|e| e.to_string())?;
133 }
134
135 self.updated_at = Utc::now();
136 self.edges.sort_by(|a, b| {
137 a.from
138 .category
139 .cmp(&b.from.category)
140 .then_with(|| a.from.key.cmp(&b.from.key))
141 .then_with(|| a.kind.as_str().cmp(b.kind.as_str()))
142 .then_with(|| a.to.category.cmp(&b.to.category))
143 .then_with(|| a.to.key.cmp(&b.to.key))
144 .then_with(|| b.count.cmp(&a.count))
145 .then_with(|| b.last_seen.cmp(&a.last_seen))
146 .then_with(|| b.created_at.cmp(&a.created_at))
147 });
148
149 let json = serde_json::to_string_pretty(self).map_err(|e| e.to_string())?;
150 std::fs::write(&path, json).map_err(|e| e.to_string())
151 }
152
153 pub fn upsert_edge(
154 &mut self,
155 from: KnowledgeNodeRef,
156 to: KnowledgeNodeRef,
157 kind: KnowledgeEdgeKind,
158 session_id: &str,
159 ) -> bool {
160 let now = Utc::now();
161 if let Some(e) = self
162 .edges
163 .iter_mut()
164 .find(|e| e.from == from && e.to == to && e.kind == kind)
165 {
166 e.count = e.count.saturating_add(1).max(1);
167 e.last_seen = Some(now);
168 e.source_session = session_id.to_string();
169 e.strength = (e.strength + 0.1 * (1.0 - e.strength)).min(1.0);
170 self.updated_at = now;
171 return false;
172 }
173
174 self.edges.push(KnowledgeEdge {
175 from,
176 to,
177 kind,
178 created_at: now,
179 last_seen: Some(now),
180 count: 1,
181 source_session: session_id.to_string(),
182 strength: default_strength(),
183 decay_rate: default_decay_rate(),
184 });
185 self.updated_at = now;
186 true
187 }
188
189 pub fn remove_edge(
190 &mut self,
191 from: &KnowledgeNodeRef,
192 to: &KnowledgeNodeRef,
193 kind: Option<KnowledgeEdgeKind>,
194 ) -> usize {
195 let before = self.edges.len();
196 self.edges.retain(|e| {
197 if &e.from != from || &e.to != to {
198 return true;
199 }
200 if let Some(k) = kind {
201 e.kind != k
202 } else {
203 false
204 }
205 });
206 before.saturating_sub(self.edges.len())
207 }
208
209 pub fn enforce_cap(&mut self, max_edges: usize) -> bool {
210 if max_edges == 0 || self.edges.len() <= max_edges {
211 return false;
212 }
213
214 self.edges.sort_by(|a, b| {
215 b.count
216 .cmp(&a.count)
217 .then_with(|| b.last_seen.cmp(&a.last_seen))
218 .then_with(|| b.created_at.cmp(&a.created_at))
219 .then_with(|| a.from.category.cmp(&b.from.category))
220 .then_with(|| a.from.key.cmp(&b.from.key))
221 .then_with(|| a.kind.as_str().cmp(b.kind.as_str()))
222 .then_with(|| a.to.category.cmp(&b.to.category))
223 .then_with(|| a.to.key.cmp(&b.to.key))
224 });
225
226 self.edges.truncate(max_edges);
227 true
228 }
229
230 pub fn strengthen_edge(
232 &mut self,
233 from: &KnowledgeNodeRef,
234 to: &KnowledgeNodeRef,
235 amount: f64,
236 ) -> bool {
237 if let Some(e) = self
238 .edges
239 .iter_mut()
240 .find(|e| &e.from == from && &e.to == to)
241 {
242 e.strength = (e.strength + amount * (1.0 - e.strength)).min(1.0);
243 e.last_seen = Some(Utc::now());
244 e.count = e.count.saturating_add(1);
245 return true;
246 }
247 if let Some(e) = self
248 .edges
249 .iter_mut()
250 .find(|e| &e.from == to && &e.to == from)
251 {
252 e.strength = (e.strength + amount * (1.0 - e.strength)).min(1.0);
253 e.last_seen = Some(Utc::now());
254 e.count = e.count.saturating_add(1);
255 return true;
256 }
257 false
258 }
259
260 pub fn decay_all_edges(&mut self, days_elapsed: f64) {
262 for e in &mut self.edges {
263 e.strength *= (1.0 - e.decay_rate).powf(days_elapsed);
264 e.strength = e.strength.max(0.0);
265 }
266 }
267
268 pub fn prune_weak_edges(&mut self, threshold: f64) -> usize {
270 let before = self.edges.len();
271 self.edges.retain(|e| e.strength >= threshold);
272 before - self.edges.len()
273 }
274}
275
276pub fn parse_node_ref(input: &str) -> Option<KnowledgeNodeRef> {
277 let s = input.trim();
278 if s.is_empty() {
279 return None;
280 }
281
282 if let Some((cat, key)) = s.split_once('/') {
283 let cat = cat.trim();
284 let key = key.trim();
285 if !cat.is_empty() && !key.is_empty() {
286 return Some(KnowledgeNodeRef::new(cat, key));
287 }
288 }
289 if let Some((cat, key)) = s.split_once(':') {
290 let cat = cat.trim();
291 let key = key.trim();
292 if !cat.is_empty() && !key.is_empty() {
293 return Some(KnowledgeNodeRef::new(cat, key));
294 }
295 }
296
297 None
298}
299
300pub fn format_mermaid(edges: &[KnowledgeEdge]) -> String {
301 if edges.is_empty() {
302 return "graph TD\n %% no relations".to_string();
303 }
304
305 fn id_for(n: &KnowledgeNodeRef) -> String {
306 let mut out = String::from("K_");
307 for ch in n.id().chars() {
308 if ch.is_ascii_alphanumeric() {
309 out.push(ch);
310 } else {
311 out.push('_');
312 }
313 }
314 out
315 }
316
317 let mut lines = Vec::new();
318 lines.push("graph TD".to_string());
319 for e in edges {
320 let from = id_for(&e.from);
321 let to = id_for(&e.to);
322 let from_label = e.from.id();
323 let to_label = e.to.id();
324 lines.push(format!(
325 " {from}[\"{from_label}\"] -->|{}| {to}[\"{to_label}\"]",
326 e.kind.as_str()
327 ));
328 }
329 lines.join("\n")
330}
331
332#[cfg(test)]
333mod tests {
334 use super::*;
335
336 #[test]
337 fn strengthen_edge_saturating() {
338 let mut graph = KnowledgeRelationGraph::new("test");
339 let from = KnowledgeNodeRef::new("a", "1");
340 let to = KnowledgeNodeRef::new("b", "2");
341 graph.upsert_edge(from.clone(), to.clone(), KnowledgeEdgeKind::RelatedTo, "s1");
342
343 let initial = graph.edges[0].strength;
344 assert!((initial - 0.5).abs() < 0.01);
345
346 graph.strengthen_edge(&from, &to, 0.3);
347 assert!(graph.edges[0].strength > initial);
348 assert!(graph.edges[0].strength <= 1.0);
349
350 for _ in 0..100 {
351 graph.strengthen_edge(&from, &to, 0.5);
352 }
353 assert!(graph.edges[0].strength <= 1.0);
354 assert!(graph.edges[0].strength > 0.99);
355 }
356
357 #[test]
358 fn decay_reduces_strength() {
359 let mut graph = KnowledgeRelationGraph::new("test");
360 let from = KnowledgeNodeRef::new("a", "1");
361 let to = KnowledgeNodeRef::new("b", "2");
362 graph.upsert_edge(from, to, KnowledgeEdgeKind::RelatedTo, "s1");
363
364 let initial = graph.edges[0].strength;
365 graph.decay_all_edges(10.0);
366 assert!(graph.edges[0].strength < initial);
367 assert!(graph.edges[0].strength > 0.0);
368 }
369
370 #[test]
371 fn prune_weak_edges_removes_below_threshold() {
372 let mut graph = KnowledgeRelationGraph::new("test");
373 graph.upsert_edge(
374 KnowledgeNodeRef::new("a", "1"),
375 KnowledgeNodeRef::new("b", "2"),
376 KnowledgeEdgeKind::RelatedTo,
377 "s1",
378 );
379 graph.upsert_edge(
380 KnowledgeNodeRef::new("c", "3"),
381 KnowledgeNodeRef::new("d", "4"),
382 KnowledgeEdgeKind::RelatedTo,
383 "s2",
384 );
385
386 graph.edges[1].strength = 0.01;
387
388 let removed = graph.prune_weak_edges(0.05);
389 assert_eq!(removed, 1);
390 assert_eq!(graph.edges.len(), 1);
391 }
392
393 #[test]
394 fn backward_compatible_edge_deserialization() {
395 let json = r#"{
396 "from": {"category": "a", "key": "1"},
397 "to": {"category": "b", "key": "2"},
398 "kind": "related_to",
399 "created_at": "2024-01-01T00:00:00Z",
400 "count": 1,
401 "source_session": "s1"
402 }"#;
403 let edge: KnowledgeEdge = serde_json::from_str(json).unwrap();
404 assert!((edge.strength - 0.5).abs() < 0.01);
405 assert!((edge.decay_rate - 0.02).abs() < 0.001);
406 }
407
408 #[test]
409 fn strengthen_edge_bidirectional() {
410 let mut graph = KnowledgeRelationGraph::new("test");
411 let from = KnowledgeNodeRef::new("a", "1");
412 let to = KnowledgeNodeRef::new("b", "2");
413 graph.upsert_edge(from.clone(), to.clone(), KnowledgeEdgeKind::RelatedTo, "s1");
414
415 let found = graph.strengthen_edge(&to, &from, 0.2);
416 assert!(found);
417 assert!(graph.edges[0].strength > 0.5);
418 }
419}