mentedb_cognitive/
pain.rs1use 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 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}