poimen-memory/docs/QUERY-OPTIMIZATION-COOKBOOK.md

# 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<OptimizationResult, String> {
        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<dyn OptimizerPlugin>)
    .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<OptimizationResult, String> {
        // 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::<Vec<_>>()
            .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::<Vec<_>>().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<OptimizationResult, String> {
        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<Vec<u8>, 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<OptimizationResult, String> {
        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<Vec<String>> {
    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<OptimizationResult, String> {
        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<dyn OptimizerPlugin>)
        .with_format(Arc::new(RawFormatter) as Arc<dyn FormatHandler>)
        .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)