Trim Context Without Orphaning Tool Resultsmedium

Trim Context Without Orphaning Tool Results

medium Build an Agent 20–25 min

Standard library only · runs in your browser

Long agent runs outgrow the context window. Cut the wrong message and the API rejects your request ('tool result without a matching tool call'), or the model forgets what it was asked. Trim the way production agents do.

Task

Implement trim_messages(messages, max_tokens) and return the messages to keep, in their original order.

Pinned, always kept: the first message if its role is system, plus the first user message right after it (the task). With no system message, pin only messages[0] when it's a user message.
Groups: split the rest into atomic groups. An assistant message plus all the tool messages directly after it form one group. Every other message is its own group.
Walk the groups from newest to oldest, keeping each group whole while its tokens fit the remaining budget (max_tokens minus the pinned tokens). Stop at the first group that doesn't fit, and never skip it to keep an older one.
Tokens are count_tokens(m["content"]), summed.

Function Signature

trim_messages(messages: list[dict], max_tokens: int) -> list[dict]

Examples

Example 1: Everything fits: unchanged
Input: messages = [{'role': 'system', 'content': 'You are a helpful travel agent.'}, {'role': 'user', 'content': 'Compare the weather in Paris and Tokyo for my trip.'}, {'role': 'assistant', 'content': '{"tool": "get_weather", "arguments": {"city": "Paris"}}'}, {'role': 'tool', 'content': '18C, light rain'}, {'role': 'assistant', 'content': '{"tool": "get_weather", "arguments": {"city": "Tokyo"}}'}, {'role': 'tool', 'content': '24C, clear'}, {'role': 'assistant', 'content': 'Paris is rainy and Tokyo is clear.'}, {'role': 'user', 'content': 'And Lima?'}]
max_tokens = 1000
Output: [{'role': 'system', 'content': 'You are a helpful travel agent.'}, {'role': 'user', 'content': 'Compare the weather in Paris and Tokyo for my trip.'}, {'role': 'assistant', 'content': '{"tool": "get_weather", "arguments": {"city": "Paris"}}'}, {'role': 'tool', 'content': '18C, light rain'}, {'role': 'assistant', 'content': '{"tool": "get_weather", "arguments": {"city": "Tokyo"}}'}, {'role': 'tool', 'content': '24C, clear'}, {'role': 'assistant', 'content': 'Paris is rainy and Tokyo is clear.'}, {'role': 'user', 'content': 'And Lima?'}]

Example 2: Oldest call and its result are dropped together
Input: messages = [{'role': 'system', 'content': 'You are a helpful travel agent.'}, {'role': 'user', 'content': 'Compare the weather in Paris and Tokyo for my trip.'}, {'role': 'assistant', 'content': '{"tool": "get_weather", "arguments": {"city": "Paris"}}'}, {'role': 'tool', 'content': '18C, light rain'}, {'role': 'assistant', 'content': '{"tool": "get_weather", "arguments": {"city": "Tokyo"}}'}, {'role': 'tool', 'content': '24C, clear'}, {'role': 'assistant', 'content': 'Paris is rainy and Tokyo is clear.'}, {'role': 'user', 'content': 'And Lima?'}]
max_tokens = 52
Output: [{'role': 'system', 'content': 'You are a helpful travel agent.'}, {'role': 'user', 'content': 'Compare the weather in Paris and Tokyo for my trip.'}, {'role': 'assistant', 'content': '{"tool": "get_weather", "arguments": {"city": "Tokyo"}}'}, {'role': 'tool', 'content': '24C, clear'}, {'role': 'assistant', 'content': 'Paris is rainy and Tokyo is clear.'}, {'role': 'user', 'content': 'And Lima?'}]

