Implementing Token-Budget-Aware Prompt Assembly: Dynamically Allocating Context Space Across System Instructions, Retrieved Chunks, Conversation History, and Tool Results at Runtime

Introduction

Picture this: your RAG-powered customer support assistant is humming along in production, then at 2 PM on a Tuesday, a user pastes in a 3,000-word legal document and asks a question about it. Your prompt assembly code dutifully shoves in the full system prompt, the top-10 retrieved chunks, the last 20 conversation turns, and the tool results from three API calls — and promptly throws a context_length_exceeded error. Your error handler catches it, returns a generic "something went wrong" message, and you've just lost a customer.

This isn't an edge case. It's the default failure mode of naïve RAG and agent systems. The context window is a fixed resource, and every component of your prompt is competing for space inside it. Without deliberate allocation logic, you're one verbose conversation away from a crash — or worse, silently truncated inputs that give the model incomplete information and degrade quality without any error signal at all.

By the end of this lesson, you'll know how to build a production-grade prompt assembly system that treats the token budget as a first-class runtime constraint. You'll write code that measures every component before assembly, enforces priority-based allocation, gracefully degrades when space is tight, and never accidentally blows the limit.

What you'll learn:

• How to precisely count tokens for each prompt component before assembly using model-specific tokenizers
• How to define and enforce a priority hierarchy across system instructions, retrieved chunks, conversation history, and tool results
• How to implement graceful degradation strategies — truncation, compression, and eviction — when budget is tight
• How to build a reusable PromptBudget abstraction that works across different models with different context sizes
• How to tune and monitor your allocation strategy in production using observability hooks

---

Prerequisites

You should be comfortable with the following before diving in:

• Calling LLM APIs in Python (OpenAI or Anthropic)
• Basic RAG architecture — if you need a refresher, see Retrieval-Augmented Generation Explained: How RAG Works and When to Use It
• Managing conversation history in the Chat Completions format
• Basic familiarity with Understanding Tokens: How LLMs Tokenize Text and Why It Affects Your Inputs, Outputs, and Costs

You'll need tiktoken, anthropic, and openai Python packages installed.

---

The Anatomy of a Context Window

Before you can allocate intelligently, you need to know what you're allocating. A typical production LLM request has four distinct content regions:

System instructions — Your base persona, behavioral rules, output format requirements, safety guardrails. These are relatively static but can grow large if you have elaborate few-shot examples baked in.

Retrieved chunks — Documents or passages retrieved from a vector store based on the current query. These are highly variable in size and count. A production document Q&A system might retrieve anywhere from 1 to 20 chunks per query.

Conversation history — Prior turns in the dialogue. A brand new session has zero history; a long support ticket thread might have 50+ exchanges spanning thousands of tokens.

Tool results — Outputs from function calls, API responses, database query results, code execution outputs. These can be tiny (a boolean flag) or enormous (a paginated JSON response from an ERP system).

Every one of these has to fit inside context_window_size - output_reserve. That last term is critical and often forgotten: you need to reserve tokens for the model's response. If your model has a 128K context window and you fill all 128K with input, the model has zero tokens to generate output with. In practice, you should reserve at least 1,000–4,000 tokens for generation, more for tasks that require long outputs.

# The fundamental constraint
CONTEXT_LIMIT = 128_000      # model's total context window
OUTPUT_RESERVE = 2_000       # tokens we guarantee for generation
INPUT_BUDGET = CONTEXT_LIMIT - OUTPUT_RESERVE   # 126,000 tokens for input

---

Token Counting That Actually Works

The first prerequisite of budget-aware assembly is accurate token counting. Eyeballing character counts or dividing by 4 (the "rough average") will get you into trouble. You need exact counts using the same tokenizer the model uses.

OpenAI / tiktoken

For GPT-3.5 and GPT-4 family models, tiktoken gives you exact counts:

import tiktoken

def count_tokens_openai(text: str, model: str = "gpt-4o") -> int:
    """Count tokens for OpenAI models using tiktoken."""
    try:
        enc = tiktoken.encoding_for_model(model)
    except KeyError:
        # Fall back to cl100k_base for unknown models
        enc = tiktoken.get_encoding("cl100k_base")
    return len(enc.encode(text))


def count_message_tokens_openai(messages: list[dict], model: str = "gpt-4o") -> int:
    """
    Count tokens for a full messages array, including per-message overhead.
    Overhead is 4 tokens per message (role + delimiters) + 3 for reply priming.
    See: https://platform.openai.com/docs/guides/chat/managing-tokens
    """
    try:
        enc = tiktoken.encoding_for_model(model)
    except KeyError:
        enc = tiktoken.get_encoding("cl100k_base")
    
    num_tokens = 3  # reply priming overhead
    for message in messages:
        num_tokens += 4  # role + structural tokens
        for key, value in message.items():
            if isinstance(value, str):
                num_tokens += len(enc.encode(value))
            elif isinstance(value, list):
                # Handle content arrays (multimodal messages)
                for item in value:
                    if isinstance(item, dict) and item.get("type") == "text":
                        num_tokens += len(enc.encode(item["text"]))
            if key == "name":
                num_tokens -= 1  # name field adjustment
    return num_tokens

Anthropic / Claude

Anthropic doesn't expose a public tokenizer library the way OpenAI does, but their API provides a token counting endpoint that you should use:

import anthropic

_anthropic_client = anthropic.Anthropic()

def count_tokens_anthropic(messages: list[dict], system: str = "", model: str = "claude-3-5-sonnet-20241022") -> int:
    """
    Count tokens via Anthropic's token counting API.
    This is a real API call — cache results aggressively.
    """
    response = _anthropic_client.messages.count_tokens(
        model=model,
        system=system,
        messages=messages,
    )
    return response.input_tokens

Building a Cached Token Counter

Here's a practical counter that caches by content hash:

import hashlib
import functools
from typing import Literal

ModelFamily = Literal["openai", "anthropic"]

class TokenCounter:
    """Token counter with content-hash-based caching."""
    
    def __init__(self, model: str, family: ModelFamily):
        self.model = model
        self.family = family
        self._cache: dict[str, int] = {}
    
    def _hash(self, content: str) -> str:
        return hashlib.sha256(content.encode()).hexdigest()[:16]
    
    def count(self, text: str) -> int:
        """Count tokens for a plain text string."""
        key = self._hash(text)
        if key not in self._cache:
            if self.family == "openai":
                self._cache[key] = count_tokens_openai(text, self.model)
            else:
                # For Anthropic, wrap in a single user message for counting
                self._cache[key] = count_tokens_anthropic(
                    messages=[{"role": "user", "content": text}],
                    model=self.model
                )
        return self._cache[key]
    
    def count_messages(self, messages: list[dict]) -> int:
        """Count tokens for a full messages array."""
        # Create a cache key from the full message array
        content_str = str(messages)
        key = self._hash(content_str)
        if key not in self._cache:
            if self.family == "openai":
                self._cache[key] = count_message_tokens_openai(messages, self.model)
            else:
                self._cache[key] = count_tokens_anthropic(
                    messages=messages,
                    model=self.model
                )
        return self._cache[key]
    
    def clear_cache(self):
        self._cache.clear()

---

Designing the Budget Allocation Model

Now for the interesting part: deciding how much of your input budget each component gets. This is fundamentally a prioritization problem, and the right answer depends on your application. A customer support bot has different priorities than a code analysis agent.

Here's the allocation framework we'll build:

Total Input Budget (e.g., 126,000 tokens)
├── System Instructions (FIXED priority — always included fully)
│   └── Hard minimum: actual size; no compression
├── Tool Results (HIGH priority — current task depends on them)
│   └── Allocated: up to 25% of remaining budget
├── Retrieved Chunks (HIGH priority — primary knowledge source)
│   └── Allocated: up to 40% of remaining budget
├── Conversation History (MEDIUM priority — context, but summarizable)
│   └── Allocated: whatever remains after above, up to a configured max
└── Current User Message (FIXED — always included)
    └── Hard minimum: actual size; never truncated

The PromptBudget Dataclass

Let's formalize this as a configuration object:

from dataclasses import dataclass, field
from typing import Optional

@dataclass
class BudgetConfig:
    """Configuration for token budget allocation."""
    
    # Model constraints
    context_window: int = 128_000
    output_reserve: int = 2_000
    
    # Allocation fractions (of remaining budget after fixed components)
    tool_results_fraction: float = 0.25
    retrieved_chunks_fraction: float = 0.40
    history_fraction: float = 0.30
    
    # Hard limits (token counts, not fractions)
    max_system_tokens: Optional[int] = None      # None = no limit
    max_chunks: int = 20                          # max number of chunks
    max_history_turns: int = 50                   # max conversation turns
    max_tool_result_tokens: int = 8_000           # per tool result
    
    # Safety margin (buffer beyond output_reserve)
    safety_margin: int = 200
    
    @property
    def input_budget(self) -> int:
        return self.context_window - self.output_reserve - self.safety_margin
    
    def validate(self):
        total = self.tool_results_fraction + self.retrieved_chunks_fraction + self.history_fraction
        if total > 1.0:
            raise ValueError(f"Allocation fractions sum to {total:.2f}, must be <= 1.0")

