poimen-memory/crates/mem-cli/src/chunk_optimizer.rs

/// Phase 4: LLM Call Optimization
///
/// Reduce LLM calls by:
/// 1. Score thresholding: skip chunks < 0.6
/// 2. Budget-aware selection: select top-K within byte budget
/// 3. Deduplication: remove near-duplicate chunks (shingle-based)
/// 4. Ranking by value: prioritize high-confidence results
///
/// Target: 70-80% fewer LLM calls for typical queries

use anyhow::Result;
use std::collections::{HashMap, HashSet};

/// Chunk with selection metrics
#[derive(Debug, Clone)]
pub struct OptimizableChunk {
    pub id: String,
    pub text: String,
    pub score: f32,
    pub confidence: f32,  // How confident are we in this result?
    pub size_bytes: usize,
}

/// Selection result with metrics
#[derive(Debug, Clone)]
pub struct SelectionMetrics {
    pub selected_count: usize,
    pub rejected_count: usize,
    pub total_bytes: usize,
    pub budget_used_pct: f32,
    pub avg_score: f32,
    pub dedup_removed: usize,
}

/// Score Threshold Filter
pub struct ScoreThresholdFilter {
    min_score: f32,
}

impl ScoreThresholdFilter {
    pub fn new(min_score: f32) -> Self {
        Self { min_score }
    }

    /// Filter chunks by minimum score
    pub fn filter(&self, chunks: Vec<OptimizableChunk>) -> Vec<OptimizableChunk> {
        chunks
            .into_iter()
            .filter(|c| c.score >= self.min_score)
            .collect()
    }
}

/// Budget-Aware Chunk Selector
pub struct BudgetSelector {
    max_bytes: usize,
    min_score_threshold: f32,
}

impl BudgetSelector {
    pub fn new(max_bytes: usize, min_score_threshold: f32) -> Self {
        Self {
            max_bytes,
            min_score_threshold,
        }
    }

    /// Select top chunks within byte budget (greedy: highest score first)
    pub fn select(
        &self,
        mut chunks: Vec<OptimizableChunk>,
    ) -> (Vec<OptimizableChunk>, SelectionMetrics) {
        // Sort by score descending
        chunks.sort_by(|a, b| {
            b.score
                .partial_cmp(&a.score)
                .unwrap_or(std::cmp::Ordering::Equal)
        });

        let total_count = chunks.len();
        let mut selected = Vec::new();
        let mut total_bytes = 0usize;
        let mut rejected_count = 0;

        for chunk in chunks {
            // Check threshold
            if chunk.score < self.min_score_threshold {
                rejected_count += 1;
                continue;
            }

            // Check budget
            if total_bytes + chunk.size_bytes <= self.max_bytes {
                total_bytes += chunk.size_bytes;
                selected.push(chunk);
            } else {
                rejected_count += 1;
            }
        }

        let avg_score = if selected.is_empty() {
            0.0
        } else {
            selected.iter().map(|c| c.score).sum::<f32>() / selected.len() as f32
        };

        let metrics = SelectionMetrics {
            selected_count: selected.len(),
            rejected_count,
            total_bytes,
            budget_used_pct: (total_bytes as f32 / self.max_bytes as f32 * 100.0).min(100.0),
            avg_score,
            dedup_removed: 0,
        };

        (selected, metrics)
    }
}

/// Shingle-based Deduplication
pub struct ShingleDeduplicator {
    min_shingle_overlap: f32,
    shingle_size: usize,
}

impl ShingleDeduplicator {
    pub fn new(min_shingle_overlap: f32, shingle_size: usize) -> Self {
        Self {
            min_shingle_overlap,
            shingle_size,
        }
    }

    /// Extract k-shingles (word-level) from text
    fn get_shingles(&self, text: &str) -> HashSet<String> {
        let text_lower = text.to_lowercase();
        let words: Vec<&str> = text_lower
            .split_whitespace()
            .collect();

        let mut shingles = HashSet::new();
        for window in words.windows(self.shingle_size) {
            shingles.insert(window.join(" "));
        }
        shingles
    }

    /// Calculate Jaccard similarity between two texts
    fn jaccard_similarity(&self, text_a: &str, text_b: &str) -> f32 {
        let shingles_a = self.get_shingles(text_a);
        let shingles_b = self.get_shingles(text_b);

        if shingles_a.is_empty() || shingles_b.is_empty() {
            return 0.0;
        }

        let intersection = shingles_a.intersection(&shingles_b).count();
        let union = shingles_a.union(&shingles_b).count();

        intersection as f32 / union as f32
    }

    /// Deduplicate chunks by shingle overlap
    pub fn deduplicate(&self, mut chunks: Vec<OptimizableChunk>) -> (Vec<OptimizableChunk>, usize) {
        // Sort by score descending (keep highest-scoring duplicates)
        chunks.sort_by(|a, b| {
            b.score
                .partial_cmp(&a.score)
                .unwrap_or(std::cmp::Ordering::Equal)
        });

        let mut kept = Vec::new();
        let mut removed = 0;

        for chunk in chunks {
            let is_duplicate = kept.iter().any(|kept_chunk: &OptimizableChunk| {
                let sim = self.jaccard_similarity(&chunk.text, &kept_chunk.text);
                sim >= self.min_shingle_overlap
            });

            if !is_duplicate {
                kept.push(chunk);
            } else {
                removed += 1;
            }
        }

        (kept, removed)
    }
}

/// Full Chunk Optimization Pipeline
pub struct ChunkOptimizer {
    threshold_filter: ScoreThresholdFilter,
    budget_selector: BudgetSelector,
    deduplicator: ShingleDeduplicator,
}

impl ChunkOptimizer {
    pub fn new(
        min_score: f32,
        max_bytes: usize,
        min_dedup_overlap: f32,
    ) -> Self {
        Self {
            threshold_filter: ScoreThresholdFilter::new(min_score),
            budget_selector: BudgetSelector::new(max_bytes, min_score),
            deduplicator: ShingleDeduplicator::new(min_dedup_overlap, 3),
        }
    }

