RAG and Context Engineering

A retrieval-augmented generation (RAG) pipeline for question answering over the RapidFire AI documentation, built as a thin orchestration layer over RapidFire AI's RFLangChainRagSpec and RFOpenAIAPIModelConfig. The system retrieves, reranks, and packs evidence from a .rst documentation corpus under a strict 2,000-token per-query context budget, then generates grounded, citation-bearing answers.

The repository is designed around a single question: given a fixed context budget, which knobs actually move answer quality? It ships the final pipeline, a pre-registered ablation sweep that answers that question empirically, a grounded golden-Q&A dataset generator, and an LLM-as-a-judge evaluation harness.

Results at a glance

The final configuration (C6) reaches a 0.855 leaderboard proxy, up from a 0.735 off-the-shelf baseline — a gain driven almost entirely by retrieval quality, with Precision@5 rising from 0.467 to 0.852 (an 82% relative improvement) while faithfulness stays saturated.

Config	Change (single knob)	P@5	R@5	F1@5	Retrieval	Correct.	Faith.	Comp./5	LB proxy
C0	Off-the-shelf baseline	0.467	0.861	0.584	0.637	0.844	1.000	0.653	0.735
C1	Smaller chunks (256/64)	0.474	0.883	0.593	0.650	0.756	1.000	0.582	0.715
C2	Smallest chunks (128/32) — predicted loser	0.482	0.906	0.597	0.661	0.622	1.000	0.427	0.672
C3	Multi-resolution chunking	0.800	0.894	0.815	0.836	0.844	1.000	0.600	0.826
C4	Hybrid BM25 + FAISS retrieval	0.807	0.917	0.831	0.852	0.844	1.000	0.613	0.836
C5	bge-reranker-v2-m3 cross-encoder	0.852	0.894	0.848	0.865	0.844	1.000	0.644	0.847
C6	gpt-oss-120b generator — final	0.852	0.894	0.848	0.865	0.889	0.978	0.671	0.855

Bold marks the best value in each column. LB proxy = 0.5·Retrieval + 0.5·Gen, where Gen averages the three released judges. Full narrative and per-config logs are in report.pdf and logs/.

Architecture

The pipeline is deliberately a thin layer: every heavy component (embeddings, reranking, generation, judging) is a hosted model behind the UCSD Triton API gateway, so the code owns only the context engineering — chunking, fusion, budgeting, and prompt assembly.

%%{init: {'flowchart': {'nodeSpacing': 22, 'rankSpacing': 24}, 'themeVariables': {'fontSize': '11px'}}}%%
flowchart TD
    A[".rst corpus"] --> B["Load + multi-resolution split<br/>128 / 256 / 512 tok, 25% overlap"]
    B --> C["FAISS · dense"]
    B --> D["BM25 · lexical"]
    C --> E["Ensemble 50/50 → top 10"]
    D --> E
    E --> F["Rerank · bge-reranker-v2-m3 → top 3"]
    F --> G["Budget pack · 2,000-tok GPT-2"]
    G --> H["Generate · api-gpt-oss-120b"]
    H --> I["Answer + sources + context"]

    classDef hosted fill:#e8eefc,stroke:#5b7fd4,color:#111;
    class C,F,H hosted;

Loading

Blue = models hosted on the UCSD Triton gateway; the code owns the context engineering between them.

Design decisions worth calling out:

• Multi-resolution indexing over single-resolution tuning. Shrinking chunks at a single resolution traded recall for completeness and never helped (C1/C2). Indexing three resolutions (128/256/512) together nearly doubled Precision@5 — the cross-encoder consistently prefers the most focused of the three overlapping views of each gold neighborhood.
• Hybrid retrieval. API-name queries (RFGridSearch, RFList) are exact-keyword targets the embedder tends to paraphrase; a 50/50 BM25 + FAISS ensemble recovers them.
• Skip-don't-stop budget packing. The 2,000-token budget is allocated dynamically (prompt + scaffold + question + 40-token safety margin subtracted first); chunks are admitted in rerank-rank order and any that would overflow are skipped rather than terminating the pack, so the budget is filled more completely. All token counts use the GPT-2 encoding to keep chunk sizing and budget arithmetic on the same ruler.
• Grounded generation with explicit abstention. The system prompt requires every claim to be grounded in retrieved context and emits <answer>I don't know</answer> when the context is insufficient; a DOTALL post-processor extracts the last <answer>…</answer> span, which discards any in-context example the model echoes.

Repository layout

Path	Purpose
main.py	Single-config runner that reproduces the final (C6) pipeline end to end.
configs.py	The C0–C7 config matrix: chunking, retrieval, reranker, and generator variants.
prompts.py	System / instruction prompt templates.
utils.py	Multi-resolution splitter, GPT-2 token counting, budget packing, span dedup, pre/post-processing.
rapidfireai_datahub_compat.py	RapidFire AI datahub compatibility shim.
experiment.ipynb	Drives the C0–C7 sweep via Experiment.run_evals + RFGridSearch across Ray actors.
metrics/	Retrieval metrics (evaluate_retrieval.py), the LLM judge (run_judge.py, judge_prompt.txt), and RapidFire integration (project1_eval.py, rapidfire_integration_example.py).
qa_gold_generator/	Standalone documentation-grounded golden-Q&A generator (see below).
Project1/sourcedocs/	The RapidFire AI .rst documentation corpus the pipeline runs against.
Project1/validation-set-golden-qa-pairs.json	Evaluation set consumed by the experiment sweep.
golden_qa_curated.json	Hand-curated golden Q&A set consumed by main.py.
output.json	Sample pipeline output in the submission schema.
report.pdf	Full methodology, results, and analysis.

