RAG Evaluation with Contextual Retrieval

Evaluate and compare retrieval strategies for RAG systems, with built-in support for Contextual Retrieval and reranking.

Overview

This repository provides a complete pipeline for:

Ingesting documents into a vector database (Qdrant), with optional Contextual Retrieval enrichment
Retrieving relevant chunks for queries, with optional reranking (Cohere or vLLM-hosted models)
Evaluating retrieval quality using character-span coverage metrics

The core idea behind Contextual Retrieval is simple: chunks extracted from documents often lose important context when isolated. A table row showing weapon stats might not mention it's from the "Martial Weapons" section. By prepending a short, LLM-generated context to each chunk before embedding, retrieval accuracy improves significantly.

We created a complete D&D 5e SRD evaluation dataset with 56 high-quality questions (easy + medium difficulty) and domain-specific prompts. Learn about our methodology in our blog post. The dataset and generation code are available on GitHub and HuggingFace. Easily adaptable to your own datasets.

📖 For a detailed explanation of why Contextual Retrieval works and our experimental results, see our blog post.

Features

Contextual Retrieval: Enrich chunks with LLM-generated context before embedding
Flexible Ingestion: From markdown files or pre-chunked JSON
Reranking: Cohere Rerank v4 Pro or self-hosted Qwen3-Reranker via vLLM
Character-span Evaluation: Precise coverage metrics (fuzzy and rigid recall)
Config-driven Workflow: Run full pipelines via YAML configuration
Visualization: Publication-quality comparison plots
Built on datapizza-ai: Modular, extensible pipeline architecture

Requirements

Python 3.13
Qdrant instance (cloud or self-hosted)
API keys for:
- Cohere (embeddings, optional reranking)
- Google (Gemini for contextual enrichment, if enabled)
- vLLM endpoint (optional, for self-hosted reranking)

Installation

# Recommended (uses uv for fast dependency resolution)
uv sync

# Or with pip
pip install -e .

Environment Variables

Set the following environment variables based on which features you use:

# Required for all operations
export COHERE_API_KEY=...
export QDRANT_HOST=...          # e.g., https://your-cluster.qdrant.io
export QDRANT_API_KEY=...

# Required if using Contextual Retrieval (default: enabled)
export GOOGLE_API_KEY=...

# Optional: Cohere reranker (uses same key as embedder by default)
export COHERE_RERANKER_API_KEY=...   # Falls back to COHERE_API_KEY

# Optional: vLLM reranker (e.g., Qwen3-Reranker-8B)
# See run_self_hosted_reranker.md for setup instructions
export RERANKER_ENDPOINT=...    # e.g., http://your-gpu-server:8000/score
export RERANKER_API_KEY=...     # If your vLLM endpoint requires auth

# Optional: Custom Cohere endpoint (e.g., Azure Foundry)
export COHERE_ENDPOINT=...

Quick Start

The fastest way to run experiments is via YAML configuration:

# Run ingestion + retrieval + evaluation
uv run src/contextual_retrieval_experiment/main.py --config config/contextual/ingestion_contextual.yaml

# Run retrieval + evaluation on existing collection
uv run src/contextual_retrieval_experiment/main.py --config config/contextual/retrieval_eval_dnd_medium_cohere.yaml

Or run individual stages via CLI:

# 1. Ingest documents with contextual enrichment
uv run src/contextual_retrieval_experiment/ingestion/ingestion.py \
  --file-path dataset/raw_data/DnD5eSRD_md/DND5eSRD_001-018.md \
  --file-path dataset/raw_data/DnD5eSRD_md/DND5eSRD_019-035.md \
  --source-dir dataset/raw_data/DnD5eSRD_md \
  --collection-name my_contextual_collection \
  --prompt-template-path src/contextual_retrieval_experiment/ingestion/constants/templates/contextual_retrieval_prompt_dnd.jinja2 \
  --output-chunks-json data/chunks_with_positions.json

# 2. Run retrieval with reranking
uv run src/contextual_retrieval_experiment/retrieval/retrieval.py \
  --collection-name my_contextual_collection \
  --dataset-json dataset/qa_sets/medium.json \
  --use-reranker \
  --reranker-type cohere \
  --retrieve-k 128 \
  --reranker-top-n 20

