// 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:

  1. 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.
  2. 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.
  3. 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: cosine tells 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 - distance converts 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 pts

Implement 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 k results; fewer if there aren't k entries.
  • If k <= 0, return [].

Log in to submit a solution and earn points.