853 lines
29 KiB
Python
853 lines
29 KiB
Python
"""Data service for the AI image studio."""
|
||
|
||
from __future__ import annotations
|
||
|
||
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"
|
||
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"}
|
||
SELECTION_TYPES = {"main", "detail"}
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class ImageStudioProject:
|
||
id: int
|
||
account_alias: str
|
||
account_slug: str
|
||
account_name: Optional[str]
|
||
item_id: str
|
||
target_main_count: int
|
||
target_detail_count: int
|
||
draft_prompt: 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]
|
||
attempts: int
|
||
points_cost: Optional[int]
|
||
points_balance: Optional[int]
|
||
call_id: Optional[str]
|
||
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
|
||
|
||
|
||
class ImageStudioError(RuntimeError):
|
||
"""Raised when the AI image studio service cannot complete an operation."""
|
||
|
||
|
||
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 _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,
|
||
target_main_count, target_detail_count, draft_prompt,
|
||
status, created_at, updated_at)
|
||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||
""",
|
||
(
|
||
alias,
|
||
slug,
|
||
name,
|
||
item,
|
||
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 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 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 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(_get(project, "item_id"), "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 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 sync_original_asset_urls(project_id, image_urls, path=None, conn=None):
|
||
"""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"]
|
||
}
|
||
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 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",
|
||
path=None,
|
||
conn=None,
|
||
):
|
||
now = _now()
|
||
task_key = str(task_key or _task_key(project_id))
|
||
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, 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 ""),
|
||
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 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',
|
||
call_id = ?,
|
||
points_cost = ?,
|
||
points_balance = ?,
|
||
submitted_at = ?,
|
||
updated_at = ?
|
||
WHERE id = ?
|
||
""",
|
||
(str(task_id), 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,
|
||
increment_attempts=False,
|
||
path=None,
|
||
conn=None,
|
||
):
|
||
if str(status) not in JOB_STATUSES:
|
||
raise db.DbError("AI工场任务状态无效")
|
||
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 = ?,
|
||
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,
|
||
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):
|
||
clauses = ["status IN (?, ?)", "task_id IS NOT NULL"]
|
||
params = ["submitted", "running"]
|
||
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,
|
||
):
|
||
"""Create/open an AI studio project and read Shopee main image URLs without editing."""
|
||
|
||
cfg = appconfig.load_config() if config is None else config
|
||
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)
|
||
readiness = _ensure_account_ready_for_read(
|
||
account,
|
||
path=database_path,
|
||
config=cfg,
|
||
login_timeout=login_timeout,
|
||
on_step=on_step,
|
||
)
|
||
cdp = None
|
||
try:
|
||
_notify_step(on_step, "open_product", "start", f"商品 {item}")
|
||
cdp = editor.open_product(account, item, on_step=on_step, bring_to_front=False)
|
||
_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 not images:
|
||
raise ImageStudioError("未读取到蝦皮商品主图 URL")
|
||
assets = sync_original_asset_urls(project.id, images, path=database_path)
|
||
_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)
|