""" Context Optimization Utilities — compaction, masking, budgeting, and cache optimization. Public API ---------- Functions: estimate_token_count(text) -> int estimate_message_tokens(messages) -> int categorize_messages(messages) -> dict summarize_content(content, category, max_length) -> str design_stable_prompt(template, dynamic_values) -> str calculate_cache_metrics(requests, cache) -> dict Classes: ObservationStore — Store and mask verbose tool outputs with retrievable references. ContextBudget — Token budget allocation and optimization trigger detection. PRODUCTION NOTES: - Token estimation uses simplified heuristics (~4 chars/token for English). Production systems should use model-specific tokenizers: - OpenAI: tiktoken library - Anthropic: anthropic tokenizer - Local models: HuggingFace tokenizers - Summarization functions use simple heuristics for demonstration. Production systems should use: - LLM-based summarization for high-quality compression - Domain-specific summarization models - Schema-based summarization for structured outputs - Cache metrics are illustrative. Production systems should integrate with actual inference infrastructure metrics. """ from typing import List, Dict, Optional, Tuple import hashlib import re import time __all__ = [ "estimate_token_count", "estimate_message_tokens", "categorize_messages", "summarize_content", "summarize_tool_output", "summarize_conversation", "summarize_document", "summarize_general", "ObservationStore", "ContextBudget", "design_stable_prompt", "calculate_cache_metrics", "generate_cache_recommendations", ] # --------------------------------------------------------------------------- # Token estimation # --------------------------------------------------------------------------- def estimate_token_count(text: str) -> int: """ Estimate token count for text. Use when: a quick token budget check is needed and a model-specific tokenizer is unavailable or too slow for the hot path. Uses approximation: ~4 characters per token for English. WARNING: This is a rough estimate. Actual tokenization varies by: - Model (GPT-5.2, Claude 4.5, Gemini 3 have different tokenizers) - Content type (code typically has higher token density) - Language (non-English may have 2-3x higher token/char ratio) Production usage:: import tiktoken enc = tiktoken.encoding_for_model("gpt-4") token_count = len(enc.encode(text)) """ return len(text) // 4 def estimate_message_tokens(messages: List[Dict[str, str]]) -> int: """ Estimate token count for a message list. Use when: checking whether the current conversation is approaching the context budget threshold before deciding to compact or mask. """ total = 0 for msg in messages: content = msg.get("content", "") total += estimate_token_count(content) # Add overhead for role/formatting total += 10 return total # --------------------------------------------------------------------------- # Compaction functions # --------------------------------------------------------------------------- def categorize_messages(messages: List[Dict]) -> Dict[str, List[Dict]]: """ Categorize messages for selective compaction. Use when: preparing to compact context and needing to apply different summarization strategies per category (tool outputs first, then old conversation turns, then retrieved documents — never the system prompt). Returns a dict mapping category name to list of messages. """ categories: Dict[str, List[Dict]] = { "system_prompt": [], "tool_definition": [], "tool_output": [], "conversation": [], "retrieved_document": [], "other": [], } for msg in messages: role = msg.get("role", "user") content = msg.get("content", "") if role == "system": categories["system_prompt"].append({**msg, "category": "system_prompt"}) elif "tool_use" in msg.get("type", ""): categories["tool_output"].append({**msg, "category": "tool_output"}) elif role == "user": categories["conversation"].append({**msg, "category": "conversation"}) elif "retrieved" in msg.get("tags", []): categories["retrieved_document"].append({**msg, "category": "retrieved_document"}) else: categories["other"].append({**msg, "category": "other"}) return categories def summarize_content(content: str, category: str, max_length: int = 500) -> str: """ Summarize content for compaction, dispatching by category. Use when: compacting context and needing category-aware summarization (tool outputs get metric extraction, conversations get decision extraction, documents get lead-paragraph extraction). """ if category == "tool_output": return summarize_tool_output(content, max_length) elif category == "conversation": return summarize_conversation(content, max_length) elif category == "retrieved_document": return summarize_document(content, max_length) else: return summarize_general(content, max_length) def summarize_tool_output(content: str, max_length: int = 500) -> str: """ Summarize tool output by extracting metrics and key findings. Use when: a tool output has served its immediate purpose and needs to be compacted while preserving actionable data points. """ # Look for metrics (numbers with context) metrics = re.findall(r'(\w+):\s*([\d.,]+)', content) # Look for key findings (lines with important keywords) keywords = ["result", "found", "total", "success", "error", "value"] findings = [] for line in content.split('\n'): if any(kw in line.lower() for kw in keywords): findings.append(line.strip()) summary_parts = [] if metrics: summary_parts.append(f"Metrics: {', '.join([f'{k}={v}' for k, v in metrics])}") if findings: summary_parts.append("Key findings: " + "; ".join(findings[:3])) result = " | ".join(summary_parts) if summary_parts else "[Tool output summarized]" return result[:max_length] def summarize_conversation(content: str, max_length: int = 500) -> str: """ Summarize conversational content by extracting decisions and questions. Use when: older conversation turns need compaction and the key decisions/commitments must survive while filler is removed. """ decisions = re.findall(r'(?i)(?:decided|decision|chose|chosen)[:\s]+([^.]+)', content) questions = re.findall(r'(?:\?|question)[:\s]+([^.]+)', content) summary_parts = [] if decisions: decision_texts = [d.strip() for d in decisions[:5]] summary_parts.append(f"Decisions: {'; '.join(decision_texts)}") if questions: question_texts = [q.strip() for q in questions[:3]] summary_parts.append(f"Open questions: {'; '.join(question_texts)}") if not summary_parts: # Fallback: extract the first few substantive sentences sentences = [s.strip() for s in content.split('.') if len(s.strip()) > 20] if sentences: summary_parts.append('. '.join(sentences[:3]) + '.') result = " | ".join(summary_parts) if summary_parts else "[Conversation summarized]" return result[:max_length] def summarize_document(content: str, max_length: int = 500) -> str: """ Summarize document content using lead-paragraph extraction. Use when: a retrieved document has been consumed for reasoning and only a brief reference needs to remain in context. """ paragraphs = content.split('\n\n') if paragraphs: first_para = paragraphs[0].strip() sentences = first_para.split('. ') if len(sentences) > 2: first_para = '. '.join(sentences[:2]) + '.' return first_para[:max_length] return "[Document summarized]" def summarize_general(content: str, max_length: int = 500) -> str: """ General-purpose summarization via truncation. Use when: content does not fit a specific category and a simple truncation with ellipsis is acceptable. """ return content[:max_length] + "..." if len(content) > max_length else content # --------------------------------------------------------------------------- # Observation masking # --------------------------------------------------------------------------- class ObservationStore: """ Store and mask verbose tool outputs with retrievable references. Use when: tool outputs dominate context (>50% of tokens) and older observations have already served their reasoning purpose. Stores the full content externally and replaces it with a compact reference containing a key-point summary. Example:: store = ObservationStore(max_size=500) masked, ref_id = store.mask(long_tool_output, max_length=200) # masked: "[Obs:a1b2c3d4 elided. Key: ... Full content retrievable.]" # Later retrieval: original = store.retrieve(ref_id) """ def __init__(self, max_size: int = 1000) -> None: self.observations: Dict[str, Dict] = {} self.order: List[str] = [] self.max_size = max_size def store(self, content: str, metadata: Optional[Dict] = None) -> str: """Store observation and return reference ID.""" ref_id = self._generate_ref_id(content) self.observations[ref_id] = { "content": content, "metadata": metadata or {}, "stored_at": time.time(), "last_accessed": time.time(), } self.order.append(ref_id) # Evict oldest if over limit if len(self.order) > self.max_size: oldest = self.order.pop(0) del self.observations[oldest] return ref_id def retrieve(self, ref_id: str) -> Optional[str]: """Retrieve observation by reference ID.""" if ref_id in self.observations: self.observations[ref_id]["last_accessed"] = time.time() return self.observations[ref_id]["content"] return None def mask(self, content: str, max_length: int = 200) -> Tuple[str, Optional[str]]: """ Mask observation if longer than max_length. Use when: deciding per-observation whether to keep inline or replace with a compact reference. Returns (masked_content, ref_id) where ref_id is None if the content was short enough to keep. """ if len(content) <= max_length: return content, None ref_id = self.store(content) key_point = self._extract_key_point(content) masked = f"[Obs:{ref_id} elided. Key: {key_point}. Full content retrievable.]" return masked, ref_id def _generate_ref_id(self, content: str) -> str: """Generate unique reference ID.""" hash_input = f"{content[:100]}{time.time()}" return hashlib.md5(hash_input.encode()).hexdigest()[:8] def _extract_key_point(self, content: str) -> str: """Extract key point from observation.""" lines = [line for line in content.split('\n') if len(line) > 20] if lines: return lines[0][:50] + "..." sentences = content.split('. ') if sentences: return sentences[0][:50] + "..." return content[:50] + "..." # --------------------------------------------------------------------------- # Context budget management # --------------------------------------------------------------------------- class ContextBudget: """ Token budget allocation and optimization trigger detection. Use when: building an agent loop that needs to monitor context usage across categories and trigger compaction/masking at the right thresholds rather than waiting until the window overflows. Example:: budget = ContextBudget(total_limit=128_000) budget.allocate("system_prompt", 1500) budget.allocate("tool_definitions", 3000) # ... after each agent turn: should_act, reasons = budget.should_optimize(current_usage) if should_act: # apply masking or compaction based on reasons pass """ def __init__(self, total_limit: int) -> None: self.total_limit = total_limit self.allocated: Dict[str, int] = { "system_prompt": 0, "tool_definitions": 0, "retrieved_docs": 0, "message_history": 0, "tool_outputs": 0, "other": 0, } self.reserved = 5000 # Reserved buffer self.reservation_limit = total_limit - self.reserved def allocate(self, category: str, amount: int) -> bool: """ Allocate budget to category. Returns True on success, False if the allocation would exceed the reservation limit. """ if category not in self.allocated: category = "other" current = sum(self.allocated.values()) proposed = current + amount if proposed > self.reservation_limit: return False self.allocated[category] += amount return True def remaining(self) -> int: """Get remaining unallocated budget.""" current = sum(self.allocated.values()) return self.reservation_limit - current def get_usage(self) -> Dict[str, object]: """ Get current usage breakdown. Use when: logging or displaying context budget state for monitoring dashboards or debug output. """ total = sum(self.allocated.values()) return { "total_used": total, "total_limit": self.total_limit, "remaining": self.remaining(), "by_category": dict(self.allocated), "utilization_ratio": total / self.total_limit, } def should_optimize( self, current_usage: int, metrics: Optional[Dict[str, float]] = None ) -> Tuple[bool, List[Tuple[str, object]]]: """ Determine if optimization should trigger. Use when: called at the end of each agent loop iteration to decide whether to apply compaction, masking, or both before the next model call. Returns (should_optimize, list_of_reasons). """ reasons: List[Tuple[str, object]] = [] # Check utilization utilization = current_usage / self.total_limit if utilization > 0.8: reasons.append(("high_utilization", utilization)) # Check degradation metrics if provided if metrics: if metrics.get("attention_degradation", 0) > 0.3: reasons.append(("attention_degradation", True)) if metrics.get("quality_score", 1.0) < 0.8: reasons.append(("quality_degradation", True)) return len(reasons) > 0, reasons # --------------------------------------------------------------------------- # Cache optimization # --------------------------------------------------------------------------- def design_stable_prompt(template: str, dynamic_values: Optional[Dict] = None) -> str: """ Stabilize a prompt template for maximum KV-cache hit rate. Use when: constructing system prompts or few-shot prefixes that will be reused across many requests. Replaces dynamic content (timestamps, session IDs, counters) with stable placeholders so the prefix hash remains constant. """ result = template # Replace timestamps date_pattern = r'\d{4}-\d{2}-\d{2}' result = re.sub(date_pattern, '[DATE_STABLE]', result) # Replace session IDs session_pattern = r'Session \d+' result = re.sub(session_pattern, 'Session [STABLE]', result) # Replace counters counter_pattern = r'\d+/\d+' result = re.sub(counter_pattern, '[COUNTER_STABLE]', result) return result def calculate_cache_metrics( requests: List[Dict], cache: Dict[str, Dict] ) -> Dict[str, object]: """ Calculate KV-cache hit metrics for a request sequence. Use when: evaluating whether prompt restructuring improved cache utilization. Feed in the request log and current cache state to get hit/miss rates and actionable recommendations. """ hits = 0 misses = 0 for req in requests: prefix = req.get("prefix_hash", "") token_count = req.get("token_count", 0) if prefix in cache: hits += token_count * cache[prefix].get("hit_ratio", 0) else: misses += token_count total = hits + misses return { "hit_rate": hits / total if total > 0 else 0, "cache_hits": hits, "cache_misses": misses, "recommendations": generate_cache_recommendations(hits, misses), } def generate_cache_recommendations(hits: int, misses: int) -> List[str]: """ Generate recommendations for cache optimization based on hit/miss ratio. Use when: cache metrics indicate sub-optimal hit rates and concrete next steps are needed. """ recommendations: List[str] = [] hit_rate = hits / (hits + misses) if (hits + misses) > 0 else 0 if hit_rate < 0.5: recommendations.append("Consider stabilizing system prompts") recommendations.append("Reduce variation in request prefixes") if hit_rate < 0.8: recommendations.append("Group similar requests together") recommendations.append("Use consistent formatting across requests") return recommendations # --------------------------------------------------------------------------- # Demo / smoke test # --------------------------------------------------------------------------- if __name__ == "__main__": print("=== Context Optimization Utilities — Demo ===\n") # 1. Token estimation sample_text = "The quick brown fox jumps over the lazy dog. " * 20 tokens = estimate_token_count(sample_text) print(f"1. Token estimate for {len(sample_text)}-char text: ~{tokens} tokens\n") # 2. Observation masking store = ObservationStore(max_size=100) long_output = ( "Result: 42 items found\n" "Total processing time: 3.2s\n" "Details:\n" + "\n".join([f" Item {i}: value={i*10}" for i in range(20)]) ) masked, ref_id = store.mask(long_output, max_length=100) print(f"2. Masked observation:\n {masked}") print(f" Ref ID: {ref_id}") retrieved = store.retrieve(ref_id) print(f" Retrievable: {retrieved is not None}\n") # 3. Context budget budget = ContextBudget(total_limit=128_000) budget.allocate("system_prompt", 1500) budget.allocate("tool_definitions", 3000) budget.allocate("message_history", 95_000) usage = budget.get_usage() print(f"3. Budget utilization: {usage['utilization_ratio']:.1%}") should_opt, reasons = budget.should_optimize( current_usage=int(128_000 * 0.85) ) print(f" Should optimize: {should_opt}, reasons: {reasons}\n") # 4. Cache-stable prompt raw_prompt = "Session 42 started on 2025-12-20. Progress: 3/10 tasks." stable = design_stable_prompt(raw_prompt) print(f"4. Original prompt: {raw_prompt}") print(f" Stabilized: {stable}\n") # 5. Summarization tool_out = "count: 150\nstatus: success\nFound 3 errors in module A." summary = summarize_content(tool_out, "tool_output", max_length=200) print(f"5. Tool output summary: {summary}\n") print("=== Demo complete ===")