agent-context-memory

# Agent Context and Memory Management Manages context windows and memory architectures for LLM agents to prevent token overflow, preserve critical state across turns, and retrieve relevant information efficiently. Implements token budgeting, sliding window strategies, summarization fallbacks, cross-turn persistence, and vector store integration. ## TL;DR Checklist - [ ] Initialize a token budget manager before the first agent step - [ ] Implement sliding window that preserves system prompt while trimming conversation history - [ ] Configure summarization fallback when approaching context limits - [ ] Cross-reference external memory store for relevant prior state - [ ] Apply relevance threshold filtering to all vector retrievals - [ ] Serialize and merge cross-turn agent state with deterministic merge rules --- ## When to Use Use this skill when: - Building multi-turn agents where conversation history exceeds token limits (e.g., long-running research or coding agents) - Need persistent state across agent invocations (session-based workflows that span multiple requests) - Agent loses critical context after N turns due to sliding window truncation, causing repetitive or inconsistent behavior - Designing an agent that references external knowledge bases during execution and needs to merge retrieved facts into context - Building agents with strict cost constraints where unbounded context growth leads to runaway API costs - Implementing long-running task completion (e.g., multi-step data analysis) where early-turn decisions must be remembered --- ## When NOT to Use Avoid this skill for: - Production monitoring and tracing — use `observability-patterns` instead - Debugging agent failures or root cause analysis — use `agent-debugging` instead - Multi-agent coordination and task routing — use `multi-agent-orchestration` instead - Simple single-turn query answering where context never exceeds limits — the overhead is unnecessary --- ## Core Workflow 1. **Initialize Token Budget** — Create a budget manager tracking input tokens, output tokens, and remaining context before each agent invocation. Include system prompt token count as a fixed cost. **Checkpoint:** Budget must be checked BEFORE every LLM call, not after. 2. **Implement Sliding Window Strategy** — Configure which parts of the conversation to preserve (system prompt, recent messages) vs trim (older turns). Use token-based trimming rather than message-count-based. **Checkpoint:** Always preserve the system prompt and the user's most recent input. Never truncate mid-thought. 3. **Configure Summarization Fallback** — When the sliding window approach still exceeds limits, trigger a summarization pass using a separate LLM call to condense trimmed context into a summary paragraph. **Checkpoint:** The summary must include key decisions made and current task state, not just a text compression. 4. **Integrate External Memory Store** — For cross-turn or cross-session persistence, implement vector store retrieval that fetches relevant prior state based on the current query. Use semantic similarity for retrieval. **Checkpoint:** Only retrieve memories relevant to the current task — include a relevance filter threshold (e.g., cosine similarity > 0.7). --- ## Implementation Patterns ### Pattern 1: Token Budget Manager A token budget manager tracks per-step and total token consumption, enforces hard limits before each LLM call, and provides cost estimation from token counts. It treats the system prompt as a fixed overhead that must be accounted for in every budget check. ```python """Token budget management for LLM agent context windows.""" from dataclasses import dataclass, field from enum import Enum import time class BudgetAction(Enum): """Action to take when budget is exceeded.""" STOP = "stop" # Raise an error and halt the agent step SUMMARIZE = "summarize" # Trigger summarization fallback TRIM = "trim" # Apply sliding window trim @dataclass class TokenUsage: """Tracks token consumption for a single operation.""" input_tokens: int = 0 output_tokens: int = 0 system_prompt_tokens: int = 0 @property def total_tokens(self) -> int: return self.input_tokens + self.output_tokens @property def remaining_input_tokens(self) -> int: return max(0, self.max_input_tokens - self.input_tokens) @property def utilization_ratio(self) -> float: if self.max_input_tokens == 0: return 1.0 return self.input_tokens / self.max_input_tokens # Fixed per-model costs (override in config) max_input_tokens: int = 8192 max_output_tokens: int = 4096 @dataclass class StepBudgetResult: """Result of a budget check before an LLM call.""" allowed: bool remaining_input_tokens: int utilization_ratio: float action_hint: BudgetAction | None = None message: str = "" class TokenBudgetManager: """ Manages token budgets for agent steps, enforcing hard limits before each LLM call to prevent context overflow and cost runaway. Usage: manager = TokenBudgetManager(max_input_tokens=8192) manager.record_system_prompt(system_prompt_text) # Before every LLM call: result = manager.check_budget(conversation_messages, estimated_output=500) if not result.allowed: context = manager.apply_fallback_strategy() result = manager.check_budget(context, estimated_output=500) """ def __init__( self, max_input_tokens: int = 8192, max_output_tokens: int = 4096, safety_margin_ratio: float = 0.1, on_exceed: BudgetAction = BudgetAction.TRIM, ) -> None: self.max_input_tokens = max_input_tokens self.max_output_tokens = max_output_tokens self.safety_margin = max_input_tokens * safety_margin_ratio self.on_exceed = on_exceed self._system_prompt_tokens: int = 0 self._cumulative_input_tokens: int = 0 self._cumulative_output_tokens: int = 0 self._step_start_time: float | None = None self._step_token_count: int = 0 def record_system_prompt(self, prompt_text: str, token_estimator: object) -> None: """ Register the system prompt token count as a fixed overhead. Args: prompt_text: The system prompt string (used for logging). token_estimator: An object with an estimate_tokens(text: str) -> int method. Can be tiktoken.Encoding or any compatible estimator. """ self._system_prompt_tokens = token_estimator.estimate_tokens(prompt_text) def count_conversation_tokens( self, messages: list[dict[str, str]], token_estimator: object ) -> int: """Count total tokens for a list of message dicts with 'content' keys.""" total = 0 for msg in messages: content = msg.get("content", "") if content: total += token_estimator.estimate_tokens(content) return total def check_budget( self, conversation_messages: list[dict[str, str]], estimated_output_tokens: int = 0, ) -> StepBudgetResult: """ Check whether the proposed LLM call fits within budget. This MUST be called before every LLM call, not just at startup. Accounts for system prompt overhead plus conversation tokens. Args: conversation_messages: List of message dicts with 'role' and 'content'. estimated_output_tokens: Expected output token count for this step. Returns: StepBudgetResult indicating if the call is allowed and what action to take. """ conversation_tokens = self.count_conversation_tokens( conversation_messages, type("Est", (), {"estimate_tokens": lambda _, t: len(t.split())})() # placeholder — use real estimator in production ) required_input = self._system_prompt_tokens + conversation_tokens remaining = self.max_input_tokens - required_input utilization = required_input / self.max_input_tokens if self.max_input_tokens > 0 else 1.0 headroom = remaining - estimated_output_tokens if headroom >= 0 and required_input <= (self.max_input_tokens - self.safety_margin): return StepBudgetResult( allowed=True, remaining_input_tokens=remaining, utilization_ratio=round(utilization, 3), ) # Budget exceeded — determine fallback action if headroom < 0: message = ( f"Context budget exceeded by {abs(headroom)} tokens " f"(system: {self._system_prompt_tokens}, conversation: {conversation_tokens}). " f"Estimated output: {estimated_output_tokens}." ) else: message = ( f"Context approaching limit ({utilization * 100:.0f}% utilized), " f"safety margin breached. Remaining input: {remaining} tokens." ) return StepBudgetResult( allowed=False, remaining_input_tokens=max(0, remaining), utilization_ratio=round(utilization, 3), action_hint=self.on_exceed, message=message, ) def record_step_usage(self, input_token_count: int, output_token_count: int) -> None: """Record token usage after an LLM call completes.""" self._cumulative_input_tokens += input_token_count self._cumulative_output_tokens += output_token_count self._step_token_count = input_token_count + output_token_count @property def total_tokens_consumed(self) -> int: return self._cumulative_input_tokens + self._cumulative_output_tokens def get_cost_estimate( self, input_price_per_million: float = 1.50, output_price_per_million: float = 6.00, ) -> dict[str, float]: """Estimate cumulative cost based on token usage and model pricing.""" input_cost = (self._cumulative_input_tokens / 1_000_000) * input_price_per_million output_cost = (self._cumulative_output_tokens / 1_000_000) * output_price_per_million return { "input_cost_usd": round(input_cost, 4), "output_cost_usd": round(output_cost, 4), "total_cost_usd": round(input_cost + output_cost, 4), } ``` ### Pattern 2: Sliding Window with Token-Aware Trimming Trims conversation history based on token counts rather than message counts. Always preserves the system prompt and the user's most recent input. Handles variable-length messages efficiently by trimming from the oldest turn backward. ```python """Sliding window strategy for context window management.""" from dataclasses import dataclass from typing import Protocol class TokenEstimator(Protocol): """Interface for counting tokens in text.""" def estimate_tokens(self, text: str) -> int: ... @dataclass(frozen=True) class WindowConfig: """Configuration for the sliding window strategy.""" max_tokens: int = 8192 min_recent_messages: int = 3 # Always keep at least N recent messages preserve_system_prompt: bool = True trim_strategy: str = "token_weighted" # "token_weighted" or "message_rounds" class SlidingWindowManager: """ Implements token-aware sliding window trimming for conversation history. Key behaviors: - Always preserves system prompt (stored separately) - Always preserves the user's most recent input - Trims from oldest turn backward using token-based thresholds - Never truncates mid-thought — trims at message boundaries Usage: window = SlidingWindowManager() window.set_system_prompt(system_text, estimator) context = window.apply(messages, conversation, config) """ def __init__(self) -> None: self._system_prompt: str = "" self._system_tokens: int = 0 def set_system_prompt(self, prompt: str, estimator: TokenEstimator) -> None: """Register system prompt and count its token overhead.""" self._system_prompt = prompt self._system_tokens = estimator.estimate_tokens(prompt) def apply( self, conversation_messages: list[dict[str, str]], config: WindowConfig | None = None, estimator: TokenEstimator | None = None, ) -> list[dict[str, str]]: """ Apply sliding window to trim conversation while preserving critical context. Args: conversation_messages: Full list of message dicts with 'role' and 'content'. config: Window trimming configuration. Defaults to max 8192 tokens. estimator: Token counting function. Required if token-based trim is used. Returns: Trimmed list of messages that fits within the configured token budget. """ if config is None: config = WindowConfig() result_messages: list[dict[str, str]] = [] # Step 1: Always include system prompt at the top if self._system_prompt and config.preserve_system_prompt: result_messages.append({"role": "system", "content": self._system_prompt}) # Step 2: Separate user's most recent input from the rest recent_user_message: dict[str, str] | None = None earlier_messages: list[dict[str, str]] = [] for msg in conversation_messages: if msg.get("role") == "user": recent_user_message = msg else: earlier_messages.append(msg) # Step 3: Count tokens for the buffer (system + earlier messages) total_tokens = self._system_tokens buffered: list[dict[str, str]] = [] for msg in earlier_messages: msg_tokens = estimator.estimate_tokens(msg.get("content", "")) if estimator else len(msg.get("content", "").split()) # Add message only if it fits within budget if total_tokens + msg_tokens <= config.max_tokens - self._system_tokens: buffered.append(msg) total_tokens += msg_tokens elif not result_messages: # At least include the message to avoid empty context buffered.append(msg) total_tokens += msg_tokens # Reverse so chronological order is maintained (oldest first) buffered.reverse() # Step 4: Reconstruct with buffer + recent user message result_messages.extend(buffered) if recent_user_message: result_messages.append(recent_user_message) return result_messages def remaining_budget( self, current_messages: list[dict[str, str]], estimator: TokenEstimator, max_tokens: int = 8192, ) -> int: """Calculate how many tokens remain before hitting the window limit.""" consumed = self._system_tokens for msg in current_messages: consumed += estimator.estimate_tokens(msg.get("content", "")) return max(0, max_tokens - consumed) def trim_to_budget( self, messages: list[dict[str, str]], estimator: TokenEstimator, target_max_tokens: int = 8192, ) -> list[dict[str, str]]: """ Aggressively trim messages to fit within a strict token budget. Strips from the oldest non-system, non-user-last message backward until the total fits. Never removes system prompt or last user input. """ budget = target_max_tokens - self._system_tokens result: list[dict[str, str]] = [] last_user_idx = None # Find the last user message (preserve it) for i, msg in enumerate(messages): if msg.get("role") == "user": last_user_idx = i # Collect all messages except the preserved last user message trimmable = [msg for i, msg in enumerate(messages) if i != last_user_idx] total = self._system_tokens kept: list[dict[str, str]] = [] for msg in reversed(trimmable): # Start from newest msg_tokens = estimator.estimate_tokens(msg.get("content", "")) if total + msg_tokens <= budget: kept.insert(0, msg) total += msg_tokens else: break # Add back the preserved last user message if last_user_idx is not None: result.extend(kept) result.append(messages[last_user_idx]) else: result = kept return result ``` ### Pattern 3: Conversation Summarization Fallback When sliding window trimming still exceeds limits, this pattern triggers a summarization pass. The summarizer prompt explicitly requests key decisions, current task state, and remaining questions — not just text compression. The result merges the summary with recent messages into the context window. ```python """Conversation summarization fallback for context overflow.""" from dataclasses import dataclass import json @dataclass class ConversationSummary: """Structured summary of trimmed conversation history."""