Skip to main content

mentedb_consolidation/
compression.rs

1use mentedb_core::MemoryNode;
2use mentedb_core::types::MemoryId;
3use serde::{Deserialize, Serialize};
4
5/// A compressed representation of a memory.
6#[derive(Debug, Clone, Serialize, Deserialize)]
7pub struct CompressedMemory {
8    pub original_id: MemoryId,
9    pub compressed_content: String,
10    pub compression_ratio: f32,
11    pub key_facts: Vec<String>,
12}
13
14const KEY_WORDS: &[&str] = &[
15    "decided",
16    "uses",
17    "prefers",
18    "switched",
19    "chose",
20    "selected",
21    "important",
22    "critical",
23    "must",
24    "should",
25    "always",
26    "never",
27];
28
29const FILLER_WORDS: &[&str] = &[
30    "actually",
31    "basically",
32    "honestly",
33    "really",
34    "very",
35    "quite",
36    "just",
37    "simply",
38    "perhaps",
39    "maybe",
40    "somewhat",
41    "rather",
42    "kind of",
43    "sort of",
44    "you know",
45    "i think",
46    "i mean",
47];
48
49/// Compresses verbose memories into token-efficient forms.
50pub struct MemoryCompressor;
51
52impl MemoryCompressor {
53    pub fn new() -> Self {
54        Self
55    }
56
57    /// Compress a single memory by extracting key sentences and removing filler.
58    pub fn compress(&self, memory: &MemoryNode) -> CompressedMemory {
59        let paragraphs: Vec<&str> = memory.content.split("\n\n").collect();
60        let mut key_sentences: Vec<String> = Vec::new();
61
62        for para in &paragraphs {
63            let sentences: Vec<&str> = para
64                .split('.')
65                .map(|s| s.trim())
66                .filter(|s| !s.is_empty())
67                .collect();
68
69            // Always include first sentence of each paragraph
70            if let Some(&first) = sentences.first() {
71                let cleaned = remove_filler(first);
72                if !cleaned.is_empty() && !key_sentences.contains(&cleaned) {
73                    key_sentences.push(cleaned);
74                }
75            }
76
77            // Include sentences with keywords
78            for &sentence in &sentences[1..] {
79                let lower = sentence.to_lowercase();
80                if KEY_WORDS.iter().any(|kw| lower.contains(kw)) {
81                    let cleaned = remove_filler(sentence);
82                    if !cleaned.is_empty() && !key_sentences.contains(&cleaned) {
83                        key_sentences.push(cleaned);
84                    }
85                }
86            }
87        }
88
89        let compressed_content = key_sentences.join(". ");
90        let original_len = memory.content.len().max(1) as f32;
91        let compressed_len = compressed_content.len() as f32;
92        let compression_ratio = compressed_len / original_len;
93
94        CompressedMemory {
95            original_id: memory.id,
96            compressed_content,
97            compression_ratio,
98            key_facts: key_sentences,
99        }
100    }
101
102    /// Compress a batch of memories.
103    pub fn compress_batch(&self, memories: &[MemoryNode]) -> Vec<CompressedMemory> {
104        memories.iter().map(|m| self.compress(m)).collect()
105    }
106
107    /// Estimate token count for text (word_count * 1.3).
108    pub fn estimate_tokens(text: &str) -> usize {
109        let word_count = text.split_whitespace().count();
110        (word_count as f64 * 1.3).ceil() as usize
111    }
112}
113
114impl Default for MemoryCompressor {
115    fn default() -> Self {
116        Self::new()
117    }
118}
119
120/// Remove filler words from text.
121fn remove_filler(text: &str) -> String {
122    let mut result = text.to_string();
123    for &filler in FILLER_WORDS {
124        // Case-insensitive removal
125        let lower = result.to_lowercase();
126        if let Some(pos) = lower.find(filler) {
127            let end = pos + filler.len();
128            let actual = &result[pos..end];
129            result = result.replacen(actual, "", 1);
130        }
131    }
132    // Collapse whitespace
133    result.split_whitespace().collect::<Vec<_>>().join(" ")
134}
135
136#[cfg(test)]
137mod tests {
138    use super::*;
139    use crate::test_helpers::make_memory;
140
141    #[test]
142    fn test_compress_extracts_key_sentences() {
143        let compressor = MemoryCompressor::new();
144        let m = make_memory(
145            "The team met today. They discussed various topics. They decided on Rust for the backend. It was a long meeting.",
146            vec![1.0],
147        );
148        let compressed = compressor.compress(&m);
149        assert!(compressed.compressed_content.contains("decided"));
150        assert!(compressed.compression_ratio <= 1.0);
151    }
152
153    #[test]
154    fn test_compress_removes_filler() {
155        let compressor = MemoryCompressor::new();
156        let m = make_memory("We basically decided to actually use Postgres", vec![1.0]);
157        let compressed = compressor.compress(&m);
158        assert!(!compressed.compressed_content.contains("basically"));
159    }
160
161    #[test]
162    fn test_estimate_tokens() {
163        assert_eq!(MemoryCompressor::estimate_tokens("hello world"), 3); // 2 * 1.3 = 2.6 -> 3
164        assert_eq!(MemoryCompressor::estimate_tokens(""), 0);
165    }
166
167    #[test]
168    fn test_compress_batch() {
169        let compressor = MemoryCompressor::new();
170        let memories = vec![
171            make_memory("First memory content", vec![1.0]),
172            make_memory("Second memory content", vec![1.0]),
173        ];
174        let results = compressor.compress_batch(&memories);
175        assert_eq!(results.len(), 2);
176    }
177}