Complete RAG Tutorial

"""minirag.py -- a complete, dependency-free RAG pipeline.

Pipeline: chunk -> index -> retrieve -> prompt -> generate -> evaluate.
Python 3.9+ stdlib only. No API keys, no model downloads, no vector database.
"""

import math
import re
from collections import Counter, defaultdict

CORPUS = {
    "d1": "Hotel Bellamar is a beachfront property in Alanya. Check-in starts at 14:00 and check-out is at 11:00. The hotel has three outdoor pools and a spa with a hammam.",
    "d2": "Guests of Bellamar report that the buffet dinner is long at peak hours. The a la carte seafood restaurant requires a reservation one day in advance.",
    "d3": "Casa Marina is an adults only boutique hotel in Fethiye. It has two bars, one pool, and a small gym. Free wifi is available in all rooms.",
    "d4": "Reviews for Casa Marina mention noisy harbour music until midnight on weekends. Soundproofing was added to rooms 20 to 34 in March.",
    "d5": "Pension Yamac is a family run guesthouse outside Oludeniz. Breakfast is included and dinner is served on request. The nearest beach is a fifteen minute walk away.",
    "d6": "Pension Yamac offers airport transfers for a fixed fee of 40 euro per car. The transfer must be booked at least twenty four hours before arrival.",
    "d7": "Hotel Bellamar introduced a new loyalty programme in June. Members get late check-out until 14:00 and a twenty percent discount on spa treatments.",
    "d8": "Casa Marina is closed for maintenance every year from mid November to mid January. The swimming pool is unheated and closes in October.",
}

K1, B, TOP_K = 1.2, 0.75, 3

def tokenize(text):
    return re.findall(r"[a-z0-9']+", text.lower())

def chunk(doc_id, text, max_words=25):
    """Split on sentence boundaries, pack sentences into chunks under max_words."""
    sentences = [s.strip().rstrip(".") + "." for s in re.split(r"(?<=\.)\s+", text) if s.strip()]
    chunks, buf, words = [], [], 0
    for s in sentences:
        n = len(tokenize(s))
        if buf and words + n > max_words:
            chunks.append(" ".join(buf))
            buf, words = [], 0
        buf.append(s)
        words += n
    if buf:
        chunks.append(" ".join(buf))
    return [{"id": f"{doc_id}#{i}", "doc": doc_id, "text": c} for i, c in enumerate(chunks)]

CHUNKS = [c for doc_id in CORPUS for c in chunk(doc_id, CORPUS[doc_id])]

POSTINGS = defaultdict(list)
DF = Counter()
for i, c in enumerate(CHUNKS):
    for term, tf in Counter(tokenize(c["text"])).items():
        POSTINGS[term].append((i, tf))
        DF[term] += 1
DL = [len(tokenize(c["text"])) for c in CHUNKS]
AVGDL = sum(DL) / len(DL)
N = len(CHUNKS)

def idf(term):
    df = DF.get(term, 0)
    return math.log(1 + (N - df + 0.5) / (df + 0.5)) if df else 0.0

def bm25(query):
    scores = defaultdict(float)
    for term in set(tokenize(query)):
        if not DF.get(term):
            continue
        term_idf = idf(term)
        for i, tf in POSTINGS[term]:
            scores[i] += term_idf * (tf * (K1 + 1)) / (tf + K1 * (1 - B + B * DL[i] / AVGDL))
    return sorted(scores.items(), key=lambda kv: (-kv[1], CHUNKS[kv[0]]["id"]))

def retrieve(query, k=TOP_K):
    return [(CHUNKS[i], round(s, 4)) for i, s in bm25(query)[:k]]

def build_prompt(query, hits):
    context = "\n".join(f"[{c['id']}] {c['text']}" for c, _ in hits)
    return (
        "You are a hotel support assistant. Answer only from the context. "
        "Cite the chunk ids you used. If the context does not answer the question, say so.\n\n"
        f"Context:\n{context}\n\nQuestion: {query}\nAnswer:"
    )

def generate_offline(query, hits):
    """Deterministic extractive stand-in for the LLM: highest scoring sentence + citation."""
    qterms = set(tokenize(query))
    best, best_score = None, 0.0
    for c, _ in hits:
        for s in re.split(r"(?<=\.)\s+", c["text"]):
            terms = set(tokenize(s))
            if not terms:
                continue
            score = sum(idf(t) for t in qterms & terms) / math.sqrt(len(terms))
            if score > best_score:
                best, best_score = (s.strip(), c["id"]), score
    if not best or best_score == 0:
        return "The retrieved context does not answer this question.", []
    return best

GOLD = [
    ("what time is check in at hotel bellamar", "d1", "14:00"),
    ("must the transfer be booked in advance", "d6", "twenty four hours"),
    ("is casa marina open in december", "d8", "closed"),
    ("how much does the airport transfer cost", "d6", "40 euro"),
    ("which rooms are soundproofed at casa marina", "d4", "20 to 34"),
]

def evaluate(gold, k=TOP_K):
    rows, rr, rc, pc, ac, em = [], 0.0, 0, 0.0, 0, 0
    n = len(gold)
    for query, gold_doc, phrase in gold:
        hits = retrieve(query, k)
        docs = [c["doc"] for c, _ in hits]
        ranked = list(dict.fromkeys(docs))
        rank = ranked.index(gold_doc) + 1 if gold_doc in ranked else 0
        rr += 1 / rank if rank else 0
        rc += 1 if rank else 0
        pc += sum(1 for d in docs if d == gold_doc) / k
        context_hit = any(phrase in c["text"].lower() for c, _ in hits)
        answer, cites = generate_offline(query, hits)
        em += phrase in answer.lower()
        ac += context_hit
        rows.append((query, gold_doc, ",".join(ranked), rank or "-",
                     "yes" if context_hit else "no", phrase, answer[:46]))
    return rows, n, rr / n, rc / n, pc / n, ac / n, em / n

