AI Engineer Portal

Your personal operating system for career transition.

Private mode

Exercise

Build a context window assembler for RAG prompts

Build a Context Window Assembler for RAG Prompts

One of the most underestimated problems in production RAG is fitting retrieved chunks into the LLM context window without silently truncating relevant content or wasting tokens on duplicated passages. A naive implementation concatenates all chunks up to a token limit — but this ignores relevance order, fails to handle near-duplicate passages, and leaves no room for the system prompt and expected output.

What you are building

Create a ContextAssembler class that:

Accepts a list of retrieved chunks with relevance scores, text, and metadata (source, section).
Enforces a token budget — a configurable max_context_tokens that reserves space for the system prompt and expected output. Use the approximation: tokens ≈ len(text.split()) * 1.3.
Prioritizes by relevance score — highest-scoring chunks appear first in the assembled context.
Deduplicates near-identical chunks — if two chunks share more than 80% of their words (Jaccard similarity), drop the lower-scored one.
Formats each chunk with a citation header: [Source: {source}, Section: {section}] followed by the chunk text, separated by ---.
Returns an AssembledContext dataclass with: the formatted context string, included chunks, dropped chunks (with reason: "truncated" or "duplicate"), and the total token estimate.

Why this matters

Context assembly is the last step before the LLM sees your data. Getting it wrong means the model misses the most relevant content (truncation) or wastes tokens on repeated information (duplication). A proper assembler makes the context window a managed resource, not an overflow buffer.

Constraints

Standard library only. No external tokenizer.
Type-hint everything. Deterministic: same input always produces the same output.

Retrieval / medium / Step 9 of 15

Practice stage

Retrieval quality and ranking

Hint

Do not trust a single score blindly. Combine ranking logic with metadata and think about why a candidate deserves to rise.

Next drill

Implement RAG answer generation with inline citations

Success criteria

- Ranking rule is explainable
- Supports future iteration without rewriting everything
- Feels tied to product trust rather than pure math trivia

Review checklist

- Did I account for both relevance and metadata?
- Would I be able to explain the ranking rule in an interview?
- Does the code make future tuning easy?

from __future__ import annotations

from dataclasses import dataclass, field

@dataclass
class RetrievedChunk:
    id: str
    text: str
    score: float  # higher is better
    metadata: dict = field(default_factory=dict)

@dataclass
class AssembledContext:
    context_str: str
    included_chunks: list[RetrievedChunk]
    dropped_chunks: list[dict]  # {"chunk": RetrievedChunk, "reason": str}
    token_estimate: int

def token_estimate(text: str) -> int:
    """Approximate token count: word count * 1.3."""
    # TODO: implement
    raise NotImplementedError

def jaccard_similarity(a: str, b: str) -> float:
    """Word-level Jaccard similarity between two strings."""
    # TODO: implement intersection / union of word sets
    raise NotImplementedError

class ContextAssembler:
    def __init__(
        self,
        max_context_tokens: int = 2000,
        dedup_threshold: float = 0.8,
        system_prompt_tokens: int = 300,
        output_reserve_tokens: int = 500,
    ):
        self.max_context_tokens = max_context_tokens
        self.dedup_threshold = dedup_threshold
        self.effective_budget = max_context_tokens - system_prompt_tokens - output_reserve_tokens

def assemble(self, chunks: list[RetrievedChunk]) -> AssembledContext:
        """
        Sort by score, deduplicate near-identical chunks, fill budget greedily,
        format with citation headers separated by ---.
        """
        # TODO: implement
        raise NotImplementedError

Practice

Generate a variation

Generate a new exercise variation to deepen understanding or practice a related concept.

Attempt history

Recent submissions

Before you submit, decide what a strong answer should make obvious to the reviewer.

No attempts yet.