lambdagent/agentpaas/engine/queue.py

"""
engine.queue — Job Queue abstraction (memory / Redis).

Dev: asyncio.PriorityQueue (in-memory)
Prod: Redis Sorted Set (ZADD score=priority*1e12+timestamp)
"""
from __future__ import annotations
import asyncio
import json
import time
from typing import Any, Dict, Optional
from dataclasses import dataclass, field
from agentpaas.db.models import gen_id, now_utc


@dataclass
class Job:
    id: str = ""
    tenant_id: str = ""
    agent_id: str = ""
    agent_version: int = 1
    config: Dict[str, Any] = field(default_factory=dict)
    input: str = ""
    priority: int = 5
    status: str = "queued"
    worker_id: str = ""
    attempts: int = 0
    max_attempts: int = 3
    result: str = ""
    error: str = ""
    created_at: str = ""

    def __post_init__(self):
        if not self.id:
            self.id = gen_id("job_")
        if not self.created_at:
            self.created_at = now_utc()


class MemoryQueue:
    """In-memory priority queue for development."""

    def __init__(self):
        self._queue: asyncio.PriorityQueue = asyncio.PriorityQueue()
        self._jobs: Dict[str, Job] = {}

    async def enqueue(self, job: Job) -> str:
        self._jobs[job.id] = job
        score = job.priority * 1_000_000_000_000 + time.time()
        await self._queue.put((score, job.id))
        return job.id

    async def dequeue(self, timeout: float = 30.0) -> Optional[Job]:
        try:
            _, job_id = await asyncio.wait_for(self._queue.get(), timeout=timeout)
            job = self._jobs.get(job_id)
            if job:
                job.status = "running"
            return job
        except asyncio.TimeoutError:
            return None

    def get_job(self, job_id: str) -> Optional[Job]:
        return self._jobs.get(job_id)

    def update_job(self, job_id: str, **kwargs):
        job = self._jobs.get(job_id)
        if job:
            for k, v in kwargs.items():
                setattr(job, k, v)

    @property
    def depth(self) -> int:
        return self._queue.qsize()


class RedisQueue:
    """Redis-backed priority queue for production."""

    def __init__(self, redis_url: str, prefix: str = "agentpaas:queue:"):
        import redis
        self.r = redis.from_url(redis_url)
        self.prefix = prefix
        self.queue_key = f"{prefix}jobs"
        self.data_key = f"{prefix}data:"

    async def enqueue(self, job: Job) -> str:
        data = json.dumps({
            "id": job.id, "tenant_id": job.tenant_id, "agent_id": job.agent_id,
            "agent_version": job.agent_version, "config": job.config,
            "input": job.input, "priority": job.priority, "status": "queued",
            "max_attempts": job.max_attempts, "created_at": job.created_at,
        })
        score = job.priority * 1_000_000_000_000 + time.time()
        self.r.zadd(self.queue_key, {job.id: score})
        self.r.set(f"{self.data_key}{job.id}", data, ex=86400)
        return job.id

    async def dequeue(self, timeout: float = 30.0) -> Optional[Job]:
        result = self.r.zpopmin(self.queue_key, count=1)
        if not result:
            return None
        job_id = result[0][0].decode() if isinstance(result[0][0], bytes) else result[0][0]
        data = self.r.get(f"{self.data_key}{job_id}")
        if not data:
            return None
        d = json.loads(data)
        job = Job(**{k: v for k, v in d.items() if k in Job.__dataclass_fields__})
        job.status = "running"
        self.r.set(f"{self.data_key}{job_id}", json.dumps({**d, "status": "running"}), ex=86400)
        return job

    def get_job(self, job_id: str) -> Optional[Job]:
        data = self.r.get(f"{self.data_key}{job_id}")
        if not data:
            return None
        d = json.loads(data)
        return Job(**{k: v for k, v in d.items() if k in Job.__dataclass_fields__})

    def update_job(self, job_id: str, **kwargs):
        data = self.r.get(f"{self.data_key}{job_id}")
        if data:
            d = json.loads(data)
            d.update(kwargs)
            self.r.set(f"{self.data_key}{job_id}", json.dumps(d), ex=86400)

    @property
    def depth(self) -> int:
        return self.r.zcard(self.queue_key)


def create_queue(redis_url: Optional[str] = None):
    """Factory: Redis if URL provided, else in-memory."""
    if redis_url:
        try:
            return RedisQueue(redis_url)
        except Exception:
            pass
    return MemoryQueue()