feat: add cmhub AI backend

This commit is contained in:
chengma
2026-07-04 15:13:15 +08:00
parent c8a5e9ada8
commit 1efb095767
11 changed files with 1061 additions and 43 deletions
+516
View File
@@ -1,17 +1,22 @@
"""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
@@ -20,6 +25,18 @@ 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),
@@ -36,11 +53,23 @@ def gen_title(
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")
@@ -88,6 +117,8 @@ def gen_cover(
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."""
@@ -100,6 +131,19 @@ def gen_cover(
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"))
@@ -200,6 +244,7 @@ def generate_batch(tasks, prompts, ai_cfg=None, on_progress=None, should_stop=No
)
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")
@@ -253,6 +298,7 @@ def generate_batch(tasks, prompts, ai_cfg=None, on_progress=None, should_stop=No
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)
@@ -278,6 +324,8 @@ def generate_batch(tasks, prompts, ai_cfg=None, on_progress=None, should_stop=No
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):
@@ -366,6 +414,8 @@ def generate_batch(tasks, prompts, ai_cfg=None, on_progress=None, should_stop=No
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:
@@ -415,6 +465,467 @@ def generate_batch(tasks, prompts, ai_cfg=None, on_progress=None, should_stop=No
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 _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 = str(error.get("code") or _cmhub_code_for_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 _cmhub_code_for_status(status):
return {
400: "bad_request",
401: "unauthorized",
402: "insufficient_points",
403: "account_disabled",
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 请求过于频繁,请稍后重试",
}
default = defaults.get(str(code))
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)
@@ -450,6 +961,8 @@ def _runtime_config(runtime):
"resolution",
"generate_cover",
"resolution_timeouts",
"backend",
"cmhub",
}
}
if ai_updates:
@@ -510,6 +1023,7 @@ def _emit_generation_event(
level="info",
attempt=None,
attempts=None,
metadata=None,
):
if callback is None:
return
@@ -526,6 +1040,8 @@ def _emit_generation_event(
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:
+93
View File
@@ -14,8 +14,10 @@ import urllib.request
CONFIG_PATH = "config.json"
AI_MODELS_PATH = os.path.join("config", "ai_models.json")
CMHUB_CONFIG_PATH = os.path.join("config", "cmhub.json")
CATEGORIES = {"text", "image"}
API_TYPES = {"chat", "images_edits", "auto"}
AI_BACKENDS = {"direct", "cmhub"}
DEFAULT_CONFIG = {
"chrome_path": r"C:\Program Files\Google\Chrome\Application\chrome.exe",
@@ -29,6 +31,14 @@ DEFAULT_CONFIG = {
"default_text_model": "GPT-5.5 文本",
"default_image_model": "Nano Banana 2",
"generate_cover": False,
"backend": "direct",
"cmhub": {
"base_url": "",
"title_alias": "",
"image_alias": "",
"connect_timeout": 10,
"check_balance_before_batch": False,
},
"title_concurrency": 4,
"image_concurrency": 4,
"retry": 2,
@@ -82,6 +92,10 @@ DEFAULT_AI_MODELS_CONFIG = {
]
}
DEFAULT_CMHUB_CONFIG = {
"api_key": "",
}
SECRET_FIELD_NAMES = {"api_key", "apikey", "key", "token", "password"}
@@ -183,6 +197,52 @@ def redact_secrets(text, secret_values=None) -> str:
return redacted
def default_cmhub_config() -> dict:
"""Return a new copy of the default cmhub key config."""
return copy.deepcopy(DEFAULT_CMHUB_CONFIG)
def _normalize_cmhub_config(config):
if config is None:
config = {}
if not isinstance(config, dict):
raise ConfigError("config/cmhub.json 必须是对象")
return {"api_key": str(config.get("api_key", "") or "")}
def load_cmhub_config(path=CMHUB_CONFIG_PATH) -> dict:
"""Load cmhub API key config. Missing file means key is not configured."""
if not os.path.exists(path):
return default_cmhub_config()
with open(path, "r", encoding="utf-8") as fh:
try:
loaded = json.load(fh)
except json.JSONDecodeError as exc:
raise ConfigError(f"cmhub 配置不是有效 JSON: {path}") from exc
return _normalize_cmhub_config(loaded)
def save_cmhub_config(config, path=CMHUB_CONFIG_PATH) -> dict:
"""Persist cmhub API key config, including the local plaintext key."""
normalized = _normalize_cmhub_config(config)
directory = os.path.dirname(os.path.abspath(path))
if directory:
os.makedirs(directory, exist_ok=True)
with open(path, "w", encoding="utf-8") as fh:
json.dump(normalized, fh, ensure_ascii=False, indent=2)
fh.write("\n")
return normalized
def get_cmhub_api_key(path=CMHUB_CONFIG_PATH, masked=False) -> str:
key = load_cmhub_config(path).get("api_key", "")
return mask_secret(key) if masked else key
def save_config(config, path=CONFIG_PATH) -> dict:
"""Persist config to JSON and return the normalized config."""
@@ -258,6 +318,39 @@ def ai_config(config=None) -> dict:
return copy.deepcopy(_config_or_load(config).get("ai", DEFAULT_CONFIG["ai"]))
def ai_backend(config=None) -> str:
ai = ai_config(config)
backend = str(ai.get("backend", "direct") or "direct").strip().lower()
if backend not in AI_BACKENDS:
raise ConfigError("AI backend 必须是 direct 或 cmhub")
return backend
def cmhub_config(config=None) -> dict:
ai = ai_config(config)
value = ai.get("cmhub", {})
if not isinstance(value, dict):
raise ConfigError("ai.cmhub 必须是对象")
merged = _deep_merge(DEFAULT_CONFIG["ai"]["cmhub"], value)
merged["base_url"] = str(merged.get("base_url", "") or "").strip()
merged["title_alias"] = str(merged.get("title_alias", "") or "").strip()
merged["image_alias"] = str(merged.get("image_alias", "") or "").strip()
merged["connect_timeout"] = int(merged.get("connect_timeout", 10) or 10)
merged["check_balance_before_batch"] = bool(merged.get("check_balance_before_batch", False))
if merged["connect_timeout"] <= 0:
raise ConfigError("ai.cmhub.connect_timeout 必须大于 0")
return merged
def cmhub_request_url(base_url, endpoint) -> str:
base = str(base_url or "").strip().rstrip("/")
path = "/" + str(endpoint or "").strip().lstrip("/")
if not base:
return path
return base + path
def response_timeout(config=None) -> int:
ai = ai_config(config)
resolution = str(ai.get("resolution", DEFAULT_CONFIG["ai"]["resolution"]))