Project: Query Optimizer System

Project overview

The Module 3 capstone project: implement 4 query optimization techniques (Expansion, Rewriting, Decomposition, HyDE), compare them with benchmarks, pick the optimal technique, and integrate it into the baseline RAG.

Goal: improve recall by +15-25% vs direct queries, document the trade-offs (latency, cost), and end up with production-ready code.

---

🎯 Project goals

1. ✅ Implement 4 query optimization techniques
2. ✅ Compare them with benchmarks (recall, precision, latency, cost)
3. ✅ Pick the optimal technique based on the requirements
4. ✅ Integrate it into the baseline RAG (replace the direct queries)
5. ✅ Document the gain (+15-25% recall expected)

---

📁 Project structure

rag_baseline_project/  (from Modules 1-2)
├── src/
│   ├── query_optimization/          # NEW
│   │   ├── __init__.py
│   │   ├── base_optimizer.py        # Base interface
│   │   ├── query_expansion.py       # Technique 1
│   │   ├── query_rewriting.py       # Technique 2
│   │   ├── query_decomposition.py   # Technique 3
│   │   ├── hyde.py                  # Technique 4
│   │   └── optimizer_comparator.py  # Benchmark tool
│   ├── retrieval.py                 # Update with optimization
│   └── ...
└── notebooks/
    └── query_optimization_demo.ipynb

---

💻 Implementation

Step 1: Base Optimizer Interface

# src/query_optimization/base_optimizer.py
from abc import ABC, abstractmethod

class BaseQueryOptimizer(ABC):
    """Base interface for query optimization"""
    
    @abstractmethod
    def optimize(self, query: str) -> str | list[str]:
        """Optimize the query"""
        pass
    
    @abstractmethod
    def name(self) -> str:
        """Name of the technique"""
        pass

---

Step 2: Implement the techniques

# src/query_optimization/query_expansion.py
from openai import OpenAI
from .base_optimizer import BaseQueryOptimizer

class QueryExpansion(BaseQueryOptimizer):
    def __init__(self, num_expansions: int = 5):
        self.num_expansions = num_expansions
        self.client = OpenAI()
    
    def optimize(self, query: str) -> list[str]:
        """Generate several related queries"""
        
        prompt = f"""Generate {self.num_expansions} diverse search queries related to: "{query}"

Output format (one query per line):
1. [query]
2. [query]
..."""

        response = self.client.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=[
                {"role": "system", "content": "You are a search query expert."},
                {"role": "user", "content": prompt}
            ],
            temperature=0.7
        )
        
        # Parse the queries
        lines = response.choices[0].message.content.strip().split('\n')
        queries = [line.split('. ', 1)[1] if '. ' in line else line for line in lines if line]
        
        return queries[:self.num_expansions]
    
    def name(self) -> str:
        return "query_expansion"

(Implement something similar for query_rewriting.py, query_decomposition.py, hyde.py)

---

Step 3: Query Optimizer Comparator

# src/query_optimization/optimizer_comparator.py
import time
import statistics
from typing import Dict

class QueryOptimizerComparator:
    """Compare query optimization techniques with benchmarks"""
    
    def __init__(self, optimizers: list, rag_system):
        self.optimizers = optimizers
        self.rag_system = rag_system
    
    def compare(
        self,
        test_queries: list[str],
        ground_truth: dict
    ) -> Dict:
        """Compare all the techniques"""
        results = {}
        
        for optimizer in self.optimizers:
            print(f"\nBenchmarking {optimizer.name()}...")
            
            recalls = []
            precisions = []
            latencies = []
            
            for query in test_queries:
                # Measure latency
                start = time.time()
                
                # Optimize the query
                optimized = optimizer.optimize(query)
                
                # Search (handle a single query or a list)
                if isinstance(optimized, list):
                    # Multiple queries (expansion, decomposition)
                    all_results = []
                    for q in optimized:
                        res = self.rag_system.retrieve(q, top_k=10)
                        all_results.append(res)
                    # Merge with RRF
                    final_results = self._rrf_merge(all_results)[:5]
                else:
                    # Single query (rewriting, hyde)
                    final_results = self.rag_system.retrieve(optimized, top_k=5)
                
                latency = (time.time() - start) * 1000  # ms
                latencies.append(latency)
                
                # Compute precision/recall
                retrieved = set(final_results['ids'])
                relevant = set(ground_truth[query]['relevant_docs'])
                
                relevant_retrieved = retrieved & relevant
                precision = len(relevant_retrieved) / len(retrieved) if retrieved else 0
                recall = len(relevant_retrieved) / len(relevant) if relevant else 0
                
                precisions.append(precision)
                recalls.append(recall)
            
            # Aggregates
            results[optimizer.name()] = {
                "recall_avg": statistics.mean(recalls),
                "precision_avg": statistics.mean(precisions),
                "latency_p50": statistics.median(latencies),
                "latency_p95": statistics.quantiles(latencies, n=20)[18] if len(latencies) >= 20 else max(latencies)
            }
        
        return results
    
    def _rrf_merge(self, results_list: list, k: int = 60) -> list:
        """Reciprocal Rank Fusion for the merge"""
        scores = {}
        for results in results_list:
            for rank, doc_id in enumerate(results['ids'], 1):
                if doc_id not in scores:
                    scores[doc_id] = 0
                scores[doc_id] += 1 / (k + rank)
        
        ranked = sorted(scores.items(), key=lambda x: x[1], reverse=True)
        return {"ids": [doc_id for doc_id, score in ranked]}

---

Step 4: Run the benchmark

