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

aiecs/domain/agent/persistence.py

@@ -0,0 +1,287 @@
+"""
+Agent Persistence
+
+Interfaces and implementations for saving/loading agent state.
+"""
+
+import logging
+import json
+from typing import Dict, Any, Optional, Protocol
+from abc import ABC, abstractmethod
+from datetime import datetime
+
+from .base_agent import BaseAIAgent
+from .exceptions import SerializationError
+
+logger = logging.getLogger(__name__)
+
+
+class AgentPersistence(Protocol):
+    """Protocol for agent persistence implementations."""
+
+    async def save(self, agent: BaseAIAgent) -> None:
+        """
+        Save agent state.
+
+        Args:
+            agent: Agent to save
+        """
+        ...
+
+    async def load(self, agent_id: str) -> Dict[str, Any]:
+        """
+        Load agent state.
+
+        Args:
+            agent_id: Agent identifier
+
+        Returns:
+            Agent state dictionary
+        """
+        ...
+
+    async def exists(self, agent_id: str) -> bool:
+        """
+        Check if agent state exists.
+
+        Args:
+            agent_id: Agent identifier
+
+        Returns:
+            True if exists
+        """
+        ...
+
+    async def delete(self, agent_id: str) -> None:
+        """
+        Delete agent state.
+
+        Args:
+            agent_id: Agent identifier
+        """
+        ...
+
+
+class InMemoryPersistence:
+    """In-memory agent persistence (for testing/development)."""
+
+    def __init__(self):
+        """Initialize in-memory storage."""
+        self._storage: Dict[str, Dict[str, Any]] = {}
+        logger.info("InMemoryPersistence initialized")
+
+    async def save(self, agent: BaseAIAgent) -> None:
+        """Save agent state to memory."""
+        try:
+            state = agent.to_dict()
+            # Convert any remaining datetime objects to ISO strings
+            state = self._serialize_datetimes(state)
+            self._storage[agent.agent_id] = {
+                "state": state,
+                "saved_at": datetime.utcnow().isoformat(),
+            }
+            logger.debug(f"Agent {agent.agent_id} saved to memory")
+        except Exception as e:
+            logger.error(f"Failed to save agent {agent.agent_id}: {e}")
+            raise SerializationError(f"Failed to save agent: {str(e)}")
+
+    def _serialize_datetimes(self, obj: Any) -> Any:
+        """Recursively serialize datetime objects to ISO strings."""
+        import json
+        if isinstance(obj, dict):
+            return {k: self._serialize_datetimes(v) for k, v in obj.items()}
+        elif isinstance(obj, list):
+            return [self._serialize_datetimes(item) for item in obj]
+        elif isinstance(obj, datetime):
+            return obj.isoformat()
+        else:
+            return obj
+
+    async def load(self, agent_id: str) -> Dict[str, Any]:
+        """Load agent state from memory."""
+        if agent_id not in self._storage:
+            raise KeyError(f"Agent {agent_id} not found in storage")
+        
+        data = self._storage[agent_id]
+        logger.debug(f"Agent {agent_id} loaded from memory")
+        return data["state"]
+
+    async def exists(self, agent_id: str) -> bool:
+        """Check if agent exists in memory."""
+        return agent_id in self._storage
+
+    async def delete(self, agent_id: str) -> None:
+        """Delete agent from memory."""
+        if agent_id in self._storage:
+            del self._storage[agent_id]
+            logger.debug(f"Agent {agent_id} deleted from memory")
+
+    def clear(self) -> None:
+        """Clear all stored agents."""
+        self._storage.clear()
+        logger.info("InMemoryPersistence cleared")
+
+
+class FilePersistence:
+    """File-based agent persistence."""
+
+    def __init__(self, base_path: str = "./agent_states"):
+        """
+        Initialize file-based storage.
+
+        Args:
+            base_path: Base directory for agent states
+        """
+        import os
+        
+        self.base_path = base_path
+        os.makedirs(base_path, exist_ok=True)
+        logger.info(f"FilePersistence initialized with base_path: {base_path}")
+
+    def _get_file_path(self, agent_id: str) -> str:
+        """Get file path for agent."""
+        import os
+        # Sanitize agent_id for filesystem
+        safe_id = agent_id.replace("/", "_").replace("\\", "_")
+        return os.path.join(self.base_path, f"{safe_id}.json")
+
+    async def save(self, agent: BaseAIAgent) -> None:
+        """Save agent state to file."""
+        try:
+            state = agent.to_dict()
+            # Convert any remaining datetime objects to ISO strings for JSON serialization
+            state = self._serialize_datetimes(state)
+            file_path = self._get_file_path(agent.agent_id)
+            
+            data = {
+                "state": state,
+                "saved_at": datetime.utcnow().isoformat(),
+            }
+            
+            with open(file_path, 'w') as f:
+                json.dump(data, f, indent=2, default=str)  # default=str handles any remaining non-serializable objects
+            
+            logger.debug(f"Agent {agent.agent_id} saved to {file_path}")
+        except Exception as e:
+            logger.error(f"Failed to save agent {agent.agent_id}: {e}")
+            raise SerializationError(f"Failed to save agent: {str(e)}")
+
+    def _serialize_datetimes(self, obj: Any) -> Any:
+        """Recursively serialize datetime objects to ISO strings."""
+        if isinstance(obj, dict):
+            return {k: self._serialize_datetimes(v) for k, v in obj.items()}
+        elif isinstance(obj, list):
+            return [self._serialize_datetimes(item) for item in obj]
+        elif isinstance(obj, datetime):
+            return obj.isoformat()
+        else:
+            return obj
+
+    async def load(self, agent_id: str) -> Dict[str, Any]:
+        """Load agent state from file."""
+        file_path = self._get_file_path(agent_id)
+        
+        try:
+            with open(file_path, 'r') as f:
+                data = json.load(f)
+            
+            logger.debug(f"Agent {agent_id} loaded from {file_path}")
+            return data["state"]
+        except FileNotFoundError:
+            raise KeyError(f"Agent {agent_id} not found in storage")
+        except Exception as e:
+            logger.error(f"Failed to load agent {agent_id}: {e}")
+            raise SerializationError(f"Failed to load agent: {str(e)}")
+
+    async def exists(self, agent_id: str) -> bool:
+        """Check if agent file exists."""
+        import os
+        file_path = self._get_file_path(agent_id)
+        return os.path.exists(file_path)
+
+    async def delete(self, agent_id: str) -> None:
+        """Delete agent file."""
+        import os
+        file_path = self._get_file_path(agent_id)
+        
+        try:
+            if os.path.exists(file_path):
+                os.remove(file_path)
+                logger.debug(f"Agent {agent_id} deleted from {file_path}")
+        except Exception as e:
+            logger.error(f"Failed to delete agent {agent_id}: {e}")
+            raise
+
+
+class AgentStateSerializer:
+    """
+    Helper class for serializing/deserializing agent state.
+    
+    Handles complex types that need special serialization.
+    """
+
+    @staticmethod
+    def serialize(agent: BaseAIAgent) -> Dict[str, Any]:
+        """
+        Serialize agent to dictionary.
+
+        Args:
+            agent: Agent to serialize
+
+        Returns:
+            Serialized state dictionary
+        """
+        return agent.to_dict()
+
+    @staticmethod
+    def deserialize(data: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Deserialize agent state.
+
+        Args:
+            data: Serialized state
+
+        Returns:
+            Deserialized state dictionary
+        
+        Note: This returns a state dictionary, not an agent instance.
+        Agent reconstruction requires the appropriate agent class.
+        """
+        # In the future, this could handle type conversion, validation, etc.
+        return data
+
+
+# Global persistence instance
+_global_persistence: Optional[AgentPersistence] = None
+
+
+def get_global_persistence() -> AgentPersistence:
+    """
+    Get or create global persistence instance.
+
+    Returns:
+        Global persistence instance (defaults to InMemoryPersistence)
+    """
+    global _global_persistence
+    if _global_persistence is None:
+        _global_persistence = InMemoryPersistence()
+    return _global_persistence
+
+
+def set_global_persistence(persistence: AgentPersistence) -> None:
+    """
+    Set global persistence instance.
+
+    Args:
+        persistence: Persistence implementation to use
+    """
+    global _global_persistence
+    _global_persistence = persistence
+    logger.info(f"Global persistence set to {type(persistence).__name__}")
+
+
+def reset_global_persistence() -> None:
+    """Reset global persistence (primarily for testing)."""
+    global _global_persistence
+    _global_persistence = None
+

aiecs/domain/agent/prompts/__init__.py

@@ -0,0 +1,25 @@
+"""
+Prompt Templates Module
+
+Native prompt template system replacing LangChain templates.
+"""
+
+from .template import PromptTemplate, ChatPromptTemplate, MessageTemplate, TemplateMissingVariableError
+from .builder import MessageBuilder
+from .formatters import (
+    format_conversation_history,
+    format_tool_result,
+    truncate_context,
+)
+
+__all__ = [
+    "PromptTemplate",
+    "ChatPromptTemplate",
+    "MessageTemplate",
+    "TemplateMissingVariableError",
+    "MessageBuilder",
+    "format_conversation_history",
+    "format_tool_result",
+    "truncate_context",
+]
+

aiecs/domain/agent/prompts/builder.py

@@ -0,0 +1,164 @@
+"""
+Message Builder
+
+Helper for constructing LLMMessage lists.
+"""
+
+from typing import List, Dict, Any, Optional
+from aiecs.llm import LLMMessage
+
+
+class MessageBuilder:
+    """
+    Builder for constructing LLM message sequences.
+
+    Example:
+        builder = MessageBuilder()
+        builder.add_system("You are a helpful assistant")
+        builder.add_user("What is AI?")
+        messages = builder.build()
+    """
+
+    def __init__(self):
+        """Initialize message builder."""
+        self.messages: List[LLMMessage] = []
+
+    def add_system(self, content: str) -> "MessageBuilder":
+        """
+        Add system message.
+
+        Args:
+            content: Message content
+
+        Returns:
+            Self for chaining
+        """
+        self.messages.append(LLMMessage(role="system", content=content))
+        return self
+
+    def add_user(self, content: str) -> "MessageBuilder":
+        """
+        Add user message.
+
+        Args:
+            content: Message content
+
+        Returns:
+            Self for chaining
+        """
+        self.messages.append(LLMMessage(role="user", content=content))
+        return self
+
+    def add_assistant(self, content: str) -> "MessageBuilder":
+        """
+        Add assistant message.
+
+        Args:
+            content: Message content
+
+        Returns:
+            Self for chaining
+        """
+        self.messages.append(LLMMessage(role="assistant", content=content))
+        return self
+
+    def add_message(self, role: str, content: str) -> "MessageBuilder":
+        """
+        Add message with custom role.
+
+        Args:
+            role: Message role
+            content: Message content
+
+        Returns:
+            Self for chaining
+        """
+        self.messages.append(LLMMessage(role=role, content=content))
+        return self
+
+    def add_messages(self, messages: List[LLMMessage]) -> "MessageBuilder":
+        """
+        Add multiple messages.
+
+        Args:
+            messages: List of messages to add
+
+        Returns:
+            Self for chaining
+        """
+        self.messages.extend(messages)
+        return self
+
+    def add_context(self, context: Dict[str, Any], prefix: str = "Context:") -> "MessageBuilder":
+        """
+        Add context as a system message.
+
+        Args:
+            context: Context dictionary
+            prefix: Prefix for context message
+
+        Returns:
+            Self for chaining
+        """
+        context_str = self._format_context(context)
+        if context_str:
+            self.add_system(f"{prefix}\n{context_str}")
+        return self
+
+    def add_conversation_history(
+        self,
+        history: List[Dict[str, str]],
+        max_messages: Optional[int] = None
+    ) -> "MessageBuilder":
+        """
+        Add conversation history.
+
+        Args:
+            history: List of {role, content} dicts
+            max_messages: Optional limit on number of messages
+
+        Returns:
+            Self for chaining
+        """
+        if max_messages:
+            history = history[-max_messages:]
+
+        for msg in history:
+            self.add_message(msg.get("role", "user"), msg.get("content", ""))
+
+        return self
+
+    def clear(self) -> "MessageBuilder":
+        """
+        Clear all messages.
+
+        Returns:
+            Self for chaining
+        """
+        self.messages.clear()
+        return self
+
+    def build(self) -> List[LLMMessage]:
+        """
+        Build and return message list.
+
+        Returns:
+            List of LLMMessage instances
+        """
+        return self.messages.copy()
+
+    def _format_context(self, context: Dict[str, Any]) -> str:
+        """Format context dictionary as string."""
+        lines = []
+        for key, value in context.items():
+            if not key.startswith('_') and value is not None:
+                lines.append(f"{key}: {value}")
+        return "\n".join(lines) if lines else ""
+
+    def __len__(self) -> int:
+        """Get number of messages."""
+        return len(self.messages)
+
+    def __repr__(self) -> str:
+        return f"MessageBuilder(messages={len(self.messages)})"
+

aiecs/domain/agent/prompts/formatters.py

@@ -0,0 +1,192 @@
+"""
+Prompt Formatters
+
+Utilities for formatting prompts with context.
+"""
+
+from typing import List, Dict, Any, Optional
+from aiecs.llm import LLMMessage
+
+
+def format_conversation_history(
+    history: List[LLMMessage],
+    max_messages: Optional[int] = None,
+    format_style: str = "compact"
+) -> str:
+    """
+    Format conversation history as string.
+
+    Args:
+        history: List of LLMMessage instances
+        max_messages: Optional limit on number of messages
+        format_style: "compact" or "detailed"
+
+    Returns:
+        Formatted conversation string
+    """
+    if max_messages:
+        history = history[-max_messages:]
+
+    if format_style == "compact":
+        lines = []
+        for msg in history:
+            lines.append(f"{msg.role.upper()}: {msg.content}")
+        return "\n".join(lines)
+
+    elif format_style == "detailed":
+        lines = []
+        for i, msg in enumerate(history):
+            lines.append(f"[{i+1}] {msg.role.upper()}")
+            lines.append(msg.content)
+            lines.append("")  # Empty line between messages
+        return "\n".join(lines)
+
+    else:
+        raise ValueError(f"Unknown format_style: {format_style}")
+
+
+def format_tool_result(
+    tool_name: str,
+    result: Any,
+    success: bool = True,
+    error: Optional[str] = None
+) -> str:
+    """
+    Format tool execution result.
+
+    Args:
+        tool_name: Tool name
+        result: Tool result (if successful)
+        success: Whether execution succeeded
+        error: Error message (if failed)
+
+    Returns:
+        Formatted tool result string
+    """
+    if success:
+        return f"Tool '{tool_name}' returned:\n{result}"
+    else:
+        return f"Tool '{tool_name}' failed: {error}"
+
+
+def truncate_context(
+    text: str,
+    max_length: int,
+    strategy: str = "middle",
+    placeholder: str = "..."
+) -> str:
+    """
+    Truncate text to fit within max_length.
+
+    Args:
+        text: Text to truncate
+        max_length: Maximum length
+        strategy: "start", "middle", or "end"
+        placeholder: Placeholder for truncated content
+
+    Returns:
+        Truncated text
+    """
+    if len(text) <= max_length:
+        return text
+
+    if strategy == "end":
+        # Keep start, truncate end
+        return text[:max_length - len(placeholder)] + placeholder
+
+    elif strategy == "start":
+        # Truncate start, keep end
+        return placeholder + text[-(max_length - len(placeholder)):]
+
+    elif strategy == "middle":
+        # Keep start and end, truncate middle
+        half = (max_length - len(placeholder)) // 2
+        return text[:half] + placeholder + text[-half:]
+
+    else:
+        raise ValueError(f"Unknown strategy: {strategy}")
+
+
+def format_list_items(items: List[str], style: str = "bullets") -> str:
+    """
+    Format list items.
+
+    Args:
+        items: List of items
+        style: "bullets", "numbered", or "compact"
+
+    Returns:
+        Formatted list string
+    """
+    if style == "bullets":
+        return "\n".join(f"• {item}" for item in items)
+
+    elif style == "numbered":
+        return "\n".join(f"{i+1}. {item}" for i, item in enumerate(items))
+
+    elif style == "compact":
+        return ", ".join(items)
+
+    else:
+        raise ValueError(f"Unknown style: {style}")
+
+
+def format_key_value_pairs(
+    data: Dict[str, Any],
+    separator: str = ": ",
+    exclude_keys: Optional[List[str]] = None
+) -> str:
+    """
+    Format dictionary as key-value pairs.
+
+    Args:
+        data: Dictionary to format
+        separator: Separator between key and value
+        exclude_keys: Keys to exclude
+
+    Returns:
+        Formatted string
+    """
+    exclude_keys = exclude_keys or []
+    lines = []
+
+    for key, value in data.items():
+        if key in exclude_keys or key.startswith('_'):
+            continue
+        lines.append(f"{key}{separator}{value}")
+
+    return "\n".join(lines)
+
+
+def inject_context_in_prompt(
+    prompt: str,
+    context: Dict[str, Any],
+    context_marker: str = "{context}"
+) -> str:
+    """
+    Inject context into prompt at marker position.
+
+    Args:
+        prompt: Prompt template with context marker
+        context: Context dictionary
+        context_marker: Marker to replace with context
+
+    Returns:
+        Prompt with context injected
+    """
+    context_str = format_key_value_pairs(context)
+    return prompt.replace(context_marker, context_str)
+
+
+def estimate_token_count(text: str) -> int:
+    """
+    Rough estimation of token count.
+
+    Args:
+        text: Text to estimate
+
+    Returns:
+        Estimated token count (4 chars ≈ 1 token)
+    """
+    return len(text) // 4
+