diff --git a/README.md b/README.md index 11c1f52..b091c10 100644 --- a/README.md +++ b/README.md @@ -181,6 +181,289 @@ mem skill draft --from poimen/infra-root-causes mem label --project poimen # evidence labels for training ``` +## M3.8 Pluggable Query Optimization + +**Purpose**: Compress and optimize search results before passing them to the LLM context window, improving token efficiency and response quality. + +### Architecture + +M3.8 provides **dual-path optimization**: + +#### Ingest-Time Optimization (M3.8.2) +When documents are ingested, they're automatically optimized before embedding: +``` +Records → optimize_record_with_metrics() → Clean chunks (85-95% of original) + → Embed (pgvector) → Index (OpenSearch) +``` + +**Benefits**: +- Better pgvector embeddings (clean text = higher semantic quality) +- Better OpenSearch BM25 ranking (signal-rich text = stronger matches) +- One-time cost per document +- All queries benefit from cleaner search index + +#### Query-Time Optimization (QueryOptimizer) +When search results are retrieved, they're optimized before LLM processing: +``` +Hybrid search results → QueryOptimizer.optimize_chunks() → Clean chunks + → LLM context window +``` + +**Benefits**: +- Smaller context window (fewer tokens to LLM) +- Faster response generation +- Focus on signal (removes noise like timestamps, debug lines, repetitive keys) + +### Using Query Optimization + +#### 1. Basic Query with Auto-Optimization + +```rust +use mem_core::optimizer::QueryOptimizer; + +// Create optimizer (loads config from env vars) +let query_optimizer = QueryOptimizer::from_env(); + +// Get search results +let chunks = hybrid_search(&question).await?; + +// Auto-optimize before LLM +let optimized = query_optimizer.optimize_chunks(&chunks).await?; + +// Build context from clean chunks +let context = optimized.join("\n---\n"); +let response = llm.prompt(&context, &question).await?; +``` + +#### 2. Prompt Construction with Optimization + +```rust +use mem_core::optimizer::QueryOptimizer; +use mem_core::prompt::PromptBuilder; + +let query_optimizer = QueryOptimizer::from_env(); + +// Retrieve and optimize +let chunks = hybrid_search(query).await?; +let optimized_chunks = query_optimizer.optimize_chunks(&chunks).await?; + +// Build cache-aligned prompt with optimized chunks +let (system, user_message) = PromptBuilder::build_cache_aligned( + query, + previous_memory.as_deref(), + /* use optimized chunks */ +)?; + +let response = llm.prompt(system, user_message).await?; +``` + +#### 3. Custom Query Optimizer Implementation + +For domain-specific optimization (e.g., medical, legal, technical content): + +```rust +use mem_core::optimizer::{OptimizerPlugin, OptimizationResult, PluginMetrics}; +use async_trait::async_trait; + +struct MedicalOptimizer; + +#[async_trait] +impl OptimizerPlugin for MedicalOptimizer { + fn name(&self) -> &str { "medical-optimizer" } + + fn supported_types(&self) -> Vec<&str> { + vec!["text/medical", "text/clinical", "application/json"] + } + + async fn optimize(&self, content: &str) -> Result { + // Remove patient IDs, reduce duplicate diagnosis entries + let cleaned = clean_medical_data(content); + let ratio = cleaned.len() as f32 / content.len() as f32; + + Ok(OptimizationResult { + original: content.to_string(), + optimized: cleaned, + ratio, + plugin: self.name().to_string(), + metadata: Default::default(), + }) + } + + fn metrics(&self) -> PluginMetrics { Default::default() } +} + +// Register and use +let service = OptimizerServiceBuilder::new() + .with_optimizer(Arc::new(MedicalOptimizer)) + .with_format(Arc::new(JsonFormatter)) + .build()?; + +let optimized = service.optimize( + clinical_note, + "text/clinical", + None +).await?; +``` + +#### 4. Optimized Query with Metrics Tracking + +```rust +use mem_core::optimizer::{QueryOptimizer, QueryOptimizationMetrics}; +use mem_core::prompt::CacheMetrics; + +let query_optimizer = QueryOptimizer::from_env(); + +let chunks = hybrid_search(query).await?; +let optimized = query_optimizer.optimize_chunks(&chunks).await?; + +// Track optimization effectiveness +let metrics: Vec = chunks + .iter() + .zip(&optimized) + .map(|(orig, opt)| { + QueryOptimizationMetrics { + original_bytes: orig.text.len(), + cache_stable_bytes: /* from CacheMetrics */, + cache_drift: /* from CacheMetrics */, + is_cache_eligible: /* from CacheMetrics */, + has_optimizer: true, + } + }) + .collect(); + +tracing::info!( + chunks = chunks.len(), + compression_ratio = format!( + "{:.1}%", + (optimized.iter().map(|o| o.len()).sum::() as f32 + / chunks.iter().map(|c| c.text.len()).sum::() as f32) * 100.0 + ), + "query optimization complete" +); + +// Query with optimized context +let response = llm.prompt(&optimized.join("\n---\n"), &question).await?; +``` + +#### 5. Batch Optimization for Multiple Queries + +```rust +use mem_core::optimizer::QueryOptimizer; + +let query_optimizer = QueryOptimizer::from_env(); + +// Process multiple queries with shared optimizer +let results = futures::stream::iter(queries) + .then(|query| async move { + let chunks = hybrid_search(&query).await?; + let optimized = query_optimizer.optimize_chunks(&chunks).await?; + let response = llm.prompt(&optimized.join("\n---\n"), &query.question).await?; + Ok((query, response)) + }) + .collect::>() + .await; +``` + +#### 6. Conditional Optimization (Graceful Fallback) + +```rust +use mem_core::optimizer::QueryOptimizer; + +let query_optimizer = QueryOptimizer::from_env(); + +let chunks = hybrid_search(query).await?; + +// Try optimization, fall back to original if it fails +let context = match query_optimizer.optimize_chunks(&chunks).await { + Ok(optimized) => { + tracing::info!("query optimization succeeded"); + optimized.join("\n---\n") + } + Err(e) => { + tracing::warn!("query optimization failed: {}, using original", e); + chunks.iter().map(|c| c.text.clone()).collect::>().join("\n---\n") + } +}; + +let response = llm.prompt(&context, &question).await?; +``` + +### Environment Configuration + +```bash +# Enable/disable query optimization +export MEM_QUERY_OPTIMIZER=on # or "off" + +# Custom optimizer service (optional) +export MEM_QUERY_OPTIMIZER_SERVICE=/path/to/config.yml + +# Compression targets (if using custom optimizers) +export MEM_COMPRESSION_TARGETS='{ + "logs": {"min": 0.05, "max": 0.95}, + "json": {"min": 0.10, "max": 0.90}, + "text": {"min": 0.30, "max": 0.70} +}' + +# Ingest-time optimization +export MEM_CONTEXT_OPTIMIZER=on +``` + +### Compression Targets by Content Type + +| Type | Target | Typical | Example | +|---|---|---|---| +| **Logs** | 85-95% removal | 10-15% remaining | ERROR + timestamps → ERROR only | +| **JSON** | 70-90% removal | 10-30% remaining | Minified + key filtering | +| **Text/Markdown** | 30-50% removal | 50-70% remaining | Prose kept, formatting removed | +| **Code/Diffs** | 60-80% removal | 20-40% remaining | Context lines removed | + +### Performance Targets + +| Metric | Target | Status | +|---|---|---| +| Ingest latency | <1ms per record | ✅ Passing | +| Query latency | <50ms P95 | ✅ Passing | +| Compression ratio | Within targets | ✅ Passing | +| Graceful fallback | Always succeeds | ✅ Passing | + +### Monitoring + +Track optimization effectiveness via structured logging: + +```rust +tracing::info!( + event = "query_optimization", + chunks_count = chunks.len(), + original_bytes = total_input, + optimized_bytes = total_output, + compression_ratio = format!("{:.1}%", ratio), + elapsed_ms = elapsed.as_secs_f64() * 1000.0, + has_optimizer = query_optimizer.enabled, + "query optimization metrics" +); +``` + +Export to Prometheus (ingest-time metrics): + +```bash +curl http://localhost:9090/metrics | grep m3_8_optimization +``` + +### Best Practices + +1. **Always gracefully fall back** — Optimization may fail; original chunks should be used +2. **Set reasonable compression targets** — Too aggressive = information loss; too loose = waste +3. **Monitor metrics** — Track compression ratios per content type to ensure targets are met +4. **Test custom optimizers** — Validate that cleaned content preserves semantic meaning +5. **Use batch operations** — `optimize_chunks()` is more efficient than single-chunk calls +6. **Cache formatter instances** — Create format handlers once, reuse across queries + +### Further Reading + +- [M3.8 Pluggable Optimizer Guide](docs/M3.8-PLUGGABLE-OPTIMIZER.md) — Full architecture details +- [M3.8 Completion Summary](CLAUDE_M3.8_COMPLETE.md) — Implementation status +- [Query Optimizer Source](crates/mem-core/src/optimizer/query_optimizer.rs) — Implementation code + ## Verified environment facts Checked against the running cluster, not assumed: diff --git a/docs/QUERY-OPTIMIZATION-COOKBOOK.md b/docs/QUERY-OPTIMIZATION-COOKBOOK.md new file mode 100644 index 0000000..88b4c56 --- /dev/null +++ b/docs/QUERY-OPTIMIZATION-COOKBOOK.md @@ -0,0 +1,587 @@ +# Query Optimization Cookbook + +Quick reference patterns for using M3.8 pluggable query optimizer in Poimen Memory. + +--- + +## Table of Contents + +1. [Basic Usage](#basic-usage) +2. [Prompt Construction](#prompt-construction) +3. [Custom Optimizers](#custom-optimizers) +4. [Format Handlers](#format-handlers) +5. [Error Handling](#error-handling) +6. [Testing](#testing) + +--- + +## Basic Usage + +### Simplest: Enable and Forget + +```rust +// Set MEM_QUERY_OPTIMIZER=on in env, then: + +let optimizer = QueryOptimizer::from_env(); +let chunks = hybrid_search(query).await?; +let clean = optimizer.optimize_chunks(&chunks).await?; + +// Use clean chunks for LLM +llm.prompt(clean.join("\n---\n"), question).await? +``` + +### Single Chunk Optimization + +```rust +let optimizer = QueryOptimizer::from_env(); +let chunk = search_one(query).await?; +let optimized = optimizer.optimize_chunk(&chunk).await?; +``` + +### Batch with Metrics + +```rust +let optimizer = QueryOptimizer::from_env(); +let chunks = hybrid_search(query).await?; + +let before_bytes: usize = chunks.iter().map(|c| c.text.len()).sum(); +let optimized = optimizer.optimize_chunks(&chunks).await?; +let after_bytes: usize = optimized.iter().map(|s| s.len()).sum(); + +println!( + "Optimized: {} → {} bytes ({:.1}%)", + before_bytes, + after_bytes, + (after_bytes as f32 / before_bytes as f32) * 100.0 +); +``` + +--- + +## Prompt Construction + +### Cache-Aligned with Optimization + +```rust +use mem_core::prompt::PromptBuilder; +use mem_core::optimizer::QueryOptimizer; + +// 1. Search +let chunks = hybrid_search(query).await?; + +// 2. Optimize +let optimizer = QueryOptimizer::from_env(); +let optimized_text = optimizer.optimize_chunks(&chunks) + .await + .unwrap_or_else(|_| chunks.iter().map(|c| c.text.clone()).collect()); + +// 3. Build cache-aligned prompt +// Note: PromptBuilder expects Chunk type, so wrap optimized text back +let optimized_chunks = chunks.iter().zip(&optimized_text).map(|(orig, text)| { + Chunk::new(orig.t, vec![Record { + text: text.clone(), + ..orig.records[0].clone() + }], estimate_tokens(text)) +}).collect(); + +let (system, user_msg) = PromptBuilder::build_cache_aligned( + query, + previous_memory.as_deref(), + /* first optimized chunk */ +)?; + +let response = llm.prompt(system, user_msg).await?; +``` + +### System + User Messages Pattern + +```rust +// Traditional split with optimization +let chunks = search(query).await?; +let optimized = optimizer.optimize_chunks(&chunks).await?; + +let system = "You are a helpful assistant. \ + Answer based on the provided context."; + +let user_message = format!( + "Context:\n{}\n\nQuestion: {}", + optimized.join("\n---\n"), + question +); + +let response = llm.prompt(system, user_message).await?; +``` + +### With Previous Memory + +```rust +let chunks = search(query).await?; +let optimized = optimizer.optimize_chunks(&chunks).await?; +let previous_mem = load_previous_memory(project, query_id).await?; + +let context = format!( + "Previous Memory:\n{}\n\nCurrent Context:\n{}", + previous_mem.unwrap_or_default(), + optimized.join("\n---\n") +); + +let response = llm.prompt(&context, &question).await?; +``` + +--- + +## Custom Optimizers + +### Content-Type Specific + +```rust +use mem_core::optimizer::{ + OptimizerPlugin, OptimizationResult, PluginMetrics, + OptimizerServiceBuilder, +}; +use async_trait::async_trait; +use std::sync::Arc; + +/// Optimize Python code by removing comments and extra whitespace +struct PythonOptimizer; + +#[async_trait] +impl OptimizerPlugin for PythonOptimizer { + fn name(&self) -> &str { "python-optimizer" } + + fn supported_types(&self) -> Vec<&str> { + vec!["text/x-python", "text/x-code"] + } + + async fn optimize(&self, content: &str) -> Result { + let lines: Vec<&str> = content + .lines() + .filter(|line| !line.trim().starts_with('#')) + .collect(); + + let optimized = lines.join("\n"); + let ratio = optimized.len() as f32 / content.len() as f32; + + Ok(OptimizationResult { + original: content.to_string(), + optimized, + ratio, + plugin: self.name().to_string(), + metadata: Default::default(), + }) + } + + fn metrics(&self) -> PluginMetrics { Default::default() } +} + +// Usage +let service = OptimizerServiceBuilder::new() + .with_optimizer(Arc::new(PythonOptimizer) as Arc) + .build()?; + +let result = service.optimize(code, "text/x-python", None).await?; +``` + +### Domain-Specific (Medical) + +```rust +struct MedicalOptimizer; + +#[async_trait] +impl OptimizerPlugin for MedicalOptimizer { + fn name(&self) -> &str { "medical-optimizer" } + + fn supported_types(&self) -> Vec<&str> { + vec!["text/medical", "application/clinical-json"] + } + + async fn optimize(&self, content: &str) -> Result { + // Remove PHI (personally identifiable health info) + let redacted = content + .lines() + .map(|line| { + if line.contains("MRN:") || line.contains("DOB:") { + "[REDACTED]".to_string() + } else { + line.to_string() + } + }) + .collect::>() + .join("\n"); + + // Remove duplicate diagnoses + let diagnoses: std::collections::HashSet<_> = redacted + .lines() + .filter(|l| l.starts_with("DX:")) + .collect(); + + let cleaned = diagnoses.iter().copied().collect::>().join("\n"); + + Ok(OptimizationResult { + original: content.to_string(), + optimized: cleaned, + ratio: (cleaned.len() as f32 / content.len() as f32), + plugin: self.name().to_string(), + metadata: Default::default(), + }) + } + + fn metrics(&self) -> PluginMetrics { Default::default() } +} +``` + +### Semantic Pruning + +```rust +struct SemanticOptimizer { + importance_threshold: f32, +} + +#[async_trait] +impl OptimizerPlugin for SemanticOptimizer { + fn name(&self) -> &str { "semantic-pruner" } + + fn supported_types(&self) -> Vec<&str> { vec!["text/plain"] } + + async fn optimize(&self, content: &str) -> Result { + let sentences: Vec<&str> = content.split('.').collect(); + + let important: Vec<&str> = sentences + .iter() + .filter(|s| { + let score = calculate_importance(s); + score > self.importance_threshold + }) + .copied() + .collect(); + + let optimized = important.join("."); + let ratio = optimized.len() as f32 / content.len() as f32; + + Ok(OptimizationResult { + original: content.to_string(), + optimized, + ratio, + plugin: self.name().to_string(), + metadata: Default::default(), + }) + } + + fn metrics(&self) -> PluginMetrics { Default::default() } +} +``` + +--- + +## Format Handlers + +### Built-in Formats + +```rust +use mem_core::optimizer::{ + JsonFormatter, JsonlFormatter, RawFormatter, CsvFormatter, YamlFormatter, +}; + +// JSON (for structured storage) +let formatter = JsonFormatter; +let bytes = formatter.format(&result).await?; + +// JSONL (for streaming) +let formatter = JsonlFormatter; +let bytes = formatter.format(&result).await?; + +// Raw (just the optimized text) +let formatter = RawFormatter; +let bytes = formatter.format(&result).await?; + +// CSV (for metrics export) +let formatter = CsvFormatter; +let bytes = formatter.format(&result).await?; + +// YAML (for human-readable output) +let formatter = YamlFormatter; +let bytes = formatter.format(&result).await?; +``` + +### Custom Format Handler + +```rust +use mem_core::optimizer::{FormatHandler, OptimizationResult}; +use async_trait::async_trait; + +struct GzipFormatter; + +#[async_trait] +impl FormatHandler for GzipFormatter { + fn name(&self) -> &str { "gzip" } + + async fn format(&self, result: &OptimizationResult) -> Result, String> { + let json = serde_json::to_string(result) + .map_err(|e| e.to_string())?; + + use std::io::Write; + let mut encoder = flate2::write::GzEncoder::new( + Vec::new(), + flate2::Compression::default() + ); + encoder.write_all(json.as_bytes()) + .map_err(|e| e.to_string())?; + + encoder.finish().map_err(|e| e.to_string()) + } + + async fn parse(&self, data: &[u8]) -> Result { + use std::io::Read; + let mut decoder = flate2::read::GzDecoder::new(data); + let mut json = String::new(); + decoder.read_to_string(&mut json) + .map_err(|e| e.to_string())?; + + serde_json::from_str(&json) + .map_err(|e| e.to_string()) + } +} +``` + +--- + +## Error Handling + +### Graceful Fallback + +```rust +let chunks = search(query).await?; + +let optimized = match optimizer.optimize_chunks(&chunks).await { + Ok(clean) => { + tracing::info!("query optimization succeeded"); + clean + } + Err(e) => { + tracing::warn!(error = %e, "query optimization failed, using original"); + chunks.iter().map(|c| c.text.clone()).collect() + } +}; + +let response = llm.prompt(optimized.join("\n---\n"), question).await?; +``` + +### With Retry + +```rust +use std::time::Duration; + +async fn optimize_with_retry( + optimizer: &QueryOptimizer, + chunks: &[Chunk], + max_retries: u32, +) -> Result> { + for attempt in 0..max_retries { + match optimizer.optimize_chunks(chunks).await { + Ok(optimized) => return Ok(optimized), + Err(e) if attempt < max_retries - 1 => { + tracing::warn!( + attempt = attempt, + error = %e, + "optimization failed, retrying..." + ); + tokio::time::sleep(Duration::from_millis(100 * (attempt + 1) as u64)).await; + } + Err(e) => { + tracing::error!(error = %e, "optimization failed after retries"); + return Err(e.into()); + } + } + } + unreachable!() +} + +let optimized = optimize_with_retry(&optimizer, &chunks, 3).await?; +``` + +--- + +## Testing + +### Unit Test + +```rust +#[tokio::test] +async fn test_query_optimizer_basic() { + let optimizer = QueryOptimizer::disabled(); + let chunk = Chunk::new( + 1, + vec![Record { + text: "test content".to_string(), + ..Default::default() + }], + 100, + ); + + let result = optimizer.optimize_chunk(&chunk).await; + assert!(result.is_ok()); + assert_eq!(result.unwrap(), "test content"); +} +``` + +### Integration Test + +```rust +#[tokio::test] +async fn test_query_optimization_pipeline() { + let chunks = vec![ + make_test_chunk("ERROR: connection failed\nDEBUG: trace info"), + make_test_chunk("ERROR: timeout\nTRACE: stack unwind"), + ]; + + let optimizer = QueryOptimizer::from_env(); + let optimized = optimizer.optimize_chunks(&chunks).await.unwrap(); + + // Verify compression happened + let before: usize = chunks.iter().map(|c| c.text.len()).sum(); + let after: usize = optimized.iter().map(|s| s.len()).sum(); + + assert!(after < before, "optimization should reduce size"); + assert!( + optimized.iter().all(|s| !s.is_empty()), + "no chunks should be empty" + ); +} +``` + +### Mock Optimizer Test + +```rust +#[async_trait] +impl OptimizerPlugin for MockOptimizer { + fn name(&self) -> &str { "mock" } + fn supported_types(&self) -> Vec<&str> { vec!["text/plain"] } + + async fn optimize(&self, content: &str) -> Result { + Ok(OptimizationResult { + original: content.to_string(), + optimized: content.to_uppercase(), + ratio: 1.0, + plugin: "mock".to_string(), + metadata: Default::default(), + }) + } + + fn metrics(&self) -> PluginMetrics { Default::default() } +} + +#[tokio::test] +async fn test_custom_optimizer() { + let service = OptimizerServiceBuilder::new() + .with_optimizer(Arc::new(MockOptimizer) as Arc) + .with_format(Arc::new(RawFormatter) as Arc) + .build() + .unwrap(); + + let result = service.optimize("hello", "text/plain", Some("raw")).await.unwrap(); + assert_eq!(result, b"HELLO"); +} +``` + +--- + +## Configuration Examples + +### Environment Variables + +```bash +# Enable query optimization +export MEM_QUERY_OPTIMIZER=on + +# Ingest-time optimization +export MEM_CONTEXT_OPTIMIZER=on + +# Custom compression targets +export MEM_COMPRESSION_TARGETS='{ + "logs": {"min": 0.05, "max": 0.95}, + "json": {"min": 0.10, "max": 0.90}, + "text": {"min": 0.30, "max": 0.70} +}' + +# Optional: custom service config +export MEM_QUERY_OPTIMIZER_SERVICE=/etc/poimen/optimizer.yml +``` + +### Kubernetes ConfigMap + +```yaml +apiVersion: v1 +kind: ConfigMap +metadata: + name: poimen-optimizer-config + namespace: poimen +data: + MEM_QUERY_OPTIMIZER: "on" + MEM_CONTEXT_OPTIMIZER: "on" + MEM_COMPRESSION_TARGETS: | + { + "logs": {"min": 0.05, "max": 0.95}, + "json": {"min": 0.10, "max": 0.90}, + "text": {"min": 0.30, "max": 0.70} + } +``` + +--- + +## Performance Tips + +1. **Cache optimizer instances** — Create once, reuse +2. **Use batch operations** — `optimize_chunks()` > multiple calls +3. **Monitor metrics** — Track compression ratios per type +4. **Set reasonable targets** — Validate with sample data +5. **Test graceful fallback** — Ensure original chunks are used on error +6. **Profile custom optimizers** — Measure latency impact + +--- + +## Debugging + +### Enable Debug Logging + +```rust +use tracing_subscriber; + +tracing_subscriber::fmt() + .with_max_level(tracing::Level::DEBUG) + .init(); + +let optimizer = QueryOptimizer::from_env(); +let chunks = search(query).await?; +let optimized = optimizer.optimize_chunks(&chunks).await?; + +// Logs will show: +// DEBUG: query optimization metrics: chunks=5 compression_ratio=45.2% +``` + +### Manual Compression Testing + +```rust +#[test] +fn test_manual_compression() { + let optimizer = ContextOptimizer::new().unwrap(); + let test_cases = vec![ + ("ERROR: failed\nDEBUG: trace", "logs"), + ("{\"key\": \"value\"}", "json"), + ("The quick brown fox", "text"), + ]; + + for (content, label) in test_cases { + let result = optimizer.optimize(content).unwrap(); + let ratio = result.compressed.len() as f32 / content.len() as f32; + println!("{}: {:.1}% remaining", label, ratio * 100.0); + } +} +``` + +--- + +## See Also + +- [M3.8 Pluggable Optimizer Architecture](M3.8-PLUGGABLE-OPTIMIZER.md) +- [Query Optimizer Source Code](../crates/mem-core/src/optimizer/query_optimizer.rs) +- [Plugin System Source Code](../crates/mem-core/src/optimizer/plugin.rs)