//! 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; pub mod plugin; pub mod builtin; pub mod query_optimizer; 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; pub use plugin::{ OptimizerPlugin, FormatHandler, Registry, OptimizationResult, PluginMetrics, SimpleRegistry, PluginLocator, DefaultLocator, OptimizerService, OptimizerServiceBuilder, }; pub use builtin::{ BuiltinOptimizer, JsonFormatter, JsonlFormatter, RawFormatter, CsvFormatter, YamlFormatter, }; pub use query_optimizer::{QueryOptimizer, QueryOptimizationMetrics}; #[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, } #[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::with_config(ContextOptimizerConfig::default()) } /// Create optimizer from environment config pub fn from_env() -> Result { Self::with_config(ContextOptimizerConfig::from_env()) } /// Create with custom config pub fn with_config(config: ContextOptimizerConfig) -> Result { 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 { 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); } }