lean_ctx/core/
knowledge_relations.rs

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

lean_ctx/core/
knowledge_relations.rs
