Context Budgeting in Agent RAG: What to Include, What to Drop

Your context window is a finite resource. Every token you spend on irrelevant context is a token stolen from quality answers. Context budgeting is the discipline of allocating your context budget across retrieved documents, conversation history, system prompts, and tool outputs.

The Context Budget Problem

A typical agent RAG request uses context for:

• System prompt: 500-2,000 tokens
• Conversation history: 2,000-20,000 tokens
• Retrieved documents: 5,000-30,000 tokens
• Tool outputs: 1,000-10,000 tokens
• Available for response: 2,000-4,000 tokens

With Gemini’s 1M+ token context windows, you might think budgeting doesn’t matter. But larger context doesn’t mean better answers—it means more noise for the model to sift through.

Budget Allocation Strategies

Fixed Budget Allocation

from dataclasses import dataclass

@dataclass
class ContextBudget:
    system_prompt: int = 2000
    conversation_history: int = 8000
    retrieved_documents: int = 15000
    tool_outputs: int = 5000
    response_reserve: int = 3000
    
    @property
    def total(self) -> int:
        return (self.system_prompt + self.conversation_history + 
                self.retrieved_documents + self.tool_outputs + 
                self.response_reserve)
    
    def remaining_for_retrieval(self, used: dict) -> int:
        return max(0, self.retrieved_documents - used.get("retrieved", 0))

Dynamic Budget Allocation

class DynamicBudgetAllocator:
    def __init__(self, total_budget: int = 128000):
        self.total_budget = total_budget
    
    def allocate(self, query: str, history_length: int, num_tools: int) -> dict:
        # Estimate system prompt size
        system_tokens = 1500
        
        # Allocate for history based on complexity
        history_ratio = min(0.3, history_length * 0.02)
        history_budget = int(self.total_budget * history_ratio)
        
        # Reserve for tool outputs
        tool_budget = num_tools * 3000
        
        # Response reserve
        response_budget = 4000
        
        # Everything else goes to retrieval
        retrieval_budget = (
            self.total_budget - 
            system_tokens - 
            history_budget - 
            tool_budget - 
            response_budget
        )
        
        return {
            "system": system_tokens,
            "history": history_budget,
            "retrieval": retrieval_budget,
            "tools": tool_budget,
            "response": response_budget,
        }

Filtering Retrieved Content

Relevance Scoring

import numpy as np

class RelevanceScorer:
    def __init__(self, embedder, llm):
        self.embedder = embedder
        self.llm = llm
    
    def score_and_filter(
        self, 
        query: str, 
        documents: list[str], 
        budget_tokens: int,
        min_score: float = 0.3
    ) -> list[dict]:
        # Embed query and documents
        query_embedding = self.embedder.embed_query(query)
        doc_embeddings = self.embedder.embed_documents(documents)
        
        # Compute similarity scores
        scores = np.dot(doc_embeddings, query_embedding)
        
        # Score each document
        scored_docs = []
        for i, (doc, score) in enumerate(zip(documents, scores)):
            token_count = estimate_tokens(doc)
            relevance = float(score)
            
            # Value = relevance / cost
            value = relevance / max(token_count / 1000, 0.1)
            
            scored_docs.append({
                "content": doc,
                "relevance": relevance,
                "tokens": token_count,
                "value": value,
                "index": i,
            })
        
        # Sort by value and fit within budget
        scored_docs.sort(key=lambda x: x["value"], reverse=True)
        
        selected = []
        tokens_used = 0
        for doc in scored_docs:
            if doc["relevance"] < min_score:
                continue
            if tokens_used + doc["tokens"] > budget_tokens:
                # Try to truncate
                remaining = budget_tokens - tokens_used
                if remaining > 200:
                    doc["content"] = truncate_to_tokens(doc["content"], remaining)
                    doc["tokens"] = remaining
                    selected.append(doc)
                break
            selected.append(doc)
            tokens_used += doc["tokens"]
        
        return selected

Progressive Disclosure

