feat(product-suite): persist current generation rounds

2026-07-16 23:27:25 +08:00
parent c146c0b41d
commit f3defdeb95
11 changed files with 806 additions and 16 deletions
@@ -302,6 +302,7 @@ CREATE TABLE IF NOT EXISTS image_studio_projects (
  target_detail_count INTEGER NOT NULL DEFAULT 12,
  draft_prompt       TEXT,
  suite_settings_json TEXT NOT NULL DEFAULT '{}',
+  current_generation_round_key TEXT,
  status             TEXT NOT NULL DEFAULT 'active',
  created_at         TEXT NOT NULL,
  updated_at         TEXT NOT NULL,
@@ -351,6 +352,8 @@ CREATE TABLE IF NOT EXISTS image_studio_jobs (
  points_cost        INTEGER,
  points_balance     INTEGER,
  call_id            TEXT,
+  generation_round_key TEXT,
+  generation_slot_index INTEGER,
  created_at         TEXT NOT NULL,
  updated_at         TEXT NOT NULL,
  submitted_at       TEXT,
@@ -484,6 +487,7 @@ def init_db(path=None, conn=None) -> None:
            _ensure_image_studio_project_suite_columns(database)
            _ensure_image_studio_project_draft_columns(database)
            _ensure_image_studio_job_recovery_columns(database)
+            _ensure_image_studio_generation_round_columns(database)


 def _ensure_batch_delete_columns(database):
@@ -578,6 +582,36 @@ def _ensure_image_studio_job_recovery_columns(database):
        """
    )

+
+def _ensure_image_studio_generation_round_columns(database):
+    project_columns = {
+        row["name"]
+        for row in database.execute("PRAGMA table_info(image_studio_projects)").fetchall()
+    }
+    if "current_generation_round_key" not in project_columns:
+        database.execute(
+            "ALTER TABLE image_studio_projects "
+            "ADD COLUMN current_generation_round_key TEXT"
+        )
+
+    job_columns = {
+        row["name"]
+        for row in database.execute("PRAGMA table_info(image_studio_jobs)").fetchall()
+    }
+    if "generation_round_key" not in job_columns:
+        database.execute(
+            "ALTER TABLE image_studio_jobs ADD COLUMN generation_round_key TEXT"
+        )
+    if "generation_slot_index" not in job_columns:
+        database.execute(
+            "ALTER TABLE image_studio_jobs ADD COLUMN generation_slot_index INTEGER"
+        )
+    database.execute(
+        "CREATE INDEX IF NOT EXISTS idx_image_studio_jobs_generation_round "
+        "ON image_studio_jobs("
+        "project_id, generation_round_key, generation_slot_index, id)"
+    )
+
 def create_batch(file_paths: Iterable[str], note=None, path=None, conn=None) -> str:
    batch_id = datetime.now().strftime("%Y%m%d_%H%M%S_") + uuid.uuid4().hex[:8]
    files = [os.path.abspath(file_path) for file_path in file_paths]
@@ -696,10 +696,12 @@ class SuiteTaskState:
    last_saved_prompt: str = ""
    settings: dict = field(default_factory=product_suite.default_suite_settings)
    current_job_ids: list = field(default_factory=list)
+    current_generation_round_key: str = ""
    show_history: bool = False
    generation_job_ids: list = field(default_factory=list)
    generation_mode: str = "batch"
    generation_retry_job_id: int = None
+    generation_round_key: str = ""
    worker: object = None
    thread: object = None
    generation_run_token: str = ""
@@ -1457,6 +1459,7 @@ class ProductSuiteTab(QWidget):

    def _load_state(self, state):
        self._sync_state_project_binding(state)
+        self._restore_current_generation_results(state)
        self._loading = True
        try:
            self.custom_category_edit.hide()
@@ -1723,6 +1726,8 @@ class ProductSuiteTab(QWidget):
            return None
        state.project_id = int(project.id)
        state.project_binding_state = project.binding_state
+        if previous_id != state.project_id:
+            self._restore_current_generation_results(state, force=True)
        stored_prompt = str(project.draft_prompt or "")
        if load_existing and previous_id != state.project_id:
            state.prompt = stored_prompt
@@ -3101,6 +3106,25 @@ class ProductSuiteTab(QWidget):
        if not specs:
            self._message("生成数量为0", "请至少把一个套图分类的数量设为1。")
            return False
+        generation_round_key = ""
+        if retrying:
+            try:
+                original_job = image_studio.get_job(retry_job_id, path=self.db_path)
+            except Exception as exc:
+                self._message("读取重试图片失败", _user_error(exc))
+                return False
+            if original_job is None or int(original_job.project_id) != int(state.project_id):
+                self._message("重试图片无效", "该图片不属于当前商品套图任务。")
+                return False
+            for spec in specs:
+                spec["generation_round_key"] = original_job.generation_round_key
+                spec["generation_slot_index"] = original_job.generation_slot_index
+            generation_round_key = str(original_job.generation_round_key or "")
+        else:
+            generation_round_key = uuid.uuid4().hex
+            for slot_index, spec in enumerate(specs):
+                spec["generation_round_key"] = generation_round_key
+                spec["generation_slot_index"] = slot_index
        if confirm_batch and not self._confirm(
            "确认生成商品套图",
            self._generation_confirmation_message(
@@ -3118,6 +3142,7 @@ class ProductSuiteTab(QWidget):
            state.project_id,
            specs,
            run_token=run_token,
+            generation_round_key=generation_round_key or None,
            aspect_ratio=state.settings["ratio"],
            db_path=self.db_path,
            config=self.config,
@@ -3130,6 +3155,7 @@ class ProductSuiteTab(QWidget):
        state.generation_job_ids = []
        state.generation_mode = "retry" if retrying else "batch"
        state.generation_retry_job_id = retry_job_id
+        state.generation_round_key = generation_round_key
        state.done = 0
        state.failed = 0
        state.total = len(specs)
@@ -3406,6 +3432,30 @@ class ProductSuiteTab(QWidget):
                    current.append(job_id)
        state.current_job_ids = current

+    def _restore_current_generation_results(self, state, *, force=False, allow_running=False):
+        if state.project_id is None or (
+            state.generation_running() and not allow_running
+        ):
+            return False
+        if state.current_job_ids and not force:
+            return False
+        try:
+            round_key = image_studio.get_current_generation_round(
+                state.project_id,
+                path=self.db_path,
+            )
+            jobs = image_studio.list_generation_round_current_jobs(
+                state.project_id,
+                round_key,
+                path=self.db_path,
+            )
+        except Exception as exc:
+            self._status("当前生成结果恢复失败：%s" % _user_error(exc), "danger")
+            return False
+        state.current_generation_round_key = str(round_key or "")
+        state.current_job_ids = [int(job.id) for job in jobs]
+        return True
+
    def _generation_job_snapshot(self, state):
        job_ids = self._generation_job_ids(state)
        counts = {
@@ -3501,12 +3551,40 @@ class ProductSuiteTab(QWidget):
            if state.started_at
            else 0
        )
+        promoted = False
+        if not retrying and state.generation_round_key and success:
+            try:
+                promoted = image_studio.promote_generation_round_if_success(
+                    state.project_id,
+                    state.generation_round_key,
+                    path=self.db_path,
+                )
+            except Exception as exc:
+                self._status("当前生成轮次保存失败：%s" % _user_error(exc), "danger")
+            self._restore_current_generation_results(
+                state,
+                force=True,
+                allow_running=True,
+            )
+        elif not retrying:
+            self._restore_current_generation_results(
+                state,
+                force=True,
+                allow_running=True,
+            )
+        elif state.generation_round_key:
+            self._restore_current_generation_results(
+                state,
+                force=True,
+                allow_running=True,
+            )
        self._generation_run_states.pop(run_token, None)
        state.generation_run_token = ""
        state.generation_stop_requested = False
        state.generation_terminal_streak = 0
        state.generation_job_ids = []
        state.generation_retry_job_id = None
+        state.generation_round_key = ""
        state.worker = None
        state.thread = None
        state.done = success + failed + cancelled
@@ -3530,6 +3608,7 @@ class ProductSuiteTab(QWidget):
                "active": active,
                "elapsed_seconds": elapsed,
                "mode": "retry" if retrying else "batch",
+                "current_round_promoted": promoted,
            },
            level="WARNING" if active or result.get("ok") is False else "INFO",
        )
@@ -334,6 +334,7 @@ class ProductSuiteGenerateWorker(BaseWorker):
        job_specs,
        *,
        run_token="",
+        generation_round_key=None,
        aspect_ratio="1:1",
        db_path=None,
        config=None,
@@ -343,6 +344,7 @@ class ProductSuiteGenerateWorker(BaseWorker):
        self.project_id = int(project_id)
        self.job_specs = [dict(spec) for spec in (job_specs or [])]
        self.run_token = str(run_token or "")
+        self.generation_round_key = str(generation_round_key or "").strip()
        self.aspect_ratio = str(aspect_ratio or "1:1")
        self.db_path = db_path
        self.config = config
@@ -366,6 +368,8 @@ class ProductSuiteGenerateWorker(BaseWorker):
                    prompt=spec.get("prompt") or "",
                    generation_source="cmhub",
                    provider="cmhub",
+                    generation_round_key=spec.get("generation_round_key") or self.generation_round_key or None,
+                    generation_slot_index=spec.get("generation_slot_index"),
                    path=self.db_path,
                )
            )
@@ -410,6 +414,7 @@ class ProductSuiteGenerateWorker(BaseWorker):
        summary["job_ids"] = list(self.job_ids)
        summary["cancelled_count"] = int(summary.get("cancelled", 0) or 0)
        summary["run_token"] = self.run_token
+        summary["generation_round_key"] = self.generation_round_key or None
        return summary


@@ -49,6 +49,7 @@ class ImageStudioProject:
    target_detail_count: int
    draft_prompt: Optional[str]
    suite_settings_json: str
+    current_generation_round_key: Optional[str]
    status: str
    created_at: str
    updated_at: str
@@ -91,6 +92,8 @@ class ImageStudioJob:
    points_cost: Optional[int]
    points_balance: Optional[int]
    call_id: Optional[str]
+    generation_round_key: Optional[str]
+    generation_slot_index: Optional[int]
    created_at: str
    updated_at: str
    submitted_at: Optional[str]
@@ -108,6 +111,23 @@ class ImageStudioSelection:
    updated_at: str


+@dataclass(frozen=True)
+class ImageStudioGenerationRound:
+    project_id: int
+    generation_round_key: Optional[str]
+    created_at: Optional[str]
+    updated_at: Optional[str]
+    job_count: int
+    slot_count: int
+    succeeded_count: int
+    failed_count: int
+    cancelled_count: int
+    active_count: int
+    retry_count: int
+    is_current: bool
+    is_legacy: bool
+
+
 class ImageStudioError(RuntimeError):
    """Raised when the AI image studio service cannot complete an operation."""

@@ -1059,11 +1079,24 @@ def create_job(
    task_key=None,
    generation_source="cmhub",
    provider="cmhub",
+    generation_round_key=None,
+    generation_slot_index=None,
    path=None,
    conn=None,
 ):
    now = _now()
    task_key = str(task_key or _task_key(project_id))
+    if generation_round_key is not None:
+        generation_round_key = str(generation_round_key).strip()
+        if not generation_round_key:
+            raise db.DbError("AI工场生成轮次标识不能为空")
+    if generation_slot_index is not None:
+        try:
+            generation_slot_index = int(generation_slot_index)
+        except (TypeError, ValueError) as exc:
+            raise db.DbError("AI工场生成槽位编号无效") from exc
+        if generation_slot_index < 0:
+            raise db.DbError("AI工场生成槽位编号不能小于0")
    with _connection(conn, path) as database:
        try:
            with database:
@@ -1072,8 +1105,9 @@ def create_job(
                    """
                    INSERT INTO image_studio_jobs
                      (project_id, source_asset_id, generation_source, provider,
-                       job_type, task_key, status, prompt, created_at, updated_at)
-                    VALUES (?, ?, ?, ?, ?, ?, 'pending', ?, ?, ?)
+                       job_type, task_key, status, prompt, generation_round_key,
+                       generation_slot_index, created_at, updated_at)
+                    VALUES (?, ?, ?, ?, ?, ?, 'pending', ?, ?, ?, ?, ?)
                    """,
                    (
                        int(project_id),
@@ -1083,6 +1117,8 @@ def create_job(
                        str(job_type),
                        task_key,
                        str(prompt or ""),
+                        generation_round_key,
+                        generation_slot_index,
                        now,
                        now,
                    ),
@@ -1120,6 +1156,198 @@ def list_jobs(project_id, *, statuses=None, path=None, conn=None):
        return _fetch_all(database, sql, params, ImageStudioJob)


+def get_current_generation_round(project_id, path=None, conn=None):
+    project = get_project(project_id, path=path, conn=conn)
+    if project is None:
+        return None
+    value = project.current_generation_round_key
+    return str(value).strip() if value else None
+
+
+def set_current_generation_round(project_id, generation_round_key, path=None, conn=None):
+    project_id = int(project_id)
+    round_key = None
+    if generation_round_key is not None:
+        round_key = str(generation_round_key).strip()
+        if not round_key:
+            raise db.DbError("当前生成轮次标识不能为空")
+    with _connection(conn, path) as database:
+        with database:
+            project = get_project(project_id, conn=database)
+            if project is None:
+                raise db.DbError("AI工场商品项目不存在或已删除")
+            if round_key is not None:
+                row = database.execute(
+                    """
+                    SELECT 1 FROM image_studio_jobs
+                    WHERE project_id = ? AND generation_round_key = ?
+                    LIMIT 1
+                    """,
+                    (project_id, round_key),
+                ).fetchone()
+                if row is None:
+                    raise db.DbError("当前生成轮次不属于该商品项目")
+            database.execute(
+                """
+                UPDATE image_studio_projects
+                SET current_generation_round_key = ?, updated_at = ?
+                WHERE id = ? AND deleted_at IS NULL
+                """,
+                (round_key, _now(), project_id),
+            )
+        return get_project(project_id, conn=database)
+
+
+def promote_generation_round_if_success(project_id, generation_round_key, path=None, conn=None):
+    project_id = int(project_id)
+    round_key = str(generation_round_key or "").strip()
+    if not round_key:
+        raise db.DbError("当前生成轮次标识不能为空")
+    with _connection(conn, path) as database:
+        with database:
+            project = get_project(project_id, conn=database)
+            if project is None:
+                raise db.DbError("AI工场商品项目不存在或已删除")
+            success = database.execute(
+                """
+                SELECT 1 FROM image_studio_jobs
+                WHERE project_id = ?
+                  AND generation_round_key = ?
+                  AND status = 'succeeded'
+                LIMIT 1
+                """,
+                (project_id, round_key),
+            ).fetchone()
+            if success is None:
+                return False
+            database.execute(
+                """
+                UPDATE image_studio_projects
+                SET current_generation_round_key = ?, updated_at = ?
+                WHERE id = ? AND deleted_at IS NULL
+                """,
+                (round_key, _now(), project_id),
+            )
+        return True
+
+
+def list_generation_rounds(project_id, *, limit=None, offset=0, path=None, conn=None):
+    project_id = int(project_id)
+    try:
+        offset = max(0, int(offset or 0))
+    except (TypeError, ValueError) as exc:
+        raise db.DbError("生成轮次偏移量无效") from exc
+    if limit is not None:
+        try:
+            limit = int(limit)
+        except (TypeError, ValueError) as exc:
+            raise db.DbError("生成轮次数量无效") from exc
+        if limit <= 0:
+            return []
+    current_round_key = get_current_generation_round(project_id, path=path, conn=conn)
+    sql = """
+        SELECT generation_round_key,
+               MIN(created_at) AS created_at,
+               MAX(updated_at) AS updated_at,
+               COUNT(*) AS job_count,
+               COUNT(DISTINCT generation_slot_index) AS slot_count,
+               SUM(CASE WHEN status = 'succeeded' THEN 1 ELSE 0 END) AS succeeded_count,
+               SUM(CASE WHEN status IN ('failed', 'expired') THEN 1 ELSE 0 END) AS failed_count,
+               SUM(CASE WHEN status = 'cancelled' THEN 1 ELSE 0 END) AS cancelled_count,
+               SUM(CASE WHEN status IN ('pending', 'submitted', 'running') THEN 1 ELSE 0 END) AS active_count
+        FROM image_studio_jobs
+        WHERE project_id = ?
+        GROUP BY generation_round_key
+        ORDER BY CASE WHEN generation_round_key IS NULL THEN 1 ELSE 0 END,
+                 MAX(id) DESC
+    """
+    params = [project_id]
+    if limit is not None:
+        sql += " LIMIT ? OFFSET ?"
+        params.extend([limit, offset])
+    elif offset:
+        sql += " LIMIT -1 OFFSET ?"
+        params.append(offset)
+    with _connection(conn, path) as database:
+        rows = database.execute(sql, params).fetchall()
+    rounds = []
+    for row in rows:
+        round_key = row["generation_round_key"]
+        job_count = int(row["job_count"] or 0)
+        slot_count = int(row["slot_count"] or 0)
+        rounds.append(
+            ImageStudioGenerationRound(
+                project_id=project_id,
+                generation_round_key=round_key,
+                created_at=row["created_at"],
+                updated_at=row["updated_at"],
+                job_count=job_count,
+                slot_count=slot_count,
+                succeeded_count=int(row["succeeded_count"] or 0),
+                failed_count=int(row["failed_count"] or 0),
+                cancelled_count=int(row["cancelled_count"] or 0),
+                active_count=int(row["active_count"] or 0),
+                retry_count=max(0, job_count - slot_count) if round_key else 0,
+                is_current=bool(round_key and round_key == current_round_key),
+                is_legacy=round_key is None,
+            )
+        )
+    return rounds
+
+
+def list_generation_round_current_jobs(project_id, generation_round_key, path=None, conn=None):
+    project_id = int(project_id)
+    if generation_round_key is None:
+        sql = """
+            SELECT * FROM image_studio_jobs
+            WHERE project_id = ? AND generation_round_key IS NULL
+            ORDER BY created_at ASC, id ASC
+        """
+        params = (project_id,)
+    else:
+        round_key = str(generation_round_key).strip()
+        if not round_key:
+            raise db.DbError("生成轮次标识不能为空")
+        sql = """
+            SELECT jobs.*
+            FROM image_studio_jobs AS jobs
+            INNER JOIN (
+                SELECT generation_slot_index, MAX(id) AS latest_id
+                FROM image_studio_jobs
+                WHERE project_id = ? AND generation_round_key = ?
+                GROUP BY generation_slot_index
+            ) AS latest ON latest.latest_id = jobs.id
+            WHERE jobs.project_id = ? AND jobs.generation_round_key = ?
+            ORDER BY jobs.generation_slot_index ASC, jobs.id ASC
+        """
+        params = (project_id, round_key, project_id, round_key)
+    with _connection(conn, path) as database:
+        return _fetch_all(database, sql, params, ImageStudioJob)
+
+
+def list_generation_round_attempts(project_id, generation_round_key, path=None, conn=None):
+    project_id = int(project_id)
+    if generation_round_key is None:
+        sql = """
+            SELECT * FROM image_studio_jobs
+            WHERE project_id = ? AND generation_round_key IS NULL
+            ORDER BY created_at ASC, id ASC
+        """
+        params = (project_id,)
+    else:
+        round_key = str(generation_round_key).strip()
+        if not round_key:
+            raise db.DbError("生成轮次标识不能为空")
+        sql = """
+            SELECT * FROM image_studio_jobs
+            WHERE project_id = ? AND generation_round_key = ?
+            ORDER BY generation_slot_index ASC, id ASC
+        """
+        params = (project_id, round_key)
+    with _connection(conn, path) as database:
+        return _fetch_all(database, sql, params, ImageStudioJob)
+
+
 def set_job_submitted(job_id, task_id, *, call_id=None, points_cost=None, points_balance=None, path=None, conn=None):
    now = _now()
    with _connection(conn, path) as database: