aiecs 1.3.8__py3-none-any.whl → 1.4.0__py3-none-any.whl

aiecs/domain/agent/integration/context_compressor.py

@@ -0,0 +1,219 @@
+"""
+Context Compression
+
+Smart context compression for token limits.
+"""
+
+import logging
+from typing import List, Dict, Any, Optional
+from enum import Enum
+
+from aiecs.llm import LLMMessage
+
+logger = logging.getLogger(__name__)
+
+
+class CompressionStrategy(Enum):
+    """Context compression strategies."""
+    TRUNCATE_MIDDLE = "truncate_middle"
+    TRUNCATE_START = "truncate_start"
+    PRESERVE_RECENT = "preserve_recent"
+    SUMMARIZE = "summarize"
+
+
+class ContextCompressor:
+    """
+    Smart context compression for managing token limits.
+
+    Example:
+        compressor = ContextCompressor(max_tokens=4000)
+        compressed = compressor.compress_messages(messages)
+    """
+
+    def __init__(
+        self,
+        max_tokens: int = 4000,
+        strategy: CompressionStrategy = CompressionStrategy.PRESERVE_RECENT,
+        preserve_system: bool = True,
+    ):
+        """
+        Initialize context compressor.
+
+        Args:
+            max_tokens: Maximum token limit
+            strategy: Compression strategy
+            preserve_system: Always preserve system messages
+        """
+        self.max_tokens = max_tokens
+        self.strategy = strategy
+        self.preserve_system = preserve_system
+
+    def compress_messages(
+        self,
+        messages: List[LLMMessage],
+        priority_indices: Optional[List[int]] = None
+    ) -> List[LLMMessage]:
+        """
+        Compress message list to fit within token limit.
+
+        Args:
+            messages: List of messages
+            priority_indices: Optional indices of high-priority messages
+
+        Returns:
+            Compressed message list
+        """
+        # Estimate tokens
+        total_tokens = self._estimate_tokens(messages)
+
+        if total_tokens <= self.max_tokens:
+            return messages
+
+        logger.debug(
+            f"Compressing {len(messages)} messages from ~{total_tokens} to ~{self.max_tokens} tokens"
+        )
+
+        # Apply compression strategy
+        if self.strategy == CompressionStrategy.PRESERVE_RECENT:
+            return self._compress_preserve_recent(messages, priority_indices)
+        elif self.strategy == CompressionStrategy.TRUNCATE_MIDDLE:
+            return self._compress_truncate_middle(messages, priority_indices)
+        elif self.strategy == CompressionStrategy.TRUNCATE_START:
+            return self._compress_truncate_start(messages)
+        else:
+            # Default: preserve recent
+            return self._compress_preserve_recent(messages, priority_indices)
+
+    def _compress_preserve_recent(
+        self,
+        messages: List[LLMMessage],
+        priority_indices: Optional[List[int]]
+    ) -> List[LLMMessage]:
+        """Preserve recent messages and priority messages."""
+        priority_indices = set(priority_indices or [])
+        compressed = []
+
+        # Always include system messages if enabled
+        if self.preserve_system:
+            system_msgs = [msg for msg in messages if msg.role == "system"]
+            compressed.extend(system_msgs)
+
+        # Calculate remaining budget
+        remaining_tokens = self.max_tokens - self._estimate_tokens(compressed)
+
+        # Add priority messages
+        for idx in priority_indices:
+            if idx < len(messages) and messages[idx] not in compressed:
+                msg_tokens = self._estimate_tokens([messages[idx]])
+                if msg_tokens <= remaining_tokens:
+                    compressed.append(messages[idx])
+                    remaining_tokens -= msg_tokens
+
+        # Add recent messages (from end)
+        for msg in reversed(messages):
+            if msg not in compressed:
+                msg_tokens = self._estimate_tokens([msg])
+                if msg_tokens <= remaining_tokens:
+                    compressed.insert(len(compressed), msg)
+                    remaining_tokens -= msg_tokens
+                else:
+                    break
+
+        return compressed
+
+    def _compress_truncate_middle(
+        self,
+        messages: List[LLMMessage],
+        priority_indices: Optional[List[int]]
+    ) -> List[LLMMessage]:
+        """Keep start and end messages, truncate middle."""
+        if len(messages) <= 4:
+            return messages
+
+        # Keep first 2 and last 2 by default
+        keep_start = 2
+        keep_end = 2
+
+        # Adjust based on token budget
+        start_msgs = messages[:keep_start]
+        end_msgs = messages[-keep_end:]
+
+        compressed = start_msgs + [
+            LLMMessage(role="system", content="[... conversation history compressed ...]")
+        ] + end_msgs
+
+        return compressed
+
+    def _compress_truncate_start(self, messages: List[LLMMessage]) -> List[LLMMessage]:
+        """Keep recent messages, truncate start."""
+        compressed = []
+        remaining_tokens = self.max_tokens
+
+        # Process from end
+        for msg in reversed(messages):
+            msg_tokens = self._estimate_tokens([msg])
+            if msg_tokens <= remaining_tokens:
+                compressed.insert(0, msg)
+                remaining_tokens -= msg_tokens
+            else:
+                break
+
+        return compressed
+
+    def _estimate_tokens(self, messages: List[LLMMessage]) -> int:
+        """
+        Estimate token count for messages.
+
+        Args:
+            messages: List of messages
+
+        Returns:
+            Estimated token count
+        """
+        total_chars = sum(len(msg.content) for msg in messages)
+        # Rough estimate: 4 chars ≈ 1 token
+        return total_chars // 4
+
+    def compress_text(self, text: str, max_tokens: int) -> str:
+        """
+        Compress text to fit within token limit.
+
+        Args:
+            text: Text to compress
+            max_tokens: Maximum tokens
+
+        Returns:
+            Compressed text
+        """
+        estimated_tokens = len(text) // 4
+
+        if estimated_tokens <= max_tokens:
+            return text
+
+        # Truncate to fit
+        max_chars = max_tokens * 4
+        if len(text) <= max_chars:
+            return text
+
+        return text[:max_chars - 20] + "... [truncated]"
+
+
+def compress_messages(
+    messages: List[LLMMessage],
+    max_tokens: int = 4000,
+    strategy: CompressionStrategy = CompressionStrategy.PRESERVE_RECENT
+) -> List[LLMMessage]:
+    """
+    Convenience function for compressing messages.
+
+    Args:
+        messages: List of messages
+        max_tokens: Maximum token limit
+        strategy: Compression strategy
+
+    Returns:
+        Compressed message list
+    """
+    compressor = ContextCompressor(max_tokens=max_tokens, strategy=strategy)
+    return compressor.compress_messages(messages)
+

aiecs/domain/agent/integration/context_engine_adapter.py

@@ -0,0 +1,258 @@
+"""
+ContextEngine Adapter
+
+Adapter for integrating agent persistence with AIECS ContextEngine.
+"""
+
+import logging
+import json
+import uuid
+from typing import Dict, Any, Optional, List, TYPE_CHECKING
+from datetime import datetime
+
+if TYPE_CHECKING:
+    from aiecs.domain.context.context_engine import ContextEngine
+
+from aiecs.domain.agent.persistence import AgentPersistence
+from aiecs.domain.agent.base_agent import BaseAIAgent
+
+logger = logging.getLogger(__name__)
+
+
+class ContextEngineAdapter:
+    """
+    Adapter for persisting agent state to ContextEngine.
+
+    Uses ContextEngine's checkpoint system for versioned state storage
+    and TaskContext for session-based state management.
+    """
+
+    def __init__(self, context_engine: "ContextEngine", user_id: str = "system"):
+        """
+        Initialize adapter.
+
+        Args:
+            context_engine: ContextEngine instance
+            user_id: User identifier for session management
+        """
+        if context_engine is None:
+            raise ValueError("ContextEngine instance is required")
+        
+        self.context_engine = context_engine
+        self.user_id = user_id
+        self._agent_state_prefix = "agent_state"
+        self._agent_conversation_prefix = "agent_conversation"
+        logger.info("ContextEngineAdapter initialized")
+
+    async def save_agent_state(
+        self,
+        agent_id: str,
+        state: Dict[str, Any],
+        version: Optional[str] = None
+    ) -> str:
+        """
+        Save agent state to ContextEngine using checkpoint system.
+
+        Args:
+            agent_id: Agent identifier
+            state: Agent state dictionary
+            version: Optional version identifier (auto-generated if None)
+
+        Returns:
+            Version identifier
+        """
+        if version is None:
+            version = str(uuid.uuid4())
+
+        checkpoint_data = {
+            "agent_id": agent_id,
+            "state": state,
+            "timestamp": datetime.utcnow().isoformat(),
+            "version": version
+        }
+
+        # Store as checkpoint (thread_id = agent_id)
+        await self.context_engine.store_checkpoint(
+            thread_id=agent_id,
+            checkpoint_id=version,
+            checkpoint_data=checkpoint_data,
+            metadata={"type": "agent_state", "agent_id": agent_id}
+        )
+
+        logger.debug(f"Saved agent {agent_id} state version {version} to ContextEngine")
+        return version
+
+    async def load_agent_state(
+        self,
+        agent_id: str,
+        version: Optional[str] = None
+    ) -> Optional[Dict[str, Any]]:
+        """
+        Load agent state from ContextEngine.
+
+        Args:
+            agent_id: Agent identifier
+            version: Optional version identifier (loads latest if None)
+
+        Returns:
+            Agent state dictionary or None
+        """
+        checkpoint = await self.context_engine.get_checkpoint(
+            thread_id=agent_id,
+            checkpoint_id=version
+        )
+
+        if checkpoint and "data" in checkpoint:
+            checkpoint_data = checkpoint["data"]
+            if isinstance(checkpoint_data, dict) and "state" in checkpoint_data:
+                logger.debug(f"Loaded agent {agent_id} state version {version or 'latest'}")
+                return checkpoint_data["state"]
+
+        logger.debug(f"No state found for agent {agent_id} version {version or 'latest'}")
+        return None
+
+    async def list_agent_versions(self, agent_id: str) -> List[Dict[str, Any]]:
+        """
+        List all versions of an agent's state.
+
+        Args:
+            agent_id: Agent identifier
+
+        Returns:
+            List of version metadata dictionaries
+        """
+        checkpoints = await self.context_engine.list_checkpoints(thread_id=agent_id)
+        if not checkpoints:
+            return []
+
+        versions = []
+        for checkpoint in checkpoints:
+            # list_checkpoints returns dicts with "data" key containing checkpoint_data
+            if isinstance(checkpoint, dict):
+                data = checkpoint.get("data", {})
+                if isinstance(data, dict) and "version" in data:
+                    versions.append({
+                        "version": data["version"],
+                        "timestamp": data.get("timestamp"),
+                        "metadata": checkpoint.get("metadata", {})
+                    })
+
+        # Sort by timestamp descending
+        versions.sort(key=lambda v: v.get("timestamp", ""), reverse=True)
+        return versions
+
+    async def save_conversation_history(
+        self,
+        session_id: str,
+        messages: List[Dict[str, Any]]
+    ) -> None:
+        """
+        Save conversation history to ContextEngine.
+
+        Args:
+            session_id: Session identifier
+            messages: List of message dictionaries with 'role' and 'content'
+        """
+        # Ensure session exists
+        session = await self.context_engine.get_session(session_id)
+        if not session:
+            await self.context_engine.create_session(
+                session_id=session_id,
+                user_id=self.user_id,
+                metadata={"type": "agent_conversation"}
+            )
+
+        # Store messages using ContextEngine's conversation API
+        for msg in messages:
+            role = msg.get("role", "user")
+            content = msg.get("content", "")
+            metadata = msg.get("metadata", {})
+            
+            await self.context_engine.add_conversation_message(
+                session_id=session_id,
+                role=role,
+                content=content,
+                metadata=metadata
+            )
+
+        logger.debug(f"Saved {len(messages)} messages to session {session_id}")
+
+    async def load_conversation_history(
+        self,
+        session_id: str,
+        limit: int = 50
+    ) -> List[Dict[str, Any]]:
+        """
+        Load conversation history from ContextEngine.
+
+        Args:
+            session_id: Session identifier
+            limit: Maximum number of messages to retrieve
+
+        Returns:
+            List of message dictionaries
+        """
+        messages = await self.context_engine.get_conversation_history(
+            session_id=session_id,
+            limit=limit
+        )
+
+        # Convert ConversationMessage objects to dictionaries
+        result = []
+        for msg in messages:
+            result.append({
+                "role": msg.role,
+                "content": msg.content,
+                "timestamp": msg.timestamp.isoformat() if hasattr(msg.timestamp, 'isoformat') else str(msg.timestamp),
+                "metadata": msg.metadata
+            })
+
+        logger.debug(f"Loaded {len(result)} messages from session {session_id}")
+        return result
+
+    async def delete_agent_state(
+        self,
+        agent_id: str,
+        version: Optional[str] = None
+    ) -> None:
+        """
+        Delete agent state from ContextEngine.
+
+        Args:
+            agent_id: Agent identifier
+            version: Optional version identifier (deletes all if None)
+        """
+        # Note: ContextEngine doesn't have explicit delete for checkpoints
+        # We'll store a tombstone checkpoint or rely on TTL
+        if version:
+            # Store empty state as deletion marker
+            await self.context_engine.store_checkpoint(
+                thread_id=agent_id,
+                checkpoint_id=f"{version}_deleted",
+                checkpoint_data={"deleted": True, "original_version": version},
+                metadata={"type": "deletion_marker"}
+            )
+        logger.debug(f"Marked agent {agent_id} state version {version or 'all'} for deletion")
+
+    # AgentPersistence Protocol implementation
+    async def save(self, agent: BaseAIAgent) -> None:
+        """Save agent state (implements AgentPersistence protocol)."""
+        state = agent.to_dict()
+        await self.save_agent_state(agent.agent_id, state)
+
+    async def load(self, agent_id: str) -> Dict[str, Any]:
+        """Load agent state (implements AgentPersistence protocol)."""
+        state = await self.load_agent_state(agent_id)
+        if state is None:
+            raise KeyError(f"Agent {agent_id} not found in storage")
+        return state
+
+    async def exists(self, agent_id: str) -> bool:
+        """Check if agent state exists (implements AgentPersistence protocol)."""
+        state = await self.load_agent_state(agent_id)
+        return state is not None
+
+    async def delete(self, agent_id: str) -> None:
+        """Delete agent state (implements AgentPersistence protocol)."""
+        await self.delete_agent_state(agent_id)
+

aiecs/domain/agent/integration/retry_policy.py

@@ -0,0 +1,228 @@
+"""
+Enhanced Retry Policy
+
+Sophisticated retry logic with exponential backoff and error classification.
+"""
+
+import asyncio
+import random
+import logging
+from typing import Optional, Callable, Any
+from datetime import datetime
+from enum import Enum
+
+logger = logging.getLogger(__name__)
+
+
+class ErrorType(Enum):
+    """Error types for classification."""
+    RATE_LIMIT = "rate_limit"
+    TIMEOUT = "timeout"
+    SERVER_ERROR = "server_error"
+    CLIENT_ERROR = "client_error"
+    NETWORK_ERROR = "network_error"
+    UNKNOWN = "unknown"
+
+
+class ErrorClassifier:
+    """Classifies errors for retry strategy."""
+
+    @staticmethod
+    def classify(error: Exception) -> ErrorType:
+        """
+        Classify error type.
+
+        Args:
+            error: Exception to classify
+
+        Returns:
+            ErrorType
+        """
+        error_str = str(error).lower()
+        error_type_name = type(error).__name__.lower()
+
+        # Rate limit errors
+        if "rate limit" in error_str or "429" in error_str:
+            return ErrorType.RATE_LIMIT
+
+        # Timeout errors
+        if "timeout" in error_str or "timed out" in error_str:
+            return ErrorType.TIMEOUT
+
+        # Server errors (5xx)
+        if any(code in error_str for code in ["500", "502", "503", "504"]):
+            return ErrorType.SERVER_ERROR
+
+        # Client errors (4xx)
+        if any(code in error_str for code in ["400", "401", "403", "404"]):
+            return ErrorType.CLIENT_ERROR
+
+        # Network errors
+        if any(term in error_type_name for term in ["connection", "network", "socket"]):
+            return ErrorType.NETWORK_ERROR
+
+        return ErrorType.UNKNOWN
+
+    @staticmethod
+    def is_retryable(error_type: ErrorType) -> bool:
+        """
+        Determine if error type should be retried.
+
+        Args:
+            error_type: Error type
+
+        Returns:
+            True if retryable
+        """
+        retryable_types = {
+            ErrorType.RATE_LIMIT,
+            ErrorType.TIMEOUT,
+            ErrorType.SERVER_ERROR,
+            ErrorType.NETWORK_ERROR,
+        }
+        return error_type in retryable_types
+
+
+class EnhancedRetryPolicy:
+    """
+    Enhanced retry policy with exponential backoff and jitter.
+
+    Example:
+        policy = EnhancedRetryPolicy(max_retries=5, base_delay=1.0)
+        result = await policy.execute_with_retry(my_async_function, arg1, arg2)
+    """
+
+    def __init__(
+        self,
+        max_retries: int = 3,
+        base_delay: float = 1.0,
+        max_delay: float = 60.0,
+        exponential_base: float = 2.0,
+        jitter: bool = True,
+    ):
+        """
+        Initialize retry policy.
+
+        Args:
+            max_retries: Maximum number of retry attempts
+            base_delay: Base delay in seconds
+            max_delay: Maximum delay in seconds
+            exponential_base: Base for exponential backoff
+            jitter: Whether to add random jitter
+        """
+        self.max_retries = max_retries
+        self.base_delay = base_delay
+        self.max_delay = max_delay
+        self.exponential_base = exponential_base
+        self.jitter = jitter
+
+    def calculate_delay(self, attempt: int, error_type: ErrorType) -> float:
+        """
+        Calculate delay for retry attempt.
+
+        Args:
+            attempt: Retry attempt number (0-indexed)
+            error_type: Type of error
+
+        Returns:
+            Delay in seconds
+        """
+        # Base exponential backoff
+        delay = min(
+            self.base_delay * (self.exponential_base ** attempt),
+            self.max_delay
+        )
+
+        # Adjust for error type
+        if error_type == ErrorType.RATE_LIMIT:
+            # Longer delay for rate limits
+            delay *= 2
+
+        # Add jitter to prevent thundering herd
+        if self.jitter:
+            delay *= (0.5 + random.random())
+
+        return delay
+
+    async def execute_with_retry(
+        self,
+        func: Callable,
+        *args,
+        **kwargs
+    ) -> Any:
+        """
+        Execute function with retry logic.
+
+        Args:
+            func: Async function to execute
+            *args: Function arguments
+            **kwargs: Function keyword arguments
+
+        Returns:
+            Function result
+
+        Raises:
+            Exception: If all retries exhausted
+        """
+        last_error = None
+
+        for attempt in range(self.max_retries + 1):
+            try:
+                result = await func(*args, **kwargs)
+                
+                # Log success after retries
+                if attempt > 0:
+                    logger.info(f"Succeeded after {attempt} retries")
+                
+                return result
+
+            except Exception as e:
+                last_error = e
+                
+                # Classify error
+                error_type = ErrorClassifier.classify(e)
+                
+                # Check if we should retry
+                if attempt >= self.max_retries:
+                    logger.error(f"Max retries ({self.max_retries}) exhausted")
+                    break
+
+                if not ErrorClassifier.is_retryable(error_type):
+                    logger.error(f"Non-retryable error: {error_type.value}")
+                    break
+
+                # Calculate delay and wait
+                delay = self.calculate_delay(attempt, error_type)
+                logger.warning(
+                    f"Attempt {attempt + 1} failed ({error_type.value}). "
+                    f"Retrying in {delay:.2f}s..."
+                )
+                await asyncio.sleep(delay)
+
+        # All retries exhausted
+        raise last_error
+
+
+async def with_retry(
+    func: Callable,
+    max_retries: int = 3,
+    base_delay: float = 1.0,
+    *args,
+    **kwargs
+) -> Any:
+    """
+    Convenience function for executing with retry.
+
+    Args:
+        func: Async function to execute
+        max_retries: Maximum number of retries
+        base_delay: Base delay in seconds
+        *args: Function arguments
+        **kwargs: Function keyword arguments
+
+    Returns:
+        Function result
+    """
+    policy = EnhancedRetryPolicy(max_retries=max_retries, base_delay=base_delay)
+    return await policy.execute_with_retry(func, *args, **kwargs)
+