Context Engineering for Agents¶

class ContextBudget:
    def __init__(self, max_tokens=128000):
        self.max_tokens = max_tokens
        self.allocations = {
            "system": 0.10,        # 10% for system prompt + tools
            "state": 0.10,         # 10% for current state/plan
            "history": 0.30,       # 30% for conversation history
            "tool_results": 0.30,  # 30% for tool outputs
            "generation": 0.20,    # 20% reserved for model output
        }

    def budget_for(self, category):
        return int(self.max_tokens * self.allocations[category])

    def fits(self, category, token_count):
        return token_count <= self.budget_for(category)

class AgentState:
    def __init__(self):
        self.goal = ""
        self.plan = []
        self.completed_steps = []
        self.current_step = None
        self.key_findings = []
        self.failed_approaches = []
        self.open_questions = []

    def to_context(self) -> str:
        """Serialize for injection into LLM context."""
        return f"""
## Current Goal
{self.goal}

## Plan
{self._format_plan()}

## Key Findings So Far
{chr(10).join(f'- {f}' for f in self.key_findings[-10:])}

## Failed Approaches (DO NOT retry)
{chr(10).join(f'- {f}' for f in self.failed_approaches[-5:])}

## Open Questions
{chr(10).join(f'- {q}' for q in self.open_questions[-5:])}
"""

    def _format_plan(self):
        lines = []
        for i, step in enumerate(self.plan):
            if i < len(self.completed_steps):
                lines.append(f"[DONE] {step}")
            elif i == len(self.completed_steps):
                lines.append(f"[CURRENT] {step}")
            else:
                lines.append(f"[ ] {step}")
        return "\n".join(lines)

# Write state to disk, read back when needed
import json

STATE_FILE = ".agent/state.json"
PLAN_FILE = ".agent/PLAN.md"
FINDINGS_FILE = ".agent/findings.md"

def save_state(state: AgentState):
    with open(STATE_FILE, 'w') as f:
        json.dump(state.__dict__, f, indent=2)

def load_state() -> AgentState:
    with open(STATE_FILE) as f:
        data = json.load(f)
    state = AgentState()
    state.__dict__.update(data)
    return state

# Agent reads state from file at start of each step
# This survives context compaction and session restarts

def compact_history(messages, target_tokens=2000):
    """Summarize conversation history to fit budget."""
    # Keep system message and last 3 exchanges intact
    preserved = messages[:1] + messages[-6:]

    # Summarize everything in between
    middle = messages[1:-6]
    if not middle:
        return messages

    summary_prompt = f"""
    Summarize this conversation history, preserving:
    1. Key decisions made
    2. Important findings
    3. Current state of the task
    4. What was tried and failed

    History to summarize:
    {format_messages(middle)}
    """
    summary = llm.complete(summary_prompt, max_tokens=target_tokens)

    return [messages[0],
            {"role": "system", "content": f"[Previous conversation summary]\n{summary}"},
            *preserved]

def truncate_tool_result(result: str, max_tokens: int = 1000) -> str:
    """Truncate tool output to fit budget."""
    token_count = count_tokens(result)
    if token_count <= max_tokens:
        return result

    # Try to extract key information
    summary = llm.complete(
        f"Extract only the key facts from this tool output (max {max_tokens} tokens):\n{result}",
        max_tokens=max_tokens
    )
    return f"[Summarized from {token_count} tokens]\n{summary}"

def post_compaction_reinject(state: AgentState) -> str:
    """After context compaction, re-inject critical state."""
    return f"""
[Context was compacted. Here is the current state to continue from:]

{state.to_context()}

[Continue from where you left off. The plan and findings above are accurate.]
"""

class ContextRetriever:
    def __init__(self, vector_store):
        self.vector_store = vector_store
        self.cache = {}

    def get_relevant_context(self, query: str, k: int = 5) -> str:
        """Fetch relevant documents/history for current query."""
        results = self.vector_store.search(query, top_k=k)
        return "\n\n".join([
            f"[Relevant context {i+1}]:\n{r.text}"
            for i, r in enumerate(results)
        ])

    def index_interaction(self, query: str, response: str, metadata: dict):
        """Index completed interactions for future retrieval."""
        text = f"Q: {query}\nA: {response}"
        self.vector_store.add(text, metadata=metadata)