Skip to content

Repository files navigation

MedRAG: Production-Grade Medical Literature RAG System

Python License FastAPI Streamlit

Evidence-based medical Q&A grounded in PubMed abstracts. Hybrid retrieval (vector + BM25), cross-encoder reranking, RAGAs evaluation, and full observability via Langfuse -- all running locally with Ollama.

Disclaimer: This system is for educational and research purposes only. It does NOT constitute medical advice.


Architecture

graph TB
    subgraph Ingestion
        A[PubMed API] -->|Biopython Entrez| B[Raw Abstracts]
        B -->|LangChain Splitter| C[Chunked Documents]
    end

    subgraph Indexing
        C -->|Ollama nomic-embed-text| D[ChromaDB Vectors]
        C -->|rank_bm25| E[BM25 Index]
    end

    subgraph Retrieval
        F[User Query] --> G[Hybrid Retriever]
        G -->|Vector Search| D
        G -->|Keyword Search| E
        G -->|RRF Fusion k=60| H[Merged Candidates]
        H -->|Cross-Encoder| I[Reranked Top-N]
    end

    subgraph Generation
        I --> J[Prompt Builder]
        J --> K[LLM - Ollama / Groq]
        K --> L[Cited Answer with PMIDs]
    end

    subgraph Evaluation
        L --> M[RAGAs Metrics]
        I --> N[Custom Retrieval Metrics]
    end

    subgraph Observability
        K -.->|traces| O[Langfuse]
    end
Loading
User Question
     |
     v
+--------------------+     +------------------+
| ChromaDB (vector)  |     | BM25 (keyword)   |
| nomic-embed-text   |     | rank_bm25        |
+--------+-----------+     +--------+---------+
         |                          |
         +--------- RRF Fusion ----+
                     |
                     v
         +---------------------+
         | Cross-Encoder       |
         | ms-marco-MiniLM     |
         +----------+----------+
                    |
                    v
         +---------------------+
         | LLM Generation      |
         | llama3.1:8b (Ollama)|
         +----------+----------+
                    |
                    v
         Cited Answer [PMID: xxxxx]

Features

  • Hybrid Search -- vector similarity (ChromaDB) + BM25 keyword matching fused with Reciprocal Rank Fusion
  • Cross-Encoder Reranking -- ms-marco-MiniLM-L-6-v2 reranks candidates for high precision
  • PMID Citations -- every claim in the answer links back to a specific PubMed paper
  • Dual LLM Backends -- Ollama (local, private) with Groq cloud fallback
  • RAGAs Evaluation -- faithfulness, answer relevancy, context precision, context recall
  • Custom Retrieval Metrics -- hit rate, MRR, average rerank score
  • Langfuse Observability -- every query traced end-to-end (retrieval + generation)
  • FastAPI Backend -- /query, /health, /ingest, /stats endpoints
  • Streamlit UI -- 4-page app with query interface, evaluation dashboard, data management, architecture overview
  • Fully Local -- runs entirely on your machine with Ollama (no data leaves your network)

Quick Start

Prerequisites

Requirement Purpose
Python 3.11+ Runtime
Ollama Local LLM + embeddings
Docker + Compose Langfuse observability (optional)
Git Clone the repo

1. Clone and install

git clone <your-repo-url> medrag
cd medrag

# Option A: use Make (recommended)
make install

# Option B: manual
python -m venv .venv
# Linux/macOS:
source .venv/bin/activate
# Windows:
.venv\Scripts\activate

pip install -r requirements.txt
pip install -e .

2. Pull Ollama models

make pull-models

# or manually:
ollama pull llama3.1:8b
ollama pull nomic-embed-text

3. Configure environment

cp .env.example .env
# Edit .env:
#   ENTREZ_EMAIL=your_email@example.com   (required by NCBI)
#   GROQ_API_KEY=...                      (optional -- for cloud fallback)
#   LANGFUSE_PUBLIC_KEY=...               (optional -- for observability)
#   LANGFUSE_SECRET_KEY=...

