RAG Step by Step: Chat with Your Documents in Python

👁 120 reads

RAG (retrieval-augmented generation) lets you chat with your own documents: the program first searches for relevant fragments, then the model answers based on them and cites the source. We built this from scratch in about 70 lines of Python: three legal documents in PDF, embeddings computed locally in Ollama, vector database in numpy (and for comparison, in Chroma). Out of 11 questions asked in plain language, the model answered correctly in 8 cases, honestly admitted it couldn’t find the answer in 2, and never hallucinated. Below you’ll find all the code, all results, and two improvements we tested. This article is part of the AI Agents from Scratch series.

How RAG Works in Four Steps

1. Split: you cut documents into chunks (in our case: one article of law, long articles split into parts with overlap).
2. Embeddings: you convert each chunk into a vector of numbers that captures its meaning. You do this once.
3. Search: you convert the question into a vector too and search for chunks with the closest vectors.
4. Answer: you attach the 5 best chunks to the prompt and tell the model to answer only based on them, citing the source.

We explain what embeddings and vector databases are, without code, in the article Embeddings and Vector Databases Explained. A theoretical introduction to RAG and its weak points is in the older article RAG: How It Works.

Step 1: Setup

1. Install Ollama and download the embedding model: ollama pull bge-m3 (1.2 GB on disk, 664 MB in GPU memory for us).
2. In your Python environment: pip install numpy pypdf requests openai python-dotenv.
3. Put PDFs in the dane/ustawy folder. Legal texts are not protected by copyright (Article 4 of the Copyright Act in Poland), so you can freely use them for testing.

You can also compute embeddings via API (OpenAI, Google, Mistral), but locally you pay nothing and documents never leave your computer. Without a GPU, Ollama will compute them on the CPU, just slower.

Step 2: Complete Code

"""04 - RAG step by step: chat with your own documents (PDF), embeddings locally in Ollama, database = numpy.
Step 1 (once):  python 04_rag.py indeksuj dane/ustawy
Step 2 (ask):   python 04_rag.py pytaj "How many vacation days does an employee get after 12 years?"
Requires: Ollama with embedding model (ollama pull bge-m3) and OpenRouter key for answers.
"""
import json, os, pathlib, re, sys
import numpy as np
import requests
from dotenv import load_dotenv
from openai import OpenAI
from pypdf import PdfReader

load_dotenv()
EMBED = os.getenv("EMBED_MODEL", "bge-m3")          # or nomic-embed-text
MODEL = os.getenv("MODEL", "openai/gpt-5.6-luna")
INDEKS = pathlib.Path(f"indeks-{EMBED.replace(':', '_')}.npz")
OLLAMA = os.getenv("OLLAMA_URL", "http://localhost:11434")

# --- 1. Text from PDF and split into chunks (here: by articles of law) ------------
def kawalki(plik: pathlib.Path, max_zn=1200):
    tekst = "n".join(p.extract_text() or "" for p in PdfReader(plik).pages)
    tekst = re.sub(r"Dziennik Ustaws*Ws*d+s*Ws*Poz. d+", " ", tekst)    # page footers
    czesci = re.split(r"s(?=Art. d+[a-z]?(?:[d+])?. )", tekst)              # new article = new chunk
    for cz in czesci:
        cz = re.sub(r"s+", " ", cz).strip()
        art = (re.match(r"Art. (d+[a-z]?(?:[d+])?).", cz) or [None, "intro"])[1]
        for i in range(0, len(cz), max_zn - 200):    # long articles split with 200-char overlap
            yield {"plik": plik.name, "art": art, "tekst": cz[i:i + max_zn]}
            if i + max_zn >= len(cz): break

# --- 2. Embeddings: text -> vector of numbers (locally, via Ollama) -------------
def embed(teksty, prefiks=""):
    wektory = []
    for i in range(0, len(teksty), 32):
        r = requests.post(f"{OLLAMA}/api/embed", json={"model": EMBED, "input": [prefiks + t for t in teksty[i:i + 32]]}, timeout=600)
        wektory += r.json()["embeddings"]
    v = np.array(wektory, dtype=np.float32)
    return v / np.linalg.norm(v, axis=1, keepdims=True)   # normalization: dot product = cosine similarity

PREF_DOK, PREF_PYT = ("search_document: ", "search_query: ") if "nomic" in EMBED else ("", "")

