import asyncio import json import logging import os from datetime import datetime from sqlalchemy import select from app.models.base import async_session from app.models.generation_record import GenerationRecord from app.services.video_gen import get_active_engine, poll_task_status, download_video, _log_video_response from app.services.image_gen import get_active_image_engine, poll_image_task_status, download_image from app.config import settings logger = logging.getLogger("videogen") POLL_INTERVAL = 30 # seconds between polls MAX_POLLS = 60 # max 30 minutes total class TaskQueue: def __init__(self): self.queue: asyncio.Queue[str] = asyncio.Queue() self.running = False self._active: dict[str, int] = {} # record_id -> poll count async def enqueue(self, record_id: str): """Add a record to the polling queue.""" await self.queue.put(record_id) async def recover(self): """Recover in-progress tasks from DB on startup.""" async with async_session() as db: result = await db.execute( select(GenerationRecord).where( GenerationRecord.status == "generating", GenerationRecord.seedance_task_id.isnot(None), ) ) records = result.scalars().all() for record in records: await self.queue.put(record.id) logger.info(f"Recovered task: {record.id} (seedance: {record.seedance_task_id})") async def run(self): """Main polling loop.""" self.running = True logger.info("Video queue started") while self.running: try: record_id = await asyncio.wait_for(self.queue.get(), timeout=5.0) except asyncio.TimeoutError: continue try: await self._process(record_id) except Exception as e: logger.error(f"Error processing {record_id}: {e}") finally: self.queue.task_done() logger.info("Video queue stopped") async def _process(self, record_id: str): """Process a single record: poll status and update DB.""" async with async_session() as db: result = await db.execute( select(GenerationRecord).where(GenerationRecord.id == record_id) ) record = result.scalar_one_or_none() if not record or record.status != "generating": return if record.gen_type == "video": await self._process_video(db, record) else: await self._process_image(db, record) async def _process_video(self, db, record): """Process video generation task.""" record_id = record.id if not record.seedance_task_id: record.status = "failed" record.error_message = "缺少外部任务ID" await db.commit() return try: engine = await get_active_engine(db) poll_result = await poll_task_status(engine, record.seedance_task_id) except Exception as e: logger.error(f"Poll error for {record_id}: {e}") count = self._active.get(record_id, 0) + 1 self._active[record_id] = count if count >= MAX_POLLS: record.status = "failed" record.error_message = f"轮询超时: {e}" await db.commit() del self._active[record_id] else: await asyncio.sleep(POLL_INTERVAL) await self.queue.put(record_id) return status = poll_result["status"] try: resp_data = json.loads(poll_result.get("response_data", "{}")) except (json.JSONDecodeError, TypeError): resp_data = {} _log_video_response(record_id, resp_data, poll_result.get("error")) if status == "succeeded": file_url = poll_result.get("video_url", "") if settings.STORAGE_TYPE == "local" and file_url: try: date_dir = datetime.now().strftime("%Y/%m/%d") dest_dir = os.path.join(settings.STORAGE_LOCAL_PATH, date_dir) os.makedirs(dest_dir, exist_ok=True) dest = os.path.join(dest_dir, f"{record_id}.mp4") await download_video(file_url, dest) record.video_url = f"/videos/{date_dir}/{record_id}.mp4" except Exception as e: logger.warning(f"Download failed, using remote URL: {e}") record.video_url = file_url else: record.video_url = file_url record.video_tokens_used = poll_result.get("video_tokens", 0) record.status = "completed" record.generated_at = datetime.now() self._active.pop(record_id, None) await db.commit() logger.info(f"Video task completed: {record_id}") elif status == "failed": record.status = "failed" record.error_message = poll_result.get("error", "视频生成失败") self._active.pop(record_id, None) await db.commit() logger.info(f"Video task failed: {record_id}") else: count = self._active.get(record_id, 0) + 1 self._active[record_id] = count if count >= MAX_POLLS: record.status = "failed" record.error_message = "视频生成超时" self._active.pop(record_id, None) await db.commit() logger.info(f"Video task timed out: {record_id}") else: await db.commit() await asyncio.sleep(POLL_INTERVAL) await self.queue.put(record_id) async def _process_image(self, db, record): """Process image generation task - calls API directly.""" record_id = record.id from app.services.image_gen import submit_image_task, download_image, _log_image_response try: engine = await get_active_image_engine(db) poll_result = await asyncio.to_thread(submit_image_task, db, engine, record) if poll_result["error"] == "": if settings.STORAGE_TYPE == "local" and poll_result.get("image_url"): try: date_dir = datetime.now().strftime("%Y/%m/%d") dest_dir = os.path.join(settings.STORAGE_IMAGE_LOCAL_PATH, date_dir) os.makedirs(dest_dir, exist_ok=True) dest = os.path.join(dest_dir, f"{record_id}.png") await download_image(poll_result.get("image_url"), dest) record.image_url = f"/images/{date_dir}/{record_id}.png" except Exception as e: logger.warning(f"Download failed, using remote URL: {e}") record.image_url = poll_result.get("image_url") else: record.image_url = poll_result.get("image_url") record.image_tokens_used = poll_result.get("image_tokens", 0) record.status = "completed" record.generated_at = datetime.now() await db.commit() logger.info(f"Image task completed: {record_id}") else: record.status = "failed" record.error_message = poll_result.get("error", "图片生成失败") await db.commit() logger.info(f"Image task failed: {record_id}") _log_image_response(record_id, poll_result) except Exception as e: record.status = "failed" record.error_message = str(e) _log_image_response(record_id, {}, str(e)) await db.commit() logger.error(f"Image task failed: {record_id}, error: {e}") def stop(self): """Signal the queue to stop.""" self.running = False task_queue = TaskQueue()