alita-sdk 0.3.379__py3-none-any.whl → 0.3.627__py3-none-any.whl

alita_sdk/cli/context/message.py

@@ -0,0 +1,285 @@
+"""
+CLI Message wrapper with context tracking.
+
+Provides a message class that tracks inclusion status, token counts,
+and supports conversion to/from LangChain message formats.
+"""
+
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import Any, Dict, List, Optional, Union
+
+from .token_estimation import estimate_message_tokens
+
+
+@dataclass
+class CLIMessage:
+    """
+    Chat message with context management metadata.
+    
+    Tracks whether the message is included in context, its token count,
+    and provides conversion to various message formats.
+    
+    Attributes:
+        role: Message role (user, assistant, system)
+        content: Message content text
+        index: Position in the full conversation history
+        token_count: Cached token count for this message
+        included: Whether message is included in LLM context
+        created_at: Timestamp when message was created
+        priority: Priority weight for importance-based pruning
+        weight: Additional weight factor for pruning decisions
+    """
+    role: str
+    content: str
+    index: int
+    token_count: int = 0
+    included: bool = True
+    created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
+    priority: float = 1.0
+    weight: float = 1.0
+    
+    def __post_init__(self):
+        """Calculate token count if not provided."""
+        if self.token_count == 0 and self.content:
+            self.token_count = estimate_message_tokens(self.role, self.content)
+    
+    @classmethod
+    def from_dict(cls, msg_dict: Dict[str, Any], index: int, model: str = 'gpt-4') -> 'CLIMessage':
+        """
+        Create CLIMessage from a simple dict format.
+        
+        Args:
+            msg_dict: Dictionary with 'role' and 'content' keys
+            index: Position in conversation
+            model: Model name for token estimation
+            
+        Returns:
+            CLIMessage instance
+        """
+        role = msg_dict.get('role', 'user')
+        content = msg_dict.get('content', '')
+        token_count = estimate_message_tokens(role, content, model)
+        
+        return cls(
+            role=role,
+            content=content,
+            index=index,
+            token_count=token_count,
+            included=msg_dict.get('included', True),
+            priority=msg_dict.get('priority', 1.0),
+            weight=msg_dict.get('weight', 1.0),
+        )
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """
+        Convert to simple dict format for LLM calls.
+        
+        Returns:
+            Dictionary with 'role' and 'content' keys
+        """
+        return {
+            'role': self.role,
+            'content': self.content,
+        }
+    
+    def to_langchain_message(self) -> Any:
+        """
+        Convert to LangChain message format.
+        
+        Returns:
+            Appropriate LangChain message type (HumanMessage, AIMessage, SystemMessage)
+        """
+        from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
+        
+        if self.role == 'user':
+            return HumanMessage(content=self.content)
+        elif self.role == 'assistant':
+            return AIMessage(content=self.content)
+        elif self.role == 'system':
+            return SystemMessage(content=self.content)
+        else:
+            # Default to HumanMessage for unknown roles
+            return HumanMessage(content=self.content)
+    
+    @classmethod
+    def from_langchain_message(cls, message: Any, index: int) -> 'CLIMessage':
+        """
+        Create CLIMessage from LangChain message.
+        
+        Args:
+            message: LangChain message (HumanMessage, AIMessage, etc.)
+            index: Position in conversation
+            
+        Returns:
+            CLIMessage instance
+        """
+        from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
+        
+        content = message.content if hasattr(message, 'content') else str(message)
+        
+        if isinstance(message, HumanMessage):
+            role = 'user'
+        elif isinstance(message, AIMessage):
+            role = 'assistant'
+        elif isinstance(message, SystemMessage):
+            role = 'system'
+        else:
+            role = 'user'
+        
+        return cls(
+            role=role,
+            content=content,
+            index=index,
+        )
+    
+    def to_state_dict(self) -> Dict[str, Any]:
+        """
+        Convert to state dictionary for persistence.
+        
+        Returns:
+            Dictionary with all fields for saving to context_state.json
+        """
+        return {
+            'index': self.index,
+            'role': self.role,
+            'content': self.content,
+            'token_count': self.token_count,
+            'included': self.included,
+            'created_at': self.created_at.isoformat(),
+            'priority': self.priority,
+            'weight': self.weight,
+        }
+    
+    @classmethod
+    def from_state_dict(cls, state: Dict[str, Any]) -> 'CLIMessage':
+        """
+        Restore CLIMessage from state dictionary.
+        
+        Args:
+            state: Dictionary from to_state_dict()
+            
+        Returns:
+            CLIMessage instance
+        """
+        created_at = state.get('created_at')
+        if isinstance(created_at, str):
+            created_at = datetime.fromisoformat(created_at)
+        elif created_at is None:
+            created_at = datetime.now(timezone.utc)
+        
+        return cls(
+            role=state['role'],
+            content=state['content'],
+            index=state['index'],
+            token_count=state.get('token_count', 0),
+            included=state.get('included', True),
+            created_at=created_at,
+            priority=state.get('priority', 1.0),
+            weight=state.get('weight', 1.0),
+        )
+    
+    @property
+    def meta(self) -> Dict[str, Any]:
+        """
+        Get metadata in context_manager compatible format.
+        
+        Used for compatibility with pruning strategies from context_manager.
+        """
+        return {
+            'context': {
+                'token_count': self.token_count,
+                'weight': self.weight,
+                'priority': self.priority,
+                'included': self.included,
+            }
+        }
+    
+    @property
+    def reply_to_id(self) -> Optional[int]:
+        """
+        Get reply-to ID (for thread awareness).
+        
+        For CLI messages, we use simple sequential ordering.
+        User messages reply to previous assistant, and vice versa.
+        """
+        if self.index > 0:
+            return self.index - 1
+        return None
+    
+    @property
+    def author_participant(self) -> Any:
+        """
+        Get author participant for importance scoring.
+        
+        Returns a simple object with entity_name attribute.
+        """
+        class Participant:
+            def __init__(self, role: str):
+                self.entity_name = role
+        
+        return Participant(self.role)
+
+
+def messages_to_cli_messages(
+    messages: List[Dict[str, Any]],
+    model: str = 'gpt-4'
+) -> List[CLIMessage]:
+    """
+    Convert a list of message dicts to CLIMessage objects.
+    
+    Args:
+        messages: List of dicts with 'role' and 'content' keys
+        model: Model name for token estimation
+        
+    Returns:
+        List of CLIMessage objects
+    """
+    return [
+        CLIMessage.from_dict(msg, index=i, model=model)
+        for i, msg in enumerate(messages)
+    ]
+
+
+def cli_messages_to_dicts(
+    messages: List[CLIMessage],
+    include_only: bool = True
+) -> List[Dict[str, str]]:
+    """
+    Convert CLIMessage objects to simple dicts for LLM calls.
+    
+    Args:
+        messages: List of CLIMessage objects
+        include_only: If True, only include messages where included=True
+        
+    Returns:
+        List of dicts with 'role' and 'content' keys
+    """
+    result = []
+    for msg in messages:
+        if include_only and not msg.included:
+            continue
+        result.append(msg.to_dict())
+    return result
+
+
+def cli_messages_to_langchain(
+    messages: List[CLIMessage],
+    include_only: bool = True
+) -> List[Any]:
+    """
+    Convert CLIMessage objects to LangChain messages.
+    
+    Args:
+        messages: List of CLIMessage objects
+        include_only: If True, only include messages where included=True
+        
+    Returns:
+        List of LangChain message objects
+    """
+    result = []
+    for msg in messages:
+        if include_only and not msg.included:
+            continue
+        result.append(msg.to_langchain_message())
+    return result

alita_sdk/cli/context/strategies.py

@@ -0,0 +1,289 @@
+"""
+Pruning strategies for CLI context management.
+
+Implements various strategies for selecting which messages to include
+in the LLM context when the token limit is exceeded.
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Type, Union
+
+from .message import CLIMessage
+
+
+@dataclass
+class PruningConfig:
+    """Configuration for pruning strategies."""
+    
+    max_context_tokens: int = 8000
+    preserve_recent_messages: int = 5
+    pruning_method: str = 'oldest_first'
+    weights: Dict[str, float] = field(default_factory=lambda: {
+        'recency': 1.0,
+        'importance': 1.0,
+        'user_messages': 1.2,
+        'thread_continuity': 1.0,
+    })
+    
+    @classmethod
+    def from_strategy(cls, strategy_config: Dict) -> 'PruningConfig':
+        """Create PruningConfig from strategy configuration dict."""
+        return cls(
+            max_context_tokens=strategy_config.get('max_context_tokens', 8000),
+            preserve_recent_messages=strategy_config.get('preserve_recent_messages', 5),
+            pruning_method=strategy_config.get('pruning_method', 'oldest_first'),
+            weights=strategy_config.get('weights', {
+                'recency': 1.0,
+                'importance': 1.0,
+                'user_messages': 1.2,
+                'thread_continuity': 1.0,
+            }),
+        )
+
+
+class PruningStrategy(ABC):
+    """Abstract base class for pruning strategies."""
+
+    def __init__(self, config: PruningConfig):
+        self.config = config
+
+    @abstractmethod
+    def select_messages(
+        self,
+        messages: List[CLIMessage],
+        available_tokens: int
+    ) -> List[CLIMessage]:
+        """
+        Select messages within the available token budget.
+        
+        Args:
+            messages: List of messages to select from
+            available_tokens: Maximum tokens available for selection
+            
+        Returns:
+            Selected messages within token budget
+        """
+        pass
+
+    @staticmethod
+    def get_token_count(message: CLIMessage) -> int:
+        """Extract token count from message."""
+        return message.token_count
+
+
+class OldestFirstStrategy(PruningStrategy):
+    """
+    Select messages starting from newest until token limit is reached.
+    
+    This effectively drops the oldest messages first when context is full.
+    """
+
+    def select_messages(
+        self,
+        messages: List[CLIMessage],
+        available_tokens: int
+    ) -> List[CLIMessage]:
+        selected = []
+        current_tokens = 0
+
+        # Sort by index descending (newest first)
+        sorted_messages = sorted(messages, key=lambda x: x.index, reverse=True)
+
+        for message in sorted_messages:
+            msg_tokens = self.get_token_count(message)
+            if current_tokens + msg_tokens <= available_tokens:
+                selected.append(message)
+                current_tokens += msg_tokens
+            else:
+                break
+
+        # Return in original order (oldest first)
+        selected.sort(key=lambda x: x.index)
+        return selected
+
+
+class ImportanceBasedStrategy(PruningStrategy):
+    """Select messages based on calculated importance scores."""
+
+    def select_messages(
+        self,
+        messages: List[CLIMessage],
+        available_tokens: int
+    ) -> List[CLIMessage]:
+        # Calculate importance scores
+        scored_messages = []
+        for message in messages:
+            score = self._calculate_importance_score(message)
+            scored_messages.append((score, message))
+
+        # Sort by importance (highest first)
+        scored_messages.sort(key=lambda x: x[0], reverse=True)
+
+        # Select until token limit
+        selected = []
+        current_tokens = 0
+
+        for score, message in scored_messages:
+            msg_tokens = self.get_token_count(message)
+            if current_tokens + msg_tokens <= available_tokens:
+                selected.append(message)
+                current_tokens += msg_tokens
+
+        # Return in original order
+        selected.sort(key=lambda x: x.index)
+        return selected
+
+    def _calculate_importance_score(self, message: CLIMessage) -> float:
+        """Calculate importance score for a message."""
+        base_score = message.priority
+        weight = message.weight
+
+        # Factor in message length (longer messages might be more important)
+        token_count = message.token_count
+        length_factor = min(1.5, token_count / 100)  # Cap at 1.5x
+
+        # Factor in replies (messages with replies might be more important)
+        reply_factor = 1.2 if message.reply_to_id else 1.0
+
+        # Factor in user vs assistant messages
+        weights = self.config.weights
+        role_factor = 1.0
+        if message.role == 'user':
+            role_factor = weights.get('user_messages', 1.0)
+        elif message.role == 'system':
+            role_factor = 1.5  # System messages are usually important
+
+        # Factor in recency (newer messages get higher scores)
+        recency_weight = weights.get('recency', 1.0)
+        # Simple recency factor based on index
+        recency_factor = 1.0 + (message.index * 0.01 * recency_weight)
+
+        return base_score * weight * length_factor * reply_factor * role_factor * recency_factor
+
+
+class PruningStrategyFactory:
+    """Factory for creating pruning strategy instances."""
+
+    _strategies: Dict[str, Type[PruningStrategy]] = {
+        'oldest_first': OldestFirstStrategy,
+        'importance_based': ImportanceBasedStrategy,
+    }
+
+    @classmethod
+    def create(cls, strategy_name: str, config: PruningConfig) -> PruningStrategy:
+        """
+        Create a pruning strategy instance.
+        
+        Args:
+            strategy_name: Name of the strategy ('oldest_first', 'importance_based')
+            config: Configuration for the strategy
+            
+        Returns:
+            PruningStrategy instance
+            
+        Raises:
+            ValueError: If strategy_name is not recognized
+        """
+        strategy_class = cls._strategies.get(strategy_name)
+        if not strategy_class:
+            available = list(cls._strategies.keys())
+            raise ValueError(
+                f"Unknown pruning strategy: {strategy_name}. "
+                f"Available strategies: {available}"
+            )
+        return strategy_class(config)
+
+    @classmethod
+    def register_strategy(cls, name: str, strategy_class: Type[PruningStrategy]):
+        """Register a custom pruning strategy."""
+        cls._strategies[name] = strategy_class
+
+    @classmethod
+    def available_strategies(cls) -> List[str]:
+        """Get list of available strategy names."""
+        return list(cls._strategies.keys())
+
+
+class PruningOrchestrator:
+    """
+    Orchestrates the pruning process by applying strategies with token budgets.
+    
+    Handles preserving recent messages, calculating available token budgets,
+    and coordinating strategy application.
+    """
+
+    def __init__(self, strategy_config: Dict):
+        """
+        Initialize the orchestrator with strategy configuration.
+        
+        Args:
+            strategy_config: Configuration dictionary containing:
+                - pruning_method: Strategy name to use
+                - preserve_recent_messages: Number of recent messages to preserve
+                - max_context_tokens: Maximum tokens allowed
+                - weights: Weighting configuration for strategies
+        """
+        self.strategy_config = strategy_config
+        self.pruning_config = PruningConfig.from_strategy(strategy_config)
+
+    def apply_pruning(
+        self,
+        messages: List[CLIMessage],
+        summaries: Optional[List[dict]] = None,
+        summary_tokens: int = 0,
+    ) -> List[CLIMessage]:
+        """
+        Apply pruning strategy to reduce context size.
+        
+        Args:
+            messages: List of messages to prune
+            summaries: List of summaries (for token calculation)
+            summary_tokens: Pre-calculated summary token count
+            
+        Returns:
+            List of messages to include in context (with included=True)
+        """
+        strategy_method = self.strategy_config.get('pruning_method', 'oldest_first')
+        preserve_recent = self.strategy_config.get('preserve_recent_messages', 5)
+        max_tokens = self.strategy_config.get('max_context_tokens', 8000)
+
+        # Split into recent (always preserved) and older messages
+        if preserve_recent > 0 and len(messages) > preserve_recent:
+            recent_messages = messages[-preserve_recent:]
+            older_messages = messages[:-preserve_recent]
+        else:
+            recent_messages = messages
+            older_messages = []
+
+        # Calculate tokens used by preserved recent messages and summaries
+        preserved_tokens = sum(m.token_count for m in recent_messages)
+        preserved_tokens += summary_tokens
+
+        available_tokens = max_tokens - preserved_tokens
+
+        if available_tokens <= 0 or not older_messages:
+            # Mark older messages as excluded
+            for msg in older_messages:
+                msg.included = False
+            return recent_messages
+
+        # Apply pruning strategy to older messages
+        try:
+            strategy = PruningStrategyFactory.create(strategy_method, self.pruning_config)
+            selected_older = strategy.select_messages(older_messages, available_tokens)
+        except ValueError:
+            # Fallback to oldest_first
+            strategy = PruningStrategyFactory.create('oldest_first', self.pruning_config)
+            selected_older = strategy.select_messages(older_messages, available_tokens)
+
+        # Mark messages as included/excluded
+        selected_indices = {m.index for m in selected_older}
+        for msg in older_messages:
+            msg.included = msg.index in selected_indices
+
+        # Combine and sort by index
+        final_selection = selected_older + recent_messages
+        final_selection.sort(key=lambda x: x.index)
+
+        return final_selection

alita_sdk/cli/context/token_estimation.py

@@ -0,0 +1,127 @@
+"""
+Token estimation utilities for CLI context management.
+
+Uses tiktoken for accurate token counting with fallback to character-based estimation.
+"""
+
+from typing import List, Optional, TYPE_CHECKING
+from functools import lru_cache
+
+try:
+    import tiktoken
+    TIKTOKEN_AVAILABLE = True
+except ImportError:
+    TIKTOKEN_AVAILABLE = False
+
+if TYPE_CHECKING:
+    from .message import CLIMessage
+
+
+@lru_cache(maxsize=8)
+def get_encoding_for_model(model: str = 'gpt-4') -> Optional[object]:
+    """
+    Get the appropriate tiktoken encoding for a given model.
+    Defaults to cl100k_base (used by most modern models) and only specifies exceptions.
+    
+    Args:
+        model: Model name to get encoding for
+        
+    Returns:
+        Tiktoken encoding object or None if not available
+    """
+    if not TIKTOKEN_AVAILABLE:
+        return None
+
+    try:
+        # Get model encoding map from tiktoken
+        model_encoding_map = tiktoken.model.MODEL_TO_ENCODING
+        # Default to cl100k_base for unknown models (most modern models use this)
+        encoding_name = model_encoding_map.get(model.lower(), 'cl100k_base')
+        return tiktoken.get_encoding(encoding_name)
+    except Exception:
+        return None
+
+
+def estimate_tokens(text: str, model: str = 'gpt-4') -> int:
+    """
+    Accurate token estimation using tiktoken.
+    Falls back to character-based estimation if tiktoken is not available.
+    
+    Args:
+        text: Text to estimate tokens for
+        model: Model name for encoding selection
+        
+    Returns:
+        Estimated token count
+    """
+    if not text or not isinstance(text, str):
+        return 0
+
+    if TIKTOKEN_AVAILABLE:
+        encoder = get_encoding_for_model(model)
+        if encoder:
+            try:
+                return len(encoder.encode(text))
+            except Exception:
+                pass
+
+    # Fallback: Simple approximation (~4 characters per token for GPT models)
+    return max(1, len(text) // 4)
+
+
+def calculate_total_tokens(
+    messages: List['CLIMessage'],
+    summaries: Optional[List[dict]] = None,
+    include_only: bool = True
+) -> int:
+    """
+    Calculate total tokens from messages and summaries.
+    
+    Args:
+        messages: List of CLIMessage objects
+        summaries: Optional list of summary dictionaries with 'token_count' key
+        include_only: If True, only count messages where included=True
+        
+    Returns:
+        Total token count
+    """
+    message_tokens = 0
+    for msg in messages:
+        if include_only and not msg.included:
+            continue
+        message_tokens += msg.token_count
+    
+    summary_tokens = 0
+    if summaries:
+        for summary in summaries:
+            summary_tokens += summary.get('token_count', 0)
+    
+    return message_tokens + summary_tokens
+
+
+def estimate_message_tokens(role: str, content: str, model: str = 'gpt-4') -> int:
+    """
+    Estimate tokens for a chat message including role overhead.
+    
+    Chat messages have overhead for role tokens and message formatting.
+    This provides a more accurate estimate for chat-style messages.
+    
+    Args:
+        role: Message role (user, assistant, system)
+        content: Message content
+        model: Model name for encoding
+        
+    Returns:
+        Estimated token count including overhead
+    """
+    # Base content tokens
+    content_tokens = estimate_tokens(content, model)
+    
+    # Add overhead for message formatting (role, separators, etc.)
+    # Most chat models add ~4 tokens per message for formatting
+    overhead = 4
+    
+    # Role name tokens (typically 1-2 tokens)
+    role_tokens = estimate_tokens(role, model)
+    
+    return content_tokens + overhead + role_tokens