Skip to content

Repository files navigation

grounded-rag

A compact RAG implementation built to demonstrate document ingestion, embeddings, pgvector retrieval, and grounded generation with citations. Deliberately small — see Scope below.

How it works

Documents (text)
      │
      ▼
   Chunking            common/chunking.py — fixed-size, overlapping
      │
      ▼
   Embedding           common/embeddings.py — sentence-transformers,
      │                all-MiniLM-L6-v2, local, no API key
      ▼
 Postgres + pgvector   schema.sql — documents + chunks tables
      │
      ▼
 Similarity search     retrieve.py — cosine distance, top-k
      │
      ▼
   Generation          generate.py — OpenRouter/DeepSeek, answers
      │                only from retrieved context, cites sources
      ▼
 Answer + citations

Exposed as a small FastAPI service (main.py): POST /ingest, POST /query, GET /health.

Setup

pip install -r requirements.txt
cp .env.example .env   # fill in OPENROUTER_API_KEY
docker compose up -d db
psql -h localhost -U postgres -f schema.sql
uvicorn main:app --reload

Or run the whole thing (app + db) via docker compose up.

curl -X POST localhost:8000/ingest -H 'Content-Type: application/json' -d '{
  "source": "example.md", "title": "Example Doc", "text": "..."
}'

curl -X POST localhost:8000/query -H 'Content-Type: application/json' -d '{
  "question": "What does this document say about X?"
}'

Testing

pip install -r requirements.txt pytest
pytest tests/ -v

28 tests: chunking (boundary/overlap correctness), embeddings (real model, dimension + determinism + similarity ordering — no mocking needed since it's local and deterministic), DB layer (mocked psycopg), ingest/retrieve orchestration, generation (mocked OpenRouter, schema validation, grounded/ ungrounded paths), and every FastAPI endpoint path via TestClient.

CI runs two jobs: unit tests, and a real integration check — applies schema.sql against an actual pgvector/pgvector Postgres service container, then runs eval/run_eval.py, a small retrieval regression check (14 questions against 6 real source documents drawn from my other repos' actual documentation, checking whether the expected source shows up in the top-3 retrieved chunks). It fails the build if hit rate drops below 80%. This is a hit-rate check, not a full Recall@K/Precision@K/MRR suite — deliberately small, just enough to prove retrieval works and to catch regressions, not a research benchmark.

A real bug this eval caught during development: the schema originally created an IVFFlat approximate index (lists = 100) on the chunks table. IVFFlat partitions rows into clusters and only probes one by default — with a corpus this small, most of those 100 clusters were empty, so queries would intermittently return zero results by probing an empty cluster instead of ranking real matches lower. Hit rate dropped to 21%. Removed the approximate index; a sequential scan is both exact and fast at this scale. See the comment in schema.sql for when it'd be worth adding back.

Scope — what this deliberately skips

Built as a scoped exercise, not a platform. On purpose, this does not include:

  • Hybrid search or reranking — pure vector similarity only.
  • A knowledge graph or structured entity memory.
  • Multiple coordinating agents — this is one retrieval step and one generation step, not an agent system.
  • Authentication, rate limiting, or a full observability stack.
  • Integrations with Notion, Slack, or any external knowledge source — ingestion takes plain text in, that's it.
  • A large-scale evaluation framework — the eval set above is 14 questions, not hundreds.

Those are real upgrades, deliberately left out to keep this a 2–4 day exercise rather than a new flagship project. See the "what I'd add next" list below if scaling this up.

What I'd add next, if I extended this

  • Hybrid search (BM25 + vector) and cross-encoder reranking.
  • Metadata filtering (document type, recency, approval status).
  • Faithfulness/groundedness scoring via an LLM-as-judge pass, not just citation presence.
  • Real production deployment with structured logging, metrics, and rate limiting.

License

MIT.

About

Compact RAG implementation: document ingestion, embeddings, pgvector retrieval, and grounded generation with citations. Scoped exercise, not a platform.

Topics

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages