A Mini RAG Pipeline in 90 Lines of Python: Retrieve, Augment, Generate
Build retrieval-augmented generation from scratch with embeddings, cosine similarity and a ranked list, then see where the toy version breaks at scale.
This article is published in English.
Retrieval-augmented generation isn't a feature of any model. It's a packing problem. The context window is small, the documents are huge, and something has to decide which pieces get in. That something can be written in about 90 lines of plain Python, with no LangChain and no vector database service. It has three steps:
- Retrieve. Embed the question, rank chunks by similarity and keep the top k.
- Augment. Stitch the winners into the system message.
- Generate. Make the same chat call used in any basic LLM script.
Every abstraction a framework like LangChain offers later maps onto one of those three boxes, so seeing the bare version first makes the framework easier to read.
Embeddings turn meaning into geometry
A model that creates embeddings turns a piece of text into a vector, meaning a long run of floating-point numbers, arranged so that texts with related meaning sit near each other. Distance is judged by cosine similarity: 1.0 for identical direction, 0.0 for no relationship. That's the whole trick. Once meaning becomes geometry, finding relevant documents turns into finding nearby vectors.
# pip install openai
# export OPENAI_API_KEY=sk-...
import math
from openai import OpenAI
client = OpenAI()
def embed(texts):
resp = client.embeddings.create(
model="text-embedding-3-small",
input=texts,
)
return [d.embedding for d in resp.data]
def cosine(a, b):
dot = sum(x * y for x, y in zip(a, b))
na = math.sqrt(sum(x * x for x in a))
nb = math.sqrt(sum(y * y for y in b))
return dot / (na * nb)
q, hit, miss = embed([
"How do I get my money back?",
"Refunds are possible within 14 days "
"of purchase.",
"Data is stored in Frankfurt and "
"encrypted at rest.",
])
print(f"refund doc: {cosine(q, hit):.3f}")
print(f"storage doc: {cosine(q, miss):.3f}")
Typical output from that script is a score of 0.531 for the refund document and 0.126 for the storage document. Look at how the match happens: the question says "money back" while the document says "refunds." No keywords are shared, yet the similarity is high. That's why you embed instead of grepping.
A vector store is a list you can rank
"Vector store" sounds like infrastructure. At small scale it's a list of (vector, text) pairs and a sort. Pinecone, Chroma and pgvector earn their keep at millions of vectors, but what matters is the interface, and that fits in one class:
class VectorStore:
def __init__(self):
self.texts = []
self.vectors = []
def add(self, texts):
self.texts.extend(texts)
self.vectors.extend(embed(texts))
def top_k(self, query, k=3):
q = embed([query])[0]
scored = [
(cosine(q, v), t)
for v, t in zip(
self.vectors, self.texts
)
]
scored.sort(
key=lambda s: s[0], reverse=True
)
return scored[:k]
store = VectorStore()
store.add([
"Refunds are possible within 14 days.",
"The Pro plan costs 12 EUR per month.",
"Data is stored in Frankfurt.",
])
for score, text in store.top_k(
"can I cancel and get a refund?", k=2
):
print(f"{score:.3f} {text}")
top_k is the decision about what makes the cut, now executable: rank everything, keep k, drop the rest. The value of k is your budget knob, trading recall against context-window space and billed tokens.
Putting retrieve, augment and generate together
Only one function is new here, build_messages. It numbers the winning chunks, drops them into the system message and tells the model to answer from them alone. The user message carries just the question. The messages list is the same one any chat call uses, so RAG changes only its contents. This file runs as is:
# examples/part07_rag.py
# Mini RAG in ~90 lines: embed -> store -> top-k
# -> stitch into messages -> generate.
#
# Setup:
# pip install openai
# export OPENAI_API_KEY=sk-...
#
# Run:
# python part07_rag.py
import math
from openai import OpenAI
client = OpenAI()
EMBED_MODEL = "text-embedding-3-small"
CHAT_MODEL = "gpt-4o-mini"
# A tiny corpus. In real life these are chunks
# of your docs, not whole documents.
DOCS = [
"Support hours are 9am to 6pm CET, "
"Monday to Friday.",
"Refunds are possible within 14 days "
"of purchase, no questions asked.",
"The Pro plan costs 12 EUR per month "
"and includes API access.",
"The Free plan allows 3 projects and "
"community support only.",
"Data is stored in Frankfurt and "
"encrypted at rest with AES-256.",
"You can export all your data as "
"JSON from the settings page.",
]
def embed(texts):
resp = client.embeddings.create(
model=EMBED_MODEL,
input=texts,
)
return [d.embedding for d in resp.data]
def cosine(a, b):
dot = sum(x * y for x, y in zip(a, b))
na = math.sqrt(sum(x * x for x in a))
nb = math.sqrt(sum(y * y for y in b))
return dot / (na * nb)
class VectorStore:
"""A list of (vector, text) pairs. That is
all a vector store is, before you need
scale."""
def __init__(self):
self.texts = []
self.vectors = []
def add(self, texts):
self.texts.extend(texts)
self.vectors.extend(embed(texts))
def top_k(self, query, k=3):
q = embed([query])[0]
scored = [
(cosine(q, v), t)
for v, t in zip(
self.vectors, self.texts
)
]
scored.sort(key=lambda s: s[0], reverse=True)
return scored[:k]
def build_messages(question, hits):
context = "\n\n".join(
f"[{i + 1}] {text}"
for i, (_, text) in enumerate(hits)
)
system = (
"Answer using only the context below. "
"If the answer is not in the context, "
"say you do not know.\n\n"
"Context:\n" + context
)
return [
{"role": "system", "content": system},
{"role": "user", "content": question},
]
def ask(store, question):
hits = store.top_k(question, k=3)
print(f"Q: {question}")
for score, text in hits:
print(f" {score:.3f} {text[:44]}...")
messages = build_messages(question, hits)
resp = client.chat.completions.create(
model=CHAT_MODEL,
messages=messages,
)
return resp.choices[0].message.content
if __name__ == "__main__":
store = VectorStore()
store.add(DOCS)
answer = ask(
store,
"Where is my data stored, and can "
"I get it out?",
)
print(f"A: {answer}")
A sample run asks where the user's data is stored and whether it can be exported. The top three scores are 0.588 for a chunk about storage in Frankfurt with encryption, 0.472 for a chunk about exporting all data as JSON and 0.244 for the refund policy. The model answers that the data lives in Frankfurt, encrypted at rest with AES-256, and that the settings page offers a full JSON export.
Both needed chunks, storage and export, made it in. The refund chunk landed in third place and the model ignored it. Retrieval can be imperfect as long as the right material ends up above the line.
Where the toy version breaks
These 90 lines have honest limits, and each one points at something real systems add:
- Chunking. The demo stores whole sentences. Real documents need splitting with overlap, and that's a preprocessing decision, not a model one.
- Re-embedding. Embedding the query on every call is fine. Re-embedding the whole corpus on every restart isn't, so persist the vectors.
- Linear scan. Cosine similarity against every vector is O(n). Past roughly 100k chunks an ANN index becomes worthwhile, and that's where vector databases pay for themselves.
- Ranking quality. Plain cosine top-k is a blunt tool, and rerankers or hybrid keyword search make it sharper.
LangChain wraps each of these as a swappable component. Once the raw loop has been written by hand, those components look like what they are, conveniences and nothing mystical.