126 lines
7.3 KiB
Python
126 lines
7.3 KiB
Python
from datetime import datetime
|
||
|
||
from sqlalchemy import Boolean, CheckConstraint, DateTime, Float, ForeignKey, Index, Integer, String, Text, text
|
||
from sqlalchemy.orm import Mapped, mapped_column
|
||
|
||
from app.models.base import Base, TimestampMixin, SoftDeleteMixin
|
||
|
||
|
||
class ApiGenerationTask(Base, TimestampMixin, SoftDeleteMixin):
|
||
"""对外开放 API 的生成任务表。
|
||
|
||
该表设计满足 ProviderGenerationRecordLike 协议，
|
||
使现有的 Volcano Ark SDK 封装函数可以直接复用。
|
||
"""
|
||
|
||
__tablename__ = "api_generation_tasks"
|
||
__table_args__ = (
|
||
# 幂等键唯一索引
|
||
Index(
|
||
"uq_api_generation_tasks_key_idempotency",
|
||
"api_key_id",
|
||
"external_idempotency_key",
|
||
unique=True,
|
||
postgresql_where=text("deleted_at IS NULL AND external_idempotency_key IS NOT NULL"),
|
||
),
|
||
# 视频轮询调度索引
|
||
Index(
|
||
"idx_api_generation_tasks_next_poll_at",
|
||
"next_poll_at",
|
||
postgresql_where=text(
|
||
"deleted_at IS NULL "
|
||
"AND status = 'generating' "
|
||
"AND gen_type = 'video' "
|
||
"AND next_poll_at IS NOT NULL"
|
||
),
|
||
),
|
||
Index("idx_api_generation_tasks_api_key_created", "api_key_id", "created_at"),
|
||
Index("idx_api_generation_tasks_provider_task_id", "provider_task_id"),
|
||
Index("idx_api_generation_tasks_status", "status"),
|
||
CheckConstraint("generation_count BETWEEN 1 AND 5", name="ck_api_generation_tasks_generation_count"),
|
||
)
|
||
|
||
id: Mapped[str] = mapped_column(String(32), primary_key=True)
|
||
api_key_id: Mapped[str] = mapped_column(
|
||
String(32), ForeignKey("api_keys.id", ondelete="CASCADE"), nullable=False, index=True
|
||
)
|
||
external_idempotency_key: Mapped[str | None] = mapped_column(String(64), nullable=True)
|
||
|
||
# === ProviderGenerationRecordLike 协议字段 ===
|
||
original_prompt: Mapped[str] = mapped_column(Text, nullable=False)
|
||
optimized_prompt: Mapped[str | None] = mapped_column(Text, nullable=True)
|
||
gen_type: Mapped[str] = mapped_column(String(16), default="video", nullable=False)
|
||
duration: Mapped[int | None] = mapped_column(Integer, nullable=True)
|
||
aspect_ratio: Mapped[str | None] = mapped_column(String(8), nullable=True)
|
||
resolution: Mapped[str | None] = mapped_column(String(8), nullable=True)
|
||
provider_generation_resolution: Mapped[str | None] = mapped_column(String(16), nullable=True)
|
||
image_size: Mapped[str | None] = mapped_column(String(16), nullable=True)
|
||
image_proportion: Mapped[str | None] = mapped_column(String(8), nullable=True)
|
||
image_px: Mapped[str | None] = mapped_column(String(16), nullable=True)
|
||
generation_count: Mapped[int] = mapped_column(Integer, default=1, server_default="1", nullable=False)
|
||
engine_id: Mapped[str | None] = mapped_column(String(32), nullable=True)
|
||
model_name: Mapped[str] = mapped_column(String(128), nullable=False, comment="模型名称")
|
||
media_references: Mapped[str | None] = mapped_column(Text, nullable=True, comment="用户原始上传的媒体URL")
|
||
local_media_json: Mapped[str | None] = mapped_column(Text, nullable=True, comment="下载到本地的媒体文件路径JSON")
|
||
engine_snapshot_json: Mapped[str | None] = mapped_column(Text, nullable=True)
|
||
|
||
# === 请求参数快照 ===
|
||
request_params_json: Mapped[str | None] = mapped_column(Text, nullable=True, comment="完整原始请求参数")
|
||
|
||
# === 流水线状态（镜像 ChatGenerationTask） ===
|
||
status: Mapped[str] = mapped_column(String(32), default="pending", nullable=False)
|
||
pipeline_stage: Mapped[str | None] = mapped_column(String(32), nullable=True)
|
||
generation_attempt_no: Mapped[int] = mapped_column(Integer, default=1, server_default="1", nullable=False)
|
||
resource_generation_started_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
||
deadline_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
||
|
||
# === 供应商交互 ===
|
||
provider_task_id: Mapped[str | None] = mapped_column(String(128), nullable=True)
|
||
remote_result_url: Mapped[str | None] = mapped_column(Text, nullable=True)
|
||
provider_response_json: Mapped[str | None] = mapped_column(Text, nullable=True)
|
||
|
||
# === 结果 ===
|
||
image_url: Mapped[str | None] = mapped_column(String(512), nullable=True)
|
||
video_url: Mapped[str | None] = mapped_column(String(512), nullable=True)
|
||
video_cover_url: Mapped[str | None] = mapped_column(String(512), nullable=True)
|
||
error_message: Mapped[str | None] = mapped_column(Text, nullable=True)
|
||
generated_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
||
|
||
# === 超分 ===
|
||
video_upscale_enabled_snapshot: Mapped[bool] = mapped_column(
|
||
Boolean, nullable=False, default=False, server_default="false"
|
||
)
|
||
video_upscale_snapshot_json: Mapped[str | None] = mapped_column(Text, nullable=True)
|
||
|
||
# === 配额消耗 ===
|
||
credits_cost: Mapped[float] = mapped_column(Float, default=0.0, server_default="0.0")
|
||
video_tokens_used: Mapped[int] = mapped_column(Integer, default=0, server_default="0")
|
||
image_tokens_used: Mapped[int] = mapped_column(Integer, default=0, server_default="0")
|
||
|
||
# === 轮询控制 ===
|
||
next_poll_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
||
poll_interval_seconds: Mapped[int] = mapped_column(Integer, default=30, server_default="30")
|
||
poll_count: Mapped[int] = mapped_column(Integer, default=0, server_default="0")
|
||
last_poll_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
||
|
||
# === Celery 执行租约（镜像 ChatGenerationTask） ===
|
||
provider_create_claim_token: Mapped[str | None] = mapped_column(String(64), nullable=True)
|
||
provider_create_lease_until: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
||
provider_create_started_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
||
poll_started_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
||
poll_claim_token: Mapped[str | None] = mapped_column(String(64), nullable=True)
|
||
poll_lease_until: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
||
poll_error_count: Mapped[int] = mapped_column(Integer, default=0, server_default="0")
|
||
download_celery_task_id: Mapped[str | None] = mapped_column(String(160), nullable=True)
|
||
download_enqueued_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
||
download_started_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
||
download_claim_token: Mapped[str | None] = mapped_column(String(64), nullable=True)
|
||
download_lease_until: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
||
download_next_retry_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
||
download_attempt_count: Mapped[int] = mapped_column(Integer, default=0, server_default="0")
|
||
download_last_error: Mapped[str | None] = mapped_column(Text, nullable=True)
|
||
download_storage_date_dir: Mapped[str | None] = mapped_column(String(16), nullable=True)
|
||
|
||
# === 存储 ===
|
||
local_path: Mapped[str | None] = mapped_column(Text, nullable=True)
|