Retrieve the k most similar documents
The retrieval step of RAG: rank documents by similarity to the query and keep the top k, with stable tie-breaking.
Write topK(query, docs, k) that returns an array with the id of the k documents most similar to query, most similar first. Each document is { id, vector }.
Use cosine(a, b), which is already defined. When two documents tie on similarity, the one with the alphabetically smaller id comes first. Do not modify docs. If k is larger than the number of documents, return all of them; if it is 0, return [].
DOCS and QUERY hold an example index with 3-dimensional embeddings.
Challenges 0/4
- Returns the 3 most similar in order
- Breaks ties by id (guia-ollama and ollama-api have the same similarity)
- Does not change the order of DOCS
- With k above the index size returns everything; with k = 0, nothing
function topK(query, docs, k) {
// score each doc with cosine(query, doc.vector), sort high to low, ties by id, take k
return docs.slice(0, k).map((d) => d.id);
}
console.log(topK(QUERY, DOCS, 3));Console output appears here (console.log).
Go deeper: the related guide →
This in production, with your data? Let's talk for 15 minutes →