💻 Exercise: Exercise: Trim history to a token budget

📝 Instructions

Using the provided count_tokens proxy (a whitespace word count), write context_budget_trim(messages, budget). Keep a leading 'system' message if present, then keep the most recent messages whose combined token count fits within budget. Return the kept messages in their original order. Return [] if budget <= 0. Standard library only.

🧪 Initial Code / Tests 1 Files

📄 evaluate.py

from unittest import TestCase
from exercise import context_budget_trim, count_tokens


class Evaluate(TestCase):
    def test_keeps_system_message(self):
        msgs = [{'role': 'system', 'content': 'sys'}, {'role': 'user', 'content': 'a b c'}]
        self.assertEqual(context_budget_trim(msgs, 5)[0]['role'], 'system')

    def test_prefers_recent_messages(self):
        msgs = [{'role': 'user', 'content': 'one two three'}, {'role': 'user', 'content': 'four'}]
        self.assertEqual(context_budget_trim(msgs, 1)[-1]['content'], 'four')

    def test_respects_budget(self):
        msgs = [{'role': 'user', 'content': 'a b'}, {'role': 'user', 'content': 'c d'}]
        kept = context_budget_trim(msgs, 2)
        self.assertLessEqual(sum(count_tokens(m['content']) for m in kept), 2)

    def test_zero_budget_returns_empty(self):
        self.assertEqual(context_budget_trim([{'role': 'user', 'content': 'x'}], 0), [])

✅ Solutions 1 Files

📄 exercise.py

def count_tokens(text):
    return len(text.split())


def context_budget_trim(messages, budget):
    if budget <= 0:
        return []
    system = []
    rest = list(messages)
    if rest and rest[0].get('role') == 'system':
        system = [rest.pop(0)]
    remaining = budget - sum(count_tokens(m.get('content', '')) for m in system)
    kept = []
    for m in reversed(rest):
        cost = count_tokens(m.get('content', ''))
        if cost <= remaining:
            kept.append(m)
            remaining -= cost
        else:
            break
    return system + list(reversed(kept))