Trim a Chat to Fit the Context
Problem
A chat app must fit each request into the model's context window of budget tokens. The request always carries the system prompt (system tokens) and must leave reserve tokens free for the reply. turns lists the conversation's token counts oldest first; it alternates user and assistant and ends with the new user message, so its length is odd. Keep the newest user message, then keep older (user, assistant) pairs from newest to oldest while they fit, never dropping half of a pair and never leaving a gap in the history. Return how many of the oldest turns are dropped, or None if even the system prompt, the new message and the reserve do not fit.
Examples
Input: budget = 100, system = 20, reserve = 25, turns = [10, 30, 10, 20, 15]
Output: 2
Why: 55 tokens are free: 15 for the new message and 30 for the newest pair, and the oldest pair of 40 does not fit in the 10 leftInput: budget = 200, system = 20, reserve = 25, turns = [10, 30, 10, 20, 15]
Output: 0
Why: the whole conversation fitsInput: budget = 50, system = 20, reserve = 10, turns = [10, 30, 40]
Output: None
Why: edge case, the new message alone is larger than the 20 tokens leftHints
0 / 3
Take the fixed costs off the budget first: the system prompt and the reserved reply space are paid no matter what.
Recent context matters most, so fill the remaining space from the newest end. Because pairs may not be split and gaps are not allowed, you stop at the first pair that does not fit.
Subtract system and reserve from the budget, then the last turn; return None if that goes negative. Walk backwards over the turns two at a time, subtracting each (user, assistant) pair while it fits, and count what you kept. The answer is the number of turns minus that count.
Solution
The system prompt and the reply reserve are fixed costs, so they come off the budget first, and the new user message is next, since a request without it is pointless; if that already overflows there is no valid request and the answer is None. Older history is kept from the newest end, a pair at a time, because the most recent exchanges are what the model most needs and an answer without its question confuses it. The walk stops at the first pair that does not fit rather than skipping to a smaller older one, since a conversation with a hole in the middle reads as if the user had said something different. Time is O(t) for t turns and space is O(1).
def fit_context(budget, system, turns, reserve):
left = budget - system - reserve - turns[-1] # the newest user turn always stays
if left < 0:
return None # not even the new message fits
kept = 1
for i in range(len(turns) - 2, 0, -2): # older (user, assistant) pairs, newest first
pair = turns[i - 1] + turns[i]
if pair > left:
break # stop at the first misfit: no holes in history
left -= pair
kept += 2
return len(turns) - kept # how many of the oldest turns are dropped
print(fit_context(100, 20, [10, 30, 10, 20, 15], 25)) # -> 2
print(fit_context(200, 20, [10, 30, 10, 20, 15], 25)) # -> 0
print(fit_context(50, 20, [10, 30, 40], 10)) # -> NoneStuck on the idea rather than the code? Context Windows covers it.