Prerequisites

• Python 3.12+
• Access to the UCSD Triton API gateway (https://tritonai-api.ucsd.edu) and an API key. The gateway hosts every model the pipeline depends on:

Role	Model
Embeddings	api-tgpt-embeddings
Reranker	BAAI/bge-reranker-v2-m3
Generator	api-gpt-oss-120b
LLM judge	claude-sonnet-4-6

Installation

pip install -e .                 # full pipeline + generator (from pyproject.toml)
# or, for the standalone generator only:
pip install -r requirements.txt

Configuration

The pipeline reads its API key from a file; the generator reads OpenAI-compatible environment variables (falling back to the Triton variables). Never commit real keys — api-key.txt and .env are gitignored.

# For the pipeline / experiment sweep:
echo "<your-triton-api-key>" > ~/api-key.txt

# For the generator:
cp .env.example .env             # then set TRITON_API_KEY, GENERATOR_MODEL, VALIDATOR_MODEL, EMBEDDING_MODEL

Usage

Run the final pipeline

python main.py \
  --input golden_qa_curated.json \
  --output output.json \
  --corpus-dir Project1/sourcedocs \
  --apikey-txt ~/api-key.txt \
  --generation-model api-gpt-oss-120b

Output is a list of {question_id, answer, sources, retrieved_context} records, with sources reported as {file, lines} using file basenames.

Run the ablation sweep (C0–C7)

Open experiment.ipynb. It builds the config grid from configs.py, dispatches evals with RFGridSearch + Experiment.run_evals (num_shards=4, seed=42) across Ray actors, and scores each config with the retrieval and judge metrics in metrics/. It reads the corpus from Project1/sourcedocs and the evaluation set from Project1/validation-set-golden-qa-pairs.json.

Generate a golden Q&A set

python -m qa_gold_generator.cli \
  --docs_dir Project1/sourcedocs \
  --out golden_qa.json \
  --num_questions 100

Experiment design

The config matrix is organized as an arc: each row changes exactly one knob relative to a named predecessor and tests a hypothesis committed to before the run. C2 is pre-registered as an expected loser (smallest chunks under-fill the budget) and C6/C7 as uncertain — including configs whose predicted outcome is negative or uncertain is deliberate, so the sweep tests hypotheses rather than fishing for wins.

• C0 → C2 (chunk size): single-resolution shrinking is the wrong lever; smaller chunks buy a little recall but starve the generator of context, dropping completeness. Faithfulness holds at 1.000 throughout — the model never fabricates to compensate for thin context.
• C3 (multi-resolution): the structural retrieval win — Precision@5 nearly doubles.
• C4 → C5 (hybrid retrieval, stronger reranker): lexical fusion and the bge-reranker-v2-m3 cross-encoder push Precision@5 to 0.852; Correctness stays pinned at 0.844 across four retrieval configs, isolating the residual failures as not retrieval-shaped.
• C6 (generator swap): replacing mistral-small with gpt-oss-120b is the only knob that moves Correctness (0.844 → 0.889), at a small faithfulness cost — the final pipeline.

Evaluation methodology

• Retrieval: Precision@5, Recall@5, and F1@5, computed by inclusive line-overlap on the file basename against the source-evidence spans. This is why the splitter stamps start_line/ end_line on every chunk.
• Generation: Correctness, Faithfulness, and Completeness scored by a claude-sonnet-4-6 LLM judge over the validation set.
• Leaderboard proxy: 0.5·Retrieval + 0.5·Gen.

The golden Q&A generator

qa_gold_generator/ builds documentation-grounded QA pairs for RAG evaluation. It loads local .rst files (preserving relative filenames and line numbers), chunks by reStructuredText sections, generates candidates with an OpenAI-compatible chat model, validates every candidate against its exact cited evidence with a separate LLM pass, deduplicates semantically (embeddings → sentence-transformers → TF-IDF fallback), and selects a type-balanced final set across factual_lookup, conceptual_explanation, procedural, comparative, feature_enumeration, and multi_file_synthesis. It caches LLM and embedding calls in .qa_gold_cache/ and writes golden_qa.json plus a generation_report.json. Deliberate overgeneration (hundreds of candidates) leaves room to filter aggressively while preserving diversity across question types, source files, and difficulty.

Reproducibility

The final pipeline is reproduced deterministically by main.py; the full sweep is reproduced by experiment.ipynb. The only external requirements are a Triton API key at ~/api-key.txt and gateway access to the models listed above — the corpus, the evaluation sets, the configs, the prompts, and the dependency manifest (pyproject.toml, Python 3.12) are all checked in.