From toy RAG to a real one: embeddings, a vector store and an LLM

# requirements.txt: exact pins, upgraded on purpose (as of Oct 2026)
openai==3.20.0
anthropic==1.13.0
chromadb==1.5.5
pytest==9.0.2

python3 -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt
export OPENAI_API_KEY="..."      # from your secrets manager, never committed
export ANTHROPIC_API_KEY="..."

"""Offline, stdlib-only helpers for the production RAG port. Run: python3 -I rag_core.py"""
import hashlib
import random
import re
import time


def chunk(doc_id, text, max_words=200, overlap=40):
    """Sentence-aware packing, same idea as the toy, with production-sized defaults."""
    sentences = re.split(r"(?<=[.!?])\s+", text.strip())
    chunks, current = [], []
    for sentence in sentences:
        words = sentence.split()
        if current and len(current) + len(words) > max_words:
            chunks.append(" ".join(current))
            current = current[-overlap:]
        current.extend(words)
    if current:
        chunks.append(" ".join(current))
    return [
        {"id": chunk_id(doc_id, c), "doc_id": doc_id, "n": i, "text": c}
        for i, c in enumerate(chunks)
    ]


def chunk_id(doc_id, text, model="text-embedding-3-small"):
    """Stable id: same doc + same text + same embedding model -> same id, so re-runs skip work."""
    digest = hashlib.sha256(f"{model}|{doc_id}|{text}".encode()).hexdigest()
    return digest[:32]


def batched(items, size):
    """Yield lists of at most `size` items (itertools.batched needs 3.12+)."""
    for start in range(0, len(items), size):
        yield items[start:start + size]


def with_retries(fn, retryable, attempts=5, base=0.5, cap=20.0, sleep=time.sleep):
    """Call fn(); on a retryable exception wait base*2^n (+ jitter, capped) and try again."""
    for attempt in range(attempts):
        try:
            return fn()
        except retryable:
            if attempt == attempts - 1:
                raise
            delay = min(cap, base * 2 ** attempt)
            sleep(delay / 2 + random.uniform(0, delay / 2))


def select_hits(ids, docs, distances, max_distance=0.6):
    """Keep only hits closer than the threshold. Cosine distance = 1 - cosine similarity."""
    return [
        {"id": i, "text": d, "distance": round(dist, 3)}
        for i, d, dist in zip(ids, docs, distances)
        if dist <= max_distance
    ]


def build_prompt(question, hits):
    if not hits:
        return None
    sources = "\n".join(f"[{h['id'][:8]}] {h['text']}" for h in hits)
    return (
        "Answer using ONLY the sources below. Cite source ids in brackets. "
        "If the sources do not contain the answer, say you don't know.\n\n"
        f"Sources:\n{sources}\n\nQuestion: {question}"
    )


if __name__ == "__main__":
    text = ("The API rate limit is 600 requests per minute per key. " * 3
            + "Keys are rotated every 90 days by the platform team.")
    parts = chunk("dev-docs/limits.md", text, max_words=20, overlap=5)
    print(len(parts), "chunks; first id", parts[0]["id"][:8])
    print([len(b) for b in batched(list(range(250)), 100)])

    calls = {"n": 0}

    def flaky():
        calls["n"] += 1
        if calls["n"] < 3:
            raise TimeoutError("simulated timeout")
        return "ok"

    print(with_retries(flaky, (TimeoutError,), sleep=lambda s: None), "after", calls["n"], "calls")
    hits = select_hits(["a" * 32, "b" * 32], ["Rate limit is 600 rpm.", "Office hours."], [0.21, 0.74])
    print(build_prompt("What is the rate limit?", hits))

4 chunks; first id 40006139
[100, 100, 50]
ok after 3 calls
Answer using ONLY the sources below. Cite source ids in brackets. If the sources do not contain the answer, say you don't know.

Sources:
[aaaaaaaa] Rate limit is 600 rpm.

Question: What is the rate limit?

"""Production port of the toy pipeline. Needs OPENAI_API_KEY and ANTHROPIC_API_KEY set."""
import os
import pathlib
import sys

import anthropic
import chromadb
import openai

from rag_core import batched, build_prompt, chunk, select_hits, with_retries

EMBED_MODEL = os.environ.get("RAG_EMBED_MODEL", "text-embedding-3-small")
LLM_MODEL = os.environ.get("RAG_LLM_MODEL", "claude-opus-5-5")
INDEX_DIR = os.environ.get("RAG_INDEX_DIR", ".rag_index")
EMBED_BATCH = 128

