A summariser concatenated whatever the user attached after a 4k system prompt and 3k of tool definitions, until a user attached forty documents and the request was rejected at the API. The team now describes every request as a list of sections — a name, a size in tokens, a priority — and wants one function to decide which sections go in.
Priority 0 is mandatory: the system prompt, the tool definitions, the user's question. Priority 1 and up are optional, 1 the most wanted. The request has a budget — the token count the engineer allows this feature, below the window, as tokens and the context window sets it — and a reserve held back for the reply, which counts against the same budget. Sections keep their prompt order when sent; the planner chooses, it never reorders and never sends part of a section. Sizes are counted with the tokenizer, not estimated, so the function can trust them.
def solve(sections, budget, reserve):
"""sections: [(name, tokens, priority), ...] in prompt order; priority 0
is mandatory, 1 is the most wanted optional section. budget is the tokens
this request may use in total; reserve is held back for the output.
Returns the names to send, in prompt order, or None when the mandatory
sections alone do not fit in budget - reserve."""
In the first example 64,000 − 6,000 = 58,000 is available; the mandatory 7,000 leaves 51,000; retrieved at priority 1 fits and leaves 21,000; history fits and leaves 9,000; examples at 10,000 does not. In the second, 5,500 − 1,000 = 4,500 is less than the 5,000 the mandatory section needs, so the planner refuses.
A strong answer · Rubric · Where answers go wrong · Solution
Subtract the reserve first: the reply is paid for before any document is. Walk the sections once and take every priority-0 section, subtracting its size; if the remaining budget is negative, return None — the caller must shorten the system prompt or compact the history, and a planner that silently dropped a mandatory section would hide exactly the failure it exists to surface.
Then sort the optional sections by priority, with the original index as the tie-breaker so two sections of equal priority keep prompt order. Walk them in that order and take each one whose size fits in what remains; a section that does not fit is skipped, not truncated, and the walk continues, because a smaller section further down the list may still fit. Collect the indices taken and return the names sorted by index — prompt order, never priority order, since the planner chooses and the prompt layout decides placement.
Two things the function never does: it never includes part of a section, because half a document is worse than no document and leaves no record of the cut; and it never prefers a smaller low-priority section over a larger high-priority one that fits, because priority is the caller's statement of what matters. This is first-fit by priority, not the optimal packing a knapsack solver would find, and that is deliberate — the order of wants is the contract, not the token count. One sort, two passes, O(n log n).
"Sort by size and pack the smallest first." The budget fills with the cheapest sections regardless of what the caller said mattered; a priority-1 document of 30k is dropped in favour of three priority-3 examples. Priority is the contract.
"Truncate the section that does not fit to the remaining space." The model then reads the first 60% of a document and summarises it as if whole, and nothing in the trace says a cut was made — the production failure this exercise started from, in a quieter form.
"Pack against the budget and forget the reserve." Every section fits and the reply hits the output cap; reading the response is the lesson on what a truncated reply does to the parser downstream.
Python 3 · standard library32 lines · 7 of 7 tests passing
"""Token budget planner — pack prioritised sections into one request budget."""
def solve(sections, budget, reserve):
"""sections: [(name, tokens, priority), ...] in prompt order; priority 0
is mandatory, 1 is the most wanted optional section. budget is the tokens
this request may use in total; reserve is held back for the output.
Returns the names to send, in prompt order, or None when the mandatory
sections alone do not fit in budget - reserve."""
available = budget - reserve # the reply is paid for before any document
chosen = [] # indices into sections, in any order for now
for i, (_, tokens, priority) in enumerate(sections):
if priority == 0:
chosen.append(i)
available -= tokens
if available < 0:
return None # refuse loudly; never drop a mandatory section
# First fit by priority; the index as tie-breaker keeps prompt order among equals.
optional = sorted((priority, i) for i, (_, _, priority) in enumerate(sections) if priority > 0)
for _, i in optional:
tokens = sections[i][1]
if tokens <= available: # invariant: available >= 0 after every step
chosen.append(i)
available -= tokens # a section that does not fit is skipped, never cut
return [sections[i][0] for i in sorted(chosen)] # prompt order, never priority order
SECTIONS = [("system", 4000, 0), ("tools", 3000, 0), ("history", 12000, 2),
("retrieved", 30000, 1), ("examples", 10000, 3)]
The cases that ran
TESTS = [
# the worked example: examples (priority 3, 10k) is the one section left out
((SECTIONS, 64000, 6000), ["system", "tools", "history", "retrieved"]),
# mandatory alone exceeds budget - reserve: refuse
(([("sys", 5000, 0), ("doc", 100, 1)], 5500, 1000), None),
# a priority-1 section that does not fit is skipped; smaller priority-2 ones still go in
(([("sys", 1000, 0), ("doc_a", 9000, 1), ("doc_b", 3000, 2), ("doc_c", 2000, 2)], 8000, 1000),
["sys", "doc_b", "doc_c"]),
# an exact fit is a fit
(([("sys", 2000, 0), ("doc", 3000, 1)], 6000, 1000), ["sys", "doc"]),
# equal priority: the earlier section wins the space
(([("sys", 500, 0), ("x", 2000, 1), ("y", 2000, 1)], 3000, 500), ["sys", "x"]),
# nothing to plan
(([], 1000, 100), []),
# the reserve consumes the whole budget: even a tiny mandatory section cannot fit
(([("sys", 10, 0)], 100, 100), None),
]
