06: Retrieval Quality Testing

Response Quality ≤ min(Retrieval Quality, Generation Quality)

User Query → Intent Classifier → Query Rewriter → Titan Embeddings v2 →
OpenSearch Serverless (k-NN + BM25 hybrid) → Re-ranker → Top-5 Chunks → LLM

graph TD
    subgraph "Test Data"
        A[Golden Retrieval Dataset<br>300 queries with labeled chunks] --> B[Test Runner]
    end

    subgraph "Retrieval Pipeline Under Test"
        B --> C[Query Rewriter]
        C --> D[Embedding Generator<br>Titan Embeddings v2]
        D --> E[OpenSearch Serverless<br>Hybrid k-NN + BM25]
        E --> F[Re-ranker]
        F --> G[Top-K Retrieved Chunks]
    end

    subgraph "Quality Metrics Engine"
        G --> H[Relevance Scorer<br>Per-chunk relevance]
        G --> I[Precision@K / Recall@K<br>Against ground truth]
        G --> J[MRR / nDCG@K<br>Ranking quality]
        G --> K[Latency Tracker<br>P50, P95, P99]
    end

    subgraph "Stage-Level Diagnostics"
        C --> L[Query Rewrite Quality<br>Intent preservation check]
        D --> M[Embedding Quality<br>Cosine similarity distribution]
        E --> N[Search Quality<br>k-NN vs BM25 contribution]
        F --> O[Re-ranker Quality<br>Position changes analysis]
    end

    subgraph "Outputs"
        H --> P[Retrieval Quality Report]
        I --> P
        J --> P
        K --> P
        L --> P
        M --> P
        N --> P
        O --> P
        P --> Q[CloudWatch Dashboard]
        P --> R[Degradation Alerts]
    end

graph LR
    subgraph "Experiment Setup"
        A[Same Query Set] --> B{A/B Router}
        B --> C[Config A: Titan v2<br>k=5, hybrid 0.7/0.3]
        B --> D[Config B: Titan v2<br>k=7, hybrid 0.5/0.5]
    end

    subgraph "Evaluation"
        C --> E[Retrieve & Score]
        D --> E
        E --> F[Paired Comparison<br>Same query, both configs]
    end

    subgraph "Decision"
        F --> G{Config B better<br>on MRR + Latency?}
        G -->|Yes, significant| H[Adopt Config B]
        G -->|No| I[Keep Config A]
    end

sequenceDiagram
    participant QR as Query Rewriter
    participant EM as Embedding Generator
    participant OS as OpenSearch
    participant RR as Re-ranker
    participant Total as Total Latency

    Note over QR,Total: Target: < 500ms total

    QR->>QR: Rewrite query<br>~50ms
    QR->>EM: Rewritten query
    EM->>EM: Generate embedding<br>~30ms (Titan v2)
    EM->>OS: 1536-dim vector
    OS->>OS: k-NN search + BM25<br>~80ms (warm) / ~300ms (cold)
    OS->>RR: Top-20 candidates
    RR->>RR: Re-rank to top-5<br>~40ms
    RR->>Total: Final chunks

    Note over Total: Warm: ~200ms | Cold: ~420ms | P99: ~600ms

from dataclasses import dataclass, field
from enum import Enum
from typing import Optional
import time
import uuid


@dataclass
class RetrievalTestCase:
    """A test case for retrieval quality evaluation."""
    test_id: str = field(default_factory=lambda: str(uuid.uuid4()))
    query: str = ""
    intent: str = ""
    expected_chunk_ids: list[str] = field(default_factory=list)    # Ground truth chunks
    expected_key_facts: list[str] = field(default_factory=list)    # Facts that must appear
    min_relevance_threshold: float = 0.6                           # Minimum acceptable relevance
    difficulty: str = "medium"                                      # easy, medium, hard, adversarial