Example 3: A group that doesn't fit is dropped whole, never split
Input: messages = [{'role': 'system', 'content': 'You are a helpful travel agent.'}, {'role': 'user', 'content': 'Compare the weather in Paris and Tokyo for my trip.'}, {'role': 'assistant', 'content': '{"tool": "get_weather", "arguments": {"city": "Paris"}}'}, {'role': 'tool', 'content': '18C, light rain'}, {'role': 'assistant', 'content': '{"tool": "get_weather", "arguments": {"city": "Tokyo"}}'}, {'role': 'tool', 'content': '24C, clear'}, {'role': 'assistant', 'content': 'Paris is rainy and Tokyo is clear.'}, {'role': 'user', 'content': 'And Lima?'}]
max_tokens = 40
Output: [{'role': 'system', 'content': 'You are a helpful travel agent.'}, {'role': 'user', 'content': 'Compare the weather in Paris and Tokyo for my trip.'}, {'role': 'assistant', 'content': 'Paris is rainy and Tokyo is clear.'}, {'role': 'user', 'content': 'And Lima?'}]

Example 4: Budget smaller than the pinned messages
Input: messages = [{'role': 'system', 'content': 'You are a helpful travel agent.'}, {'role': 'user', 'content': 'Compare the weather in Paris and Tokyo for my trip.'}, {'role': 'assistant', 'content': '{"tool": "get_weather", "arguments": {"city": "Paris"}}'}, {'role': 'tool', 'content': '18C, light rain'}, {'role': 'assistant', 'content': '{"tool": "get_weather", "arguments": {"city": "Tokyo"}}'}, {'role': 'tool', 'content': '24C, clear'}, {'role': 'assistant', 'content': 'Paris is rainy and Tokyo is clear.'}, {'role': 'user', 'content': 'And Lima?'}]
max_tokens = 5
Output: [{'role': 'system', 'content': 'You are a helpful travel agent.'}, {'role': 'user', 'content': 'Compare the weather in Paris and Tokyo for my trip.'}]

Constraints

Never return a `tool` message without the `assistant` message that requested it.
The kept history is one contiguous suffix of the conversation.
If the pinned messages alone exceed the budget, return just the pinned messages.

Python3Saved ✓

import json

import math

import re

def trim_messages(messages, max_tokens):

"""Keep pinned messages + the newest whole groups that fit max_tokens."""

# Your code here

pass

# =============================================================================

# PROVIDED BELOW -- the sandbox and your earlier stages. Read it, don't edit it.

# =============================================================================

# ---------------------------------------------------------------------------

# A scripted stand-in for the model (provided). It replies with `script`

# in order, but it still READS the conversation: placeholders in a reply

# are filled from the messages it was actually sent, so a loop that forgets

# to feed results back gets a wrong answer -- just like a real model would.

# {last_tool_result} -> content of the most recent "tool" message

# {last_feedback} -> content of the most recent non-assistant message

# {context_size} -> how many messages the model was sent this call

# ---------------------------------------------------------------------------

SYSTEM_PROMPT = (

"You are a helpful agent. To use a tool, reply with JSON: "

'{"tool": <name>, "arguments": {...}}. Otherwise reply with your final answer.'

)

def make_llm(script):

replies = list(script)

def llm(messages):

if not replies:

raise RuntimeError("model script exhausted -- your loop called the model too many times")

reply = replies.pop(0)

Case 1: Everything fits: unchanged

Input: messages = [{'role': 'system', 'content': 'You are a helpful travel agent.'}, {'role': 'user', 'content': 'Compare the weather in Paris and Tokyo for my trip.'}, {'role': 'assistant', 'content': '{"tool": "get_weather", "arguments": {"city": "Paris"}}'}, {'role': 'tool', 'content': '18C, light rain'}, {'role': 'assistant', 'content': '{"tool": "get_weather", "arguments": {"city": "Tokyo"}}'}, {'role': 'tool', 'content': '24C, clear'}, {'role': 'assistant', 'content': 'Paris is rainy and Tokyo is clear.'}, {'role': 'user', 'content': 'And Lima?'}]
max_tokens = 1000

