Your first RAG pipeline in pure Python

"""A complete RAG retrieval pipeline in standard-library Python (3.10+)."""
from __future__ import annotations

import math
import re
from collections import Counter
from dataclasses import dataclass, field

DOCS: dict[str, str] = {
    "rate-limits.md": (
        "Each API key may send 600 requests per minute. "
        "Requests over the limit receive HTTP 429 with a Retry-After header. "
        "Batch endpoints count as one request per call, not per item."
    ),
    "webhooks.md": (
        "Webhook deliveries time out after 10 seconds. "
        "Failed deliveries are retried with exponential backoff for up to 72 hours. "
        "Each delivery carries an HMAC-SHA256 signature in the X-Signature header."
    ),
    "auth.md": (
        "API keys are scoped to a single workspace. "
        "Keys can be rotated from the dashboard without downtime. "
        "OAuth access tokens expire after 60 minutes."
    ),
}

STOPWORDS = {"a", "an", "the", "and", "or", "to", "of", "in", "is", "are", "can", "do",
             "does", "we", "you", "any", "for", "after", "how", "what", "i", "my", "per",
             "with", "not", "from", "up", "one", "each", "many", "long", "much"}


@dataclass
class Chunk:
    id: str
    text: str
    vec: dict[str, float] = field(default_factory=dict)


def tokenize(text: str) -> list[str]:
    words = re.findall(r"[a-z0-9]+", text.lower())
    return [w for w in words if w not in STOPWORDS]


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


def chunk(doc_id: str, text: str, max_words: int = 25, overlap_sentences: int = 1) -> list[Chunk]:
    """Pack whole sentences into chunks of at most ~max_words.

    Overlap is counted in sentences, not words, so every chunk starts at a
    sentence boundary. A single sentence longer than max_words becomes its own chunk.
    """
    sentences = split_sentences(text)
    pieces: list[list[str]] = []
    current: list[str] = []
    for sentence in sentences:
        words_now = sum(len(s.split()) for s in current)
        if current and words_now + len(sentence.split()) > max_words:
            pieces.append(current)
            carried = current[-overlap_sentences:] if overlap_sentences else []
            # Drop the carry if it would leave no room for the new sentence.
            if sum(len(s.split()) for s in carried) + len(sentence.split()) > max_words:
                carried = []
            current = list(carried)
        current.append(sentence)
    if current:
        pieces.append(current)
    return [Chunk(id=f"{doc_id}#{i}", text=" ".join(p)) for i, p in enumerate(pieces)]


def build_idf(token_lists: list[list[str]]) -> dict[str, float]:
    n = len(token_lists)
    df: Counter[str] = Counter()
    for tokens in token_lists:
        df.update(set(tokens))
    # Smoothed IDF: a word in every chunk still gets weight 1, never 0.
    return {term: math.log((1 + n) / (1 + count)) + 1 for term, count in df.items()}


def tfidf_vector(tokens: list[str], idf: dict[str, float]) -> dict[str, float]:
    counts = Counter(tokens)
    return {t: c * idf[t] for t, c in counts.items() if t in idf}


def cosine(a: dict[str, float], b: dict[str, float]) -> float:
    dot = sum(v * b.get(t, 0.0) for t, v in a.items())
    norm_a = math.sqrt(sum(v * v for v in a.values()))
    norm_b = math.sqrt(sum(v * v for v in b.values()))
    if norm_a == 0 or norm_b == 0:
        return 0.0
    return dot / (norm_a * norm_b)


@dataclass
class Index:
    chunks: list[Chunk]
    idf: dict[str, float]


def build_index(docs: dict[str, str], max_words: int = 25) -> Index:
    """Indexing: runs once, ahead of time."""
    chunks = [c for doc_id, text in docs.items() for c in chunk(doc_id, text, max_words)]
    token_lists = [tokenize(c.text) for c in chunks]
    idf = build_idf(token_lists)
    for c, tokens in zip(chunks, token_lists):
        c.vec = tfidf_vector(tokens, idf)
    return Index(chunks, idf)


def retrieve(index: Index, question: str, k: int = 2, min_score: float = 0.05) -> list[tuple[float, Chunk]]:
    """Query time: score every chunk, keep the best k above a floor."""
    q_vec = tfidf_vector(tokenize(question), index.idf)
    scored = sorted(((cosine(q_vec, c.vec), c) for c in index.chunks),
                    key=lambda pair: pair[0], reverse=True)
    return [(s, c) for s, c in scored[:k] if s >= min_score]


def build_prompt(question: str, hits: list[tuple[float, Chunk]]) -> str | None:
    if not hits:
        return None
    sources = "\n".join(f"[{c.id}] {c.text}" for _, c in hits)
    return (
        "Answer using ONLY the sources below. Cite the source id in brackets after each claim. "
        "If the sources do not contain the answer, reply exactly: I don't know.\n\n"
        f"Sources:\n{sources}\n\nQuestion: {question}\nAnswer:"
    )


def call_llm(prompt: str) -> str:
    """Stub. The second listing in this post replaces it with a real model call."""
    return "(the model's cited answer would appear here)"


