Your personal operating system for career transition.

Token Counting and Context Budget Management in Python

Context window overflow is a runtime surprise you can prevent. A well-designed AI application counts tokens at assembly time — before the API call — and enforces budgets per component. This article covers the mechanics of token counting in Python and the architecture of a reusable budget manager.

Why token counting matters beyond "avoid overflow"

Token counting has three separate uses in a production AI application:

1. Overflow prevention — avoid 400 errors from requests that exceed the model's context window.
2. Cost estimation — calculate approximate request cost before sending it. Input tokens and output tokens are priced separately.
3. Component budget allocation — enforce how many tokens each component (system prompt, retrieved context, history) is allowed to consume so one component cannot crowd out others.

Getting any of these wrong silently degrades quality or creates runtime failures.

tiktoken: accurate counting for OpenAI models

The tiktoken library is OpenAI's official tokenizer. For OpenAI models, it gives exact counts:

import tiktoken

def count_tokens(text: str, model: str = "gpt-4o") -> int:
    enc = tiktoken.encoding_for_model(model)
    return len(enc.encode(text))

# For message arrays (includes per-message overhead):
def count_messages_tokens(messages: list[dict], model: str = "gpt-4o") -> int:
    enc = tiktoken.encoding_for_model(model)
    tokens = 0
    for msg in messages:
        tokens += 4  # per-message overhead: role + content formatting
        tokens += len(enc.encode(msg.get("content", "")))
    tokens += 2  # reply priming
    return tokens

Each message in the array has a small overhead (roughly 4 tokens for role formatting) beyond the content tokens. For large batches, this overhead is negligible. For short messages, it can be meaningful.

Anthropic and other providers

Anthropic does not publish a Python tokenizer for Claude. Use the token-counting API endpoint for exact counts:

import anthropic

client = anthropic.Anthropic()

def count_tokens_anthropic(messages: list[dict], model: str = "claude-3-5-haiku-20241022") -> int:
    response = client.messages.count_tokens(
        model=model,
        messages=messages,
    )
    return response.input_tokens

For non-OpenAI, non-Anthropic providers where no tokenizer is available, use a character-based approximation:

def estimate_tokens(text: str) -> int:
    # English technical text: ~4 chars per token on average
    # Use math.ceil to err on the side of over-counting
    import math
    return math.ceil(len(text) / 4)

This approximation is accurate to within ±15% for typical English text. Use it for budgeting; do not rely on it for billing calculations.

The PromptBudget pattern

The core pattern is a stateful budget manager that enforces constraints at assembly time:

import math
from dataclasses import dataclass, field


@dataclass
class BudgetAllocation:
    label: str
    tokens: int
    included: bool


class PromptBudget:
    def __init__(
        self,
        total_tokens: int,
        reserve_output: int = 1024,
        count_fn=None,
    ) -> None:
        self._available = total_tokens - reserve_output
        self._allocations: list[BudgetAllocation] = []
        self._count = count_fn or (lambda t: math.ceil(len(t.split()) * 1.3))

    def consume(self, label: str, text: str, required: bool = False) -> bool:
        tokens = self._count(text)
        if tokens > self._available:
            self._allocations.append(BudgetAllocation(label, tokens, included=False))
            if required:
                raise ValueError(f"Required component '{label}' exceeds budget ({tokens} > {self._available})")
            return False
        self._available -= tokens
        self._allocations.append(BudgetAllocation(label, tokens, included=True))
        return True

    def remaining(self) -> int:
        return self._available

    def report(self) -> dict:
        return {
            a.label: {"tokens": a.tokens, "included": a.included}
            for a in self._allocations
        } | {"remaining": self._available}

Usage in a RAG prompt assembly function:

def assemble_rag_prompt(
    system: str,
    user_query: str,
    chunks: list[str],
    history: list[dict],
    model_context_window: int = 16000,
) -> list[dict]:
    budget = PromptBudget(model_context_window, reserve_output=2000)

    # Required components first
    budget.consume("system", system, required=True)
    budget.consume("user", user_query, required=True)

    # History: newest first, drop oldest if needed
    included_history = []
    for msg in reversed(history):
        if budget.consume(f"history_{msg['role']}", msg["content"]):
            included_history.insert(0, msg)

    # Retrieved context: add as many chunks as fit
    included_chunks = []
    for i, chunk in enumerate(chunks):
        if budget.consume(f"chunk_{i}", chunk):
            included_chunks.append(chunk)

    # Assemble
    context_text = "\n\n".join(f"<source>\n{c}\n</source>" for c in included_chunks)
    system_with_context = system + (f"\n\n## Context\n{context_text}" if context_text else "")

    messages = [{"role": "system", "content": system_with_context}]
    messages.extend(included_history)
    messages.append({"role": "user", "content": user_query})
    return messages

Common mistakes

No token check before the API call. The API returns a 400 if you exceed the context window. This is a runtime exception that should be a pre-flight check. Budget management turns a runtime exception into a graceful drop.

Counting only the user message. The system prompt, history, and retrieved context all consume tokens. Count the full messages array, not just the latest user turn.

Using character count as a proxy for tokens without a multiplier. Raw character count is always lower than token count for code and technical text. Use the 1/4 approximation or tiktoken for realistic estimates.

Forgetting the output reservation. A model with a 16K context window and a 14K-token input leaves only 2K tokens for output. If your feature needs verbose output, reserve more. Track P95 output token length and use that as your reservation.

Practical rule

Count tokens at the point where you are assembling the prompt array. If any component exceeds its allocation, decide at that point: truncate, drop, or raise. Never pass the overflow responsibility to the API call.