Context Engineering for LLMs — How a Missing `max_tokens` Caused a $12k/Week Token Waste Incident

import tiktoken
from typing import List, Dict

class TokenBudgetManager:
    def __init__(self, model: str = "gpt-4", max_context: int = 128000):
        self.encoding = tiktoken.encoding_for_model(model)
        self.max_context = max_context
        # Allocate budgets: system (1K), tools (2K), history (4K), RAG (3K)
        self.budgets = {
            "system": 1024,
            "tools": 2048,
            "history": 4096,
            "rag": 3072,
        }
        # Reserve 10% for the model's response
        self.reserved = int(max_context * 0.1)
        self.available = max_context - self.reserved

    def count_tokens(self, text: str) -> int:
        return len(self.encoding.encode(text))

    def trim_to_budget(self, text: str, budget: int) -> str:
        tokens = self.encoding.encode(text)
        if len(tokens) <= budget:
            return text
        # Truncate from the middle to preserve beginning and end
        half_budget = budget // 2
        trimmed = tokens[:half_budget] + tokens[-half_budget:]
        return self.encoding.decode(trimmed)

    def build_context(self, system_prompt: str, tools: List[Dict], history: List[str], rag_docs: List[str]) -> str:
        # Enforce budgets
        system_prompt = self.trim_to_budget(system_prompt, self.budgets["system"])
        tools_str = "\n".join([str(t) for t in tools])
        tools_str = self.trim_to_budget(tools_str, self.budgets["tools"])
        history_str = "\n".join(history[-10:])  # Last 10 messages
        history_str = self.trim_to_budget(history_str, self.budgets["history"])
        rag_str = "\n".join(rag_docs[:5])  # Top 5 docs
        rag_str = self.trim_to_budget(rag_str, self.budgets["rag"])

        context = f"{system_prompt}\n\nTools:\n{tools_str}\n\nHistory:\n{history_str}\n\nContext:\n{rag_str}"
        total_tokens = self.count_tokens(context)
        if total_tokens > self.available:
            # Emergency trim: reduce RAG and history
            rag_str = self.trim_to_budget(rag_str, self.budgets["rag"] // 2)
            history_str = self.trim_to_budget(history_str, self.budgets["history"] // 2)
            context = f"{system_prompt}\n\nTools:\n{tools_str}\n\nHistory:\n{history_str}\n\nContext:\n{rag_str}"
        return context

# Usage example
manager = TokenBudgetManager()
context = manager.build_context(
    system_prompt="You are a helpful assistant.",
    tools=[{"name": "search", "description": "Search the web"}],
    history=["User: What's the weather?", "Assistant: It's sunny."],
    rag_docs=["The weather today is sunny with a high of 25°C."]
)
print(f"Context token count: {manager.count_tokens(context)}")

import os
from typing import List, Dict, Optional
from langchain_openai import ChatOpenAI
from langchain.agents import AgentExecutor, create_openai_tools_agent
from langchain.tools import tool
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
from langchain.memory import ConversationSummaryBufferMemory

# Set your API key
os.environ["OPENAI_API_KEY"] = "sk-your-key-here"

# Define a simple tool
@tool
def get_weather(city: str) -> str:
    """Get the current weather for a city."""
    # Simulate API call
    return f"The weather in {city} is sunny, 25°C."

# Context manager with token budgeting
class ProductionContextManager:
    def __init__(self, max_token_limit: int = 8000):
        self.max_token_limit = max_token_limit
        self.memory = ConversationSummaryBufferMemory(
            llm=ChatOpenAI(model="gpt-3.5-turbo", temperature=0),
            max_token_limit=max_token_limit,
            return_messages=True,
        )

    def add_message(self, message):
        self.memory.chat_memory.add_message(message)
        # Check if we need to summarize
        if self.memory.chat_memory.messages[-1].token_count > self.max_token_limit:
            self.memory.prune()

    def get_context(self) -> List:
        return self.memory.chat_memory.messages

# Build the agent
llm = ChatOpenAI(model="gpt-4", temperature=0)
tools = [get_weather]
prompt = ChatPromptTemplate.from_messages([
    ("system", "You are a helpful assistant. Use the tools provided to answer questions."),
    MessagesPlaceholder(variable_name="chat_history"),
    ("human", "{input}"),
    MessagesPlaceholder(variable_name="agent_scratchpad"),
])

agent = create_openai_tools_agent(llm, tools, prompt)
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True, handle_parsing_errors=True)

# Run the agent with context management
context_manager = ProductionContextManager(max_token_limit=4000)

user_inputs = ["What's the weather in Paris?", "And in London?", "What was the first city I asked about?"]
for user_input in user_inputs:
    # Add user input to context
    context_manager.add_message(HumanMessage(content=user_input))
    # Get context (history)
    chat_history = context_manager.get_context()
    # Run agent
    result = agent_executor.invoke({"input": user_input, "chat_history": chat_history})
    # Add agent response to context
    context_manager.add_message(AIMessage(content=result["output"]))
    print(f"User: {user_input}")
    print(f"Agent: {result['output']}")
    print("---")

# Anti-pattern: trying to force a 100K document into context
# Instead, use RAG to retrieve only the relevant chunks

from langchain_community.vectorstores import Chroma
from langchain_openai import OpenAIEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter

# Bad: injecting entire document
# context = open("large_document.txt").read()  # 100K tokens

# Good: split and retrieve only relevant chunks
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=500,
    chunk_overlap=50,
)
chunks = text_splitter.split_text(open("large_document.txt").read())