def main() -> None:
    index = build_index(DOCS)
    for question in ["What happens if I exceed the rate limit?",
                     "How long are failed webhook deliveries retried?",
                     "Do you support SAML?"]:
        hits = retrieve(index, question)
        print("Q:", question)
        for score, c in hits:
            print(f"  {score:.2f}  {c.id:<16} {c.text[:52]}...")
        prompt = build_prompt(question, hits)
        print("  ->", "I don't know (no source cleared the floor)." if prompt is None else call_llm(prompt))
        print()


if __name__ == "__main__":
    main()

Q: What happens if I exceed the rate limit?
  0.23  rate-limits.md#1 Requests over the limit receive HTTP 429 with a Retr...
  0.23  rate-limits.md#0 Each API key may send 600 requests per minute. Reque...
  -> (the model's cited answer would appear here)

Q: How long are failed webhook deliveries retried?
  0.62  webhooks.md#0    Webhook deliveries time out after 10 seconds. Failed...
  0.30  webhooks.md#1    Failed deliveries are retried with exponential backo...
  -> (the model's cited answer would appear here)

Q: Do you support SAML?
  -> I don't know (no source cleared the floor).

# test_rag.py: run with "pytest -q" next to rag.py (or "python3 test_rag.py").
from rag import DOCS, build_index, build_prompt, chunk, retrieve, split_sentences


def test_every_chunk_starts_on_a_sentence_boundary() -> None:
    for doc_id, text in DOCS.items():
        starts = {s.split()[0] for s in split_sentences(text)}
        for c in chunk(doc_id, text, max_words=25):
            assert c.text.split()[0] in starts, c.id


def test_overlap_repeats_exactly_one_whole_sentence() -> None:
    first, second = chunk("rate-limits.md", DOCS["rate-limits.md"], max_words=25)[:2]
    shared = split_sentences(first.text)[-1]
    assert second.text.startswith(shared)


def test_oversized_sentence_becomes_its_own_chunk() -> None:
    long_sentence = " ".join(["word"] * 40) + "."
    pieces = chunk("x.md", long_sentence + " Short one.", max_words=25)
    assert [p.text for p in pieces] == [long_sentence, "Short one."]


def test_known_question_retrieves_its_source() -> None:
    hits = retrieve(build_index(DOCS), "How long are failed webhook deliveries retried?")
    assert hits and hits[0][1].id.startswith("webhooks.md")


def test_unanswerable_question_abstains() -> None:
    hits = retrieve(build_index(DOCS), "Do you support SAML?")
    assert hits == [] and build_prompt("Do you support SAML?", hits) is None


if __name__ == "__main__":
    for name, fn in list(globals().items()):
        if name.startswith("test_"):
            fn()
            print("ok", name)

# rag_real.py: needs network, an API key and two packages. Not runnable offline.
#   pip install "anthropic==1.13.0" "sentence-transformers==3.3.1"
#   export ANTHROPIC_API_KEY=...        (never hard-code it)
from __future__ import annotations

import anthropic
from sentence_transformers import SentenceTransformer

from rag import DOCS, Chunk, build_prompt, chunk

EMBEDDER = SentenceTransformer("BAAI/bge-small-en-v1.5")  # 33M params, 384 dims, runs on CPU
QUERY_PREFIX = "Represent this sentence for searching relevant passages: "  # BGE's query instruction
CLIENT = anthropic.Anthropic(timeout=30.0, max_retries=3)  # reads ANTHROPIC_API_KEY; retries 429/5xx with backoff
MODEL = "claude-opus-5-5"


def build_dense_index(docs: dict[str, str]) -> tuple[list[Chunk], list[list[float]]]:
    chunks = [c for doc_id, text in docs.items() for c in chunk(doc_id, text)]
    vectors = EMBEDDER.encode([c.text for c in chunks], normalize_embeddings=True)
    return chunks, vectors.tolist()


def retrieve_dense(chunks: list[Chunk], vectors: list[list[float]], question: str,
                   k: int = 2, min_score: float = 0.55) -> list[tuple[float, Chunk]]:
    q = EMBEDDER.encode([QUERY_PREFIX + question], normalize_embeddings=True)[0].tolist()
    # Unit-length vectors, so the dot product IS the cosine similarity.
    scored = sorted(((sum(a * b for a, b in zip(q, v)), c) for v, c in zip(vectors, chunks)),
                    key=lambda pair: pair[0], reverse=True)
    return [(s, c) for s, c in scored[:k] if s >= min_score]  # tune the floor on a golden set


def call_llm(prompt: str) -> str:
    response = CLIENT.messages.create(
        model=MODEL,
        max_tokens=4096,
        output_config={"effort": "low"},  # short grounded answers do not need deep thinking
        messages=[{"role": "user", "content": prompt}],
    )
    if response.stop_reason == "max_tokens":
        raise RuntimeError("answer truncated: raise max_tokens rather than ship half an answer")
    if response.stop_reason == "refusal":
        return "I can't help with that one."
    return "".join(block.text for block in response.content if block.type == "text")


def main() -> None:
    chunks, vectors = build_dense_index(DOCS)
    for question in ["What happens if I go over my request quota?", "Do you support SAML?"]:
        hits = retrieve_dense(chunks, vectors, question)
        prompt = build_prompt(question, hits)
        print("Q:", question)
        print("  ->", "I don't know." if prompt is None else call_llm(prompt))


if __name__ == "__main__":
    main()