/// 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) -> Vec { 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, ) -> (Vec, 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::() / 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 { 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) -> (Vec, 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) -> (Vec, 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); } }