ReflectRAG

A self-correcting RAG (Retrieval-Augmented Generation) system that actually checks its own work before giving you an answer.

I built this to understand how AI agents really work under the hood. No LangChain, no LangGraph, and no frameworks hiding the interesting parts. Just Python, FAISS, and the OpenAI SDK, wired together by hand.

Why this exists

Most RAG tutorials follow the same pattern: chunk a PDF, embed it, retrieve top-K results, throw them at GPT, and call it a day. The problem? FAISS always returns K results even when none of them are relevant. And LLMs will hallucinate even when you tell them not to. There's no quality gate anywhere.

ReflectRAG solves this with two self-correction loops:

Amber Loop - Relevance re-routing. After retrieval, an LLM grades each chunk for relevance. If nothing passes, the system rewrites the query and retries retrieval (up to 2 attempts). This catches the "FAISS returned garbage" failure mode.

Red Loop - Hallucination retry. After generating an answer, a separate LLM call verifies that every claim is grounded in the source context. If it catches unsupported claims, the system regenerates (up to 3 retries). This catches the "GPT made stuff up" failure mode.

The whole orchestrator is ~60 lines of Python while loops with if/break routing. That's it. That's what LangGraph's add_conditional_edges() does under the hood.

How it works

Each box in that diagram is a Python function in src/nodes.py with the same signature: takes an AgentState, returns an AgentState. The orchestrator in src/orchestrator.py just calls them in order with conditional branching.

Tech stack

What	How
Orchestration	Custom state machine (src/orchestrator.py)
LLM calls	OpenAI SDK — gpt-4o for generation, gpt-4o-mini for grading
Vector search	FAISS (faiss-cpu) — raw library, no wrappers
Embeddings	OpenAI text-embedding-3-small (1536-dim)
Data validation	Pydantic v2
PDF parsing	pypdf
API	FastAPI + Uvicorn
Tests	pytest (all LLM calls mocked)

Project structure

├── main.py                 # CLI — index PDFs or ask questions
├── src/
│   ├── config.py           # All tuneable params (chunk size, models, retries)
│   ├── models.py           # Pydantic models: Document, AgentState
│   ├── indexer.py          # PDF loading, recursive chunking, FAISS index build
│   ├── embeddings.py       # OpenAI embedding wrapper with batching
│   ├── llm.py              # LLM call wrapper (standard + structured JSON output)
│   ├── nodes.py            # Pipeline nodes: rewrite, retrieve, grade, generate, check
│   ├── orchestrator.py     # The state machine that connects everything
│   ├── api.py              # FastAPI REST endpoints
│   ├── exceptions.py       # Custom exception classes
│   └── logger.py           # Structured logging setup
├── tests/                  # Unit tests for every module (mocked LLM, no API calls)
├── data/                   # Drop your PDFs here
├── requirements.txt
└── pyproject.toml

Getting started

Prerequisites

• Python 3.11+
• An OpenAI API key

Setup

# Clone and enter the project
git clone https://github.com/dcs-soni/reflectRAG.git
cd reflectRAG

# Create a virtual environment
python -m venv venv
source venv/bin/activate        # On Windows: venv\Scripts\activate

# Install dependencies
pip install -r requirements.txt

# Set up your API key
cp .env.example .env
# Open .env and paste your OpenAI key

Index a PDF

Drop any PDF into the data/ folder and build the FAISS index:

python main.py --index data/your-document.pdf

This reads the PDF, chunks the text (800 chars with 100 char overlap), generates embeddings via OpenAI, and saves the FAISS index + chunk metadata to faiss_index/.

Ask questions

python main.py --ask "What are the key terms of the agreement?"

The CLI will run the full pipeline — rewrite, retrieve, grade, generate, check — and print the answer along with pipeline metrics (grounding status, retry count, rewrite count, etc.).

Run the API server

uvicorn src.api:app --reload --port 8000

API endpoints:

Method	Endpoint	What it does
GET	/api/health	Liveness check
GET	/api/index/status	Check if an index exists + chunk count
POST	/api/index	Upload a PDF and build the index
POST	/api/ask	Ask a question, get answer + step trace

Interactive docs at http://localhost:8000/docs.

Run tests

pytest tests/ -v

All tests mock the OpenAI API — no real API calls, no cost, runs completely offline.

Design decisions worth knowing

• Two-model strategy. gpt-4o for answer generation (quality matters), gpt-4o-mini for grading and checking (it's just binary classification, cheap and fast).
• Fail-open vs fail-safe. The relevance grader defaults to including a document if the API fails (better noisy context than missing info). The hallucination checker defaults to rejecting the answer if the API fails (better to retry than pass through an unverified answer).
• Structured outputs. Grading and hallucination checks use OpenAI's beta.parse to force responses into strict Pydantic schemas. No regex parsing, no "I think so" answers — just {"relevant": "yes"} or {"grounded": "no"}.
• Chunk size 800, overlap 100. Small enough for precise retrieval, large enough to carry context. Configurable in src/config.py.

What I learned building this

The biggest takeaway: most of the "framework magic" in LangChain/LangGraph is surprisingly thin abstraction over basic Python patterns. A state machine is a while loop. Conditional routing is an if statement. Structured output is a Pydantic model. Once you see it, you can't unsee it.

---

Built by Divyanshu Soni