poimen-memory/crates/mem-core/src/optimizer/text.rs

//! TextCompressor — Token importance scoring for plain text
//!
//! Strategy:
//! - Keep: high-entropy tokens (IDs, hashes, error codes, numbers, symbols)
//! - Drop: low-information prose (filler words, common phrases)
//! - Reuses M3.7.8 stop words list for detection

use lazy_static::lazy_static;
use std::collections::HashSet;

lazy_static! {
    /// Common low-information words (expanded from M3.7.8 stop words)
    static ref STOP_WORDS: HashSet<&'static str> = {
        let words = vec![
            // Common articles, prepositions, conjunctions
            "the", "a", "an", "and", "or", "but", "in", "on", "at", "to", "for",
            "of", "with", "by", "from", "is", "are", "was", "were", "be", "been",
            "have", "has", "do", "does", "did", "will", "would", "could", "should",
            "may", "might", "must", "can", "shall",
            
            // Filler words
            "the", "this", "that", "these", "those", "it", "its", "as", "also",
            "very", "just", "only", "even", "still", "again", "about", "while",
            "where", "when", "why", "what", "which", "who", "whom",
            
            // Common verbs (low signal when standalone)
            "go", "get", "put", "make", "take", "come", "see", "say", "know",
            "think", "want", "use", "find", "give", "tell", "work", "call",
            "try", "ask", "need", "feel", "become", "leave", "show",
            
            // Weak modifiers
            "good", "bad", "new", "old", "big", "small", "first", "last",
            "some", "any", "no", "own", "other", "more", "most", "less", "least",
            
            // Pronouns
            "i", "me", "we", "us", "you", "he", "him", "she", "her", "they", "them",
            "my", "our", "your", "his", "her", "their",
            
            // Numbers and common phrases (context-dependent, lower priority)
            "one", "two", "three", "four", "five", "etc", "etc.",
        ];
        words.into_iter().collect()
    };
}

pub struct TextCompressor;

impl TextCompressor {
    pub fn new() -> Self {
        Self
    }

    /// Compress text by scoring tokens and keeping high-entropy ones
    pub fn compress(&self, content: &str, target_ratio: f32) -> String {
        if content.is_empty() {
            return content.to_string();
        }

        let tokens: Vec<&str> = content.split_whitespace().collect();
        if tokens.is_empty() {
            return content.to_string();
        }

        // Score each token
        let mut scored: Vec<(usize, &str, f32)> = tokens
            .iter()
            .enumerate()
            .map(|(idx, token)| (idx, *token, Self::score_token(token)))
            .collect();

        // Calculate target count
        let target_count = ((tokens.len() as f32) * target_ratio).ceil() as usize;
        let target_count = target_count.max(1).min(tokens.len()); // At least 1, at most all

        // Sort by score descending, then by original index to preserve order
        scored.sort_by(|a, b| {
            b.2.partial_cmp(&a.2)
                .unwrap_or(std::cmp::Ordering::Equal)
                .then_with(|| a.0.cmp(&b.0))
        });

        // Take top tokens and re-sort by original index
        let mut kept: Vec<(usize, &str)> = scored
            .into_iter()
            .take(target_count)
            .map(|(idx, token, _)| (idx, token))
            .collect();
        kept.sort_by_key(|a| a.0);

        kept.iter().map(|(_, token)| *token).collect::<Vec<_>>().join(" ")
    }

    /// Score a token for importance
    fn score_token(token: &str) -> f32 {
        let lower = token.to_lowercase();

        // High-entropy tokens
        let mut score = 0.0f32;

        // IDs, hashes, hex
        if is_id_like(token) {
            score += 10.0;
        }

        // Numbers
        if token.chars().any(|c| c.is_numeric()) {
            score += 3.0;
        }

        // Error codes, markers
        if is_error_marker(&lower) {
            score += 8.0;
        }

        // Symbols (punctuation often marks structure)
        if token.chars().any(|c| !c.is_alphanumeric()) {
            score += 2.0;
        }

        // Stop words (negative score)
        if STOP_WORDS.contains(lower.as_str()) {
            score -= 5.0;
        }

        // Length (longer tokens usually more informative)
        if token.len() > 10 {
            score += 1.0;
        }

        // Capitalization (usually proper nouns or emphatic)
        if token.chars().next().is_some_and(|c| c.is_uppercase()) && token.len() > 1 {
            score += 1.0;
        }

        score
    }
}

impl Default for TextCompressor {
    fn default() -> Self {
        Self::new()
    }
}

