poimen-memory/docs/QUERY_OPTIMIZATION_ENGINE.md

Query Optimization Engine: Design & Implementation Guide

Executive Summary

This document specifies the Query Optimization Engine and Query Context Constructor for Poimen Memory's hybrid search. We are implementing Approach A (Parallel RRF) for:

• ✅ Highest accuracy (mission-critical for agent reasoning)
• ✅ Fault tolerance (hybrid + semantic-only fallback)
• ✅ No false negatives (semantic catches synonyms lexical misses)
• ✅ Transparent scoring (debug + optimize)
• ✅ Decoupled systems (embedding model changes don't break architecture)

---

Architecture Overview

User Query
  │
  ├─ QueryOptimizer (Query Context Constructor)
  │  ├─ Normalize
  │  ├─ Tokenize
  │  ├─ Extract entities (years, names, keywords)
  │  ├─ Analyze characteristics (special syntax, dates, negation)
  │  ├─ Classify question type (procedural, factual, troubleshooting, etc)
  │  └─ Route to search strategy (Hybrid, Semantic, Lexical, or Cascading)
  │
  └─ HybridQueryWorker (Parallel Retrieval Orchestration)
     ├─ Generate embedding (LLM)
     │
     ├─ Execute Strategy
     │  ├─ HYBRID: Parallel pgvector + OpenSearch
     │  ├─ SEMANTIC: pgvector only (fallback)
     │  ├─ LEXICAL: OpenSearch only (fallback)
     │  └─ CASCADING: OpenSearch (narrow) → pgvector (rerank)
     │
     ├─ Fuse Results (RRF Algorithm)
     │  ├─ Normalize scores to [0-1] range
     │  ├─ Apply RRF formula: 1 / (k + rank)
     │  └─ Merge and re-rank
     │
     └─ Return Response
        └─ Top-10 results with score breakdown + metrics

---

Part 1: Query Context Constructor (QueryOptimizer)

Why This Matters

Before executing a query against the hybrid system, we need to:

1. Understand the query intent
2. Extract relevant context (years, entities, exact phrases)
3. Classify the question type
4. Route to the best search strategy

This prevents "garbage in, garbage out" retrieval.

Stage 1: Query Normalization

Input:  "  HOW do I FIX  kubernetes PORT 8080 ???  "
Output: "how do i fix kubernetes port 8080 ???"

Purpose:

• Lowercase for consistent matching
• Trim whitespace
• Remove double spaces

---

Stage 2: Tokenization

Input:  "how do i fix kubernetes port 8080"
Output: ["how", "do", "i", "fix", "kubernetes", "port", "8080"]

Tokens used for:

• Entity extraction
• Date filter detection
• Negation detection
• Token count heuristic

---

Stage 3: Entity Extraction

Extract structured information from query:

{
  "year": "2024",
  "exact_phrase": "kubernetes port forwarding",
  "dates": ["january", "2024"],
  "tags": ["#networking", "@devops"]
}

Entities detected:

• Years: YYYY format (2000-2100)
• Quoted phrases: "exact text" (preserve for lexical search)
• Date keywords: month names, relative dates ("this month", "last week")
• Tags: #hashtags, @mentions (preserve for special syntax routing)

---

Stage 4: Query Characteristic Analysis

pub struct QueryContext {
    pub has_special_syntax: bool,    // #tag, @mention, "phrase"
    pub has_date_filters: bool,      // years, month names
    pub has_negation: bool,          // -word, "NOT", "no"
    pub token_count: usize,          // ["how", "do", "i", ...].len()
    pub question_type: QuestionType, // Factual, Procedural, etc
    pub search_strategy: SearchStrategy, // Hybrid, Semantic, etc
}

Importance:

• Short queries (< 3 tokens) → lexical better
• Special syntax (#tag) → preserve exact matches
• Date filters → use cascading (lexical to narrow, then semantic)
• Negation → more complex semantic reasoning needed

---

Stage 5: Question Type Classification

Classify into one of 6 types:

Type	Indicators	Best Strategy	Example
Factual	"What", "Define", "When"	Hybrid	"What is Kubernetes?"
Procedural	"How", "Steps", "Guide"	Hybrid ⭐⭐	"How do I deploy to K8s?"
Comparative	"Compare", "Difference", "vs"	Hybrid	"Compare Docker vs Kubernetes"
Troubleshooting	"Fix", "Error", "Debug", "Broken"	Hybrid ⭐⭐	"Fix port 8080 conflict"
Navigational	"Where", "Find", "Show"	Lexical First	"Where is the deployment guide?"
Open	General, conversational	Hybrid	"Tell me about networking"

Classification algorithm:

fn classify_question(raw_query: &str, tokens: &[String]) -> QuestionType {
    match first_token {
        "how" => Procedural,
        "what" => {
            if contains("difference") { Comparative } else { Factual }
        }
        "fix" | "error" => Troubleshooting,
        "where" | "find" => Navigational,
        _ => Open,
    }
}

---

Stage 6: Search Strategy Routing

Decision tree to choose optimal search strategy:

Query Characteristics
  │
  ├─ Token count < 3?
  │  └─ YES → LEXICAL_ONLY (keywords better than embeddings)
  │
  ├─ Has special syntax (#tag, "phrase")?
  │  ├─ YES + date filters → LEXICAL_FIRST (narrow by keywords, rerank by semantic)
  │  └─ YES alone → LEXICAL_ONLY (preserve exact syntax)
  │
  ├─ Has date filters?
  │  └─ YES → LEXICAL_FIRST (OpenSearch to filter by date, pgvector reranks)
  │
  └─ Question Type?
     ├─ Procedural → HYBRID ⭐⭐ (need both exact steps + understanding)
     ├─ Troubleshooting → HYBRID ⭐⭐ (need errors + semantic understanding)
     ├─ Navigational → LEXICAL_FIRST (find specific docs, then semantic rank)
     └─ Others → HYBRID (good default)

Output:

SearchStrategy::Hybrid  // or Semantic, Lexical, LexicalFirst
confidence: 0.95       // 0.0-1.0, how confident in routing decision

Confidence scores:

• Date filters + cascading: 0.9 (high confidence)
• Procedural questions: 0.95 (very high confidence)
• Special syntax: 0.8-0.85 (moderate, but clear signal)
• Generic queries: 0.8 (default, reasonable)

---

Part 2: Hybrid Query Worker (Parallel Orchestration)

Overview

The HybridQueryWorker executes the query against selected engines and fuses results.

Flow: HYBRID Strategy

Input: "How do I fix kubernetes port conflict?"
       Strategy: HYBRID, Confidence: 0.95

       ↓ Generate Embedding (LLM)
       ↓
    ┌──────────────────────────────────────┐
    │   PARALLEL Execution                 │
    ├──────────────────────────────────────┤
    │                                      │
    │ Thread 1: pgvector                   │
    │ ├─ Query: embedding <-> vector       │
    │ ├─ Filter: project_id, level, date   │
    │ └─ Retrieve: Top-50 chunks           │
    │    Results: [(doc1, 0.95), ...]      │
    │                                      │
    │ Thread 2: OpenSearch + JWT           │
    │ ├─ Query: multi_match BM25           │
    │ ├─ Filter: project_id, level, date   │
    │ └─ Retrieve: Top-50 documents        │
    │    Results: [(doc1, 8.5), ...]       │
    │                                      │
    └──────────────────────────────────────┘
       ↓ ~200-300ms total
       ↓
    Normalize Scores
    ├─ pgvector: already [0.0, 1.0]
    └─ OpenSearch: min-max to [0.0, 1.0]
       ↓
    RRF Fusion
    ├─ doc1: 1/(60+1) + 1/(60+1) = 0.033
    ├─ doc2: 1/(60+2) + 0 = 0.016
    └─ doc3: 0 + 1/(60+1) = 0.016
       ↓ Sort by score
       ↓
    Return Top-10: [doc1, doc2, doc3, ...]

Flow: CASCADING Strategy (Lexical First)

Input: "Fix kubernetes networking in 2024"
       Strategy: LEXICAL_FIRST, Confidence: 0.9

       ↓ Generate Embedding (LLM)
       ↓
    Stage 1: OpenSearch (Narrow)
    ├─ Query: multi_match + filter(year=2024)
    ├─ Retrieve: Top-200 documents
    └─ Extract chunk_ids: [id1, id2, id3, ...]
       ↓
    Stage 2: pgvector (Rerank)
    ├─ Query: embedding <-> vector WHERE chunk_id IN (top-200)
    ├─ Retrieve: Top-10 from narrowed set
    └─ Results: [(id1, 0.95), (id3, 0.92), ...]
       ↓
    Return Top-10 Results

Benefit of cascading:

• OpenSearch quickly filters by date/keywords (100ms)
• pgvector only ranks the 200 most relevant (not all 10,000+)
• Much faster than hybrid (150ms vs 300ms)
• No false negatives (semantic still sees all matches)

---

Part 3: Reciprocal Rank Fusion (RRF)

Why RRF?

OpenSearch BM25 and pgvector cosine distance use completely different mathematical distributions. You cannot simply average them.

Example:

pgvector scores:     [0.95, 0.88, 0.82, ...]  (0.0 to 1.0)
OpenSearch scores:   [8.5, 7.2, 6.1, ...]     (0.0 to 50+, unbounded)

Naive average would bias towards OpenSearch (much larger numbers)

RRF Solution: Convert both to ranks (positions), then fuse ranks:

pgvector results:
  Rank 1: doc1 (score 0.95)
  Rank 2: doc2 (score 0.88)
  Rank 3: doc3 (score 0.82)

OpenSearch results:
  Rank 1: doc1 (score 8.5)
  Rank 2: doc4 (score 7.2)
  Rank 3: doc2 (score 6.8)

RRF Fusion:
  For each document, calculate: 1 / (k + rank)
  where k = 60 (constant)
  
  doc1: 1/(60+1) + 1/(60+1) = 0.0164 + 0.0164 = 0.0328
  doc2: 1/(60+2) + 1/(60+3) = 0.0159 + 0.0158 = 0.0317
  doc4: 1/(60+2) = 0.0159
  doc3: 1/(60+3) = 0.0158
  
  Sorted: [doc1 (0.0328), doc2 (0.0317), doc4 (0.0159), doc3 (0.0158)]

RRF Algorithm (Rust)

pub fn fuse_rrf(
    semantic_results: Vec<(String, f32)>,  // [(id, score), ...]
    lexical_results: Vec<(String, f32)>,
    k: f32,                                 // Usually 60
    final_k: usize,                         // Return top-k (usually 10)
) -> Vec<(String, f32)> {
    let mut fused_scores: HashMap<String, f32> = HashMap::new();
    
    // Add semantic ranks
    for (rank, (id, _)) in semantic_results.into_iter().enumerate() {
        let rrf_score = 1.0 / (k + (rank as f32) + 1.0);
        fused_scores.insert(id, rrf_score);
    }
    
    // Add lexical ranks (combine if already present)
    for (rank, (id, _)) in lexical_results.into_iter().enumerate() {
        let rrf_score = 1.0 / (k + (rank as f32) + 1.0);
        *fused_scores.entry(id).or_insert(0.0) += rrf_score;
    }
    
    // Sort by combined score
    let mut results: Vec<_> = fused_scores.into_iter().collect();
    results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
    results.truncate(final_k);
    
    results
}

Why RRF Over Weighted Linear?

Factor	RRF	Weighted Linear (0.6/0.4)
Stability	✅ Rank-based, not score-value-dependent	⚠️ Sensitive to score magnitude differences
Parameter tuning	❌ None needed	✅ Can tune 0.6/0.4 weights
Academic backing	✅ Proven in information retrieval	⚠️ Arbitrary without data
Robustness	✅ Works if embedding model changes	⚠️ May need re-tuning
Debugging	✅ Clear: "in top 1, top 3"	⚠️ Harder: "score 0.88 vs 0.96"

Recommendation: Use RRF by default. If A/B testing shows 0.6/0.4 weighted linear works better, switch to that.

---

Part 4: Response Format

API Response with Score Breakdown

{
  "query": "how do i fix kubernetes port 8080",
  "project": "poimen",
  "search_strategy": "Hybrid",
  "strategy_confidence": 0.95,
  
  "results": [
    {
      "id": "chunk-123",
      "rank": 1,
      "final_score": 0.0328,
      "semantic_score": 0.95,
      "lexical_score": 8.5,
      "fusion_method": "rrf",
      "retrieval_engine": "Hybrid",
      
      "text": "kubectl port-forward service port:8080...",
      "source": "runbooks/kubernetes/networking.md",
      "level": "L1",
      "breadcrumb": ["runbooks", "kubernetes", "troubleshooting"],
      
      "score_breakdown": {
        "semantic_rank": 1,
        "lexical_rank": 1,
        "rrf_components": {
          "semantic_contribution": 0.0164,
          "lexical_contribution": 0.0164
        }
      }
    }
  ],
  
  "metrics": {
    "total_time_ms": 245,
    "semantic_time_ms": 120,
    "lexical_time_ms": 118,
    "fusion_time_ms": 7,
    "semantic_results_count": 50,
    "lexical_results_count": 50,
    "final_results_count": 10
  }
}

---

Part 5: Integration with Existing API

Current Endpoint

GET /memory/query?query=...&project=...&limit=...

Updated Implementation

Replace:

state.query_worker.query(&project, &question, Some(limit))

With:

state.hybrid_query_worker.query(
    &project,
    &question,
    limit,
    &jwt_token,  // Forward to OpenSearch
)

Fallback Strategy

async fn query_handler(...) -> HttpResponse {
    // Try hybrid
    match state.hybrid_query_worker.query(...).await {
        Ok(response) => return HttpResponse::Ok().json(response),
        Err(e) => {
            tracing::warn!("Hybrid query failed: {}, falling back to semantic", e);
            
            // Fallback: semantic only
            match state.query_worker.query(...).await {
                Ok(results) => return HttpResponse::Ok().json(results),
                Err(e2) => return HttpResponse::InternalServerError().json(...)
            }
        }
    }
}

This ensures availability:

• Hybrid works → Use hybrid (best accuracy)
• OpenSearch down → Fallback to semantic (still good)
• Both down → Error (clear signal)

---

Part 6: Implementation Phases

Phase 1: Query Optimization Engine (Week 1)

Deliverables:

• ✅ QueryOptimizer (6-stage pipeline)
• ✅ QueryContext data structure
• ✅ Question classification
• ✅ Search strategy routing
• ✅ RRFusion algorithm
• ✅ Unit tests (15+ tests)

Files:

• crates/mem-cli/src/query_optimizer.rs (16KB)
• crates/mem-cli/src/hybrid_query_worker.rs (13KB)

Status: ✅ COMPLETE (code committed)

---

Phase 2: Hybrid Query Integration (Week 2)

Tasks:

1. Update /memory/query handler to use HybridQueryWorker
2. Add fallback strategy (hybrid → semantic → error)
3. Forward JWT token to OpenSearch
4. Return score breakdown in response
5. Add metrics to response
6. Integration tests (10+ tests)

Files to modify:

• crates/mem-cli/src/http_server.rs
• tests/it_hybrid_query.rs (new)

Timeline: 2-3 days

---

Phase 3: Performance Optimization (Week 3)

Tasks:

1. Measure baseline latency (semantic vs lexical vs hybrid)
2. Optimize pgvector query (HNSW index tuning)
3. Optimize OpenSearch query (field boosts, analyzers)
4. Add query caching (1hr TTL)
5. Benchmark with 1000-query test set

Metrics to track:

• Semantic search: 80-120ms
• Lexical search: 60-100ms
• Hybrid search: 150-250ms (parallel)
• Cascading search: 100-180ms

---

Phase 4: Testing & Validation (Week 4)

Test Fixtures: Create query dataset with ground truth

queries:
  - query: "How do I fix kubernetes port 8080?"
    expected_docs: ["runbooks/networking.md", "docs/troubleshooting.md"]
    min_ndcg: 0.85
    question_type: Procedural
    expected_strategy: Hybrid
  
  - query: "#networking @devops"
    expected_docs: ["docs/network-policies.md"]
    min_ndcg: 0.9
    question_type: Navigational
    expected_strategy: LexicalOnly
  
  - query: "fix port 2024"
    expected_docs: ["runbooks/deployment-2024.md"]
    min_ndcg: 0.8
    question_type: Troubleshooting
    expected_strategy: LexicalFirst

Accuracy metrics:

• NDCG@10 (ranking quality)
• MRR (how early is first correct result)
• Precision@5, @10
• Recall@10

A/B Testing:

• Hybrid vs semantic-only
• RRF (k=60) vs weighted linear (0.6/0.4)
• Different question types

---

Part 7: Configuration & Tuning

Environment Variables

# Query optimization
QUERY_OPTIMIZATION_ENABLED=true
QUERY_QUESTION_CLASSIFICATION=true
QUERY_ENTITY_EXTRACTION=true

# RRF configuration
RRF_K_CONSTANT=60
RRF_RETRIEVE_K=50        # Top-50 from each engine
RRF_FINAL_K=10           # Return top-10

# Cascading strategy
CASCADING_LEXICAL_MULTIPLIER=4  # Fetch 4x results in stage 1

# Query caching
QUERY_CACHE_TTL_SECS=3600
QUERY_CACHE_MAX_SIZE=10000

# Fallback strategy
ENABLE_SEMANTIC_FALLBACK=true

RRF Tuning

k parameter (constant):

• Lower k → Earlier ranks weighted more
• k=60 is standard (academic consensus)
• k=20-30 → Aggressively weights top results
• k=100-200 → More uniform weighting

retrieve_k (top-K from each engine):

• Default: 50
• Can increase to 100 for more diversity
• Trade-off: 50 is good balance (speed vs coverage)

final_k (return top-K):

• Default: 10
• Agents typically use 5-10
• Can lower to 5 for faster inference

---

Part 8: Testing Checklist

Unit Tests

#[tokio::test]
async fn test_query_optimizer_procedural() { }

#[tokio::test]
async fn test_query_optimizer_short_query() { }

#[tokio::test]
async fn test_query_optimizer_special_syntax() { }

#[tokio::test]
async fn test_rrf_fusion_basic() { }

#[tokio::test]
async fn test_rrf_fusion_single_engine() { }

#[tokio::test]
async fn test_hybrid_query_integration() { }

#[tokio::test]
async fn test_cascading_query_integration() { }

#[tokio::test]
async fn test_fallback_semantic_only() { }

Integration Tests

#[tokio::test]
async fn test_end_to_end_hybrid_query() {
    // Setup: ingest test documents
    // Execute: /memory/query with hybrid strategy
    // Verify: top result is expected document
    // Assert: NDCG >= 0.85
}

#[tokio::test]
async fn test_cascading_vs_hybrid() {
    // Compare latency and accuracy
    // Cascading should be faster
    // Accuracy should be similar
}

A/B Testing Queries

test_dataset:
  - category: "Procedural"
    queries:
      - "How do I deploy to Kubernetes?"
      - "Steps to fix port conflicts"
      - "Deploy application guide"
    
  - category: "Troubleshooting"
    queries:
      - "Fix OOMKilled error"
      - "Debug networking issue"
      - "Resolve timeout errors"
  
  - category: "Navigational"
    queries:
      - "Where is the deployment runbook?"
      - "Find kubernetes best practices"
      - "#networking documents"
  
  - category: "Factual"
    queries:
      - "What is a StatefulSet?"
      - "Define PVC"
      - "What does idempotency mean?"

---

Summary: Architecture Decision

Aspect	Approach A (Parallel RRF)
Retrieval Method	Parallel pgvector + OpenSearch
Fusion Algorithm	RRF (Reciprocal Rank Fusion)
Accuracy	⭐⭐⭐⭐⭐ Highest
Latency	⭐⭐⭐ Moderate (150-250ms for hybrid)
Complexity	Moderate (RRF logic, parallel orchestration)
Fault Tolerance	✅ Fallback to semantic if OpenSearch down
Debugging	✅ Clear score breakdown
Recommended for	Mission-critical LLM agent reasoning

Status: ✅ DESIGN COMPLETE, CODE IMPLEMENTED

Next: Phase 2 (Integrate into /memory/query endpoint)