# benchmark_query_optimization.py
from src.query_optimization.query_expansion import QueryExpansion
from src.query_optimization.query_rewriting import QueryRewriting
from src.query_optimization.query_decomposition import QueryDecomposition
from src.query_optimization.hyde import HyDE
from src.query_optimization.optimizer_comparator import QueryOptimizerComparator
from src.rag_system import BaselineRAGSystem

# Set up the optimizers
optimizers = [
    QueryExpansion(num_expansions=5),
    QueryRewriting(),
    QueryDecomposition(),
    HyDE()
]

# RAG system
rag_system = BaselineRAGSystem()

# Test queries with ground truth
test_queries = [
    "How to implement authentication in FastAPI?",
    "What are FastAPI performance optimizations?",
    "Compare FastAPI and Flask",
    # ... 20 more
]

ground_truth = {
    "How to implement authentication in FastAPI?": {
        "relevant_docs": ["doc_123", "doc_456", "doc_789"]
    },
    # ... for all the queries
}

# Compare
comparator = QueryOptimizerComparator(optimizers, rag_system)
results = comparator.compare(test_queries, ground_truth)

# Show the results
print("\n" + "="*60)
print("QUERY OPTIMIZATION BENCHMARK RESULTS")
print("="*60)

for optimizer_name, metrics in results.items():
    print(f"\n{optimizer_name.upper()}:")
    print(f"  Recall: {metrics['recall_avg']:.2%}")
    print(f"  Precision: {metrics['precision_avg']:.2%}")
    print(f"  Latency P50: {metrics['latency_p50']:.0f}ms")
    print(f"  Latency P95: {metrics['latency_p95']:.0f}ms")

Expected output:

============================================================
QUERY OPTIMIZATION BENCHMARK RESULTS
============================================================

QUERY_EXPANSION:
  Recall: 72%
  Precision: 80%
  Latency P50: 820ms
  Latency P95: 1,050ms

QUERY_REWRITING:
  Recall: 60%
  Precision: 88%
  Latency P50: 680ms
  Latency P95: 850ms

QUERY_DECOMPOSITION:
  Recall: 65%
  Precision: 85%
  Latency P50: 1,180ms
  Latency P95: 1,450ms

HYDE:
  Recall: 77%
  Precision: 82%
  Latency P50: 880ms
  Latency P95: 1,120ms

---

Step 5: Integrate into the baseline RAG

# src/retrieval.py (updated)
from src.query_optimization.query_expansion import QueryExpansion

class OptimizedRetrievalPipeline:
    def __init__(self, optimization_technique: str = "expansion"):
        # Pick the technique
        if optimization_technique == "expansion":
            self.optimizer = QueryExpansion()
        elif optimization_technique == "rewriting":
            self.optimizer = QueryRewriting()
        elif optimization_technique == "hyde":
            self.optimizer = HyDE()
        else:
            self.optimizer = None  # Fall back to the direct query
        
        # ... rest of the setup
    
    def retrieve(self, query: str, top_k: int = 5):
        """Retrieval with query optimization"""
        
        if self.optimizer:
            # Optimize the query
            optimized = self.optimizer.optimize(query)
            
            if isinstance(optimized, list):
                # Multiple queries → RRF merge
                all_results = []
                for q in optimized:
                    results = self._search(q, top_k=10)
                    all_results.append(results)
                return self._rrf_merge(all_results)[:top_k]
            else:
                # Single query
                return self._search(optimized, top_k=top_k)
        else:
            # No optimization (baseline)
            return self._search(query, top_k=top_k)

---

📊 Measuring the gain

Before (Modules 1-2 baseline):

# Direct queries + optimized chunking
baseline_results = {
    "recall@50": 0.57,
    "precision@5": 0.78
}

After (Module 3 optimized):

# Query expansion + optimized chunking
optimized_results = {
    "recall@50": 0.72,  # +15%
    "precision@5": 0.80  # +2%
}

Gain achieved: ✅ +15% recall (target hit)

---

📝 Documenting the decision

Updated README.md:

# RAG System - Module 3: Query Optimization

## Query Optimization Technique Selected: Query Expansion

### Rationale:
- **Gain:** +15% recall, +2% precision vs direct queries
- **Cost:** +$0.001/query (acceptable)
- **Latency:** +650ms (within the 1s budget)
- **Rationale:** The best recall/cost/latency balance

### Alternatives Considered:
- **HyDE:** +20% recall but +$0.0015/query → Dropped on cost
- **Rewriting:** +10% precision but +3% recall → Not enough for the recall target
- **Decomposition:** +20% precision on complex ones but +1s latency → Too slow

### Configuration:
```python
QueryExpansion(num_expansions=5)

Cumulative Performance (Modules 1-3)

Metric	Module 1	Module 2	Module 3	Total Improvement
Precision@5	68%	78% (+10%)	80% (+2%)	+12%
Recall@50	52%	57% (+5%)	72% (+15%)	+20%

---

## 🎯 Success criteria

✅ **The project is complete if:**

1. 4 techniques implemented with working code
2. The comparative benchmark has been run
3. A technique selected, with justification
4. Integration into the baseline RAG is complete
5. The +15-25% recall gain vs Module 2 has been achieved

---

## 🎯 Recap

**Query Optimizer project:**

- ✅ Implement 4 techniques (Expansion, Rewriting, Decomposition, HyDE)
- ✅ Compare with quantitative benchmarks
- ✅ Pick the optimal technique (typically: Expansion)
- ✅ Improve the baseline by +15-25% recall
- ✅ Document the trade-offs (latency, cost)

**Expected gain:** Recall 57% → 72% (+15%)

**Next module:** Module 4 optimizes retrieval with re-ranking (+20-25% precision after the initial retrieval).

---

**Created:** February 6, 2026  
**Version:** 1.0