What is RAG

import hashlib
import re

import numpy as np

STOPWORDS = {"a", "an", "the", "is", "are", "be", "can", "do", "does", "i", "my", "to",
             "of", "in", "on", "for", "and", "or", "what", "how", "it", "you", "your", "at"}


def words(text: str) -> list[str]:
    """Lowercase, drop stopwords, and strip a few suffixes (refunded -> refund)."""
    out = []
    for w in re.findall(r"[a-z0-9-]+", text.lower()):
        if w in STOPWORDS:
            continue
        for suffix in ("ing", "ed", "s"):
            if w.endswith(suffix) and len(w) > len(suffix) + 3:
                w = w[: -len(suffix)]
                break
        out.append(w)
    return out


def embed(text: str, dim: int = 256) -> np.ndarray:
    """Toy embedding: hashed bag of words, L2-normalised. Real models capture meaning."""
    vec = np.zeros(dim, dtype=np.float32)
    for word in words(text):
        bucket = int(hashlib.md5(word.encode()).hexdigest(), 16) % dim
        vec[bucket] += 1.0
    norm = np.linalg.norm(vec)
    return vec / norm if norm else vec


# 1. INDEX: documents -> chunks -> vectors
documents = {
    "refunds.md": "Annual plans can be refunded within 30 days of purchase. Monthly plans are not refunded.",
    "backups.md": "Backups run every night at 2 AM. You can restore any file from the last 90 days.",
    "security.md": "All backups are encrypted with AES-256. Two-factor authentication is required for admins.",
    "billing.md": "Invoices are emailed on the first day of each month. You can change your card in Settings.",
}
chunks = [{"source": name, "text": text} for name, text in documents.items()]
matrix = np.stack([embed(c["text"]) for c in chunks])          # shape (4, 256)


# 2. RETRIEVE: question -> vector -> most similar chunks
def retrieve(question: str, k: int = 2) -> list[dict]:
    scores = matrix @ embed(question)                          # cosine similarity
    best = np.argsort(-scores)[:k]
    return [{**chunks[i], "score": round(float(scores[i]), 3)} for i in best]


# 3. AUGMENT: build a grounded prompt
def build_prompt(question: str, hits: list[dict]) -> str:
    context = "\n".join(f"[{h['source']}] {h['text']}" for h in hits)
    return (
        "Answer using only the context. Cite sources in brackets. "
        "If the answer is not in the context, say you don't know.\n\n"
        f"Context:\n{context}\n\nQuestion: {question}"
    )


question = "Can I get a refund on an annual plan?"
hits = retrieve(question)
for h in hits:
    print(h["score"], h["source"])
print()
print(build_prompt(question, hits))