The Retrieval Pipeline
22 min
This module wires together two things you've already built: Embeddings' cosine similarity (rank candidates by relevance) and LLM Fundamentals' context-window awareness (respect a hard token limit). A real RAG pipeline does exactly this — embed the query, score every candidate chunk by similarity, then greedily pack the most relevant chunks into the prompt until the token budget available for retrieved context runs out.
def cosine_similarity(a, b):
dot = sum(x * y for x, y in zip(a, b))
mag_a = sum(x ** 2 for x in a) ** 0.5
mag_b = sum(x ** 2 for x in b) ** 0.5
return dot / (mag_a * mag_b)
def retrieve_within_budget(query, candidates, token_budget):
scored = sorted(candidates, key=lambda c: -cosine_similarity(query, c[1]))
selected = []
used = 0
for chunk_id, vector, token_count in scored:
if used + token_count <= token_budget:
selected.append(chunk_id)
used += token_count
return selected
This is a GREEDY selection, not an optimal one -- it always takes the highest-similarity chunk that still fits, skipping over ones that don't, rather than solving for the absolute best-scoring COMBINATION that fits the budget (which would be a proper knapsack problem). Real RAG systems use this same greedy approach because it's simple, fast, and 'most relevant first' is usually good enough in practice.
Using the provided `cosine_similarity`, write `retrieve_within_budget(query, candidates, token_budget)`: `candidates` is a list of `(chunk_id, vector, token_count)`. Rank candidates by similarity to `query` (highest first), then greedily select chunk IDs in that order as long as adding one doesn't exceed `token_budget`. Return the selected chunk IDs in selection order.
Why does a RAG pipeline need a token BUDGET when selecting retrieved chunks, not just 'take the top-k most similar'?
You can build a budget-aware retrieval step that combines similarity ranking with a hard token limit, the core of a real RAG pipeline's context-selection stage.