cmhub/apps/api/models.py

from __future__ import annotations

import uuid

from django.conf import settings
from django.db import models
from django.db.models import Q


class ImageGenerationTask(models.Model):
    class Status(models.TextChoices):
        QUEUED = "queued", "排队中"
        RUNNING = "running", "处理中"
        SUCCEEDED = "succeeded", "成功"
        FAILED = "failed", "失败"
        EXPIRED = "expired", "已过期"

    task_id = models.UUIDField("任务 ID", default=uuid.uuid4, unique=True, editable=False)
    user = models.ForeignKey(
        settings.AUTH_USER_MODEL,
        verbose_name="用户",
        on_delete=models.PROTECT,
        related_name="image_generation_tasks",
    )
    api_key = models.ForeignKey(
        "users.ApiKey",
        verbose_name="API 密钥",
        on_delete=models.PROTECT,
        related_name="image_generation_tasks",
    )
    call_record = models.OneToOneField(
        "billing.CallRecord",
        verbose_name="调用记录",
        on_delete=models.PROTECT,
        related_name="image_generation_task",
    )
    status = models.CharField(
        "状态",
        max_length=20,
        choices=Status.choices,
        default=Status.QUEUED,
    )
    idempotency_key = models.CharField("幂等键", max_length=128, blank=True)
    idempotency_key_hash = models.CharField(
        "幂等键哈希",
        max_length=64,
        null=True,
        blank=True,
        editable=False,
    )
    request_hash = models.CharField("请求哈希", max_length=64)
    request_payload = models.JSONField("请求快照", default=dict, blank=True)
    input_image = models.FileField(
        "输入图片",
        upload_to="generated/task_inputs/%Y/%m/%d/",
        blank=True,
    )
    result_url = models.TextField("结果 URL", blank=True)
    error_code = models.CharField("错误码", max_length=64, blank=True)
    error_message = models.TextField("错误信息", blank=True)
    points_balance_after_charge = models.BigIntegerField("扣费后余额", default=0)
    started_at = models.DateTimeField("开始时间", null=True, blank=True)
    finished_at = models.DateTimeField("完成时间", null=True, blank=True)
    expires_at = models.DateTimeField("任务元数据过期时间", null=True, blank=True)
    next_attempt_at = models.DateTimeField("下次重试时间", null=True, blank=True)
    locked_at = models.DateTimeField("锁定时间", null=True, blank=True)
    lease_expires_at = models.DateTimeField("租约过期时间", null=True, blank=True)
    heartbeat_at = models.DateTimeField("心跳时间", null=True, blank=True)
    worker_id = models.CharField("Worker ID", max_length=128, blank=True)
    attempt_count = models.PositiveIntegerField("尝试次数", default=0)
    created_at = models.DateTimeField("创建时间", auto_now_add=True)
    updated_at = models.DateTimeField("更新时间", auto_now=True)

    class Meta:
        db_table = "image_generation_task"
        verbose_name = "图片生成任务"
        verbose_name_plural = "图片生成任务"
        ordering = ("-created_at", "-id")
        constraints = [
            models.UniqueConstraint(
                fields=("api_key", "idempotency_key_hash"),
                name="unique_image_task_idempotency_hash_per_api_key",
            ),
            models.CheckConstraint(
                condition=Q(points_balance_after_charge__gte=0),
                name="image_task_balance_after_charge_non_negative",
            ),
        ]
        indexes = [
            models.Index(fields=("user", "created_at")),
            models.Index(fields=("api_key", "created_at")),
            models.Index(fields=("status", "created_at")),
            models.Index(fields=("status", "next_attempt_at")),
            models.Index(fields=("status", "lease_expires_at")),
            models.Index(fields=("worker_id", "status")),
            models.Index(fields=("expires_at",)),
        ]

    def __str__(self) -> str:
        return f"{self.task_id} {self.status}"


class ImageGenerationTaskInput(models.Model):
    task = models.ForeignKey(
        ImageGenerationTask,
        verbose_name="图片生成任务",
        on_delete=models.CASCADE,
        related_name="input_images",
    )
    ordinal = models.PositiveSmallIntegerField("输入顺序")
    image = models.FileField(
        "输入图片",
        upload_to="generated/task_inputs/%Y/%m/%d/",
    )
    mime_type = models.CharField("MIME 类型", max_length=100)
    filename = models.CharField("原始文件名", max_length=255)
    created_at = models.DateTimeField("创建时间", auto_now_add=True)

    class Meta:
        db_table = "image_generation_task_input"
        verbose_name = "图片生成任务输入"
        verbose_name_plural = "图片生成任务输入"
        ordering = ("ordinal", "id")
        constraints = [
            models.UniqueConstraint(
                fields=("task", "ordinal"),
                name="unique_image_task_input_ordinal",
            )
        ]

    def __str__(self) -> str:
        return f"{self.task.task_id} #{self.ordinal + 1}"