PraisonAI 3.0.0__py3-none-any.whl

praisonai/cli/features/queue/manager.py

@@ -0,0 +1,395 @@
+"""
+Queue Manager for PraisonAI.
+
+High-level interface that coordinates scheduler, workers, and persistence.
+"""
+
+import asyncio
+import logging
+import os
+import time
+from typing import Any, Callable, Coroutine, Dict, List, Optional
+
+from .models import QueuedRun, RunState, RunPriority, QueueConfig, QueueStats, QueueEvent
+from .scheduler import QueueScheduler, QueueFullError
+from .worker import WorkerPool
+from .persistence import QueuePersistence
+
+logger = logging.getLogger(__name__)
+
+
+class QueueManager:
+    """
+    High-level queue manager that coordinates all queue components.
+    
+    Provides a unified interface for:
+    - Submitting and managing runs
+    - Starting/stopping workers
+    - Persistence and crash recovery
+    - Event handling
+    """
+    
+    def __init__(
+        self,
+        config: Optional[QueueConfig] = None,
+        on_output: Optional[Callable[[str, str], Coroutine[Any, Any, None]]] = None,
+        on_complete: Optional[Callable[[str, QueuedRun], Coroutine[Any, Any, None]]] = None,
+        on_error: Optional[Callable[[str, Exception], Coroutine[Any, Any, None]]] = None,
+        on_event: Optional[Callable[[QueueEvent], Coroutine[Any, Any, None]]] = None,
+    ):
+        """
+        Initialize queue manager.
+        
+        Args:
+            config: Queue configuration.
+            on_output: Callback for streaming output.
+            on_complete: Callback when run completes.
+            on_error: Callback when run fails.
+            on_event: Callback for queue events.
+        """
+        self.config = config or QueueConfig()
+        
+        # Initialize components
+        self.scheduler = QueueScheduler(self.config)
+        self.persistence = QueuePersistence(self.config.db_path) if self.config.enable_persistence else None
+        
+        # Store callbacks
+        self._on_output = on_output
+        self._on_complete = on_complete
+        self._on_error = on_error
+        self._on_event = on_event
+        
+        # Worker pool (created on start)
+        self.workers: Optional[WorkerPool] = None
+        
+        # Autosave task
+        self._autosave_task: Optional[asyncio.Task] = None
+        
+        # Running state
+        self._running = False
+        
+        # Session tracking
+        self._current_session_id: Optional[str] = None
+    
+    async def start(self, recover: bool = True) -> None:
+        """
+        Start the queue manager.
+        
+        Args:
+            recover: If True, recover pending runs from persistence.
+        """
+        if self._running:
+            return
+        
+        logger.info("Starting queue manager")
+        
+        # Initialize persistence
+        if self.persistence:
+            self.persistence.initialize()
+            
+            if recover:
+                await self._recover_runs()
+        
+        # Create and start worker pool
+        self.workers = WorkerPool(
+            scheduler=self.scheduler,
+            on_output=self._on_output,
+            on_complete=self._wrap_complete_callback(),
+            on_error=self._wrap_error_callback(),
+            on_event=self._on_event,
+            max_workers=self.config.max_concurrent_global,
+            poll_interval=self.config.worker_poll_interval,
+            stream_buffer_size=self.config.stream_buffer_size,
+        )
+        
+        await self.workers.start()
+        
+        # Start autosave
+        if self.persistence and self.config.autosave_interval_seconds > 0:
+            self._autosave_task = asyncio.create_task(
+                self._autosave_loop(),
+                name="queue-autosave"
+            )
+        
+        self._running = True
+        logger.info("Queue manager started")
+    
+    async def stop(self, timeout: float = 10.0) -> None:
+        """
+        Stop the queue manager gracefully.
+        
+        Args:
+            timeout: Seconds to wait for workers to finish.
+        """
+        if not self._running:
+            return
+        
+        logger.info("Stopping queue manager")
+        
+        # Stop autosave
+        if self._autosave_task:
+            self._autosave_task.cancel()
+            try:
+                await self._autosave_task
+            except asyncio.CancelledError:
+                pass
+            self._autosave_task = None
+        
+        # Stop workers
+        if self.workers:
+            await self.workers.stop(timeout=timeout)
+            self.workers = None
+        
+        # Final save
+        if self.persistence:
+            await self._save_all_runs()
+            self.persistence.close()
+        
+        self._running = False
+        logger.info("Queue manager stopped")
+    
+    async def submit(
+        self,
+        input_content: str,
+        agent_name: str = "Assistant",
+        priority: RunPriority = RunPriority.NORMAL,
+        session_id: Optional[str] = None,
+        workspace: Optional[str] = None,
+        config: Optional[Dict[str, Any]] = None,
+        run_id: Optional[str] = None,
+    ) -> str:
+        """
+        Submit a new run to the queue.
+        
+        Args:
+            input_content: The input/prompt for the agent.
+            agent_name: Name of the agent to use.
+            priority: Run priority.
+            session_id: Session ID for grouping.
+            workspace: Workspace path.
+            config: Additional configuration.
+            run_id: Optional specific run ID.
+            
+        Returns:
+            The run ID.
+        """
+        run = QueuedRun(
+            agent_name=agent_name,
+            input_content=input_content,
+            priority=priority,
+            session_id=session_id or self._current_session_id,
+            workspace=workspace or os.getcwd(),
+            config=config or {},
+            max_retries=self.config.default_max_retries,
+        )
+        
+        if run_id:
+            run.run_id = run_id
+        
+        # Submit to scheduler
+        result_id = await self.scheduler.submit(run)
+        
+        # Persist
+        if self.persistence:
+            self.persistence.save_run(run)
+        
+        return result_id
+    
+    async def cancel(self, run_id: str) -> bool:
+        """Cancel a run."""
+        result = await self.scheduler.cancel(run_id)
+        
+        if result and self.persistence:
+            run = self.scheduler.get_run(run_id)
+            if run:
+                self.persistence.save_run(run)
+        
+        return result
+    
+    async def retry(self, run_id: str) -> Optional[str]:
+        """Retry a failed run."""
+        new_id = await self.scheduler.retry(run_id)
+        
+        if new_id and self.persistence:
+            run = self.scheduler.get_run(new_id)
+            if run:
+                self.persistence.save_run(run)
+        
+        return new_id
+    
+    async def pause(self, run_id: str) -> bool:
+        """Pause a running run."""
+        return await self.scheduler.pause(run_id)
+    
+    async def resume(self, run_id: str) -> bool:
+        """Resume a paused run."""
+        return await self.scheduler.resume(run_id)
+    
+    def get_run(self, run_id: str) -> Optional[QueuedRun]:
+        """Get a run by ID."""
+        return self.scheduler.get_run(run_id)
+    
+    def get_queued(self) -> List[QueuedRun]:
+        """Get all queued runs."""
+        return self.scheduler.get_queued()
+    
+    def get_running(self) -> List[QueuedRun]:
+        """Get all running runs."""
+        return self.scheduler.get_running()
+    
+    def list_runs(
+        self,
+        state: Optional[RunState] = None,
+        session_id: Optional[str] = None,
+        limit: int = 100,
+    ) -> List[QueuedRun]:
+        """
+        List runs with optional filters.
+        
+        For active runs, uses scheduler. For historical, uses persistence.
+        """
+        if state and state.is_terminal() and self.persistence:
+            return self.persistence.list_runs(
+                state=state,
+                session_id=session_id,
+                limit=limit,
+            )
+        
+        # Get from scheduler
+        runs = self.scheduler.get_all()
+        
+        if state:
+            runs = [r for r in runs if r.state == state]
+        if session_id:
+            runs = [r for r in runs if r.session_id == session_id]
+        
+        return runs[:limit]
+    
+    async def clear_queue(self) -> int:
+        """Clear all queued runs."""
+        count = await self.scheduler.clear_queue()
+        return count
+    
+    def get_stats(self) -> QueueStats:
+        """Get queue statistics."""
+        if self.persistence:
+            return self.persistence.get_stats(self._current_session_id)
+        
+        # Calculate from scheduler
+        queued = self.scheduler.get_queued()
+        running = self.scheduler.get_running()
+        
+        return QueueStats(
+            queued_count=len(queued),
+            running_count=len(running),
+        )
+    
+    # Session management
+    
+    def set_session(self, session_id: str) -> None:
+        """Set the current session ID."""
+        self._current_session_id = session_id
+    
+    def save_session_state(self, state: Dict[str, Any]) -> None:
+        """Save session state."""
+        if self.persistence and self._current_session_id:
+            self.persistence.save_session(
+                self._current_session_id,
+                state=state,
+            )
+    
+    def load_session_state(self, session_id: str) -> Optional[Dict[str, Any]]:
+        """Load session state."""
+        if self.persistence:
+            return self.persistence.load_session(session_id)
+        return None
+    
+    def list_sessions(self, limit: int = 50) -> List[Dict[str, Any]]:
+        """List recent sessions."""
+        if self.persistence:
+            return self.persistence.list_sessions(limit)
+        return []
+    
+    # Internal methods
+    
+    async def _recover_runs(self) -> None:
+        """Recover pending runs from persistence."""
+        if not self.persistence:
+            return
+        
+        # Mark interrupted runs as failed
+        interrupted = self.persistence.mark_interrupted_as_failed()
+        if interrupted:
+            logger.warning(f"Marked {interrupted} interrupted runs as failed")
+        
+        # Load pending runs
+        pending = self.persistence.load_pending_runs()
+        if pending:
+            logger.info(f"Recovering {len(pending)} pending runs")
+            self.scheduler.load_runs(pending)
+            
+            for run in pending:
+                self.persistence.mark_recovered(run.run_id)
+    
+    async def _save_all_runs(self) -> None:
+        """Save all active runs to persistence."""
+        if not self.persistence:
+            return
+        
+        for run in self.scheduler.get_all():
+            self.persistence.save_run(run)
+    
+    async def _autosave_loop(self) -> None:
+        """Periodically save runs to persistence."""
+        while True:
+            try:
+                await asyncio.sleep(self.config.autosave_interval_seconds)
+                await self._save_all_runs()
+                logger.debug("Autosaved queue state")
+            except asyncio.CancelledError:
+                break
+            except Exception as e:
+                logger.error(f"Autosave error: {e}")
+    
+    def _wrap_complete_callback(self):
+        """Wrap complete callback to include persistence."""
+        async def wrapper(run_id: str, run: QueuedRun):
+            # Persist
+            if self.persistence:
+                self.persistence.save_run(run)
+            
+            # Call user callback
+            if self._on_complete:
+                await self._on_complete(run_id, run)
+        
+        return wrapper
+    
+    def _wrap_error_callback(self):
+        """Wrap error callback to include persistence."""
+        async def wrapper(run_id: str, error: Exception):
+            # Persist
+            if self.persistence:
+                run = self.scheduler.get_run(run_id)
+                if run:
+                    self.persistence.save_run(run)
+            
+            # Call user callback
+            if self._on_error:
+                await self._on_error(run_id, error)
+        
+        return wrapper
+    
+    @property
+    def is_running(self) -> bool:
+        """Whether the queue manager is running."""
+        return self._running
+    
+    @property
+    def queued_count(self) -> int:
+        """Number of queued runs."""
+        return self.scheduler.queued_count
+    
+    @property
+    def running_count(self) -> int:
+        """Number of running runs."""
+        return self.scheduler.running_count

