Problem 125 · LLM Systems

Token-Budget Context Packer with Dedupe

HardAI engineering

01Problem

Assemble a context window. Items have `id`, `text`, `score`, `tokens` (> 0) and `pinned` (system prompt, user question: always included). Write `pack(items, budget, near_dup=0.8)`: raise if pinned items exceed the budget; drop candidates whose normalised text duplicates, or whose word-trigram Jaccard similarity is at least `near_dup` with, a higher-scored kept item or a pinned item; then choose candidates to maximise total score within the remaining budget using greedy-by-density (score per token) and the classic fix of also trying 'best single item first, then greedy fill', keeping whichever set scores higher. Return pinned items, then chosen items in their original order.

02Key insight

Context assembly is a 0/1 knapsack: exact DP is pseudo-polynomial in the budget, too slow per request at 100k tokens, so you use greedy by value density. Pure density greedy has an unbounded worst case (many tiny low-value items crowd out one big valuable one), and comparing against the best single item restores a 1/2 approximation guarantee for one extra pass. Dedupe before packing, because near-duplicate chunks (the same paragraph in two document versions) burn budget and over-weight one source. Preserving original order keeps the retrieval or document order the model expects.

03Solution

solution.pypython63 lines

import re
from dataclasses import dataclass

_WORD = re.compile(r"[a-z0-9]+")


@dataclass(frozen=True)
class Item:
    id: str
    text: str
    score: float
    tokens: int
    pinned: bool = False


def _norm(text: str) -> list[str]:
    return _WORD.findall(text.lower())


def _shingles(words: list[str], n: int = 3) -> set[tuple[str, ...]]:
    if len(words) < n:
        return {tuple(words)}
    return {tuple(words[i:i + n]) for i in range(len(words) - n + 1)}


def _jaccard(a: set, b: set) -> float:
    return len(a & b) / len(a | b) if a or b else 1.0


def _greedy(cands: list[Item], room: int, start: list[Item]) -> list[Item]:
    chosen, used = list(start), sum(i.tokens for i in start)
    for it in sorted(cands, key=lambda i: (-i.score / i.tokens, -i.score, i.id)):
        if it not in chosen and used + it.tokens <= room:
            chosen.append(it)
            used += it.tokens
    return chosen


def pack(items: list[Item], budget: int, near_dup: float = 0.8) -> list[Item]:
    if any(i.tokens <= 0 for i in items):
        raise ValueError("tokens must be positive")
    pinned = [i for i in items if i.pinned]
    room = budget - sum(i.tokens for i in pinned)
    if room < 0:
        raise ValueError("pinned items exceed the budget")
    kept: list[Item] = []
    sigs = [(_norm(p.text), _shingles(_norm(p.text))) for p in pinned]
    for it in sorted((i for i in items if not i.pinned), key=lambda i: (-i.score, i.id)):
        words = _norm(it.text)
        sh = _shingles(words)
        if any(words == w or _jaccard(sh, s) >= near_dup for w, s in sigs):
            continue
        kept.append(it)
        sigs.append((words, sh))
    fits = [i for i in kept if i.tokens <= room]
    best_set = _greedy(fits, room, [])
    if fits:
        best_single = max(fits, key=lambda i: (i.score, -i.tokens))
        alt = _greedy(fits, room, [best_single])
        if sum(i.score for i in alt) > sum(i.score for i in best_set):
            best_set = alt
    chosen = {i.id for i in best_set}
    return pinned + [i for i in items if not i.pinned and i.id in chosen]

Try it first.

TimeO(m^2 * L) for pairwise dedupe over m candidates of length L, plus O(m log m) packing

SpaceO(m * L) shingles

tests.py · 8 passing

tests.pypython27 lines

import random
sys_p = Item("sys", "You are a helpful assistant.", 0, 10, pinned=True)
A = Item("A", "the refund window is thirty days from delivery for all orders", 10, 100)
B = Item("B", "shipping is free", 2, 1)
assert [i.id for i in pack([sys_p, A, B], 110)] == ["sys", "A"]
assert [i.id for i in pack([sys_p, A, B], 111)] == ["sys", "A", "B"]
C = Item("C", "warranty lasts one year", 3, 5)
D = Item("D", "returns need a receipt", 3, 5)
E = Item("E", "a long but only mildly useful policy appendix", 4, 10)
assert [i.id for i in pack([C, E, D], 10)] == ["C", "D"]
dup = Item("A2", "The refund window is thirty days from delivery for all orders!", 9, 5)
near = Item("A3", "the refund window is thirty days from delivery for all orders today", 8, 5)
out = [i.id for i in pack([A, dup, near, B], 200)]
assert out == ["A", "B"], out
assert [i.id for i in pack([sys_p, Item("echo", "You are a helpful assistant", 5, 3)], 50)] == ["sys"]
try:
    pack([sys_p], 5)
    assert False
except ValueError:
    pass
rng = random.Random(0)
pool = [Item(f"i{n}", f"chunk {n} " + " ".join(rng.choice("abcdefghij") for _ in range(12)),
             rng.uniform(0, 1), rng.randint(5, 60)) for n in range(40)]
for budget in (0, 50, 200, 500):
    res = pack([sys_p] + pool, budget + 10)
    assert res[0] is sys_p and sum(i.tokens for i in res) <= budget + 10
    assert [i.id for i in res[1:]] == [i.id for i in pool if i in res]

04Staff-level follow-up

Replace pairwise dedupe with MinHash plus LSH once candidates number in the thousands. Then question the objective: retrieval scores are not additive utility, so diminishing returns per source (submodular objectives, MMR) model it better, and a reranker's calibrated scores beat raw cosine. Finally, place the most important evidence at the start and end of the context, since long-context models attend less reliably to the middle.