drkrag — Neurology RAG & Evaluation

A retrieval-augmented generation (RAG) pipeline over neurology medical textbooks, with a rigorous LLM-as-judge evaluation harness and an experimental agentic (self-correcting) RAG variant built on LangGraph.

Ask a clinical question → the pipeline retrieves the most relevant passages from the textbook corpus → a language model answers grounded in those passages, keeping source metadata and a similarity score for each retrieved chunk. Multiple generator models are then scored head-to-head on the same questions so you can see which one actually answers best.

---

What's inside

• RAG pipeline — local embeddings + ChromaDB vector store + retrieve-then-generate.
• Pluggable generators — compare Mistral Small, Gemini Flash, and local Ollama models (e.g. qwen2.5:14b) on the same pipeline.
• Evaluation harness — four LLM-as-judge metrics (correctness, groundedness, answer relevance, retrieval relevance) logged to LangSmith, with a local, neutral judge (llama3.1:8b) so scoring is free, unlimited, and unbiased.
• Agentic RAG — a LangGraph graph that grades its own retrieved context and verifies its own answer, re-retrieving / rewriting / regenerating with a loop budget.
• Parallel ingestion — PyMuPDF + ProcessPoolExecutor for fast multi-PDF loading.

---

Architecture

flowchart LR
    PDFs["PDFs (pdfs/)"] --> CHUNK["chunk<br/>3000 / 500 overlap"]
    CHUNK --> EMB["embed<br/>all-MiniLM-L6-v2 · 384-d"]
    EMB --> STORE[("ChromaDB<br/>data/vector_store · cosine")]
    Q(["Question"]) --> RET["retrieve top-k=6"]
    STORE --> RET
    RET --> GEN["generate<br/>(Mistral / Gemini / qwen)"]
    GEN --> A(["Answer + sources + similarity"])

Loading

Ingestion is cached (fingerprinted by filename/size/mtime), and the ChromaDB store is persistent — so re-runs skip re-parsing and re-embedding.

The agentic variant wraps the same retrieve → generate steps in a decision loop; see notebooks/agentic_rag_security.md and the diagram in notebooks/agentic_rag.ipynb.

---

Repository structure

drkrag/
├── notebooks/          # all notebooks (run these) — see notebooks/README.md
│   ├── rag_evaluation.ipynb        # main eval harness (generators vs judges, LangSmith)
│   ├── agentic_rag.ipynb           # self-correcting LangGraph RAG + A/B vs linear
│   ├── notebook_parallelized.ipynb # parallel PDF loading (PyMuPDF)
│   ├── notebook.ipynb, rag_semantic_extended.ipynb, simple_example_loaders.ipynb
│   ├── agentic_rag_security.md     # security hardening reference
│   └── README.md                   # the bootstrap convention + prerequisites
├── src/                # importable package
│   ├── document_loader.py          # process_all_pdfs(), split_documents() (+ caching)
│   ├── embeddings.py               # EmbeddingManager, VectorStore (ChromaDB)
│   └── check_ocr_needed.py
├── utils/              # importable package
│   └── pdf_parallel_utils.py       # parallel PyMuPDF page-range loading
├── pdfs/               # source PDFs           (gitignored — bring your own)
├── data/               # vector store + caches (gitignored — auto-generated)
├── .env                # API keys              (gitignored — never commit)
└── requirements.txt

Note: pdfs/, data/, and .env are gitignored on purpose — the corpus is copyrighted, the vector store is large and regenerable, and .env holds secrets.

---

Setup

1. Python environment

python3 -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

2. Ollama (local models)

Install Ollama, make sure it's running (http://localhost:11434), and pull the models used by the notebooks:

ollama pull qwen2.5:14b     # local generator
ollama pull llama3.1:8b     # local LLM-as-judge (neutral, unlimited)

3. .env (project root)

Create a .env with the keys you use — only the ones you actually need:

