284 lines
11 KiB
Python
284 lines
11 KiB
Python
"""gpt_image: 调 unifyllm 网关的 OpenAI Images API 生图 / 改图。
|
||
|
||
第二图像后端(第一个是豆包 seedream):模型 ID + 单价全在 `config/media/unifyllm.yaml`,
|
||
本 tool 只装配:
|
||
- 文生图走 /images/generations JSON,gpt-image-2 支持自定义尺寸和质量档;
|
||
- 单图改图走 /images/edits multipart/form-data;
|
||
- 响应直接返 b64_json,无需二次下载;
|
||
- 复杂图片可能需 ~2min。
|
||
完成后:
|
||
- 图片落 `<working_dir>/figures/<YYYYMMDD-HHMMSS>-<rand6>.png`，技术元数据进隐藏 `.meta/`
|
||
- usage_events 写 kind="image" 一行(model_profile="unifyllm.<variant>",
|
||
usage tokens 记进 units → 网关价目公布后可回填对账)
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import base64
|
||
import json
|
||
import re
|
||
import struct
|
||
import time
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
from typing import Optional
|
||
from uuid import UUID
|
||
|
||
from core.ark_client import ArkClient, ArkConfig, ArkError
|
||
from core.artifacts import ArtifactRef, ToolExecutionResult, resolve_artifact_path
|
||
from core.storage.usage import record_image_usage
|
||
|
||
from .base import Tool
|
||
from .image_ref import load_image_as_data_url
|
||
from .media_common import quota_gate, record_usage_safe, stamped_path, write_meta
|
||
|
||
|
||
class GptImageTool(Tool):
|
||
name = "gpt_image"
|
||
description = (
|
||
"Generate an image via GPT Image, saved to working_dir/figures/. Supports custom size "
|
||
"and quality, plus single-reference image editing. For editing, pass reference_images "
|
||
"with the existing image path; the original is preserved and the result is saved as a "
|
||
"new image. Complex images may take up to ~2 minutes. Don't generate decoratively — "
|
||
"only when the user actually wants an image. Returns the saved relative path."
|
||
)
|
||
parameters = {
|
||
"type": "object",
|
||
"properties": {
|
||
"prompt": {
|
||
"type": "string",
|
||
"description": "中文或英文都行,详尽描述画面(主体/风格/光线/构图)。",
|
||
},
|
||
"reference_images": {
|
||
"type": "array",
|
||
"items": {"type": "string"},
|
||
"description": (
|
||
"改图(image-to-image):传 1 张 task_dir 内已存在图片的相对路径。"
|
||
"不传 = 从零文生图;当前仅支持单张参考图。"
|
||
),
|
||
},
|
||
"size": {
|
||
"type": "string",
|
||
"description": (
|
||
"输出尺寸。默认 auto;也可传 WIDTHxHEIGHT,如 1024x1024、1536x864、"
|
||
"2048x1152。宽高须为 16 的倍数,比例不超过 3:1,总像素 "
|
||
"655360-8294400,单边不超过 3840。"
|
||
),
|
||
},
|
||
"quality": {
|
||
"type": "string",
|
||
"enum": ["auto", "low", "medium", "high"],
|
||
"description": "生成质量。默认 auto;草稿用 low,正式产物用 medium 或 high。",
|
||
},
|
||
},
|
||
"required": ["prompt"],
|
||
}
|
||
|
||
def __init__(
|
||
self,
|
||
*,
|
||
gw_cfg: ArkConfig,
|
||
image_variant_cfg: dict,
|
||
variant_key: str,
|
||
working_dir: Path,
|
||
task_id: UUID,
|
||
user_id: UUID,
|
||
base_dir: Optional[Path] = None,
|
||
user_root: Optional[Path] = None,
|
||
daily_limit: int = 0,
|
||
) -> None:
|
||
super().__init__(base_dir, user_root=user_root)
|
||
self.gw_cfg = gw_cfg
|
||
self.cfg = image_variant_cfg
|
||
self.variant_key = variant_key # 'gpt_image' → usage_events.model_profile = "unifyllm.gpt_image"
|
||
self.working_dir = Path(working_dir)
|
||
self.task_id = task_id
|
||
self.user_id = user_id
|
||
self.daily_limit = int(daily_limit) # 0 / 负 = 不限;与 seedream 共享 kind="image" 每日配额
|
||
|
||
def execute(
|
||
self,
|
||
prompt: str,
|
||
reference_images: list | None = None,
|
||
size: Optional[str] = None,
|
||
quality: Optional[str] = None,
|
||
) -> str:
|
||
if not (prompt or "").strip():
|
||
return "[Error] prompt 不能为空"
|
||
|
||
refs = [str(r).strip() for r in (reference_images or []) if str(r).strip()]
|
||
if len(refs) > 1:
|
||
return (
|
||
"[Error] reference_images 当前仅支持单张参考图(传了 "
|
||
f"{len(refs)} 张)。请只传 1 张。"
|
||
)
|
||
ref_bytes = b""
|
||
ref_mime = ""
|
||
ref_disp = ""
|
||
if refs:
|
||
data_url, ref_disp, ref_err = load_image_as_data_url(
|
||
refs[0],
|
||
working_dir=self.working_dir,
|
||
user_root=self.user_root,
|
||
display_fn=self._display,
|
||
)
|
||
if ref_err:
|
||
return ref_err
|
||
header, encoded = data_url.split(",", 1)
|
||
ref_mime = header.removeprefix("data:").removesuffix(";base64")
|
||
ref_bytes = base64.b64decode(encoded)
|
||
is_i2i = bool(ref_bytes)
|
||
|
||
cfg = self.cfg
|
||
chosen_size, size_err = self._normalize_size(
|
||
size if size is not None else cfg.get("default_size", "auto")
|
||
)
|
||
if size_err:
|
||
return size_err
|
||
chosen_quality = str(
|
||
quality if quality is not None else cfg.get("default_quality", "auto")
|
||
).strip().lower()
|
||
if chosen_quality not in {"auto", "low", "medium", "high"}:
|
||
return "[Error] quality 必须是 auto / low / medium / high"
|
||
|
||
# 每账号每日配额(kind="image" 与 seedream 同口径合计)
|
||
quota_err = quota_gate(
|
||
self.user_id, kind="image", limit=self.daily_limit, what="图片生成", noun="张",
|
||
)
|
||
if quota_err:
|
||
return quota_err
|
||
|
||
model_id = cfg["model_id"]
|
||
endpoint = (
|
||
cfg.get("edit_endpoint", "/images/edits")
|
||
if is_i2i
|
||
else cfg.get("endpoint", "/images/generations")
|
||
)
|
||
timeout_s = float(cfg.get("request_timeout_s", 300))
|
||
price = float(cfg.get("price_cny_per_image", 0))
|
||
|
||
body = {
|
||
"model": model_id,
|
||
"prompt": prompt,
|
||
"n": 1,
|
||
"size": chosen_size,
|
||
"quality": chosen_quality,
|
||
}
|
||
|
||
t0 = time.monotonic()
|
||
try:
|
||
with ArkClient(self.gw_cfg, timeout_s=timeout_s) as client:
|
||
if is_i2i:
|
||
form = {
|
||
key: str(value).lower() if isinstance(value, bool) else str(value)
|
||
for key, value in body.items()
|
||
}
|
||
form["response_format"] = "b64_json"
|
||
resp = client.post_multipart(
|
||
endpoint,
|
||
form,
|
||
{"image": (Path(refs[0]).name or "reference.png", ref_bytes, ref_mime)},
|
||
timeout_s=timeout_s,
|
||
)
|
||
else:
|
||
resp = client.post_json(endpoint, body, timeout_s=timeout_s)
|
||
except ArkError as e:
|
||
return f"[Error] gpt_image API: {e}"
|
||
|
||
data = resp.get("data")
|
||
b64 = ""
|
||
if isinstance(data, list) and data and isinstance(data[0], dict):
|
||
b64 = data[0].get("b64_json") or ""
|
||
if not b64:
|
||
image_url = str(data[0].get("url") or "")
|
||
if image_url.startswith("data:") and ";base64," in image_url:
|
||
b64 = image_url.split(",", 1)[1]
|
||
if not b64:
|
||
return f"[Error] gpt_image response 缺 b64_json: {json.dumps(resp, ensure_ascii=False)[:300]}"
|
||
try:
|
||
img_bytes = base64.b64decode(b64)
|
||
except Exception as e:
|
||
return f"[Error] gpt_image b64 解码失败: {e}"
|
||
|
||
dest_png = stamped_path(self.working_dir / "figures", ".png")
|
||
dest_png.write_bytes(img_bytes)
|
||
|
||
elapsed = time.monotonic() - t0
|
||
# 网关回填的实际尺寸/质量/tokens(价目未知期成本记 price snapshot,tokens 留对账)
|
||
actual_size = self._png_size(img_bytes) or str(resp.get("size") or chosen_size)
|
||
actual_quality = str(resp.get("quality") or chosen_quality)
|
||
usage = resp.get("usage") or {}
|
||
output_tokens = int(usage.get("output_tokens") or 0)
|
||
|
||
meta = {
|
||
"prompt": prompt,
|
||
"model_id": model_id,
|
||
"size": actual_size,
|
||
"requested_size": chosen_size,
|
||
"quality": actual_quality,
|
||
"requested_quality": chosen_quality,
|
||
"mode": "i2i" if is_i2i else "t2i",
|
||
"reference_images": [ref_disp] if is_i2i else [],
|
||
"cost_cny": price,
|
||
"output_tokens": output_tokens,
|
||
"elapsed_s": round(elapsed, 2),
|
||
"ts": datetime.now().isoformat(timespec="seconds"),
|
||
}
|
||
write_meta(dest_png, meta)
|
||
|
||
record_usage_safe(
|
||
"gpt_image", record_image_usage,
|
||
task_id=self.task_id,
|
||
user_id=self.user_id,
|
||
model_profile=f"unifyllm.{self.variant_key}",
|
||
n_images=1,
|
||
size=actual_size,
|
||
price_cny_per_image=price,
|
||
extra_units={"output_tokens": output_tokens, "quality": actual_quality},
|
||
)
|
||
|
||
disp = self._display(dest_png)
|
||
# 首行 banner 协议同 seedream(`key=value · ` 分隔,前端 extractMediaBanner 解析);
|
||
# 价格未知(price=0)时不放 cost 段,避免"¥0.00 = 免费"的误导。
|
||
cost_seg = f" · cost=¥{price:.2f}" if price > 0 else ""
|
||
mode_seg = " · mode=i2i" if is_i2i else ""
|
||
ref_line = f"\nreference={ref_disp}" if is_i2i else ""
|
||
result = (
|
||
f"[gpt_image] model={model_id} · size={actual_size} · quality={actual_quality}"
|
||
f"{cost_seg} · elapsed={elapsed:.1f}s{mode_seg}\n"
|
||
f"saved: {disp}{ref_line}\n"
|
||
f"prompt={prompt!r}"
|
||
)
|
||
if self.user_root is None:
|
||
return result
|
||
_, rel = resolve_artifact_path(
|
||
str(dest_png), working_dir=self.working_dir, user_root=self.user_root,
|
||
)
|
||
return ToolExecutionResult(content=result, artifacts=(ArtifactRef(path=rel),))
|
||
|
||
@staticmethod
|
||
def _normalize_size(raw: object) -> tuple[str, str]:
|
||
value = str(raw or "auto").strip().lower().replace("×", "x")
|
||
if value == "auto":
|
||
return "auto", ""
|
||
match = re.fullmatch(r"([1-9]\d*)x([1-9]\d*)", value)
|
||
if not match:
|
||
return "", "[Error] size 必须是 auto 或 WIDTHxHEIGHT,例如 1536x864"
|
||
width, height = (int(match.group(1)), int(match.group(2)))
|
||
pixels = width * height
|
||
if width % 16 or height % 16:
|
||
return "", "[Error] size 的宽和高都必须是 16 的倍数"
|
||
if max(width, height) > 3840:
|
||
return "", "[Error] size 单边不能超过 3840"
|
||
if max(width, height) > min(width, height) * 3:
|
||
return "", "[Error] size 宽高比不能超过 3:1"
|
||
if not 655_360 <= pixels <= 8_294_400:
|
||
return "", "[Error] size 总像素必须在 655360 到 8294400 之间"
|
||
return f"{width}x{height}", ""
|
||
|
||
@staticmethod
|
||
def _png_size(data: bytes) -> str:
|
||
if len(data) >= 24 and data[:8] == b"\x89PNG\r\n\x1a\n" and data[12:16] == b"IHDR":
|
||
width, height = struct.unpack(">II", data[16:24])
|
||
return f"{width}x{height}"
|
||
return ""
|