Optimizing Context Windows: Token Budgeting Best Practices

The model’s context window keeps getting bigger — 128K, 1M, 10M. That size is a trap: every token you stuff in costs money, adds latency, and can degrade the answer as the model’s attention dilutes. Treat the context window like RAM, not like a trash can. Budget it.

In this post, you will learn how to budget tokens deliberately: measure your inputs, allocate a per-section budget, manage long conversations with sliding and compaction, and use caching so repeated context stops being repeat-billed.

Measure before you budget

You can’t budget what you can’t count. Every SDK exposes token counts — and you should always log them (it doubles as cost monitoring):

from google import genai

client = genai.Client()
response = client.models.generate_content(
    model="gemini-2.5-pro", contents="Summarize this."
)
meta = response.usage_metadata
print(meta.prompt_token_count, meta.candidates_token_count)

Log prompt vs. completion tokens per request and you’ll immediately see which feature is leaking context.

The classic budget model

Aim for a prompt budget that leaves headroom. A common split for a 128K window:

Section	Budget	Rule
System prompt + instructions	~2K	Trim mercilessly; this is billed on every request
Retrieved context (RAG)	~32K	Only the top-k chunks that passed the reranker
Conversation history	~16K	Sliding window, most recent wins
Working memory / tools	~8K	Tool schemas only for tools likely to be called
Output	up to max_output_tokens	Reserve explicitly
Headroom	~25%	Never fill the window — the model reasons better with slack

The exact numbers depend on your window; the shape — explicit per-section budgets plus headroom — is the practice.

Enforce the budget in code

Build a tiny TokenBudget that tracks and hard-limits what you send:

class TokenBudget:
    def __init__(self, limit: int, reserve_output: int):
        self.limit = limit
        self.remaining = limit - reserve_output

    def add(self, label: str, text: str, count_tokens) -> bool:
        n = count_tokens(text)
        if n > self.remaining:
            print(f"dropping '{label}' ({n} tokens > {self.remaining} left)")
            return False
        self.remaining -= n
        return True

Then assemble the prompt and let the budget decide what fits — nothing exceeds, and the expensive sections get cut first by design.

Manage long conversations: sliding + compaction

For chat apps, the naive “send all history” path is how bills explode. Two mechanisms, in order:

1. Sliding window. Keep the last N messages, always including the latest. Cheap and fine for most apps:

messages = system_prompt + history[-12:]   # most recent 12 turns

2. Compaction. When history is long, have the model compress old turns into a summary once, then use the summary as the oldest message. This is lossy — store the full transcript separately if you need it later.

def compact(history) -> str:
    resp = client.models.generate_content(
        model="gemini-2.5-pro",
        contents="Compress this conversation into a dense summary "
                 "preserving all facts, decisions, and pending actions:\n"
                 + json.dumps(history, indent=2),
    )
    return resp.text

Cache the stable prefix

In agent loops, most of the prompt is identical across turns: system prompt, tool schemas, and the fixed corpus context. Context caching bills those cached tokens at a fraction of the normal input price and skips re-processing them — huge savings in long loops:

# Gemini: mark the stable prefix as cacheable
cached_contents = [
    "system", system_prompt,
    "tool-schemas", json.dumps(tools),
    "documentation", fixed_docs,     # the stable RAG corpus
]
resp = client.models.generate_content(
    model="gemini-2.5-pro",
    contents=[*cached_contents, *live_turn],
    config={
        "cached_content": cached_contents,   # cache the stable prefix
        "temperature": 0,
    },
)

Cache the prefix that never changes; keep the live turn outside the cache. Monitor your cache hit rate — if it drops, someone changed the prefix construction and your bill just jumped.

Diagnose context bloat

When latency or cost climbs, check the obvious culprits in order:

1. Unbounded history → add the sliding window.
2. Full RAG dump instead of top-k → pass the reranker’s top-k, not everything.
3. Entire tool schema for every call → pass schemas only for tools the model actually has.
4. Duplicated system prompts → once, at the top.

Putting It All Together

The complete budget assembler — sectioned budget, sliding window, compaction, and cache-aware prompt building — is in this gist. Point it at a real chat/agent flow and watch prompt tokens drop 40–70% on the first run.

Conclusion & Next Steps

You now treat the context window as a budget: measured, sectioned, capped, and cached. Next steps: log token counts per section so you can see which section is bloating over time, add a semantic cache so near-identical queries reuse whole responses, and set a hard alert when average prompt tokens cross your budget line.

References / Sources

• Gemini API context caching. https://ai.google.dev/gemini-api/docs/caching
• Gemini API token counting and usage metadata. https://ai.google.dev/gemini-api/docs
• LlamaIndex token-counting callback. https://docs.llamaindex.ai