// lesson: retrieval
Retrieval and the Vector Database
Last lesson's brute-force loop is retrieval โ score every chunk, take the best. So why does every RAG stack ship a vector database? Three reasons:
- Persistence. Embedding your corpus costs time (or API dollars). A DB stores the vectors so you index once and query forever; the brute-force script re-embedded everything on every run.
- Scale. Scoring all vectors is fine at 1,000 chunks and hopeless at 100 million. Vector DBs build approximate-nearest-neighbor indexes (HNSW graphs, IVF cells) that find the near-best matches while touching a tiny fraction of the data.
- Filters and metadata. "Top 5 chunks, but only from docs tagged
runbook, modified this year" โ databases are good at that part.
We'll use Chroma because it's an embedded library โ no server to run,
the index is just a folder. FAISS (a raw index library, closer to the metal)
and Qdrant (a client-server DB you talk to over HTTP) occupy the same seat
in the pipeline; everything below maps one-to-one. Add to rag.py:
import chromadb
_db = chromadb.PersistentClient(path="./index")
_collection = _db.get_or_create_collection(
"docs", metadata={"hnsw:space": "cosine"}
)
def index_corpus(corpus, vectors):
_collection.add(
ids=[chunk_id for chunk_id, _ in corpus],
documents=[chunk for _, chunk in corpus],
embeddings=vectors,
)
def retrieve(question, k=4):
qvec = embed_texts([question])[0]
res = _collection.query(query_embeddings=[qvec], n_results=k)
return list(zip(res["ids"][0], res["documents"][0], res["distances"][0]))
Two details worth understanding rather than cargo-culting:
hnsw:space: cosinetells Chroma to rank by cosine; its default is squared L2 distance. Rule one of vector search: know which metric your index uses.- Chroma returns cosine distance (
1 - similarity, lower is better), not similarity.score = 1.0 - distanceconverts back.
Index your corpus and query it:
if __name__ == "__main__":
if _collection.count() == 0:
corpus = build_corpus(load_documents("docs"))
index_corpus(corpus, embed_texts([chunk for _, chunk in corpus]))
print(f"indexed {_collection.count()} chunks")
for chunk_id, chunk, dist in retrieve("how do I restore a backup"):
print(f"{1.0 - dist:.3f} {chunk_id} {chunk[:70]!r}")
First run indexes; every run after that starts instantly from the ./index
folder โ persistence, reason one, felt firsthand. The scores should match
your brute-force numbers from last lesson, because .query() is not
magic: score the query vector against stored vectors, return the k best,
with an index structure so it doesn't have to touch all of them. Strip away
the acceleration and what remains is one small pure function โ score
everything, sort descending, slice. That's your challenge, and when you can
write it, n_results=4 stops being an incantation.
Two behaviors your version should pin down (the real ones do too): asking for more results than exist returns what exists, and ties are stable โ equal scores keep their original order instead of shuffling between runs. Determinism makes retrieval debuggable.
โบ Top K
15 ptsImplement top_k(query_vec, entries, k), where entries is a list of
(id, vector) tuples. A working cosine_similarity is provided in the
starter โ this challenge is about the selection logic:
- Score every entry:
cosine_similarity(query_vec, vector). - Return a list of
(id, score)tuples sorted by score, highest first. - Entries with equal scores keep their original input order (stable sort).
- Return at most
kresults; fewer if there aren'tkentries. - If
k <= 0, return[].
Log in to submit a solution and earn points.