zcbot/tools/gpt_image.py

"""gpt_image: 调 unifyllm 网关的 OpenAI Images API 生图,产物落 working_dir/figures/。

第二图像后端(第一个是豆包 seedream):模型 ID + 单价全在 `config/media/unifyllm.yaml`,
本 tool 只装配。与 seedream 的差异(网关实测,见 yaml 头注释):
- 只支持文生图,**不支持改图(i2i)/ 自定尺寸 / 质量档**(网关忽略 size/quality 参数,
  实际尺寸由上游模型按画面自动决定),所以参数只有 prompt;
- 响应直接返 b64_json,无需二次下载;
- 单张 ~35-40s(慢于 seedream 3-5s)。
完成后:
- 图片落 `<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 secrets
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.storage.usage import check_daily_quota, record_image_usage

from .base import Tool


class GptImageTool(Tool):
    name = "gpt_image"
    description = (
        "Generate an image (text-to-image only) via the GPT image backend, "
        "saved to working_dir/figures/. Slower than seedream (~40-60s/image); output size is "
        "chosen automatically by the model and CANNOT be specified. Does NOT support "
        "image-to-image editing — if the user wants to modify an existing image or needs an "
        "exact size/aspect ratio, tell them to switch the image model back to 豆包 Seedream "
        "(顶栏图像模型下拉). 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": "中文或英文都行,详尽描述画面(主体/风格/光线/构图)。尺寸/比例不可控,别在 prompt 里承诺具体分辨率。",
            },
        },
        "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) -> str:
        if not (prompt or "").strip():
            return "[Error] prompt 不能为空"

        # 每账号每日配额(kind="image" 与 seedream 同口径合计;失败不落库故 retry 不计)
        if self.daily_limit > 0:
            used, over = check_daily_quota(user_id=self.user_id, kind="image", limit=self.daily_limit)
            if over:
                return (
                    f"[Error] 已达每日图片生成上限({used}/{self.daily_limit} 张),"
                    f"次日 00:00 重置。"
                )

        cfg = self.cfg
        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}

        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}"

        ts = datetime.now().strftime("%Y%m%d-%H%M%S")
        short = secrets.token_hex(3)
        dest_png = self.working_dir / "figures" / f"{ts}-{short}.png"
        dest_png.parent.mkdir(parents=True, exist_ok=True)
        dest_png.write_bytes(img_bytes)

        elapsed = time.monotonic() - t0
        # 网关回填的实际尺寸/质量/tokens(价目未知期成本记 price snapshot,tokens 留对账)
        actual_size = str(resp.get("size") or "")
        quality = str(resp.get("quality") or "")
        usage = resp.get("usage") or {}
        output_tokens = int(usage.get("output_tokens") or 0)

        meta = {
            "prompt": prompt,
            "model_id": model_id,
            "size": actual_size,
            "quality": quality,
            "mode": "t2i",
            "cost_cny": price,
            "output_tokens": output_tokens,
            "elapsed_s": round(elapsed, 2),
            "ts": datetime.now().isoformat(timespec="seconds"),
        }
        dest_png.with_suffix(".meta.json").write_text(
            json.dumps(meta, ensure_ascii=False, indent=2), encoding="utf-8"
        )

        try:
            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": quality},
            )
        except Exception as e:
            print(f"[gpt_image] record_image_usage failed: {type(e).__name__}: {e}", flush=True)

        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 ""
        return (
            f"[gpt_image] model={model_id} · size={actual_size}{cost_seg} · elapsed={elapsed:.1f}s\n"
            f"saved: {disp}\n"
            f"prompt={prompt!r}"
        )