if __name__ == "__main__":
    print(f"corpus: {len(CORPUS)} docs -> {len(CHUNKS)} chunks, N={N}, avgdl={AVGDL:.2f} terms")
    demo_q = "what time is check in at hotel bellamar"
    print(f"\nretrieve({demo_q!r}) top-{TOP_K}:")
    for c, s in retrieve(demo_q):
        print(f"  {s:8.4f}  {c['id']}  {c['text'][:64]}")
    term, i = "bellamar", next(x for x, c in enumerate(CHUNKS) if c["doc"] == "d1")
    tf = dict(POSTINGS[term])[i]
    sat = (tf * (K1 + 1)) / (tf + K1 * (1 - B + B * DL[i] / AVGDL))
    print(f"\nBM25 worked example, term={term!r} chunk={CHUNKS[i]['id']}: df={DF[term]} "
          f"idf=ln(1+({N}-{DF[term]}+0.5)/({DF[term]}+0.5))={idf(term):.4f} tf={tf} "
          f"len={DL[i]} saturation={sat:.4f} contribution={idf(term) * sat:.4f}")
    rows, n, mrr, rec, prec, actx, emsc = evaluate(GOLD)
    print(f"\nevaluation set (n={n}):")
    for q, gd, docs, rk, ch, ph, ans in rows:
        print(f"  {gd:<4} rank={rk:<3} ctx={ch:<4} {ph:<19} {q}")
    print(f"\nrecall@{TOP_K}={rec:.2f}  precision@{TOP_K}={prec:.2f}  MRR={mrr:.3f}")
    print(f"answer-in-context={actx:.2f}  answer-match={emsc:.2f}")

corpus: 8 docs -> 13 chunks, N=13, avgdl=16.08 terms

retrieve('what time is check in at hotel bellamar') top-3:
    7.7569  d1#0  Hotel Bellamar is a beachfront property in Alanya. Check-in star
    3.9674  d7#0  Hotel Bellamar introduced a new loyalty programme in June. Membe
    3.6853  d2#0  Guests of Bellamar report that the buffet dinner is long at peak

BM25 worked example, term='bellamar' chunk=d1#0: df=3 idf=ln(1+(13-3+0.5)/(3+0.5))=1.3863 tf=1 len=21 saturation=0.8887 contribution=1.2320

evaluation set (n=5):
  d1   rank=1   ctx=yes  14:00               what time is check in at hotel bellamar
  d6   rank=1   ctx=yes  twenty four hours   must the transfer be booked in advance
  d8   rank=2   ctx=yes  closed              is casa marina open in december
  d6   rank=1   ctx=yes  40 euro             how much does the airport transfer cost
  d4   rank=1   ctx=yes  20 to 34            which rooms are soundproofed at casa marina

recall@3=1.00  precision@3=0.40  MRR=0.900
answer-in-context=1.00  answer-match=0.40

BM25 worked example, term='bellamar' chunk=d1#0: df=3 idf=ln(1+(13-3+0.5)/(3+0.5))=1.3863 tf=1 len=21 saturation=0.8887 contribution=1.2320

d6#1 bm25=3.3482  The transfer must be booked at least twenty four hours before arrival.
d6#0 bm25=2.3582  Pension Yamac offers airport transfers for a fixed fee of 40 euro per car.

sentence scores: 0.8663 (wrong sentence) vs 0.5970 (right sentence)
term idf: airport=2.2336  transfer=2.2336  transfers=2.2336  the=0.7673  cost=0.0000 (df=0)

python3 -m venv .venv && source .venv/bin/activate
pip install sentence-transformers qdrant-client openai
export OPENAI_API_KEY="sk-..."   # or GOOGLE_API_KEY for Vertex/Gemini

# embed: dense vectors instead of BM25  (requires: pip install sentence-transformers)
from sentence_transformers import SentenceTransformer
model = SentenceTransformer("BAAI/bge-small-en-v1.5")
vectors = model.encode([c["text"] for c in CHUNKS], normalize_embeddings=True)

# index and retrieve: a real vector store  (requires: pip install qdrant-client)
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct
db = QdrantClient(path="./qdrant_data")
db.create_collection("chunks", vectors_config=VectorParams(size=384, distance=Distance.COSINE))
db.upsert("chunks", points=[PointStruct(id=i, vector=vectors[i].tolist(), payload=CHUNKS[i]) for i in range(len(CHUNKS))])
hits = db.query_points("chunks", query=model.encode([query], normalize_embeddings=True)[0].tolist(), limit=5)

# hybrid: fuse lexical and dense rankings, then rerank with a cross-encoder
def rrf(rankings, k=60):
    fused = {}
    for ranking in rankings:
        for pos, doc in enumerate(ranking):
            fused[doc] = fused.get(doc, 0) + 1 / (k + pos + 1)
    return sorted(fused, key=fused.get, reverse=True)

# generate: an actual LLM over the assembled prompt  (requires: OPENAI_API_KEY)
from openai import OpenAI
client = OpenAI()
resp = client.chat.completions.create(model="gpt-4o-mini", messages=[{"role": "user", "content": build_prompt(query, hits)}])
answer = resp.choices[0].message.content