    /// End-to-end optimization pipeline
    pub fn optimize(&self, chunks: Vec<OptimizableChunk>) -> (Vec<OptimizableChunk>, SelectionMetrics) {
        // Step 1: Filter by threshold
        let filtered = self.threshold_filter.filter(chunks.clone());

        // Step 2: Deduplicate
        let (deduplicated, dedup_removed) = self.deduplicator.deduplicate(filtered);

        // Step 3: Select within budget
        let (selected, mut metrics) = self.budget_selector.select(deduplicated);

        metrics.dedup_removed = dedup_removed;

        (selected, metrics)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn test_chunk(id: &str, text: &str, score: f32, size: usize) -> OptimizableChunk {
        OptimizableChunk {
            id: id.to_string(),
            text: text.to_string(),
            score,
            confidence: score * 0.9,
            size_bytes: size,
        }
    }

    #[test]
    fn test_score_threshold_filter() {
        let filter = ScoreThresholdFilter::new(0.6);
        let chunks = vec![
            test_chunk("doc1", "high score", 0.9, 100),
            test_chunk("doc2", "low score", 0.3, 100),
            test_chunk("doc3", "medium score", 0.65, 100),
        ];

        let filtered = filter.filter(chunks);
        assert_eq!(filtered.len(), 2);
        assert!(filtered.iter().all(|c| c.score >= 0.6));
    }

    #[test]
    fn test_budget_selector_within_budget() {
        let selector = BudgetSelector::new(500, 0.5);
        let chunks = vec![
            test_chunk("doc1", "text1", 0.9, 100),
            test_chunk("doc2", "text2", 0.8, 100),
            test_chunk("doc3", "text3", 0.7, 100),
        ];

        let (selected, metrics) = selector.select(chunks);
        assert_eq!(selected.len(), 3);
        assert_eq!(metrics.total_bytes, 300);
        assert!(metrics.budget_used_pct < 100.0);
    }

    #[test]
    fn test_budget_selector_over_budget() {
        let selector = BudgetSelector::new(150, 0.5);
        let chunks = vec![
            test_chunk("doc1", "text1", 0.9, 100),
            test_chunk("doc2", "text2", 0.8, 100),
            test_chunk("doc3", "text3", 0.7, 100),
        ];

        let (selected, metrics) = selector.select(chunks);
        assert!(selected.len() < 3);
        assert!(metrics.total_bytes <= 150);
    }

    #[test]
    fn test_budget_selector_threshold() {
        let selector = BudgetSelector::new(500, 0.7);
        let chunks = vec![
            test_chunk("doc1", "text1", 0.9, 100),
            test_chunk("doc2", "text2", 0.5, 100), // Below threshold
            test_chunk("doc3", "text3", 0.8, 100),
        ];

        let (selected, metrics) = selector.select(chunks);
        assert_eq!(selected.len(), 2);
        assert_eq!(metrics.rejected_count, 1);
    }

    #[test]
    fn test_shingle_deduplicator_identical() {
        let dedup = ShingleDeduplicator::new(0.8, 3);
        let chunks = vec![
            test_chunk("doc1", "the quick brown fox", 0.9, 100),
            test_chunk("doc2", "the quick brown fox", 0.8, 100), // Identical
        ];

        let (kept, removed) = dedup.deduplicate(chunks);
        assert_eq!(kept.len(), 1);
        assert_eq!(removed, 1);
        assert_eq!(kept[0].id, "doc1"); // Kept highest score
    }

    #[test]
    fn test_shingle_deduplicator_different() {
        let dedup = ShingleDeduplicator::new(0.8, 3);
        let chunks = vec![
            test_chunk("doc1", "kubernetes pod debugging", 0.9, 100),
            test_chunk("doc2", "docker container deployment", 0.8, 100),
        ];

        let (kept, removed) = dedup.deduplicate(chunks);
        assert_eq!(kept.len(), 2);
        assert_eq!(removed, 0);
    }

    #[test]
    fn test_shingle_deduplicator_partial_overlap() {
        let dedup = ShingleDeduplicator::new(0.3, 2);
        let chunks = vec![
            test_chunk("doc1", "kubernetes pod debugging", 0.9, 100),
            test_chunk("doc2", "kubernetes deployment guide", 0.8, 100),
        ];

        let (kept, removed) = dedup.deduplicate(chunks);
        // Both share "kubernetes" shingle, but not enough overlap at 0.3
        assert!(kept.len() <= 2);
    }

    #[test]
    fn test_chunk_optimizer_full_pipeline() {
        let optimizer = ChunkOptimizer::new(0.6, 200, 0.8);
        let chunks = vec![
            test_chunk("doc1", "high score chunk", 0.9, 100),
            test_chunk("doc2", "low score chunk", 0.3, 100),
            test_chunk("doc3", "medium score chunk", 0.7, 100),
        ];

        let (selected, metrics) = optimizer.optimize(chunks);
        assert!(selected.len() > 0);
        assert!(metrics.avg_score >= 0.6);
        assert!(metrics.budget_used_pct <= 100.0);
    }

    #[test]
    fn test_selection_metrics_calculation() {
        let selector = BudgetSelector::new(500, 0.5);
        let chunks = vec![
            test_chunk("doc1", "text1", 0.9, 100),
            test_chunk("doc2", "text2", 0.8, 100),
        ];

        let (selected, metrics) = selector.select(chunks);
        assert_eq!(metrics.selected_count, 2);
        assert_eq!(metrics.total_bytes, 200);
        assert!(metrics.budget_used_pct > 0.0);
        assert!(metrics.avg_score > 0.0);
    }
}