RAG: Retrieval-Augmented Generation

OFFLINE - ingestion (run when documents change)
  documents -> parse/clean -> CHUNK -> EMBED -> store in vector DB (+ BM25 index) with metadata

ONLINE - answering (run for every question)
  question
     -> (rewrite query)
     -> RETRIEVE (hybrid: vector + BM25)  -> top 30
     -> RERANK                             -> top 5
     -> build PROMPT = instructions + retrieved chunks + question
     -> LLM generates answer (with citations)
     -> (check faithfulness) -> answer to user

import math, re
from collections import Counter, defaultdict

# ---------- OFFLINE: the knowledge base ----------
DOCUMENTS = {
    "returns.md": "Laptops can be returned within 30 days of delivery for a full refund. Opened software is not refundable. Refunds go back to the original payment method within five business days.",
    "shipping.md": "Shipping is free on orders over 50 dollars. Express shipping takes two days and costs 12 dollars. We ship to the US and Canada only.",
    "warranty.md": "Every laptop includes a one year warranty that covers manufacturing defects. Batteries are covered for six months. Water damage is not covered.",
    "support.md": "Support is available Monday to Friday from 9am to 6pm Eastern. Error ERR-4471 means the payment gateway timed out. Retry after five minutes.",
}

def chunk(text):
    return [s.strip() for s in re.split(r"(?<=[.!?])\s+", text) if s.strip()]

chunks = []                                   # (id, source, text)
for source, body in DOCUMENTS.items():
    for n, piece in enumerate(chunk(body)):
        chunks.append((f"{source}#{n}", source, piece))

# ---------- retrieval ----------
SYN = {"money": "refund", "back": "refund", "return": "returned", "cost": "costs", "deliver": "shipping",
       "broken": "defects", "hours": "available", "open": "available", "time": "available", "canada": "canada"}
tok = lambda t: [SYN.get(w, w) for w in re.findall(r"[a-z0-9-]+", t.lower())]
STOP = {"the", "a", "an", "is", "are", "to", "of", "for", "on", "in", "and", "do", "i", "my", "can", "how", "what", "much", "does", "it"}
content = lambda t: [w for w in tok(t) if w not in STOP]

T = [content(c[2]) for c in chunks]
N = len(T); avg = sum(map(len, T)) / N
df = Counter(w for t in T for w in set(t))

def bm25(q, i, k1=1.5, b=0.75):
    tf = Counter(T[i]); s = 0.0
    for w in content(q):
        if w in tf:
            s += math.log((N - df[w] + .5) / (df[w] + .5) + 1) * tf[w] * (k1 + 1) / (tf[w] + k1 * (1 - b + b * len(T[i]) / avg))
    return s

def retrieve(question, k=2):
    ranked = sorted(range(N), key=lambda i: -bm25(question, i))
    return [(chunks[i], round(bm25(question, i), 2)) for i in ranked[:k] if bm25(question, i) > 0]

# ---------- generation ----------
PROMPT = """Answer using ONLY the context. If it is not there, say "I don't have that information."
Cite the source in [brackets].

<context>
{context}
</context>

Question: {question}"""

def answer(question):
    hits = retrieve(question)
    if not hits:
        return "I don't have that information.", None
    context = "\n".join(f"[{c[1]}] {c[2]}" for c, _ in hits)
    prompt = PROMPT.format(context=context, question=question)
    best = hits[0][0]                                     # stand-in for the LLM
    return f"{best[2]} [{best[1]}]", prompt

for q in ["How long do I have to return a laptop?", "What does ERR-4471 mean?", "Who won the World Cup?"]:
    a, prompt = answer(q)
    print("Q:", q)
    print("A:", a)
    print()

PROMPT = """Answer using ONLY the context. If it is not there, say "I don't have that information."
Cite the source in [brackets].

<context>
{context}
</context>

Question: {question}"""

context = "[returns.md] Laptops can be returned within 30 days of delivery for a full refund.\n[returns.md] Opened software is not refundable."
print(PROMPT.format(context=context, question="How long do I have to return a laptop?"))

def answer(question, history):
    standalone = llm.rewrite(question, history)             # "and for laptops?" -> full question
    candidates = hybrid_search(standalone, k=40)            # BM25 + vectors, fused with RRF
    top = reranker.rerank(standalone, candidates, top_n=5)  # cross-encoder
    context = "\n\n".join(f"[{c.source}] {c.text}" for c in top)

    response = llm.chat(
        system="Answer ONLY from the context. Cite sources. If unsure, say you do not know.",
        user=f"<context>{context}</context>\n\nQuestion: {standalone}",
        temperature=0,
        stream=True,                                         # SSE to the browser
    )
    return response