@dataclass
class RetrievalResult:
    """Result of a single retrieval operation."""
    test_id: str = ""
    query: str = ""
    retrieved_chunks: list[dict] = field(default_factory=list)     # {chunk_id, text, score, source}
    retrieval_latency_ms: float = 0.0
    embedding_latency_ms: float = 0.0
    search_latency_ms: float = 0.0
    rerank_latency_ms: float = 0.0
    total_latency_ms: float = 0.0


@dataclass
class RetrievalQualityMetrics:
    """Aggregated retrieval quality metrics."""
    precision_at_k: dict[int, float] = field(default_factory=dict)   # k -> precision
    recall_at_k: dict[int, float] = field(default_factory=dict)      # k -> recall
    mrr: float = 0.0                                                  # Mean Reciprocal Rank
    ndcg_at_k: dict[int, float] = field(default_factory=dict)        # k -> nDCG
    avg_relevance_score: float = 0.0
    fact_coverage_rate: float = 0.0                                   # % of expected facts found
    latency_p50_ms: float = 0.0
    latency_p95_ms: float = 0.0
    latency_p99_ms: float = 0.0
    total_queries: int = 0
    timestamp: float = field(default_factory=time.time)

import json
import logging
import math
import time
from statistics import median, quantiles
from typing import Optional

import boto3

logger = logging.getLogger(__name__)


