Retrieval augmented generation (RAG)

You’ll build everything in two scripts. make_docs.py writes three short IT policy files into a policies folder. rag.py chunks them, searches them, and (when you switch it on) asks llama3.2:1b for an answer. Steps 1 and 2 use keyword search in plain Python, so their output is exact.

Steps 3 and 4 use the embedding model, so the scores and words will differ a little. On macOS or Linux, use python3 and nano.

1Create the documents and search by keyword

Run this:

mkdir C:\ai-lab\rag
cd C:\ai-lab\rag
notepad make_docs.py

Paste this whole block into make_docs.py, then save:

import os

DOCS = {
    "passwords.txt": [
        "Passwords must be at least 14 characters long. Use a passphrase of four or more random words, such as river-lamp-cotton-seven. Never reuse a password from a personal account.",
        "Your password expires every 90 days. Windows reminds you 14 days before that. You cannot reuse any of your last 10 passwords.",
        "After five failed sign-in attempts, your account is locked for 30 minutes. If you need access sooner, call the service desk on extension 4400 and confirm your employee ID.",
        "Never share your password, not even with the IT team. The service desk will never ask for it by phone, email or chat.",
    ],
    "vpn.txt": [
        "Staff working outside the office must connect through the company VPN before opening any internal system.",
        "To connect, open the VPN client, enter vpn.example.com as the portal, and sign in with your normal account and the six-digit code from the authenticator app.",
        "If the VPN connects but internal sites do not open, disconnect, restart your laptop and try again. If it still fails, check that the date and time on your laptop are correct, because the six-digit code depends on them.",
        "The guest Wi-Fi in the office does not need the VPN. Its password changes every Monday and is shown on the screen at reception.",
    ],
    "laptops.txt": [
        "Every new employee receives a company laptop on their first day. Laptops are replaced every four years, or sooner if a repair would cost more than half the price of a new one.",
        "Do not install software yourself. Request it in the Software Center app; most requests are approved within one working day.",
        "If your laptop is lost or stolen, report it to the service desk within 24 hours so it can be locked and wiped remotely.",
        "Before you travel abroad with a laptop, tell the service desk at least five working days in advance, so software licences and VPN access can be checked.",
    ],
}

os.makedirs("policies", exist_ok=True)
for name, paragraphs in DOCS.items():
    with open(os.path.join("policies", name), "w", encoding="utf-8") as f:
        f.write("\n\n".join(paragraphs) + "\n")
    print(f"Wrote policies/{name}: {len(paragraphs)} paragraphs")

Run this:

python make_docs.py

You should see this:

Wrote policies/passwords.txt: 4 paragraphs
Wrote policies/vpn.txt: 4 paragraphs
Wrote policies/laptops.txt: 4 paragraphs

Run this:

notepad rag.py

Paste this whole block into rag.py, then save:

import json
import math
import os
import re
import urllib.request

QUESTION = "How long must my password be?"
SEARCH = "keyword"      # "keyword" or "embed"
TOP_K = 3               # how many chunks to retrieve
CHUNK_WORDS = 0         # 0 = one chunk per paragraph; N = cut every N words
ANSWER = False          # True = ask llama3.2:1b to answer from the chunks
STRICT = True           # True = only answer from the chunks
EMBED_MODEL = "nomic-embed-text"
CHAT_MODEL = "llama3.2:1b"
STOP = {"a", "an", "the", "is", "are", "be", "do", "does", "i", "my", "me", "to", "of", "and",
        "or", "in", "on", "for", "it", "what", "how", "who", "can", "should", "must", "after", "with"}


def post(path, body):
    req = urllib.request.Request("http://localhost:11434" + path, data=json.dumps(body).encode(),
                                 headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req) as r:
        return json.load(r)


def load_chunks():
    chunks = []
    for name in sorted(os.listdir("policies")):
        text = open(os.path.join("policies", name), encoding="utf-8").read()
        if CHUNK_WORDS:
            words = text.split()
            parts = [" ".join(words[i:i + CHUNK_WORDS]) for i in range(0, len(words), CHUNK_WORDS)]
        else:
            parts = [p.strip() for p in text.split("\n\n") if p.strip()]
        for n, part in enumerate(parts, 1):
            chunks.append({"source": f"{name}#{n}", "text": part})
    return chunks


def words(text):
    return {w for w in re.findall(r"[a-z0-9]+", text.lower()) if w not in STOP}


def keyword_scores(question, chunks):
    q = words(question)
    return [len(q & words(c["text"])) for c in chunks]


def embed_scores(question, chunks):
    # nomic-embed-text expects these prefixes on documents and on the question
    docs = post("/api/embed", {"model": EMBED_MODEL,
                               "input": ["search_document: " + c["text"] for c in chunks]})["embeddings"]
    q = post("/api/embed", {"model": EMBED_MODEL, "input": ["search_query: " + question]})["embeddings"][0]
    return [sum(a * b for a, b in zip(q, d)) for d in docs]   # vectors are unit length: dot = cosine


chunks = load_chunks()
print(f"Indexed {len(chunks)} chunks from {len(os.listdir('policies'))} files"
      + (f" ({CHUNK_WORDS}-word chunks)" if CHUNK_WORDS else " (one per paragraph)"))
