""" 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()