massgen 0.1.4__py3-none-any.whl → 0.1.5__py3-none-any.whl

massgen/memory/_compression.py

@@ -0,0 +1,237 @@
+# -*- coding: utf-8 -*-
+"""
+Context Window Compression
+
+Automatically compresses conversation history when context window fills up.
+Since messages are already recorded to persistent memory after each turn,
+compression simply removes old messages from active context while keeping
+them accessible via semantic search.
+"""
+
+from typing import Any, Callable, Dict, List, Optional
+
+from ..logger_config import logger
+from ..token_manager.token_manager import TokenCostCalculator
+from ._conversation import ConversationMemory
+from ._persistent import PersistentMemoryBase
+
+
+class CompressionStats:
+    """Statistics about a compression operation."""
+
+    def __init__(
+        self,
+        messages_removed: int = 0,
+        tokens_removed: int = 0,
+        messages_kept: int = 0,
+        tokens_kept: int = 0,
+    ):
+        self.messages_removed = messages_removed
+        self.tokens_removed = tokens_removed
+        self.messages_kept = messages_kept
+        self.tokens_kept = tokens_kept
+
+
+class ContextCompressor:
+    """
+    Compresses conversation history when context window fills up.
+
+    Strategy:
+    - Messages are already recorded to persistent_memory after each turn
+    - Compression removes old messages from conversation_memory
+    - Recent messages stay in active context
+    - Old messages remain accessible via semantic retrieval
+
+    Features:
+    - Token-aware compression (not just message count)
+    - Preserves system messages
+    - Keeps most recent messages
+    - Detailed compression logging
+
+    Example:
+        >>> compressor = ContextCompressor(
+        ...     token_calculator=TokenCostCalculator(),
+        ...     conversation_memory=conversation_memory,
+        ...     persistent_memory=persistent_memory
+        ... )
+        >>>
+        >>> stats = await compressor.compress_if_needed(
+        ...     messages=messages,
+        ...     current_tokens=96000,
+        ...     target_tokens=51200
+        ... )
+    """
+
+    def __init__(
+        self,
+        token_calculator: TokenCostCalculator,
+        conversation_memory: ConversationMemory,
+        persistent_memory: Optional[PersistentMemoryBase] = None,
+        on_compress: Optional[Callable[[CompressionStats], None]] = None,
+    ):
+        """
+        Initialize context compressor.
+
+        Args:
+            token_calculator: Calculator for token estimation
+            conversation_memory: Conversation memory to compress
+            persistent_memory: Optional persistent memory (for logging purposes)
+            on_compress: Optional callback called after compression
+        """
+        self.token_calculator = token_calculator
+        self.conversation_memory = conversation_memory
+        self.persistent_memory = persistent_memory
+        self.on_compress = on_compress
+
+        # Stats tracking
+        self.total_compressions = 0
+        self.total_messages_removed = 0
+        self.total_tokens_removed = 0
+
+    async def compress_if_needed(
+        self,
+        messages: List[Dict[str, Any]],
+        current_tokens: int,
+        target_tokens: int,
+        should_compress: bool = None,
+    ) -> Optional[CompressionStats]:
+        """
+        Compress messages if needed.
+
+        Args:
+            messages: Current conversation messages
+            current_tokens: Current token count
+            target_tokens: Target token count after compression
+            should_compress: Optional explicit compression flag
+                           If None, compresses only if current_tokens > target_tokens
+
+        Returns:
+            CompressionStats if compression occurred, None otherwise
+        """
+        # Determine if we need to compress
+        if should_compress is None:
+            should_compress = current_tokens > target_tokens
+
+        if not should_compress:
+            return None
+
+        # Select messages to keep
+        messages_to_keep = self._select_messages_to_keep(
+            messages=messages,
+            target_tokens=target_tokens,
+        )
+
+        if len(messages_to_keep) >= len(messages):
+            # No compression needed (already under target)
+            logger.debug("All messages fit within target, skipping compression")
+            return None
+
+        # Calculate stats
+        messages_removed = len(messages) - len(messages_to_keep)
+        messages_to_remove = [msg for msg in messages if msg not in messages_to_keep]
+        tokens_removed = self.token_calculator.estimate_tokens(messages_to_remove)
+        tokens_kept = self.token_calculator.estimate_tokens(messages_to_keep)
+
+        # Update conversation memory
+        try:
+            await self.conversation_memory.clear()
+            await self.conversation_memory.add(messages_to_keep)
+        except Exception as e:
+            logger.error(f"Failed to update conversation memory during compression: {e}")
+            return None
+
+        # Log compression result
+        if self.persistent_memory:
+            logger.info(
+                f"📦 Context compressed: Removed {messages_removed} old messages "
+                f"({tokens_removed:,} tokens) from active context.\n"
+                f"   Kept {len(messages_to_keep)} recent messages ({tokens_kept:,} tokens).\n"
+                f"   Old messages remain accessible via semantic search.",
+            )
+        else:
+            logger.warning(
+                f"⚠️  Context compressed: Removed {messages_removed} old messages "
+                f"({tokens_removed:,} tokens) from active context.\n"
+                f"   Kept {len(messages_to_keep)} recent messages ({tokens_kept:,} tokens).\n"
+                f"   No persistent memory - old messages NOT retrievable.",
+            )
+
+        # Update stats
+        self.total_compressions += 1
+        self.total_messages_removed += messages_removed
+        self.total_tokens_removed += tokens_removed
+
+        # Create stats object
+        stats = CompressionStats(
+            messages_removed=messages_removed,
+            tokens_removed=tokens_removed,
+            messages_kept=len(messages_to_keep),
+            tokens_kept=tokens_kept,
+        )
+
+        # Trigger callback if provided
+        if self.on_compress:
+            self.on_compress(stats)
+
+        return stats
+
+    def _select_messages_to_keep(
+        self,
+        messages: List[Dict[str, Any]],
+        target_tokens: int,
+    ) -> List[Dict[str, Any]]:
+        """
+        Select which messages to keep in active context.
+
+        Strategy:
+        1. Always keep system messages at the start
+        2. Keep most recent messages that fit in target_tokens
+        3. Remove everything in between
+
+        Args:
+            messages: All messages in conversation
+            target_tokens: Target token budget for kept messages
+
+        Returns:
+            List of messages to keep in conversation_memory
+        """
+        if not messages:
+            return []
+
+        # Separate system messages from others
+        system_messages = []
+        non_system_messages = []
+
+        for msg in messages:
+            if msg.get("role") == "system":
+                system_messages.append(msg)
+            else:
+                non_system_messages.append(msg)
+
+        # Start with system messages in kept list
+        messages_to_keep = system_messages.copy()
+        tokens_so_far = self.token_calculator.estimate_tokens(system_messages)
+
+        # Work backwards from most recent, adding messages until we hit target
+        recent_messages_to_keep = []
+        for msg in reversed(non_system_messages):
+            msg_tokens = self.token_calculator.estimate_tokens([msg])
+            if tokens_so_far + msg_tokens <= target_tokens:
+                tokens_so_far += msg_tokens
+                recent_messages_to_keep.insert(0, msg)  # Maintain order
+            else:
+                # Hit token limit, stop here
+                break
+
+        # Combine: system messages + recent messages
+        messages_to_keep.extend(recent_messages_to_keep)
+
+        return messages_to_keep
+
+    def get_stats(self) -> Dict[str, Any]:
+        """Get compression statistics."""
+        return {
+            "total_compressions": self.total_compressions,
+            "total_messages_removed": self.total_messages_removed,
+            "total_tokens_removed": self.total_tokens_removed,
+        }

