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

//! Context Optimizer — Pre-LLM compression pipeline
//!
//! Sits between hybrid search retrieval and LLM gateway.
//! Compresses evidence chunks to reduce token costs and stabilize KV cache hits.
//! Search indexes remain untouched at full fidelity.

pub mod router;
pub mod log;
pub mod json;
pub mod diff;
pub mod text;
pub mod cache_align;
pub mod ccr;

use anyhow::Result;
use serde::{Deserialize, Serialize};

pub use router::ContentRouter;
pub use log::LogCompressor;
pub use json::JsonCrusher;
pub use diff::DiffCompressor;
pub use text::TextCompressor;
pub use cache_align::{CacheAligner, AlignedContent};
pub use ccr::CcrStore;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OptimizedChunk {
    /// Compressed content
    pub compressed: String,
    /// Original token count (estimated)
    pub original_tokens: usize,
    /// Compressed token count (estimated)
    pub compressed_tokens: usize,
    /// Detected content type
    pub content_type: ContentType,
    /// CCR hash for retrieving original (if compressed)
    pub ccr_hash: Option<String>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ContentType {
    Json,
    Code,
    Log,
    Diff,
    Config,
    Text,
}

#[derive(Debug, Clone)]
pub struct ContextOptimizerConfig {
    /// Enable/disable optimizer entirely
    pub enabled: bool,
    /// Enable Magika ML detection
    pub use_magika: bool,
    /// Confidence threshold for Magika (0.0-1.0)
    pub magika_threshold: f32,
    /// Per-compressor toggles
    pub compress_json: bool,
    pub compress_logs: bool,
    pub compress_code: bool,
    pub compress_diff: bool,
    pub compress_text: bool,
    /// Target token budget per chunk (0 = no budget)
    pub token_budget: usize,
    /// Enable CCR (Compress-Cache-Retrieve)
    pub ccr_enabled: bool,
}

impl ContextOptimizerConfig {
    /// Load from environment variables
    pub fn from_env() -> Self {
        Self {
            enabled: std::env::var("MEM_CONTEXT_OPTIMIZER")
                .map(|v| v.to_lowercase() == "on" || v == "true")
                .unwrap_or(true),
            use_magika: std::env::var("MEM_MAGIKA_ENABLED")
                .map(|v| v.to_lowercase() == "on" || v == "true")
                .unwrap_or(true),
            magika_threshold: std::env::var("MEM_MAGIKA_THRESHOLD")
                .ok()
                .and_then(|v| v.parse().ok())
                .unwrap_or(0.7),
            compress_json: std::env::var("MEM_COMPRESS_JSON")
                .map(|v| v.to_lowercase() == "on" || v == "true")
                .unwrap_or(true),
            compress_logs: std::env::var("MEM_COMPRESS_LOGS")
                .map(|v| v.to_lowercase() == "on" || v == "true")
                .unwrap_or(true),
            compress_code: std::env::var("MEM_COMPRESS_CODE")
                .map(|v| v.to_lowercase() == "on" || v == "true")
                .unwrap_or(false), // opt-in
            compress_diff: std::env::var("MEM_COMPRESS_DIFF")
                .map(|v| v.to_lowercase() == "on" || v == "true")
                .unwrap_or(true),
            compress_text: std::env::var("MEM_COMPRESS_TEXT")
                .map(|v| v.to_lowercase() == "on" || v == "true")
                .unwrap_or(true),
            token_budget: std::env::var("MEM_TOKEN_BUDGET")
                .ok()
                .and_then(|v| v.parse().ok())
                .unwrap_or(0),
            ccr_enabled: std::env::var("MEM_CCR_ENABLED")
                .map(|v| v.to_lowercase() == "on" || v == "true")
                .unwrap_or(true),
        }
    }
}

impl Default for ContextOptimizerConfig {
    fn default() -> Self {
        Self {
            enabled: true,
            use_magika: true,
            magika_threshold: 0.7,
            compress_json: true,
            compress_logs: true,
            compress_code: false,
            compress_diff: true,
            compress_text: true,
            token_budget: 0,
            ccr_enabled: true,
        }
    }
}

/// Main context optimizer
pub struct ContextOptimizer {
    router: ContentRouter,
    log_compressor: LogCompressor,
    json_crusher: JsonCrusher,
    diff_compressor: DiffCompressor,
    text_compressor: TextCompressor,
    ccr_store: CcrStore,
    config: ContextOptimizerConfig,
}

impl ContextOptimizer {
    /// Create a new optimizer with default config
    pub fn new() -> Result<Self> {
        Self::with_config(ContextOptimizerConfig::default())
    }

    /// Create optimizer from environment config
    pub fn from_env() -> Result<Self> {
        Self::with_config(ContextOptimizerConfig::from_env())
    }

    /// Create with custom config
    pub fn with_config(config: ContextOptimizerConfig) -> Result<Self> {
        let router = ContentRouter::new()?;
        let log_compressor = LogCompressor::new();
        let json_crusher = JsonCrusher::new();
        let diff_compressor = DiffCompressor::new();
        let text_compressor = TextCompressor::new();
        let ccr_store = CcrStore::new();

        Ok(Self {
            router,
            log_compressor,
            json_crusher,
            diff_compressor,
            text_compressor,
            ccr_store,
            config,
        })
    }

    /// Get reference to CCR store for retrieval
    pub fn ccr_store(&self) -> &CcrStore {
        &self.ccr_store
    }

    /// Optimize a chunk of content
    pub fn optimize(&self, content: &str) -> Result<OptimizedChunk> {
        if !self.config.enabled {
            // Passthrough mode
            let token_estimate = estimate_tokens(content);
            return Ok(OptimizedChunk {
                compressed: content.to_string(),
                original_tokens: token_estimate,
                compressed_tokens: token_estimate,
                content_type: ContentType::Text,
                ccr_hash: None,
            });
        }

        // Detect content type
        let content_type = self.router.detect(content)?;

        // Compress based on type
        let compressed = match content_type {
            ContentType::Log if self.config.compress_logs => {
                self.log_compressor.compress(content)?
            }
            ContentType::Json if self.config.compress_json => {
                self.json_crusher.compress(content)?
            }
            ContentType::Diff if self.config.compress_diff => {
                self.diff_compressor.compress(content)?
            }
            ContentType::Text if self.config.compress_text => {
                self.text_compressor.compress(content, 0.4) // Keep 40% of tokens
            }
            _ => content.to_string(), // Passthrough for other types
        };

        let original_tokens = estimate_tokens(content);
        let compressed_tokens = estimate_tokens(&compressed);

        // Store original in CCR if compression happened and CCR is enabled
        let ccr_hash = if self.config.ccr_enabled && compressed != content {
            self.ccr_store.store(content).ok()
        } else {
            None
        };

        Ok(OptimizedChunk {
            compressed,
            original_tokens,
            compressed_tokens,
            content_type,
            ccr_hash,
        })
    }
}

impl Default for ContextOptimizer {
    fn default() -> Self {
        Self::new().expect("failed to create optimizer")
    }
}

/// Simple token estimation (1 token ~= 4 chars, 1 space-separated word)
fn estimate_tokens(text: &str) -> usize {
    (text.len() / 4).max(text.split_whitespace().count())
}

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

    #[test]
    fn test_token_estimate() {
        let text = "hello world this is a test";
        let tokens = estimate_tokens(text);
        assert!(tokens > 0);
    }

    #[test]
    fn test_optimizer_passthrough_when_disabled() {
        let config = ContextOptimizerConfig {
            enabled: false,
            ..Default::default()
        };
        let optimizer = ContextOptimizer::with_config(config).unwrap();
        let chunk = optimizer.optimize("test content").unwrap();
        assert_eq!(chunk.compressed, "test content");
    }

    #[test]
    fn test_optimizer_with_ccr() {
        let optimizer = ContextOptimizer::new().unwrap();
        let log_content = "ERROR: failed\nWARN: ignored";
        let chunk = optimizer.optimize(log_content).unwrap();

        // If compression happened and CCR enabled, should have hash
        if chunk.compressed != log_content {
            assert!(chunk.ccr_hash.is_some());
        }
    }

    #[test]
    fn test_optimizer_cache_align() {
        let content = "System prompt\nAt 2026-08-28T09:15:00Z query was:\nContext";
        let aligned = CacheAligner::align(content);

        assert!(aligned.stable_prefix.contains("System"));
        assert!(aligned.stable_prefix.contains("Context"));
        assert!(aligned.dynamic_tail.contains("2026-08-28"));
    }

    #[test]
    fn test_optimizer_compresses_text() {
        let config = ContextOptimizerConfig {
            enabled: true,
            compress_text: true,
            ..Default::default()
        };
        let optimizer = ContextOptimizer::with_config(config).unwrap();

        let text = "The very important and critical error message is present in the system";
        let chunk = optimizer.optimize(text).unwrap();

        // Should detect as text
        assert_eq!(chunk.content_type, ContentType::Text);
        // Should compress
        assert!(chunk.compressed_tokens < chunk.original_tokens);
    }

    #[test]
    fn test_config_from_env() {
        // Just verify it doesn't panic
        let config = ContextOptimizerConfig::from_env();
        assert!(config.enabled);
    }
}