Introduction to Context Window Optimization with Pydantic AI
Large language models have transformed how developers build intelligent applications, but every model has a fundamental constraint: the context window. The context window defines how many tokens a model can process in a single request, encompassing both the input prompt and the generated output. When you exceed this limit, requests fail, costs skyrocket, and response quality degrades. Pydantic AI, a typed framework for building AI agents, provides powerful tools to manage and optimize context windows effectively.
Context window optimization is the practice of structuring, compressing, and managing the information you send to an LLM so that every token contributes meaningful value. This tutorial walks you through the concepts, techniques, and practical implementations using Pydantic AI to build efficient, production-ready AI agents.
Why Context Window Optimization Matters
Before diving into implementation, it is essential to understand why context window optimization deserves attention in your AI application architecture.
Cost Efficiency
Most LLM providers charge based on token usage. Every token in your context window costs money, whether it is a system prompt, conversation history, or retrieved documents. Inefficient context management can inflate your API bills dramatically. A well-optimized context window ensures you only pay for tokens that genuinely contribute to the model's understanding.
Response Quality
Models do not treat all tokens in their context window equally. Research shows that LLMs suffer from attention dilution when context windows are filled with irrelevant or redundant information. Important instructions buried in lengthy histories can be overlooked. By optimizing context, you ensure the model focuses on what matters most.
Latency Reduction
Processing time scales with input size. A request with 30,000 tokens takes noticeably longer than one with 3,000 tokens. For real-time applications like chatbots or coding assistants, this latency difference directly impacts user experience.
Avoiding Hard Failures
Exceeding the context window is not a graceful degradation. It is a hard error. Requests get rejected, conversations break, and users encounter failures. Proactive optimization prevents these scenarios entirely.
Understanding Pydantic AI's Context Management
Pydantic AI is a framework that brings type safety and structured validation to AI agent development. Built by the team behind Pydantic, it integrates seamlessly with popular LLM providers and emphasizes correctness through Python type hints. One of its standout features is the ability to manage conversation state and context systematically.
In Pydantic AI, context management revolves around several key concepts: agent state, message history, dependency injection, and structured outputs. Each of these plays a role in how much information ends up in the context window.
Setting Up Your Environment
Let us start by installing Pydantic AI and setting up a basic project structure.
# Install Pydantic AI and supporting packages
pip install pydantic-ai
pip install pydantic-ai[openai]
pip install tiktoken
# For token counting utilities
pip install transformersCreate a new Python file and configure your environment variables:
import os
from pydantic_ai import Agent
from pydantic_ai.models.openai import OpenAIModel
# Set your API key
os.environ["OPENAI_API_KEY"] = "your-api-key-here"
# Initialize the model
model = OpenAIModel("gpt-4o")
# Create a basic agent
agent = Agent(
model=model,
system_prompt="You are a helpful coding assistant."
)Basic Context Window Management
The simplest form of context optimization is being mindful of what you include in your prompts. Pydantic AI agents maintain a conversation history automatically, but you have full control over how that history is managed.
Understanding Token Limits
Different models have different context window sizes. Here is a quick reference for common models:
- GPT-4o: 128,000 tokens
- GPT-4 Turbo: 128,000 tokens
- GPT-3.5 Turbo: 16,385 tokens
- Claude 3.5 Sonnet: 200,000 tokens
- Llama 3.1 70B: 128,000 tokens
Even with large windows, optimization remains critical because cost and latency scale with usage. Let us build a utility to count tokens before sending requests.
import tiktoken
from pydantic_ai.messages import ModelMessage
def count_tokens(messages: list[ModelMessage], model: str = "gpt-4o") -> int:
"""Count the total tokens in a list of messages."""
try:
encoding = tiktoken.encoding_for_model(model)
except KeyError:
encoding = tiktoken.get_encoding("cl100k_base")
total = 0
for message in messages:
# Each message has overhead tokens for formatting
total += 4 # message wrapper tokens
if hasattr(message, "content"):
content = message.content
if isinstance(content, str):
total += len(encoding.encode(content))
elif isinstance(content, list):
for part in content:
if isinstance(part, str):
total += len(encoding.encode(part))
return total
def get_context_usage(messages: list[ModelMessage], max_tokens: int = 128000) -> dict:
"""Get context window usage statistics."""
used = count_tokens(messages)
return {
"tokens_used": used,
"tokens_remaining": max_tokens - used,
"usage_percentage": round((used / max_tokens) * 100, 2),
"max_tokens": max_tokens
}Building a Context-Aware Agent
Now let us create an agent that tracks its own context usage and warns when approaching limits.
from pydantic_ai import Agent
from pydantic_ai.messages import ModelMessage
from dataclasses import dataclass, field
@dataclass
class ContextState:
messages: list[ModelMessage] = field(default_factory=list)
max_context_tokens: int = 128000
warning_threshold: float = 0.75
agent = Agent(
"openai:gpt-4o",
system_prompt="You are a helpful assistant that provides concise answers.",
deps_type=ContextState
)
async def chat_with_context_management(
user_input: str,
state: ContextState
) -> str:
"""Run agent with context window monitoring."""
# Check current context usage before adding new message
usage = get_context_usage(state.messages, state.max_context_tokens)
if usage["usage_percentage"] > state.warning_threshold * 100:
print(f"WARNING: Context window at {usage['usage_percentage']}% capacity")
print(f"Consider pruning or summarizing conversation history.")
# Run the agent
result = await agent.run(user_input, message_history=state.messages)
# Update state with new messages
state.messages = result.all_messages()
# Report usage after the call
new_usage = get_context_usage(state.messages, state.max_context_tokens)
print(f"Context usage: {new_usage['tokens_used']}/{new_usage['max_tokens']} tokens "
f"({new_usage['usage_percentage']}%)")
return result.dataAdvanced Techniques for Context Optimization
Message Pruning Strategies
Message pruning is the process of removing older messages from conversation history to make room for new content. The key is deciding which messages to keep and which to discard. A common approach is to always preserve the system prompt and the most recent N messages while removing older ones.
from pydantic_ai.messages import ModelMessage, SystemPromptPart, UserPromptPart, ModelResponse
def prune_messages(
messages: list[ModelMessage],
max_tokens: int = 100000,
min_messages_to_keep: int = 6
) -> list[ModelMessage]:
"""
Prune message history to stay within token limits.
Always keeps system prompts and the most recent messages.
"""
if count_tokens(messages) <= max_tokens:
return messages
# Separate system messages from conversation messages
system_messages = []
conversation_messages = []
for msg in messages:
if isinstance(msg, ModelMessage) and any(
isinstance(part, SystemPromptPart)
for part in getattr(msg, "parts", [])
):
system_messages.append(msg)
else:
conversation_messages.append(msg)
# Keep removing oldest conversation messages until under limit
while conversation_messages and (
count_tokens(system_messages + conversation_messages) > max_tokens or
len(conversation_messages) > min_messages_to_keep * 2
):
if len(conversation_messages) <= min_messages_to_keep:
break
conversation_messages.pop(0)
return system_messages + conversation_messagesConversation Summarization
Instead of simply discarding old messages, you can summarize them. This preserves important context while dramatically reducing token count. Pydantic AI makes this straightforward with its structured output capabilities.
from pydantic import BaseModel, Field
class ConversationSummary(BaseModel):
"""Structured summary of a conversation."""
key_topics: list[str] = Field(description="Main topics discussed")
decisions_made: list[str] = Field(description="Decisions or conclusions reached")
important_context: str = Field(description="Critical context to preserve")
user_preferences: list[str] = Field(default_factory=list, description="User preferences noted")
# Dedicated summarization agent
summarizer_agent = Agent(
"openai:gpt-4o",
output_type=ConversationSummary,
system_prompt=(
"Summarize the following conversation, extracting key information "
"that would be important for continuing the discussion. "
"Be concise but preserve all critical details."
)
)
async def summarize_and_compress(
messages: list[ModelMessage],
keep_recent: int = 4
) -> list[ModelMessage]:
"""
Summarize older messages and keep recent ones intact.
Returns a compressed message list.
"""
if len(messages) <= keep_recent + 2:
return messages
# Separate system messages
system_messages = [m for m in messages if any(
isinstance(part, SystemPromptPart)
for part in getattr(m, "parts", [])
)]
# Split into old messages (to summarize) and recent (to keep)
conversation = [m for m in messages if m not in system_messages]
to_summarize = conversation[:-keep_recent]
to_keep = conversation[-keep_recent:]
# Build a text representation of old messages for summarization
conversation_text = ""
for msg in to_summarize:
for part in getattr(msg, "parts", []):
if isinstance(part, (UserPromptPart,)):
conversation_text += f"User: {part.content}\n"
elif isinstance(part, ModelResponse):
conversation_text += f"Assistant: {part.content}\n"
if not conversation_text.strip():
return messages
# Generate structured summary
summary_result = await summarizer_agent.run(
f"Summarize this conversation:\n\n{conversation_text}"
)
summary = summary_result.data
# Create a compressed context message
summary_text = (
f"[Previous conversation summary]\n"
f"Key topics: {', '.join(summary.key_topics)}\n"
f"Decisions: {', '.join(summary.decisions_made)}\n"
f"Important context: {summary.important_context}\n"
)
if summary.user_preferences:
summary_text += f"User preferences: {', '.join(summary.user_preferences)}\n"
# Build new message list with summary as context
# In practice, you would construct proper ModelMessage objects here
print(f"Summarized {len(to_summarize)} messages into {count_tokens([summary_text])} tokens")
return system_messages + to_keep # In production, prepend summary as a messageToken Budgeting with Dynamic Allocation
A sophisticated approach is to allocate token budgets across different components of your context. For example, you might reserve tokens for the system prompt, conversation history, retrieved documents, and the expected output.
from dataclasses import dataclass
@dataclass
class TokenBudget:
"""Allocates token budget across context components."""
total_budget: int = 128000
# Reserve tokens for model output
output_reserve: int = 4096
@property
def system_prompt_budget(self) -> int:
return min(2000, int(self.total_budget * 0.02))
@property
def conversation_history_budget(self) -> int:
return int((self.total_budget - self.output_reserve - self.system_prompt_budget) * 0.4)
@property
def retrieved_context_budget(self) -> int:
return int((self.total_budget - self.output_reserve - self.system_prompt_budget) * 0.5)
@property
def user_input_budget(self) -> int:
return int((self.total_budget - self.output_reserve - self.system_prompt_budget) * 0.1)
def validate_allocation(self) -> bool:
"""Ensure allocations do not exceed total budget."""
allocated = (
self.system_prompt_budget +
self.conversation_history_budget +
self.retrieved_context_budget +
self.user_input_budget +
self.output_reserve
)
return allocated <= self.total_budget
# Usage
budget = TokenBudget(total_budget=128000)
print(f"System prompt budget: {budget.system_prompt_budget} tokens")
print(f"Conversation history budget: {budget.conversation_history_budget} tokens")
print(f"Retrieved context budget: {budget.retrieved_context_budget} tokens")
print(f"User input budget: {budget.user_input_budget} tokens")
print(f"Output reserve: {budget.output_reserve} tokens")
print(f"Allocation valid: {budget.validate_allocation()}")Integrating Budgeting with Retrieval-Augmented Generation
When using RAG, the retrieved documents can quickly consume your context window. Here is how to integrate token budgeting with document retrieval to ensure you never exceed limits.
from pydantic_ai import Agent
from pydantic import BaseModel
class DocumentChunk(BaseModel):
content: str
relevance_score: float
token_count: int = 0
class RAGContextManager:
"""Manages context for RAG applications with token budgeting."""
def __init__(self, budget: TokenBudget):
self.budget = budget
def select_chunks(
self,
chunks: list[DocumentChunk],
conversation_tokens: int = 0
) -> list[DocumentChunk]:
"""
Select document chunks that fit within the retrieved context budget.
Prioritizes by relevance score.
"""
available = self.budget.retrieved_context_budget - conversation_tokens
if available <= 0:
return []
# Sort by relevance score (descending)
sorted_chunks = sorted(chunks, key=lambda c: c.relevance_score, reverse=True)
selected = []
used_tokens = 0
for chunk in sorted_chunks:
if used_tokens + chunk.token_count <= available:
selected.append(chunk)
used_tokens += chunk.token_count
else:
# Try to fit a truncated version
remaining = available - used_tokens
if remaining > 100: # Only include if meaningful space remains
# Truncate content to fit
words = chunk.content.split()
truncated = " ".join(words[:remaining // 2])
selected.append(DocumentChunk(
content=truncated + "...[truncated]",
relevance_score=chunk.relevance_score,
token_count=remaining
))
used_tokens += remaining
break
return selected
def build_context_string(self, chunks: list[DocumentChunk]) -> str:
"""Build a context string from selected chunks."""
if not chunks:
return ""
parts = ["[Retrieved Context]"]
for i, chunk in enumerate(chunks, 1):
parts.append(f"\n--- Document {i} (relevance: {chunk.relevance_score:.2f}) ---")
parts.append(chunk.content)
return "\n".join(parts)
# Example usage
budget = TokenBudget(total_budget=128000)
rag_manager = RAGContextManager(budget)
# Simulate retrieved chunks
chunks = [
DocumentChunk(content="Python is a high-level programming language...", relevance_score=0.95, token_count=500),
DocumentChunk(content="Pydantic AI provides typed AI agent framework...", relevance_score=0.89, token_count=300),
DocumentChunk(content="Context windows limit how much text...", relevance_score=0.82, token_count=450),
]
selected = rag_manager.select_chunks(chunks, conversation_tokens=5000)
context_string = rag_manager.build_context_string(selected)
print(f"Selected {len(selected)} chunks")
print(f"Context string length: {len(context_string)} characters")Building a Complete Optimized Agent
Let us now combine all the techniques into a single, production-ready agent that handles context optimization automatically.
import asyncio
from pydantic_ai import Agent
from pydantic_ai.messages import ModelMessage, SystemPromptPart, UserPromptPart
from dataclasses import dataclass, field
from typing import Optional
@dataclass
class OptimizedAgentConfig:
model_name: str = "openai:gpt-4o"
max_context_tokens: int = 120000 # Leave buffer below model limit
max_conversation_messages: int = 20
summarization_threshold: int = 80000
warning_threshold: float = 0.80
system_prompt: str = "You are a helpful, concise assistant."
@dataclass
class AgentState:
messages: list[ModelMessage] = field(default_factory=list)
total_tokens_used: int = 0
summarization_count: int = 0
class OptimizedAgent:
"""A production-ready agent with automatic context window optimization."""
def __init__(self, config: OptimizedAgentConfig):
self.config = config
self.agent = Agent(
config.model_name,
system_prompt=config.system_prompt
)
self.summarizer = Agent(
config.model_name,
system_prompt=(
"Create a concise summary of the conversation that preserves "
"all important context, decisions, and user preferences. "
"Keep it under 500 words."
)
)
self.state = AgentState()
async def chat(self, user_input: str) -> str:
"""Process user input with automatic context optimization."""
# Step 1: Check if we need to optimize before processing
current_tokens = count_tokens(self.state.messages)
if current_tokens > self.config.summarization_threshold:
await self._compress_history()
# Step 2: Prune if still over limit
self.state.messages = prune_messages(
self.state.messages,
max_tokens=self.config.max_context_tokens,
min_messages_to_keep=6
)
# Step 3: Check final context usage
usage = get_context_usage(
self.state.messages,
self.config.max_context_tokens
)
if usage["usage_percentage"] > self.config.warning_threshold * 100:
print(f"[OPTIMIZATION] Context at {usage['usage_percentage']}% - pruning applied")
# Step 4: Run the agent
try:
result = await self.agent.run(
user_input,
message_history=self.state.messages
)
# Update state
self.state.messages = result.all_messages()
self.state.total_tokens_used = count_tokens(self.state.messages)
return result.data
except Exception as e:
# Handle context overflow errors gracefully
if "context_length" in str(e).lower() or "maximum context" in str(e).lower():
print("[ERROR] Context overflow detected. Forcing aggressive pruning.")
self.state.messages = prune_messages(
self.state.messages,
max_tokens=self.config.max_context_tokens // 2,
min_messages_to_keep=4
)
# Retry with reduced context
result = await self.agent.run(
user_input,
message_history=self.state.messages
)
self.state.messages = result.all_messages()
return result.data
raise
async def _compress_history(self) -> None:
"""Summarize older conversation history."""
if len(self.state.messages) <= 8:
return
# Keep system messages and last 4 conversation messages
system_msgs = []
conversation = []
for msg in self.state.messages:
if any(isinstance(p, SystemPromptPart) for p in getattr(msg, "parts", [])):
system_msgs.append(msg)
else:
conversation.append(msg)
if len(conversation) <= 4:
return
to_summarize = conversation[:-4]
to_keep = conversation[-4:]
# Build conversation text
conv_text = ""
for msg in to_summarize:
for part in getattr(msg, "parts", []):
if isinstance(part, UserPromptPart):
conv_text += f"User: {part.content}\n"
if not conv_text.strip():
return
# Generate summary
summary_result = await self.summarizer.run(conv_text)
summary = summary_result.data
# Log optimization
old_tokens = count_tokens(to_summarize)
new_tokens = count_tokens([summary])
print(f"[OPTIMIZATION] Summarized {old_tokens} tokens -> {new_tokens} tokens")
self.state.summarization_count += 1
# In production, construct a proper ModelMessage with the summary
self.state.messages = system_msgs + to_keep
def get_stats(self) -> dict:
"""Return current agent statistics."""
usage = get_context_usage(
self.state.messages,
self.config.max_context_tokens
)
return {
"messages_count": len(self.state.messages),
"tokens_used": usage["tokens_used"],
"usage_percentage": usage["usage_percentage"],
"summarizations_performed": self.state.summarization_count,
"tokens_remaining": usage["tokens_remaining"]
}
# Example usage
async def main():
config = OptimizedAgentConfig(
model_name="openai:gpt-4o",
max_context_tokens=120000,
summarization_threshold=80000,
system_prompt="You are a helpful coding assistant. Be concise."
)
agent = OptimizedAgent(config)
# Simulate a conversation
questions = [
"What is Python?",
"How do I create a class in Python?",
"Can you show me inheritance?",
"What about async/await?",
"How does error handling work?",
]
for question in questions:
print(f"\nUser: {question}")
response = await agent.chat(question)
print(f"Assistant: {response}")
print(f"Stats: {agent.get_stats()}")
if __name__ == "__main__":
asyncio.run(main())Best Practices for Context Window Optimization
Based on the techniques covered, here are the key best practices to follow when building Pydantic AI applications.
Design Concise System Prompts
Your system prompt is included in every single request. Keep it focused and eliminate redundancy. Use structured formatting like bullet points instead of verbose paragraphs. Every token saved here compounds across thousands of requests.
Implement Progressive Compression
Do not wait until you hit the context limit to act. Implement a multi-stage approach: warn at 70% capacity, summarize at 80%, and aggressively prune at 90%. This prevents hard failures and maintains conversation quality.
Use Structured Outputs to Reduce Tokens
Pydantic AI's structured output feature ensures model responses follow a defined schema. This can reduce token usage by eliminating verbose formatting and ensuring responses are concise and predictable.
from pydantic import BaseModel, Field
class ConciseAnswer(BaseModel):
answer: str = Field(description="Direct answer, max 2 sentences")
confidence: float = Field(ge=0, le=1, description="Confidence score")
sources: list[str] = Field(default_factory=list, max_length=3)
concise_agent = Agent(
"openai:gpt-4o",
output_type=ConciseAnswer,
system_prompt="Answer questions concisely and accurately."
)Cache and Reuse Computed Context
If your agent uses static reference material, system prompts, or tool definitions, cache these rather than recomputing them. Pydantic AI agents cache system prompts by default, but be mindful of dynamic prompt construction.
Monitor and Log Token Usage
Always track token usage in production. This helps identify inefficiencies and unexpected cost spikes. Build dashboards that show average context size, summarization frequency, and cost per conversation.
Choose the Right Model for the Task
Not every request needs a 128K context window. For simple tasks, use models with smaller, cheaper context windows. Reserve large-context models for tasks that genuinely require processing extensive information.
Handle Edge Cases Gracefully
Always implement fallback behavior for context overflow errors. Your application should degrade gracefully, perhaps by summarizing more aggressively or informing the user that the conversation is getting long, rather than crashing.
Conclusion
Context window optimization is not an optional refinement but a fundamental requirement for building production-grade AI applications with Pydantic AI. By implementing the techniques covered in this guide, including token counting, message pruning, conversation summarization, and token budgeting, you can build agents that are cost-efficient, responsive, and reliable. The key is to treat the context window as a finite resource that must be managed proactively rather than reactively. Start with the basic monitoring utilities, progressively add compression strategies as your conversations grow longer, and always measure the impact of your optimizations on both cost and response quality. With Pydantic AI's typed framework and the patterns demonstrated here, you have everything needed to build AI agents that scale gracefully while keeping token usage and costs under control.