class RetrievalQualityEvaluator:
    """Evaluates retrieval quality across multiple dimensions.

    Metrics computed:
    - Precision@K: fraction of retrieved chunks that are relevant
    - Recall@K: fraction of relevant chunks that are retrieved
    - MRR: average reciprocal rank of the first relevant chunk
    - nDCG@K: ranking quality considering position
    - Fact coverage: does the retrieved context contain the key facts?
    - Latency breakdown: per-stage timing
    """

    def __init__(self, opensearch_endpoint: str = "", embedding_model_id: str = "amazon.titan-embed-text-v2:0"):
        self.bedrock = boto3.client("bedrock-runtime", region_name="us-east-1")
        self.embedding_model_id = embedding_model_id
        self.opensearch_endpoint = opensearch_endpoint
        self.cloudwatch = boto3.client("cloudwatch", region_name="us-east-1")

    def evaluate_retrieval_batch(
        self,
        test_cases: list[RetrievalTestCase],
        retrieval_fn=None,  # Callable that takes query -> RetrievalResult
        k_values: list[int] = None,
    ) -> RetrievalQualityMetrics:
        """Evaluate retrieval quality on a batch of test cases."""
        if k_values is None:
            k_values = [1, 3, 5, 10]

        all_results: list[tuple[RetrievalTestCase, RetrievalResult]] = []

        for tc in test_cases:
            if retrieval_fn:
                result = retrieval_fn(tc.query)
            else:
                result = self._default_retrieve(tc.query)
            all_results.append((tc, result))

        # Compute metrics
        precision_at_k = {k: self._compute_precision_at_k(all_results, k) for k in k_values}
        recall_at_k = {k: self._compute_recall_at_k(all_results, k) for k in k_values}
        mrr = self._compute_mrr(all_results)
        ndcg_at_k = {k: self._compute_ndcg_at_k(all_results, k) for k in k_values}
        avg_relevance = self._compute_avg_relevance(all_results)
        fact_coverage = self._compute_fact_coverage(all_results)

        # Latency statistics
        latencies = [r.total_latency_ms for _, r in all_results if r.total_latency_ms > 0]
        sorted_latencies = sorted(latencies) if latencies else [0]

        return RetrievalQualityMetrics(
            precision_at_k=precision_at_k,
            recall_at_k=recall_at_k,
            mrr=mrr,
            ndcg_at_k=ndcg_at_k,
            avg_relevance_score=avg_relevance,
            fact_coverage_rate=fact_coverage,
            latency_p50_ms=self._percentile(sorted_latencies, 50),
            latency_p95_ms=self._percentile(sorted_latencies, 95),
            latency_p99_ms=self._percentile(sorted_latencies, 99),
            total_queries=len(all_results),
        )

    def _compute_precision_at_k(
        self, results: list[tuple[RetrievalTestCase, RetrievalResult]], k: int
    ) -> float:
        """Precision@K: of the top-K retrieved, how many are relevant?"""
        precisions = []
        for tc, result in results:
            top_k = result.retrieved_chunks[:k]
            if not top_k:
                precisions.append(0.0)
                continue
            relevant = sum(
                1 for chunk in top_k
                if chunk.get("chunk_id") in tc.expected_chunk_ids
            )
            precisions.append(relevant / len(top_k))
        return sum(precisions) / len(precisions) if precisions else 0.0

    def _compute_recall_at_k(
        self, results: list[tuple[RetrievalTestCase, RetrievalResult]], k: int
    ) -> float:
        """Recall@K: of all relevant chunks, how many are in top-K?"""
        recalls = []
        for tc, result in results:
            if not tc.expected_chunk_ids:
                continue
            top_k_ids = {
                chunk.get("chunk_id") for chunk in result.retrieved_chunks[:k]
            }
            relevant_retrieved = len(top_k_ids & set(tc.expected_chunk_ids))
            recalls.append(relevant_retrieved / len(tc.expected_chunk_ids))
        return sum(recalls) / len(recalls) if recalls else 0.0

    def _compute_mrr(
        self, results: list[tuple[RetrievalTestCase, RetrievalResult]]
    ) -> float:
        """Mean Reciprocal Rank: average of 1/rank of first relevant chunk."""
        reciprocal_ranks = []
        for tc, result in results:
            rr = 0.0
            for i, chunk in enumerate(result.retrieved_chunks):
                if chunk.get("chunk_id") in tc.expected_chunk_ids:
                    rr = 1.0 / (i + 1)
                    break
            reciprocal_ranks.append(rr)
        return sum(reciprocal_ranks) / len(reciprocal_ranks) if reciprocal_ranks else 0.0

    def _compute_ndcg_at_k(
        self, results: list[tuple[RetrievalTestCase, RetrievalResult]], k: int
    ) -> float:
        """Normalized Discounted Cumulative Gain at K."""
        ndcg_scores = []
        for tc, result in results:
            # Relevance: 1 if chunk is expected, 0 otherwise
            relevance = [
                1 if chunk.get("chunk_id") in tc.expected_chunk_ids else 0
                for chunk in result.retrieved_chunks[:k]
            ]

            # DCG
            dcg = sum(
                rel / math.log2(i + 2)  # i+2 because log2(1) = 0
                for i, rel in enumerate(relevance)
            )

            # Ideal DCG (all relevant chunks ranked first)
            ideal_relevance = sorted(relevance, reverse=True)
            idcg = sum(
                rel / math.log2(i + 2)
                for i, rel in enumerate(ideal_relevance)
            )

            ndcg = dcg / idcg if idcg > 0 else 0.0
            ndcg_scores.append(ndcg)

        return sum(ndcg_scores) / len(ndcg_scores) if ndcg_scores else 0.0

    def _compute_avg_relevance(
        self, results: list[tuple[RetrievalTestCase, RetrievalResult]]
    ) -> float:
        """Average relevance score of top-K retrieved chunks."""
        scores = []
        for _, result in results:
            for chunk in result.retrieved_chunks:
                score = chunk.get("score", chunk.get("relevance_score", 0))
                scores.append(float(score))
        return sum(scores) / len(scores) if scores else 0.0

    def _compute_fact_coverage(
        self, results: list[tuple[RetrievalTestCase, RetrievalResult]]
    ) -> float:
        """What percentage of expected key facts appear in retrieved chunks?"""
        total_facts = 0
        found_facts = 0
        for tc, result in results:
            if not tc.expected_key_facts:
                continue
            combined_text = " ".join(
                chunk.get("text", "") for chunk in result.retrieved_chunks
            ).lower()
            for fact in tc.expected_key_facts:
                total_facts += 1
                if fact.lower() in combined_text:
                    found_facts += 1
        return found_facts / total_facts if total_facts > 0 else 0.0

    def _percentile(self, sorted_values: list[float], pct: int) -> float:
        """Compute percentile from sorted values."""
        if not sorted_values:
            return 0.0
        idx = int(len(sorted_values) * pct / 100)
        return sorted_values[min(idx, len(sorted_values) - 1)]

    def _default_retrieve(self, query: str) -> RetrievalResult:
        """Default retrieval using MangaAssist's OpenSearch Serverless pipeline."""
        total_start = time.time()

        # Stage 1: Generate embedding
        embed_start = time.time()
        embedding = self._generate_embedding(query)
        embed_latency = (time.time() - embed_start) * 1000

        # Stage 2: Hybrid search (k-NN + BM25)
        search_start = time.time()
        raw_results = self._hybrid_search(query, embedding, k=20)
        search_latency = (time.time() - search_start) * 1000

        # Stage 3: Re-rank
        rerank_start = time.time()
        reranked = self._rerank(query, raw_results, top_k=5)
        rerank_latency = (time.time() - rerank_start) * 1000

        total_latency = (time.time() - total_start) * 1000

        return RetrievalResult(
            query=query,
            retrieved_chunks=reranked,
            embedding_latency_ms=embed_latency,
            search_latency_ms=search_latency,
            rerank_latency_ms=rerank_latency,
            total_latency_ms=total_latency,
        )

    def _generate_embedding(self, text: str) -> list[float]:
        """Generate embedding using Amazon Titan Embeddings v2."""
        body = json.dumps({"inputText": text})
        response = self.bedrock.invoke_model(modelId=self.embedding_model_id, body=body)
        return json.loads(response["body"].read())["embedding"]

    def _hybrid_search(self, query: str, embedding: list[float], k: int = 20) -> list[dict]:
        """Hybrid k-NN + BM25 search on OpenSearch Serverless."""
        # In production, this calls OpenSearch Serverless
        # Simplified for illustration
        return []

    def _rerank(self, query: str, candidates: list[dict], top_k: int = 5) -> list[dict]:
        """Re-rank candidates using cross-encoder or LLM-based re-ranker."""
        # In production, this uses a cross-encoder model
        return candidates[:top_k]

