Top-k Cosine Retriever

Build the retrieval half of a RAG pipeline: given a query embedding and a set of pre-computed document embeddings, return the k documents whose embeddings are most similar to the query. This is the “R” in retrieval-augmented generation, and it is a standard warm-up in an agent-coding round.

Cosine similarity

The cosine similarity between two vectors a and b is their dot product divided by the product of their L2 norms:

cos(a, b) = (a . b) / (||a|| * ||b||)

It measures the angle between the vectors and ignores their magnitude, which is why it is the default relevance score for embeddings. If you L2-normalize every vector first, cosine similarity reduces to a plain dot product.

Task

Complete top_k_cosine(query, doc_embeddings, doc_texts, k=3):

1. Compute the cosine similarity between query and each document embedding.
2. Pair each score with its document text.
3. Return the min(k, n) highest-scoring (text, score) tuples, sorted by score descending.

The embeddings are passed in — there is no model call. NumPy is allowed; so is pure Python.

Example

query = [1.0, 0.0]
doc_embeddings = [
    [1.0, 0.0],   # same direction as query  -> cos 1.0
    [0.0, 1.0],   # orthogonal               -> cos 0.0
    [2.0, 0.0],   # same direction, longer   -> cos 1.0
]
doc_texts = ["exact match", "unrelated", "scaled match"]

top_k_cosine(query, doc_embeddings, doc_texts, k=2)
# -> [("exact match", 1.0), ("scaled match", 1.0)]
# (both cosine 1.0; "unrelated" at 0.0 is excluded)

Constraints

• doc_embeddings and doc_texts have the same length n, and every vector has the same dimension d as query.
• Handle k larger than n by returning all n documents.
• A zero vector (norm 0) must not raise; treat its cosine similarity as 0.0.
• Do not call any embedding model or network — the embeddings are provided.

Approach

Stack the document embeddings into an (n, d) matrix and L2-normalize both it and the query. Once everything is unit-length, cosine similarity is a single matrix-vector product, giving one score per document. Pair each score with its text, sort descending, and slice the first k. A small epsilon on the norms keeps zero vectors from dividing by zero.

Solution

from typing import List, Tuple
import numpy as np

def top_k_cosine(
    query: List[float],
    doc_embeddings: List[List[float]],
    doc_texts: List[str],
    k: int = 3,
) -> List[Tuple[str, float]]:
    if not doc_embeddings:
        return []

    q = np.asarray(query, dtype=float)
    M = np.asarray(doc_embeddings, dtype=float)   # shape (n, d)

    # L2-normalize; clip the norms so a zero vector -> zero score, not NaN.
    q_norm = q / max(np.linalg.norm(q), 1e-12)
    row_norms = np.linalg.norm(M, axis=1, keepdims=True)
    M_norm = M / np.clip(row_norms, 1e-12, None)

    scores = M_norm @ q_norm                      # cosine similarity per doc, shape (n,)

    # Rank by score descending; ties broken by original order via stable sort.
    order = np.argsort(-scores, kind="stable")[:k]
    return [(doc_texts[i], float(scores[i])) for i in order]

Walkthrough

On the example, query = [1, 0] and the three docs are [1,0], [0,1], [2,0]:

1. Normalizing: q_norm = [1, 0], and the rows become [1,0], [0,1], [1,0] (the length-2 vector [2,0] normalizes to the unit [1,0]).
2. M_norm @ q_norm gives scores = [1.0, 0.0, 1.0].
3. argsort(-scores) with a stable sort yields index order [0, 2, 1]; the two tied 1.0 scores keep their original relative order (doc 0 before doc 2).
4. Slicing k=2 returns indices [0, 2] -> [("exact match", 1.0), ("scaled match", 1.0)], dropping the orthogonal “unrelated” doc.

Complexity & notes

• Time is O(n·d) for the dot products plus O(n log n) for the sort; space is O(n·d) for the normalized matrix.
• Normalizing once and reusing it is the key trick — cosine over unit vectors is just a dot product, so the whole scoring pass is a single BLAS-backed matmul.
• The 1e-12 clamp is the pitfall interviewers probe: without it a zero-norm (or all-zero) embedding produces NaN, which then sorts unpredictably. Clamping the denominator sends its score cleanly to 0.0.
• For very large corpora you would swap the full argsort for np.argpartition(-scores, k) (O(n) selection) and reserve sorting for just the k survivors; a full sort is fine at interview scale and less error-prone.
• This is retrieval only. In a full RAG loop the returned texts are concatenated into the prompt context passed to the (stubbed) model — scoring correctness here is what determines whether the generator sees relevant evidence.