133 lines
5.0 KiB
Python
133 lines
5.0 KiB
Python
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}"
|