Luxx/luxx/agents/core.py

"""Agent core models"""
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional
from enum import Enum
from datetime import datetime

from luxx.tools.core import CommandPermission


class AgentType(str, Enum):
    """Agent type enumeration"""
    SUPERVISOR = "supervisor"
    WORKER = "worker"


class AgentStatus(str, Enum):
    """Agent status enumeration"""
    IDLE = "idle"
    PLANNING = "planning"
    EXECUTING = "executing"
    WAITING = "waiting"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


@dataclass
class AgentConfig:
    """Agent configuration"""
    name: str
    agent_type: AgentType
    description: str = ""
    max_permission: CommandPermission = CommandPermission.EXECUTE
    max_turns: int = 10  # Context window: sliding window size
    model: str = "deepseek-chat"
    temperature: float = 0.7
    max_tokens: int = 4096
    system_prompt: str = ""
    tools: List[str] = field(default_factory=list)  # Tool names available to this agent
    metadata: Dict[str, Any] = field(default_factory=dict)


@dataclass
class Agent:
    """
    Agent entity
    
    Represents an AI agent with its configuration, state, and context.
    """
    id: str
    config: AgentConfig
    status: AgentStatus = AgentStatus.IDLE
    user_id: Optional[int] = None
    conversation_id: Optional[str] = None
    workspace: Optional[str] = None
    
    # Runtime state
    created_at: datetime = field(default_factory=datetime.utcnow)
    updated_at: datetime = field(default_factory=datetime.utcnow)
    
    # Context management
    context_window: List[Dict[str, Any]] = field(default_factory=list)
    accumulated_result: Dict[str, Any] = field(default_factory=dict)
    
    # Progress tracking
    current_task_id: Optional[str] = None
    progress: float = 0.0  # 0.0 - 1.0
    
    # Permission (effective permission = min(user_permission, agent.max_permission))
    effective_permission: CommandPermission = field(default_factory=lambda: CommandPermission.EXECUTE)
    
    def __post_init__(self):
        """Post-initialization processing"""
        if self.config.system_prompt and not self.context_window:
            # Initialize with system prompt
            self.context_window = [
                {"role": "system", "content": self.config.system_prompt}
            ]
    
    def add_message(self, role: str, content: str) -> None:
        """
        Add message to context window with sliding window management
        
        Args:
            role: Message role (user/assistant/system)
            content: Message content
        """
        self.context_window.append({"role": role, "content": content})
        self._trim_context()
        self.updated_at = datetime.utcnow()
    
    def _trim_context(self) -> None:
        """
        Trim context window using sliding window strategy
        
        Keeps system prompt and the most recent N turns (user+assistant pairs)
        """
        max_items = 1 + (self.config.max_turns * 2)  # system + (max_turns * 2)
        
        # Always keep system prompt at index 0
        if len(self.context_window) > max_items and len(self.context_window) > 1:
            # Keep system prompt and the most recent messages
            system_prompt = self.context_window[0]
            remaining = self.context_window[1:]
            trimmed = remaining[-(max_items - 1):]
            self.context_window = [system_prompt] + trimmed
    
    def get_context(self) -> List[Dict[str, Any]]:
        """Get current context window"""
        return self.context_window.copy()
    
    def set_user_permission(self, user_permission: CommandPermission) -> None:
        """
        Set effective permission based on user and agent limits
        
        Effective permission = min(user_permission, agent.max_permission)
        
        Args:
            user_permission: User's permission level
        """
        self.effective_permission = min(user_permission, self.config.max_permission)
    
    def store_result(self, key: str, value: Any) -> None:
        """
        Store result for supervisor's result-based context management
        
        Args:
            key: Result key
            value: Result value
        """
        self.accumulated_result[key] = value
        self.updated_at = datetime.utcnow()
    
    def get_result(self, key: str) -> Optional[Any]:
        """Get stored result by key"""
        return self.accumulated_result.get(key)
    
    def clear_context(self) -> None:
        """Clear context but keep system prompt"""
        if self.context_window and self.context_window[0]["role"] == "system":
            system_prompt = self.context_window[0]
            self.context_window = [system_prompt]
        else:
            self.context_window = []
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for serialization"""
        return {
            "id": self.id,
            "name": self.config.name,
            "type": self.config.agent_type.value,
            "status": self.status.value,
            "user_id": self.user_id,
            "conversation_id": self.conversation_id,
            "workspace": self.workspace,
            "created_at": self.created_at.isoformat(),
            "updated_at": self.updated_at.isoformat(),
            "current_task_id": self.current_task_id,
            "progress": self.progress,
            "effective_permission": self.effective_permission.name,
        }