"""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 appconfig, db from .config import make_slug PROJECT_STATUS_ACTIVE = "active" 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 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 _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 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 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)