xiaoxia-saas/packages/domain/generation_task.py

from __future__ import annotations

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"


@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 = ""
    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 = ""
    is_preview: bool = True
    source_task_id: str = ""
    output_width: int = 1280
    output_height: int = 720
    cover_url: str = ""
    custom_title: 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 = "",
        is_preview: bool = True,
        source_task_id: str = "",
        output_width: int = 1280,
        output_height: int = 720,
        cover_url: str = "",
        custom_title: str = "",
    ) -> "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,
            is_preview=is_preview,
            source_task_id=source_task_id,
            output_width=output_width,
            output_height=output_height,
            cover_url=cover_url,
            custom_title=custom_title,
        )