diff --git a/crates/mem-core/src/optimizer/mod.rs b/crates/mem-core/src/optimizer/mod.rs index b000f74..7cc08b3 100644 --- a/crates/mem-core/src/optimizer/mod.rs +++ b/crates/mem-core/src/optimizer/mod.rs @@ -13,6 +13,7 @@ pub mod cache_align; pub mod ccr; pub mod plugin; pub mod builtin; +pub mod query_optimizer; use anyhow::Result; use serde::{Deserialize, Serialize}; @@ -31,6 +32,7 @@ pub use plugin::{ pub use builtin::{ BuiltinOptimizer, JsonFormatter, JsonlFormatter, RawFormatter, CsvFormatter, YamlFormatter, }; +pub use query_optimizer::{QueryOptimizer, QueryOptimizationMetrics}; #[derive(Debug, Clone, Serialize, Deserialize)] pub struct OptimizedChunk { diff --git a/crates/mem-core/src/optimizer/query_optimizer.rs b/crates/mem-core/src/optimizer/query_optimizer.rs new file mode 100644 index 0000000..34b7493 --- /dev/null +++ b/crates/mem-core/src/optimizer/query_optimizer.rs @@ -0,0 +1,213 @@ +//! Query-time optimization: compress context before sending to LLM +//! +//! Uses pluggable service to apply custom optimization strategies +//! during query execution (in addition to ingest-time optimization). + +use super::plugin::OptimizerService; +use crate::prompt::CacheMetrics; +use crate::domain::{Chunk, Record}; +use anyhow::Result; + +/// Query optimizer: compresses chunks before LLM processing +/// Operates on search results AFTER retrieval but BEFORE LLM context window +pub struct QueryOptimizer { + service: Option, + enabled: bool, +} + +impl QueryOptimizer { + /// Create optimizer from environment + pub fn from_env() -> Self { + let enabled = std::env::var("MEM_QUERY_OPTIMIZER") + .map(|v| v.to_lowercase() == "on") + .unwrap_or(false); + + if !enabled { + return Self { + service: None, + enabled: false, + }; + } + + // Try to initialize service (will use defaults if available) + let service = std::env::var("MEM_QUERY_OPTIMIZER_SERVICE") + .ok() + .and_then(|_| { + // Would load custom plugins from config here + None + }); + + Self { + service, + enabled: true, + } + } + + /// Create disabled optimizer + pub fn disabled() -> Self { + Self { + service: None, + enabled: false, + } + } + + /// Set service (for testing/custom initialization) + pub fn with_service(mut self, service: OptimizerService) -> Self { + self.service = Some(service); + self.enabled = true; + self + } + + /// Get text from chunk by concatenating records + fn chunk_text(chunk: &Chunk) -> String { + chunk + .records + .iter() + .map(|r| r.text.as_str()) + .collect::>() + .join("\n") + } + + /// Optimize chunk for LLM context (async) + /// Returns optimized text (or original if optimization disabled/failed) + pub async fn optimize_chunk(&self, chunk: &Chunk) -> Result { + if !self.enabled || self.service.is_none() { + return Ok(Self::chunk_text(chunk)); + } + + let service = self.service.as_ref().unwrap(); + let chunk_text = Self::chunk_text(chunk); + let content_type = Self::infer_content_type(&chunk_text); + + // Optimize with default format (raw = just the text) + match service.optimize(&chunk_text, &content_type, Some("raw")).await { + Ok(bytes) => { + let text = String::from_utf8(bytes) + .unwrap_or_else(|_| chunk_text); + Ok(text) + } + Err(_) => { + // Fail gracefully: return original if optimization fails + tracing::warn!("Query optimization failed, using original chunk"); + Ok(chunk_text) + } + } + } + + /// Optimize multiple chunks (batch processing) + pub async fn optimize_chunks(&self, chunks: &[Chunk]) -> Result> { + let mut optimized = Vec::with_capacity(chunks.len()); + for chunk in chunks { + optimized.push(self.optimize_chunk(chunk).await?); + } + Ok(optimized) + } + + /// Calculate metrics for a chunk (cache alignment + compression ratio) + pub fn metrics(&self, chunk: &Chunk, cache_metrics: &CacheMetrics) -> QueryOptimizationMetrics { + QueryOptimizationMetrics { + original_bytes: Self::chunk_text(chunk).len(), + cache_stable_bytes: cache_metrics.stable_prefix_bytes, + cache_drift: cache_metrics.drift_metric, + is_cache_eligible: cache_metrics.cache_eligible, + has_optimizer: self.enabled, + } + } + + /// Infer content type from chunk text + fn infer_content_type(text: &str) -> String { + // Simple heuristics (real implementation would use ContentRouter) + if text.contains('{') && text.contains('}') { + "application/json".to_string() + } else if text.contains("ERROR:") || text.contains("WARN:") { + "text/x-log".to_string() + } else if text.starts_with("---") || text.contains("---") { + "text/x-diff".to_string() + } else { + "text/plain".to_string() + } + } +} + +/// Metrics for query-time optimization +#[derive(Debug, Clone)] +pub struct QueryOptimizationMetrics { + pub original_bytes: usize, + pub cache_stable_bytes: usize, + pub cache_drift: f32, + pub is_cache_eligible: bool, + pub has_optimizer: bool, +} + +#[cfg(test)] +mod tests { + use super::*; + use time::macros::datetime; + use crate::domain::{Provenance, Role, Record}; + + fn make_test_chunk(text: &str) -> Chunk { + Chunk::new( + 1, + vec![Record { + role: Role::User, + text: text.to_string(), + timestamp: datetime!(2024-08-28 12:00:00 UTC), + provenance: Provenance { + source_id: "test".to_string(), + offset: 0, + }, + }], + 100, + ) + } + + #[test] + fn test_optimizer_disabled() { + let optimizer = QueryOptimizer::disabled(); + assert!(!optimizer.enabled); + } + + #[test] + fn test_optimizer_from_env() { + std::env::set_var("MEM_QUERY_OPTIMIZER", "off"); + let optimizer = QueryOptimizer::from_env(); + assert!(!optimizer.enabled); + } + + #[tokio::test] + async fn test_optimize_chunk_disabled() { + let optimizer = QueryOptimizer::disabled(); + let chunk = make_test_chunk("test content"); + + let result = optimizer.optimize_chunk(&chunk).await; + assert_eq!(result.unwrap(), "test content"); + } + + #[test] + fn test_infer_content_type_json() { + let json = "{\"key\": \"value\"}"; + let content_type = QueryOptimizer::infer_content_type(json); + assert_eq!(content_type, "application/json"); + } + + #[test] + fn test_infer_content_type_log() { + let log = "ERROR: something failed"; + let content_type = QueryOptimizer::infer_content_type(log); + assert_eq!(content_type, "text/x-log"); + } + + #[test] + fn test_infer_content_type_diff() { + let diff = "---\n+++ file.txt"; + let content_type = QueryOptimizer::infer_content_type(diff); + assert_eq!(content_type, "text/x-diff"); + } + + #[test] + fn test_infer_content_type_text() { + let text = "Plain text content"; + let content_type = QueryOptimizer::infer_content_type(text); + assert_eq!(content_type, "text/plain"); + } +}