aimem_core/
extractor.rs

1//! Memory extractor — heuristic keyword extraction, no LLM required.
2//!
3//! Identifies 5 memory types from text:
4//!  1. DECISIONS    — choices made, reasoning
5//!  2. PREFERENCES  — always/never/prefer patterns
6//!  3. MILESTONES   — breakthroughs, things that finally worked
7//!  4. PROBLEMS     — what broke, root causes, fixes
8//!  5. EMOTIONAL    — feelings, relationships
9//!
10//! Includes a narrow multilingual rule set (EN / ZH / CAN / JP) while trying
11//! hard to avoid obvious false positives from overly broad keywords.
12
13use regex::Regex;
14use std::sync::OnceLock;
15
16#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
17pub enum MemoryType {
18    Decision,
19    Preference,
20    Milestone,
21    Problem,
22    Emotional,
23    General,
24}
25
26impl std::fmt::Display for MemoryType {
27    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
28        let s = match self {
29            Self::Decision => "decision",
30            Self::Preference => "preference",
31            Self::Milestone => "milestone",
32            Self::Problem => "problem",
33            Self::Emotional => "emotional",
34            Self::General => "general",
35        };
36        write!(f, "{s}")
37    }
38}
39
40#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
41pub struct ExtractedMemory {
42    pub content: String,
43    pub memory_type: MemoryType,
44    pub chunk_index: usize,
45}
46
47// ── Pattern sets ─────────────────────────────────────────────────────────────
48
49fn decision_patterns() -> &'static [Regex] {
50    static P: OnceLock<Vec<Regex>> = OnceLock::new();
51    P.get_or_init(|| {
52        compile_patterns(&[
53            r"(?i)\blet'?s (use|go with|try|pick|choose|switch to)\b",
54            r"(?i)\bwe (should|decided|chose|went with|picked|settled on)\b",
55            r"(?i)\bi'?m going (to|with)\b",
56            r"(?i)\bbetter (to|than|approach|option|choice)\b",
57            r"(?i)\binstead of\b",
58            r"(?i)\brather than\b",
59            r"(?i)\btrade-?off\b",
60            r"决定",
61            r"选择了",
62            r"决定使用",
63            r"決めました",
64            r"選びました",
65            r"にしました",
66        ])
67    })
68}
69
70fn preference_patterns() -> &'static [Regex] {
71    static P: OnceLock<Vec<Regex>> = OnceLock::new();
72    P.get_or_init(|| {
73        compile_patterns(&[
74            r"(?i)\bi prefer\b",
75            r"(?i)\balways use\b",
76            r"(?i)\bnever use\b",
77            r"(?i)\bdon'?t (ever |like to )?(use|do|mock|stub|import)\b",
78            r"(?i)\bi like (to|when|how)\b",
79            r"(?i)\bi hate (when|how|it when)\b",
80            r"(?i)\bstop using\b",
81            r"我喜欢",
82            r"我中意",
83            r"我鍾意",
84            r"不喜欢",
85            r"唔钟意",
86            r"讨厌",
87            r"が好き",
88            r"が嫌い",
89        ])
90    })
91}
92
93fn milestone_patterns() -> &'static [Regex] {
94    static P: OnceLock<Vec<Regex>> = OnceLock::new();
95    P.get_or_init(|| {
96        compile_patterns(&[
97            r"(?i)\bfinally (works?|working|got it)\b",
98            r"(?i)\bit'?s (working|alive|done)\b",
99            r"(?i)\bbreakthrough\b",
100            r"(?i)\bcracked it\b",
101            r"(?i)\bwe did it\b",
102            r"(?i)\bcompleted\b",
103            r"(?i)\bshipped\b",
104            r"(?i)\blaunched\b",
105            r"(?i)\bsuccessfully\b",
106            r"(?i)\bfixed it\b",
107            r"终于搞定了",
108            r"终于好了",
109            r"成功しました",
110            r"できました",
111        ])
112    })
113}
114
115fn problem_patterns() -> &'static [Regex] {
116    static P: OnceLock<Vec<Regex>> = OnceLock::new();
117    P.get_or_init(|| {
118        compile_patterns(&[
119            r"(?i)\bbug\b",
120            r"(?i)\berror\b",
121            r"(?i)\bcrash\b",
122            r"(?i)\bfailed\b",
123            r"(?i)\bbroken\b",
124            r"(?i)\bissue\b",
125            r"(?i)\bproblem\b",
126            r"(?i)\broot cause\b",
127            r"(?i)\bfix\b",
128            r"(?i)\bworkaround\b",
129            r"(?i)\bdebugging\b",
130            r"报错",
131            r"崩溃",
132            r"失败",
133            r"バグ",
134            r"エラー",
135            r"失敗",
136        ])
137    })
138}
139
140fn emotional_patterns() -> &'static [Regex] {
141    static P: OnceLock<Vec<Regex>> = OnceLock::new();
142    P.get_or_init(|| {
143        compile_patterns(&[
144            r"(?i)\bscared\b",
145            r"(?i)\bafraid\b",
146            r"(?i)\bworried\b",
147            r"(?i)\bhappy\b",
148            r"(?i)\bsad\b",
149            r"(?i)\blove\b",
150            r"(?i)\bcry\b",
151            r"(?i)\btears\b",
152            r"(?i)\bvulnerable\b",
153            r"(?i)\bgrateful\b",
154            r"开心",
155            r"难过",
156            r"嬉しい",
157            r"悲しい",
158            r"心配",
159        ])
160    })
161}
162
163fn compile_patterns(pats: &[&str]) -> Vec<Regex> {
164    pats.iter().filter_map(|p| Regex::new(p).ok()).collect()
165}
166
167// ── Classifier ────────────────────────────────────────────────────────────────
168
169fn classify(text: &str) -> MemoryType {
170    let score = |pats: &[Regex]| -> usize { pats.iter().filter(|p| p.is_match(text)).count() };
171
172    let d = score(decision_patterns());
173    let p = score(preference_patterns());
174    let m = score(milestone_patterns());
175    let pr = score(problem_patterns());
176    let e = score(emotional_patterns());
177
178    let max = d.max(p).max(m).max(pr).max(e);
179    if max == 0 {
180        return MemoryType::General;
181    }
182
183    if e == max {
184        MemoryType::Emotional
185    } else if m == max {
186        MemoryType::Milestone
187    } else if d == max {
188        MemoryType::Decision
189    } else if p == max {
190        MemoryType::Preference
191    } else {
192        MemoryType::Problem
193    }
194}
195
196// ── Public API ────────────────────────────────────────────────────────────────
197
198/// Extract typed memories from a text block.
199///
200/// Returns only chunks whose type is *not* `General` (i.e., matches at least
201/// one typed pattern).  Use `extract_all` to include general chunks.
202pub fn extract_memories(text: &str) -> Vec<ExtractedMemory> {
203    extract_all(text)
204        .into_iter()
205        .filter(|m| m.memory_type != MemoryType::General)
206        .collect()
207}
208
209/// Extract and classify all paragraph-level chunks.
210pub fn extract_all(text: &str) -> Vec<ExtractedMemory> {
211    text.split("\n\n")
212        .enumerate()
213        .filter_map(|(i, chunk)| {
214            let chunk = chunk.trim();
215            if chunk.len() < 30 {
216                return None;
217            }
218            Some(ExtractedMemory {
219                content: chunk.to_string(),
220                memory_type: classify(chunk),
221                chunk_index: i,
222            })
223        })
224        .collect()
225}
226
227#[cfg(test)]
228mod tests {
229    use super::*;
230
231    #[test]
232    fn test_classify_decision() {
233        let t = "We decided to go with Rust instead of Python because of performance.";
234        assert_eq!(classify(t), MemoryType::Decision);
235    }
236
237    #[test]
238    fn test_classify_preference() {
239        let t = "I prefer to always use turso for local databases.";
240        assert_eq!(classify(t), MemoryType::Preference);
241    }
242
243    #[test]
244    fn test_classify_milestone() {
245        let t = "Finally works! The vector search is returning correct results.";
246        assert_eq!(classify(t), MemoryType::Milestone);
247    }
248
249    #[test]
250    fn test_classify_problem() {
251        let t = "There was a bug in the embedding generation that caused crashes.";
252        assert_eq!(classify(t), MemoryType::Problem);
253    }
254
255    #[test]
256    fn test_does_not_false_positive_on_generic_architecture_words() {
257        let t = "This architecture document lists several possible approaches for future work.";
258        assert_eq!(classify(t), MemoryType::General);
259    }
260
261    #[test]
262    fn test_does_not_false_positive_on_generic_i_am_statement() {
263        let t = "I am working on the backend service this week.";
264        assert_eq!(classify(t), MemoryType::General);
265    }
266
267    #[test]
268    fn test_does_not_false_positive_on_generic_i_want_statement() {
269        let t = "I want to understand the system before changing anything.";
270        assert_eq!(classify(t), MemoryType::General);
271    }
272
273    #[test]
274    fn test_classify_multilingual_preference() {
275        let t = "我喜欢寿司，也喜欢拉面。";
276        assert_eq!(classify(t), MemoryType::Preference);
277    }
278
279    #[test]
280    fn test_classify_multilingual_problem() {
281        let t = "服务启动时报错，然后直接崩溃。";
282        assert_eq!(classify(t), MemoryType::Problem);
283    }
284
285    #[test]
286    fn test_classify_multilingual_decision() {
287        let t = "我们决定使用 Rust 重写后端。";
288        assert_eq!(classify(t), MemoryType::Decision);
289    }
290
291    #[test]
292    fn test_extract_all_keeps_paragraph_chunking_for_short_fragments() {
293        let t = "我喜欢寿司。\n\n下雨了。";
294        let memories = extract_all(t);
295        assert_eq!(memories.len(), 0);
296    }
297}
aimem_core/extractor.rs

aimem_core/
extractor.rs
