RAG

def chunk_text(text: str, size: int = 500, overlap: int = 50) -> list[str]:
    """Split text into overlapping windows of `size` characters."""
    chunks = []
    step = size - overlap
    for start in range(0, max(len(text) - overlap, 1), step):
        chunk = text[start : start + size].strip()
        if chunk:
            chunks.append(chunk)
    return chunks


with open("handbook.md") as f:
    chunks = chunk_text(f.read())

import chromadb

client = chromadb.PersistentClient(path="./index")
collection = client.get_or_create_collection(
    name="handbook",
    configuration={"hnsw": {"space": "cosine"}},
)

collection.add(
    ids=[f"chunk_{i}" for i in range(len(chunks))],
    documents=chunks,
    embeddings=embeddings.tolist(),
)

QUERY_PREFIX = "Represent this sentence for searching relevant passages: "


def retrieve(question: str, k: int = 3) -> list[str]:
    """Return the k chunks closest in meaning to the question."""
    q = model.encode(QUERY_PREFIX + question, normalize_embeddings=True)
    results = collection.query(
        query_embeddings=[q.tolist()],
        n_results=k,
    )
    return results["documents"][0]


contexts = retrieve("What's the refund window for annual plans?")

def build_prompt(question: str, contexts: list[str]) -> str:
    """Number the chunks and put them in front of the question."""
    sources = "\n\n".join(
        f"[{i}] {chunk}" for i, chunk in enumerate(contexts, start=1)
    )
    return (
        "Answer the question using only the sources below.\n"
        "Cite every claim like [1]. If the sources don't contain\n"
        "the answer, say you don't know.\n\n"
        f"Sources:\n{sources}\n\n"
        f"Question: {question}"
    )


prompt = build_prompt("What's the refund window for annual plans?", contexts)