@bastani/atomic 0.5.11 → 0.5.12-0

package/.agents/skills/context-optimization/scripts/compaction.py

@@ -0,0 +1,562 @@
+"""
+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 ===")