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;

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

pub use router::ContentRouter;
pub use log::LogCompressor;
pub use json::JsonCrusher;
pub use diff::DiffCompressor;

#[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 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,
    config: ContextOptimizerConfig,
}

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

    /// 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();

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

    /// 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)?
            }
            _ => content.to_string(), // Passthrough for other types
        };

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

        Ok(OptimizedChunk {
            compressed,
            original_tokens,
            compressed_tokens,
            content_type,
            ccr_hash: None, // TODO: Implement CCR
        })
    }
}

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");
    }
}