agent-runtime-core 0.1.0__py3-none-any.whl

agent_runtime/queue/base.py

@@ -0,0 +1,167 @@
+"""
+Abstract base class for run queue implementations.
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import Optional
+from uuid import UUID
+
+
+@dataclass
+class QueuedRun:
+    """A run waiting in the queue."""
+    
+    run_id: UUID
+    agent_key: str
+    input: dict
+    metadata: dict = field(default_factory=dict)
+    priority: int = 0
+    attempt: int = 1
+    created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
+    scheduled_at: Optional[datetime] = None
+
+
+class RunQueue(ABC):
+    """
+    Abstract interface for run queue implementations.
+    
+    Queues handle:
+    - Enqueueing new runs
+    - Claiming runs for processing
+    - Lease management
+    - Retries
+    """
+    
+    @abstractmethod
+    async def enqueue(
+        self,
+        run_id: UUID,
+        agent_key: str,
+        input: dict,
+        metadata: Optional[dict] = None,
+        priority: int = 0,
+    ) -> None:
+        """
+        Add a run to the queue.
+        
+        Args:
+            run_id: Unique run identifier
+            agent_key: Agent to handle the run
+            input: Input data for the run
+            metadata: Optional metadata
+            priority: Priority (higher = more urgent)
+        """
+        ...
+    
+    @abstractmethod
+    async def claim(
+        self,
+        worker_id: str,
+        lease_seconds: int = 60,
+    ) -> Optional[QueuedRun]:
+        """
+        Claim the next available run.
+        
+        Args:
+            worker_id: ID of the claiming worker
+            lease_seconds: How long to hold the lease
+            
+        Returns:
+            QueuedRun if one is available, None otherwise
+        """
+        ...
+    
+    @abstractmethod
+    async def release(
+        self,
+        run_id: UUID,
+        worker_id: str,
+        success: bool,
+        output: Optional[dict] = None,
+        error: Optional[dict] = None,
+    ) -> None:
+        """
+        Release a claimed run.
+        
+        Args:
+            run_id: Run to release
+            worker_id: Worker releasing the run
+            success: Whether the run succeeded
+            output: Output data (if success)
+            error: Error info (if failure)
+        """
+        ...
+    
+    @abstractmethod
+    async def extend_lease(
+        self,
+        run_id: UUID,
+        worker_id: str,
+        lease_seconds: int,
+    ) -> bool:
+        """
+        Extend the lease on a run.
+        
+        Args:
+            run_id: Run to extend
+            worker_id: Worker holding the lease
+            lease_seconds: New lease duration
+            
+        Returns:
+            True if extended, False if lease was lost
+        """
+        ...
+    
+    @abstractmethod
+    async def is_cancelled(self, run_id: UUID) -> bool:
+        """
+        Check if a run has been cancelled.
+        
+        Args:
+            run_id: Run to check
+            
+        Returns:
+            True if cancelled
+        """
+        ...
+    
+    @abstractmethod
+    async def cancel(self, run_id: UUID) -> bool:
+        """
+        Cancel a run.
+        
+        Args:
+            run_id: Run to cancel
+            
+        Returns:
+            True if cancelled, False if not found or already complete
+        """
+        ...
+    
+    @abstractmethod
+    async def requeue_for_retry(
+        self,
+        run_id: UUID,
+        worker_id: str,
+        error: dict,
+        delay_seconds: int = 0,
+    ) -> bool:
+        """
+        Requeue a failed run for retry.
+        
+        Args:
+            run_id: Run to retry
+            worker_id: Worker releasing the run
+            error: Error information
+            delay_seconds: Delay before retry
+            
+        Returns:
+            True if requeued, False if max retries exceeded
+        """
+        ...
+    
+    async def close(self) -> None:
+        """Close any connections. Override if needed."""
+        pass

agent_runtime/queue/memory.py

@@ -0,0 +1,184 @@
+"""
+In-memory queue implementation.
+
+Good for:
+- Unit testing
+- Local development
+- Simple single-process scripts
+"""
+
+import asyncio
+from collections import deque
+from dataclasses import dataclass, field
+from datetime import datetime, timezone, timedelta
+from typing import Optional
+from uuid import UUID
+
+from agent_runtime.queue.base import RunQueue, QueuedRun
+
+
+@dataclass
+class QueueEntry:
+    """Internal queue entry with lease info."""
+    run: QueuedRun
+    lease_owner: Optional[str] = None
+    lease_expires_at: Optional[datetime] = None
+    cancelled: bool = False
+    completed: bool = False
+    output: Optional[dict] = None
+    error: Optional[dict] = None
+
+
+class InMemoryQueue(RunQueue):
+    """
+    In-memory queue implementation.
+    
+    Stores runs in memory. Data is lost when the process exits.
+    """
+    
+    def __init__(self, max_retries: int = 3):
+        self._entries: dict[UUID, QueueEntry] = {}
+        self._queue: deque[UUID] = deque()
+        self._lock = asyncio.Lock()
+        self._max_retries = max_retries
+    
+    async def enqueue(
+        self,
+        run_id: UUID,
+        agent_key: str,
+        input: dict,
+        metadata: Optional[dict] = None,
+        priority: int = 0,
+    ) -> None:
+        """Add a run to the queue."""
+        async with self._lock:
+            run = QueuedRun(
+                run_id=run_id,
+                agent_key=agent_key,
+                input=input,
+                metadata=metadata or {},
+                priority=priority,
+            )
+            self._entries[run_id] = QueueEntry(run=run)
+            self._queue.append(run_id)
+    
+    async def claim(
+        self,
+        worker_id: str,
+        lease_seconds: int = 60,
+    ) -> Optional[QueuedRun]:
+        """Claim the next available run."""
+        async with self._lock:
+            now = datetime.now(timezone.utc)
+            
+            # Find an available run
+            for _ in range(len(self._queue)):
+                run_id = self._queue.popleft()
+                entry = self._entries.get(run_id)
+                
+                if entry is None or entry.completed or entry.cancelled:
+                    continue
+                
+                # Check if lease expired
+                if entry.lease_owner and entry.lease_expires_at:
+                    if entry.lease_expires_at > now:
+                        # Still leased, put back
+                        self._queue.append(run_id)
+                        continue
+                
+                # Claim it
+                entry.lease_owner = worker_id
+                entry.lease_expires_at = now + timedelta(seconds=lease_seconds)
+                return entry.run
+            
+            return None
+    
+    async def release(
+        self,
+        run_id: UUID,
+        worker_id: str,
+        success: bool,
+        output: Optional[dict] = None,
+        error: Optional[dict] = None,
+    ) -> None:
+        """Release a claimed run."""
+        async with self._lock:
+            entry = self._entries.get(run_id)
+            if entry is None:
+                return
+            
+            if entry.lease_owner != worker_id:
+                return
+            
+            entry.completed = True
+            entry.lease_owner = None
+            entry.lease_expires_at = None
+            entry.output = output
+            entry.error = error
+    
+    async def extend_lease(
+        self,
+        run_id: UUID,
+        worker_id: str,
+        lease_seconds: int,
+    ) -> bool:
+        """Extend the lease on a run."""
+        async with self._lock:
+            entry = self._entries.get(run_id)
+            if entry is None:
+                return False
+            
+            if entry.lease_owner != worker_id:
+                return False
+            
+            entry.lease_expires_at = datetime.now(timezone.utc) + timedelta(seconds=lease_seconds)
+            return True
+    
+    async def is_cancelled(self, run_id: UUID) -> bool:
+        """Check if a run has been cancelled."""
+        entry = self._entries.get(run_id)
+        return entry.cancelled if entry else False
+    
+    async def cancel(self, run_id: UUID) -> bool:
+        """Cancel a run."""
+        async with self._lock:
+            entry = self._entries.get(run_id)
+            if entry is None or entry.completed:
+                return False
+            
+            entry.cancelled = True
+            return True
+    
+    async def requeue_for_retry(
+        self,
+        run_id: UUID,
+        worker_id: str,
+        error: dict,
+        delay_seconds: int = 0,
+    ) -> bool:
+        """Requeue a failed run for retry."""
+        async with self._lock:
+            entry = self._entries.get(run_id)
+            if entry is None:
+                return False
+            
+            if entry.lease_owner != worker_id:
+                return False
+            
+            if entry.run.attempt >= self._max_retries:
+                return False
+            
+            # Increment attempt and requeue
+            entry.run.attempt += 1
+            entry.lease_owner = None
+            entry.lease_expires_at = None
+            entry.error = error
+            
+            # Add back to queue
+            self._queue.append(run_id)
+            return True
+    
+    def clear(self) -> None:
+        """Clear all entries. Useful for testing."""
+        self._entries.clear()
+        self._queue.clear()