Skip to main content

mneme/
compress.rs

1use std::collections::HashSet;
2
3use serde::{Deserialize, Serialize};
4
5use crate::store::memory::Memory;
6
7/// Estrategia de compresión.
8#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)]
9#[serde(rename_all = "snake_case")]
10pub enum CompressionStrategy {
11    /// Truncación simple: solo primeros N chars.
12    Truncate,
13    /// Resumen inteligente: primer párrafo + oraciones clave.
14    SmartSummary,
15    /// Solo extraer keywords del contenido.
16    KeywordsOnly,
17    /// Mínimo: solo título + tipo + keywords.
18    Minimal,
19}
20
21impl std::fmt::Display for CompressionStrategy {
22    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
23        let s = match self {
24            CompressionStrategy::Truncate => "truncate",
25            CompressionStrategy::SmartSummary => "smart_summary",
26            CompressionStrategy::KeywordsOnly => "keywords_only",
27            CompressionStrategy::Minimal => "minimal",
28        };
29        write!(f, "{}", s)
30    }
31}
32
33impl std::str::FromStr for CompressionStrategy {
34    type Err = crate::error::MnemeError;
35    fn from_str(s: &str) -> Result<Self, Self::Err> {
36        match s.to_lowercase().as_str() {
37            "truncate" => Ok(CompressionStrategy::Truncate),
38            "smart_summary" | "smart-summary" | "smartsummary" => {
39                Ok(CompressionStrategy::SmartSummary)
40            }
41            "keywords_only" | "keywords-only" | "keywordsonly" => {
42                Ok(CompressionStrategy::KeywordsOnly)
43            }
44            "minimal" => Ok(CompressionStrategy::Minimal),
45            other => Err(crate::error::MnemeError::Config(format!(
46                "Invalid compression strategy: {}",
47                other
48            ))),
49        }
50    }
51}
52
53/// Resultado de la compresión.
54#[derive(Debug, Clone, Serialize, Deserialize)]
55pub struct CompressedMemory {
56    /// ID de la memoria original.
57    pub memory_id: String,
58    /// Título original (sin comprimir).
59    pub title: String,
60    /// Contenido comprimido.
61    pub compressed_content: String,
62    /// Estrategia utilizada.
63    pub strategy: String,
64    /// Ratio de compresión (0.0 - 1.0), donde 1.0 = 100% reducción.
65    pub compression_ratio: f64,
66    /// Keywords extraídas (todas las estrategias).
67    pub keywords: Vec<String>,
68    /// ¿Se puede revertir?
69    pub reversible: bool,
70}
71
72/// Pipeline de compresión reversible para contenido de memorias.
73/// Inspirado por Headroom (CCR — Context Compression & Retrieval).
74pub struct CompressionPipeline;
75
76impl CompressionPipeline {
77    /// Comprime una memoria usando la estrategia especificada.
78    pub fn compress(memory: &Memory, strategy: CompressionStrategy) -> CompressedMemory {
79        let original_len = memory.content.len() as f64;
80
81        let (compressed_content, keywords) = match strategy {
82            CompressionStrategy::Truncate => Self::truncate_compress(&memory.content, 200),
83            CompressionStrategy::SmartSummary => Self::smart_summary(&memory.content),
84            CompressionStrategy::KeywordsOnly => Self::keywords_only(&memory.content),
85            CompressionStrategy::Minimal => Self::minimal_compress(memory),
86        };
87
88        let compressed_len = compressed_content.len() as f64;
89        let compression_ratio = if original_len > 0.0 {
90            1.0 - (compressed_len / original_len)
91        } else {
92            0.0
93        };
94
95        CompressedMemory {
96            memory_id: memory.id.to_string(),
97            title: memory.title.clone(),
98            compressed_content,
99            strategy: strategy.to_string(),
100            compression_ratio,
101            keywords,
102            reversible: true,
103        }
104    }
105
106    /// Truncación simple a N caracteres.
107    fn truncate_compress(content: &str, max_chars: usize) -> (String, Vec<String>) {
108        if content.len() <= max_chars {
109            let keywords = Self::extract_keywords(content);
110            return (content.to_string(), keywords);
111        }
112
113        let mut result = String::with_capacity(max_chars + 3);
114        // Keep first `max_chars` chars
115        result.push_str(&content[..max_chars]);
116        result.push_str("...");
117
118        let keywords = Self::extract_keywords(&content[..max_chars]);
119        (result, keywords)
120    }
121
122    /// Resumen inteligente: primer párrafo + oraciones clave.
123    fn smart_summary(content: &str) -> (String, Vec<String>) {
124        let mut parts: Vec<String> = Vec::new();
125
126        // First paragraph
127        if let Some(first_para) = content.split("\n\n").next() {
128            if !first_para.is_empty() {
129                parts.push(format!("[Intro] {}", first_para));
130            }
131        }
132
133        // Extract key sentences (sentences with important keywords or patterns)
134        let sentences: Vec<&str> = content
135            .split(['.', '!', '?'])
136            .map(|s| s.trim())
137            .filter(|s| !s.is_empty())
138            .collect();
139
140        let important_patterns = [
141            "important",
142            "critical",
143            "key",
144            "must",
145            "should",
146            "never",
147            "always",
148            "architecture",
149            "decision",
150            "chose",
151            "selected",
152            "implemented",
153            "importante",
154            "crítico",
155            "clave",
156            "debe",
157            "nunca",
158            "siempre",
159            "importante",
160            "decisión",
161            "arquitectura",
162            "seleccionó",
163        ];
164
165        let mut key_sentences: Vec<&str> = sentences
166            .iter()
167            .filter(|s| {
168                let lower = s.to_lowercase();
169                important_patterns.iter().any(|p| lower.contains(p))
170            })
171            .copied()
172            .collect();
173
174        // Deduplicate and limit
175        key_sentences.sort();
176        key_sentences.dedup();
177        let max_sentences = 5.min(key_sentences.len());
178
179        if max_sentences > 0 {
180            parts.push("[Key points]".to_string());
181            for sent in key_sentences.iter().take(max_sentences) {
182                parts.push(format!("- {}", sent.trim()));
183            }
184        }
185
186        let keywords = Self::extract_keywords(content);
187        let result = parts.join("\n");
188        (result, keywords)
189    }
190
191    /// Solo extraer keywords.
192    fn keywords_only(content: &str) -> (String, Vec<String>) {
193        let keywords = Self::extract_keywords(content);
194        (format!("Keywords: {}", keywords.join(", ")), keywords)
195    }
196
197    /// Mínimo: título + tipo + keywords.
198    fn minimal_compress(memory: &Memory) -> (String, Vec<String>) {
199        let keywords = Self::extract_keywords(&memory.content);
200        let mut parts = vec![format!("[{}] {}", memory.memory_type, memory.title)];
201
202        if let Some(ref what) = memory.what {
203            let truncated: String = what.chars().take(100).collect();
204            parts.push(format!("What: {}", truncated));
205        }
206
207        if !keywords.is_empty() {
208            parts.push(format!("Keywords: {}", keywords.join(", ")));
209        }
210
211        (parts.join(" | "), keywords)
212    }
213
214    /// Extrae keywords relevantes del contenido.
215    pub fn extract_keywords(content: &str) -> Vec<String> {
216        let stopwords: HashSet<&str> = HashSet::from([
217            "the", "a", "an", "is", "are", "was", "were", "be", "been", "being", "have", "has",
218            "had", "do", "does", "did", "will", "would", "could", "should", "may", "might",
219            "shall", "can", "need", "dare", "ought", "i", "you", "he", "she", "it", "we", "they",
220            "me", "him", "her", "us", "them", "my", "your", "his", "its", "our", "their", "this",
221            "that", "these", "those", "some", "any", "each", "every", "all", "both", "few",
222            "several", "many", "much", "no", "not", "only", "own", "same", "so", "than", "too",
223            "very", "just", "because", "as", "until", "while", "of", "at", "by", "for", "with",
224            "about", "against", "between", "into", "through", "during", "before", "after", "above",
225            "below", "to", "from", "up", "down", "in", "out", "on", "off", "over", "under",
226            "again", "further", "then", "once", "here", "there", "when", "where", "why", "how",
227            "el", "la", "los", "las", "un", "una", "y", "e", "o", "u", "de", "del", "en", "al",
228            "por", "para", "con", "sin", "sobre", "entre", "como", "que", "es", "se", "su", "lo",
229            "le", "ha", "está", "esta", "este", "ese", "eso", "era", "ser", "han",
230        ]);
231
232        let mut word_counts: std::collections::HashMap<String, u32> =
233            std::collections::HashMap::new();
234
235        // Normalize and split
236        for word in content.split(|c: char| !c.is_alphanumeric() && c != '-' && c != '_') {
237            let lower = word.trim().to_lowercase();
238            if lower.len() > 3
239                && !stopwords.contains(lower.as_str())
240                && !lower.chars().all(|c| c.is_numeric())
241            {
242                *word_counts.entry(lower).or_insert(0) += 1;
243            }
244        }
245
246        // Also check for CamelCase/PascalCase identifiers
247        for word in content.split_whitespace() {
248            if word.len() > 4 {
249                let has_upper = word.chars().any(|c| c.is_uppercase());
250                let has_lower = word.chars().any(|c| c.is_lowercase());
251                if has_upper && has_lower && !word.contains('_') {
252                    let lower = word.to_lowercase();
253                    if lower.len() > 3 && !stopwords.contains(lower.as_str()) {
254                        *word_counts.entry(lower).or_insert(0) += 2;
255                    }
256                }
257            }
258        }
259
260        // Sort by frequency
261        let mut sorted: Vec<(String, u32)> = word_counts.into_iter().collect();
262        sorted.sort_by(|a, b| b.1.cmp(&a.1));
263
264        sorted.into_iter().take(10).map(|(w, _)| w).collect()
265    }
266
267    /// Genera un bloque de contexto comprimido para inyección en prompts.
268    /// Versión comprimida de `MemoryStore::inject_context`.
269    pub fn compress_context_block(
270        memories: &[Memory],
271        strategy: CompressionStrategy,
272        max_memories: usize,
273    ) -> String {
274        let mut lines = vec![
275            "## Contexto comprimido del proyecto".to_string(),
276            String::new(),
277        ];
278
279        for memory in memories.iter().take(max_memories) {
280            let compressed = Self::compress(memory, strategy);
281            lines.push(format!(
282                "- **{}** [{}] ({}): {}",
283                compressed.title,
284                memory.memory_type,
285                compressed.strategy,
286                compressed
287                    .compressed_content
288                    .chars()
289                    .take(150)
290                    .collect::<String>()
291            ));
292        }
293
294        lines.push(String::new());
295        lines.push(format!(
296            "_Contexto comprimido con estrategia '{}'. Usa mem_expand para ver contenido completo._",
297            strategy
298        ));
299
300        lines.join("\n")
301    }
302}
303
304#[cfg(test)]
305mod tests {
306    use super::*;
307    use crate::store::memory::{Importance, Scope};
308
309    #[test]
310    fn test_truncate_compress_short() {
311        let (compressed, keywords) = CompressionPipeline::truncate_compress("Short text", 200);
312        assert_eq!(compressed, "Short text");
313    }
314
315    #[test]
316    fn test_truncate_compress_long() {
317        let long = "A".repeat(500);
318        let (compressed, _) = CompressionPipeline::truncate_compress(&long, 200);
319        assert!(compressed.len() < 210);
320        assert!(compressed.ends_with("..."));
321    }
322
323    #[test]
324    fn test_extract_keywords() {
325        let text = "rust is a systems programming language focused on safety and performance";
326        let keywords = CompressionPipeline::extract_keywords(text);
327        assert!(keywords.contains(&"rust".to_string()));
328        assert!(keywords.contains(&"systems".to_string()));
329        assert!(keywords.contains(&"programming".to_string()));
330    }
331
332    #[test]
333    fn test_minimal_compress() {
334        let memory = Memory {
335            id: uuid::Uuid::new_v4(),
336            project: "test".to_string(),
337            scope: crate::store::memory::Scope::Project,
338            title: "Test Memory".to_string(),
339            content: "This is a test content with some important keywords for testing purpose"
340                .to_string(),
341            what: Some("What was done".to_string()),
342            why: None,
343            context: None,
344            learned: None,
345            memory_type: crate::store::memory::MemoryType::Architecture,
346            importance: crate::store::memory::Importance::High,
347            tags: vec![],
348            topic_key: None,
349            access_count: 0,
350            revision_count: 0,
351            duplicate_count: 0,
352            normalized_hash: None,
353            created_at: chrono::DateTime::UNIX_EPOCH,
354            updated_at: chrono::DateTime::UNIX_EPOCH,
355            last_accessed_at: None,
356            last_seen_at: None,
357            deleted_at: None,
358            deprecated_at: None,
359            deprecated_reason: None,
360            supersedes_id: None,
361            context_inject_count: 0,
362            origin_peer: None,
363            is_encrypted: false,
364            encrypted_for: None,
365            valid_from: None,
366            valid_until: None,
367            provenance: None,
368        };
369
370        let result = CompressionPipeline::compress(&memory, CompressionStrategy::Minimal);
371        assert!(result.compressed_content.contains("architecture"));
372        assert!(result.compressed_content.contains("Test Memory"));
373        assert_eq!(result.reversible, true);
374    }
375}