/// Check if token looks like an ID (UUID, hash, etc.)
fn is_id_like(token: &str) -> bool {
    // UUIDs
    if token.len() == 36 && token.matches('-').count() == 4 {
        return true;
    }

    // SHA256 / hashes
    if token.len() == 64 && token.chars().all(|c| c.is_ascii_hexdigit()) {
        return true;
    }

    // Short hex strings
    if token.len() > 8 && token.len() < 20 && token.chars().all(|c| c.is_ascii_hexdigit()) {
        return true;
    }

    // Alphanumeric with underscores (typical ID pattern)
    if token.len() > 6 && token.contains('_') && token.chars().all(|c| c.is_alphanumeric() || c == '_') {
        return true;
    }

    false
}

/// Check if token is an error marker
fn is_error_marker(token: &str) -> bool {
    token.contains("error")
        || token.contains("err")
        || token.contains("fail")
        || token.contains("exception")
        || token.contains("panic")
        || token.contains("warn")
        || token.contains("critical")
        || token.contains("fatal")
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_score_high_entropy_tokens() {
        let uuid = "550e8400-e29b-41d4-a716-446655440000";
        let hash = "abc123def456abc123def456abc123def456abc123def456abc123def456abc1";
        let error = "ConnectionError";

        assert!(TextCompressor::score_token(uuid) > 5.0);
        assert!(TextCompressor::score_token(hash) > 5.0);
        assert!(TextCompressor::score_token(error) > 5.0);
    }

    #[test]
    fn test_score_low_value_tokens() {
        let the = "the";
        let and = "and";
        let very = "very";

        assert!(TextCompressor::score_token(the) < 0.0);
        assert!(TextCompressor::score_token(and) < 0.0);
        assert!(TextCompressor::score_token(very) < 0.0);
    }

    #[test]
    fn test_compress_keeps_identifiers() {
        let compressor = TextCompressor::new();
        let text = "The request with ID abc123def456abc123def456abc123def456abc1 failed";

        let compressed = compressor.compress(text, 0.5); // Keep 50%

        // Should keep ID
        assert!(compressed.contains("abc123def456"));
        // Should keep failed (error marker)
        assert!(compressed.contains("failed"));
    }

    #[test]
    fn test_compress_drops_filler() {
        let compressor = TextCompressor::new();
        let text = "The very important and critical error message is here";

        let compressed = compressor.compress(text, 0.4); // Keep 40%

        // Should keep important, error, message
        assert!(compressed.contains("important"));
        assert!(compressed.contains("error"));
        assert!(compressed.contains("message"));

        // Should drop filler
        assert!(!compressed.contains("very") || compressed.split_whitespace().count() < 6);
    }

    #[test]
    fn test_compress_preserves_numbers() {
        let compressor = TextCompressor::new();
        let text = "Exit code 127 indicates command not found after 5 seconds";

        let compressed = compressor.compress(text, 0.6);

        // Should keep numbers
        assert!(compressed.contains("127"));
        assert!(compressed.contains("5"));
    }

    #[test]
    fn test_is_id_like() {
        assert!(is_id_like("550e8400-e29b-41d4-a716-446655440000")); // UUID
        assert!(is_id_like("abc123def456abc123def456abc123def456abc123def456abc123def456abc1")); // SHA256
        assert!(is_id_like("session_id_12345"));
        assert!(!is_id_like("the"));
        assert!(!is_id_like("and"));
    }

    #[test]
    fn test_is_error_marker() {
        assert!(is_error_marker("error"));
        assert!(is_error_marker("exception"));
        assert!(is_error_marker("connectionerror"));
        assert!(is_error_marker("fatalpanic"));
        assert!(!is_error_marker("message"));
        assert!(!is_error_marker("data"));
    }

    #[test]
    fn test_compress_ratio() {
        let compressor = TextCompressor::new();
        let text = "This is a very long piece of text with many filler words that should be compressed significantly while keeping important identifiers like abc123def456abc123def456abc123def456abc1 and error codes like 404";

        let compressed = compressor.compress(text, 0.3); // Keep 30%
        let original_tokens = text.split_whitespace().count();
        let compressed_tokens = compressed.split_whitespace().count();

        // Should be significantly smaller
        assert!(compressed_tokens <= original_tokens);
        assert!(compressed_tokens as f32 / original_tokens as f32 <= 0.35);
    }

    #[test]
    fn test_compress_empty() {
        let compressor = TextCompressor::new();
        assert_eq!(compressor.compress("", 0.5), "");
    }

    #[test]
    fn test_compress_single_token() {
        let compressor = TextCompressor::new();
        let result = compressor.compress("hello", 0.5);
        assert_eq!(result, "hello");
    }
}