praisonai/cli/features/queue/models.py

@@ -0,0 +1,286 @@
+"""
+Data models for the PraisonAI Queue System.
+
+Defines QueuedRun, RunState, RunPriority, and QueueConfig.
+"""
+
+from dataclasses import dataclass, field
+from enum import Enum, IntEnum
+from typing import Any, Dict, List, Optional
+import time
+import uuid
+
+
+class RunState(str, Enum):
+    """State of a queued run."""
+    QUEUED = "queued"
+    RUNNING = "running"
+    PAUSED = "paused"
+    SUCCEEDED = "succeeded"
+    FAILED = "failed"
+    CANCELLED = "cancelled"
+    
+    def is_terminal(self) -> bool:
+        """Check if this is a terminal state."""
+        return self in (RunState.SUCCEEDED, RunState.FAILED, RunState.CANCELLED)
+    
+    def is_active(self) -> bool:
+        """Check if this is an active (non-terminal) state."""
+        return not self.is_terminal()
+
+
+class RunPriority(IntEnum):
+    """Priority levels for queued runs. Higher value = higher priority."""
+    LOW = 0
+    NORMAL = 1
+    HIGH = 2
+    URGENT = 3
+    
+    @classmethod
+    def from_string(cls, s: str) -> "RunPriority":
+        """Parse priority from string."""
+        mapping = {
+            "low": cls.LOW,
+            "normal": cls.NORMAL,
+            "high": cls.HIGH,
+            "urgent": cls.URGENT,
+        }
+        return mapping.get(s.lower(), cls.NORMAL)
+
+
+@dataclass
+class QueuedRun:
+    """A single queued agent run."""
+    
+    # Identity
+    run_id: str = field(default_factory=lambda: str(uuid.uuid4())[:8])
+    agent_name: str = ""
+    input_content: str = ""
+    
+    # State
+    state: RunState = RunState.QUEUED
+    priority: RunPriority = RunPriority.NORMAL
+    
+    # Attribution / tracing
+    session_id: Optional[str] = None
+    trace_id: Optional[str] = None
+    workspace: Optional[str] = None
+    user_id: Optional[str] = None
+    
+    # Timing
+    created_at: float = field(default_factory=time.time)
+    started_at: Optional[float] = None
+    ended_at: Optional[float] = None
+    
+    # Results
+    output_content: Optional[str] = None
+    error: Optional[str] = None
+    metrics: Dict[str, Any] = field(default_factory=dict)
+    
+    # Retry tracking
+    retry_count: int = 0
+    max_retries: int = 3
+    parent_run_id: Optional[str] = None  # Link to original run if this is a retry
+    
+    # Configuration
+    config: Dict[str, Any] = field(default_factory=dict)
+    
+    # Streaming state
+    output_chunks: List[str] = field(default_factory=list)
+    
+    def __post_init__(self):
+        """Ensure proper types after initialization."""
+        if isinstance(self.state, str):
+            self.state = RunState(self.state)
+        if isinstance(self.priority, int) and not isinstance(self.priority, RunPriority):
+            self.priority = RunPriority(self.priority)
+    
+    @property
+    def duration_seconds(self) -> Optional[float]:
+        """Get run duration in seconds."""
+        if self.started_at is None:
+            return None
+        end = self.ended_at or time.time()
+        return end - self.started_at
+    
+    @property
+    def wait_seconds(self) -> float:
+        """Get time spent waiting in queue."""
+        start = self.started_at or time.time()
+        return start - self.created_at
+    
+    def can_retry(self) -> bool:
+        """Check if this run can be retried."""
+        return (
+            self.state == RunState.FAILED and
+            self.retry_count < self.max_retries
+        )
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary for persistence."""
+        return {
+            "run_id": self.run_id,
+            "agent_name": self.agent_name,
+            "input_content": self.input_content,
+            "state": self.state.value,
+            "priority": int(self.priority),
+            "session_id": self.session_id,
+            "trace_id": self.trace_id,
+            "workspace": self.workspace,
+            "user_id": self.user_id,
+            "created_at": self.created_at,
+            "started_at": self.started_at,
+            "ended_at": self.ended_at,
+            "output_content": self.output_content,
+            "error": self.error,
+            "metrics": self.metrics,
+            "retry_count": self.retry_count,
+            "max_retries": self.max_retries,
+            "parent_run_id": self.parent_run_id,
+            "config": self.config,
+        }
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "QueuedRun":
+        """Create from dictionary."""
+        # Handle state conversion
+        state = data.get("state", "queued")
+        if isinstance(state, str):
+            state = RunState(state)
+        
+        # Handle priority conversion
+        priority = data.get("priority", 1)
+        if isinstance(priority, int):
+            priority = RunPriority(priority)
+        
+        return cls(
+            run_id=data.get("run_id", str(uuid.uuid4())[:8]),
+            agent_name=data.get("agent_name", ""),
+            input_content=data.get("input_content", ""),
+            state=state,
+            priority=priority,
+            session_id=data.get("session_id"),
+            trace_id=data.get("trace_id"),
+            workspace=data.get("workspace"),
+            user_id=data.get("user_id"),
+            created_at=data.get("created_at", time.time()),
+            started_at=data.get("started_at"),
+            ended_at=data.get("ended_at"),
+            output_content=data.get("output_content"),
+            error=data.get("error"),
+            metrics=data.get("metrics", {}),
+            retry_count=data.get("retry_count", 0),
+            max_retries=data.get("max_retries", 3),
+            parent_run_id=data.get("parent_run_id"),
+            config=data.get("config", {}),
+        )
+
+
+@dataclass
+class QueueConfig:
+    """Configuration for the queue system."""
+    
+    # Concurrency limits
+    max_concurrent_global: int = 4
+    max_concurrent_per_agent: int = 2
+    max_concurrent_per_workspace: int = 4
+    
+    # Queue limits
+    max_queue_size: int = 100
+    
+    # Defaults
+    default_priority: RunPriority = RunPriority.NORMAL
+    default_max_retries: int = 3
+    
+    # Persistence
+    enable_persistence: bool = True
+    db_path: str = ".praison/queue.db"
+    
+    # Autosave
+    autosave_interval_seconds: float = 30.0
+    
+    # Backpressure
+    stream_buffer_size: int = 1000
+    drop_strategy: str = "oldest"  # "oldest" or "newest"
+    
+    # Timeouts
+    run_timeout_seconds: Optional[float] = None  # None = no timeout
+    worker_poll_interval: float = 0.1
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary."""
+        return {
+            "max_concurrent_global": self.max_concurrent_global,
+            "max_concurrent_per_agent": self.max_concurrent_per_agent,
+            "max_concurrent_per_workspace": self.max_concurrent_per_workspace,
+            "max_queue_size": self.max_queue_size,
+            "default_priority": int(self.default_priority),
+            "default_max_retries": self.default_max_retries,
+            "enable_persistence": self.enable_persistence,
+            "db_path": self.db_path,
+            "autosave_interval_seconds": self.autosave_interval_seconds,
+            "stream_buffer_size": self.stream_buffer_size,
+            "drop_strategy": self.drop_strategy,
+            "run_timeout_seconds": self.run_timeout_seconds,
+            "worker_poll_interval": self.worker_poll_interval,
+        }
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "QueueConfig":
+        """Create from dictionary."""
+        priority = data.get("default_priority", 1)
+        if isinstance(priority, int):
+            priority = RunPriority(priority)
+        
+        return cls(
+            max_concurrent_global=data.get("max_concurrent_global", 4),
+            max_concurrent_per_agent=data.get("max_concurrent_per_agent", 2),
+            max_concurrent_per_workspace=data.get("max_concurrent_per_workspace", 4),
+            max_queue_size=data.get("max_queue_size", 100),
+            default_priority=priority,
+            default_max_retries=data.get("default_max_retries", 3),
+            enable_persistence=data.get("enable_persistence", True),
+            db_path=data.get("db_path", ".praison/queue.db"),
+            autosave_interval_seconds=data.get("autosave_interval_seconds", 30.0),
+            stream_buffer_size=data.get("stream_buffer_size", 1000),
+            drop_strategy=data.get("drop_strategy", "oldest"),
+            run_timeout_seconds=data.get("run_timeout_seconds"),
+            worker_poll_interval=data.get("worker_poll_interval", 0.1),
+        )
+
+
+@dataclass
+class QueueStats:
+    """Statistics about the queue."""
+    queued_count: int = 0
+    running_count: int = 0
+    succeeded_count: int = 0
+    failed_count: int = 0
+    cancelled_count: int = 0
+    total_runs: int = 0
+    avg_wait_seconds: float = 0.0
+    avg_duration_seconds: float = 0.0
+    
+    @property
+    def active_count(self) -> int:
+        """Count of active (non-terminal) runs."""
+        return self.queued_count + self.running_count
+
+
+@dataclass 
+class StreamChunk:
+    """A chunk of streaming output."""
+    run_id: str
+    content: str
+    timestamp: float = field(default_factory=time.time)
+    chunk_index: int = 0
+    is_final: bool = False
+    
+    
+@dataclass
+class QueueEvent:
+    """An event from the queue system."""
+    event_type: str  # "run_started", "run_completed", "run_failed", "output_chunk", "state_changed"
+    run_id: str
+    timestamp: float = field(default_factory=time.time)
+    data: Dict[str, Any] = field(default_factory=dict)