151 lines
8.8 KiB
Python
151 lines
8.8 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 ChatGenerationTask(Base, TimestampMixin, SoftDeleteMixin):
|
|
"""Project-independent AI chat/image/video generation task.
|
|
|
|
This table is intentionally NOT linked to projects. It is used by the
|
|
/generation-ai Celery pipeline so chat-style generation does not touch the
|
|
legacy generation_records -> projects foreign-key chain.
|
|
"""
|
|
|
|
__tablename__ = "chat_generation_tasks"
|
|
__table_args__ = (
|
|
# 防止前端按钮连点/网络重试时同一个 idempotency_key 并发创建多条任务。
|
|
# nullable unique 兼容不传 idempotency_key 的普通请求。
|
|
Index(
|
|
"uq_chat_generation_tasks_user_mode_idempotency",
|
|
"user_id",
|
|
"generation_mode",
|
|
"idempotency_key",
|
|
unique=True,
|
|
postgresql_where=text("deleted_at IS NULL AND idempotency_key IS NOT NULL"),
|
|
),
|
|
# AI 创作顶层任务在 chatapi_async/chatapi_main 之间切换时，
|
|
# 同一个前端幂等键也只能创建一组任务。
|
|
Index(
|
|
"uq_chat_generation_tasks_user_chat_idempotency",
|
|
"user_id",
|
|
"idempotency_key",
|
|
unique=True,
|
|
postgresql_where=text(
|
|
"deleted_at IS NULL "
|
|
"AND idempotency_key IS NOT NULL "
|
|
"AND generation_mode IN ('chatapi_async', 'chatapi_main')"
|
|
),
|
|
),
|
|
# 视频 24 小时降频轮询调度使用。
|
|
Index(
|
|
"idx_chat_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(
|
|
"uq_chat_generation_tasks_parent_index",
|
|
"parent_task_id",
|
|
"generation_index",
|
|
unique=True,
|
|
postgresql_where=text("parent_task_id IS NOT NULL AND generation_index IS NOT NULL"),
|
|
),
|
|
Index("idx_chat_generation_tasks_parent", "parent_task_id"),
|
|
Index("idx_chat_generation_tasks_user_mode_created", "user_id", "generation_mode", "created_at"),
|
|
CheckConstraint("generation_count BETWEEN 1 AND 5", name="ck_chat_generation_tasks_generation_count"),
|
|
CheckConstraint("generation_index IS NULL OR generation_index > 0", name="ck_chat_generation_tasks_generation_index"),
|
|
)
|
|
|
|
|
|
id: Mapped[str] = mapped_column(String(32), primary_key=True)
|
|
user_id: Mapped[str] = mapped_column(
|
|
String(32), ForeignKey("users.id", ondelete="CASCADE"), index=True, nullable=False
|
|
)
|
|
|
|
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", index=True)
|
|
|
|
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)
|
|
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)
|
|
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)
|
|
|
|
status: Mapped[str] = mapped_column(String(32), default="generating", index=True)
|
|
pipeline_stage: Mapped[str | None] = mapped_column(String(32), nullable=True, index=True)
|
|
generation_mode: Mapped[str] = mapped_column(String(32), default="chatapi_async", index=True)
|
|
parent_task_id: Mapped[str | None] = mapped_column(
|
|
String(32), ForeignKey("chat_generation_tasks.id", ondelete="RESTRICT"), nullable=True
|
|
)
|
|
generation_count: Mapped[int] = mapped_column(Integer, default=1, server_default="1", nullable=False)
|
|
generation_index: Mapped[int | None] = mapped_column(Integer, nullable=True)
|
|
generation_attempt_no: Mapped[int] = mapped_column(Integer, nullable=False, default=1, server_default="1")
|
|
resource_generation_started_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True, index=True)
|
|
|
|
# 图片主任务同步调用供应商时的分布式执行租约。
|
|
# 防止重复 Celery 消息或恢复任务同时触发多次组图请求。
|
|
provider_create_claim_token: Mapped[str | None] = mapped_column(String(64), nullable=True, index=True)
|
|
provider_create_lease_until: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True, index=True)
|
|
provider_create_started_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
|
|
|
media_references: Mapped[str | None] = mapped_column(Text, nullable=True)
|
|
provider_task_id: Mapped[str | None] = mapped_column(String(128), nullable=True, index=True)
|
|
# Kept as alias-compatible storage for existing frontend/service naming.
|
|
seedance_task_id: Mapped[str | None] = mapped_column(String(128), nullable=True, index=True)
|
|
|
|
remote_result_url: 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)
|
|
|
|
engine_id: Mapped[str | None] = mapped_column(String(32), nullable=True, index=True)
|
|
engine_snapshot_json: Mapped[str | None] = mapped_column(Text, nullable=True)
|
|
provider_response_json: Mapped[str | None] = mapped_column(Text, nullable=True)
|
|
|
|
credits_cost: Mapped[float] = mapped_column(Float, default=0.0)
|
|
text_credits_cost: Mapped[float] = mapped_column(Float, default=0.0)
|
|
text_tokens_used: Mapped[int] = mapped_column(Integer, default=0)
|
|
video_tokens_used: Mapped[int] = mapped_column(Integer, default=0)
|
|
image_tokens_used: Mapped[int] = mapped_column(Integer, default=0)
|
|
|
|
# retry_count is retained as a compatibility mirror of manual retries.
|
|
retry_count: Mapped[int] = mapped_column(Integer, default=0)
|
|
manual_retry_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0, server_default="0")
|
|
poll_error_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0, server_default="0")
|
|
poll_count: Mapped[int] = mapped_column(Integer, default=0)
|
|
last_poll_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
|
# 视频降频轮询调度字段。
|
|
# 图片同步生成仍沿用原超时逻辑；这些字段主要给 video + provider poll 使用。
|
|
poll_started_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
|
next_poll_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
|
poll_interval_seconds: Mapped[int] = mapped_column(Integer, default=0)
|
|
poll_claim_token: Mapped[str | None] = mapped_column(String(64), nullable=True, index=True)
|
|
poll_lease_until: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True, index=True)
|
|
deadline_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
|
generated_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
|
error_message: Mapped[str | None] = mapped_column(Text, nullable=True)
|
|
idempotency_key: Mapped[str | None] = mapped_column(String(64), nullable=True, index=True)
|
|
|
|
# Celery 下载容灾字段。
|
|
download_celery_task_id: Mapped[str | None] = mapped_column(String(160), nullable=True, index=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, index=True)
|
|
download_lease_until: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True, index=True)
|
|
download_next_retry_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True, index=True)
|
|
download_attempt_count: Mapped[int] = mapped_column(Integer, 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)
|