MISTRAL_API_KEY=...      # Mistral generator (free tier)
GOOGLE_API_KEY=...       # Gemini generator (free tier; quota-limited)
LANGSMITH_API_KEY=...    # experiment tracking / tracing
LANGSMITH_TRACING=true   # optional: populate the LangSmith Traces tab
GROQ_API_KEY=...         # optional alternate hosted model

Local generation + local judging need no keys — only LangSmith needs one for tracking.

4. Add your PDFs

Drop your source PDFs into pdfs/. The first run embeds them into data/vector_store; subsequent runs reuse the cache.

---

Usage

Open a notebook in notebooks/ and run cells top-to-bottom. Run the bootstrap cell first — it resolves the project root, fixes imports, and loads .env (see notebooks/README.md for the convention).

• rag_evaluation.ipynb — builds the store, defines generators + judges, and runs a head-to-head evaluation against a LangSmith dataset. Results render side-by-side under Datasets & Experiments in smith.langchain.com.
• agentic_rag.ipynb — the self-correcting graph, plus an A/B comparison against the linear pipeline on the same dataset and judges.

---

How it works

Ingestion & embedding (src/)

• process_all_pdfs() loads PDFs (cached by a directory fingerprint) and split_documents() chunks them (default 3000 chars, 500 overlap).
• EmbeddingManager embeds chunks with SentenceTransformers — all-MiniLM-L6-v2 by default (384-dim, cosine), pluggable to larger models (e.g. Nemotron, with bf16 and asymmetric query/passage prefixes). Device auto-detects mps/cuda/cpu.
• VectorStore persists them in ChromaDB (cosine distance) at data/vector_store.

Retrieval & generation

• retrieve(query, k=6) embeds the query, pulls the top-6 chunks, and returns clean dicts with content, metadata (source file / page), and a similarity score.
• The generator answers using only the retrieved context, via LangChain's init_chat_model (swap the model string to change providers).

---

Evaluation

Each answer is graded by an LLM-as-judge on four metrics:

Metric	Needs a reference?	Question it answers
Correctness	yes	Is the answer factually right vs. the ground-truth answer?
Answer relevance	no	Does the answer address the question?
Groundedness	no	Is every claim supported by the retrieved docs (no hallucination)?
Retrieval relevance	no	Were the retrieved chunks actually relevant?

Separating them is diagnostic: high groundedness + low correctness ⇒ a retrieval problem; good retrieval + low groundedness ⇒ a generation problem.

The judge is llama3.1:8b running locally — chosen because it's neutral (not one of the generators, so no self-preference bias), unlimited, and never rate-limited. Verdicts use constrained JSON decoding with a graceful fallback so a malformed grade never crashes a run or silently produces a null score.

---

Agentic RAG

agentic_rag.ipynb turns the fixed retrieve → generate chain into a LangGraph graph with two gates where the model steers its own flow:

• grade documents — if the retrieved context is weak, rewrite the query and re-retrieve.
• verify — if the answer hallucinates, regenerate; if it's incomplete, fetch more context; otherwise return it.

Loop counters (MAX_LOOPS) plus a recursion_limit guarantee it always terminates with a best-effort answer. Because this lets untrusted document text influence decisions, it carries a larger attack surface — see the hardening reference in notebooks/agentic_rag_security.md.

---

Models at a glance

Role	Options
Generator	mistral-small-latest, gemini-flash-latest, qwen2.5:14b (local Ollama)
Judge	llama3.1:8b (local Ollama)
Embeddings	all-MiniLM-L6-v2 (default), pluggable
Vector store	ChromaDB (cosine, persistent)
Orchestration	LangChain · LangGraph · LangSmith

---

Data & privacy

The corpus is copyrighted reference material and is not included — bring your own PDFs. Two paths send content off-machine: LangSmith tracing (questions, retrieved passages, answers) and hosted generators (Mistral/Gemini). Keep the fully-local path (Ollama generator + local judge) for anything sensitive, and never put patient data (PHI) in this system. See the security doc for the full threat model and hardenings.