Skip to main content

lean_ctx/core/
knowledge_relations.rs

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    /// Hebbian strengthening: saturating formula so strength approaches but never exceeds 1.0
231    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    /// Time-based exponential decay on all edge strengths
261    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    /// Remove edges whose strength has fallen below `threshold`
269    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}