Context Window Management at Scale: Token Budgets and RAG

Total Context Window (e.g., 128K)
├── System prompt .............. 500-2,000 tokens
├── Tool definitions ............. 1,000-5,000 tokens
├── Conversation history ......... 2,000-50,000+ tokens (grows unbounded!)
├── Retrieved RAG context ........ 2,000-30,000 tokens
├── Tool call results ............ 500-10,000 tokens per call
├── User query ................... 100-2,000 tokens
└── Response reserve ............. 1,000-4,000 tokens

from dataclasses import dataclass
from enum import Enum

class UseCase(str, Enum):
    SUPPORT_CHAT = "support_chat"
    DOCUMENT_QA = "document_qa"
    CODE_REVIEW = "code_review"
    AGENT_WORKFLOW = "agent_workflow"

@dataclass(frozen=True)
class ContextBudget:
    """Fixed token allocation per context component."""
    total_window: int
    system_prompt: int
    tool_definitions: int
    conversation_history: int
    retrieved_context: int
    tool_results: int
    user_query: int
    response_reserve: int

    @property
    def max_input(self) -> int:
        return self.total_window - self.response_reserve

    def validate(self) -> None:
        allocated = (
            self.system_prompt + self.tool_definitions +
            self.conversation_history + self.retrieved_context +
            self.tool_results + self.user_query
        )
        if allocated > self.max_input:
            raise ValueError(
                f"Budget overflow: {allocated} tokens allocated, "
                f"max input is {self.max_input}"
            )

BUDGETS = {
    UseCase.SUPPORT_CHAT: ContextBudget(
        total_window=128_000,
        system_prompt=800,
        tool_definitions=2_000,
        conversation_history=4_000,
        retrieved_context=3_000,
        tool_results=2_000,
        user_query=500,
        response_reserve=2_000,
    ),
    UseCase.DOCUMENT_QA: ContextBudget(
        total_window=128_000,
        system_prompt=500,
        tool_definitions=0,
        conversation_history=2_000,
        retrieved_context=8_000,
        tool_results=0,
        user_query=1_000,
        response_reserve=4_000,
    ),
    UseCase.AGENT_WORKFLOW: ContextBudget(
        total_window=128_000,
        system_prompt=1_000,
        tool_definitions=4_000,
        conversation_history=3_000,
        retrieved_context=4_000,
        tool_results=6_000,
        user_query=500,
        response_reserve=3_000,
    ),
}

import tiktoken

enc = tiktoken.encoding_for_model("gpt-4o")

def count_tokens(text: str) -> int:
    return len(enc.encode(text))

def truncate_to_token_limit(text: str, max_tokens: int, keep: str = "end") -> str:
    """Truncate text to fit token budget. Keep 'end' for chat, 'start' for documents."""
    tokens = enc.encode(text)
    if len(tokens) <= max_tokens:
        return text
    if keep == "end":
        return enc.decode(tokens[-max_tokens:])
    return enc.decode(tokens[:max_tokens])

@dataclass
class ContextComponents:
    system_prompt: str = ""
    tool_definitions: str = ""
    conversation_history: str = ""
    retrieved_context: str = ""
    tool_results: str = ""
    user_query: str = ""

class ContextAssembler:
    """Assemble context within token budget."""

    def __init__(self, budget: ContextBudget):
        self.budget = budget

    def assemble(self, components: ContextComponents) -> list[dict]:
        """Build messages array within budget, prioritizing user query and recent history."""
        messages = []
        system = truncate_to_token_limit(
            components.system_prompt, self.budget.system_prompt, keep="start"
        )
        if components.tool_definitions:
            tools = truncate_to_token_limit(
                components.tool_definitions, self.budget.tool_definitions, keep="start"
            )
            system += f"\n\nAvailable tools:\n{tools}"
        messages.append({"role": "system", "content": system})

        # Retrieved context — high priority for QA
        if components.retrieved_context:
            rag = truncate_to_token_limit(
                components.retrieved_context, self.budget.retrieved_context, keep="start"
            )
            messages.append({
                "role": "system",
                "content": f"Relevant context:\n{rag}",
            })

        # Conversation history — keep recent (from the end)
        if components.conversation_history:
            history = truncate_to_token_limit(
                components.conversation_history,
                self.budget.conversation_history,
                keep="end",
            )
            messages.append({"role": "assistant", "content": history})

        # Tool results
        if components.tool_results:
            results = truncate_to_token_limit(
                components.tool_results, self.budget.tool_results, keep="end"
            )
            messages.append({"role": "system", "content": f"Tool results:\n{results}"})

        # User query — never truncate (highest priority)
        messages.append({"role": "user", "content": components.user_query})

        # Verify total budget
        total = sum(count_tokens(m["content"]) for m in messages)
        if total > self.budget.max_input:
            raise ContextOverflowError(
                f"Context exceeds budget: {total} > {self.budget.max_input}"
            )

        return messages

class ContextOverflowError(Exception):
    pass

from dataclasses import dataclass
from datetime import datetime

@dataclass
class Message:
    role: str
    content: str
    timestamp: datetime
    token_count: int = 0

class SlidingWindowMemory:
    """Keep the most recent messages within a token budget."""

    def __init__(self, max_tokens: int = 4000):
        self.max_tokens = max_tokens
        self.messages: list[Message] = []

    def add(self, role: str, content: str) -> None:
        tokens = count_tokens(content)
        self.messages.append(Message(
            role=role,
            content=content,
            timestamp=datetime.utcnow(),
            token_count=tokens,
        ))
        self._trim()

    def _trim(self) -> None:
        """Remove oldest messages until within budget."""
        while self.total_tokens > self.max_tokens and len(self.messages) > 1:
            self.messages.pop(0)

    @property
    def total_tokens(self) -> int:
        return sum(m.token_count for m in self.messages)

    def to_messages(self) -> list[dict]:
        return [{"role": m.role, "content": m.content} for m in self.messages]

    def to_text(self) -> str:
        return "\n".join(f"{m.role}: {m.content}" for m in self.messages)

class PriorityMemory(SlidingWindowMemory):
    """Sliding window that preserves high-priority messages."""

    PRIORITY_PATTERNS = [
        r"ORD-\d+",          # Order IDs
        r"user_id:\s*\S+",   # User identifiers
        r"preference:",       # User preferences
    ]

    def _is_priority(self, content: str) -> bool:
        import re
        return any(re.search(p, content) for p in self.PRIORITY_PATTERNS)

    def _trim(self) -> None:
        priority = [m for m in self.messages if self._is_priority(m.content)]
        non_priority = [m for m in self.messages if not self._is_priority(m.content)]

        priority_tokens = sum(m.token_count for m in priority)

        # Trim non-priority from the front until within budget
        while (
            priority_tokens + sum(m.token_count for m in non_priority)
            > self.max_tokens
            and len(non_priority) > 0
        ):
            non_priority.pop(0)

        self.messages = priority + non_priority

class RollingSummarizer:
    """Summarize old conversation turns to save context budget."""

    def __init__(
        self,
        client,
        model: str = "gpt-4o-mini",
        summarize_every: int = 6,
        keep_recent: int = 4,
    ):
        self.client = client
        self.model = model
        self.summarize_every = summarize_every
        self.keep_recent = keep_recent
        self.summary: str = ""
        self.recent_messages: list[Message] = []
        self.turn_count: int = 0

    async def add_message(self, role: str, content: str) -> None:
        self.recent_messages.append(Message(
            role=role, content=content,
            timestamp=datetime.utcnow(),
            token_count=count_tokens(content),
        ))
        self.turn_count += 1

        if self.turn_count % self.summarize_every == 0:
            await self._summarize_old_turns()

    async def _summarize_old_turns(self) -> None:
        if len(self.recent_messages) <= self.keep_recent:
            return

        to_summarize = self.recent_messages[:-self.keep_recent]
        self.recent_messages = self.recent_messages[-self.keep_recent:]

        conversation_text = "\n".join(
            f"{m.role}: {m.content}" for m in to_summarize
        )
        existing = f"Previous summary: {self.summary}\n\n" if self.summary else ""

        response = await self.client.chat.completions.create(
            model=self.model,
            messages=[{
                "role": "user",
                "content": f"""{existing}Summarize this conversation segment.
                Preserve: names, IDs, decisions, preferences, unresolved issues.
                Omit: greetings, filler, repeated information.
                Max 300 words.

                Conversation:
                {conversation_text}
                """,
            }],
            max_tokens=400,
        )
        self.summary = response.choices[0].message.content

    def build_context(self) -> str:
        parts = []
        if self.summary:
            parts.append(f"[Conversation summary: {self.summary}]")
        parts.extend(f"{m.role}: {m.content}" for m in self.recent_messages)
        return "\n".join(parts)