class ProgressiveDisclosure:
    """Load more context only when needed."""
    
    def __init__(self, vector_store, llm):
        self.vector_store = vector_store
        self.llm = llm
    
    def retrieve_with_budget(self, query: str, budget_tokens: int) -> dict:
        # Stage 1: Get summaries (low token cost)
        summaries = self.vector_store.similarity_search(query, k=10, search_type="mmr")
        summary_text = "\n".join([f"- {s.metadata['title']}: {s.page_content[:100]}" for s in summaries])
        
        # Ask LLM which documents are relevant
        selection_prompt = f"""Given the query: "{query}"
        
        Which of these documents are most relevant? Return indices (comma-separated).
        
        {summary_text}
        
        Relevant indices:"""
        
        selected_indices = self.llm.invoke([HumanMessage(content=selection_prompt)])
        indices = parse_indices(selected_indices.content)
        
        # Stage 2: Load full content only for selected documents
        full_docs = []
        tokens_used = estimate_tokens(summary_text)
        
        for idx in indices:
            if idx < len(summaries):
                doc = summaries[idx]
                full_content = self.vector_store.get_full_document(doc.id)
                doc_tokens = estimate_tokens(full_content)
                
                if tokens_used + doc_tokens <= budget_tokens:
                    full_docs.append({
                        "title": doc.metadata["title"],
                        "content": full_content,
                        "tokens": doc_tokens,
                    })
                    tokens_used += doc_tokens
        
        return {
            "documents": full_docs,
            "total_tokens": tokens_used,
            "budget": budget_tokens,
            "utilization": tokens_used / budget_tokens,
        }

Compression Strategies

Extractive Compression

def extractive_compress(document: str, query: str, max_tokens: int) -> str:
    """Keep only the most relevant sentences."""
    sentences = split_into_sentences(document)
    
    # Score each sentence by relevance to query
    scored = []
    for sent in sentences:
        # Simple keyword overlap scoring
        query_words = set(query.lower().split())
        sent_words = set(sent.lower().split())
        overlap = len(query_words & sent_words) / max(len(query_words), 1)
        scored.append((sent, overlap))
    
    # Select top sentences that fit budget
    scored.sort(key=lambda x: x[1], reverse=True)
    selected = []
    tokens_used = 0
    for sent, score in scored:
        sent_tokens = estimate_tokens(sent)
        if tokens_used + sent_tokens <= max_tokens:
            selected.append(sent)
            tokens_used += sent_tokens
    
    return " ".join(selected)

LLM-Based Summarization

def summarize_for_context(document: str, query: str, max_tokens: int) -> str:
    """Use LLM to create a query-focused summary."""
    
    prompt = f"""Summarize the following document, focusing on information relevant to: "{query}"
    
    Keep the summary under {max_tokens} tokens.
    
    Document:
    {document}
    
    Query-focused summary:"""
    
    response = llm.invoke([HumanMessage(content=prompt)])
    return response.content

History Windowing

class ConversationHistoryManager:
    def __init__(self, max_tokens: int = 8000):
        self.max_tokens = max_tokens
    
    def manage_history(self, messages: list) -> list:
        if estimate_tokens(messages) <= self.max_tokens:
            return messages
        
        # Strategy 1: Keep system message + recent messages
        system_msgs = [m for m in messages if m.role == "system"]
        other_msgs = [m for m in messages if m.role != "system"]
        
        # Always keep first user message (original intent)
        if other_msgs:
            first_user = other_msgs[0]
            recent = other_msgs[1:]
            
            # Add recent messages that fit
            kept = [first_user]
            tokens_used = estimate_tokens(system_msgs + [first_user])
            
            for msg in reversed(recent):
                msg_tokens = estimate_tokens([msg])
                if tokens_used + msg_tokens <= self.max_tokens:
                    kept.insert(1, msg)
                    tokens_used += msg_tokens
                else:
                    break
            
            return system_msgs + kept
        
        return messages[-10:]  # Fallback: last 10 messages

Monitoring Context Usage

class ContextBudgetMonitor:
    def __init__(self):
        self.metrics = []
    
    def record(self, query: str, allocation: dict, actual_usage: dict):
        self.metrics.append({
            "timestamp": time.time(),
            "query_length": estimate_tokens([query]),
            "allocation": allocation,
            "actual_usage": actual_usage,
            "utilization": actual_usage.get("total", 0) / allocation.get("total", 1),
            "retrieval_utilization": actual_usage.get("retrieved", 0) / max(allocation.get("retrieval", 1), 1),
        })
    
    def get_insights(self) -> dict:
        recent = self.metrics[-100:]
        return {
            "avg_utilization": sum(m["utilization"] for m in recent) / len(recent),
            "avg_retrieval_utilization": sum(m["retrieval_utilization"] for m in recent) / len(recent),
            "over_budget_rate": sum(1 for m in recent if m["utilization"] > 1.0) / len(recent),
            "under_budget_rate": sum(1 for m in recent if m["utilization"] < 0.5) / len(recent),
        }

Context budgeting is about quality, not just fitting tokens. An agent with 80% well-chosen context outperforms one with 100% stuffed context. Budget wisely, compress aggressively, and always reserve room for the model to think.