baodan/api/insurance/models/generation_task.py

"""统一生成任务模型。"""
import uuid
from sqlalchemy import Column, String, Text, Integer, TIMESTAMP, func
from insurance.db.compat import db


class GenerationTask(db.Model):
    """统一生成任务表（PPT + 海报）。

    职责分离：
    - 工作区（PptSession/PosterRecord）保存可编辑数据
    - 任务表保存一次执行的不可变快照、状态和结果
    """
    __tablename__ = "insurance_generation_tasks"

    id = Column(String(36), primary_key=True, default=lambda: uuid.uuid4().hex)
    user_id = Column(String(64), nullable=False, comment="用户 ID")
    artifact_type = Column(String(10), nullable=False, comment="ppt/poster")
    operation = Column(String(10), nullable=False, comment="parse/generate")
    workspace_id = Column(String(36), nullable=False, comment="PptSession.id 或 PosterRecord.id")
    title_snapshot = Column(String(200), default="", comment="提交时任务名称")
    status = Column(String(20), nullable=False, default="queued",
                    comment="queued/running/done/failed/cancelled")
    stage = Column(String(30), default="", comment="当前执行阶段")
    progress = Column(Integer, default=0, comment="进度 0-100")
    message = Column(String(500), default="", comment="用户可读进度")
    error_code = Column(String(50), default="", comment="结构化错误代码")
    error_message = Column(Text, nullable=True, comment="用户可读错误")
    input_revision = Column(Integer, default=1, comment="本次生成使用的草稿版本")
    input_snapshot_json = Column(Text, nullable=True, comment="不可变输入快照")
    output_json = Column(Text, nullable=True, comment="输出路径、页数、预览等")
    idempotency_key = Column(String(100), nullable=True, comment="防止重复提交")
    celery_task_id = Column(String(200), nullable=True, comment="Celery ID")
    attempt_count = Column(Integer, default=0, comment="重试次数")
    heartbeat_at = Column(TIMESTAMP, nullable=True, comment="Worker 心跳")
    started_at = Column(TIMESTAMP, nullable=True, comment="开始时间")
    finished_at = Column(TIMESTAMP, nullable=True, comment="完成时间")
    viewed_at = Column(TIMESTAMP, nullable=True, comment="用户查看时间")
    dock_hidden_at = Column(TIMESTAMP, nullable=True, comment="从任务坞隐藏时间")
    created_at = Column(TIMESTAMP, server_default=func.now())
    updated_at = Column(TIMESTAMP, server_default=func.now(), onupdate=func.now())

    def to_dict(self):
        import json
        return {
            "id": self.id,
            "userId": self.user_id,
            "artifactType": self.artifact_type,
            "operation": self.operation,
            "workspaceId": self.workspace_id,
            "titleSnapshot": self.title_snapshot,
            "status": self.status,
            "stage": self.stage,
            "progress": self.progress,
            "message": self.message,
            "errorCode": self.error_code,
            "errorMessage": self.error_message,
            "inputRevision": self.input_revision,
            "inputSnapshot": json.loads(self.input_snapshot_json) if self.input_snapshot_json else None,
            "output": json.loads(self.output_json) if self.output_json else None,
            "idempotencyKey": self.idempotency_key,
            "celeryTaskId": self.celery_task_id,
            "attemptCount": self.attempt_count,
            "heartbeatAt": self.heartbeat_at.isoformat() if self.heartbeat_at else None,
            "startedAt": self.started_at.isoformat() if self.started_at else None,
            "finishedAt": self.finished_at.isoformat() if self.finished_at else None,
            "viewedAt": self.viewed_at.isoformat() if self.viewed_at else None,
            "dockHiddenAt": self.dock_hidden_at.isoformat() if self.dock_hidden_at else None,
            "createdAt": self.created_at.isoformat() if self.created_at else None,
            "updatedAt": self.updated_at.isoformat() if self.updated_at else None,
        }

    @staticmethod
    def create_task(user_id: str, artifact_type: str, operation: str,
                    workspace_id: str, title: str = "",
                    input_snapshot: dict = None, idempotency_key: str = None,
                    input_revision: int = None):
        """创建新任务（幂等检查）。"""
        import json
        if idempotency_key:
            existing = GenerationTask.query.filter_by(
                idempotency_key=idempotency_key,
            ).filter(GenerationTask.status.in_(["queued", "running", "done"])).first()
            if existing:
                return existing

        task = GenerationTask(
            user_id=user_id,
            artifact_type=artifact_type,
            operation=operation,
            workspace_id=workspace_id,
            title_snapshot=title,
            input_snapshot_json=json.dumps(input_snapshot, ensure_ascii=False) if input_snapshot else None,
            input_revision=input_revision or 1,
            idempotency_key=idempotency_key,
        )
        db.session.add(task)
        db.session.commit()
        return task