# 3. Evaluate results
uv run src/contextual_retrieval_experiment/evaluation/evaluation.py \
  --retrieval-results-path results/retrieval_my_contextual_collection_cohere_p128_q20.json \
  --chunks-path data/chunks_with_positions.json \
  --dataset-json dataset/qa_sets/medium.json \
  --k "5,10,20"

Data Preparation

Using the Included D&D 5e SRD Dataset

This repository includes a ready-to-use dataset:

Source documents: dataset/raw_data/DnD5eSRD_md/ (20 markdown files)
QA sets: dataset/qa_sets/easy.json (25 questions), dataset/qa_sets/medium.json (31 questions)
Pre-built chunks: dataset/chunks/chunks_with_positions.json

Building Your Own Dataset (Optional)

To create your own evaluation dataset from PDFs:

Use rag-dataset-builder to:
- Convert PDFs to Markdown
- Generate questions and answers
- Extract ground-truth passages with character positions
Place your markdown files in a directory and create a QA JSON following the format below.

Pipeline Stages

1. Ingestion

Ingestion loads documents, splits them into chunks, optionally enriches them with context, embeds them, and stores them in Qdrant.

With Contextual Retrieval (default):

For each chunk, the LLM receives the full source document plus the chunk
It generates a short context paragraph explaining where the chunk fits
The context is prepended as CONTEXT: ... CONTENT: ... before embedding

Baseline (no enrichment):

Chunks are embedded directly without context
Use --skip-contextual-retrieval flag

Chunking parameters:

Default chunk size: 4000 characters
Default overlap: 100 characters

2. Retrieval

Retrieval queries the vector database and optionally reranks results.

Two-stage retrieval:

Vector search: Retrieve retrieve_k candidates (default: 128)
Reranking: Score and re-order using a cross-encoder, keep top reranker_top_n (default: 20)

Supported rerankers:

Cohere Rerank v4 Pro: Hosted API, state-of-the-art quality
Qwen3-Reranker-8B via vLLM: Self-hosted, cost-effective for high volume (see run_self_hosted_reranker.md for setup instructions)

3. Evaluation

Evaluation computes how well retrieved chunks cover the ground-truth passages.

Metrics:

Fuzzy coverage (recall): Fraction of required chunks retrieved
Rigid coverage: 1 if all required chunks retrieved, else 0

Modes:

Offline: Evaluate from a saved retrieval JSON file
Live: Query the vectorstore directly during evaluation

CLI Reference

Ingestion

uv run src/contextual_retrieval_experiment/ingestion/ingestion.py [OPTIONS]

Option	Description
`--file-path PATH`	File(s) to ingest (can specify multiple times)
`--json-chunks-path PATH`	Load from pre-chunked JSON instead of files
`--source-dir PATH`	Directory containing source files (required for contextual retrieval)
`--collection-name NAME`	Qdrant collection name (auto-generated if omitted)
`--prompt-template-path PATH`	Jinja2 template for context generation
`--output-chunks-json PATH`	Save chunks lookup JSON for evaluation
`--skip-contextual-retrieval`	Disable context enrichment (baseline mode)
`--chunk-size INT`	Characters per chunk (default: 4000)
`--chunk-overlap INT`	Overlap between chunks (default: 100)
`--model NAME`	LLM for context generation (default: gemini-2.5-pro)
`--max-retries INT`	Max retries for LLM calls (default: 3)

Retrieval

uv run src/contextual_retrieval_experiment/retrieval/retrieval.py [OPTIONS]

Option	Description
`--collection-name NAME`	Qdrant collection name (required)
`--query TEXT`	Single query (mutually exclusive with --dataset-json)
`--dataset-json PATH`	QA dataset for batch processing
`--use-reranker / --no-reranker`	Enable/disable reranking (default: enabled)
`--reranker-type {vllm,cohere}`	Reranker backend (default: vllm)
`--retrieve-k INT`	Vector search candidates (default: 20)
`--reranker-top-n INT`	Results after reranking (default: 5)
`--score-threshold FLOAT`	Minimum retrieval score
`--reranker-threshold FLOAT`	Minimum reranker score
`--embedding-model NAME`	Embedding model (default: embed-v4.0)
`--vllm-reranker-model NAME`	vLLM model (default: Qwen/Qwen3-Reranker-8B)
`--cohere-reranker-model NAME`	Cohere model (default: rerank-v4.0-pro)
`--output PATH`	Output JSON file