---

The Prompt Assembler: Core Implementation

With the counting infrastructure in place, let's build the assembler itself. The key principle is measure before you commit: count every component's tokens before adding it to the final messages list.

from dataclasses import dataclass
from typing import Optional
import logging

logger = logging.getLogger(__name__)

@dataclass
class AssemblyResult:
    """Output from the prompt assembler, including metadata for observability."""
    messages: list[dict]
    total_tokens: int
    budget_used: int
    budget_available: int
    
    # Component breakdown
    system_tokens: int = 0
    history_tokens: int = 0
    chunks_tokens: int = 0
    tool_tokens: int = 0
    query_tokens: int = 0
    
    # Truncation/eviction info
    history_turns_dropped: int = 0
    chunks_dropped: int = 0
    tool_results_truncated: int = 0
    
    @property
    def utilization_pct(self) -> float:
        return (self.budget_used / self.budget_available) * 100


class TokenBudgetAssembler:
    """
    Runtime prompt assembler that enforces token budget constraints
    across all content regions.
    """
    
    def __init__(
        self,
        counter: TokenCounter,
        config: BudgetConfig,
    ):
        self.counter = counter
        self.config = config
        config.validate()
    
    def assemble(
        self,
        *,
        system_prompt: str,
        user_query: str,
        retrieved_chunks: list[str],
        conversation_history: list[dict],
        tool_results: list[dict],
    ) -> AssemblyResult:
        """
        Assemble a prompt that fits within the token budget.
        
        Priority order:
        1. System prompt (fixed, always included)
        2. Current user query (fixed, always included)
        3. Tool results (high priority, up to allocation)
        4. Retrieved chunks (high priority, up to allocation)
        5. Conversation history (fills remaining space)
        """
        budget = self.config.input_budget
        remaining = budget
        result = AssemblyResult(
            messages=[],
            total_tokens=0,
            budget_used=0,
            budget_available=budget,
        )
        
        # === STEP 1: Fixed components (system + query) ===
        system_tokens = self.counter.count(system_prompt)
        query_tokens = self.counter.count(user_query)
        fixed_tokens = system_tokens + query_tokens
        
        if fixed_tokens > budget:
            raise ValueError(
                f"Fixed components alone ({fixed_tokens} tokens) exceed "
                f"input budget ({budget} tokens). Reduce system prompt or "
                f"output_reserve."
            )
        
        remaining -= fixed_tokens
        result.system_tokens = system_tokens
        result.query_tokens = query_tokens
        
        # === STEP 2: Tool results ===
        tool_budget = int(remaining * self.config.tool_results_fraction)
        assembled_tools, tool_tokens, tools_truncated = self._fit_tool_results(
            tool_results, tool_budget
        )
        remaining -= tool_tokens
        result.tool_tokens = tool_tokens
        result.tool_results_truncated = tools_truncated
        
        # === STEP 3: Retrieved chunks ===
        chunk_budget = int(remaining * self.config.retrieved_chunks_fraction)
        assembled_chunks, chunk_tokens, chunks_dropped = self._fit_chunks(
            retrieved_chunks, chunk_budget
        )
        remaining -= chunk_tokens
        result.chunks_tokens = chunk_tokens
        result.chunks_dropped = chunks_dropped
        
        # === STEP 4: Conversation history ===
        # History gets whatever's left, subject to configured max
        history_budget = min(
            remaining - 100,  # small buffer
            int(budget * self.config.history_fraction)
        )
        assembled_history, history_tokens, turns_dropped = self._fit_history(
            conversation_history, history_budget
        )
        result.history_tokens = history_tokens
        result.history_turns_dropped = turns_dropped
        
        # === STEP 5: Build final messages array ===
        messages = []
        
        # Inject retrieved chunks and tool results into system prompt
        # (or as separate messages — depends on model and design)
        enriched_system = self._build_enriched_system(
            system_prompt, assembled_chunks, assembled_tools
        )
        
        # Build messages: history + current query
        messages = list(assembled_history)  # already oldest-first
        messages.append({"role": "user", "content": user_query})
        
        total_used = self.counter.count_messages(messages) + self.counter.count(enriched_system)
        
        result.messages = messages
        result.total_tokens = total_used
        result.budget_used = total_used
        
        logger.info(
            "Prompt assembled: %d/%d tokens (%.1f%%) | "
            "sys=%d hist=%d chunks=%d tools=%d query=%d | "
            "dropped: %d history turns, %d chunks, %d tool results truncated",
            total_used, budget, result.utilization_pct,
            result.system_tokens, result.history_tokens,
            result.chunks_tokens, result.tool_tokens, result.query_tokens,
            result.history_turns_dropped, result.chunks_dropped,
            result.tool_results_truncated
        )
        
        return result, enriched_system

---

Implementing Graceful Degradation Strategies

The real craft is in the _fit_* methods. Each content type requires a different degradation strategy.

Fitting Tool Results

Tool results often come from function calling and tool use integrations. They can be enormous — imagine a tool that queries a database and returns 500 rows of JSON. Strategies here: truncate from the tail (most API responses have the most relevant data at the top), or summarize structured data.

    def _fit_tool_results(
        self, 
        tool_results: list[dict],
        budget: int
    ) -> tuple[list[dict], int, int]:
        """
        Fit tool results into budget.
        Strategy: include all results but truncate content of oversized ones.
        """
        if not tool_results:
            return [], 0, 0
        
        fitted = []
        total_tokens = 0
        truncated_count = 0
        per_result_limit = min(
            self.config.max_tool_result_tokens,
            budget // max(len(tool_results), 1)
        )
        
        for result in tool_results:
            content = result.get("content", "")
            content_tokens = self.counter.count(content)
            
            if content_tokens <= per_result_limit:
                fitted.append(result)
                total_tokens += content_tokens
            else:
                # Truncate content to fit
                truncated_content = self._truncate_to_budget(
                    content, per_result_limit
                )
                truncated_result = {
                    **result,
                    "content": truncated_content + "\n[... truncated to fit context window]"
                }
                fitted.append(truncated_result)
                total_tokens += per_result_limit
                truncated_count += 1
            
            if total_tokens >= budget:
                break
        
        return fitted, total_tokens, truncated_count
    
    def _truncate_to_budget(self, text: str, token_budget: int) -> str:
        """
        Truncate text to fit within token_budget tokens.
        Uses binary search on character count for efficiency.
        """
        if self.counter.count(text) <= token_budget:
            return text
        
        # Binary search for the right character cutoff
        lo, hi = 0, len(text)
        while lo < hi - 1:
            mid = (lo + hi) // 2
            if self.counter.count(text[:mid]) <= token_budget:
                lo = mid
            else:
                hi = mid
        return text[:lo]

Fitting Retrieved Chunks

Retrieved chunks from your vector store have an important property: they're already ranked by relevance. This means your eviction strategy should always drop from the bottom of the ranked list first. Never evict the top chunk to make room for more lower-ranked chunks.

    def _fit_chunks(
        self, 
        chunks: list[str],
        budget: int
    ) -> tuple[list[str], int, int]:
        """
        Fit retrieved chunks into budget.
        Chunks assumed to be in descending relevance order.
        Strategy: include top chunks greedily until budget is exhausted.
        """
        if not chunks:
            return [], 0, 0
        
        fitted = []
        total_tokens = 0
        dropped = 0
        
        # Enforce hard limit on chunk count
        candidate_chunks = chunks[:self.config.max_chunks]
        dropped += max(0, len(chunks) - self.config.max_chunks)
        
        for chunk in candidate_chunks:
            chunk_tokens = self.counter.count(chunk)
            
            if total_tokens + chunk_tokens <= budget:
                fitted.append(chunk)
                total_tokens += chunk_tokens
            else:
                # Can we fit a truncated version of this chunk?
                remaining_budget = budget - total_tokens
                if remaining_budget > 200:  # worth including a partial chunk
                    truncated = self._truncate_to_budget(chunk, remaining_budget - 20)
                    fitted.append(truncated + "\n[... truncated]")
                    total_tokens += remaining_budget
                dropped += len(candidate_chunks) - len(fitted)
                break
        
        return fitted, total_tokens, dropped

Fitting Conversation History

Conversation history is the trickiest component. If you're building a conversational memory system with many turns, you have several options, each with different quality trade-offs:

1. Sliding window — Keep only the N most recent turns. Simple, but can lose critical early context.
2. Token-budget sliding window — Keep as many recent turns as fit in the budget. Better.
3. Compress-then-slide — Summarize old turns into a compact paragraph, then keep recent turns verbatim.
4. Importance scoring — Assign relevance scores to past turns and keep the highest-scoring ones regardless of recency.

We'll implement option 2 (token-budget sliding window) with option 3 (compression) as an optional enhancement:

    def _fit_history(
        self,
        history: list[dict],
        budget: int,
    ) -> tuple[list[dict], int, int]:
        """
        Fit conversation history into budget by keeping the most recent turns.
        Drops oldest turns first when over budget.
        
        history: list of {"role": ..., "content": ...} dicts, oldest first.
        """
        if not history or budget <= 0:
            return [], 0, len(history)
        
        # Enforce hard turn limit first
        capped_history = history[-self.config.max_history_turns:]
        turns_dropped = len(history) - len(capped_history)
        
        # Now greedily include from the end (most recent first),
        # then reverse back to chronological order
        fitted = []
        total_tokens = 0
        
        for turn in reversed(capped_history):
            turn_tokens = self.counter.count(turn.get("content", "")) + 4  # overhead
            if total_tokens + turn_tokens <= budget:
                fitted.append(turn)
                total_tokens += turn_tokens
            else:
                turns_dropped += 1
        
        # Restore chronological order
        fitted.reverse()
        
        return fitted, total_tokens, turns_dropped

---

Assembling the Final Prompt Structure

Now we need to think about how the assembled components are structured in the actual API call. There are two main patterns:

Pattern A: Everything in the system prompt — Inject retrieved chunks and tool results into the system message as XML-tagged sections. Works well with Claude's XML-native style.

Pattern B: Injected as separate messages — Add retrieved context as a synthetic user or assistant message before the conversation history. Works well with OpenAI's chat format.

Here's a flexible implementation of pattern A (which tends to perform better in our experience for knowledge-intensive tasks):

    def _build_enriched_system(
        self,
        base_system: str,
        chunks: list[str],
        tool_results: list[dict],
    ) -> str:
        """
        Build enriched system prompt by appending retrieved context
        and tool results as clearly delimited sections.
        """
        parts = [base_system]
        
        if chunks:
            chunk_sections = "\n\n".join(
                f"<document index=\"{i+1}\">\n{chunk}\n</document>"
                for i, chunk in enumerate(chunks)
            )
            parts.append(
                f"\n\n<retrieved_context>\n"
                f"The following documents were retrieved as relevant to the user's query. "
                f"Use them to inform your response.\n\n"
                f"{chunk_sections}\n"
                f"</retrieved_context>"
            )
        
        if tool_results:
            tool_sections = "\n\n".join(
                f"<tool_result tool_name=\"{r.get('tool_name', 'unknown')}\" "
                f"call_id=\"{r.get('call_id', '')}\">\n"
                f"{r.get('content', '')}\n"
                f"</tool_result>"
                for r in tool_results
            )
            parts.append(
                f"\n\n<tool_outputs>\n"
                f"{tool_sections}\n"
                f"</tool_outputs>"
            )
        
        return "".join(parts)

---

Putting It All Together: A Working Example

Let's see this system in action with a realistic scenario: a financial analyst assistant that retrieves relevant earnings report sections, has an ongoing conversation about a company, and has just called a tool to get current stock data.

import os
from openai import OpenAI

# Configuration for GPT-4o (128K context)
config = BudgetConfig(
    context_window=128_000,
    output_reserve=3_000,  # generous — we expect detailed financial analysis
    tool_results_fraction=0.15,   # stock data is compact
    retrieved_chunks_fraction=0.50,  # earnings reports need lots of space
    history_fraction=0.25,
    max_chunks=15,
    max_history_turns=30,
    max_tool_result_tokens=4_000,
    safety_margin=500,
)

counter = TokenCounter(model="gpt-4o", family="openai")
assembler = TokenBudgetAssembler(counter=counter, config=config)

# Realistic system prompt
system_prompt = """You are a financial analysis assistant specializing in 
equity research. You help analysts understand earnings reports, financial 
metrics, and market trends. You cite specific passages from documents when 
making claims. Your responses are precise, data-driven, and you flag 
uncertainty explicitly. When you reference data from tools, note the 
timestamp of when it was retrieved."""

# Retrieved chunks from vector search over earnings report corpus
retrieved_chunks = [
    """Q3 2024 Revenue: $18.2B, up 12% YoY. Operating margin expanded 
    by 180bps to 23.4%. Management cited strong enterprise demand and 
    favorable FX tailwinds as primary drivers...""",
    
    """Free cash flow for the quarter reached $4.1B, representing a 
    22.5% FCF margin. Capital expenditure was $1.8B, primarily directed 
    toward data center expansion in APAC...""",
    
    # ... imagine 10 more chunks here
]

