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