class RetrievalStageDiagnostics:
    """Diagnoses retrieval quality issues at each pipeline stage.

    When overall retrieval quality drops, this class helps isolate
    which stage is the bottleneck:
    - Query rewriting losing intent?
    - Embeddings not capturing semantics?
    - Search ranking poorly?
    - Re-ranker demoting correct results?
    """

    def diagnose_query_rewriter(
        self, original_query: str, rewritten_query: str, intent: str
    ) -> dict:
        """Evaluate whether the query rewriter preserved intent and improved searchability."""
        # Check intent preservation: key terms from original should appear in rewritten
        original_terms = set(original_query.lower().split())
        rewritten_terms = set(rewritten_query.lower().split())
        overlap = original_terms & rewritten_terms
        preservation_rate = len(overlap) / len(original_terms) if original_terms else 0

        # Check expansion: did the rewriter add useful terms?
        added_terms = rewritten_terms - original_terms
        expansion_rate = len(added_terms) / max(len(rewritten_terms), 1)

        return {
            "stage": "query_rewriter",
            "intent_preservation": preservation_rate,
            "query_expansion_rate": expansion_rate,
            "original_length": len(original_query),
            "rewritten_length": len(rewritten_query),
            "added_terms": list(added_terms)[:10],
            "healthy": preservation_rate >= 0.5,
        }

    def diagnose_embedding_quality(
        self, query_embedding: list[float], expected_chunk_embeddings: list[list[float]]
    ) -> dict:
        """Evaluate embedding quality by checking cosine similarity distribution."""
        similarities = []
        for chunk_emb in expected_chunk_embeddings:
            sim = self._cosine_similarity(query_embedding, chunk_emb)
            similarities.append(sim)

        if not similarities:
            return {"stage": "embedding", "healthy": False, "reason": "No embeddings to compare"}

        avg_sim = sum(similarities) / len(similarities)
        max_sim = max(similarities)
        min_sim = min(similarities)

        return {
            "stage": "embedding",
            "avg_similarity": round(avg_sim, 4),
            "max_similarity": round(max_sim, 4),
            "min_similarity": round(min_sim, 4),
            "similarity_spread": round(max_sim - min_sim, 4),
            "healthy": avg_sim >= 0.65,
            "recommendation": (
                "Embedding model captures query-chunk similarity well"
                if avg_sim >= 0.65
                else "Consider fine-tuning embeddings or adjusting chunk size"
            ),
        }

    def diagnose_search_components(
        self,
        knn_results: list[dict],
        bm25_results: list[dict],
        hybrid_results: list[dict],
        expected_chunk_ids: list[str],
    ) -> dict:
        """Compare k-NN and BM25 contributions to identify search issues."""
        knn_set = {r.get("chunk_id") for r in knn_results[:10]}
        bm25_set = {r.get("chunk_id") for r in bm25_results[:10]}
        hybrid_set = {r.get("chunk_id") for r in hybrid_results[:10]}
        expected_set = set(expected_chunk_ids)

        return {
            "stage": "search",
            "knn_recall@10": len(knn_set & expected_set) / max(len(expected_set), 1),
            "bm25_recall@10": len(bm25_set & expected_set) / max(len(expected_set), 1),
            "hybrid_recall@10": len(hybrid_set & expected_set) / max(len(expected_set), 1),
            "knn_unique": len(knn_set - bm25_set),
            "bm25_unique": len(bm25_set - knn_set),
            "overlap": len(knn_set & bm25_set),
            "recommendation": self._search_recommendation(
                knn_set, bm25_set, hybrid_set, expected_set
            ),
        }

    def diagnose_reranker(
        self,
        pre_rerank: list[dict],
        post_rerank: list[dict],
        expected_chunk_ids: list[str],
    ) -> dict:
        """Evaluate whether the re-ranker improved or degraded ranking quality."""
        expected_set = set(expected_chunk_ids)

        # Compare positions of relevant chunks pre and post rerank
        pre_positions = {}
        for i, chunk in enumerate(pre_rerank):
            if chunk.get("chunk_id") in expected_set:
                pre_positions[chunk["chunk_id"]] = i

        post_positions = {}
        for i, chunk in enumerate(post_rerank):
            if chunk.get("chunk_id") in expected_set:
                post_positions[chunk["chunk_id"]] = i

        improvements = 0
        degradations = 0
        for cid in expected_set:
            pre = pre_positions.get(cid, len(pre_rerank))
            post = post_positions.get(cid, len(post_rerank))
            if post < pre:
                improvements += 1
            elif post > pre:
                degradations += 1

        return {
            "stage": "reranker",
            "improvements": improvements,
            "degradations": degradations,
            "no_change": len(expected_set) - improvements - degradations,
            "healthy": improvements >= degradations,
            "recommendation": (
                "Re-ranker is improving ranking"
                if improvements > degradations
                else "Re-ranker may be degrading results — review cross-encoder model"
            ),
        }

    def _cosine_similarity(self, a: list[float], b: list[float]) -> float:
        """Compute cosine similarity between two vectors."""
        dot_product = sum(x * y for x, y in zip(a, b))
        norm_a = math.sqrt(sum(x * x for x in a))
        norm_b = math.sqrt(sum(x * x for x in b))
        return dot_product / (norm_a * norm_b) if norm_a > 0 and norm_b > 0 else 0.0

    def _search_recommendation(self, knn, bm25, hybrid, expected) -> str:
        knn_recall = len(knn & expected) / max(len(expected), 1)
        bm25_recall = len(bm25 & expected) / max(len(expected), 1)
        if knn_recall > bm25_recall + 0.2:
            return "k-NN dominates — consider increasing k-NN weight in hybrid"
        elif bm25_recall > knn_recall + 0.2:
            return "BM25 dominates — queries may be keyword-heavy, check embeddings"
        return "Balanced contribution — hybrid search is working well"