# Conversation history (20 prior turns)
conversation_history = [
    {"role": "user", "content": "Let's analyze Contoso Corp's recent performance."},
    {"role": "assistant", "content": "Happy to help. What aspect would you like to focus on first?"},
    {"role": "user", "content": "Start with revenue trends over the last 4 quarters."},
    # ... 17 more turns
]

# Tool result from a market data API call
tool_results = [
    {
        "tool_name": "get_stock_quote",
        "call_id": "call_abc123",
        "content": """{
            "symbol": "CNTS",
            "price": 142.37,
            "change_pct": 2.14,
            "volume": 8_420_000,
            "market_cap": "89.2B",
            "retrieved_at": "2024-11-15T14:32:00Z"
        }"""
    }
]

user_query = "How does the FCF margin compare to their 5-year historical average, and what's driving the improvement?"

# Assemble the prompt
result, enriched_system = assembler.assemble(
    system_prompt=system_prompt,
    user_query=user_query,
    retrieved_chunks=retrieved_chunks,
    conversation_history=conversation_history,
    tool_results=tool_results,
)

print(f"Budget utilization: {result.utilization_pct:.1f}%")
print(f"History turns dropped: {result.history_turns_dropped}")
print(f"Chunks dropped: {result.chunks_dropped}")

# Make the API call
client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])
response = client.chat.completions.create(
    model="gpt-4o",
    messages=[
        {"role": "system", "content": enriched_system},
        *result.messages,
    ],
    max_tokens=3_000,
    temperature=0.2,
)

print(response.choices[0].message.content)

---

Advanced Pattern: Dynamic Priority Adjustment

Static fractions work well for most cases, but sophisticated systems benefit from adjusting allocation priorities based on runtime context. Here are two high-value examples:

Query-Type-Aware Allocation

from enum import Enum

class QueryType(Enum):
    FACTUAL_LOOKUP = "factual_lookup"      # heavy on retrieval
    CONVERSATIONAL = "conversational"       # heavy on history
    TOOL_DEPENDENT = "tool_dependent"       # heavy on tool results
    ANALYTICAL = "analytical"              # balanced

def build_config_for_query_type(
    base_config: BudgetConfig,
    query_type: QueryType
) -> BudgetConfig:
    """
    Adjust budget allocation fractions based on detected query type.
    """
    from dataclasses import replace
    
    adjustments = {
        QueryType.FACTUAL_LOOKUP: dict(
            retrieved_chunks_fraction=0.60,
            history_fraction=0.15,
            tool_results_fraction=0.10,
        ),
        QueryType.CONVERSATIONAL: dict(
            retrieved_chunks_fraction=0.20,
            history_fraction=0.55,
            tool_results_fraction=0.10,
        ),
        QueryType.TOOL_DEPENDENT: dict(
            retrieved_chunks_fraction=0.20,
            history_fraction=0.20,
            tool_results_fraction=0.45,
        ),
        QueryType.ANALYTICAL: dict(
            retrieved_chunks_fraction=0.40,
            history_fraction=0.30,
            tool_results_fraction=0.15,
        ),
    }
    
    return replace(base_config, **adjustments[query_type])

Adaptive History Compression

When history keeps getting truncated, consider adding a compression layer. This requires a secondary LLM call but can dramatically improve coherence:

async def compress_history(
    history: list[dict],
    target_token_count: int,
    counter: TokenCounter,
    client: OpenAI,
) -> list[dict]:
    """
    Summarize old conversation turns into a compact system-injected summary.
    Keeps the most recent N turns verbatim, summarizes the rest.
    """
    if not history:
        return history
    
    # Keep last 6 turns verbatim (3 exchanges)
    verbatim_turns = history[-6:]
    old_turns = history[:-6]
    
    if not old_turns:
        return history
    
    # Format old turns for summarization
    history_text = "\n".join(
        f"{t['role'].upper()}: {t['content']}"
        for t in old_turns
    )
    
    summary_response = client.chat.completions.create(
        model="gpt-4o-mini",  # cheaper model for summarization
        messages=[
            {
                "role": "system", 
                "content": "Summarize the following conversation history into a compact paragraph. "
                          "Preserve key facts, decisions made, and important context. "
                          "Be concise — under 200 words."
            },
            {"role": "user", "content": history_text}
        ],
        max_tokens=300,
        temperature=0.0,
    )
    
    summary = summary_response.choices[0].message.content
    
    # Inject summary as a synthetic system-context message
    summary_message = {
        "role": "user",
        "content": f"[Earlier conversation summary: {summary}]"
    }
    
    return [summary_message] + verbatim_turns