4. Start Langfuse (optional)

make langfuse
# Open http://localhost:3000, create a project, copy API keys to .env

5. Ingest PubMed data

make ingest

# or with a custom query:
make ingest-query QUERY="hypertension treatment guidelines" MAX=500

6. Start the application

# Terminal 1: API
make api

# Terminal 2: UI
make ui

# Or run everything together:
make all

7. Query via API

curl -X POST http://localhost:8000/query \
  -H "Content-Type: application/json" \
  -d '{"question": "What is the mechanism of action of metformin?", "top_k": 3}'

Screenshots

TODO: Add screenshots of the Streamlit UI pages

Page Description
Query Chat interface with cited answers and source cards
Evaluation RAGAs metrics dashboard with per-question breakdown
Data Management Collection stats, health checks, ingestion controls
About Architecture diagram and tech stack

Evaluation Results

TODO: Fill in after running make eval

Metric Score
Faithfulness --
Answer Relevancy --
Context Precision --
Context Recall --
Retrieval Hit Rate --
MRR --
Avg Rerank Score --

Run the evaluation suite:

make eval

Results are saved to data/eval/results_<timestamp>.csv and viewable in the Streamlit evaluation dashboard.

Tech Stack

Layer Technology Why
LLM (primary) Ollama llama3.1:8b Runs locally, no API costs, data stays private
LLM (fallback) Groq llama-3.1-8b-instant Fast cloud inference when Ollama is unavailable
Embeddings Ollama nomic-embed-text (768d) High-quality open embeddings, local inference
Vector Store ChromaDB >= 1.0 (persistent) Simple, embedded, ships prebuilt wheels
Sparse Retrieval rank_bm25 Lightweight BM25 that catches exact-match terms vectors miss
Fusion Reciprocal Rank Fusion (k=60) Parameter-free merging of vector + BM25 ranked lists
Reranker cross-encoder/ms-marco-MiniLM-L-6-v2 ~80 MB model, strong reranking accuracy, fast on CPU
Chunking LangChain RecursiveCharacterTextSplitter Respects sentence boundaries, configurable overlap
Data Source PubMed via Biopython Entrez Authoritative biomedical literature, free API
API FastAPI + Uvicorn Async, auto-docs, Pydantic validation
UI Streamlit Rapid prototyping, built-in charts, session state
Evaluation RAGAs + custom metrics Standard RAG evaluation framework + retrieval-specific checks
Observability Langfuse (self-hosted) Open-source, full trace visibility, Docker deploy
Config YAML + Pydantic + python-dotenv Typed, validated, secrets separated from config

Design Decisions

Why hybrid search (vector + BM25)?

Vector search excels at semantic similarity ("heart attack" matches "myocardial infarction") but can miss exact terms that matter in medicine -- drug names, gene symbols, lab values. BM25 catches these lexical matches. Combining both with Reciprocal Rank Fusion consistently outperforms either retriever alone on medical benchmarks, without requiring weight tuning (RRF is parameter-free at k=60).

Why cross-encoder reranking?

Bi-encoder retrieval (ChromaDB) is fast but approximate -- it computes query and document embeddings independently. A cross-encoder jointly attends to the query-document pair, giving much more accurate relevance scores. We retrieve a broad set (top-10 hybrid candidates) and rerank down to top-3, getting the precision of a cross-encoder with the speed of bi-encoder retrieval. ms-marco-MiniLM-L-6-v2 was chosen because it is only ~80 MB, runs on CPU in <100 ms per query, and performs well on passage reranking tasks.

Why 512-token chunks with 128-token overlap?

PubMed abstracts are typically 200-400 words. A 512-token chunk usually fits an entire abstract or a substantial section. The 128-token overlap ensures that sentences split across chunk boundaries are captured in at least one chunk, preventing information loss at boundaries. This is especially important for medical text where a finding and its statistical significance may span a chunk boundary.

Why Ollama-first?