# max_retries=0: retries happen in with_retries, once, where we can see and test them.
oai = openai.OpenAI(timeout=30.0, max_retries=0)
llm = anthropic.Anthropic(timeout=60.0, max_retries=0)
store = chromadb.PersistentClient(path=INDEX_DIR)
collection = store.get_or_create_collection(
    name=f"docs-{EMBED_MODEL}", metadata={"hnsw:space": "cosine"}
)

OPENAI_RETRYABLE = (openai.RateLimitError, openai.APITimeoutError,
                    openai.APIConnectionError, openai.InternalServerError)
ANTHROPIC_RETRYABLE = (anthropic.RateLimitError, anthropic.APITimeoutError,
                       anthropic.APIConnectionError, anthropic.InternalServerError)


def embed(texts):
    vectors = []
    for batch in batched(texts, EMBED_BATCH):
        resp = with_retries(
            lambda: oai.embeddings.create(model=EMBED_MODEL, input=batch),
            OPENAI_RETRYABLE,
        )
        vectors.extend(item.embedding for item in sorted(resp.data, key=lambda d: d.index))
    return vectors


def index_folder(folder, embed_fn=embed):
    pieces = []
    for path in sorted(pathlib.Path(folder).glob("**/*.md")):
        pieces.extend(chunk(str(path), path.read_text(encoding="utf-8")))
    if not pieces:
        return 0
    known = set(collection.get(ids=[p["id"] for p in pieces])["ids"])
    new = [p for p in pieces if p["id"] not in known]
    for batch in batched(new, EMBED_BATCH):
        collection.upsert(
            ids=[p["id"] for p in batch],
            embeddings=embed_fn([p["text"] for p in batch]),
            documents=[p["text"] for p in batch],
            metadatas=[{"doc_id": p["doc_id"], "n": p["n"]} for p in batch],
        )
    return len(new)


def retrieve(question, k=5, embed_fn=embed):
    res = collection.query(query_embeddings=embed_fn([question]), n_results=k)
    return select_hits(res["ids"][0], res["documents"][0], res["distances"][0])


def answer(question):
    prompt = build_prompt(question, retrieve(question))
    if prompt is None:
        return "I don't know: nothing in the index is close enough to answer that."
    msg = with_retries(
        lambda: llm.messages.create(
            model=LLM_MODEL,
            max_tokens=1024,
            messages=[{"role": "user", "content": prompt}],
        ),
        ANTHROPIC_RETRYABLE,
    )
    return "".join(b.text for b in msg.content if b.type == "text")


if __name__ == "__main__":
    if sys.argv[1:2] == ["index"]:
        print("embedded", index_folder(sys.argv[2]), "new chunks")
    else:
        print(answer(" ".join(sys.argv[1:])))

# Needs both API keys. Index a folder of Markdown files, then ask.
python rag_prod.py index ./dev-docs
python rag_prod.py "What is the API rate limit per key?"

import pytest

from rag_core import batched, build_prompt, chunk, chunk_id, select_hits, with_retries

DOC = "Keys rotate every 90 days. " * 30


def test_chunks_overlap_and_ids_are_stable():
    a = chunk("wiki/keys.md", DOC, max_words=40, overlap=10)
    b = chunk("wiki/keys.md", DOC, max_words=40, overlap=10)
    assert len(a) > 1
    assert [c["id"] for c in a] == [c["id"] for c in b]
    assert a[0]["text"].split()[-10:] == a[1]["text"].split()[:10]


def test_id_changes_with_embedding_model():
    assert chunk_id("d", "t", model="m1") != chunk_id("d", "t", model="m2")


def test_batched_never_exceeds_size():
    assert [len(b) for b in batched(list(range(257)), 128)] == [128, 128, 1]


def test_retry_recovers_from_transient_errors():
    calls = []

    def flaky():
        calls.append(1)
        if len(calls) < 3:
            raise TimeoutError
        return "ok"

    assert with_retries(flaky, (TimeoutError,), sleep=lambda s: None) == "ok"
    assert len(calls) == 3


def test_retry_does_not_swallow_permanent_errors():
    def bad_request():
        raise ValueError("400: bad input")

    with pytest.raises(ValueError):
        with_retries(bad_request, (TimeoutError,), sleep=lambda s: None)


def test_retry_gives_up_after_attempts():
    def always_down():
        raise TimeoutError

    with pytest.raises(TimeoutError):
        with_retries(always_down, (TimeoutError,), attempts=3, sleep=lambda s: None)


def test_far_hits_are_dropped_and_empty_means_abstain():
    hits = select_hits(["x" * 32], ["unrelated"], [0.9], max_distance=0.6)
    assert hits == []
    assert build_prompt("q", hits) is None

python -m pytest -q test_rag_core.py
# 7 passed