Embed & Retrieve Top-k

Roadmaps › AI Engineering

Scenario

Your support bot is asked: 'How do I reset my password?' Your docs have 4,000 paragraphs. You can't paste all of them into the prompt — that's a million tokens. You need to find the 3 or 4 paragraphs that actually answer the question, and only send those.

How does the system know WHICH paragraphs are relevant, without you writing a single keyword rule?

Mental model

Retrieval is 'find the nearest neighbours' — turn the question into a point in space, and grab the chunks that sit closest to it.

Every chunk of your docs was already turned into a vector (a list of ~1,500 numbers) at ingest time. At query time you turn the user's question into a vector the SAME way, then ask: which stored vectors point in nearly the same direction? The closest k of them are your answer. No keyword matching — meaning is encoded in the direction of the vector, so 'reset my password' lands near 'recover your account' even though they share no words.

Deep dive

Picture every paragraph in your docs as a dot floating in space, where dots that MEAN similar things sit close together — all your password paragraphs in one corner, billing in another, API docs in a third. Retrieval is just this: drop the user's question in as a new dot, and grab the handful of dots nearest to it.

Retrieval = find the chunks nearest your question in embedding space.

Why not just match keywords? Language hides meaning behind different words. 'How do I reset my password?' and 'Steps to recover your account' share almost no words, so keyword search misses the second. But they mean nearly the same thing — so their dots land right next to each other. Semantic closeness, not word overlap, is what we search on.

Embeddings capture meaning, so retrieval finds matches keyword search can't.

The dots come from an embedding model: feed it text, get back a vector positioned so similar meanings sit close. Two phases: at INGEST you chunk the docs, embed each chunk, and store every (vector, text) pair. At QUERY time you embed the question — with the EXACT SAME model — to get one query vector. If you mix models with the SAME dimensionality (like ada-002 and 3-small), retrieval silently returns garbage. If dimensions differ, it hard-crashes.

Same embedding model on both sides, or the vectors aren't comparable.

Now: one query vector, thousands of chunk vectors — which are closest? Not by straight-line distance but by DIRECTION. Two vectors pointing the same way mean the same thing; cosine similarity measures exactly that (1 = identical, 0 = unrelated). You don't need the formula — just hold the picture. Watch the query land nearest its cluster and the top-k light up:

Closeness = direction (cosine); the nearest k chunks become the context.

That k is usually 4 to 6 — a real tradeoff. Too small and one slightly-off chunk leaves the model nothing, so it hallucinates. Too large and the real answer gets buried among marginal chunks (the 'lost in the middle' effect) while you pay for extra tokens. Scanning every vector is fine for hundreds of chunks; at millions it's too slow — which is why vector databases keep an ANN index that finds the nearest k in milliseconds.

Tune k to 4–6; a vector DB's ANN index makes nearest-k fast at scale.

The catch: the vector store never understands English — it only compares numbers. ALL the meaning lives in the embedding model; the store is just fast geometry. And it ALWAYS returns k results, even when nothing is relevant — the least-bad matches, not the correct ones. Nearest ≠ good, which is why reranking and thresholds come next.

Pipeline: Docs → Chunk → Embed → Store, then Query → Embed → Retrieve → Prompt → LLM. This is the Embed → Retrieve hinge.

Code examples

Embed the query with the same model used for chunks

from openai import OpenAI
client = OpenAI()

def embed(text: str) -> list[float]:
    resp = client.embeddings.create(
        model="text-embedding-3-small",   # MUST match the model used at ingest
        input=text,
    )
    return resp.data[0].embedding   # e.g. 1536 floats

query_vec = embed("How do I reset my password?")

The question goes through the identical embedding model the chunks used. Mixing models (chunks on model A, query on model B) produces vectors in different spaces — the similarity numbers become meaningless.

Score every chunk by cosine similarity, keep the top k

import numpy as np

def cosine(a, b):
    a, b = np.array(a), np.array(b)
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

# top-k by brute force over stored (vector, text) pairs
def retrieve(query_vec, store, k=4):
    scored = [(cosine(query_vec, vec), text) for vec, text in store]
    scored.sort(reverse=True)          # highest similarity first
    return [text for _, text in scored[:k]]

This brute-force scan is exactly what a vector DB does for you — just optimised with an ANN index so it doesn't compare against all N vectors. Conceptually it is this loop: score every chunk, sort, take the top k.

In production: let the vector store do the search

# pgvector / Postgres — the <=> operator is cosine distance
rows = db.execute(
    """
    SELECT chunk_text
    FROM documents
    ORDER BY embedding <=> :qvec   -- nearest first
    LIMIT :k
    """,
    {"qvec": query_vec, "k": 4},
).fetchall()
top_chunks = [r.chunk_text for r in rows]

You almost never hand-roll the scan. The vector store (pgvector here) keeps an index and returns the nearest k in milliseconds even over millions of chunks. `<=>` is distance, so ORDER BY ascending = most similar first.

Key points

Common mistakes

Glossary

vector
A list of numbers that represents the meaning of a piece of text, allowing computers to understand what the text is actually about.
cosine similarity
A mathematical way to measure how closely related two pieces of text are by comparing the direction of their vectors.
top-k
The practice of asking the database to return only the single best, or k most relevant, pieces of information to keep the AI focused.

Recall questions

Questions & answers

In a RAG pipeline, a user reports the bot gives confident but wrong answers when asked about topics NOT covered in the docs. What is happening at the retrieval layer and how do you fix it?

Similarity search always returns the k nearest chunks, so even off-topic questions get 'least-bad' chunks injected, which the LLM then answers from. Fix: apply a minimum similarity threshold — if even the top match is below it, retrieve nothing and let the model refuse — and/or add reranking to filter weak matches.

Approach: Name the root cause: 'nearest ≠ relevant.' The store can't return empty by default. The fix is a threshold/guardrail on the retrieval score, not a bigger k.

Your retrieval quality is poor: relevant chunks exist but aren't being returned. The chunks were embedded with text-embedding-ada-002; the query is embedded with text-embedding-3-small. What's the bug?

Model mismatch. Chunks and query are in different embedding spaces, so cosine similarity between them is meaningless and the true matches don't surface. Re-embed the query (or the chunks) so both sides use the same model.

Approach: Spot the two different model names. Same-model-both-sides is the first invariant to check when retrieval 'silently' fails.

Continue learning

Previous: Embeddings

Previous: Chunking Strategies

Next: Reranking

Next: Context Injection & Prompt Assembly

Return to AI Engineering Roadmap