"""Data service for the AI image studio.""" from __future__ import annotations import json import os import uuid from contextlib import contextmanager from dataclasses import dataclass from datetime import datetime from typing import Iterable, Optional from . import accounts, appconfig, chrome, db, editor from .config import make_slug PROJECT_STATUS_ACTIVE = "active" PROJECT_BINDING_DRAFT = "draft" PROJECT_BINDING_BOUND = "bound" PROJECT_BINDING_STATES = {PROJECT_BINDING_DRAFT, PROJECT_BINDING_BOUND} TEMPORARY_ITEM_PREFIX = "draft_" ASSET_KIND_ORIGINAL = "original" ASSET_STATUS_AVAILABLE = "available" ASSET_STATUS_MISSING = "missing" ASSET_STATUSES = {ASSET_STATUS_AVAILABLE, ASSET_STATUS_MISSING} JOB_STATUSES = {"pending", "submitted", "running", "succeeded", "failed", "expired", "cancelled"} JOB_RESUMABLE_STATUSES = {"submitted", "running"} JOB_RECOVERY_NONE = "none" JOB_RECOVERY_RESUME = "resume" JOB_RECOVERY_REGENERATE = "regenerate" JOB_RECOVERY_ACTIONS = { JOB_RECOVERY_NONE, JOB_RECOVERY_RESUME, JOB_RECOVERY_REGENERATE, } SELECTION_TYPES = {"main", "detail"} @dataclass(frozen=True) class ImageStudioProject: id: int account_alias: str account_slug: str account_name: Optional[str] item_id: str storage_key: str binding_state: str target_main_count: int 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 deleted_at: Optional[str] deleted_reason: Optional[str] @dataclass(frozen=True) class ImageStudioAsset: id: int project_id: int kind: str remote_url: Optional[str] local_path: Optional[str] aspect_ratio: Optional[str] parent_asset_id: Optional[int] prompt: Optional[str] status: str source_order: int created_at: str updated_at: str @dataclass(frozen=True) class ImageStudioJob: id: int project_id: int source_asset_id: Optional[int] output_asset_id: Optional[int] generation_source: str provider: str job_type: str task_key: str task_id: Optional[str] status: str prompt: Optional[str] error: Optional[str] recovery_action: str attempts: int 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] finished_at: Optional[str] @dataclass(frozen=True) class ImageStudioSelection: id: int project_id: int selection_type: str position: int asset_id: int created_at: str 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.""" class ImageStudioProjectConflictError(ImageStudioError): """Raised when a draft cannot be bound because the formal project already exists.""" class ImageStudioPullCancelled(ImageStudioError): """Raised when a read-only Shopee image pull stops at a safe boundary.""" def __init__(self, project=None, assets=None): super().__init__("用户停止拉取蝦皮主图") self.project = project self.assets = list(assets or []) def _now() -> str: return datetime.now().isoformat(timespec="seconds") @contextmanager def _connection(conn=None, path=None): with db._connection(conn, path) as database: yield database def _row_to_dataclass(row, cls): return None if row is None else cls(**dict(row)) def _fetch_one(conn, sql, params, cls): return _row_to_dataclass(conn.execute(sql, params).fetchone(), cls) def _fetch_all(conn, sql, params, cls): return [_row_to_dataclass(row, cls) for row in conn.execute(sql, params).fetchall()] def _get(obj, name, default=None): if obj is None: return default if isinstance(obj, dict): return obj.get(name, default) return getattr(obj, name, default) def _safe_component(value, default): text = str(value or "").strip() if not text: text = str(default) safe = "".join(ch if ch.isalnum() or ch in "_-" else "_" for ch in text).strip("_") return safe or str(default) def _account_fields(account=None, account_alias=None, account_name=None, account_slug=None): alias = str(account_alias or _get(account, "alias") or "").strip() if not alias: raise db.DbError("AI工场项目缺少账号别名") name = str(account_name or _get(account, "account_name") or "").strip() or None slug = str(account_slug or _get(account, "slug") or "").strip() or make_slug(alias) return alias, name, _safe_component(slug, "account") def _task_key(project_id): return f"image-studio-{int(project_id)}-{uuid.uuid4().hex}" def _db_path(path=None, config=None) -> str: return path or appconfig.db_path(config) def _normalize_item_id(item_id) -> str: text = str(item_id or "").strip() if not text: raise ImageStudioError("商品ID不能为空") return text def is_formal_item_id(item_id) -> bool: return str(item_id or "").strip().isdigit() def is_draft_project(project) -> bool: return str(_get(project, "binding_state", "")).strip() == PROJECT_BINDING_DRAFT def project_storage_key(project) -> str: return str(_get(project, "storage_key") or _get(project, "item_id") or "").strip() def _new_draft_item_id() -> str: return TEMPORARY_ITEM_PREFIX + uuid.uuid4().hex def _notify_step(callback, step, result="start", detail=None): if callback is None: return payload = {"step": step, "result": result} if detail: payload["detail"] = str(detail) try: callback(payload) except Exception: pass def _is_definitive_logged_out(status): reason = str((status or {}).get("reason") or "").strip() url = str((status or {}).get("url") or "").lower() return reason.startswith("LOGIN_PAGE") or ( "accounts.shopee." in url and "/seller/login" in url ) def _login_status_detail(status): status = status or {} reason = status.get("reason") or "未知原因" url = status.get("url") or "未知URL" cookie_names = [str(name) for name in (status.get("cookie_names") or []) if name] cookie_text = ",".join(sorted(cookie_names)) if cookie_names else "未读到登录Cookie" return f"原因={reason}，URL={url}，Cookie名称={cookie_text}" def _ensure_account_ready_for_read(account, *, path=None, config=None, login_timeout=8, on_step=None): port = accounts.normalize_debug_port(_get(account, "debug_port")) alias = str(_get(account, "alias") or "").strip() _notify_step(on_step, "ensure_chrome", "start", f"账号 {alias} debug_port={port}") if chrome.is_running(port): chrome_info = { "action": "reused", "reused": True, "launched": False, "debug_port": port, } _notify_step(on_step, "ensure_chrome", "reused", f"账号 {alias} Chrome 已打开") else: try: chrome_info = accounts.launch_for_login(account, path=path, config=config) except Exception as exc: raise ImageStudioError(f"账号 {alias} Chrome 启动失败：{exc}") from exc action = "launched" if chrome_info.get("launched") else "reused" _notify_step(on_step, "ensure_chrome", action, f"账号 {alias} Chrome 已{ '启动' if action == 'launched' else '复用' }") _notify_step(on_step, "login_check", "start", f"账号 {alias}") try: status = accounts.detect_login(account, timeout=login_timeout, path=path, config=config) except Exception as exc: status = {"logged_in": False, "reason": f"LOGIN_CHECK_FAILED: {exc}"} if status.get("logged_in"): _notify_step(on_step, "login_check", "success", f"账号 {alias} 已登录") return {**chrome_info, "login_status": status, "login_uncertain": False} if _is_definitive_logged_out(status): detail = _login_status_detail(status) _notify_step(on_step, "login_check", "blocked", detail) raise ImageStudioError(f"账号 {alias} 未登录：{detail}") detail = _login_status_detail(status) _notify_step(on_step, "login_check", "uncertain", f"账号 {alias} 登录状态暂不稳定，继续尝试读取商品：{detail}") return {**chrome_info, "login_status": status, "login_uncertain": True} def _ensure_assets_belong_to_project(database, project_id, asset_ids): ids = [int(asset_id) for asset_id in asset_ids if asset_id is not None] if not ids: return placeholders = ",".join("?" for _ in ids) rows = database.execute( f"SELECT id FROM image_studio_assets WHERE project_id = ? AND id IN ({placeholders})", [int(project_id), *ids], ).fetchall() found = {int(row["id"]) for row in rows} missing = sorted(set(ids) - found) if missing: raise db.DbError("AI工场资产不属于当前项目") def get_project(project_id, path=None, conn=None, include_deleted=False): clauses = ["id = ?"] params = [int(project_id)] if not include_deleted: clauses.append("deleted_at IS NULL") sql = "SELECT * FROM image_studio_projects WHERE " + " AND ".join(clauses) with _connection(conn, path) as database: return _fetch_one(database, sql, params, ImageStudioProject) def get_project_by_account_item(account_alias, item_id, path=None, conn=None, include_deleted=False): clauses = ["account_alias = ?", "item_id = ?"] params = [str(account_alias).strip(), str(item_id).strip()] if not include_deleted: clauses.append("deleted_at IS NULL") sql = "SELECT * FROM image_studio_projects WHERE " + " AND ".join(clauses) with _connection(conn, path) as database: return _fetch_one(database, sql, params, ImageStudioProject) def create_or_get_project( account=None, item_id=None, *, account_alias=None, account_name=None, account_slug=None, target_main_count=9, target_detail_count=12, draft_prompt="", path=None, conn=None, ): alias, name, slug = _account_fields( account, account_alias=account_alias, account_name=account_name, account_slug=account_slug, ) item = str(item_id or _get(account, "item_id") or "").strip() if not item: raise db.DbError("AI工场项目缺少商品ID") now = _now() with _connection(conn, path) as database: existing = get_project_by_account_item(alias, item, conn=database, include_deleted=True) if existing is not None: if existing.deleted_at is not None: with database: database.execute( """ UPDATE image_studio_projects SET deleted_at = NULL, deleted_reason = NULL, status = ?, updated_at = ? WHERE id = ? """, (PROJECT_STATUS_ACTIVE, now, existing.id), ) return get_project(existing.id, conn=database) return existing with database: cursor = database.execute( """ INSERT INTO image_studio_projects (account_alias, account_slug, account_name, item_id, storage_key, binding_state, target_main_count, target_detail_count, draft_prompt, status, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( alias, slug, name, item, item, PROJECT_BINDING_BOUND, int(target_main_count), int(target_detail_count), str(draft_prompt or ""), PROJECT_STATUS_ACTIVE, now, now, ), ) project_id = cursor.lastrowid return get_project(project_id, conn=database) def create_draft_project( account=None, *, account_alias=None, account_name=None, account_slug=None, target_main_count=9, target_detail_count=12, draft_prompt="", path=None, conn=None, ): """Create one project-bound local draft without exposing its internal item key.""" alias, name, slug = _account_fields( account, account_alias=account_alias, account_name=account_name, account_slug=account_slug, ) item = _new_draft_item_id() now = _now() with _connection(conn, path) as database: with database: cursor = database.execute( """ INSERT INTO image_studio_projects (account_alias, account_slug, account_name, item_id, storage_key, binding_state, target_main_count, target_detail_count, draft_prompt, status, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( alias, slug, name, item, item, PROJECT_BINDING_DRAFT, int(target_main_count), int(target_detail_count), str(draft_prompt or ""), PROJECT_STATUS_ACTIVE, now, now, ), ) project_id = cursor.lastrowid return get_project(project_id, conn=database) def bind_draft_project(project_id, item_id, path=None, conn=None): """Bind one active draft to a formal numeric item ID without moving its files.""" item = str(item_id or "").strip() if not is_formal_item_id(item): raise ImageStudioError("正式商品ID必须是数字") with _connection(conn, path) as database: source = get_project(project_id, conn=database, include_deleted=True) if source is None: raise ImageStudioError("临时草稿不存在") if source.deleted_at is not None: raise ImageStudioError("临时草稿已删除，无法绑定商品") if not is_draft_project(source): if source.item_id == item: return source raise ImageStudioError("当前项目不是临时草稿，不能重新绑定商品") existing = get_project_by_account_item( source.account_alias, item, conn=database, include_deleted=True, ) if existing is not None and existing.id != source.id: raise ImageStudioProjectConflictError("该商品项目已存在，不能覆盖或合并") now = _now() with database: database.execute( """ UPDATE image_studio_projects SET item_id = ?, binding_state = ?, updated_at = ? WHERE id = ? AND deleted_at IS NULL """, (item, PROJECT_BINDING_BOUND, now, int(source.id)), ) return get_project(source.id, conn=database) def list_projects(path=None, conn=None, include_deleted=False): sql = "SELECT * FROM image_studio_projects" if not include_deleted: sql += " WHERE deleted_at IS NULL" sql += " ORDER BY updated_at DESC, id DESC" with _connection(conn, path) as database: return _fetch_all(database, sql, (), ImageStudioProject) def list_recoverable_draft_projects(path=None, conn=None): """Return non-empty drafts that should reappear as product-suite task tabs.""" sql = """ SELECT p.* FROM image_studio_projects AS p WHERE p.deleted_at IS NULL AND p.binding_state = ? AND ( TRIM(COALESCE(p.draft_prompt, '')) <> '' OR EXISTS ( SELECT 1 FROM image_studio_assets AS a WHERE a.project_id = p.id ) OR EXISTS ( SELECT 1 FROM image_studio_jobs AS j WHERE j.project_id = p.id ) ) ORDER BY p.updated_at DESC, p.id DESC """ with _connection(conn, path) as database: return _fetch_all(database, sql, (PROJECT_BINDING_DRAFT,), ImageStudioProject) def project_has_content(project_id, path=None, conn=None) -> bool: with _connection(conn, path) as database: row = database.execute( """ SELECT EXISTS( SELECT 1 FROM image_studio_assets WHERE project_id = ? ) OR EXISTS( SELECT 1 FROM image_studio_jobs WHERE project_id = ? ) OR EXISTS( SELECT 1 FROM image_studio_projects WHERE id = ? AND TRIM(COALESCE(draft_prompt, '')) <> '' ) AS has_content """, (int(project_id), int(project_id), int(project_id)), ).fetchone() return bool(row["has_content"] if row is not None else False) def discard_empty_draft_project(project_id, reason="空临时草稿", path=None, conn=None): """Soft-delete a newly created, still empty draft and leave user files untouched.""" with _connection(conn, path) as database: project = get_project(project_id, conn=database, include_deleted=True) if project is None or project.deleted_at is not None or not is_draft_project(project): return project if project_has_content(project.id, conn=database): return project return soft_delete_project(project.id, reason=reason, conn=database) def update_project_prompt(project_id, draft_prompt, path=None, conn=None): with _connection(conn, path) as database: with database: database.execute( """ UPDATE image_studio_projects SET draft_prompt = ?, updated_at = ? WHERE id = ? AND deleted_at IS NULL """, (str(draft_prompt or ""), _now(), int(project_id)), ) return get_project(project_id, conn=database) def project_suite_settings(project) -> dict: raw = _get(project, "suite_settings_json", "{}") or "{}" try: value = json.loads(str(raw)) except (TypeError, ValueError): return {} return value if isinstance(value, dict) else {} def update_project_suite_settings(project_id, settings, path=None, conn=None): if not isinstance(settings, dict): raise db.DbError("商品套图设置必须是对象") payload = json.dumps(settings, ensure_ascii=False, sort_keys=True) with _connection(conn, path) as database: with database: database.execute( """ UPDATE image_studio_projects SET suite_settings_json = ?, updated_at = ? WHERE id = ? AND deleted_at IS NULL """, (payload, _now(), int(project_id)), ) return get_project(project_id, conn=database) def soft_delete_project(project_id, reason="", path=None, conn=None): now = _now() with _connection(conn, path) as database: with database: database.execute( """ UPDATE image_studio_projects SET deleted_at = ?, deleted_reason = ?, updated_at = ? WHERE id = ? AND deleted_at IS NULL """, (now, str(reason or ""), now, int(project_id)), ) return get_project(project_id, conn=database, include_deleted=True) def project_image_dirs(image_root, project): root = os.path.abspath( os.path.join( str(image_root or "images"), "pool", _safe_component(_get(project, "account_slug"), "account"), _safe_component(project_storage_key(project), "item"), ) ) return { "root": root, "originals": os.path.join(root, "originals"), "generated": os.path.join(root, "generated"), "exports": os.path.join(root, "exports"), } def default_project_image_dirs(project, config=None): return project_image_dirs(appconfig.image_dir(config), project) def add_asset( project_id, kind, *, remote_url=None, local_path=None, aspect_ratio=None, parent_asset_id=None, prompt=None, status=ASSET_STATUS_AVAILABLE, source_order=0, path=None, conn=None, ): now = _now() abs_local_path = os.path.abspath(local_path) if local_path else None with _connection(conn, path) as database: with database: cursor = database.execute( """ INSERT INTO image_studio_assets (project_id, kind, remote_url, local_path, aspect_ratio, parent_asset_id, prompt, status, source_order, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( int(project_id), str(kind), remote_url, abs_local_path, aspect_ratio, parent_asset_id, prompt, str(status or ASSET_STATUS_AVAILABLE), int(source_order or 0), now, now, ), ) asset_id = cursor.lastrowid return get_asset(asset_id, conn=database) def get_asset(asset_id, path=None, conn=None): with _connection(conn, path) as database: return _fetch_one( database, "SELECT * FROM image_studio_assets WHERE id = ?", (int(asset_id),), ImageStudioAsset, ) def list_assets(project_id, kind=None, include_missing=True, path=None, conn=None): clauses = ["project_id = ?"] params = [int(project_id)] if kind is not None: clauses.append("kind = ?") params.append(str(kind)) if not include_missing: clauses.append("status != ?") params.append(ASSET_STATUS_MISSING) sql = "SELECT * FROM image_studio_assets WHERE " + " AND ".join(clauses) sql += " ORDER BY source_order, id" with _connection(conn, path) as database: return _fetch_all(database, sql, params, ImageStudioAsset) def reorder_original_assets(project_id, asset_ids, path=None, conn=None): ordered_ids = [int(asset_id) for asset_id in asset_ids] if len(ordered_ids) != len(set(ordered_ids)): raise db.DbError("商品原图排序不能包含重复图片") with _connection(conn, path) as database: rows = database.execute( """ SELECT id FROM image_studio_assets WHERE project_id = ? AND kind = ? ORDER BY source_order, id """, (int(project_id), ASSET_KIND_ORIGINAL), ).fetchall() existing_ids = [int(row["id"]) for row in rows] if set(existing_ids) != set(ordered_ids): raise db.DbError("商品原图排序必须包含当前项目的全部原图") now = _now() with database: for source_order, asset_id in enumerate(ordered_ids, 1): database.execute( """ UPDATE image_studio_assets SET source_order = ?, updated_at = ? WHERE id = ? AND project_id = ? AND kind = ? """, ( source_order, now, asset_id, int(project_id), ASSET_KIND_ORIGINAL, ), ) return list_assets( project_id, kind=ASSET_KIND_ORIGINAL, path=path, conn=database, ) def asset_reference_counts(asset_id, path=None, conn=None): """Return selection/job references for one asset before pool removal.""" asset_id = int(asset_id) with _connection(conn, path) as database: selection_count = database.execute( "SELECT COUNT(*) FROM image_studio_selections WHERE asset_id = ?", (asset_id,), ).fetchone()[0] source_job_count = database.execute( "SELECT COUNT(*) FROM image_studio_jobs WHERE source_asset_id = ?", (asset_id,), ).fetchone()[0] output_job_count = database.execute( "SELECT COUNT(*) FROM image_studio_jobs WHERE output_asset_id = ?", (asset_id,), ).fetchone()[0] return { "selection": int(selection_count), "source_job": int(source_job_count), "output_job": int(output_job_count), "total": int(selection_count + source_job_count + output_job_count), } def remove_asset_if_unused(asset_id, path=None, conn=None): """Remove one pool asset row only when no job or selection references it. The local image file is intentionally kept on disk. AI工场 removal is a pool-level operation, not a destructive file cleanup. """ asset_id = int(asset_id) with _connection(conn, path) as database: asset = get_asset(asset_id, conn=database) if asset is None: raise db.DbError("照片不存在") counts = asset_reference_counts(asset_id, conn=database) if counts["total"]: raise db.DbError("照片正在被生成任务或终选引用，不能移除") with database: database.execute("DELETE FROM image_studio_assets WHERE id = ?", (asset_id,)) return asset def remove_original_assets_if_unused(project_id, asset_ids, path=None, conn=None): """Atomically remove unreferenced original assets from one project. Local files are intentionally retained. If any requested asset is invalid or referenced, no asset row is removed. """ project_id = int(project_id) ordered_ids = [] seen = set() for value in asset_ids or []: asset_id = int(value) if asset_id not in seen: seen.add(asset_id) ordered_ids.append(asset_id) if not ordered_ids: raise db.DbError("请选择要删除的商品原图") placeholders = ",".join("?" for _ in ordered_ids) with _connection(conn, path) as database: with database: project = get_project(project_id, conn=database) if project is None: raise db.DbError("商品套图项目不存在或已删除") rows = database.execute( f""" SELECT * FROM image_studio_assets WHERE project_id = ? AND kind = ? AND id IN ({placeholders}) """, [project_id, ASSET_KIND_ORIGINAL, *ordered_ids], ).fetchall() by_id = {int(row["id"]): row for row in rows} if set(by_id) != set(ordered_ids): raise db.DbError("选中的商品原图不存在或不属于当前项目") referenced_ids = [ asset_id for asset_id in ordered_ids if asset_reference_counts(asset_id, conn=database)["total"] ] if referenced_ids: raise db.DbError("选中的图片正在被生成任务或终选引用，不能移除") database.execute( f""" DELETE FROM image_studio_assets WHERE project_id = ? AND kind = ? AND id IN ({placeholders}) """, [project_id, ASSET_KIND_ORIGINAL, *ordered_ids], ) remaining_rows = database.execute( """ SELECT id FROM image_studio_assets WHERE project_id = ? AND kind = ? ORDER BY source_order, id """, (project_id, ASSET_KIND_ORIGINAL), ).fetchall() now = _now() for source_order, row in enumerate(remaining_rows, 1): database.execute( """ UPDATE image_studio_assets SET source_order = ?, updated_at = ? WHERE id = ? AND project_id = ? AND kind = ? """, ( source_order, now, int(row["id"]), project_id, ASSET_KIND_ORIGINAL, ), ) return [_row_to_dataclass(by_id[asset_id], ImageStudioAsset) for asset_id in ordered_ids] def restore_original_asset_snapshot(project_id, snapshots, path=None, conn=None): """Restore status/order for original assets that existed before a pull.""" project_id = int(project_id) normalized = [] seen = set() for value in snapshots or []: asset_id = int(_get(value, "id")) if asset_id in seen: continue status = str(_get(value, "status") or ASSET_STATUS_AVAILABLE) if status not in ASSET_STATUSES: raise db.DbError("商品原图快照状态无效") normalized.append( { "id": asset_id, "status": status, "source_order": max(0, int(_get(value, "source_order") or 0)), } ) seen.add(asset_id) if not normalized: return [] ids = [item["id"] for item in normalized] placeholders = ",".join("?" for _ in ids) with _connection(conn, path) as database: with database: rows = database.execute( f""" SELECT id FROM image_studio_assets WHERE project_id = ? AND kind = ? AND id IN ({placeholders}) """, [project_id, ASSET_KIND_ORIGINAL, *ids], ).fetchall() if {int(row["id"]) for row in rows} != set(ids): raise db.DbError("拉取前商品原图快照已失效") now = _now() for item in normalized: database.execute( """ UPDATE image_studio_assets SET status = ?, source_order = ?, updated_at = ? WHERE id = ? AND project_id = ? AND kind = ? """, ( item["status"], item["source_order"], now, item["id"], project_id, ASSET_KIND_ORIGINAL, ), ) remaining_rows = database.execute( f""" SELECT id FROM image_studio_assets WHERE project_id = ? AND kind = ? AND id NOT IN ({placeholders}) ORDER BY source_order, id """, [project_id, ASSET_KIND_ORIGINAL, *ids], ).fetchall() next_order = max( [int(item["source_order"]) for item in normalized] + [0] ) + 1 for row in remaining_rows: database.execute( """ UPDATE image_studio_assets SET source_order = ?, updated_at = ? WHERE id = ? AND project_id = ? AND kind = ? """, ( next_order, now, int(row["id"]), project_id, ASSET_KIND_ORIGINAL, ), ) next_order += 1 return [ get_asset(item["id"], conn=database) for item in normalized ] def sync_original_asset_urls(project_id, image_urls, path=None, conn=None, max_assets=16): """Store the read-only Shopee main image URL snapshot as remote-only assets.""" normalized = [] seen_urls = set() for position, item in enumerate(image_urls or [], start=1): if isinstance(item, dict): src = str(item.get("src") or "").strip() source_order = int(item.get("index") or position) else: src = str(item or "").strip() source_order = position if not src or src in seen_urls: continue seen_urls.add(src) normalized.append({"src": src, "source_order": source_order}) now = _now() with _connection(conn, path) as database: with database: existing_rows = database.execute( """ SELECT * FROM image_studio_assets WHERE project_id = ? AND kind = ? """, (int(project_id), ASSET_KIND_ORIGINAL), ).fetchall() by_url = { str(row["remote_url"]): row for row in existing_rows if row["remote_url"] } local_only_count = sum( 1 for row in existing_rows if not row["remote_url"] and row["status"] != ASSET_STATUS_MISSING ) remote_limit = max(0, int(max_assets) - local_only_count) normalized = normalized[:remote_limit] active_ids = set() for item in normalized: row = by_url.get(item["src"]) if row is not None: active_ids.add(int(row["id"])) database.execute( """ UPDATE image_studio_assets SET source_order = ?, status = ?, updated_at = ? WHERE id = ? """, ( int(item["source_order"]), ASSET_STATUS_AVAILABLE, now, int(row["id"]), ), ) else: cursor = database.execute( """ INSERT INTO image_studio_assets (project_id, kind, remote_url, local_path, status, source_order, created_at, updated_at) VALUES (?, ?, ?, NULL, ?, ?, ?, ?) """, ( int(project_id), ASSET_KIND_ORIGINAL, item["src"], ASSET_STATUS_AVAILABLE, int(item["source_order"]), now, now, ), ) active_ids.add(int(cursor.lastrowid)) missing_ids = [ int(row["id"]) for row in existing_rows if row["remote_url"] and int(row["id"]) not in active_ids ] if missing_ids: placeholders = ",".join("?" for _ in missing_ids) database.execute( f""" UPDATE image_studio_assets SET status = ?, updated_at = ? WHERE id IN ({placeholders}) """, [ASSET_STATUS_MISSING, now, *missing_ids], ) return list_assets(project_id, kind=ASSET_KIND_ORIGINAL, conn=database) def mark_asset_status(asset_id, status, path=None, conn=None): if str(status) not in ASSET_STATUSES: raise db.DbError("AI工场资产状态必须是 available 或 missing") with _connection(conn, path) as database: with database: database.execute( """ UPDATE image_studio_assets SET status = ?, updated_at = ? WHERE id = ? """, (str(status), _now(), int(asset_id)), ) return get_asset(asset_id, conn=database) def update_asset_local_path(asset_id, local_path, status=ASSET_STATUS_AVAILABLE, path=None, conn=None): if str(status) not in ASSET_STATUSES: raise db.DbError("AI工场资产状态必须是 available 或 missing") abs_local_path = os.path.abspath(local_path) if local_path else None with _connection(conn, path) as database: with database: database.execute( """ UPDATE image_studio_assets SET local_path = ?, status = ?, updated_at = ? WHERE id = ? """, (abs_local_path, str(status), _now(), int(asset_id)), ) return get_asset(asset_id, conn=database) def create_job( project_id, *, source_asset_id=None, job_type="main", prompt="", 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: _ensure_assets_belong_to_project(database, project_id, [source_asset_id]) cursor = database.execute( """ INSERT INTO image_studio_jobs (project_id, source_asset_id, generation_source, provider, job_type, task_key, status, prompt, generation_round_key, generation_slot_index, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, 'pending', ?, ?, ?, ?, ?) """, ( int(project_id), source_asset_id, str(generation_source or "cmhub"), str(provider or "cmhub"), str(job_type), task_key, str(prompt or ""), generation_round_key, generation_slot_index, now, now, ), ) job_id = cursor.lastrowid except Exception as exc: raise db.DbError(f"创建AI工场生图任务失败: {exc}") from exc return get_job(job_id, conn=database) def get_job(job_id, path=None, conn=None): with _connection(conn, path) as database: return _fetch_one( database, "SELECT * FROM image_studio_jobs WHERE id = ?", (int(job_id),), ImageStudioJob, ) def list_jobs(project_id, *, statuses=None, path=None, conn=None): clauses = ["project_id = ?"] params = [int(project_id)] requested_statuses = [str(status) for status in (statuses or []) if str(status)] if requested_statuses: invalid = set(requested_statuses) - JOB_STATUSES if invalid: raise db.DbError("图片生成任务状态无效") placeholders = ",".join("?" for _ in requested_statuses) clauses.append(f"status IN ({placeholders})") params.extend(requested_statuses) sql = "SELECT * FROM image_studio_jobs WHERE " + " AND ".join(clauses) sql += " ORDER BY updated_at DESC, id DESC" with _connection(conn, path) as database: 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, MIN(created_at) DESC, 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: with database: database.execute( """ UPDATE image_studio_jobs SET task_id = ?, status = 'submitted', recovery_action = ?, call_id = ?, points_cost = ?, points_balance = ?, submitted_at = ?, updated_at = ? WHERE id = ? """, ( str(task_id), JOB_RECOVERY_RESUME, call_id, points_cost, points_balance, now, now, int(job_id), ), ) return get_job(job_id, conn=database) def update_job_status( job_id, status, *, error=None, output_asset_id=None, points_balance=None, recovery_action=None, increment_attempts=False, path=None, conn=None, ): if str(status) not in JOB_STATUSES: raise db.DbError("AI工场任务状态无效") if recovery_action is not None and str(recovery_action) not in JOB_RECOVERY_ACTIONS: raise db.DbError("AI工场任务恢复方式无效") if recovery_action is None and str(status) == "succeeded": recovery_action = JOB_RECOVERY_NONE now = _now() terminal = str(status) in {"succeeded", "failed", "expired", "cancelled"} with _connection(conn, path) as database: with database: if output_asset_id is not None: job = get_job(job_id, conn=database) if job is None: raise db.DbError("AI工场任务不存在") _ensure_assets_belong_to_project(database, job.project_id, [output_asset_id]) database.execute( """ UPDATE image_studio_jobs SET status = ?, error = ?, recovery_action = COALESCE(?, recovery_action), output_asset_id = COALESCE(?, output_asset_id), points_balance = COALESCE(?, points_balance), attempts = attempts + ?, finished_at = CASE WHEN ? THEN ? ELSE finished_at END, updated_at = ? WHERE id = ? """, ( str(status), error, recovery_action, output_asset_id, points_balance, 1 if increment_attempts else 0, 1 if terminal else 0, now, now, int(job_id), ), ) return get_job(job_id, conn=database) def list_resumable_jobs(path=None, conn=None, project_id=None, include_failed_downloads=False): if include_failed_downloads: clauses = [ "task_id IS NOT NULL", "recovery_action = ?", "status IN (?, ?, ?, ?)", ] params = [JOB_RECOVERY_RESUME, "submitted", "running", "failed", "cancelled"] else: clauses = ["status IN (?, ?)", "task_id IS NOT NULL", "recovery_action = ?"] params = ["submitted", "running", JOB_RECOVERY_RESUME] if project_id is not None: clauses.append("project_id = ?") params.append(int(project_id)) sql = "SELECT * FROM image_studio_jobs WHERE " + " AND ".join(clauses) sql += " ORDER BY updated_at, id" with _connection(conn, path) as database: return _fetch_all(database, sql, params, ImageStudioJob) def replace_selections(project_id, selection_type, asset_ids: Iterable[int], path=None, conn=None): selection = str(selection_type) if selection not in SELECTION_TYPES: raise db.DbError("终选类型必须是 main 或 detail") ids = [int(asset_id) for asset_id in asset_ids] if len(ids) != len(set(ids)): raise db.DbError("同一资产不能重复加入同一类终选") now = _now() with _connection(conn, path) as database: with database: _ensure_assets_belong_to_project(database, project_id, ids) database.execute( """ DELETE FROM image_studio_selections WHERE project_id = ? AND selection_type = ? """, (int(project_id), selection), ) database.executemany( """ INSERT INTO image_studio_selections (project_id, selection_type, position, asset_id, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?) """, [ (int(project_id), selection, index + 1, asset_id, now, now) for index, asset_id in enumerate(ids) ], ) return list_selections(project_id, selection, conn=database) def list_selections(project_id, selection_type=None, path=None, conn=None): clauses = ["project_id = ?"] params = [int(project_id)] if selection_type is not None: clauses.append("selection_type = ?") params.append(str(selection_type)) sql = "SELECT * FROM image_studio_selections WHERE " + " AND ".join(clauses) sql += " ORDER BY selection_type, position" with _connection(conn, path) as database: return _fetch_all(database, sql, params, ImageStudioSelection) def pull_remote_main_image_urls( account_or_alias, item_id, *, path=None, config=None, login_timeout=8, on_step=None, should_stop=None, ): """Create/open an AI studio project and read Shopee main image URLs without editing.""" cfg = appconfig.load_config() if config is None else config should_stop = should_stop or (lambda: False) if should_stop(): raise ImageStudioPullCancelled() database_path = _db_path(path, cfg) db.init_db(database_path) item = _normalize_item_id(item_id) try: account = accounts.resolve_account(account_or_alias, path=database_path, config=cfg) except Exception as exc: raise ImageStudioError(f"AI工场账号不可用：{exc}") from exc project = create_or_get_project(account, item_id=item, path=database_path) if should_stop(): raise ImageStudioPullCancelled(project=project) readiness = _ensure_account_ready_for_read( account, path=database_path, config=cfg, login_timeout=login_timeout, on_step=on_step, ) if should_stop(): raise ImageStudioPullCancelled(project=project) cdp = None try: if should_stop(): raise ImageStudioPullCancelled(project=project) _notify_step(on_step, "open_product", "start", f"商品 {item}") cdp = editor.open_product(account, item, on_step=on_step, bring_to_front=False) if should_stop(): raise ImageStudioPullCancelled(project=project) _notify_step(on_step, "open_product", "success", f"商品 {item}") _notify_step(on_step, "read_main_images", "start", f"商品 {item}") images = editor.read_product_image_urls(cdp) if should_stop(): raise ImageStudioPullCancelled(project=project) if not images: raise ImageStudioError("未读取到蝦皮商品主图 URL") assets = sync_original_asset_urls(project.id, images, path=database_path) if should_stop(): raise ImageStudioPullCancelled(project=project, assets=assets) _notify_step(on_step, "read_main_images", "success", f"读取 {len(images)} 张主图 URL") project = get_project(project.id, path=database_path) return { "project": project, "images": images, "assets": assets, "account": account, "readiness": readiness, } except ImageStudioError: raise except editor.EditorError as exc: raise ImageStudioError(str(exc)) from exc except Exception as exc: raise ImageStudioError(f"读取蝦皮原主图失败：{exc}") from exc finally: if cdp is not None: editor.close_readonly_product(cdp)