227 lines
8.8 KiB
Python
227 lines
8.8 KiB
Python
"""gpt_image: 调 unifyllm 网关的 OpenAI Images API 生图,产物落 working_dir/figures/。
|
||
|
||
第二图像后端(第一个是豆包 seedream):模型 ID + 单价全在 `config/media/unifyllm.yaml`,
|
||
本 tool 只装配:
|
||
- 文生图走 /images/generations JSON,gpt-image-2 支持自定义尺寸和质量档;
|
||
- 当前 unifyllm 网关的 /images/edits 不可用,暂不暴露改图参数;
|
||
- 响应直接返 b64_json,无需二次下载;
|
||
- 复杂图片可能需 ~2min。
|
||
完成后:
|
||
- 图片落 `<working_dir>/figures/<YYYYMMDD-HHMMSS>-<rand6>.png` + 同名 `.meta.json`
|
||
- 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.artifacts import ArtifactRef, ToolExecutionResult, resolve_artifact_path
|
||
from core.ark_client import ArkClient, ArkConfig, ArkError
|
||
from core.storage.usage import record_image_usage
|
||
|
||
from .base import Tool
|
||
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. Complex images may take up to ~2 minutes. Image-to-image editing is not "
|
||
"available through the current gateway; ask the user to switch to 豆包 Seedream for "
|
||
"editing an existing image. 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": "中文或英文都行,详尽描述画面(主体/风格/光线/构图)。",
|
||
},
|
||
"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,
|
||
size: Optional[str] = None,
|
||
quality: Optional[str] = None,
|
||
) -> str:
|
||
if not (prompt or "").strip():
|
||
return "[Error] prompt 不能为空"
|
||
|
||
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("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:
|
||
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:
|
||
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": "t2i",
|
||
"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 ""
|
||
result = (
|
||
f"[gpt_image] model={model_id} · size={actual_size} · quality={actual_quality}"
|
||
f"{cost_seg} · elapsed={elapsed:.1f}s\n"
|
||
f"saved: {disp}\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 ""
|