Syllabus Lesson 153 of 239 · Build a RAG Pipeline
Build a RAG Pipeline

A Tiny Vector Store

Once a document is chunked and each chunk is turned into a vector, you need somewhere to keep those vectors and a way to ask "which stored vectors point most like this query?" That is a vector store. Pinecone, Chroma, FAISS, and pgvector are production versions of this idea, but the core is small enough to build by hand, and building it once means you understand what those services are actually doing.

The retrieval metric is cosine similarity: the cosine of the angle between two vectors, a . b / (|a| |b|). It rewards vectors that point the same direction and ignores their length, which is what you want for text, where a long passage and a short one can be about the same thing. A score near 1 means "very similar direction"; near 0 means "unrelated".

You will build a class so the store has memory between calls. The gate that protects these lessons does not stub class methods, so the safeguard here is that your search must return the correct nearest-neighbour order for several different queries: only real cosine math gets all of them right.

import numpy as np

class VectorStore:
    def __init__(self):
        self.ids = []
        self.vectors = []

    def add(self, id, text, vector):
        self.ids.append(id)
        self.vectors.append(np.asarray(vector, dtype=float))

    def search(self, query_vec, k):
        q = np.asarray(query_vec, dtype=float)
        scored = []
        for vec_id, v in zip(self.ids, self.vectors):
            denom = np.linalg.norm(q) * np.linalg.norm(v)
            sim = float(q @ v / denom) if denom else 0.0
            scored.append((vec_id, sim))
        scored.sort(key=lambda pair: pair[1], reverse=True)
        return scored[:k]

What to build. Implement class VectorStore with:

  • add(id, text, vector) storing the id, the text, and the vector (store the vector as a numpy array of floats).
  • __len__ so len(store) returns how many items are stored.
  • an ids attribute (a list) so callers can see what was added.
  • search(query_vec, k) returning the top-k nearest neighbours as a list of (id, similarity) tuples sorted by similarity descending, length capped at k.

Use numpy.argsort or a plain sort on the scores. Press Run to add a few labelled vectors and see which ones a query pulls back first.

Your turn

Implement class VectorStore with add(id, text, vector) (store id, text, and the vector as a numpy float array), __len__ returning the item count, a public ids list, and search(query_vec, k) returning the top-k (id, similarity) tuples by cosine similarity, highest first, capped at k. A zero vector must score 0.0, not crash.

Spotted a problem in this lesson? Report it

Code · runs in your browser
Output