---

Observability and Monitoring in Production

Budget-aware assembly is wasted effort if you can't see what's happening in production. You need to log the AssemblyResult metadata for every request and track trends over time.

The observability patterns from implementing LLM observability apply directly here. At minimum, emit these metrics per request:

def emit_assembly_metrics(result: AssemblyResult, request_id: str):
    """
    Emit structured metrics for monitoring budget utilization.
    Compatible with any observability backend (Datadog, Prometheus, etc.)
    """
    metrics = {
        "request_id": request_id,
        "token.budget_used": result.budget_used,
        "token.budget_available": result.budget_available,
        "token.utilization_pct": result.utilization_pct,
        "token.system": result.system_tokens,
        "token.history": result.history_tokens,
        "token.chunks": result.chunks_tokens,
        "token.tools": result.tool_tokens,
        "token.query": result.query_tokens,
        "degradation.history_turns_dropped": result.history_turns_dropped,
        "degradation.chunks_dropped": result.chunks_dropped,
        "degradation.tool_results_truncated": result.tool_results_truncated,
    }
    
    # Alert if degradation is happening frequently
    if result.history_turns_dropped > 5:
        logger.warning(
            "Heavy history truncation: %d turns dropped for request %s",
            result.history_turns_dropped, request_id
        )
    
    if result.chunks_dropped > 3:
        logger.warning(
            "Significant chunk eviction: %d chunks dropped for request %s",
            result.chunks_dropped, request_id
        )
    
    # Your observability backend call here
    # statsd.gauge("llm.token.utilization_pct", result.utilization_pct)
    return metrics

Key dashboards to build:

• P50/P95/P99 budget utilization — If P99 is consistently above 90%, you're close to the edge
• Degradation rate — What % of requests have any history or chunk eviction?
• Component breakdown over time — Is history growing without bound? Are your system prompts getting longer?
• Budget utilization by query type — Some query patterns are systematically more expensive

---

Hands-On Exercise

Build a complete budget-aware assembler for a customer service chatbot scenario.

Setup: You have a chatbot that handles technical support for a software product. Each request involves:

• A system prompt (persona + guidelines): ~800 tokens
• Retrieved documentation chunks from a vector store (variable, up to 15 chunks)
• Conversation history (variable, 0 to 40 turns)
• Occasionally, tool results from a ticket management API call

Your task:

1. Implement a PromptAssemblerV2 that adds one feature the base implementation above lacks: chunk deduplication. Before fitting chunks into the budget, remove any chunk that has >80% token overlap with a previously included chunk. (Hint: you can use a simple sliding-window character n-gram similarity rather than embedding-based similarity for speed.)

2. Implement adaptive output reservation. Instead of a fixed output_reserve, detect whether the user is asking for a "short answer" or a "detailed explanation" (a simple keyword heuristic is fine) and adjust the output reserve accordingly — 500 tokens for short answers, 3000 for detailed ones.

3. Write a load test that generates 100 requests with varying history lengths (0 to 40 turns) and chunk counts (1 to 15), assembles each prompt, and produces a report showing:

   • What % of requests experienced any degradation
   • Average budget utilization
   • The request parameters that caused the most degradation

4. Tune the allocation fractions until fewer than 5% of your synthetic test requests experience chunk eviction (history eviction is acceptable).

Validation: Your assembler should never produce an AssemblyResult where budget_used > budget_available. Write a property-based test using hypothesis that verifies this invariant holds across randomly generated inputs.

---

Common Mistakes & Troubleshooting

Mistake 1: Counting Tokens After Assembly

The most common error is building the full messages array first, then counting tokens to check if it fits. By then it's too late — you've already assembled something that may be over budget, and trimming it cleanly is harder than building it right the first time.

Fix: Always count each component before deciding to include it.

Mistake 2: Forgetting the Message Format Overhead

A message with {"role": "user", "content": "hello"} does not cost 1 token. It costs the content tokens plus the structural overhead (~4 tokens per message for OpenAI models). In a conversation with 40 messages, this overhead alone adds ~160 tokens — enough to matter at the margins.

Fix: Use the count_message_tokens_openai function from this lesson, not a raw text counter.

Mistake 3: Using Character Count as a Token Proxy

The "divide by 4" heuristic fails badly for non-English text, code, URLs, and structured data. A JSON string with lots of brackets and colons tokenizes more expensively than prose. A German compound word tokenizes differently than its English equivalent.

Fix: Use proper tokenizers. The latency cost is negligible (tiktoken is fast enough for production use), and the accuracy improvement prevents budget overruns.