async def hierarchical_summarize(
    client,
    document: str,
    chunk_size: int = 4000,
    model: str = "gpt-4o-mini",
) -> str:
    """Summarize a long document in layers: chunks → section summaries → final summary."""
    chunks = split_into_chunks(document, chunk_size)

    # Layer 1: Summarize each chunk
    chunk_summaries = []
    for chunk in chunks:
        resp = await client.chat.completions.create(
            model=model,
            messages=[{
                "role": "user",
                "content": f"Summarize key facts, dates, names, and obligations:\n\n{chunk}",
            }],
            max_tokens=300,
        )
        chunk_summaries.append(resp.choices[0].message.content)

    # Layer 2: Combine chunk summaries into final summary
    combined = "\n\n".join(chunk_summaries)
    if count_tokens(combined) > 3000:
        resp = await client.chat.completions.create(
            model=model,
            messages=[{
                "role": "user",
                "content": f"Create a comprehensive summary from these section summaries:\n\n{combined}",
            }],
            max_tokens=800,
        )
        return resp.choices[0].message.content

    return combined

def split_into_chunks(text: str, max_tokens: int) -> list[str]:
    """Split text into chunks at paragraph boundaries."""
    paragraphs = text.split("\n\n")
    chunks, current = [], ""
    for para in paragraphs:
        if count_tokens(current + para) > max_tokens and current:
            chunks.append(current.strip())
            current = para
        else:
            current += "\n\n" + para if current else para
    if current.strip():
        chunks.append(current.strip())
    return chunks

from dataclasses import dataclass

@dataclass
class RetrievedChunk:
    text: str
    source: str
    score: float
    token_count: int

class RAGContextManager:
    """Retrieve and assemble relevant context within token budget."""

    def __init__(self, vector_store, budget_tokens: int = 4000, top_k: int = 10):
        self.vector_store = vector_store
        self.budget_tokens = budget_tokens
        self.top_k = top_k

    async def retrieve_context(self, query: str) -> str:
        results = await self.vector_store.search(query, limit=self.top_k)

        chunks: list[RetrievedChunk] = []
        total_tokens = 0

        for result in sorted(results, key=lambda r: r.score, reverse=True):
            tokens = count_tokens(result.text)
            if total_tokens + tokens > self.budget_tokens:
                # Try truncating this chunk to fill remaining budget
                remaining = self.budget_tokens - total_tokens
                if remaining > 100:
                    truncated = truncate_to_token_limit(result.text, remaining, keep="start")
                    chunks.append(RetrievedChunk(
                        text=truncated,
                        source=result.source,
                        score=result.score,
                        token_count=remaining,
                    ))
                break

            chunks.append(RetrievedChunk(
                text=result.text,
                source=result.source,
                score=result.score,
                token_count=tokens,
            ))
            total_tokens += tokens

        return self._format_chunks(chunks)

    def _format_chunks(self, chunks: list[RetrievedChunk]) -> str:
        formatted = []
        for i, chunk in enumerate(chunks, 1):
            formatted.append(
                f"[Source {i}: {chunk.source} (relevance: {chunk.score:.2f})]\n{chunk.text}"
            )
        return "\n\n---\n\n".join(formatted)

