A tiny vector store
A vector store is the database behind retrieval. You index documents by their embeddings up front, then ask it for the top-k most similar to a query. Production uses Pinecone, pgvector, or similar; the concept fits in an in-memory dict.
Three operations:
- add(doc_id, text) — embed the text once and store it under an id, as
{'text': ..., 'vector': ...}. You embed at index time so retrieval is just comparisons, not re-embedding. - retrieve(query, k) — embed the query, score every stored doc with cosine similarity, sort, and return the best
k. - Return ids, not text — the caller resolves them, exactly like a real store hands back row keys.
Two details make it production-grade. Sort by similarity descending so the closest match is first. And break ties deterministically — when two docs score equally, fall back to ascending doc_id so the same query always returns the same order. Non-deterministic retrieval is a nightmare to test and debug; a stable tie-break is cheap insurance.
Syntax
self.docs = {} # doc_id -> {"text", "vector"}
self.docs[doc_id] = {"text": text, "vector": embed(text)}
# retrieve: score, then sort by (-similarity, doc_id) for stable ties
scored.sort(key=lambda pair: (-pair[1], pair[0]))
return [doc_id for doc_id, _ in scored[:k]]Worked examples
The store: add and retrieve top-k
import math
VOCAB = ["cat", "dog", "kitten", "puppy", "fish", "pet", "vet", "food"]
def embed(text):
return [float(text.lower().split().count(t)) for t in VOCAB]
def cosine_similarity(a, b):
dot = sum(x * y for x, y in zip(a, b))
ma = math.sqrt(sum(x * x for x in a))
mb = math.sqrt(sum(x * x for x in b))
return 0.0 if ma == 0 or mb == 0 else dot / (ma * mb)
class VectorStore:
def __init__(self):
self.docs = {}
def add(self, doc_id, text):
self.docs[doc_id] = {"text": text, "vector": embed(text)}
def retrieve(self, query, k=2):
q = embed(query)
scored = [(doc_id, cosine_similarity(q, rec["vector"]))
for doc_id, rec in self.docs.items()]
scored.sort(key=lambda pair: (-pair[1], pair[0]))
return [doc_id for doc_id, _ in scored[:k]]
store = VectorStore()
store.add("d1", "cat kitten pet")
store.add("d2", "dog puppy pet")
store.add("d3", "fish food vet")
print(store.retrieve("cat pet", k=2))
# Output:
# ['d1', 'd2']Deterministic tie-break by ascending doc id
# (VectorStore as above)
store = VectorStore()
store.add("b", "cat")
store.add("a", "cat") # identical text -> identical score to "b"
store.add("z", "fish")
print(store.retrieve("cat", k=2))
# Output:
# ['a', 'b'] # tie broken by id, so the order is stableWhat add() actually stores
# (VectorStore as above)
store = VectorStore()
store.add("d1", "cat pet")
print(store.docs["d1"]["text"]) # cat pet
print(store.docs["d1"]["vector"]) # [1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0]Common mistakes & gotchas
Embed once at index time, not on every search
Embedding is the expensive step (a model call in production). Doing it in add() and storing the vector means retrieve() only does cheap comparisons. Re-embedding stored docs on every query works but is wasteful — and in a real system it would mean an API call per document per search.
Ties without a tie-break are non-deterministic
Two docs with equal similarity can come back in either order depending on dict insertion, which makes tests flaky and bugs un-reproducible. Sort by (-score, doc_id) so equal scores fall back to a stable secondary key. Cheap to add, painful to debug when missing.
Returning text instead of ids couples you to the store
Handing back ids (not the document text) keeps the store a pure index — the caller looks up whatever it needs by id. Real vector DBs return keys/row-ids for exactly this reason; returning full text bloats results and ties retrieval to one representation of the document.
Why it matters
Every retrieval feature — search, recommendations, RAG — needs a place to keep vectors and a fast "give me the closest k" query. A managed vector DB scales this to billions of rows, but it's doing precisely what these few lines do: store embeddings, score by similarity, sort, slice. Understanding the in-memory version means the big one holds no mystery.