Evaluation

uv run src/contextual_retrieval_experiment/evaluation/evaluation.py [OPTIONS]

Option	Description
`--retrieval-results-path PATH`	Saved retrieval JSON (offline mode)
`--chunks-path PATH`	Chunks lookup JSON (required)
`--dataset-json PATH`	QA dataset (required for live mode or passage joining)
`--collection-name NAME`	Qdrant collection (required for live mode)
`--k VALUE`	k value(s): `5`, `[5,10,20]`, or `5,10,20` (default: 5)
`--score-threshold FLOAT`	Minimum score for live retrieval (default: 0.0)
`--output-dir PATH`	Output directory (default: dataset/results)

Main Orchestrator

uv run src/contextual_retrieval_experiment/main.py --config CONFIG.yaml

Runs multiple stages from a single YAML configuration. See YAML Configuration below.

Visualization

uv run src/visualize/plot_graphs.py [OPTIONS]

Option	Description
`--csv-path PATH`	Results CSV (default: results/table.csv)
`--output-dir PATH`	Output directory (default: results/plots)
`--tier {easy,medium}`	Difficulty tier to plot (default: medium)
`--metric {fuzzy_coverage,rigid_coverage}`	Metric to plot (default: fuzzy_coverage)
`--show / --no-show`	Display plot interactively
`--font-path PATH`	Custom font file (.ttf, .otf)

YAML Configuration

The recommended way to run experiments is via YAML config files. See config/ for examples.

Full Pipeline Example

pipeline:
  stages: [ingestion, retrieval, evaluation]

ingestion:
  file_paths:
    - dataset/raw_data/DnD5eSRD_md/DND5eSRD_001-018.md
    - dataset/raw_data/DnD5eSRD_md/DND5eSRD_019-035.md
    # ... more files
  source_dir: dataset/raw_data/DnD5eSRD_md
  collection_name: dnd5e_contextual_experiment
  prompt_template_path: src/contextual_retrieval_experiment/ingestion/constants/templates/contextual_retrieval_prompt_dnd.jinja2
  skip_contextual_retrieval: false
  model: gemini-2.5-pro
  chunk_size: 4000
  chunk_overlap: 100
  output_chunks_json: data/chunks_with_positions.json

retrieval:
  dataset_json: dataset/qa_sets/medium.json
  collection_name: dnd5e_contextual_experiment
  use_reranker: true
  reranker_type: cohere
  retrieve_k: 128
  reranker_top_n: 20
  output_dir: results/contextual/medium_cohere

evaluation:
  chunks_path: dataset/chunks/chunks_with_positions.json
  dataset_json: dataset/qa_sets/medium.json
  k: [5, 10, 20]
  output_dir: results/contextual/medium_cohere

Retrieval + Evaluation Only

pipeline:
  stages: [retrieval, evaluation]

retrieval:
  dataset_json: dataset/qa_sets/medium.json
  collection_name: dnd5e_contextual_experiment
  use_reranker: true
  reranker_type: cohere
  retrieve_k: 128
  reranker_top_n: 20
  output_dir: results/experiment_1

evaluation:
  chunks_path: dataset/chunks/chunks_with_positions.json
  dataset_json: dataset/qa_sets/medium.json
  k: [5, 10, 20]
  output_dir: results/experiment_1

Baseline Ingestion (No Contextual Retrieval)

pipeline:
  stages: ingestion

ingestion:
  file_paths:
    - dataset/raw_data/DnD5eSRD_md/DND5eSRD_001-018.md
  source_dir: dataset/raw_data/DnD5eSRD_md
  collection_name: dnd5e_base_experiment
  skip_contextual_retrieval: true
  output_chunks_json: data/chunks_with_positions.json

Stage Rules