Expected: [{'role': 'system', 'content': 'You are a helpful travel agent.'}, {'role': 'user', 'content': 'Compare the weather in Paris and Tokyo for my trip.'}, {'role': 'assistant', 'content': '{"tool": "get_weather", "arguments": {"city": "Paris"}}'}, {'role': 'tool', 'content': '18C, light rain'}, {'role': 'assistant', 'content': '{"tool": "get_weather", "arguments": {"city": "Tokyo"}}'}, {'role': 'tool', 'content': '24C, clear'}, {'role': 'assistant', 'content': 'Paris is rainy and Tokyo is clear.'}, {'role': 'user', 'content': 'And Lima?'}]

Case 2: Oldest call and its result are dropped together

Input: messages = [{'role': 'system', 'content': 'You are a helpful travel agent.'}, {'role': 'user', 'content': 'Compare the weather in Paris and Tokyo for my trip.'}, {'role': 'assistant', 'content': '{"tool": "get_weather", "arguments": {"city": "Paris"}}'}, {'role': 'tool', 'content': '18C, light rain'}, {'role': 'assistant', 'content': '{"tool": "get_weather", "arguments": {"city": "Tokyo"}}'}, {'role': 'tool', 'content': '24C, clear'}, {'role': 'assistant', 'content': 'Paris is rainy and Tokyo is clear.'}, {'role': 'user', 'content': 'And Lima?'}]
max_tokens = 52

Expected: [{'role': 'system', 'content': 'You are a helpful travel agent.'}, {'role': 'user', 'content': 'Compare the weather in Paris and Tokyo for my trip.'}, {'role': 'assistant', 'content': '{"tool": "get_weather", "arguments": {"city": "Tokyo"}}'}, {'role': 'tool', 'content': '24C, clear'}, {'role': 'assistant', 'content': 'Paris is rainy and Tokyo is clear.'}, {'role': 'user', 'content': 'And Lima?'}]

Case 3: A group that doesn't fit is dropped whole, never split

Input: messages = [{'role': 'system', 'content': 'You are a helpful travel agent.'}, {'role': 'user', 'content': 'Compare the weather in Paris and Tokyo for my trip.'}, {'role': 'assistant', 'content': '{"tool": "get_weather", "arguments": {"city": "Paris"}}'}, {'role': 'tool', 'content': '18C, light rain'}, {'role': 'assistant', 'content': '{"tool": "get_weather", "arguments": {"city": "Tokyo"}}'}, {'role': 'tool', 'content': '24C, clear'}, {'role': 'assistant', 'content': 'Paris is rainy and Tokyo is clear.'}, {'role': 'user', 'content': 'And Lima?'}]
max_tokens = 40

Expected: [{'role': 'system', 'content': 'You are a helpful travel agent.'}, {'role': 'user', 'content': 'Compare the weather in Paris and Tokyo for my trip.'}, {'role': 'assistant', 'content': 'Paris is rainy and Tokyo is clear.'}, {'role': 'user', 'content': 'And Lima?'}]

Case 4: Budget smaller than the pinned messages

Input: messages = [{'role': 'system', 'content': 'You are a helpful travel agent.'}, {'role': 'user', 'content': 'Compare the weather in Paris and Tokyo for my trip.'}, {'role': 'assistant', 'content': '{"tool": "get_weather", "arguments": {"city": "Paris"}}'}, {'role': 'tool', 'content': '18C, light rain'}, {'role': 'assistant', 'content': '{"tool": "get_weather", "arguments": {"city": "Tokyo"}}'}, {'role': 'tool', 'content': '24C, clear'}, {'role': 'assistant', 'content': 'Paris is rainy and Tokyo is clear.'}, {'role': 'user', 'content': 'And Lima?'}]
max_tokens = 5

Expected: [{'role': 'system', 'content': 'You are a helpful travel agent.'}, {'role': 'user', 'content': 'Compare the weather in Paris and Tokyo for my trip.'}]