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
10use regex::Regex;
11use std::sync::OnceLock;
12
13#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
14pub enum MemoryType {
15    Decision,
16    Preference,
17    Milestone,
18    Problem,
19    Emotional,
20    General,
21}
22
23impl std::fmt::Display for MemoryType {
24    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
25        let s = match self {
26            Self::Decision => "decision",
27            Self::Preference => "preference",
28            Self::Milestone => "milestone",
29            Self::Problem => "problem",
30            Self::Emotional => "emotional",
31            Self::General => "general",
32        };
33        write!(f, "{s}")
34    }
35}
36
37#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
38pub struct ExtractedMemory {
39    pub content: String,
40    pub memory_type: MemoryType,
41    pub chunk_index: usize,
42}
43
44// ── Pattern sets ─────────────────────────────────────────────────────────────
45
46fn decision_patterns() -> &'static [Regex] {
47    static P: OnceLock<Vec<Regex>> = OnceLock::new();
48    P.get_or_init(|| {
49        compile_patterns(&[
50            r"(?i)\blet'?s (use|go with|try|pick|choose|switch to)\b",
51            r"(?i)\bwe (should|decided|chose|went with|picked|settled on)\b",
52            r"(?i)\bi'?m going (to|with)\b",
53            r"(?i)\bbetter (to|than|approach|option|choice)\b",
54            r"(?i)\binstead of\b",
55            r"(?i)\brather than\b",
56            r"(?i)\bthe reason (is|was|being)\b",
57            r"(?i)\btrade-?off\b",
58            r"(?i)\barchitecture\b",
59            r"(?i)\bapproach\b",
60            r"(?i)\bstrategy\b",
61        ])
62    })
63}
64
65fn preference_patterns() -> &'static [Regex] {
66    static P: OnceLock<Vec<Regex>> = OnceLock::new();
67    P.get_or_init(|| {
68        compile_patterns(&[
69            r"(?i)\bi prefer\b",
70            r"(?i)\balways use\b",
71            r"(?i)\bnever use\b",
72            r"(?i)\bdon'?t (ever |like to )?(use|do|mock|stub|import)\b",
73            r"(?i)\bi like (to|when|how)\b",
74            r"(?i)\bi hate (when|how|it when)\b",
75            r"(?i)\bmy preference\b",
76            r"(?i)\bi want\b",
77            r"(?i)\bstop using\b",
78        ])
79    })
80}
81
82fn milestone_patterns() -> &'static [Regex] {
83    static P: OnceLock<Vec<Regex>> = OnceLock::new();
84    P.get_or_init(|| {
85        compile_patterns(&[
86            r"(?i)\bfinally (works?|working|got it)\b",
87            r"(?i)\bit'?s (working|alive|done)\b",
88            r"(?i)\bbreakthrough\b",
89            r"(?i)\bcracked it\b",
90            r"(?i)\bwe did it\b",
91            r"(?i)\bcompleted\b",
92            r"(?i)\bshipped\b",
93            r"(?i)\blaunched\b",
94            r"(?i)\bsuccessfully\b",
95            r"(?i)\bfixed it\b",
96        ])
97    })
98}
99
100fn problem_patterns() -> &'static [Regex] {
101    static P: OnceLock<Vec<Regex>> = OnceLock::new();
102    P.get_or_init(|| {
103        compile_patterns(&[
104            r"(?i)\bbug\b",
105            r"(?i)\berror\b",
106            r"(?i)\bcrash\b",
107            r"(?i)\bfailed\b",
108            r"(?i)\bbroken\b",
109            r"(?i)\bissue\b",
110            r"(?i)\bproblem\b",
111            r"(?i)\broot cause\b",
112            r"(?i)\bfix\b",
113            r"(?i)\bworkaround\b",
114            r"(?i)\bdebugging\b",
115        ])
116    })
117}
118
119fn emotional_patterns() -> &'static [Regex] {
120    static P: OnceLock<Vec<Regex>> = OnceLock::new();
121    P.get_or_init(|| {
122        compile_patterns(&[
123            r"(?i)\bscared\b",
124            r"(?i)\bafraid\b",
125            r"(?i)\bworried\b",
126            r"(?i)\bhappy\b",
127            r"(?i)\bsad\b",
128            r"(?i)\blove\b",
129            r"(?i)\bfeel\b",
130            r"(?i)\bcry\b",
131            r"(?i)\btears\b",
132            r"(?i)\bvulnerable\b",
133            r"(?i)\bgrateful\b",
134        ])
135    })
136}
137
138fn compile_patterns(pats: &[&str]) -> Vec<Regex> {
139    pats.iter().filter_map(|p| Regex::new(p).ok()).collect()
140}
141
142// ── Classifier ────────────────────────────────────────────────────────────────
143
144fn classify(text: &str) -> MemoryType {
145    let score = |pats: &[Regex]| -> usize { pats.iter().filter(|p| p.is_match(text)).count() };
146
147    let d = score(decision_patterns());
148    let p = score(preference_patterns());
149    let m = score(milestone_patterns());
150    let pr = score(problem_patterns());
151    let e = score(emotional_patterns());
152
153    let max = d.max(p).max(m).max(pr).max(e);
154    if max == 0 {
155        return MemoryType::General;
156    }
157
158    if e == max {
159        MemoryType::Emotional
160    } else if m == max {
161        MemoryType::Milestone
162    } else if d == max {
163        MemoryType::Decision
164    } else if p == max {
165        MemoryType::Preference
166    } else {
167        MemoryType::Problem
168    }
169}
170
171// ── Public API ────────────────────────────────────────────────────────────────
172
173/// Extract typed memories from a text block.
174///
175/// Returns only chunks whose type is *not* `General` (i.e., matches at least
176/// one typed pattern).  Use `extract_all` to include general chunks.
177pub fn extract_memories(text: &str) -> Vec<ExtractedMemory> {
178    extract_all(text)
179        .into_iter()
180        .filter(|m| m.memory_type != MemoryType::General)
181        .collect()
182}
183
184/// Extract and classify all paragraph-level chunks.
185pub fn extract_all(text: &str) -> Vec<ExtractedMemory> {
186    text.split("\n\n")
187        .enumerate()
188        .filter_map(|(i, chunk)| {
189            let chunk = chunk.trim();
190            if chunk.len() < 30 {
191                return None;
192            }
193            Some(ExtractedMemory {
194                content: chunk.to_string(),
195                memory_type: classify(chunk),
196                chunk_index: i,
197            })
198        })
199        .collect()
200}
201
202#[cfg(test)]
203mod tests {
204    use super::*;
205
206    #[test]
207    fn test_classify_decision() {
208        let t = "We decided to go with Rust instead of Python because of performance.";
209        assert_eq!(classify(t), MemoryType::Decision);
210    }
211
212    #[test]
213    fn test_classify_preference() {
214        let t = "I prefer to always use turso for local databases.";
215        assert_eq!(classify(t), MemoryType::Preference);
216    }
217
218    #[test]
219    fn test_classify_milestone() {
220        let t = "Finally works! The vector search is returning correct results.";
221        assert_eq!(classify(t), MemoryType::Milestone);
222    }
223
224    #[test]
225    fn test_classify_problem() {
226        let t = "There was a bug in the embedding generation that caused crashes.";
227        assert_eq!(classify(t), MemoryType::Problem);
228    }
229}
aimem_core/extractor.rs

aimem_core/
extractor.rs
