R524809
71 lines
2.3 KiB
Python
71 lines
2.3 KiB
Python
"""内存任务注册表：套图生成任务的生命周期与进程一致（重启即新会话）。
|
|
|
|
轮询/导出只服务「当前会话正在跟踪的任务」——前端没有历史记录功能，
|
|
任务状态无需跨进程持久化；重启后轮询自然 404，前端提示任务已中断。
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import uuid
|
|
from dataclasses import dataclass, field
|
|
|
|
# 任务状态
|
|
TASK_PENDING = "pending"
|
|
TASK_RUNNING = "running"
|
|
TASK_DONE = "done"
|
|
TASK_PARTIAL = "partial"
|
|
TASK_FAILED = "failed"
|
|
|
|
# 任务内单张图状态
|
|
IMG_PENDING = "pending" # 生成中（前端据此隐藏占位格，只渲染 ok/failed 终态）
|
|
IMG_OK = "ok"
|
|
IMG_FAILED = "failed"
|
|
|
|
|
|
@dataclass
|
|
class TaskImage:
|
|
"""任务里单张生成图：完成一张追加一条（前端进度 x/y 依赖此语义）。"""
|
|
|
|
type_id: str
|
|
name: str
|
|
status: str = IMG_PENDING # 循环里先建后跑，成功改 ok、失败显式改 failed
|
|
url: str = ""
|
|
error: str | None = None
|
|
|
|
|
|
@dataclass
|
|
class Task:
|
|
"""一次套图生成任务：轮询可见字段 + 仅供 run_suite 消费的执行参数。"""
|
|
|
|
id: str
|
|
status: str = TASK_PENDING
|
|
platform: str = "cn"
|
|
lang: str = "zh"
|
|
ratio: str = "1:1"
|
|
style_set: int = 1
|
|
style_prompt: str | None = None
|
|
requirements: str | None = None
|
|
provider: str = ""
|
|
model: str | None = None
|
|
total: int = 0 # 计划总张数（进度分母）
|
|
images: list[TaskImage] = field(default_factory=list)
|
|
error: str | None = None
|
|
# ── 执行参数（不进轮询响应）──
|
|
context: dict = field(default_factory=dict) # 采集文本素材（build_context 的输入）
|
|
plan: list[dict] = field(default_factory=list) # 展开后的逐张任务
|
|
ref_images: list[dict] = field(default_factory=list) # 参考图池（main 优先）
|
|
watermark: dict | None = None # 水印选项（落盘前服务端后处理）
|
|
|
|
|
|
# 进程内任务表：asyncio 单事件循环读写，无并发问题；不做淘汰（单会话量级很小）
|
|
_TASKS: dict[str, Task] = {}
|
|
|
|
|
|
def create_task(**kwargs) -> Task:
|
|
task = Task(id=uuid.uuid4().hex, **kwargs)
|
|
_TASKS[task.id] = task
|
|
return task
|
|
|
|
|
|
def get_task(task_id: str) -> Task | None:
|
|
return _TASKS.get(task_id)
|