def indeksuj(katalog):
    k = [x for p in sorted(pathlib.Path(katalog).glob("*.pdf")) for x in kawalki(p)]
    v = embed([x["tekst"] for x in k], PREF_DOK)
    np.savez(INDEKS, wektory=v, meta=np.array([json.dumps(x, ensure_ascii=False) for x in k]))
    print(f"saved {len(k)} chunks, {v.shape[1]}-dimensional vectors -> {INDEKS}")

# --- 3. Search: question -> vector -> closest chunks -------------------
def szukaj(pytanie, k=5):
    d = np.load(INDEKS)
    podob = d["wektory"] @ embed([pytanie], PREF_PYT)[0]
    top = np.argsort(-podob)[:k]
    return [{**json.loads(d["meta"][i]), "podobienstwo": round(float(podob[i]), 3)} for i in top]

# --- 4. Model's answer only based on found fragments -------------
def pytaj(pytanie):
    fragmenty = szukaj(pytanie)
    kontekst = "nn".join(f"[{f['plik']}, art. {f['art']}]n{f['tekst']}" for f in fragmenty)
    client = OpenAI(base_url="https://openrouter.ai/api/v1", api_key=os.environ["OPENROUTER_API_KEY"])
    r = client.chat.completions.create(model=MODEL, messages=[
        {"role": "system", "content": "Answer only based on the law fragments provided below. At the end, cite the "
         "source in parentheses: [file, art.]. If the fragments don't contain the answer, write: Not in the documents."},
        {"role": "user", "content": f"FRAGMENTS:n{kontekst}nnQUESTION: {pytanie}"}])
    for f in fragmenty:
        print(f"  found: {f['plik']} art. {f['art']} (similarity {f['podobienstwo']})")
    print(f"nANSWER: {r.choices[0].message.content}n(cost: {getattr(r.usage, 'cost', 0) or 0:.5f} USD)")

if __name__ == "__main__":
    {"indeksuj": indeksuj, "pytaj": pytaj}[sys.argv[1]](sys.argv[2])

Key decisions in this code:

• Split by articles. The regular expression looks for “Art. 154.” or “Art. 167[2].” (that’s how pypdf writes superscripts from the Official Journal). One article is usually one legal concept, so this is a better split than cutting every 1000 characters in the middle of a sentence.
• Removal of footers “Dziennik Ustaw - 23 - Poz. 1245” (Official Journal - 23 - Item 1245), which otherwise end up in the middle of articles.
• Vector normalization: after this, simple matrix multiplication gives cosine similarity. The entire “vector search” is one line here: d["wektory"] @ embed([pytanie])[0].
• System prompt with fallback: “If the fragments don’t contain the answer, write: Not in the documents”. This is what prevented the model from hallucinating in our test.

Step 3: Indexing

The command python 04_rag.py indeksuj dane/ustawy cut 149 pages into 920 chunks and saved the vectors (1024 numbers per chunk) to the file indeks-bge-m3.npz. It took 155 seconds, but the GPU was heavily loaded by another program at the time, so on a free GPU it will be faster. You do this once, then just ask questions.

Step 4: Questions and Answers

We asked 11 questions in the language an ordinary person uses, not a lawyer. Ten have answers in the laws, one (about highway speed) deliberately doesn’t.

Result: 8 correct answers out of 10 questions that have answers in the documents, plus correct refusal for the off-topic question. Both failures came from search, not from the model. The model got the wrong fragments and said plainly that they don’t contain the answer. That’s how a well-configured RAG should work.

Here’s the full run for the overtime question:

> python 04_rag.py pytaj "Boss makes me stay after hours at night. What's the overtime rate?"
  found: kodeks-pracy-DU-2026-1245.pdf art. 151[8] (similarity 0.682)
  found: kodeks-pracy-DU-2026-1245.pdf art. 151[1] (similarity 0.663)
  found: kodeks-pracy-DU-2026-1245.pdf art. 151[6] (similarity 0.644)
  found: kodeks-pracy-DU-2026-1245.pdf art. 151[7] (similarity 0.621)
  found: kodeks-pracy-DU-2026-1245.pdf art. 151[4] (similarity 0.616)

ANSWER: For each overtime hour falling at night, you're entitled to a supplement of **100% of pay**. Additionally, for each hour of work during night time, you're entitled to a night supplement of **20% of the hourly rate derived from the minimum wage**. Night time covers 8 hours between 9:00 PM and 7:00 AM. [kodeks-pracy-DU-2026-1245.pdf, art. 151[1] § 1 pt 1 lit. a, art. 151[8] § 1, art. 151[7] § 1]
(cost: 0.00084 USD)

The model combined three fragments (overtime supplement, night supplement, definition of night time) and cited the provision for each. And here’s what a search failure looks like:

> python 04_rag.py pytaj "I've been working for 12 years. How many vacation days am I entitled to?"
  found: kodeks-pracy-DU-2026-1245.pdf art. 167[3] (similarity 0.693)
  found: kodeks-pracy-DU-2026-1245.pdf art. 167[2] (similarity 0.685)
  found: kodeks-pracy-DU-2026-1245.pdf art. 186 (similarity 0.669)
  found: kodeks-pracy-DU-2026-1245.pdf art. 182[3] (similarity 0.664)
  found: kodeks-pracy-DU-2026-1245.pdf art. 172 (similarity 0.66)

ANSWER: Not in the documents. The provided fragments only state that vacation on demand can be a maximum of 4 days per calendar year; they don't specify the full vacation entitlement for someone working 12 years. [kodeks-pracy-DU-2026-1245.pdf, art. 167[2]-167[3]]
(cost: 0.00063 USD)

The search found provisions about vacation on demand and parental leave, because the question “how many vacation days” is very similar to them. Article 154 begins with the words “Wymiar urlopu wynosi: 20 dni… 26 dni” (Vacation entitlement is: 20 days… 26 days), using different vocabulary than the question.

Fix 1: Rewriting the Question in Legal Language

A popular technique (query rewriting) is to have the model rewrite the question in the language of the documents before searching. We tested it on 10 questions:

The other 6 questions didn’t change position. In total: correct article in the top five in 9 out of 10 cases (previously 8). The question about mom shows the limit of this method: the law talks about “close person,” and neither the question nor its rewritten version contains that phrase. Hybrid search (keywords plus vectors) or a synonym dictionary would help here. Rewriting cost 0.0017 USD for 10 questions and adds about a second to each answer.

Fix 2: Similarity Threshold

For the highway question, the best chunk had similarity 0.43, while for questions with answers in the documents it was 0.64-0.72. That’s a natural threshold: if the best result is below about 0.5, the program can immediately answer “I don’t have that in the documents” without asking the model. The threshold must be tuned to your embedding model and documents. For nomic-embed-text on the same laws, all results, even nonsensical ones, fell between 0.69 and 0.77, so such a threshold wouldn’t make sense.

Same Thing in Chroma Database

Numpy is sufficient for tens of thousands of chunks. With larger numbers or when you want to filter by metadata (e.g., “search only in Labor Code”), a vector database is more convenient. The same RAG on Chroma 1.5.9:

"""04b - Same RAG, but with Chroma vector database instead of numpy (disk storage, metadata filtering).
python 04b_rag_chroma.py indeksuj dane/ustawy
python 04b_rag_chroma.py szukaj "What's the maximum security deposit?" [file.pdf]
"""
import importlib, pathlib, sys
import chromadb

rag = importlib.import_module("04_rag")          # we take chunking and embeddings from 04_rag.py
baza = chromadb.PersistentClient(path="chroma_db")
kolekcja = baza.get_or_create_collection(f"ustawy-{rag.EMBED.replace(':', '_')}", metadata={"hnsw:space": "cosine"})

def indeksuj(katalog):
    k = [x for p in sorted(pathlib.Path(katalog).glob("*.pdf")) for x in rag.kawalki(p)]
    v = rag.embed([x["tekst"] for x in k], rag.PREF_DOK)
    for i in range(0, len(k), 500):  # Chroma accepts data in batches
        kolekcja.upsert(ids=[f"{x['plik']}#{n}" for n, x in enumerate(k[i:i + 500], i)],
                        embeddings=v[i:i + 500].tolist(), documents=[x["tekst"] for x in k[i:i + 500]],
                        metadatas=[{"plik": x["plik"], "art": x["art"]} for x in k[i:i + 500]])
    print(f"collection has {kolekcja.count()} chunks")

def szukaj(pytanie, plik=None):
    w = kolekcja.query(query_embeddings=rag.embed([pytanie], rag.PREF_PYT).tolist(), n_results=3,
                       where={"plik": plik} if plik else None)   # metadata filter: only one law
    for meta, odl, tekst in zip(w["metadatas"][0], w["distances"][0], w["documents"][0]):
        print(f"{meta['plik']} art. {meta['art']} (similarity {1 - odl:.3f}): {tekst[:110]}...")

if __name__ == "__main__":
    indeksuj(sys.argv[2]) if sys.argv[1] == "indeksuj" else szukaj(sys.argv[2], *sys.argv[3:4])

> python 04b_rag_chroma.py indeksuj dane/ustawy
collection has 920 chunks
(indexing took 126 s)

> python 04b_rag_chroma.py szukaj "What's the maximum security deposit for a rental apartment?"
ochrona-praw-lokatorow-DU-2023-725.pdf art. 19f (similarity 0.673): security deposit covering claims from institutional rental of premises due to the owner...
ochrona-praw-lokatorow-DU-2023-725.pdf art. 6 (similarity 0.639): Art. 6. 1. Conclusion of a rental agreement may be conditioned on the tenant paying a security deposit covering...
ochrona-praw-lokatorow-DU-2023-725.pdf art. 19k (similarity 0.615): tenant's statement in which the tenant submitted to enforcement and undertook to vacate and surrender the pr...

> python 04b_rag_chroma.py szukaj "Can my mom be with me in the examination room during the exam?" prawa-pacjenta-DU-2024-581.pdf
prawa-pacjenta-DU-2024-581.pdf art. 17 (similarity 0.530): Art. 17. 1. A patient, including a minor who has reached 16 years of age, has the right to consent to an exam...
prawa-pacjenta-DU-2024-581.pdf art. 14 (similarity 0.520): - person practicing a medical profession. A person practicing a medical profession may file a motion with the court also in case of...
prawa-pacjenta-DU-2024-581.pdf art. 18 (similarity 0.518): Art. 18. 1. In case of a surgical procedure or application of a treatment or diagnostic method posing risk...

Results and similarities are identical to numpy, because these are the same vectors. A vector database doesn’t improve search quality, only convenience: disk storage (11 MB for us), adding new documents without recalculating everything, and filters. The second query shows a file filter: even limited to the Patient Rights Act, the search didn’t hit art. 21, confirming that the problem is vocabulary, not mixing laws.

What Breaks RAG Most Often

• Bad text splitting. Chunks cut mid-sentence or too large (entire chapter) are the most common cause of poor search.
• Embedding model that doesn’t know Polish well. On our laws, bge-m3 found the correct article in the top five in 8 out of 10 cases, while nomic-embed-text only in 4. The comparison is in the article about embeddings.
• Vocabulary difference between the questioner and the document (“mom” vs. “close person”).
• No fallback in the prompt. Without the sentence “if not there, say so,” the model happily fills in from its own knowledge.
• Outdated documents. RAG answers based on what you give it. The consolidated text of the Tenant Rights Protection Act from 2023 doesn’t include later amendments.

Next Step: RAG as an Agent Tool

You can give the szukaj() function to an agent as a tool. Then the agent decides whether to search, in which law, and can do several searches before answering (this is called agentic RAG). How to connect a function as a tool is shown in the article How to Build an AI Agent in Python, and how to expose it to all AI clients at once in the article Your Own MCP Server.

Frequently Asked Questions

What is RAG in simple words?

A cheat sheet for the model. Before the model answers, the program searches your documents for matching fragments and gives them to it to read. The model answers based on them, not from memory.

RAG or fine-tuning?

For answering questions from documents, almost always RAG: it’s cheaper, you update it by dropping in a new file, and you can see the source of the answer. Fine-tuning changes the model’s style and behavior, but is poor at cramming facts.

Does RAG work in Polish?

Yes, if the embedding model knows Polish. Our test on Polish laws with bge-m3 gave 8 correct answers out of 10 questions with answers in the documents.

Can I do RAG completely locally, without internet?

Yes. We computed embeddings locally anyway, and instead of a model from OpenRouter you can use a model from Ollama (in the code you change the address and model name). With weaker local models, answers will be simpler and more often imprecise.

How much does RAG cost?

For us, embeddings nothing (local), and 11 answers 0.0072 USD, or about 0.00065 USD per question. The biggest cost is input tokens: 5 fragments of 1200 characters each with every question.

What vector database to start with?

For learning, numpy (zero installation). For a small project, Chroma or SQLite with vector extension. For large collections and many users: pgvector in PostgreSQL, Qdrant, or cloud services.

Google · Your sourcesQuantslant higher in your Google - one clickAdd to preferred sources →

Sources and Verification Date

• Legal texts: Official Journal 2026 item 1245, Official Journal 2023 item 725, Official Journal 2024 item 581 (downloaded October 4, 2026).
• Model bge-m3 in Ollama library, Chroma documentation.
• Measurements: our computer (RTX 4090, Windows 11), October 4, 2026, 2:30-2:55 PM. We manually assessed the model’s answers based on the text of the laws. This is not legal advice.

Kinetyra