async def retrieve_with_reranking(
    query: str,
    vector_store,
    reranker,
    budget_tokens: int = 4000,
    initial_k: int = 20,
    final_k: int = 5,
) -> str:
    # Stage 1: Broad retrieval
    candidates = await vector_store.search(query, limit=initial_k)

    # Stage 2: Rerank for precision
    reranked = await reranker.rerank(
        query=query,
        documents=[c.text for c in candidates],
        top_k=final_k,
    )

    # Stage 3: Assemble within budget
    manager = RAGContextManager(vector_store, budget_tokens)
    selected = [candidates[r.index] for r in reranked]
    return manager._format_chunks([
        RetrievedChunk(text=c.text, source=c.source, score=c.score, token_count=count_tokens(c.text))
        for c in selected
    ])

User Query
    ↓
[Sliding Window] Recent 4 turns (2K tokens)
    +
[Rolling Summary] Compressed older turns (500 tokens)
    +
[RAG Retrieval] Top-5 relevant KB articles (3K tokens)
    +
[System Prompt] Agent instructions (800 tokens)
    = ~6.3K tokens total (vs 50K+ unmanaged)

async def map_reduce_analysis(client, document: str, question: str) -> str:
    """Analyze documents too long for a single context window."""
    chunks = split_into_chunks(document, max_tokens=6000)

    # MAP: Analyze each chunk independently
    chunk_answers = []
    for i, chunk in enumerate(chunks):
        resp = await client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{
                "role": "user",
                "content": f"Document section {i+1}/{len(chunks)}:\n{chunk}\n\nQuestion: {question}\n\nAnswer based on this section only. Say 'NOT FOUND' if not in this section.",
            }],
            max_tokens=500,
        )
        answer = resp.choices[0].message.content
        if "NOT FOUND" not in answer:
            chunk_answers.append(f"Section {i+1}: {answer}")

    # REDUCE: Synthesize chunk answers into final response
    if not chunk_answers:
        return "The answer was not found in the document."

    resp = await client.chat.completions.create(
        model="gpt-4o",
        messages=[{
            "role": "user",
            "content": f"Question: {question}\n\nPartial answers from document sections:\n" +
                       "\n".join(chunk_answers) +
                       "\n\nSynthesize a complete, coherent answer.",
        }],
        max_tokens=1000,
    )
    return resp.choices[0].message.content

class AgentContextManager:
    """Provide minimal context per agent step."""

    def __init__(self, budget: ContextBudget):
        self.budget = budget
        self.step_outputs: dict[str, str] = {}
        self.summary: str = ""

    def record_step_output(self, step_name: str, output: str) -> None:
        # Store full output externally (database), keep summary in context
        self.step_outputs[step_name] = output

    async def context_for_step(self, step_name: str, required_steps: list[str]) -> str:
        """Build context from required step outputs only."""
        parts = []
        if self.summary:
            parts.append(f"Workflow summary: {self.summary}")

        for req_step in required_steps:
            if req_step in self.step_outputs:
                output = truncate_to_token_limit(
                    self.step_outputs[req_step],
                    self.budget.tool_results // len(required_steps),
                    keep="start",
                )
                parts.append(f"[{req_step} output: {output}]")

        return "\n".join(parts)

@dataclass
class ContextMetrics:
    request_id: str
    use_case: str
    total_input_tokens: int
    system_tokens: int
    history_tokens: int
    rag_tokens: int
    tool_tokens: int
    query_tokens: int
    budget_utilization: float  # total_input / max_input
    truncated_components: list[str]
    timestamp: str

class ContextMonitor:
    """Track context usage patterns for optimization."""

    def __init__(self, sink):
        self.sink = sink

    async def record(self, metrics: ContextMetrics) -> None:
        await self.sink.insert(metrics)

        if metrics.budget_utilization > 0.9:
            await self.sink.alert(
                f"Context budget >90% for {metrics.use_case}: "
                f"{metrics.total_input_tokens} tokens"
            )

        if metrics.truncated_components:
            await self.sink.increment(
                "context_truncations",
                tags={"components": ",".join(metrics.truncated_components)},
            )