Context Engineering: Token Budgets, Compaction and Running Notes for Long Chats

"""Context engineering for a long chat: fit the conversation into a token budget without losing what matters."""
import re

import harness

SYSTEM = ("You are the assistant in this conversation. Answer the user's latest question briefly. "
          "Facts from earlier in the conversation are in your notes; trust them.")
BUDGET = 1000        # tokens of conversation history kept between calls
TOTAL_BUDGET = 900   # tokens of the whole request: system prompt, notes, kept turns and the question


# ---------- Step 1: count and window ----------
def count_tokens(messages):
    """Tokens a list of messages costs: harness.tokens(content) + 4 per message for the role and separators."""
    # TODO (Step 1): the sum over messages of harness.tokens(m["content"]) + 4
    raise NotImplementedError("Step 1: write count_tokens()")


def recent_window(turns, budget=BUDGET):
    """The most recent whole user+assistant pairs that fit in the budget, oldest first. Never split a pair."""
    # TODO (Step 1): walk the pairs from the end (i = len(turns) - 2, then i - 2, ...); put each pair turns[i:i + 2]
    # in front of what you have while count_tokens() of the result stays within budget; stop at the first that does not.
    raise NotImplementedError("Step 1: write recent_window()")


# ---------- Step 2: trim bulky tool output ----------
BULKY = re.compile(r"^(Search results:|Train options|Quotes for)")


def trim_bulky(turns, keep_last=2):
    """Copies of the turns where each bulky assistant message (search results, options, quotes) older than the last
    keep_last pairs is replaced by a one-line stub: "[<its first line, without a trailing colon> (<N> options)
    shown to the user; omitted]", N being its number of lines after the first."""
    raise NotImplementedError("trim_bulky() arrives in Step 2")


# ---------- Step 3: compact into running notes ----------
NOTES_PROMPT = """You keep the running notes for a long conversation between a user and an assistant. Rewrite the notes to include what the new turns add.

Write two sections:
FACTS: people, dates and times, bookings with their reference codes, budgets, health and dietary needs, preferences. Copy names, numbers, times and codes exactly.
DECISIONS: what the user has decided, with the latest decision replacing any earlier one it contradicts (write what changed, e.g. "Hiroshima dropped; Nara instead").

Leave out search results, options the user did not choose, and suggestions. Use short bullet points. At most 180 words.

Notes so far:
{notes}

New turns:
{turns}

Updated notes:"""


def compact(turns, budget=BUDGET):
    """Replay the conversation pair by pair. Whenever the kept turns exceed the budget, fold the oldest pairs into the
    notes (one NOTES_MODEL call) until the kept turns are at most budget // 2. Returns (notes, kept)."""
    raise NotImplementedError("compact() arrives in Step 3")


# ---------- Step 4: assemble within the total budget ----------
def build_messages(notes, kept, question, total=TOTAL_BUDGET):
    """System message (SYSTEM, a blank line, "Notes on the earlier conversation:", the notes), then as many of the
    most recent kept pairs as fit within `total` tokens for the whole request, then the question as the last user
    message. The system message and the question always go in; pairs are added newest first and never split."""
    system = {"role": "system", "content": f"{SYSTEM}\n\nNotes on the earlier conversation:\n{notes}"}
    q = {"role": "user", "content": question}
    return [system] + list(kept) + [q]  # Step 4 fits this into the total budget


def prepare(turns):
    """What the app keeps between calls: trim, then compact."""
    return compact(trim_bulky(turns))


def answer(notes, kept, question):
    return harness.chat(build_messages(notes, kept, question), harness.ANSWER_MODEL, 300)