aiecs 1.0.1__py3-none-any.whl → 1.7.17__py3-none-any.whl

aiecs/domain/agent/prompts/formatters.py

@@ -0,0 +1,187 @@
+"""
+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

aiecs/domain/agent/prompts/template.py

@@ -0,0 +1,255 @@
+"""
+Prompt Templates
+
+Native template system with variable substitution.
+"""
+
+import re
+from typing import Dict, List, Optional
+from dataclasses import dataclass
+
+from aiecs.llm import LLMMessage
+
+
+class TemplateMissingVariableError(Exception):
+    """Raised when required template variable is missing."""
+
+
+class PromptTemplate:
+    """
+    String-based prompt template with {variable} substitution.
+
+    Example:
+        template = PromptTemplate(
+            "Hello {name}, you are a {role}.",
+            required_variables=["name", "role"]
+        )
+        result = template.format(name="Alice", role="developer")
+    """
+
+    def __init__(
+        self,
+        template: str,
+        required_variables: Optional[List[str]] = None,
+        defaults: Optional[Dict[str, str]] = None,
+    ):
+        """
+        Initialize prompt template.
+
+        Args:
+            template: Template string with {variable} placeholders
+            required_variables: List of required variable names
+            defaults: Default values for optional variables
+        """
+        self.template = template
+        self.required_variables = required_variables or []
+        self.defaults = defaults or {}
+
+        # Extract all variables from template
+        self._extract_variables()
+
+    def _extract_variables(self) -> None:
+        """Extract variable names from template."""
+        # Find all {variable_name} patterns
+        pattern = r"\{(\w+)\}"
+        self.variables = set(re.findall(pattern, self.template))
+
+    def format(self, **kwargs) -> str:
+        """
+        Format template with provided variables.
+
+        Args:
+            **kwargs: Variable values
+
+        Returns:
+            Formatted string
+
+        Raises:
+            TemplateMissingVariableError: If required variable missing
+        """
+        # Check required variables
+        for var in self.required_variables:
+            if var not in kwargs and var not in self.defaults:
+                raise TemplateMissingVariableError(f"Required variable '{var}' not provided")
+
+        # Merge with defaults
+        values = {**self.defaults, **kwargs}
+
+        # Format template
+        try:
+            return self.template.format(**values)
+        except KeyError as e:
+            raise TemplateMissingVariableError(f"Variable {e} not provided and has no default")
+
+    def partial(self, **kwargs) -> "PromptTemplate":
+        """
+        Create a partial template with some variables pre-filled.
+
+        Args:
+            **kwargs: Variable values to pre-fill
+
+        Returns:
+            New PromptTemplate with updated defaults
+        """
+        new_defaults = {**self.defaults, **kwargs}
+        return PromptTemplate(
+            template=self.template,
+            required_variables=self.required_variables,
+            defaults=new_defaults,
+        )
+
+    def __repr__(self) -> str:
+        return f"PromptTemplate(variables={self.variables})"
+
+
+@dataclass
+class MessageTemplate:
+    """Template for a single message."""
+
+    role: str
+    content: str
+
+
+class ChatPromptTemplate:
+    """
+    Multi-message chat template.
+
+    Example:
+        template = ChatPromptTemplate([
+            MessageTemplate("system", "You are a {role}."),
+            MessageTemplate("user", "{task}"),
+        ])
+        messages = template.format_messages(role="assistant", task="Help me")
+    """
+
+    def __init__(
+        self,
+        messages: List[MessageTemplate],
+        required_variables: Optional[List[str]] = None,
+        defaults: Optional[Dict[str, str]] = None,
+    ):
+        """
+        Initialize chat template.
+
+        Args:
+            messages: List of message templates
+            required_variables: List of required variable names
+            defaults: Default values for optional variables
+        """
+        self.messages = messages
+        self.required_variables = required_variables or []
+        self.defaults = defaults or {}
+
+        # Extract all variables
+        self._extract_variables()
+
+    def _extract_variables(self) -> None:
+        """Extract variables from all message templates."""
+        self.variables = set()
+        pattern = r"\{(\w+)\}"
+
+        for msg in self.messages:
+            vars_in_msg = set(re.findall(pattern, msg.content))
+            self.variables.update(vars_in_msg)
+
+    def format_messages(self, **kwargs) -> List[LLMMessage]:
+        """
+        Format all messages with provided variables.
+
+        Args:
+            **kwargs: Variable values
+
+        Returns:
+            List of LLMMessage instances
+
+        Raises:
+            TemplateMissingVariableError: If required variable missing
+        """
+        # Check required variables
+        for var in self.required_variables:
+            if var not in kwargs and var not in self.defaults:
+                raise TemplateMissingVariableError(f"Required variable '{var}' not provided")
+
+        # Merge with defaults
+        values = {**self.defaults, **kwargs}
+
+        # Format each message
+        formatted_messages = []
+        for msg_template in self.messages:
+            try:
+                content = msg_template.content.format(**values)
+                formatted_messages.append(LLMMessage(role=msg_template.role, content=content))
+            except KeyError as e:
+                raise TemplateMissingVariableError(f"Variable {e} not provided and has no default")
+
+        return formatted_messages
+
+    def partial(self, **kwargs) -> "ChatPromptTemplate":
+        """
+        Create a partial template with some variables pre-filled.
+
+        Args:
+            **kwargs: Variable values to pre-fill
+
+        Returns:
+            New ChatPromptTemplate with updated defaults
+        """
+        new_defaults = {**self.defaults, **kwargs}
+        return ChatPromptTemplate(
+            messages=self.messages,
+            required_variables=self.required_variables,
+            defaults=new_defaults,
+        )
+
+    def add_message(self, role: str, content: str) -> "ChatPromptTemplate":
+        """
+        Add a message to the template.
+
+        Args:
+            role: Message role
+            content: Message content template
+
+        Returns:
+            New ChatPromptTemplate with added message
+        """
+        new_messages = self.messages + [MessageTemplate(role, content)]
+        return ChatPromptTemplate(
+            messages=new_messages,
+            required_variables=self.required_variables,
+            defaults=self.defaults,
+        )
+
+    def __repr__(self) -> str:
+        return f"ChatPromptTemplate(messages={len(self.messages)}, variables={self.variables})"
+
+
+def create_system_prompt(content: str) -> ChatPromptTemplate:
+    """
+    Helper to create a chat template with system message.
+
+    Args:
+        content: System message content
+
+    Returns:
+        ChatPromptTemplate with system message
+    """
+    return ChatPromptTemplate([MessageTemplate("system", content)])
+
+
+def create_basic_chat(system: str, user: str) -> ChatPromptTemplate:
+    """
+    Helper to create a basic system + user chat template.
+
+    Args:
+        system: System message content
+        user: User message content
+
+    Returns:
+        ChatPromptTemplate with system and user messages
+    """
+    return ChatPromptTemplate(
+        [
+            MessageTemplate("system", system),
+            MessageTemplate("user", user),
+        ]
+    )

aiecs/domain/agent/registry.py

@@ -0,0 +1,253 @@
+"""
+Agent Registry
+
+Central registry for tracking and managing active agents.
+"""
+
+import logging
+from typing import Dict, List, Optional, Set
+
+from .base_agent import BaseAIAgent
+from .models import AgentState, AgentType
+from .exceptions import AgentNotFoundError, AgentAlreadyRegisteredError
+
+logger = logging.getLogger(__name__)
+
+
+class AgentRegistry:
+    """
+    Central registry for tracking and managing active agents.
+
+    Thread-safe registry for agent lifecycle management.
+    """
+
+    def __init__(self) -> None:
+        """Initialize agent registry."""
+        self._agents: Dict[str, BaseAIAgent] = {}
+        self._agents_by_type: Dict[AgentType, Set[str]] = {}
+        self._agents_by_state: Dict[AgentState, Set[str]] = {}
+
+        logger.info("AgentRegistry initialized")
+
+    def register(self, agent: BaseAIAgent) -> None:
+        """
+        Register an agent.
+
+        Args:
+            agent: Agent to register
+
+        Raises:
+            AgentAlreadyRegisteredError: If agent already registered
+        """
+        if agent.agent_id in self._agents:
+            raise AgentAlreadyRegisteredError(agent.agent_id)
+
+        # Register agent
+        self._agents[agent.agent_id] = agent
+
+        # Index by type
+        if agent.agent_type not in self._agents_by_type:
+            self._agents_by_type[agent.agent_type] = set()
+        self._agents_by_type[agent.agent_type].add(agent.agent_id)
+
+        # Index by state
+        if agent.state not in self._agents_by_state:
+            self._agents_by_state[agent.state] = set()
+        self._agents_by_state[agent.state].add(agent.agent_id)
+
+        logger.info(f"Agent registered: {agent.agent_id} ({agent.agent_type.value})")
+
+    def unregister(self, agent_id: str) -> None:
+        """
+        Unregister an agent.
+
+        Args:
+            agent_id: Agent identifier
+
+        Raises:
+            AgentNotFoundError: If agent not found
+        """
+        agent = self.get(agent_id)
+
+        # Remove from indexes
+        if agent.agent_type in self._agents_by_type:
+            self._agents_by_type[agent.agent_type].discard(agent_id)
+
+        if agent.state in self._agents_by_state:
+            self._agents_by_state[agent.state].discard(agent_id)
+
+        # Remove from main registry
+        del self._agents[agent_id]
+
+        logger.info(f"Agent unregistered: {agent_id}")
+
+    def get(self, agent_id: str) -> BaseAIAgent:
+        """
+        Get agent by ID.
+
+        Args:
+            agent_id: Agent identifier
+
+        Returns:
+            Agent instance
+
+        Raises:
+            AgentNotFoundError: If agent not found
+        """
+        if agent_id not in self._agents:
+            raise AgentNotFoundError(agent_id)
+
+        return self._agents[agent_id]
+
+    def get_optional(self, agent_id: str) -> Optional[BaseAIAgent]:
+        """
+        Get agent by ID, returning None if not found.
+
+        Args:
+            agent_id: Agent identifier
+
+        Returns:
+            Agent instance or None
+        """
+        return self._agents.get(agent_id)
+
+    def exists(self, agent_id: str) -> bool:
+        """
+        Check if agent exists.
+
+        Args:
+            agent_id: Agent identifier
+
+        Returns:
+            True if agent exists
+        """
+        return agent_id in self._agents
+
+    def list_all(self) -> List[BaseAIAgent]:
+        """
+        List all registered agents.
+
+        Returns:
+            List of all agents
+        """
+        return list(self._agents.values())
+
+    def list_by_type(self, agent_type: AgentType) -> List[BaseAIAgent]:
+        """
+        List agents by type.
+
+        Args:
+            agent_type: Agent type
+
+        Returns:
+            List of agents of specified type
+        """
+        agent_ids = self._agents_by_type.get(agent_type, set())
+        return [self._agents[aid] for aid in agent_ids if aid in self._agents]
+
+    def list_by_state(self, state: AgentState) -> List[BaseAIAgent]:
+        """
+        List agents by state.
+
+        Args:
+            state: Agent state
+
+        Returns:
+            List of agents in specified state
+        """
+        agent_ids = self._agents_by_state.get(state, set())
+        return [self._agents[aid] for aid in agent_ids if aid in self._agents]
+
+    def update_state_index(self, agent_id: str, old_state: AgentState, new_state: AgentState) -> None:
+        """
+        Update state index when agent state changes.
+
+        Args:
+            agent_id: Agent identifier
+            old_state: Previous state
+            new_state: New state
+        """
+        # Remove from old state index
+        if old_state in self._agents_by_state:
+            self._agents_by_state[old_state].discard(agent_id)
+
+        # Add to new state index
+        if new_state not in self._agents_by_state:
+            self._agents_by_state[new_state] = set()
+        self._agents_by_state[new_state].add(agent_id)
+
+    def count(self) -> int:
+        """
+        Get total number of registered agents.
+
+        Returns:
+            Number of agents
+        """
+        return len(self._agents)
+
+    def count_by_type(self, agent_type: AgentType) -> int:
+        """
+        Get count of agents by type.
+
+        Args:
+            agent_type: Agent type
+
+        Returns:
+            Number of agents of specified type
+        """
+        return len(self._agents_by_type.get(agent_type, set()))
+
+    def count_by_state(self, state: AgentState) -> int:
+        """
+        Get count of agents by state.
+
+        Args:
+            state: Agent state
+
+        Returns:
+            Number of agents in specified state
+        """
+        return len(self._agents_by_state.get(state, set()))
+
+    def get_stats(self) -> Dict:
+        """
+        Get registry statistics.
+
+        Returns:
+            Dictionary with stats
+        """
+        return {
+            "total_agents": self.count(),
+            "by_type": {agent_type.value: len(agent_ids) for agent_type, agent_ids in self._agents_by_type.items()},
+            "by_state": {state.value: len(agent_ids) for state, agent_ids in self._agents_by_state.items()},
+        }
+
+    def clear(self) -> None:
+        """Clear all agents from registry."""
+        self._agents.clear()
+        self._agents_by_type.clear()
+        self._agents_by_state.clear()
+        logger.info("AgentRegistry cleared")
+
+
+# Global registry instance
+_global_registry: Optional[AgentRegistry] = None
+
+
+def get_global_registry() -> AgentRegistry:
+    """
+    Get or create global agent registry.
+
+    Returns:
+        Global AgentRegistry instance
+    """
+    global _global_registry
+    if _global_registry is None:
+        _global_registry = AgentRegistry()
+    return _global_registry
+
+
+def reset_global_registry() -> None:
+    """Reset global registry (primarily for testing)."""
+    global _global_registry
+    _global_registry = None