Upload aco/context_budgeter.py

aco/context_budgeter.py

@@ -0,0 +1,209 @@
+"""Context Budgeter - Module 4.
+
+Decides what context is needed, what can be omitted/summarized, and what should be retrieved.
+
+Context sources:
+- system rules
+- tool descriptions
+- user preferences
+- project memory
+- retrieved docs
+- prior trace failures
+- examples
+- recent messages
+- artifacts
+- task plan
+"""
+
+from typing import Dict, List, Tuple, Optional, Any
+from dataclasses import dataclass, field
+
+from .trace_schema import TaskType
+from .config import ACOConfig
+
+
+@dataclass
+class ContextSource:
+    name: str
+    tokens: int
+    importance: float  # 0-1
+    staleness: float  # 0=current, 1=very stale
+    mutable: bool  # True if content changes per turn
+    cacheable: bool  # True if can be prefix-cached
+    summary: Optional[str] = None
+
+
+@dataclass
+class ContextBudget:
+    total_budget_tokens: int
+    allocated_sources: List[ContextSource]
+    omitted_sources: List[ContextSource]
+    summarized_sources: List[Tuple[ContextSource, str]]  # (source, summary_text)
+    retrieval_queries: List[str]
+    cache_prefix_tokens: int
+    dynamic_suffix_tokens: int
+    estimated_cost: float
+
+
+class ContextBudgeter:
+    """Intelligently budgets context window to minimize cost while preserving quality."""
+
+    # Task-specific context budgets (tokens)
+    DEFAULT_BUDGETS = {
+        TaskType.QUICK_ANSWER: 2048,
+        TaskType.UNKNOWN_AMBIGUOUS: 4096,
+        TaskType.TOOL_HEAVY: 8192,
+        TaskType.RETRIEVAL_HEAVY: 16384,
+        TaskType.DOCUMENT_DRAFTING: 8192,
+        TaskType.CODING: 12288,
+        TaskType.RESEARCH: 16384,
+        TaskType.LONG_HORIZON: 32768,
+        TaskType.LEGAL_REGULATED: 24576,
+    }
+
+    # Importance weights by task type
+    IMPORTANCE_RULES = {
+        TaskType.QUICK_ANSWER: {
+            "system_rules": 0.9,
+            "recent_messages": 0.9,
+            "user_preferences": 0.5,
+        },
+        TaskType.CODING: {
+            "system_rules": 0.8,
+            "tool_descriptions": 0.9,
+            "artifacts": 0.9,
+            "recent_messages": 0.7,
+            "examples": 0.6,
+        },
+        TaskType.RESEARCH: {
+            "retrieved_docs": 0.95,
+            "recent_messages": 0.6,
+            "system_rules": 0.5,
+            "task_plan": 0.8,
+        },
+        TaskType.LEGAL_REGULATED: {
+            "retrieved_docs": 0.95,
+            "system_rules": 0.9,
+            "user_preferences": 0.7,
+            "artifacts": 0.8,
+        },
+        TaskType.LONG_HORIZON: {
+            "task_plan": 0.95,
+            "recent_messages": 0.8,
+            "artifacts": 0.85,
+            "system_rules": 0.7,
+            "prior_trace_failures": 0.6,
+        },
+    }
+
+    def __init__(self, config: Optional[ACOConfig] = None):
+        self.config = config or ACOConfig()
+
+    def budget(
+        self,
+        task_type: TaskType,
+        available_sources: List[ContextSource],
+        model_max_context: int = 128000,
+        cost_per_1k_input: float = 0.01,
+    ) -> ContextBudget:
+        """Allocate context budget across sources."""
+        
+        budget_tokens = self.DEFAULT_BUDGETS.get(task_type, 8192)
+        # Don't exceed model limit
+        budget_tokens = min(budget_tokens, int(model_max_context * 0.8))
+        
+        # Apply importance rules
+        importance_map = self.IMPORTANCE_RULES.get(task_type, {})
+        for source in available_sources:
+            source.importance = importance_map.get(source.name, source.importance)
+        
+        # Separate stable (cacheable) vs dynamic
+        stable_sources = [s for s in available_sources if s.cacheable and not s.mutable]
+        dynamic_sources = [s for s in available_sources if s.mutable or not s.cacheable]
+        
+        # Always include stable sources in prefix (they're cache-efficient)
+        prefix_tokens = sum(s.tokens for s in stable_sources)
+        remaining = budget_tokens - prefix_tokens
+        
+        # Sort dynamic by importance / staleness ratio
+        dynamic_sources.sort(key=lambda s: s.importance / (1 + s.staleness), reverse=True)
+        
+        allocated = list(stable_sources)
+        omitted = []
+        summarized = []
+        retrieval_queries = []
+        
+        for source in dynamic_sources:
+            if source.tokens <= remaining:
+                allocated.append(source)
+                remaining -= source.tokens
+            elif source.importance > 0.7 and source.tokens > remaining * 1.5:
+                # Source is important but too big — summarize it
+                summary_tokens = min(int(remaining * 0.3), 512)
+                if summary_tokens > 50:
+                    summary = self._summarize(source, summary_tokens)
+                    summarized.append((source, summary))
+                    remaining -= summary_tokens
+            elif source.importance > 0.8:
+                # Critical but doesn't fit — mark for retrieval instead
+                omitted.append(source)
+                retrieval_queries.append(f"retrieve:{source.name}")
+            else:
+                omitted.append(source)
+        
+        dynamic_used = sum(s.tokens for s in allocated if s in dynamic_sources)
+        dynamic_used += sum(len(s[1].split()) for s in summarized)
+        
+        total_tokens = prefix_tokens + dynamic_used
+        estimated_cost = (total_tokens / 1000) * cost_per_1k_input
+        
+        return ContextBudget(
+            total_budget_tokens=budget_tokens,
+            allocated_sources=allocated,
+            omitted_sources=omitted,
+            summarized_sources=summarized,
+            retrieval_queries=retrieval_queries,
+            cache_prefix_tokens=prefix_tokens,
+            dynamic_suffix_tokens=dynamic_used,
+            estimated_cost=estimated_cost,
+        )
+
+    def _summarize(self, source: ContextSource, max_tokens: int) -> str:
+        """Produce a token-budgeted summary of a context source."""
+        # In production, use a summarization model
+        # Here we return a placeholder
+        return f"[SUMMARY:{source.name}:{max_tokens}tokens]"
+
+    def should_retrieve(self, source: ContextSource, task_type: TaskType) -> bool:
+        """Decide if a source should be retrieved on-demand vs. kept in context."""
+        if source.staleness > 0.5:
+            return True
+        if source.tokens > 4096 and source.importance < 0.8:
+            return True
+        if task_type in (TaskType.RESEARCH, TaskType.RETRIEVAL_HEAVY):
+            return True
+        return False
+
+    def compress_history(
+        self,
+        messages: List[Dict[str, str]],
+        max_messages: int = 10,
+        summarize_older: bool = True,
+    ) -> List[Dict[str, str]]:
+        """Compress message history by summarizing older messages."""
+        if len(messages) <= max_messages:
+            return messages
+        
+        keep = messages[-max_messages:]
+        older = messages[:-max_messages]
+        
+        if summarize_older and older:
+            summary = self._summarize_messages(older)
+            return [{"role": "system", "content": f"[Earlier context summary]: {summary}"}] + keep
+        
+        return keep
+
+    def _summarize_messages(self, messages: List[Dict[str, str]]) -> str:
+        """Summarize a list of messages."""
+        # In production, use a summarization model
+        return f"{len(messages)} earlier messages summarized."