from typing import Any, Dict, List import json from loguru import logger from app.core.config import settings try: import redis except Exception: # pragma: no cover - optional dependency redis = None class InMemoryTaskStore: def __init__(self) -> None: self._tasks: Dict[str, Dict[str, Any]] = {} def create(self, task_id: str, user_id: str) -> Dict[str, Any]: task = { "status": "pending", "task_id": task_id, "progress": 0, "user_id": user_id, } self._tasks[task_id] = task return task def get(self, task_id: str) -> Dict[str, Any]: return self._tasks.get(task_id, {"status": "not_found"}) def list(self) -> List[Dict[str, Any]]: return list(self._tasks.values()) def update(self, task_id: str, updates: Dict[str, Any]) -> Dict[str, Any]: task = self._tasks.get(task_id) if not task: task = {"status": "pending", "task_id": task_id} self._tasks[task_id] = task task.update(updates) return task class RedisTaskStore: def __init__(self, client: "redis.Redis") -> None: self._client = client self._index_key = "vigent:tasks:index" def _key(self, task_id: str) -> str: return f"vigent:tasks:{task_id}" def create(self, task_id: str, user_id: str) -> Dict[str, Any]: task = { "status": "pending", "task_id": task_id, "progress": 0, "user_id": user_id, } self._client.set(self._key(task_id), json.dumps(task, ensure_ascii=False)) self._client.sadd(self._index_key, task_id) return task def get(self, task_id: str) -> Dict[str, Any]: raw = self._client.get(self._key(task_id)) if not raw: return {"status": "not_found"} return json.loads(raw) def list(self) -> List[Dict[str, Any]]: task_ids = list(self._client.smembers(self._index_key) or []) if not task_ids: return [] keys = [self._key(task_id) for task_id in task_ids] raw_items = self._client.mget(keys) tasks = [] for raw in raw_items: if raw: try: tasks.append(json.loads(raw)) except Exception: continue return tasks def update(self, task_id: str, updates: Dict[str, Any]) -> Dict[str, Any]: task = self.get(task_id) if task.get("status") == "not_found": task = {"status": "pending", "task_id": task_id} task.update(updates) self._client.set(self._key(task_id), json.dumps(task, ensure_ascii=False)) self._client.sadd(self._index_key, task_id) return task def _build_task_store(): if redis is None: logger.warning("Redis not available, using in-memory task store") return InMemoryTaskStore() try: client = redis.Redis.from_url(settings.REDIS_URL, decode_responses=True) client.ping() logger.info("Using Redis task store") return RedisTaskStore(client) except Exception as e: logger.warning(f"Redis connection failed, using in-memory task store: {e}") return InMemoryTaskStore() task_store = _build_task_store() def create_task(task_id: str, user_id: str) -> Dict[str, Any]: return task_store.create(task_id, user_id) def get_task(task_id: str) -> Dict[str, Any]: return task_store.get(task_id) def list_tasks() -> List[Dict[str, Any]]: return task_store.list()