Mistake 4: Not Handling the "Budget Below Fixed Components" Case

What if the system prompt plus the current user query alone exceed the input budget? This can happen when:

• A user submits an extremely long message (pasted document, code dump)
• The system prompt has grown too large over time
• The model was changed to one with a smaller context window

# Don't do this — it silently truncates the user query
user_query = user_query[:budget - system_tokens]

# Do this — fail explicitly with a clear error
if system_tokens + query_tokens > budget:
    raise PromptBudgetError(
        f"Cannot assemble prompt: system ({system_tokens} tok) + "
        f"query ({query_tokens} tok) = {system_tokens + query_tokens} tok "
        f"exceeds budget ({budget} tok). "
        f"Consider: increasing context window, reducing system prompt size, "
        f"or rejecting requests with queries > {budget - system_tokens} tokens."
    )

Mistake 5: Fixed Allocation Fractions That Don't Sum Correctly

If your fractions for tools + chunks + history sum to 1.0, you leave no slack for component counting overhead and rounding errors. Budget 5-10% of remaining space as unallocated slack.

Mistake 6: Treating All Chunks as Equal

If your retrieval returns 15 chunks but only 8 fit in the budget, you should always drop the least relevant 7 — not arbitrarily truncate the 8th. This requires maintaining relevance scores through the assembly pipeline.

# Better chunk representation
@dataclass
class RankedChunk:
    content: str
    relevance_score: float  # from retrieval
    source: str
    
# Sort descending before fitting
chunks_sorted = sorted(chunks, key=lambda c: c.relevance_score, reverse=True)

---

Performance Considerations at Scale

When you're handling hundreds of requests per second, the token counting overhead becomes significant. Here are the optimizations that matter:

Cache system prompt token counts aggressively. Your system prompt rarely changes. Count it once at startup, store the count, and reuse it. If you're using prompt versioning, tie the cached count to the prompt version hash.

Pre-count chunks at index time. When you ingest documents into your vector store, pre-compute and store the token count of each chunk alongside its embedding. This eliminates counting overhead at query time for all chunks that were in the index when ingested.

Use fast counting for pre-screening. For extremely long inputs (>50K characters), do a character-based pre-check before invoking the tokenizer. If len(text) / 3 > budget, you can skip the tokenizer and go straight to truncation.

Batch Anthropic token count API calls. If you must count multiple components via the API, batch them into a single request where possible rather than making sequential calls.

Parallelize counting and retrieval. Token counting and vector store retrieval can happen concurrently. Don't wait for retrieval to finish before starting to count the system prompt and history.

import asyncio

async def assemble_parallel(assembler, system_prompt, query, history):
    """
    Parallelize token counting and retrieval.
    """
    # These can happen simultaneously
    system_count_task = asyncio.create_task(
        asyncio.to_thread(assembler.counter.count, system_prompt)
    )
    retrieval_task = asyncio.create_task(
        vector_store.search(query, top_k=15)
    )
    history_count_task = asyncio.create_task(
        asyncio.to_thread(assembler.counter.count_messages, history)
    )
    
    system_tokens, chunks, history_tokens = await asyncio.gather(
        system_count_task, retrieval_task, history_count_task
    )
    
    return system_tokens, chunks, history_tokens

---

Summary & Next Steps

You've built a production-grade token-budget-aware prompt assembler from the ground up. The key principles to carry forward:

Measure before you commit. Every component gets token-counted before it's added to the prompt. No surprises at the API boundary.

Priority hierarchy is the design decision. How you rank system instructions vs. retrieved content vs. history reflects what your application fundamentally cares about. Get it right for your use case, not just for convention.

Graceful degradation beats hard failure. A response built from the top 8 of 15 chunks is almost always better than an error. Design your degradation paths deliberately rather than letting the LLM API decide what to discard.

Observability makes tuning possible. Without logging AssemblyResult metadata, you're flying blind. The metrics are cheap to emit and invaluable for catching allocation regressions.

For your next steps, consider:

• Integrating this assembler with a hybrid search layer so your chunk quality improves, reducing how many chunks you need to retrieve to get good coverage
• Adding a reranking layer upstream of the assembler so that when budget forces eviction, you're always evicting the least relevant content
• Wiring the assembler into an agentic loop where tool results feed back into subsequent prompt assemblies
• Exploring cost optimization patterns — budget-aware assembly isn't just about avoiding errors, it's also about not paying for tokens you don't need

The context window is the most valuable real estate in your LLM application. Treat it that way.