Medical data can be sensitive. Running the LLM locally via Ollama means no patient-adjacent text ever leaves the user's machine. Groq is available as a cloud fallback for environments where local GPU/CPU is insufficient, but the default configuration keeps everything private.

Why dict-based pipeline (not LangChain Documents)?

The retrieval pipeline uses plain list[dict] with a text key throughout, rather than LangChain Document objects. This gives full control over metadata fields (PMID, journal, year, rerank scores, RRF scores) without fighting framework abstractions. It also simplifies serialisation, testing (dicts are easy to construct in fixtures), and debugging.

Configuration

All settings in configs/config.yaml. Secrets go in .env (never committed).

Setting Default Description
llm.model llama3.1:8b Primary Ollama model
llm.fallback.model llama-3.1-8b-instant Groq fallback
embedding.model nomic-embed-text Ollama embedding model
chunking.chunk_size 512 Token chunk size
chunking.chunk_overlap 128 Overlap between chunks
retrieval.top_k 10 Candidates from hybrid search
retrieval.rerank_top_n 3 Final docs after reranking
vectorstore.collection_name medrag_abstracts ChromaDB collection
reranker.model cross-encoder/ms-marco-MiniLM-L-6-v2 Reranker model

Project Structure

medrag/
├── configs/
│   └── config.yaml              # all settings
├── data/
│   ├── raw/                     # raw PubMed JSONL files
│   ├── processed/               # BM25 index, chunked data
│   ├── chroma_db/               # ChromaDB persistent storage
│   └── eval/                    # evaluation dataset + results
├── src/
│   ├── config.py                # typed config loader (Pydantic)
│   ├── ingestion/
│   │   ├── pubmed_fetcher.py    # PubMed search + XML parsing
│   │   ├── chunker.py           # text splitting with metadata
│   │   └── pipeline.py          # end-to-end ingestion orchestrator
│   ├── embeddings/
│   │   ├── embedder.py          # Ollama embedding manager
│   │   └── vector_store.py      # ChromaDB wrapper (direct API)
│   ├── retrieval/
│   │   ├── bm25_retriever.py    # BM25 keyword retriever
│   │   ├── hybrid_retriever.py  # RRF fusion of vector + BM25
│   │   ├── reranker.py          # cross-encoder reranker
│   │   └── retrieval_pipeline.py # orchestrator
│   ├── generation/
│   │   ├── prompt_templates.py  # RAG system prompt + formatters
│   │   ├── llm_client.py        # Ollama + Groq client abstraction
│   │   └── generator.py         # RAGGenerator (main entry-point)
│   ├── evaluation/
│   │   ├── eval_dataset.py      # 20 medical Q&A pairs
│   │   ├── ragas_eval.py        # RAGAs evaluation pipeline
│   │   ├── ragas_evaluator.py   # CLI entry-point
│   │   ├── custom_metrics.py    # hit rate, MRR, avg rerank score
│   │   └── eval_dashboard.py    # Streamlit evaluation viewer
│   ├── api/
│   │   ├── main.py              # FastAPI app + middleware
│   │   └── models.py            # Pydantic request/response models
│   └── ui/
│       └── app.py               # Streamlit 4-page UI
├── tests/                       # pytest suite (29 tests)
├── docker-compose.yml           # Langfuse + API + UI services
├── Dockerfile                   # API container
├── Dockerfile.ui                # UI container
├── Makefile                     # dev shortcuts
├── requirements.txt             # pip dependencies
├── pyproject.toml               # editable install config
└── .env.example                 # environment template

Running Tests

make test

# with coverage:
make test-cov

Docker Deployment

# Start everything (Langfuse + API + UI):
docker compose up -d

# Or just Langfuse:
make langfuse

License

MIT

About

Evidence-based medical Q&A over PubMed abstracts. Hybrid retrieval (vector + BM25), cross-encoder reranking, RAGAs evaluation, Langfuse observability, running locally with Ollama. FastAPI + Streamlit.

Topics

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages