Add optional RAG cross-encoder reranking

ai_engineering/rag_assistant/.env.example

@@ -3,3 +3,8 @@
 OPENAI_API_KEY=sk-your-key-here
 OPENAI_MODEL=gpt-4o-mini
 EMBEDDING_MODEL=sentence-transformers/all-MiniLM-L6-v2
+
+# Optional serving configuration. Leave RERANKER_MODEL unset to preserve
+# embedding-only retrieval.
+# RERANKER_MODEL=cross-encoder/ms-marco-MiniLM-L-6-v2
+RERANKER_CANDIDATES=20

ai_engineering/rag_assistant/README.md

@@ -145,6 +145,59 @@ This is a 5-case smoke eval (`sample_docs/eval_cases.json`), not a large benchma
 - **Honest evaluation**: ships an actual `recall@k` + MRR eval harness with checked-in cases, instead of just claiming "high quality".
 - **Reproducible smoke pipeline**: `python -m smoke.run_smoke` exercises ingest + retrieve + eval end-to-end, no key required, and writes a markdown report a reviewer can read in 30 seconds.
 
+## Optional Cross-Encoder Re-ranking
+
+Embedding retrieval remains the default. Add `--rerank` to `ask` or `eval` to
+retrieve a larger candidate pool and re-score it with
+`cross-encoder/ms-marco-MiniLM-L-6-v2`:
+
+```bash
+python cli.py ask \
+  --store ./index \
+  --question "What is RAG?" \
+  --rerank \
+  --candidate-k 20 \
+  -k 5
+
+python cli.py eval \
+  --store ./index \
+  --cases ./sample_docs/eval_cases.json \
+  --rerank \
+  --candidate-k 20 \
+  -k 3
+```
+
+For the Flask service, set `RERANKER_MODEL` and optionally
+`RERANKER_CANDIDATES`. The JSON response exposes both the final cross-encoder
+`score` and the original `retrieval_score`.
+
+Run the baseline comparison with:
+
+```bash
+python -m smoke.run_reranker_eval
+```
+
+The generated `reports/reranker_eval.md` records baseline and re-ranked
+recall@k and MRR, their deltas, and per-case reciprocal-rank changes. The
+current five-case suite is saturated, so the observed delta is provisional
+until the expanded benchmark in issue #18 is available.
+
+### Measured Smoke Result
+
+A CPU run using `sentence-transformers/all-MiniLM-L6-v2` and
+`cross-encoder/ms-marco-MiniLM-L-6-v2` produced:
+
+| Configuration | recall@3 | MRR |
+|---|---:|---:|
+| Embedding only | 1.000 | 0.900 |
+| Cross-encoder re-ranked | 1.000 | 1.000 |
+| Observed delta | +0.000 | +0.100 |
+
+The re-ranker moved the approximate-nearest-neighbor question from rank 2 to
+rank 1. This records the observed result on the checked-in five-case smoke
+suite; it does not establish a general quality lift. Issue #18 remains the
+required follow-up for a discriminative benchmark.
+
 ## Scope
 
 - The chunker is character-based and language-agnostic, which is good portability but slightly worse than tokenization-aware chunking for very long contexts.

ai_engineering/rag_assistant/app.py

@@ -5,6 +5,7 @@
 
 Build with `build_app("./index")` so a saved store is loaded once at
 startup and reused across requests (no per-request rebuild).
+Set RERANKER_MODEL to opt into cross-encoder re-ranking.
 """
 
 from __future__ import annotations
@@ -13,12 +14,23 @@
 
 from flask import Flask, jsonify, request
 
-from rag import RAGPipeline
+from rag import CrossEncoderReranker, RAGPipeline
 
 
 def build_app(store_path: str) -> Flask:
     app = Flask(__name__)
-    pipeline = RAGPipeline.load(store_path)
+    reranker_model = os.environ.get("RERANKER_MODEL", "").strip()
+    reranker = (
+        CrossEncoderReranker(model_name=reranker_model)
+        if reranker_model
+        else None
+    )
+    candidate_k = int(os.environ.get("RERANKER_CANDIDATES", "20"))
+    pipeline = RAGPipeline.load(
+        store_path,
+        reranker=reranker,
+        candidate_k=candidate_k,
+    )
 
     @app.route("/health", methods=["GET"])
     def health():
@@ -39,6 +51,7 @@ def ask():
             "retrieved": [
                 {
                     "score": r.score,
+                    "retrieval_score": r.retrieval_score,
                     "source": r.chunk.source,
                     "chunk_index": r.chunk.chunk_index,
                 }

ai_engineering/rag_assistant/cli.py

@@ -9,8 +9,10 @@
 Usage:
     python cli.py ingest --docs-dir ./sample_docs --store ./index
     python cli.py ask --store ./index --question "What is RAG?"
+    python cli.py ask --store ./index --question "What is RAG?" --rerank
     python cli.py serve --store ./index --port 8080
-    python cli.py eval --store ./index --cases ./sample_docs/eval_cases.json --k 3
+    python cli.py eval --store ./index --cases ./sample_docs/eval_cases.json -k 3
+    python cli.py eval --store ./index --cases ./sample_docs/eval_cases.json -k 3 --rerank
 """
 
 from __future__ import annotations
@@ -20,9 +22,18 @@
 import os
 import sys
 from pathlib import Path
-from typing import List
-
-from rag import Chunker, Document, RAGPipeline, eval_retrieval, EvalCase
+from typing import List, Optional
+
+from rag import (
+    Chunker,
+    CrossEncoderReranker,
+    DEFAULT_RERANKER_MODEL,
+    Document,
+    EvalCase,
+    RAGPipeline,
+    Reranker,
+    eval_retrieval,
+)
 from rag.embedder import make_default_embedder
 from rag.vector_store import VectorStore
 
@@ -48,6 +59,31 @@ def load_docs_from_dir(dir_path: str) -> List[Document]:
     return docs
 
 
+def _build_reranker(args: argparse.Namespace) -> Optional[Reranker]:
+    if not getattr(args, "rerank", False):
+        return None
+    return CrossEncoderReranker(model_name=args.reranker_model)
+
+
+def _add_reranker_args(parser: argparse.ArgumentParser) -> None:
+    parser.add_argument(
+        "--rerank",
+        action="store_true",
+        help="re-rank a larger candidate pool with a cross-encoder",
+    )
+    parser.add_argument(
+        "--reranker-model",
+        default=DEFAULT_RERANKER_MODEL,
+        help="sentence-transformers CrossEncoder model name",
+    )
+    parser.add_argument(
+        "--candidate-k",
+        type=int,
+        default=20,
+        help="embedding candidates passed to the re-ranker",
+    )
+
+
 def cmd_ingest(args: argparse.Namespace) -> None:
     docs = load_docs_from_dir(args.docs_dir)
     if not docs:
@@ -60,7 +96,11 @@ def cmd_ingest(args: argparse.Namespace) -> None:
 
 
 def cmd_ask(args: argparse.Namespace) -> None:
-    pipeline = RAGPipeline.load(args.store)
+    pipeline = RAGPipeline.load(
+        args.store,
+        reranker=_build_reranker(args),
+        candidate_k=args.candidate_k,
+    )
     result = pipeline.ask(args.question, k=args.k)
     print(f"\nAnswer:\n{result.answer}\n")
     if result.sources:
@@ -82,7 +122,16 @@ def cmd_eval(args: argparse.Namespace) -> None:
     if embedder.dim != store.dim:
         sys.exit(f"embedder dim {embedder.dim} != store dim {store.dim}; reindex.")
     from rag.retriever import Retriever
-    result = eval_retrieval(Retriever(embedder, store), cases, k=args.k)
+    result = eval_retrieval(
+        Retriever(
+            embedder,
+            store,
+            reranker=_build_reranker(args),
+            candidate_k=args.candidate_k,
+        ),
+        cases,
+        k=args.k,
+    )
     print(result)
     for row in result.per_case:
         print(f"  - {row['question'][:60]:60s}  recall={row['recall']:.0f}  rr={row['reciprocal_rank']:.3f}")
@@ -103,6 +152,7 @@ def main() -> None:
     p_ask.add_argument("--store", required=True)
     p_ask.add_argument("--question", required=True)
     p_ask.add_argument("-k", type=int, default=5)
+    _add_reranker_args(p_ask)
     p_ask.set_defaults(func=cmd_ask)
 
     p_serve = sub.add_parser("serve")
@@ -114,6 +164,7 @@ def main() -> None:
     p_eval.add_argument("--store", required=True)
     p_eval.add_argument("--cases", required=True)
     p_eval.add_argument("-k", type=int, default=5)
+    _add_reranker_args(p_eval)
     p_eval.set_defaults(func=cmd_eval)
 
     args = parser.parse_args()

ai_engineering/rag_assistant/rag/__init__.py