print(f"Question: {QUESTION}")
scores = keyword_scores(QUESTION, chunks) if SEARCH == "keyword" else embed_scores(QUESTION, chunks)
best = sorted(zip(scores, range(len(chunks))), reverse=True)[:TOP_K]
print(f"Top {TOP_K} by {SEARCH} search:")
for score, i in best:
    text = chunks[i]["text"]
    print(f"  {score:6.3f}  {chunks[i]['source']:<15} {text[:62]}{'...' if len(text) > 62 else ''}")

if ANSWER:
    context = "\n\n".join(f"[{chunks[i]['source']}] {chunks[i]['text']}" for _, i in best)
    rules = ("Answer the question using only the context. Cite the source you used, like [passwords.txt#2]. "
             "If the context does not contain the answer, say: I don't know, that is not in the IT policies.")
    system = rules if STRICT else "You are a helpful IT assistant."
    r = post("/api/chat", {"model": CHAT_MODEL, "stream": False, "options": {"temperature": 0},
                           "messages": [{"role": "system", "content": system},
                                        {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {QUESTION}"}]})
    print(f"Prompt tokens: {r['prompt_eval_count']}")
    print("Answer:", r["message"]["content"].strip())

Run this:

python rag.py

You should see this (keyword search is exact):

Indexed 12 chunks from 3 files (one per paragraph)
Question: How long must my password be?
Top 3 by keyword search:
   2.000  passwords.txt#1 Passwords must be at least 14 characters long. Use a passphras...
   1.000  vpn.txt#4       The guest Wi-Fi in the office does not need the VPN. Its passw...
   1.000  passwords.txt#4 Never share your password, not even with the IT team. The serv...

Twelve chunks, one per paragraph. The score is the number of important words the question and the chunk share: “long” and “password” both appear in passwords.txt#1, so it wins.

That’s the right paragraph. Keyword search works when you use the same words as the document.

✅ Checkpoint: 12 chunks, and passwords.txt#1 is the top result with a score of 2.000.

2Ask in your own words

Real users don’t talk like policy documents. Ask about being locked out, without using any of the document’s words:

Change the QUESTION line:

QUESTION = "I keep getting kicked out of my computer after too many wrong tries"

Run this:

python rag.py

You should see this (exact):

Indexed 12 chunks from 3 files (one per paragraph)
Question: I keep getting kicked out of my computer after too many wrong tries
Top 3 by keyword search:
   0.000  vpn.txt#4       The guest Wi-Fi in the office does not need the VPN. Its passw...
   0.000  vpn.txt#3       If the VPN connects but internal sites do not open, disconnect...
   0.000  vpn.txt#2       To connect, open the VPN client, enter vpn.example.com as the ...

Every score is zero. The answer is sitting in passwords.txt#3, but it says “failed sign-in attempts” and “locked”, and the user said “kicked out” and “wrong tries”. Keyword search can’t see that they mean the same thing, so it returns three random VPN paragraphs.

✅ Checkpoint: All three scores are 0.000, and passwords.txt#3 is not in the list.

3Search by meaning with embeddings

Download a small embedding model, then switch the search:

Run this:

ollama pull nomic-embed-text

Change this line:

SEARCH = "embed"

Run this:

python rag.py

You should see something like this (your scores and the model’s words will differ a little):

Indexed 12 chunks from 3 files (one per paragraph)
Question: I keep getting kicked out of my computer after too many wrong tries
Top 3 by embed search:
   0.712  passwords.txt#3 After five failed sign-in attempts, your account is locked for...
   0.583  passwords.txt#4 Never share your password, not even with the IT team. The serv...
   0.541  laptops.txt#3   If your laptop is lost or stolen, report it to the service des...

The lockout paragraph is now first, with the highest similarity. The embedding model learned from huge amounts of text that “kicked out after wrong tries” and “locked after failed sign-in attempts” describe the same thing. Notice the search_document: and search_query: prefixes in the code: nomic-embed-text was trained with them, and leaving them out makes its results worse.

✅ Checkpoint: passwords.txt#3 is the top result.

4Answer with a source

Change this line:

ANSWER = True

Run this:

python rag.py

You should see something like this (your scores and the model’s words will differ a little):

Indexed 12 chunks from 3 files (one per paragraph)
Question: I keep getting kicked out of my computer after too many wrong tries
Top 3 by embed search:
   0.712  passwords.txt#3 After five failed sign-in attempts, your account is locked for...
   0.583  passwords.txt#4 Never share your password, not even with the IT team. The serv...
   0.541  laptops.txt#3   If your laptop is lost or stolen, report it to the service des...
Prompt tokens: 200
Answer: After five failed sign-in attempts your account is locked for 30 minutes. If you need access sooner, call the service desk on extension 4400 and confirm your employee ID [passwords.txt#3].

That’s RAG end to end. The answer has the real numbers from your policy, 30 minutes and extension 4400, which the model could never have guessed, and the citation tells a person exactly where to check. The prompt was about 200 tokens: three paragraphs, not the whole policy library.

✅ Checkpoint: The answer mentions 30 minutes and cites [passwords.txt#3].