vectorstore = Chroma.from_texts(
    texts=chunks,
    embedding=OpenAIEmbeddings(),
    persist_directory="./chroma_db"
)

# Retrieve only the top 3 most relevant chunks
query = "What is the capital of France?"
retrieved_docs = vectorstore.similarity_search(query, k=3)
context = "\n".join([doc.page_content for doc in retrieved_docs])
print(f"Context token count: {len(context.split())}")  # ~1500 tokens, not 100K

import asyncio
from collections import deque
import time

class BatchedContextManager:
    def __init__(self, batch_size: int = 5, batch_interval: float = 2.0):
        self.batch_size = batch_size
        self.batch_interval = batch_interval
        self.pending_updates = deque()
        self.last_flush = time.time()
        self.context_cache = {}

    async def add_update(self, session_id: str, update: dict):
        self.pending_updates.append((session_id, update))
        if len(self.pending_updates) >= self.batch_size or \
           (time.time() - self.last_flush) >= self.batch_interval:
            await self.flush()

    async def flush(self):
        # Batch process all pending updates
        updates_by_session = {}
        while self.pending_updates:
            session_id, update = self.pending_updates.popleft()
            if session_id not in updates_by_session:
                updates_by_session[session_id] = []
            updates_by_session[session_id].append(update)

        # Apply updates (simulated)
        for session_id, updates in updates_by_session.items():
            # In production, you'd update a database or in-memory store
            self.context_cache[session_id] = updates[-1]  # Keep latest
        self.last_flush = time.time()

    def get_context(self, session_id: str) -> dict:
        return self.context_cache.get(session_id, {})

# Usage in an async agent
async def run_agent():
    manager = BatchedContextManager()
    # Simulate concurrent requests
    tasks = []
    for i in range(10):
        tasks.append(manager.add_update(f"session_{i % 3}", {"user_input": f"query_{i}"}))
    await asyncio.gather(*tasks)
    print(manager.context_cache)

asyncio.run(run_agent())

# Mistake 1: Full history injection
# Bad:
history = all_messages  # 80K tokens

# Good:
history = all_messages[-10:]  # Last 10 messages, ~2K tokens

# Mistake 2: Untrimmed tool outputs
# Bad:
tool_output = db.query("SELECT * FROM large_table")  # 10MB JSON

# Good:
tool_output = db.query("SELECT * FROM large_table LIMIT 10")  # Truncate at source
# Or truncate after:
MAX_TOOL_OUTPUT_TOKENS = 200
import tiktoken
enc = tiktoken.encoding_for_model("gpt-4")
tokens = enc.encode(str(tool_output))
if len(tokens) > MAX_TOOL_OUTPUT_TOKENS:
    tool_output = enc.decode(tokens[:MAX_TOOL_OUTPUT_TOKENS]) + "... (truncated)"

# Mistake 3: Asking the model to remember
# Bad:
system_prompt = "Remember the user's name is Alice."

# Good:
# Store in external memory
user_memory = {"name": "Alice"}
# Inject when relevant
if "name" in user_memory:
    context += f"\nThe user's name is {user_memory['name']}."

# Decision tree for choosing the right technique

def choose_technique(task_type: str, context_size: int, need_long_term_memory: bool):
    if task_type == "one-shot classification":
        return "prompt engineering"
    elif task_type == "multi-turn agent":
        if context_size > 10000:
            return "context engineering + RAG"
        else:
            return "context engineering"
    elif task_type == "skill acquisition":
        return "fine-tuning"
    elif need_long_term_memory:
        return "context engineering + external memory"
    else:
        return "context engineering"

# Examples
print(choose_technique("one-shot classification", 500, False))  # prompt engineering
print(choose_technique("customer support agent", 50000, True))  # context engineering + RAG + external memory

import logging
import hashlib
import json
from datetime import datetime

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("context_monitor")

def log_context(context: str, session_id: str, turn: int):
    """Log context metadata for debugging."""
    import tiktoken
    enc = tiktoken.encoding_for_model("gpt-4")
    token_count = len(enc.encode(context))
    context_hash = hashlib.sha256(context.encode()).hexdigest()[:8]
    
    log_entry = {
        "timestamp": datetime.utcnow().isoformat(),
        "session_id": session_id,
        "turn": turn,
        "token_count": token_count,
        "context_hash": context_hash,
    }
    logger.info(json.dumps(log_entry))
    
    # Alert if token count exceeds threshold
    if token_count > 8000:
        logger.warning(f"Context token count {token_count} exceeds threshold for session {session_id}")
        # Trigger alert (e.g., send to PagerDuty)
        # send_alert(f"High context token count: {token_count}")

# Usage in agent loop
session_id = "session_123"
turn = 0
context = "..."  # Built context
log_context(context, session_id, turn)