Skip to main content

mentedb_cognitive/
pain.rs

1use mentedb_core::types::{MemoryId, Timestamp};
2use serde::{Deserialize, Serialize};
3
4#[derive(Debug, Clone, Serialize, Deserialize)]
5pub struct PainSignal {
6    pub id: MemoryId,
7    pub memory_id: MemoryId,
8    pub intensity: f32,
9    pub trigger_keywords: Vec<String>,
10    pub description: String,
11    pub created_at: Timestamp,
12    pub decay_rate: f32,
13}
14
15pub struct PainRegistry {
16    signals: Vec<PainSignal>,
17    max_warnings: usize,
18}
19
20impl PainRegistry {
21    pub fn new(max_warnings: usize) -> Self {
22        Self {
23            signals: Vec::new(),
24            max_warnings,
25        }
26    }
27
28    pub fn record_pain(&mut self, signal: PainSignal) {
29        self.signals.push(signal);
30    }
31
32    /// Returns all recorded pain signals.
33    pub fn all_signals(&self) -> &[PainSignal] {
34        &self.signals
35    }
36
37    pub fn get_pain_for_context(&self, context_keywords: &[String]) -> Vec<&PainSignal> {
38        let ctx_lower: Vec<String> = context_keywords.iter().map(|k| k.to_lowercase()).collect();
39
40        let mut scored: Vec<(&PainSignal, f32)> = self
41            .signals
42            .iter()
43            .filter_map(|signal| {
44                let matched = signal
45                    .trigger_keywords
46                    .iter()
47                    .filter(|trigger| {
48                        ctx_lower
49                            .iter()
50                            .any(|ctx| ctx.contains(&trigger.to_lowercase()))
51                    })
52                    .count();
53
54                if matched > 0 {
55                    let relevance = matched as f32 / signal.trigger_keywords.len().max(1) as f32;
56                    let score = signal.intensity * relevance;
57                    Some((signal, score))
58                } else {
59                    None
60                }
61            })
62            .collect();
63
64        scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
65        scored.into_iter().map(|(s, _)| s).collect()
66    }
67
68    pub fn decay_all(&mut self, current_time: Timestamp) {
69        for signal in &mut self.signals {
70            signal.intensity = Self::compute_decayed(signal, current_time);
71        }
72    }
73
74    pub fn current_intensity(signal: &PainSignal, now: Timestamp) -> f32 {
75        Self::compute_decayed(signal, now)
76    }
77
78    fn compute_decayed(signal: &PainSignal, now: Timestamp) -> f32 {
79        if now <= signal.created_at {
80            return signal.intensity;
81        }
82        let time_delta = (now - signal.created_at) as f64;
83        let decayed = signal.intensity as f64 * (-signal.decay_rate as f64 * time_delta).exp();
84        decayed as f32
85    }
86
87    pub fn format_pain_warnings(&self, signals: &[&PainSignal]) -> String {
88        if signals.is_empty() {
89            return String::new();
90        }
91
92        signals
93            .iter()
94            .take(self.max_warnings)
95            .map(|s| {
96                format!(
97                    "CAUTION: {} (pain: {:.2}). Triggers: [{}]",
98                    s.description,
99                    s.intensity,
100                    s.trigger_keywords.join(", ")
101                )
102            })
103            .collect::<Vec<_>>()
104            .join("\n")
105    }
106}
107
108impl Default for PainRegistry {
109    fn default() -> Self {
110        Self::new(5)
111    }
112}
113
114#[cfg(test)]
115mod tests {
116    use super::*;
117
118    fn make_signal(keywords: Vec<&str>, intensity: f32) -> PainSignal {
119        PainSignal {
120            id: MemoryId::new(),
121            memory_id: MemoryId::new(),
122            intensity,
123            trigger_keywords: keywords.into_iter().map(String::from).collect(),
124            description: "Test pain signal".to_string(),
125            created_at: 1000,
126            decay_rate: 0.0001,
127        }
128    }
129
130    #[test]
131    fn test_pain_trigger_matching() {
132        let mut registry = PainRegistry::default();
133        registry.record_pain(make_signal(vec!["mongodb", "nosql"], 0.95));
134        registry.record_pain(make_signal(vec!["python", "flask"], 0.5));
135
136        let results =
137            registry.get_pain_for_context(&["mongodb".to_string(), "database".to_string()]);
138        assert_eq!(results.len(), 1);
139        assert!(results[0].trigger_keywords.contains(&"mongodb".to_string()));
140    }
141
142    #[test]
143    fn test_decay() {
144        let signal = make_signal(vec!["test"], 1.0);
145        let decayed = PainRegistry::current_intensity(&signal, 11000);
146        assert!(decayed < 1.0, "Expected decay, got {}", decayed);
147        assert!(
148            decayed > 0.0,
149            "Expected positive intensity, got {}",
150            decayed
151        );
152    }
153
154    #[test]
155    fn test_format_warnings() {
156        let registry = PainRegistry::default();
157        let s = make_signal(vec!["mongodb", "nosql"], 0.95);
158        let warnings = registry.format_pain_warnings(&[&s]);
159        assert!(warnings.contains("CAUTION"));
160        assert!(warnings.contains("0.95"));
161    }
162}