"""AI generation helpers backed by configurable HTTP model endpoints.""" import base64 from concurrent.futures import CancelledError, ThreadPoolExecutor, as_completed import copy import json import mimetypes import os import threading import time import urllib.error import urllib.request import uuid from . import appconfig, db, diagnostics, image_paths from . import prompts as prompt_module class AIError(RuntimeError): """Raised when AI generation cannot complete.""" _RESOLUTION_SIZES = { "512": (512, 512), "1k": (1024, 1024), "2k": (2048, 2048), "4k": (4096, 4096), } def gen_title( title_prompt, old_title, retry=None, config=None, models_path=appconfig.AI_MODELS_PATH, on_step=None, ): """Generate a new product title from a prompt and the old title.""" cfg = appconfig.load_config() if config is None else config ai_cfg = appconfig.ai_config(cfg) _notify_step(on_step, "load_text_model") model = _role_model("text", ai_cfg.get("default_text_model"), models_path) _notify_step(on_step, "title_build_request") payload = _chat_payload( model, [ {"role": "system", "content": str(title_prompt or "").strip()}, { "role": "user", "content": "旧标题：\n%s\n\n请只返回新标题，不要解释。" % str(old_title or ""), }, ], ) attempts = _attempt_count(ai_cfg, retry) _notify_step(on_step, "title_request") data = _call_with_retry( model, payload, cfg, attempts, request_kind="json", on_retry=lambda attempt, total_attempts, exc: _notify_retry( on_step, "title_request", attempt, total_attempts, exc, model, ), ) _notify_step(on_step, "title_parse_response") text = _extract_text(data).strip() if not text: raise AIError("AI 返回为空标题") return text def gen_cover( cover_prompt, old_cover_path, out_path, resolution=None, jpg_quality=None, retry=None, config=None, models_path=appconfig.AI_MODELS_PATH, on_step=None, ): """Generate a new cover image and save it as a JPEG file.""" _notify_step(on_step, "cover_validate_input") old_cover_path = os.path.abspath(str(old_cover_path)) if not os.path.exists(old_cover_path): raise FileNotFoundError("旧封面图片不存在: %s" % old_cover_path) if not out_path: raise AIError("缺少新封面输出路径") cfg = appconfig.load_config() if config is None else config ai_cfg = appconfig.ai_config(cfg) _notify_step(on_step, "load_image_model") model = _role_model("image", ai_cfg.get("default_image_model"), models_path) resolution = str(resolution or ai_cfg.get("resolution", "1k")) quality = _jpg_quality(jpg_quality if jpg_quality is not None else ai_cfg.get("jpg_quality", 90)) attempts = _attempt_count(ai_cfg, retry) api_type = model.get("api_type", "auto") _notify_step(on_step, "cover_build_request") if api_type == "images_edits": body, content_type = _image_edit_body(model, cover_prompt, old_cover_path, resolution) _notify_step(on_step, "cover_request") data = _call_with_retry( model, body, cfg, attempts, request_kind="multipart", content_type=content_type, on_retry=lambda attempt, total_attempts, exc: _notify_retry( on_step, "cover_request", attempt, total_attempts, exc, model, ), ) else: payload = _image_chat_payload(model, cover_prompt, old_cover_path, resolution) _notify_step(on_step, "cover_request") data = _call_with_retry( model, payload, cfg, attempts, request_kind="json", on_retry=lambda attempt, total_attempts, exc: _notify_retry( on_step, "cover_request", attempt, total_attempts, exc, model, ), ) _notify_step(on_step, "cover_parse_response") image_bytes = _extract_image_bytes(data, model, cfg) _notify_step(on_step, "cover_save") return _save_jpeg(image_bytes, out_path, resolution, quality) def _notify_step(callback, step): if callback is None: return try: callback(step) except Exception: pass def _notify_retry(callback, step, attempt, attempts, exc, model): if callback is None: return try: callback( { "step": step, "result": "retry", "attempt": attempt, "attempts": attempts, "detail": _redact(str(exc), model), } ) except Exception: pass def generate_batch(tasks, prompts, ai_cfg=None, on_progress=None, should_stop=None): """Generate titles first, then covers, and persist each successful task.""" runtime = dict(ai_cfg or {}) config = _runtime_config(runtime) generation_cfg = appconfig.ai_config(config) generation_cfg.update( { key: value for key, value in runtime.items() if key in { "title_concurrency", "image_concurrency", "retry", "jpg_quality", "resolution", "generate_cover", } } ) db_path = runtime.get("db_path") models_path = runtime.get("models_path", appconfig.AI_MODELS_PATH) image_root = runtime.get("image_dir") or appconfig.image_dir(config) account_by_alias = runtime.get("account_by_alias") or {} on_task_update = runtime.get("on_task_update") on_event = runtime.get("on_event") on_error = runtime.get("on_error") title_prompt = _prompt_value(prompts, "title") cover_prompt = _prompt_value(prompts, "cover") generate_cover = bool(generation_cfg.get("generate_cover", False)) should_stop = should_stop or (lambda: False) eligible = [ task for task in list(tasks) if getattr(task, "stage", None) == "collected" ] summary = { "ok": True, "total": len(eligible), "title_done": 0, "cover_done": 0, "cover_total": len(eligible) if generate_cover else 0, "generated_done": 0, "failed": 0, "cancelled": False, "generate_cover": generate_cover, } _emit_generation_progress(on_progress, summary) title_results = {} step_by_task = {} step_lock = threading.Lock() def set_step(task, step): with step_lock: step_by_task[getattr(task, "id", None)] = str(step) def get_step(task, fallback): with step_lock: return step_by_task.get(getattr(task, "id", None), fallback) def step_callback(task, phase): def callback(event): if isinstance(event, dict): step = event.get("step") or "unknown" result = event.get("result") or "start" set_step(task, step) _emit_generation_event( on_event, task, phase, step, result, detail=event.get("detail"), level=event.get("level") or ("warning" if result == "retry" else "info"), attempt=event.get("attempt"), attempts=event.get("attempts"), ) return set_step(task, event) _emit_generation_event(on_event, task, phase, event, "start") return callback with ThreadPoolExecutor( max_workers=max(1, int(generation_cfg.get("title_concurrency", 1))) ) as executor: futures = {} for task in eligible: if should_stop(): summary["cancelled"] = True break set_step(task, "title_submit") _emit_generation_event(on_event, task, "title", "title_submit", "start") futures[ executor.submit( gen_title, title_prompt, getattr(task, "old_title", "") or "", retry=generation_cfg.get("retry"), config=config, models_path=models_path, on_step=step_callback(task, "title"), ) ] = task for future in as_completed(futures): task = futures[future] if should_stop(): summary["cancelled"] = True _cancel_pending(futures) try: new_title = future.result() title_results[task.id] = new_title summary["title_done"] += 1 set_step(task, "title_done") _emit_generation_event(on_event, task, "title", "title_done", "success") if not generate_cover: set_step(task, "db_write") _emit_generation_event(on_event, task, "title", "db_write", "start") db.set_generated(task.id, new_title, None, path=db_path) summary["generated_done"] += 1 if on_task_update is not None: on_task_update( task.id, { "stage": "generated", "status": "success", "new_title": new_title, "new_cover_path": None, }, ) _emit_generation_event( on_event, task, "title", "db_write", "success", detail="仅生成标题", ) except CancelledError: summary["cancelled"] = True _emit_generation_event(on_event, task, "title", get_step(task, "title_request"), "cancelled", level="warning") except Exception as exc: summary["failed"] += 1 summary["ok"] = False step = get_step(task, "title_request") error = _mark_generate_failed(task, exc, db_path, on_task_update) _emit_generation_event(on_event, task, "title", step, "failed", detail=error, level="error") _emit_generation_error(on_error, task, "title", step, exc, error) _emit_generation_progress(on_progress, summary) if not generate_cover: if summary["cancelled"]: summary["ok"] = False return summary cover_tasks = [ task for task in eligible if task.id in title_results ] with ThreadPoolExecutor( max_workers=max(1, int(generation_cfg.get("image_concurrency", 1))) ) as executor: futures = {} for task in cover_tasks: if should_stop(): summary["cancelled"] = True break new_title = title_results[task.id] try: set_step(task, "cover_prompt_render") _emit_generation_event(on_event, task, "cover", "cover_prompt_render", "start") rendered_cover_prompt = prompt_module.render_prompt( cover_prompt, _prompt_context(task, new_title, account_by_alias), ) _emit_generation_event(on_event, task, "cover", "cover_prompt_render", "success") set_step(task, "cover_submit") _emit_generation_event(on_event, task, "cover", "cover_submit", "start") futures[ executor.submit( gen_cover, rendered_cover_prompt, getattr(task, "old_cover_path", "") or "", _new_cover_path(task, account_by_alias, image_root), resolution=generation_cfg.get("resolution"), jpg_quality=generation_cfg.get("jpg_quality"), retry=generation_cfg.get("retry"), config=config, models_path=models_path, on_step=step_callback(task, "cover"), ) ] = (task, new_title) except Exception as exc: summary["failed"] += 1 summary["ok"] = False step = get_step(task, "cover_prompt_render") error = _mark_generate_failed(task, exc, db_path, on_task_update) _emit_generation_event(on_event, task, "cover", step, "failed", detail=error, level="error") _emit_generation_error(on_error, task, "cover", step, exc, error) _emit_generation_progress(on_progress, summary) for future in as_completed(futures): task, new_title = futures[future] if should_stop(): summary["cancelled"] = True _cancel_pending(futures) try: new_cover_path = future.result() set_step(task, "db_write") _emit_generation_event(on_event, task, "cover", "db_write", "start") db.set_generated(task.id, new_title, new_cover_path, path=db_path) summary["cover_done"] += 1 summary["generated_done"] += 1 if on_task_update is not None: on_task_update( task.id, { "stage": "generated", "status": "success", "new_title": new_title, "new_cover_path": new_cover_path, }, ) _emit_generation_event(on_event, task, "cover", "db_write", "success", detail=new_cover_path) except CancelledError: summary["cancelled"] = True _emit_generation_event(on_event, task, "cover", get_step(task, "cover_request"), "cancelled", level="warning") except Exception as exc: summary["failed"] += 1 summary["ok"] = False step = get_step(task, "cover_request") error = _mark_generate_failed(task, exc, db_path, on_task_update) _emit_generation_event(on_event, task, "cover", step, "failed", detail=error, level="error") _emit_generation_error(on_error, task, "cover", step, exc, error) _emit_generation_progress(on_progress, summary) if summary["cancelled"]: summary["ok"] = False return summary def _role_model(category, name, models_path): if not name: raise AIError("未配置默认 %s 模型" % category) model = appconfig.get_model(name, path=models_path) if model.get("category") != category: raise AIError("模型 %s 不是 %s 类别" % (name, category)) if not model.get("enabled", True): raise AIError("模型已禁用: %s" % name) missing = [ field for field in ("url", "model", "api_key") if not str(model.get(field, "")).strip() ] if missing: raise AIError("模型 %s 缺少字段: %s" % (name, ", ".join(missing))) return model def _runtime_config(runtime): if runtime.get("config") is not None: return runtime["config"] config = appconfig.load_config() ai_updates = { key: value for key, value in runtime.items() if key in { "default_text_model", "default_image_model", "title_concurrency", "image_concurrency", "retry", "jpg_quality", "resolution", "generate_cover", "resolution_timeouts", } } if ai_updates: config = copy.deepcopy(config) config.setdefault("ai", {}).update(ai_updates) return config def _prompt_value(prompt_values, name): if isinstance(prompt_values, dict): return str( prompt_values.get(name) or prompt_values.get(f"{name}_prompt") or "" ) return "" def _prompt_context(task, new_title, account_by_alias): return { "old_title": getattr(task, "old_title", ""), "new_title": new_title, "item_id": getattr(task, "item_id", ""), "account_name": _account_name(task, account_by_alias), "alias": getattr(task, "alias", ""), } def _account_name(task, account_by_alias): alias = str(getattr(task, "alias", "") or "").strip() account = account_by_alias.get(alias) if account is not None: return getattr(account, "account_name", "") or alias return getattr(task, "account_name", "") or alias def _new_cover_path(task, account_by_alias, image_root): alias = str(getattr(task, "alias", "") or "").strip() account = account_by_alias.get(alias) return image_paths.task_image_path(image_root, task, account, "new") def _mark_generate_failed(task, exc, db_path, on_task_update): error = diagnostics.redact_log_text(str(exc) or exc.__class__.__name__) db.mark_failed(task.id, "generate", error, path=db_path) if on_task_update is not None: on_task_update(task.id, {"status": "failed", "last_error": error}) return error def _emit_generation_event( callback, task, phase, step, result, detail=None, level="info", attempt=None, attempts=None, ): if callback is None: return payload = { "task": task, "phase": phase, "step": str(step), "result": result, "level": level, } if detail is not None: payload["detail"] = diagnostics.redact_log_text(detail) if attempt is not None: payload["attempt"] = attempt if attempts is not None: payload["attempts"] = attempts try: callback(payload) except Exception: return def _emit_generation_error(callback, task, phase, step, exc, error): if callback is None: return payload = { "task": task, "phase": phase, "step": str(step), "error": diagnostics.redact_log_text(error), "exception": exc, } try: callback(payload) except Exception: return def _cancel_pending(futures): for future in futures: if not future.done(): future.cancel() def _emit_generation_progress(on_progress, summary): if on_progress is None: return payload = dict(summary) try: on_progress(payload) except TypeError: on_progress( payload.get("title_done", 0), payload.get("cover_done", 0), payload.get("failed", 0), ) def _attempt_count(ai_cfg, retry): retry_count = ai_cfg.get("retry", 2) if retry is None else retry return max(1, int(retry_count) + 1) def _read_timeout(model, config): return int(model.get("timeout_seconds") or appconfig.response_timeout(config)) def _connect_timeout(model): return int(model.get("connect_timeout_seconds") or 30) def _headers(model, content_type): return { "Authorization": "Bearer " + model["api_key"], "Content-Type": content_type, } def _call_with_retry( model, body, config, attempts, request_kind, content_type=None, on_retry=None, ): last_exc = None for index in range(attempts): try: return _call_once(model, body, config, request_kind, content_type=content_type) except Exception as exc: last_exc = exc if index + 1 >= attempts: break if on_retry is not None: try: on_retry(index + 1, attempts, exc) except Exception: pass time.sleep(min(2.0, 0.4 * (index + 1))) raise AIError( "AI 调用失败（已尝试 %s 次）: %s" % (attempts, _redact(str(last_exc), model)) ) from last_exc def _call_once(model, body, config, request_kind, content_type=None): if request_kind == "json": data = json.dumps(body, ensure_ascii=False).encode("utf-8") content_type = "application/json" else: data = body request = urllib.request.Request( appconfig.model_request_url(model), data=data, headers=_headers(model, content_type), method="POST", ) timeout = max(_connect_timeout(model), _read_timeout(model, config)) try: with urllib.request.urlopen(request, timeout=timeout) as response: raw = response.read() except urllib.error.HTTPError as exc: detail = _safe_http_error(exc, model) raise AIError("HTTP %s: %s" % (exc.code, detail)) from exc except urllib.error.URLError as exc: raise AIError(_redact(str(exc.reason), model)) from exc try: return json.loads(raw.decode("utf-8")) except (UnicodeDecodeError, json.JSONDecodeError) as exc: raise AIError("AI 返回不是有效 JSON") from exc def _safe_http_error(exc, model): try: body = exc.read(2048).decode("utf-8", errors="replace") except Exception: body = "" return _redact(body or str(exc), model) def _chat_payload(model, messages): payload = { "model": model["model"], "messages": messages, } payload.update(copy.deepcopy(model.get("extra_body", {}))) return payload def _image_chat_payload(model, cover_prompt, old_cover_path, resolution): prompt = "%s\n\n目标分辨率：%s。" % (str(cover_prompt or "").strip(), resolution) content = [ {"type": "text", "text": prompt.strip()}, { "type": "image_url", "image_url": {"url": _image_data_url(old_cover_path)}, }, ] return _chat_payload(model, [{"role": "user", "content": content}]) def _image_edit_body(model, cover_prompt, old_cover_path, resolution): fields = { "model": model["model"], "prompt": str(cover_prompt or ""), "size": _resolution_size_text(resolution), } fields.update(copy.deepcopy(model.get("extra_body", {}))) files = { "image": ( os.path.basename(old_cover_path), open(old_cover_path, "rb").read(), mimetypes.guess_type(old_cover_path)[0] or "application/octet-stream", ) } return _multipart_body(fields, files) def _multipart_body(fields, files): boundary = "----cmshopee-%s" % uuid.uuid4().hex chunks = [] for name, value in fields.items(): chunks.extend( [ ("--%s\r\n" % boundary).encode("utf-8"), ('Content-Disposition: form-data; name="%s"\r\n\r\n' % name).encode("utf-8"), str(value).encode("utf-8"), b"\r\n", ] ) for name, file_info in files.items(): filename, data, content_type = file_info chunks.extend( [ ("--%s\r\n" % boundary).encode("utf-8"), ( 'Content-Disposition: form-data; name="%s"; filename="%s"\r\n' % (name, filename) ).encode("utf-8"), ("Content-Type: %s\r\n\r\n" % content_type).encode("utf-8"), data, b"\r\n", ] ) chunks.append(("--%s--\r\n" % boundary).encode("utf-8")) return b"".join(chunks), "multipart/form-data; boundary=%s" % boundary def _image_data_url(path): mime = mimetypes.guess_type(path)[0] or "image/jpeg" with open(path, "rb") as fh: encoded = base64.b64encode(fh.read()).decode("ascii") return "data:%s;base64,%s" % (mime, encoded) def _extract_text(data): if isinstance(data, dict): for key in ("output_text", "text", "content"): value = data.get(key) if isinstance(value, str): return value choices = data.get("choices") if isinstance(choices, list) and choices: first = choices[0] if isinstance(first, dict): if isinstance(first.get("text"), str): return first["text"] message = first.get("message") or {} content = message.get("content") return _content_text(content) return "" def _content_text(content): if isinstance(content, str): return content if isinstance(content, list): parts = [] for item in content: if not isinstance(item, dict): continue value = item.get("text") or item.get("content") if isinstance(value, str): parts.append(value) return "".join(parts) return "" def _extract_image_bytes(data, model, config): image_ref = _find_image_ref(data) if not image_ref: raise AIError("AI 返回中没有图片数据") if image_ref.startswith("data:"): return _decode_data_url(image_ref) if _looks_base64(image_ref): return base64.b64decode(image_ref) return _download_image(image_ref, model, config) def _find_image_ref(value): if isinstance(value, dict): for key in ("b64_json", "base64", "image_base64", "image", "url"): candidate = value.get(key) if isinstance(candidate, str) and candidate.strip(): return candidate.strip() image_url = value.get("image_url") if isinstance(image_url, str): return image_url if isinstance(image_url, dict): candidate = image_url.get("url") if isinstance(candidate, str): return candidate for key in ("data", "choices", "output", "content", "images"): candidate = _find_image_ref(value.get(key)) if candidate: return candidate message = value.get("message") if message is not None: candidate = _find_image_ref(message) if candidate: return candidate elif isinstance(value, list): for item in value: candidate = _find_image_ref(item) if candidate: return candidate return None def _decode_data_url(value): if "," not in value: raise AIError("图片 data URL 格式错误") return base64.b64decode(value.split(",", 1)[1]) def _looks_base64(value): compact = value.strip() if compact.startswith(("http://", "https://")): return False if len(compact) < 32: return False allowed = set("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/=\n\r") return all(char in allowed for char in compact) def _download_image(url, model, config): request = urllib.request.Request(url, method="GET") try: with urllib.request.urlopen(request, timeout=_read_timeout(model, config)) as response: return response.read() except Exception as exc: raise AIError("下载 AI 图片失败: %s" % _redact(str(exc), model)) from exc def _save_jpeg(image_bytes, out_path, resolution, jpg_quality): try: from PIL import Image except ImportError as exc: raise AIError("缺少 Pillow，无法保存 AI 图片") from exc import io out_path = os.path.abspath(str(out_path)) directory = os.path.dirname(out_path) if directory: os.makedirs(directory, exist_ok=True) size = _resolution_size(resolution) try: with Image.open(io.BytesIO(image_bytes)) as image: image = image.convert("RGB") if size: image = image.resize(size, Image.LANCZOS) image.save(out_path, "JPEG", quality=jpg_quality, optimize=True) except Exception as exc: raise AIError("AI 图片保存失败: %s" % exc) from exc return out_path def _resolution_size(resolution): return _RESOLUTION_SIZES.get(str(resolution)) def _resolution_size_text(resolution): size = _resolution_size(resolution) if not size: return str(resolution) return "%sx%s" % size def _jpg_quality(value): value = int(value) return min(100, max(1, value)) def _redact(text, model): result = str(text) for secret in (model.get("api_key"),): if secret: result = result.replace(str(secret), "***") return result