ds-agent-cli 0.1.0

package/src/utils/token_budget.py

@@ -0,0 +1,411 @@
+"""
+Strict Token Budget Management
+
+Implements sliding window conversation history, aggressive compression,
+and emergency context truncation to prevent context window overflow.
+"""
+
+from typing import List, Dict, Any, Optional, Tuple
+import json
+import tiktoken
+from pathlib import Path
+
+
+class ConversationMessage:
+    """Represents a message with priority for history management."""
+    
+    def __init__(self, role: str, content: str, message_type: str = "normal",
+                 priority: int = 5, tokens: Optional[int] = None):
+        self.role = role
+        self.content = content
+        self.message_type = message_type  # system, tool_result, assistant, user, normal
+        self.priority = priority  # 1 (drop first) to 10 (keep last)
+        self.tokens = tokens
+        self.timestamp = None
+    
+    def to_dict(self) -> Dict[str, str]:
+        """Convert to OpenAI message format."""
+        return {"role": self.role, "content": self.content}
+
+
+class TokenBudgetManager:
+    """
+    Manages conversation history with strict token budget enforcement.
+    
+    Features:
+    - Accurate token counting using tiktoken
+    - Priority-based message dropping
+    - Sliding window with smart compression
+    - Emergency context truncation
+    - Keeps recent tool results, drops old assistant messages
+    """
+    
+    def __init__(self, model: str = "gpt-4", max_tokens: int = 128000,
+                 reserve_tokens: int = 8000):
+        """
+        Initialize token budget manager.
+        
+        Args:
+            model: Model name for token counting
+            max_tokens: Maximum context window size
+            reserve_tokens: Tokens to reserve for response
+        """
+        self.model = model
+        self.max_tokens = max_tokens
+        self.reserve_tokens = reserve_tokens
+        self.available_tokens = max_tokens - reserve_tokens
+        
+        # Initialize tokenizer
+        try:
+            self.encoding = tiktoken.encoding_for_model(model)
+        except:
+            # Fallback to cl100k_base (GPT-4/GPT-3.5)
+            self.encoding = tiktoken.get_encoding("cl100k_base")
+        
+        print(f"📊 Token Budget: {self.available_tokens:,} tokens available ({self.max_tokens:,} - {self.reserve_tokens:,} reserve)")
+    
+    def count_tokens(self, text: str) -> int:
+        """Count tokens in text using tiktoken."""
+        try:
+            return len(self.encoding.encode(text))
+        except:
+            # Fallback estimation: ~4 chars per token
+            return len(text) // 4
+    
+    def count_message_tokens(self, message) -> int:
+        """
+        Count tokens in a message (includes role overhead).
+        
+        Args:
+            message: Either a dict or a Pydantic ChatMessage object
+        """
+        # Format: <|role|>content<|endofmessage|>
+        # Approximately 4 tokens overhead per message
+        
+        # Handle both dict and Pydantic object formats
+        if isinstance(message, dict):
+            content = message.get("content", "")
+            role = message.get("role", "")
+        else:
+            # Pydantic object (like ChatMessage from Mistral SDK)
+            content = getattr(message, "content", "")
+            role = getattr(message, "role", "")
+        
+        content_tokens = self.count_tokens(str(content))
+        role_tokens = self.count_tokens(str(role))
+        return content_tokens + role_tokens + 4
+    
+    def count_messages_tokens(self, messages: List) -> int:
+        """Count total tokens in message list."""
+        return sum(self.count_message_tokens(msg) for msg in messages)
+    
+    def compress_tool_result(self, tool_result: str, max_tokens: int = 500) -> str:
+        """
+        Aggressively compress tool result while keeping key information.
+        
+        Keeps:
+        - Success/failure status
+        - Key metrics and numbers
+        - Error messages
+        
+        Drops:
+        - Verbose logs
+        - Duplicate information
+        - Large data structures
+        """
+        if self.count_tokens(tool_result) <= max_tokens:
+            return tool_result
+        
+        try:
+            # Try to parse as JSON
+            result_dict = json.loads(tool_result)
+            
+            # Extract essential fields
+            compressed = {
+                "success": result_dict.get("success", True),
+            }
+            
+            # Add error if present
+            if "error" in result_dict:
+                compressed["error"] = str(result_dict["error"])[:200]
+            
+            # Add key metrics (numbers, scores, paths)
+            for key in ["score", "accuracy", "best_score", "n_rows", "n_cols", 
+                       "output_path", "best_model", "result_summary"]:
+                if key in result_dict:
+                    compressed[key] = result_dict[key]
+            
+            # Add result if it's small
+            if "result" in result_dict:
+                result_str = str(result_dict["result"])
+                if len(result_str) < 300:
+                    compressed["result"] = result_str[:300]
+            
+            return json.dumps(compressed, indent=None)
+            
+        except json.JSONDecodeError:
+            # Not JSON - truncate intelligently
+            lines = tool_result.split('\n')
+            
+            # Keep first 5 and last 5 lines
+            if len(lines) > 15:
+                compressed_lines = lines[:5] + ["... (truncated) ..."] + lines[-5:]
+                result = '\n'.join(compressed_lines)
+            else:
+                result = tool_result
+            
+            # Hard truncate if still too long
+            token_count = self.count_tokens(result)
+            if token_count > max_tokens:
+                # Truncate to character limit (rough)
+                char_limit = max_tokens * 4
+                result = result[:char_limit] + "... (truncated)"
+            
+            return result
+    
+    def prioritize_messages(self, messages: List[ConversationMessage]) -> List[ConversationMessage]:
+        """
+        Assign priorities to messages based on type and importance.
+        
+        Priority levels:
+        - 10: System prompt, recent user messages
+        - 9: Recent tool results (last 3)
+        - 8: Recent assistant responses (last 2)
+        - 5: Normal messages
+        - 3: Old tool results
+        - 2: Old assistant responses
+        - 1: Very old messages
+        """
+        # Find recent messages (last 5)
+        recent_threshold = max(0, len(messages) - 5)
+        
+        for i, msg in enumerate(messages):
+            if msg.message_type == "system":
+                msg.priority = 10
+            elif msg.role == "user":
+                msg.priority = 10 if i >= recent_threshold else 7
+            elif msg.message_type == "tool_result":
+                msg.priority = 9 if i >= recent_threshold else 3
+            elif msg.role == "assistant":
+                msg.priority = 8 if i >= recent_threshold else 2
+            else:
+                msg.priority = 5 if i >= recent_threshold else 1
+        
+        return messages
+    
+    def apply_sliding_window(self, messages: List[ConversationMessage],
+                            target_tokens: int) -> List[ConversationMessage]:
+        """
+        Apply sliding window to fit within token budget.
+        
+        Strategy:
+        1. Always keep system prompt (first message)
+        2. Keep recent messages (last N)
+        3. Drop low-priority messages from middle
+        4. Compress tool results if needed
+        
+        Args:
+            messages: List of ConversationMessage objects
+            target_tokens: Target token count
+        
+        Returns:
+            Filtered message list within budget
+        """
+        if not messages:
+            return []
+        
+        # Always keep system prompt
+        system_msg = messages[0] if messages[0].message_type == "system" else None
+        other_messages = messages[1:] if system_msg else messages
+        
+        # Prioritize messages
+        other_messages = self.prioritize_messages(other_messages)
+        
+        # Sort by priority (high to low)
+        sorted_messages = sorted(other_messages, key=lambda m: m.priority, reverse=True)
+        
+        # Calculate tokens for each message
+        for msg in sorted_messages:
+            if msg.tokens is None:
+                msg.tokens = self.count_message_tokens(msg.to_dict())
+        
+        # Greedily add messages until budget exhausted
+        kept_messages = []
+        current_tokens = 0
+        
+        # Add system prompt first
+        if system_msg:
+            system_msg.tokens = self.count_message_tokens(system_msg.to_dict())
+            kept_messages.append(system_msg)
+            current_tokens += system_msg.tokens
+        
+        # Add other messages by priority
+        for msg in sorted_messages:
+            if current_tokens + msg.tokens <= target_tokens:
+                kept_messages.append(msg)
+                current_tokens += msg.tokens
+            elif msg.message_type == "tool_result" and msg.priority >= 8:
+                # Try compressing critical tool results
+                compressed_content = self.compress_tool_result(msg.content, max_tokens=300)
+                compressed_tokens = self.count_tokens(compressed_content)
+                
+                if current_tokens + compressed_tokens <= target_tokens:
+                    msg.content = compressed_content
+                    msg.tokens = compressed_tokens
+                    kept_messages.append(msg)
+                    current_tokens += compressed_tokens
+        
+        # Sort kept messages back to chronological order
+        # System message stays first, rest in order they appeared
+        if system_msg:
+            non_system = [m for m in kept_messages if m != system_msg]
+            # Sort by original index (approximate by content comparison)
+            original_order = []
+            for orig_msg in messages:
+                for kept in non_system:
+                    if kept.content == orig_msg.content:
+                        original_order.append(kept)
+                        break
+            
+            kept_messages = [system_msg] + original_order
+        
+        print(f"📊 Sliding window: {len(messages)} → {len(kept_messages)} messages ({current_tokens:,} tokens)")
+        
+        return kept_messages
+    
+    def emergency_truncate(self, messages: List[Dict[str, str]],
+                          max_tokens: int) -> List[Dict[str, str]]:
+        """
+        Emergency truncation when context is about to overflow.
+        
+        Aggressive strategy:
+        - Keep system prompt
+        - Keep last user message
+        - Keep last 2 messages
+        - Truncate everything else
+        
+        Args:
+            messages: Message list
+            max_tokens: Hard token limit
+        
+        Returns:
+            Truncated message list
+        """
+        if not messages:
+            return []
+        
+        print("⚠️ EMERGENCY TRUNCATION: Context overflow imminent")
+        
+        # Always keep system, last user, and last 2 messages
+        essential_messages = []
+        
+        # System prompt (first message)
+        if messages:
+            essential_messages.append(messages[0])
+        
+        # Last 2 messages
+        if len(messages) > 2:
+            essential_messages.extend(messages[-2:])
+        else:
+            essential_messages.extend(messages[1:])
+        
+        # Count tokens
+        total_tokens = self.count_messages_tokens(essential_messages)
+        
+        if total_tokens <= max_tokens:
+            return essential_messages
+        
+        # Still too large - truncate system prompt
+        print("⚠️ Truncating system prompt to fit budget")
+        system_msg = essential_messages[0]
+        
+        # Handle both dict and Pydantic object formats
+        if isinstance(system_msg, dict):
+            system_content = system_msg["content"]
+        else:
+            system_content = getattr(system_msg, "content", "")
+        
+        # Keep first 1000 chars of system prompt
+        truncated_system = {
+            "role": "system",
+            "content": str(system_content)[:1000] + "\n\n... (truncated due to context limit) ..."
+        }
+        
+        return [truncated_system] + essential_messages[1:]
+    
+    def enforce_budget(self, messages: List[Dict[str, str]],
+                      system_prompt: Optional[str] = None) -> Tuple[List[Dict[str, str]], int]:
+        """
+        Main entry point: Enforce token budget on message list.
+        
+        Args:
+            messages: List of messages
+            system_prompt: Optional new system prompt to prepend
+        
+        Returns:
+            (filtered_messages, total_tokens)
+        """
+        # Add system prompt if provided
+        if system_prompt:
+            messages = [{"role": "system", "content": system_prompt}] + messages
+        
+        # Count current tokens
+        current_tokens = self.count_messages_tokens(messages)
+        
+        print(f"📊 Token Budget Check: {current_tokens:,} / {self.available_tokens:,} tokens")
+        
+        # If within budget, return as-is
+        if current_tokens <= self.available_tokens:
+            print("✅ Within budget")
+            return messages, current_tokens
+        
+        print(f"⚠️ Over budget by {current_tokens - self.available_tokens:,} tokens")
+        
+        # Convert to ConversationMessage objects
+        conv_messages = []
+        for i, msg in enumerate(messages):
+            # Handle both dict and Pydantic object formats
+            if isinstance(msg, dict):
+                role = msg.get("role", "")
+                content = msg.get("content", "")
+            else:
+                role = getattr(msg, "role", "")
+                content = getattr(msg, "content", "")
+            
+            msg_type = "system" if i == 0 and role == "system" else "normal"
+            if "tool" in str(content).lower() or "function" in str(content).lower():
+                msg_type = "tool_result"
+            
+            conv_msg = ConversationMessage(
+                role=role,
+                content=str(content),
+                message_type=msg_type
+            )
+            conv_messages.append(conv_msg)
+        
+        # Apply sliding window
+        filtered = self.apply_sliding_window(conv_messages, self.available_tokens)
+        
+        # Convert back to dict format
+        result_messages = [msg.to_dict() for msg in filtered]
+        final_tokens = self.count_messages_tokens(result_messages)
+        
+        # Emergency truncation if still over
+        if final_tokens > self.available_tokens:
+            result_messages = self.emergency_truncate(result_messages, self.available_tokens)
+            final_tokens = self.count_messages_tokens(result_messages)
+        
+        print(f"✅ Budget enforced: {final_tokens:,} tokens ({len(result_messages)} messages)")
+        
+        return result_messages, final_tokens
+
+
+# Global token budget manager instance
+_token_manager = None
+
+def get_token_manager(model: str = "gpt-4", max_tokens: int = 128000) -> TokenBudgetManager:
+    """Get or create global token budget manager."""
+    global _token_manager
+    if _token_manager is None:
+        _token_manager = TokenBudgetManager(model=model, max_tokens=max_tokens)
+    return _token_manager