"""AI generation helpers backed by configurable HTTP model endpoints.""" import base64 import ipaddress from concurrent.futures import CancelledError, ThreadPoolExecutor, as_completed import copy import json import mimetypes import os import socket import threading import time import urllib.error import urllib.parse import urllib.request import uuid import requests from . import appconfig, db, diagnostics, image_paths from . import prompts as prompt_module class AIError(RuntimeError): """Raised when AI generation cannot complete.""" class CMHubError(AIError): """Structured cmhub gateway error.""" def __init__(self, code, message, status=None, retryable=False, retry_after=None): self.code = str(code or "unknown") self.status = status self.retryable = bool(retryable) self.retry_after = retry_after super().__init__("cmhub %s: %s" % (self.code, message)) CMHUB_IMAGE_MAX_BYTES = 20 * 1024 * 1024 _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, on_event=None, cmhub_config_path=appconfig.CMHUB_CONFIG_PATH, ): """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) if _ai_backend(ai_cfg) == "cmhub": return _gen_title_cmhub( title_prompt, old_title, retry=retry, config=cfg, cmhub_config_path=cmhub_config_path, on_step=on_step, on_event=on_event, ) _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, on_event=None, cmhub_config_path=appconfig.CMHUB_CONFIG_PATH, ): """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) if _ai_backend(ai_cfg) == "cmhub": return _gen_cover_cmhub( cover_prompt, old_cover_path, out_path, resolution=resolution, jpg_quality=jpg_quality, retry=retry, config=cfg, cmhub_config_path=cmhub_config_path, on_step=on_step, on_event=on_event, ) _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 _task_attempt_count(task, field): try: return int(getattr(task, field, 0) or 0) except (TypeError, ValueError): return 0 def is_generatable_task(task): """判断任务是否能由② AI生成执行或重试。""" stage = str(getattr(task, "stage", "") or "") status = str(getattr(task, "status", "") or "") if status in {"running", "skipped"}: return False if stage == "collected": return True if stage == "generated" and status == "failed": return _task_attempt_count(task, "apply_attempts") == 0 return False 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) cmhub_config_path = runtime.get("cmhub_config_path", appconfig.CMHUB_CONFIG_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 is_generatable_task(task) ] 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"), metadata=event.get("metadata"), ) 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"), on_event=step_callback(task, "title"), cmhub_config_path=cmhub_config_path, ) ] = 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"), on_event=step_callback(task, "cover"), cmhub_config_path=cmhub_config_path, ) ] = (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 fetch_cmhub_models(base_url, api_key, connect_timeout=10, read_timeout=30): """Fetch cmhub model aliases for settings UI.""" base_url = str(base_url or "").strip() api_key = str(api_key or "") if not base_url: raise CMHubError("cmhub_not_configured", "请去⑤设置配置 cmhub Base URL") if not api_key: raise CMHubError("cmhub_not_configured", "请去⑤设置配置 cmhub API Key") data = _cmhub_call_with_retry( "GET", appconfig.cmhub_request_url(base_url, "/api/v1/models"), api_key, payload=None, connect_timeout=max(1, int(connect_timeout or 10)), read_timeout=max(1, int(read_timeout or 30)), attempts=1, on_retry=None, ) models = data.get("models", []) if not isinstance(models, list): raise CMHubError("bad_response", "cmhub 模型列表格式错误") return [copy.deepcopy(model) for model in models if isinstance(model, dict)] def fetch_cmhub_balance(base_url, api_key, connect_timeout=10, read_timeout=30): """Fetch cmhub point balance for settings UI.""" base_url = str(base_url or "").strip() api_key = str(api_key or "") if not base_url: raise CMHubError("cmhub_not_configured", "请去⑤设置配置 cmhub Base URL") if not api_key: raise CMHubError("cmhub_not_configured", "请去⑤设置配置 cmhub API Key") data = _cmhub_call_with_retry( "GET", appconfig.cmhub_request_url(base_url, "/api/v1/balance"), api_key, payload=None, connect_timeout=max(1, int(connect_timeout or 10)), read_timeout=max(1, int(read_timeout or 30)), attempts=1, on_retry=None, ) if "points_balance" not in data: raise CMHubError("bad_response", "cmhub 余额返回格式错误") return copy.deepcopy(data) def _ai_backend(ai_cfg): backend = str(ai_cfg.get("backend", "direct") or "direct").strip().lower() if backend not in appconfig.AI_BACKENDS: raise AIError("AI backend 必须是 direct 或 cmhub") return backend def _gen_title_cmhub( title_prompt, old_title, retry, config, cmhub_config_path, on_step=None, on_event=None, ): ai_cfg = appconfig.ai_config(config) _notify_step(on_step, "load_text_model") runtime = _cmhub_runtime(config, "title", cmhub_config_path) resolution = _normalize_cmhub_resolution(ai_cfg.get("resolution", "1k")) _notify_step(on_step, "title_build_request") payload = { "prompt": _cmhub_title_prompt(title_prompt, old_title), "model": runtime["alias"], "resolution": resolution, } attempts = _attempt_count(ai_cfg, retry) _notify_step(on_step, "title_request") data = _cmhub_call_with_retry( "POST", appconfig.cmhub_request_url(runtime["base_url"], "/api/v1/generate/title"), runtime["api_key"], payload=payload, connect_timeout=runtime["connect_timeout"], read_timeout=_cmhub_read_timeout(config, ai_cfg.get("resolution", "1k")), attempts=attempts, on_retry=lambda attempt, total_attempts, exc: _notify_cmhub_retry( on_step, "title_request", attempt, total_attempts, exc, ), ) _emit_cmhub_metadata(on_event, data, "title_request") _notify_step(on_step, "title_parse_response") titles = data.get("titles") if not isinstance(titles, list) or not titles: raise AIError("AI 返回为空标题") text = str(titles[0] or "").strip() if not text: raise AIError("AI 返回为空标题") return text def _gen_cover_cmhub( cover_prompt, old_cover_path, out_path, resolution, jpg_quality, retry, config, cmhub_config_path, on_step=None, on_event=None, ): ai_cfg = appconfig.ai_config(config) _notify_step(on_step, "load_image_model") runtime = _cmhub_runtime(config, "image", cmhub_config_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)) _notify_step(on_step, "cover_build_request") payload = { "prompt": str(cover_prompt or ""), "model": runtime["alias"], "image_base64": _image_data_url(old_cover_path), "resolution": _normalize_cmhub_resolution(resolution), "aspect_ratio": "1:1", } attempts = _attempt_count(ai_cfg, retry) read_timeout = _cmhub_read_timeout(config, resolution) _notify_step(on_step, "cover_request") data = _cmhub_call_with_retry( "POST", appconfig.cmhub_request_url(runtime["base_url"], "/api/v1/generate/image"), runtime["api_key"], payload=payload, connect_timeout=runtime["connect_timeout"], read_timeout=read_timeout, attempts=attempts, on_retry=lambda attempt, total_attempts, exc: _notify_cmhub_retry( on_step, "cover_request", attempt, total_attempts, exc, ), ) _emit_cmhub_metadata(on_event, data, "cover_request") _notify_step(on_step, "cover_parse_response") image_url = str(data.get("image_url") or "").strip() if not image_url: raise AIError("AI 返回中没有图片数据") image_bytes = _download_cmhub_image( image_url, connect_timeout=runtime["connect_timeout"], read_timeout=read_timeout, ) _notify_step(on_step, "cover_save") return _save_jpeg(image_bytes, out_path, resolution, quality) def _cmhub_runtime(config, operation, cmhub_config_path): hub = appconfig.cmhub_config(config) api_key = appconfig.get_cmhub_api_key(path=cmhub_config_path) alias_key = "title_alias" if operation == "title" else "image_alias" missing = [] if not hub.get("base_url"): missing.append("Base URL") if not api_key: missing.append("API Key") if not hub.get(alias_key): missing.append("生文别名" if operation == "title" else "生图别名") if missing: raise CMHubError( "cmhub_not_configured", "请去⑤设置配置 cmhub：缺少 " + "、".join(missing), retryable=False, ) return { "base_url": hub["base_url"].rstrip("/"), "api_key": api_key, "alias": hub[alias_key], "connect_timeout": max(1, int(hub.get("connect_timeout", 10) or 10)), } def _cmhub_title_prompt(title_prompt, old_title): return "%s\n\n旧标题：\n%s\n\n请只返回新标题，不要解释。" % ( str(title_prompt or "").strip(), str(old_title or ""), ) def _normalize_cmhub_resolution(resolution): value = str(resolution or "1k").strip().lower() mapping = { "512": "512", "512x512": "512", "1k": "1K", "1K": "1K", "1024": "1K", "2k": "2K", "2K": "2K", "2048": "2K", "4k": "4K", "4K": "4K", "4096": "4K", } return mapping.get(value, str(resolution or "1K").upper()) def _cmhub_read_timeout(config, resolution): try: timeout = int(appconfig.response_timeout(config)) except Exception: timeout = 600 resolution_key = str(resolution or "").strip().lower() timeouts = appconfig.ai_config(config).get("resolution_timeouts", {}) if resolution_key in timeouts: timeout = int(timeouts[resolution_key]) return min(600, max(1, timeout)) def _cmhub_call_with_retry( method, url, api_key, payload, connect_timeout, read_timeout, attempts, on_retry=None, ): attempts = max(1, int(attempts or 1)) last_exc = None for index in range(attempts): try: return _cmhub_call_once( method, url, api_key, payload, connect_timeout=connect_timeout, read_timeout=read_timeout, ) except CMHubError as exc: last_exc = exc if not exc.retryable or index + 1 >= attempts: raise if on_retry is not None: try: on_retry(index + 1, attempts, exc) except Exception: pass time.sleep(_cmhub_retry_delay(exc, index)) raise last_exc def _cmhub_call_once(method, url, api_key, payload, connect_timeout, read_timeout): headers = { "Authorization": "Bearer " + str(api_key), "Accept": "application/json", } request_kwargs = { "headers": headers, "timeout": (max(1, int(connect_timeout)), max(1, int(read_timeout))), } if str(method).upper() != "GET": request_kwargs["json"] = payload or {} try: response = requests.request(str(method).upper(), url, **request_kwargs) except requests.exceptions.ConnectTimeout as exc: raise CMHubError("connect_timeout", "连接 cmhub 超时", retryable=True) from exc except requests.exceptions.ReadTimeout as exc: raise CMHubError("read_timeout", "等待 cmhub 返回超时", retryable=False) from exc except requests.exceptions.RequestException as exc: raise CMHubError( "network_error", _redact_cmhub(str(exc), api_key), retryable=False, ) from exc return _cmhub_response_json(response, api_key) def _cmhub_response_json(response, api_key): try: data = response.json() except ValueError: data = {} status = getattr(response, "status_code", None) if status and status >= 400: raise _cmhub_error_from_response(data, response, api_key) if isinstance(data, dict) and isinstance(data.get("error"), dict): raise _cmhub_error_from_response(data, response, api_key) if not isinstance(data, dict): raise CMHubError("bad_response", "cmhub 返回格式错误", status=status) return data def _cmhub_error_from_response(data, response, api_key): status = getattr(response, "status_code", None) error = data.get("error") if isinstance(data, dict) else None if not isinstance(error, dict): error = {} code = _normalize_cmhub_error_code(error.get("code") or _cmhub_code_for_status(status), status) raw_message = (error.get("message") or data.get("message")) if isinstance(data, dict) else "" if not raw_message: raw_message = getattr(response, "text", "")[:500] message = _cmhub_user_message(code, _redact_cmhub(raw_message or code, api_key)) retry_after = _parse_retry_after(getattr(response, "headers", {}).get("Retry-After")) return CMHubError( code, message, status=status, retryable=_cmhub_retryable(code), retry_after=retry_after, ) def _normalize_cmhub_error_code(code, status=None): normalized = str(code or "unknown").strip().lower().replace("-", "_") if status == 404 or normalized in {"notfound", "not_found"}: return "not_found" return normalized or "unknown" def _cmhub_code_for_status(status): return { 400: "bad_request", 401: "unauthorized", 402: "insufficient_points", 403: "account_disabled", 404: "not_found", 429: "rate_limited", 502: "upstream_error", }.get(status, "unknown") def _cmhub_retryable(code): return str(code) in {"upstream_error", "rate_limited", "connect_timeout"} def _cmhub_user_message(code, message): defaults = { "insufficient_points": "点数不足，请先充值", "unauthorized": "cmhub API Key 无效，请去⑤设置重填", "account_disabled": "cmhub 账号已禁用，请去网页端处理", "bad_request": "cmhub 请求参数错误", "model_not_allowed": "cmhub 模型别名无权限", "no_pricing_rule": "cmhub 模型别名未配置价格", "content_blocked": "cmhub 内容安全策略拒绝本次生成", "upstream_error": "cmhub 上游生成失败，请稍后重试", "rate_limited": "cmhub 请求过于频繁，请稍后重试", "not_found": "cmhub 接口不存在，请检查 Base URL 或该实例是否已部署 /api/v1/models", } default = defaults.get(str(code)) if str(code) == "not_found": return default if default and message and str(message) not in default: return "%s：%s" % (default, message) return default or str(message or code) def _parse_retry_after(value): try: if value is None or value == "": return None return max(0.0, float(value)) except (TypeError, ValueError): return None def _cmhub_retry_delay(exc, index): if exc.retry_after is not None: return min(2.0, max(0.1, float(exc.retry_after))) return min(2.0, 0.4 * (index + 1)) def _notify_cmhub_retry(callback, step, attempt, attempts, exc): if callback is None: return try: callback( { "step": step, "result": "retry", "attempt": attempt, "attempts": attempts, "detail": str(exc), } ) except Exception: pass def _emit_cmhub_metadata(callback, data, step): if callback is None or not isinstance(data, dict): return metadata = { key: data.get(key) for key in ("alias", "model_used", "points_cost", "points_balance", "call_id") if data.get(key) is not None } if not metadata: return try: callback( { "step": step, "result": "meta", "level": "info", "metadata": metadata, } ) except Exception: pass def _download_cmhub_image(url, connect_timeout, read_timeout, max_bytes=CMHUB_IMAGE_MAX_BYTES): _assert_public_http_url(url) try: response = requests.get( url, stream=True, timeout=(max(1, int(connect_timeout)), max(1, int(read_timeout))), ) except requests.exceptions.RequestException as exc: raise AIError("下载 cmhub 图片失败: %s" % exc) from exc status = getattr(response, "status_code", 200) if status >= 400: raise AIError("下载 cmhub 图片失败: HTTP %s" % status) chunks = [] total = 0 iterator = response.iter_content(chunk_size=65536) if hasattr(response, "iter_content") else [response.content] for chunk in iterator: if not chunk: continue total += len(chunk) if total > max_bytes: raise AIError("下载 cmhub 图片失败: 图片超过大小上限") chunks.append(chunk) return b"".join(chunks) def _assert_public_http_url(url): parts = urllib.parse.urlsplit(str(url or "")) if parts.scheme not in {"http", "https"}: raise AIError("cmhub 图片地址只允许 http/https") host = parts.hostname if not host: raise AIError("cmhub 图片地址缺少域名") if _is_local_hostname(host): raise AIError("cmhub 图片地址不能指向本机或内网") try: _assert_public_ip(host) return except ValueError: pass try: addresses = socket.getaddrinfo( host, parts.port or (443 if parts.scheme == "https" else 80), type=socket.SOCK_STREAM, ) except OSError as exc: raise AIError("cmhub 图片地址无法解析: %s" % exc) from exc if not addresses: raise AIError("cmhub 图片地址无法解析") for address in addresses: ip_text = address[4][0] _assert_public_ip(ip_text) def _is_local_hostname(host): lowered = str(host or "").strip().lower().rstrip(".") return lowered in {"localhost"} or lowered.endswith(".localhost") or lowered.endswith(".local") def _assert_public_ip(value): ip = ipaddress.ip_address(value) if ( ip.is_private or ip.is_loopback or ip.is_link_local or ip.is_multicast or ip.is_reserved or ip.is_unspecified ): raise AIError("cmhub 图片地址不能指向本机或内网") def _redact_cmhub(text, api_key): return appconfig.redact_secrets(text, [api_key]) 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", "backend", "cmhub", } } 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, metadata=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 if metadata is not None: payload["metadata"] = appconfig.sanitize_for_log(metadata) 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, } code = getattr(exc, "code", None) status = getattr(exc, "status", None) retryable = getattr(exc, "retryable", None) if code is not None: payload["code"] = str(code) if status is not None: payload["status"] = status if retryable is not None: payload["retryable"] = bool(retryable) 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