"""GenerationTask 领域模型 — 视频生成任务. 状态机: pending → running → completed ↘ failed → pending (重试) ↘ cancelled """ from __future__ import annotations import json import sys from dataclasses import dataclass, field from datetime import datetime, timezone if sys.version_info >= (3, 11): from enum import StrEnum else: from enum import Enum class StrEnum(str, Enum): pass from uuid import uuid4 class GenerationTaskStatus(StrEnum): """生成任务状态枚举。""" PENDING = "pending" """待处理(任务已创建,等待执行)""" RUNNING = "running" """运行中(正在生成视频)""" COMPLETED = "completed" """已完成(视频生成成功)""" FAILED = "failed" """失败(生成失败)""" CANCELLED = "cancelled" """已取消(用户取消或系统取消)""" # 终态集合 TERMINAL_STATUSES = frozenset( {GenerationTaskStatus.COMPLETED, GenerationTaskStatus.FAILED, GenerationTaskStatus.CANCELLED} ) # 合法状态转换 _VALID_TRANSITIONS: dict[GenerationTaskStatus, set[GenerationTaskStatus]] = { GenerationTaskStatus.PENDING: { GenerationTaskStatus.RUNNING, GenerationTaskStatus.FAILED, GenerationTaskStatus.CANCELLED, }, GenerationTaskStatus.RUNNING: { GenerationTaskStatus.COMPLETED, GenerationTaskStatus.FAILED, GenerationTaskStatus.CANCELLED, }, GenerationTaskStatus.FAILED: {GenerationTaskStatus.PENDING}, # 重试回到 pending } @dataclass(slots=True) class GenerationTask: id: str project_id: str asset_library_id: str strategy_id: str = "" voice_library_id: str = "" template_id: str = "" asset_ids: list[str] = field(default_factory=list) title_ids: list[str] = field(default_factory=list) voice_ids: list[str] = field(default_factory=list) status: GenerationTaskStatus = GenerationTaskStatus.PENDING progress: float = 0.0 result_count: int = 0 error_message: str = "" error_info: dict = field(default_factory=dict) retry_count: int = 0 auto_retry_enabled: bool = False auto_retry_max: int = 0 started_at: datetime | None = None completed_at: datetime | None = None source_edit_plan_id: str = "" created_by_user_id: str = "" asset_select_mode: str = "" batch_id: str = "" logs: str = "[]" created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc)) @classmethod def create( cls, project_id: str, asset_library_id: str, *, strategy_id: str = "", voice_library_id: str = "", template_id: str = "", asset_ids: list[str] | None = None, title_ids: list[str] | None = None, voice_ids: list[str] | None = None, created_by_user_id: str = "", source_edit_plan_id: str = "", asset_select_mode: str = "", batch_id: str = "", auto_retry_enabled: bool = False, auto_retry_max: int = 0, ) -> "GenerationTask": if not project_id.strip() and not template_id.strip(): raise ValueError("project_id 或 template_id 至少需要提供一个") if not asset_library_id.strip() and not (asset_ids or title_ids or voice_ids): raise ValueError("asset_library_id 或 asset_ids/title_ids/voice_ids 至少需要提供一个") return cls( id=uuid4().hex, project_id=project_id.strip(), asset_library_id=asset_library_id.strip(), strategy_id=strategy_id.strip(), voice_library_id=voice_library_id.strip(), template_id=template_id.strip(), asset_ids=list(asset_ids) if asset_ids else [], title_ids=list(title_ids) if title_ids else [], voice_ids=list(voice_ids) if voice_ids else [], created_by_user_id=created_by_user_id.strip(), source_edit_plan_id=source_edit_plan_id.strip(), asset_select_mode=asset_select_mode, batch_id=batch_id, auto_retry_enabled=auto_retry_enabled, auto_retry_max=auto_retry_max, ) # ── 状态查询 ──────────────────────────────────────────────────────────── @property def is_terminal(self) -> bool: """是否处于终态(completed / failed / cancelled)。""" return self.status in TERMINAL_STATUSES @property def is_completed(self) -> bool: """是否已完成。""" return self.status == GenerationTaskStatus.COMPLETED @property def is_failed(self) -> bool: """是否失败。""" return self.status == GenerationTaskStatus.FAILED @property def is_running(self) -> bool: """是否运行中。""" return self.status == GenerationTaskStatus.RUNNING # ── 状态转换 ──────────────────────────────────────────────────────────── def transition_to(self, new_status: GenerationTaskStatus | str) -> None: """执行状态转换。 Args: new_status: 目标状态 Raises: ValueError: 非法状态转换 """ if isinstance(new_status, str): try: new_status = GenerationTaskStatus(new_status) except ValueError: raise ValueError(f"无效状态: {new_status}") allowed = _VALID_TRANSITIONS.get(self.status, set()) if new_status not in allowed: raise ValueError( f"非法状态转换: {self.status.value} → {new_status.value}," f"允许: {{{', '.join(sorted(s.value for s in allowed))}}}" ) self.status = new_status def mark_processing(self) -> None: """标记为处理中(pending → running)。 设置 started_at,清除 error_message。 Raises: ValueError: 当前状态不允许转换到 running """ self.transition_to(GenerationTaskStatus.RUNNING) self.started_at = datetime.now(timezone.utc) self.error_message = "" def mark_completed(self, result_count: int = 1) -> None: """标记为已完成(running → completed)。 设置 completed_at、progress=100.0、result_count,清除 error_message。 Args: result_count: 生成的视频数量,默认为 1 Raises: ValueError: 当前状态不允许转换到 completed """ self.transition_to(GenerationTaskStatus.COMPLETED) self.completed_at = datetime.now(timezone.utc) self.progress = 100.0 self.result_count = result_count self.error_message = "" def mark_failed(self, error_message: str, error_info: dict | None = None) -> None: """标记为失败(pending / running → failed)。 设置 error_message、error_info、completed_at。 Args: error_message: 错误信息 error_info: 结构化错误信息(error_type, stack_trace, stage, failed_at等) Raises: ValueError: 当前状态不允许转换到 failed """ self.transition_to(GenerationTaskStatus.FAILED) self.error_message = error_message self.completed_at = datetime.now(timezone.utc) if error_info is not None: self.error_info = error_info else: self.error_info = { "error_type": "UnknownError", "message": error_message, "failed_at": datetime.now(timezone.utc).isoformat(), } def mark_cancelled(self) -> None: """标记为已取消(pending / running → cancelled)。 设置 completed_at。 Raises: ValueError: 当前状态不允许转换到 cancelled """ self.transition_to(GenerationTaskStatus.CANCELLED) self.completed_at = datetime.now(timezone.utc) # ── 日志辅助 ──────────────────────────────────────────────────────────── _MAX_LOGS = 200 def append_log(self, stage: str, message: str, level: str = "INFO", **kwargs) -> None: """追加一条结构化日志到 logs 字段。 Args: stage: 阶段名称(如 "接收任务"、"下载素材"、"渲染") message: 日志消息 level: 日志级别(INFO / WARN / ERROR) **kwargs: 额外字段(如 asset_id、duration 等) """ try: entries = json.loads(self.logs) if self.logs else [] except (json.JSONDecodeError, TypeError): entries = [] entry = { "ts": datetime.now(timezone.utc).isoformat(), "level": level, "stage": stage, "message": message, **kwargs, } entries.append(entry) # 限制最多保留 _MAX_LOGS 条,防止字段过大 if len(entries) > self._MAX_LOGS: entries = entries[-self._MAX_LOGS :] self.logs = json.dumps(entries, ensure_ascii=False) def get_logs(self) -> list[dict]: """解析 logs 字段为 list[dict]。""" try: return json.loads(self.logs) if self.logs else [] except (json.JSONDecodeError, TypeError): return [] def mark_pending_from_failed(self) -> None: """从失败状态重置为待处理(用于重试)。 清除 error_message、error_info、started_at、completed_at、progress, 递增 retry_count。 Raises: ValueError: 当前状态不是 failed """ if self.status != GenerationTaskStatus.FAILED: raise ValueError(f"只有 failed 状态的任务可以重置为 pending,当前状态: {self.status.value}") self.transition_to(GenerationTaskStatus.PENDING) self.error_message = "" self.error_info = {} self.started_at = None self.completed_at = None self.progress = 0.0 self.result_count = 0 self.retry_count += 1