Context Window Management for Long-Running Agents: Engineering Patterns

import tiktoken
import logging
from typing import Optional
from dataclasses import dataclass, field
from langchain_openai import ChatOpenAI
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, SystemMessage
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

logger = logging.getLogger(__name__)


@dataclass
class ContextBudget:
    """Token budget allocation for a single LLM call."""
    model_max_tokens: int  # Model's hard context limit
    system_prompt_tokens: int  # Reserved for system prompt
    tool_definitions_tokens: int  # Reserved for tool definitions
    generation_tokens: int  # Reserved for model's output
    history_budget: int = field(init=False)

    def __post_init__(self):
        self.history_budget = (
            self.model_max_tokens
            - self.system_prompt_tokens
            - self.tool_definitions_tokens
            - self.generation_tokens
        )
        if self.history_budget < 1000:
            raise ValueError(
                f"Insufficient history budget ({self.history_budget} tokens). "
                f"Reduce system prompt, tool definitions, or generation reservation."
            )


class ProductionContextManager:
    """
    Manages conversation context for long-running agents with auto-compression.

    Strategy: Keep the last KEEP_VERBATIM messages intact.
    When total history exceeds the budget, compress the oldest messages
    into a structured summary that preserves decisions, facts, and commitments.
    """

    KEEP_VERBATIM = 6  # Always keep the last 6 messages verbatim
    COMPRESSION_THRESHOLD = 0.75  # Compress when at 75% of history budget

    def __init__(
        self,
        budget: ContextBudget,
        model_name: str = "gpt-4o",
        compression_model: str = "gpt-4o-mini"  # cheaper model for compression
    ):
        self.budget = budget
        self.model_name = model_name
        self.messages: list[BaseMessage] = []
        self.compression_count = 0

        try:
            self.tokenizer = tiktoken.encoding_for_model(model_name)
        except KeyError:
            self.tokenizer = tiktoken.get_encoding("cl100k_base")

        self.compression_llm = ChatOpenAI(
            model=compression_model,
            temperature=0,
            max_tokens=800  # Summaries should be concise
        )

    def _count_tokens(self, messages: list[BaseMessage]) -> int:
        """Count tokens for a list of messages."""
        total = 0
        for msg in messages:
            # Account for message overhead (role, format tokens)
            total += 4  # approximate overhead per message
            total += len(self.tokenizer.encode(str(msg.content)))
        return total

    def _compress_messages(self, messages: list[BaseMessage]) -> AIMessage:
        """Compress a batch of messages into a structured summary."""
        if not messages:
            return AIMessage(content="[No prior context]")

        conversation_text = "\n".join(
            f"{type(m).__name__.replace('Message', '')}: {m.content[:500]}"
            for m in messages
        )

        compression_prompt = ChatPromptTemplate.from_messages([
            ("system",
             """You are compressing conversation history for an AI agent.
             Preserve ALL of the following:
             - Decisions made and their rationale
             - Facts established (numbers, names, dates, statuses)
             - Commitments or actions agreed to
             - Open questions or blockers identified
             - Any constraints or requirements specified

             Format as a structured summary under these headings:
             DECISIONS: [bullet points]
             ESTABLISHED FACTS: [bullet points]
             COMMITMENTS: [bullet points]
             OPEN ITEMS: [bullet points]

             Be specific. Preserve exact values (numbers, names, IDs) not paraphrases.
             Maximum 600 words."""),
            ("human", "Compress this conversation history:\n\n{history}")
        ])

        chain = compression_prompt | self.compression_llm | StrOutputParser()
        summary = chain.invoke({"history": conversation_text})

        self.compression_count += 1
        logger.info(
            f"context_compressed | compression_count={self.compression_count} | "
            f"messages_compressed={len(messages)} | summary_chars={len(summary)}"
        )

        return AIMessage(content=f"[COMPRESSED CONTEXT - Summary of {len(messages)} prior messages]\n{summary}")

    def add_message(self, message: BaseMessage) -> None:
        """Add a message and trigger compression if budget threshold is exceeded."""
        self.messages.append(message)
        self._maybe_compress()

    def _maybe_compress(self) -> None:
        """Check if compression is needed and compress if so."""
        current_tokens = self._count_tokens(self.messages)
        threshold_tokens = int(self.budget.history_budget * self.COMPRESSION_THRESHOLD)

        if current_tokens <= threshold_tokens:
            return

        # Protect the last KEEP_VERBATIM messages from compression
        if len(self.messages) <= self.KEEP_VERBATIM:
            logger.warning(
                f"context_budget_exceeded | tokens={current_tokens} | budget={self.budget.history_budget} | "
                f"cannot_compress: only {len(self.messages)} messages"
            )
            return

        messages_to_compress = self.messages[:-self.KEEP_VERBATIM]
        messages_to_keep = self.messages[-self.KEEP_VERBATIM:]

        summary_message = self._compress_messages(messages_to_compress)
        self.messages = [summary_message] + messages_to_keep

        new_tokens = self._count_tokens(self.messages)
        logger.info(
            f"context_after_compression | before_tokens={current_tokens} | "
            f"after_tokens={new_tokens} | reduction_pct={round((1 - new_tokens/current_tokens)*100, 1)}"
        )

    def get_messages(self) -> list[BaseMessage]:
        """Return the current managed context window."""
        return list(self.messages)

    def get_token_usage(self) -> dict:
        """Return current token usage statistics."""
        current = self._count_tokens(self.messages)
        return {
            "current_tokens": current,
            "budget_tokens": self.budget.history_budget,
            "utilisation_pct": round(current / self.budget.history_budget * 100, 1),
            "compression_count": self.compression_count,
            "message_count": len(self.messages)
        }


# ---- Usage Example ----
if __name__ == "__main__":
    # GPT-4o configuration: 128K total, reserve 2K system + 3K tools + 4K generation
    budget = ContextBudget(
        model_max_tokens=128000,
        system_prompt_tokens=2000,
        tool_definitions_tokens=3000,
        generation_tokens=4000
    )

    ctx = ProductionContextManager(budget=budget)

    # Simulate a long-running agent conversation
    for i in range(25):
        ctx.add_message(HumanMessage(
            content=f"Step {i+1}: Analyse the Q{(i%4)+1} financial data for the {['Americas','EMEA','APAC'][i%3]} region and identify variance drivers."
        ))
        ctx.add_message(AIMessage(
            content=f"Analysis for Step {i+1}: Revenue variance of ${(i+1)*120}K driven by [detailed analysis would appear here across multiple sentences covering all the key points for this step of the analysis workflow]..."
        ))

        usage = ctx.get_token_usage()
        if (i + 1) % 5 == 0:
            print(f"After turn {i+1}: {usage['current_tokens']} tokens ({usage['utilisation_pct']}% of budget), compressions: {usage['compression_count']}")

    print(f"\nFinal message count: {len(ctx.get_messages())}")
    print(f"Final token usage: {ctx.get_token_usage()}")