Files
cmshoppe/app/ai.py
T

1457 lines
47 KiB
Python
Raw Normal View History

2026-06-27 15:47:21 +08:00
"""AI generation helpers backed by configurable HTTP model endpoints."""
import base64
2026-07-04 15:13:15 +08:00
import ipaddress
2026-06-27 16:26:08 +08:00
from concurrent.futures import CancelledError, ThreadPoolExecutor, as_completed
2026-06-27 15:47:21 +08:00
import copy
import json
import mimetypes
import os
2026-07-04 15:13:15 +08:00
import socket
import threading
2026-06-27 15:47:21 +08:00
import time
import urllib.error
2026-07-04 15:13:15 +08:00
import urllib.parse
2026-06-27 15:47:21 +08:00
import urllib.request
import uuid
2026-07-04 15:13:15 +08:00
import requests
2026-07-01 08:35:46 +08:00
from . import appconfig, db, diagnostics, image_paths
2026-06-27 16:26:08 +08:00
from . import prompts as prompt_module
2026-07-01 08:35:46 +08:00
2026-06-27 15:47:21 +08:00
class AIError(RuntimeError):
"""Raised when AI generation cannot complete."""
2026-07-04 15:13:15 +08:00
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
2026-06-27 15:47:21 +08:00
_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,
2026-07-04 15:13:15 +08:00
on_event=None,
cmhub_config_path=appconfig.CMHUB_CONFIG_PATH,
2026-06-27 15:47:21 +08:00
):
"""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)
2026-07-04 15:13:15 +08:00
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")
2026-06-27 15:47:21 +08:00
model = _role_model("text", ai_cfg.get("default_text_model"), models_path)
_notify_step(on_step, "title_build_request")
2026-06-27 15:47:21 +08:00
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")
2026-06-27 15:47:21 +08:00
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,
),
2026-06-27 15:47:21 +08:00
)
_notify_step(on_step, "title_parse_response")
2026-06-27 15:47:21 +08:00
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,
2026-07-04 15:13:15 +08:00
on_event=None,
cmhub_config_path=appconfig.CMHUB_CONFIG_PATH,
2026-06-27 15:47:21 +08:00
):
"""Generate a new cover image and save it as a JPEG file."""
_notify_step(on_step, "cover_validate_input")
2026-06-27 15:47:21 +08:00
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)
2026-07-04 15:13:15 +08:00
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")
2026-06-27 15:47:21 +08:00
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")
2026-06-27 15:47:21 +08:00
if api_type == "images_edits":
body, content_type = _image_edit_body(model, cover_prompt, old_cover_path, resolution)
_notify_step(on_step, "cover_request")
2026-06-27 15:47:21 +08:00
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,
),
2026-06-27 15:47:21 +08:00
)
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,
),
)
2026-06-27 15:47:21 +08:00
_notify_step(on_step, "cover_parse_response")
2026-06-27 15:47:21 +08:00
image_bytes = _extract_image_bytes(data, model, cfg)
_notify_step(on_step, "cover_save")
2026-06-27 15:47:21 +08:00
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
2026-06-27 16:26:08 +08:00
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",
2026-07-02 15:17:11 +08:00
"generate_cover",
2026-06-27 16:26:08 +08:00
}
}
)
db_path = runtime.get("db_path")
models_path = runtime.get("models_path", appconfig.AI_MODELS_PATH)
2026-07-04 15:13:15 +08:00
cmhub_config_path = runtime.get("cmhub_config_path", appconfig.CMHUB_CONFIG_PATH)
2026-06-27 16:26:08 +08:00
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")
2026-06-27 16:26:08 +08:00
title_prompt = _prompt_value(prompts, "title")
cover_prompt = _prompt_value(prompts, "cover")
2026-07-02 15:17:11 +08:00
generate_cover = bool(generation_cfg.get("generate_cover", False))
2026-06-27 16:26:08 +08:00
should_stop = should_stop or (lambda: False)
eligible = [
task for task in list(tasks)
if is_generatable_task(task)
2026-06-27 16:26:08 +08:00
]
summary = {
"ok": True,
"total": len(eligible),
"title_done": 0,
"cover_done": 0,
2026-07-02 15:17:11 +08:00
"cover_total": len(eligible) if generate_cover else 0,
"generated_done": 0,
2026-06-27 16:26:08 +08:00
"failed": 0,
"cancelled": False,
2026-07-02 15:17:11 +08:00
"generate_cover": generate_cover,
2026-06-27 16:26:08 +08:00
}
_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"),
2026-07-04 15:13:15 +08:00
metadata=event.get("metadata"),
)
return
set_step(task, event)
_emit_generation_event(on_event, task, phase, event, "start")
return callback
2026-06-27 16:26:08 +08:00
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")
2026-06-27 16:26:08 +08:00
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"),
2026-07-04 15:13:15 +08:00
on_event=step_callback(task, "title"),
cmhub_config_path=cmhub_config_path,
2026-06-27 16:26:08 +08:00
)
] = task
for future in as_completed(futures):
task = futures[future]
if should_stop():
summary["cancelled"] = True
_cancel_pending(futures)
try:
2026-07-02 15:17:11 +08:00
new_title = future.result()
title_results[task.id] = new_title
2026-06-27 16:26:08 +08:00
summary["title_done"] += 1
set_step(task, "title_done")
_emit_generation_event(on_event, task, "title", "title_done", "success")
2026-07-02 15:17:11 +08:00
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="仅生成标题",
)
2026-06-27 16:26:08 +08:00
except CancelledError:
summary["cancelled"] = True
_emit_generation_event(on_event, task, "title", get_step(task, "title_request"), "cancelled", level="warning")
2026-06-27 16:26:08 +08:00
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)
2026-06-27 16:26:08 +08:00
_emit_generation_progress(on_progress, summary)
2026-07-02 15:17:11 +08:00
if not generate_cover:
if summary["cancelled"]:
summary["ok"] = False
return summary
2026-06-27 16:26:08 +08:00
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),
2026-06-27 16:26:08 +08:00
)
_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"),
2026-07-04 15:13:15 +08:00
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)
2026-06-27 16:26:08 +08:00
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")
2026-06-27 16:26:08 +08:00
db.set_generated(task.id, new_title, new_cover_path, path=db_path)
summary["cover_done"] += 1
2026-07-02 15:17:11 +08:00
summary["generated_done"] += 1
2026-06-27 16:26:08 +08:00
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)
2026-06-27 16:26:08 +08:00
except CancelledError:
summary["cancelled"] = True
_emit_generation_event(on_event, task, "cover", get_step(task, "cover_request"), "cancelled", level="warning")
2026-06-27 16:26:08 +08:00
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)
2026-06-27 16:26:08 +08:00
_emit_generation_progress(on_progress, summary)
if summary["cancelled"]:
summary["ok"] = False
return summary
2026-07-04 15:13:15 +08:00
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)]
2026-07-04 16:29:15 +08:00
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)
2026-07-04 15:13:15 +08:00
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)
2026-07-04 15:13:15 +08:00
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"
2026-07-04 15:13:15 +08:00
def _cmhub_code_for_status(status):
return {
400: "bad_request",
401: "unauthorized",
402: "insufficient_points",
403: "account_disabled",
404: "not_found",
2026-07-04 15:13:15 +08:00
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",
2026-07-04 15:13:15 +08:00
}
default = defaults.get(str(code))
if str(code) == "not_found":
return default
2026-07-04 15:13:15 +08:00
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])
2026-06-27 15:47:21 +08:00
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
2026-06-27 16:26:08 +08:00
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",
2026-07-02 15:17:11 +08:00
"generate_cover",
2026-06-27 16:26:08 +08:00
"resolution_timeouts",
2026-07-04 15:13:15 +08:00
"backend",
"cmhub",
2026-06-27 16:26:08 +08:00
}
}
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)
2026-07-01 08:35:46 +08:00
return image_paths.task_image_path(image_root, task, account, "new")
2026-06-27 16:26:08 +08:00
def _mark_generate_failed(task, exc, db_path, on_task_update):
error = diagnostics.redact_log_text(str(exc) or exc.__class__.__name__)
2026-06-27 16:26:08 +08:00
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,
2026-07-04 15:13:15 +08:00
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
2026-07-04 15:13:15 +08:00
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,
}
2026-07-04 17:14:29 +08:00
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
2026-06-27 16:26:08 +08:00
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),
)
2026-06-27 15:47:21 +08:00
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,
):
2026-06-27 15:47:21 +08:00
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
2026-06-27 15:47:21 +08:00
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),
2026-06-27 15:47:21 +08:00
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