Assemble Context Within a Token Budget

Retrieval gives you a ranked list of candidate chunks, but the model prompt has a hard token ceiling. You cannot send everything, so you have to choose a subset that fits. This exercise is the packing step that sits between the retriever and the prompt: take the best chunks you can afford, in order, and report how much of the budget you spent.

How the packing works

The chunks arrive already sorted by relevance, best first. You walk them in that order, keeping a running total of the token_count you have committed to. For each chunk you ask a single question: would adding it push the running total past max_tokens?

• If it still fits (used + token_count <= max_tokens), select it and add its tokens to the running total.
• If it does not fit, skip it and keep going. A later chunk may be small enough to slot into the remaining space.

You never truncate a chunk to make it fit and you never reorder — a half-passage is worse than no passage, and rank order is the relevance signal. The separator that glues chunks together is presentation only and does not count against the budget.

Task

Complete assemble_context(chunks, max_tokens, separator="\n\n"):

1. Iterate chunks in the given rank order, tracking used_tokens (start at 0) and a list of selected texts.
2. For each chunk, select it only if used_tokens + chunk["token_count"] <= max_tokens; then add its token_count to used_tokens.
3. If a chunk would overflow, skip it and continue to the next one (the skip-and-continue policy).
4. Return (context, used_tokens) where context is the selected texts joined by separator in rank order.

Example

chunks = [
    {"text": "alpha", "token_count": 40},
    {"text": "bravo", "token_count": 30},
    {"text": "charlie", "token_count": 50},   # would overflow, skipped
    {"text": "delta", "token_count": 20},
]

assemble_context(chunks, max_tokens=100)
# -> ("alpha\n\nbravo\n\ndelta", 90)
# alpha(40) + bravo(70) fit; charlie(120) overflows and is skipped;
# delta(90) fits in the leftover room.

assemble_context([], max_tokens=100)
# -> ("", 0)

Constraints

• Preserve rank order in the output; do not sort by size or reorder to pack tighter.
• Never truncate a chunk’s text — a chunk is taken whole or not at all.
• The separator does not count against max_tokens; only token_count values do.
• Handle an empty chunk list and a budget too small for any chunk by returning ("", 0).
• A single chunk larger than max_tokens is simply never selected.

Approach

Walk the ranked chunks once, keeping a running token total. A chunk is selected only when its token_count still fits under max_tokens; otherwise it is skipped and we keep testing later chunks, which may be smaller and still fit. We chose skip-and-continue over stop-at-first-overflow because a large low-priority chunk should not block a small high-priority one waiting behind it. The separator is treated as presentation glue and never charged against the budget.

Solution

from typing import List, Dict, Tuple

def assemble_context(
    chunks: List[Dict],
    max_tokens: int,
    separator: str = "\n\n",
) -> Tuple[str, int]:
    selected: List[str] = []
    used_tokens = 0

    for chunk in chunks:                     # rank order, best first
        cost = chunk["token_count"]
        if used_tokens + cost <= max_tokens:  # fits under the ceiling
            selected.append(chunk["text"])
            used_tokens += cost
        # else: overflow -> skip this chunk, keep trying later smaller ones

    context = separator.join(selected)
    return (context, used_tokens)

Walkthrough

On the example with max_tokens=100 and the four chunks:

1. used_tokens = 0. alpha costs 40; 0 + 40 = 40 <= 100, select it. used_tokens = 40.
2. bravo costs 30; 40 + 30 = 70 <= 100, select it. used_tokens = 70.
3. charlie costs 50; 70 + 50 = 120 > 100, overflow — skip it and move on. used_tokens stays 70.
4. delta costs 20; 70 + 20 = 90 <= 100, select it. used_tokens = 90.

We join ["alpha", "bravo", "delta"] with "\n\n" and return ("alpha\n\nbravo\n\ndelta", 90). Note that delta was reachable only because we skipped past charlie instead of stopping — the design choice that separates this from a plain prefix cut.

On the empty-list call the loop body never runs, so we return ("", 0).

Complexity & notes

• Time is O(n) over the chunks for the scan, plus O(total selected characters) for the final join. Space is O(k) for the k selected texts.
• Skip-and-continue is a policy choice, not the only correct one. Stopping at the first overflow is simpler and preserves a strict relevance prefix, but it wastes budget when a big chunk sits above small ones. Say which you picked and why; both are defensible.
• This is a greedy fill, not the optimal knapsack packing. Greedy respects rank order, which is usually what you want for RAG because relevance, not density, drives quality. True 0/1 knapsack would maximize tokens used but could drop a top-ranked chunk to fit two lower ones — rarely the goal.
• The separator is deliberately excluded from the budget here to keep the accounting clean. In a real prompt you would reserve headroom for separators, system instructions, and the expected completion, then pass the remaining budget in as max_tokens.
• token_count is precomputed and passed in so the function stays deterministic and framework-free. In production it would come from the model’s tokenizer; keeping it as data makes the packing logic unit-testable by hand.