massgen/memory/_context_monitor.py

@@ -0,0 +1,211 @@
+# -*- coding: utf-8 -*-
+"""
+Context Window Monitoring Utility
+
+Provides logging and tracking for context window usage during agent execution.
+Helps debug memory and token management by showing real-time context usage.
+"""
+
+from typing import Any, Dict, List, Optional
+
+from ..logger_config import logger
+from ..token_manager.token_manager import TokenCostCalculator
+
+
+class ContextWindowMonitor:
+    """Monitor and log context window usage during agent execution."""
+
+    def __init__(
+        self,
+        model_name: str,
+        provider: str = "openai",
+        trigger_threshold: float = 0.75,
+        target_ratio: float = 0.40,
+        enabled: bool = True,
+    ):
+        """
+        Initialize context window monitor.
+
+        Args:
+            model_name: Name of the model (e.g., "gpt-4o-mini")
+            provider: Provider name (e.g., "openai", "anthropic")
+            trigger_threshold: Percentage (0-1) at which to warn about context usage
+            target_ratio: Target percentage after compression
+            enabled: Whether to enable logging
+        """
+        self.model_name = model_name
+        self.provider = provider
+        self.trigger_threshold = trigger_threshold
+        self.target_ratio = target_ratio
+        self.enabled = enabled
+
+        # Get model pricing info to determine context window size
+        self.calculator = TokenCostCalculator()
+        self.pricing = self.calculator.get_model_pricing(provider, model_name)
+
+        if self.pricing and self.pricing.context_window:
+            self.context_window = self.pricing.context_window
+        else:
+            # Default fallbacks
+            self.context_window = 128000  # Common default
+            logger.warning(
+                f"Could not determine context window for {provider}/{model_name}, " f"using default {self.context_window}",
+            )
+
+        # Tracking
+        self.total_input_tokens = 0
+        self.total_output_tokens = 0
+        self.turn_count = 0
+
+    def log_context_usage(
+        self,
+        messages: List[Dict[str, Any]],
+        turn_number: Optional[int] = None,
+    ) -> Dict[str, Any]:
+        """
+        Log current context window usage.
+
+        Args:
+            messages: Current conversation messages
+            turn_number: Optional turn number to display
+
+        Returns:
+            Dict with usage stats: {
+                "current_tokens": int,
+                "max_tokens": int,
+                "usage_percent": float,
+                "should_compress": bool,
+                "target_tokens": int
+            }
+        """
+        if not self.enabled:
+            return {}
+
+        # Estimate tokens in current context
+        current_tokens = self.calculator.estimate_tokens(messages)
+        usage_percent = current_tokens / self.context_window
+        should_compress = usage_percent >= self.trigger_threshold
+        target_tokens = int(self.context_window * self.target_ratio)
+
+        # Build log message
+        turn_str = f" (Turn {turn_number})" if turn_number is not None else ""
+        status_emoji = "📊"
+
+        if usage_percent >= self.trigger_threshold:
+            status_emoji = "⚠️"
+            logger.warning(
+                f"{status_emoji} Context Window{turn_str}: " f"{current_tokens:,} / {self.context_window:,} tokens " f"({usage_percent*100:.1f}%) - Approaching limit!",
+            )
+            logger.warning(
+                f"   Compression threshold reached. Target after compression: " f"{target_tokens:,} tokens ({self.target_ratio*100:.0f}%)",
+            )
+        elif usage_percent >= 0.50:
+            logger.info(
+                f"{status_emoji} Context Window{turn_str}: " f"{current_tokens:,} / {self.context_window:,} tokens " f"({usage_percent*100:.1f}%)",
+            )
+        else:
+            logger.info(
+                f"{status_emoji} Context Window{turn_str}: " f"{current_tokens:,} / {self.context_window:,} tokens " f"({usage_percent*100:.1f}%)",
+            )
+
+        return {
+            "current_tokens": current_tokens,
+            "max_tokens": self.context_window,
+            "usage_percent": usage_percent,
+            "should_compress": should_compress,
+            "target_tokens": target_tokens,
+            "trigger_threshold": self.trigger_threshold,
+            "target_ratio": self.target_ratio,
+        }
+
+    def log_turn_summary(
+        self,
+        input_tokens: int,
+        output_tokens: int,
+        turn_number: Optional[int] = None,
+    ):
+        """
+        Log summary for a single turn.
+
+        Args:
+            input_tokens: Input tokens for this turn
+            output_tokens: Output tokens for this turn
+            turn_number: Optional turn number
+        """
+        if not self.enabled:
+            return
+
+        self.total_input_tokens += input_tokens
+        self.total_output_tokens += output_tokens
+        self.turn_count += 1
+
+        turn_str = f" {turn_number}" if turn_number is not None else f" {self.turn_count}"
+
+        logger.info(
+            f"Turn{turn_str} tokens: {input_tokens:,} input + {output_tokens:,} output = " f"{input_tokens + output_tokens:,} total",
+        )
+
+    def log_session_summary(self):
+        """Log overall session summary."""
+        if not self.enabled or self.turn_count == 0:
+            return
+
+        total_tokens = self.total_input_tokens + self.total_output_tokens
+        avg_per_turn = total_tokens / self.turn_count if self.turn_count > 0 else 0
+
+        logger.info("=" * 70)
+        logger.info("📊 Session Summary:")
+        logger.info(f"   Total turns: {self.turn_count}")
+        logger.info(f"   Total input tokens: {self.total_input_tokens:,}")
+        logger.info(f"   Total output tokens: {self.total_output_tokens:,}")
+        logger.info(f"   Total tokens: {total_tokens:,}")
+        logger.info(f"   Average per turn: {avg_per_turn:,.0f} tokens")
+        logger.info(f"   Context window: {self.context_window:,} tokens")
+        logger.info(f"   Peak usage: {(total_tokens/self.context_window)*100:.1f}%")
+        logger.info("=" * 70)
+
+    def get_stats(self) -> Dict[str, Any]:
+        """Get current monitoring stats."""
+        total_tokens = self.total_input_tokens + self.total_output_tokens
+
+        return {
+            "turn_count": self.turn_count,
+            "total_input_tokens": self.total_input_tokens,
+            "total_output_tokens": self.total_output_tokens,
+            "total_tokens": total_tokens,
+            "context_window": self.context_window,
+            "avg_tokens_per_turn": total_tokens / self.turn_count if self.turn_count > 0 else 0,
+            "peak_usage_percent": total_tokens / self.context_window if self.context_window > 0 else 0,
+        }
+
+
+def create_monitor_from_config(
+    config: Dict[str, Any],
+    model_name: str,
+    provider: str = "openai",
+) -> ContextWindowMonitor:
+    """
+    Create a context window monitor from YAML config.
+
+    Args:
+        config: Config dict (should have 'memory' section)
+        model_name: Model name
+        provider: Provider name
+
+    Returns:
+        ContextWindowMonitor instance
+    """
+    memory_config = config.get("memory", {})
+    compression_config = memory_config.get("compression", {})
+
+    trigger_threshold = compression_config.get("trigger_threshold", 0.75)
+    target_ratio = compression_config.get("target_ratio", 0.40)
+    enabled = memory_config.get("enabled", True)
+
+    return ContextWindowMonitor(
+        model_name=model_name,
+        provider=provider,
+        trigger_threshold=trigger_threshold,
+        target_ratio=target_ratio,
+        enabled=enabled,
+    )

massgen/memory/_conversation.py

@@ -0,0 +1,255 @@
+# -*- coding: utf-8 -*-
+"""
+Conversation memory implementation for MassGen.
+
+This module provides in-memory storage for conversation messages, optimized
+for quick access during active chat sessions.
+"""
+
+import uuid
+from typing import Any, Dict, Iterable, List, Optional, Union
+
+from ._base import MemoryBase
+
+
+class ConversationMemory(MemoryBase):
+    """
+    In-memory storage for conversation messages.
+
+    This memory type is designed for short-term storage of ongoing conversations.
+    It keeps messages in a simple list structure for fast access and iteration.
+
+    Features:
+    - Fast in-memory access
+    - Duplicate detection based on message IDs
+    - Index-based deletion
+    - State serialization for session persistence
+
+    Example:
+        >>> memory = ConversationMemory()
+        >>> await memory.add({"role": "user", "content": "Hello"})
+        >>> messages = await memory.get_messages()
+        >>> print(len(messages))  # 1
+    """
+
+    def __init__(self) -> None:
+        """Initialize an empty conversation memory."""
+        super().__init__()
+        self.messages: List[Dict[str, Any]] = []
+
+    def state_dict(self) -> Dict[str, Any]:
+        """
+        Serialize memory state to a dictionary.
+
+        Returns:
+            Dictionary with 'messages' key containing all stored messages
+        """
+        return {
+            "messages": [msg.copy() for msg in self.messages],
+        }
+
+    def load_state_dict(
+        self,
+        state_dict: Dict[str, Any],
+        strict: bool = True,
+    ) -> None:
+        """
+        Load memory state from a serialized dictionary.
+
+        Args:
+            state_dict: Dictionary containing 'messages' key
+            strict: If True, validates the state dictionary structure
+
+        Raises:
+            ValueError: If strict=True and state_dict is invalid
+        """
+        if strict and "messages" not in state_dict:
+            raise ValueError(
+                "State dictionary must contain 'messages' key when strict=True",
+            )
+
+        self.messages = []
+        for msg_data in state_dict.get("messages", []):
+            # Ensure each message is a proper dictionary
+            if isinstance(msg_data, dict):
+                self.messages.append(msg_data.copy())
+
+    async def size(self) -> int:
+        """
+        Get the number of messages in memory.
+
+        Returns:
+            Count of stored messages
+        """
+        return len(self.messages)
+
+    async def retrieve(self, *args: Any, **kwargs: Any) -> None:
+        """
+        Retrieve is not supported for basic conversation memory.
+
+        Use get_messages() to access all messages directly.
+
+        Raises:
+            NotImplementedError: Always, as basic retrieval is not supported
+        """
+        raise NotImplementedError(
+            f"The retrieve method is not implemented in {self.__class__.__name__}. " "Use get_messages() to access conversation history directly.",
+        )
+
+    async def delete(self, index: Union[Iterable, int]) -> None:
+        """
+        Delete message(s) by index position.
+
+        Args:
+            index: Single index or iterable of indices to delete
+
+        Raises:
+            IndexError: If any index is out of range
+
+        Example:
+            >>> await memory.delete(0)  # Delete first message
+            >>> await memory.delete([1, 3, 5])  # Delete multiple messages
+        """
+        if isinstance(index, int):
+            index = [index]
+
+        # Validate all indices first
+        invalid_indices = [i for i in index if i < 0 or i >= len(self.messages)]
+
+        if invalid_indices:
+            raise IndexError(
+                f"The following indices do not exist: {invalid_indices}. " f"Valid range is 0-{len(self.messages) - 1}",
+            )
+
+        # Create new list excluding deleted indices
+        self.messages = [msg for idx, msg in enumerate(self.messages) if idx not in index]
+
+    async def add(
+        self,
+        messages: Union[List[Dict[str, Any]], Dict[str, Any], None],
+        allow_duplicates: bool = False,
+    ) -> None:
+        """
+        Add one or more messages to the conversation memory.
+
+        Args:
+            messages: Single message dict or list of message dicts to add.
+                     Each message should have at minimum a 'role' and 'content'.
+            allow_duplicates: If False, skip messages with duplicate IDs
+
+        Raises:
+            TypeError: If messages are not in the expected format
+
+        Example:
+            >>> # Add single message
+            >>> await memory.add({"role": "user", "content": "Hello"})
+            >>>
+            >>> # Add multiple messages
+            >>> await memory.add([
+            ...     {"role": "user", "content": "Hi"},
+            ...     {"role": "assistant", "content": "Hello!"}
+            ... ])
+        """
+        if messages is None:
+            return
+
+        # Normalize to list
+        if isinstance(messages, dict):
+            messages = [messages]
+
+        if not isinstance(messages, list):
+            raise TypeError(
+                f"Messages should be a list of dicts or a single dict, " f"but got {type(messages)}",
+            )
+
+        # Validate each message
+        for msg in messages:
+            if not isinstance(msg, dict):
+                raise TypeError(
+                    f"Each message should be a dictionary, but got {type(msg)}",
+                )
+
+        # Add message IDs if not present (for duplicate detection)
+        processed_messages = []
+        for msg in messages:
+            msg_copy = msg.copy()
+            if "id" not in msg_copy:
+                msg_copy["id"] = f"msg_{uuid.uuid4().hex[:12]}"
+            processed_messages.append(msg_copy)
+
+        # Filter duplicates if needed
+        if not allow_duplicates:
+            existing_ids = {msg.get("id") for msg in self.messages if "id" in msg}
+            processed_messages = [msg for msg in processed_messages if msg.get("id") not in existing_ids]
+
+        self.messages.extend(processed_messages)
+
+    async def get_messages(self, limit: Optional[int] = None) -> List[Dict[str, Any]]:
+        """
+        Get all messages in the conversation.
+
+        Args:
+            limit: Optional limit on number of most recent messages to return
+
+        Returns:
+            List of message dictionaries (copies, not references)
+
+        Example:
+            >>> # Get all messages
+            >>> all_msgs = await memory.get_messages()
+            >>>
+            >>> # Get last 10 messages
+            >>> recent = await memory.get_messages(limit=10)
+        """
+        if limit is not None and limit > 0:
+            return [msg.copy() for msg in self.messages[-limit:]]
+        return [msg.copy() for msg in self.messages]
+
+    async def clear(self) -> None:
+        """
+        Remove all messages from memory.
+
+        Example:
+            >>> await memory.clear()
+            >>> assert await memory.size() == 0
+        """
+        self.messages = []
+
+    async def get_last_message(self) -> Optional[Dict[str, Any]]:
+        """
+        Get the most recent message.
+
+        Returns:
+            Last message dictionary, or None if memory is empty
+        """
+        if not self.messages:
+            return None
+        return self.messages[-1].copy()
+
+    async def get_messages_by_role(self, role: str) -> List[Dict[str, Any]]:
+        """
+        Filter messages by role.
+
+        Args:
+            role: Role to filter by (e.g., 'user', 'assistant', 'system')
+
+        Returns:
+            List of messages with matching role
+        """
+        return [msg.copy() for msg in self.messages if msg.get("role") == role]
+
+    async def truncate_to_size(self, max_messages: int) -> None:
+        """
+        Keep only the most recent messages up to max_messages.
+
+        This is useful for managing memory usage in long conversations.
+
+        Args:
+            max_messages: Maximum number of messages to keep
+
+        Example:
+            >>> # Keep only last 100 messages
+            >>> await memory.truncate_to_size(100)
+        """
+        if max_messages < len(self.messages):
+            self.messages = self.messages[-max_messages:]