99 lines
3.3 KiB
Python
99 lines
3.3 KiB
Python
|
"""Redis worker support for Superviseur bot."""
|
||
|
|
|
||
|
|
import json
|
||
|
|
import uuid
|
||
|
|
import asyncio
|
||
|
|
import logging
|
||
|
|
from typing import Dict, Optional, Any
|
||
|
|
|
||
|
|
logger = logging.getLogger('Superviseur')
|
||
|
|
|
||
|
|
|
||
|
|
class RedisWorker:
|
||
|
|
"""Manages Redis-based LLM job queue."""
|
||
|
|
|
||
|
|
def __init__(self, bot):
|
||
|
|
self.bot = bot
|
||
|
|
self.redis_client = None
|
||
|
|
self._pending_jobs: Dict[str, asyncio.Future] = {}
|
||
|
|
self._listener_task: Optional[asyncio.Task] = None
|
||
|
|
self._response_queue = 'llm_queue:responses'
|
||
|
|
self._job_queue = 'llm_queue:jobs'
|
||
|
|
|
||
|
|
async def initialize(self, redis_url: str) -> bool:
|
||
|
|
"""Initialize Redis connection."""
|
||
|
|
if not self.bot.redis_available:
|
||
|
|
return False
|
||
|
|
|
||
|
|
try:
|
||
|
|
self.redis_client = self.bot.aioredis.from_url(redis_url)
|
||
|
|
self._listener_task = asyncio.create_task(self._listen())
|
||
|
|
logger.info("Redis LLM worker enabled; listening for responses")
|
||
|
|
return True
|
||
|
|
except Exception:
|
||
|
|
logger.exception("Failed to initialize redis client")
|
||
|
|
return False
|
||
|
|
|
||
|
|
async def close(self) -> None:
|
||
|
|
"""Close Redis connection."""
|
||
|
|
if self._listener_task:
|
||
|
|
try:
|
||
|
|
self._listener_task.cancel()
|
||
|
|
except Exception:
|
||
|
|
pass
|
||
|
|
|
||
|
|
if self.redis_client:
|
||
|
|
try:
|
||
|
|
await self.redis_client.close()
|
||
|
|
except Exception:
|
||
|
|
pass
|
||
|
|
|
||
|
|
async def enqueue_job(self, payload: Dict[str, Any], timeout: int = 60) -> Optional[Dict[str, Any]]:
|
||
|
|
"""Enqueue an LLM job and wait for result."""
|
||
|
|
if not self.redis_client:
|
||
|
|
return None
|
||
|
|
|
||
|
|
job_id = str(uuid.uuid4())
|
||
|
|
future = self.bot.loop.create_future()
|
||
|
|
self._pending_jobs[job_id] = future
|
||
|
|
payload['job_id'] = job_id
|
||
|
|
|
||
|
|
await self.redis_client.lpush(self._job_queue, json.dumps(payload))
|
||
|
|
logger.debug(f"LLM job {job_id} enqueued.")
|
||
|
|
|
||
|
|
try:
|
||
|
|
return await asyncio.wait_for(future, timeout=timeout)
|
||
|
|
finally:
|
||
|
|
self._pending_jobs.pop(job_id, None)
|
||
|
|
|
||
|
|
async def _listen(self) -> None:
|
||
|
|
"""Listen for Redis responses."""
|
||
|
|
logger.info(f"Starting Redis response listener on {self._response_queue}")
|
||
|
|
|
||
|
|
while True:
|
||
|
|
try:
|
||
|
|
result = await self.redis_client.brpop(self._response_queue, timeout=1)
|
||
|
|
if result is None:
|
||
|
|
await asyncio.sleep(0.1)
|
||
|
|
continue
|
||
|
|
|
||
|
|
data_bytes = result[1]
|
||
|
|
response = json.loads(data_bytes.decode('utf-8'))
|
||
|
|
job_id = response.get('job_id')
|
||
|
|
|
||
|
|
if job_id in self._pending_jobs:
|
||
|
|
future = self._pending_jobs[job_id]
|
||
|
|
if not future.done():
|
||
|
|
future.set_result(response)
|
||
|
|
logger.debug(f"LLM job {job_id} completed via Redis.")
|
||
|
|
else:
|
||
|
|
logger.warning(f"Response for unknown job {job_id}")
|
||
|
|
|
||
|
|
except asyncio.CancelledError:
|
||
|
|
logger.info("Redis listener cancelled")
|
||
|
|
break
|
||
|
|
except Exception as e:
|
||
|
|
logger.exception(f"Redis listener error: {e}")
|
||
|
|
await asyncio.sleep(5)
|
||
|
|
|