@@ -4,7 +4,8 @@
     Chunker        - split documents into overlapping chunks
     Embedder       - convert chunks to dense vectors
     VectorStore    - FAISS-backed nearest-neighbor index with metadata
-    Retriever      - thin wrapper that wires Embedder + VectorStore together
+    Retriever      - embedding retrieval with optional cross-encoder re-ranking
+    Reranker       - protocol for re-ordering an initial candidate pool
     Generator      - calls the OpenAI Chat Completions API with retrieved context
     RAGPipeline    - end-to-end: ingest -> retrieve -> generate
     eval_retrieval - retrieval-quality metrics (recall@k, MRR)
@@ -16,6 +17,11 @@
 from .chunker import Chunker, Document, Chunk
 from .embedder import Embedder, HashEmbedder
 from .vector_store import VectorStore
+from .reranker import (
+    CrossEncoderReranker,
+    DEFAULT_RERANKER_MODEL,
+    Reranker,
+)
 from .retriever import Retriever, RetrievedChunk
 from .generator import Generator
 from .pipeline import RAGPipeline
@@ -25,6 +31,7 @@
     "Chunker", "Document", "Chunk",
     "Embedder", "HashEmbedder",
     "VectorStore",
+    "Reranker", "CrossEncoderReranker", "DEFAULT_RERANKER_MODEL",
     "Retriever", "RetrievedChunk",
     "Generator",
     "RAGPipeline",

ai_engineering/rag_assistant/rag/embedder.py

@@ -66,7 +66,14 @@ def __init__(self, model_name: str = "sentence-transformers/all-MiniLM-L6-v2"):
         from sentence_transformers import SentenceTransformer  # type: ignore
 
         self._model = SentenceTransformer(model_name)
-        self.dim = self._model.get_sentence_embedding_dimension()
+        dimension_getter = getattr(
+            self._model,
+            "get_embedding_dimension",
+            None,
+        )
+        if dimension_getter is None:
+            dimension_getter = self._model.get_sentence_embedding_dimension
+        self.dim = dimension_getter()
 
     def embed(self, texts: Sequence[str]) -> np.ndarray:
         if not texts:

ai_engineering/rag_assistant/rag/pipeline.py

@@ -1,7 +1,8 @@
 """End-to-end ingest + retrieve + generate pipeline.
 
 `RAGPipeline.ingest(docs)` runs the chunker + embedder + store.
-`RAGPipeline.ask(question, k=5)` runs the retriever + generator.
+`RAGPipeline.ask(question, k=5)` runs the retriever, optional re-ranker,
+and generator.
 
 `from_env()` reads `OPENAI_API_KEY`, `OPENAI_MODEL`, `EMBEDDING_MODEL`,
 so the CLI and serving layers share configuration.
@@ -18,6 +19,7 @@
 from .chunker import Chunk, Chunker, Document
 from .embedder import Embedder, make_default_embedder
 from .generator import GenerationResult, Generator
+from .reranker import Reranker
 from .retriever import RetrievedChunk, Retriever
 from .vector_store import VectorStore
 
@@ -37,10 +39,17 @@ def __init__(
         store: VectorStore,
         generator: Optional[Generator] = None,
         chunker: Optional[Chunker] = None,
+        reranker: Optional[Reranker] = None,
+        candidate_k: int = 20,
     ):
         self.embedder = embedder
         self.store = store
-        self.retriever = Retriever(embedder, store)
+        self.retriever = Retriever(
+            embedder,
+            store,
+            reranker=reranker,
+            candidate_k=candidate_k,
+        )
         self.generator = generator
         self.chunker = chunker or Chunker()
 
@@ -95,6 +104,8 @@ def load(
         dir_path: str,
         embedder: Optional[Embedder] = None,
         generator: Optional[Generator] = None,
+        reranker: Optional[Reranker] = None,
+        candidate_k: int = 20,
     ) -> "RAGPipeline":
         store = VectorStore.load(dir_path)
         emb = embedder or make_default_embedder()
@@ -103,4 +114,10 @@ def load(
                 f"embedder dim {emb.dim} does not match stored index dim {store.dim}; "
                 "did you change EMBEDDING_MODEL since indexing?"
             )
-        return cls(embedder=emb, store=store, generator=generator or Generator())
+        return cls(
+            embedder=emb,
+            store=store,
+            generator=generator or Generator(),
+            reranker=reranker,
+            candidate_k=candidate_k,
+        )