stages must be a contiguous subsequence of [ingestion, retrieval, evaluation]
Valid: ingestion, [ingestion, retrieval], [retrieval, evaluation], [ingestion, retrieval, evaluation]
Invalid: [ingestion, evaluation] (skips retrieval)

Data Formats

Chunks JSON

Used for evaluation to map chunk IDs to character positions:

{
  "DND5eSRD_001-018.md::chunk_0": {
    "content": "# Introduction\n\nThe Dungeons & Dragons...",
    "document_path": "DND5eSRD_001-018.md",
    "start_char": 0,
    "end_char": 3950
  },
  "DND5eSRD_001-018.md::chunk_1": {
    "content": "## Character Creation\n\nYour first step...",
    "document_path": "DND5eSRD_001-018.md",
    "start_char": 3850,
    "end_char": 7800
  }
}

QA Dataset

Each question includes ground-truth passages with character positions:

[
  {
    "id": 0,
    "question": "What is the range of a longbow?",
    "answer": "The longbow has a range of 150/600 feet...",
    "passages": [
      {
        "content": "| Longbow | 1d8 Piercing | ...",
        "document_path": "DND5eSRD_047-063.md",
        "start_char": 12500,
        "end_char": 13200
      }
    ]
  }
]

Retrieval Output

Saved by the retrieval stage for offline evaluation:

{
  "summary": {
    "collection_name": "dnd5e_contextual_experiment",
    "retrieve_k": 128,
    "use_reranker": true,
    "reranker_type": "cohere",
    "reranker_top_n": 20,
    "reranker_model": "rerank-v4.0-pro"
  },
  "results": [
    {
      "question": "What is the range of a longbow?",
      "retrieved": [
        {
          "id": "DND5eSRD_047-063.md::chunk_5",
          "score": 0.95,
          "document_path": "DND5eSRD_047-063.md",
          "start_char": 12000,
          "end_char": 16000
        }
      ]
    }
  ]
}

Evaluation Output

{
  "summary": {
    "collection_name": "dnd5e_contextual_experiment",
    "k": 20,
    "average_fuzzy_coverage": 0.641,
    "average_rigid_coverage": 0.615
  },
  "results": [
    {
      "question": "What is the range of a longbow?",
      "retrieved": [...],
      "per_passage": [
        {
          "document_path": "DND5eSRD_047-063.md",
          "start_char": 12500,
          "end_char": 13200,
          "required_chunks": ["DND5eSRD_047-063.md::chunk_5"],
          "fuzzy": 1.0,
          "rigid": 1
        }
      ],
      "fuzzy_coverage": 1.0,
      "rigid_coverage": 1.0
    }
  ]
}

Prompt Templates

Contextual retrieval uses Jinja2 templates for context generation. Two templates are included:

Generic Template

src/contextual_retrieval_experiment/ingestion/constants/templates/contextual_retrieval_prompt_base.jinja2

<document>
{{ whole_document }}
</document>
Here are the chunks we want to situate within the whole document.

{% for chunk in chunks %}
<chunk id="{{ chunk.id }}">
{{ chunk.text }}
</chunk>
{% endfor %}

Please give a short succinct context to situate this chunk within the overall document for the purposes of improving search retrieval of the chunk. Answer only with the succinct context and nothing else.

D&D-Specific Template

src/contextual_retrieval_experiment/ingestion/constants/templates/contextual_retrieval_prompt_dnd.jinja2

Includes domain-specific instructions for D&D 5e content (class features, spells, monsters, etc.).

Creating Custom Templates

Your template receives:

whole_document: The full source document text
chunks: List of chunks with id and text fields

Tips:

Keep generated contexts short and factual
Include structural information (section headers, categories)
Tailor instructions to your domain

Qdrant Setup

Local Development (Docker)

# docker-compose.yml
version: "3.8"
services:
  qdrant:
    image: qdrant/qdrant:latest
    ports:
      - "6333:6333"
      - "6334:6334"
    environment:
      QDRANT__SERVICE__API_KEY: "local-dev-key"
    volumes:
      - ./qdrant_storage:/qdrant/storage

docker-compose up -d

export QDRANT_HOST=http://localhost:6333
export QDRANT_API_KEY=local-dev-key

Cloud (Qdrant Cloud)

Create a cluster at cloud.qdrant.io
Get your cluster URL and API key
Set environment variables:

export QDRANT_HOST=https://your-cluster-id.us-east4-0.gcp.cloud.qdrant.io
export QDRANT_API_KEY=your-api-key

Troubleshooting

Missing Environment Variables

ValueError: GOOGLE_API_KEY is required when contextual retrieval is enabled

Solution: Export the required API key, or use --skip-contextual-retrieval for baseline mode.

Reranker Errors

ValueError: VLLM reranker requires endpoint. Set RERANKER_ENDPOINT env var.

Solution: Either set RERANKER_ENDPOINT or switch to Cohere reranker (--reranker-type cohere).

Dimension Mismatch

QdrantException: Vector dimension mismatch

Solution: Collections expect 1536 dimensions (Cohere embed-v4.0). Delete and recreate the collection if you changed embedders.

Rate Limits

If you hit API rate limits:

Reduce ingestion batch sizes
Add delays between requests
Use --max-retries for automatic retry with backoff

Evaluation Shows Zero Coverage

Check that:

chunks_path points to the correct chunks JSON
document_path in your QA dataset matches the filenames in chunks
start_char and end_char are within the document bounds

Project Structure

rag-evaluation/
├── config/                               # YAML configuration files
│   ├── base/                             # Baseline (no contextual) configs
│   │   └── k_eval/                       # K-value evaluation configs
│   ├── contextual/                       # Contextual retrieval configs
│   │   └── k_eval/                       # K-value evaluation configs
│   ├── eval_only/                        # Evaluation-only configs
│   └── ingestion_retrieval_eval.yaml     # Full pipeline example
├── dataset/
│   ├── chunks/                      # Pre-built chunks JSON
│   ├── qa_sets/                     # Question-answer datasets
│   └── raw_data/
│       ├── DnD5eSRD_md/                  # Source markdown files (20 files)
│       └── DnD5eSRD.pdf                  # Original PDF source
├── results/                              # Experiment outputs
│   ├── base/                             # Baseline results
│   ├── contextual/                       # Contextual results
│   ├── plots/                            # Generated visualizations
│   ├── table.csv                         # Aggregated results for plotting
│   └── expanded_table.csv                # Detailed results table
├── src/
│   ├── contextual_retrieval_experiment/
│   │   ├── common/                       # Shared models, settings, utilities
│   │   ├── evaluation/                   # Evaluation logic and metrics
│   │   ├── ingestion/                    # Ingestion pipeline
│   │   │   ├── constants/templates/      # Jinja2 prompt templates
│   │   │   └── custom_rag_components/    # Custom pipeline components
│   │   ├── retrieval/                    # Retrieval pipeline
│   │   │   └── custom_components/        # vLLM reranker
│   │   └── main.py                       # YAML-driven orchestrator
│   └── visualize/
│       ├── plot_graphs.py                # Main visualization CLI
│       └── plot_recall_difference.py     # Recall difference plots
├── pyproject.toml                        # Python dependencies (uv)
├── uv.lock                               # Locked dependencies
├── run_self_hosted_reranker.md           # Instructions for setting up vLLM reranker
└── README.md

Contributing

Pull requests welcome! Please:

Keep changes focused and well-scoped
Follow existing code style
Update documentation for new features

Acknowledgments

Contextual Retrieval pattern by Anthropic
Built with datapizza-ai
Embeddings by Cohere
Vector search by Qdrant
D&D 5e System Reference Document 5.2.1 by Wizards of the Coast, licensed under CC-BY 4.0
Dataset and preprocessing tools: rag-dataset-builder

License

MIT License. See LICENSE for details.

The included D&D 5e SRD dataset is licensed under CC-BY 4.0. See the SRD license for details.

Name		Name	Last commit message	Last commit date
config		config
dataset		dataset
results		results
src		src
LICENSE		LICENSE
README.md		README.md
pyproject.toml		pyproject.toml
run_self_hosted_reranker.md		run_self_hosted_reranker.md
uv.lock		uv.lock

Uh oh!

Latest commit

History

Folders and files

Repository files navigation