zcbot/tools/look_at_image.py

"""look_at_image: 给 DeepSeek V4 主模型(纯文本)借一双眼睛。

主模型无视觉,这个 tool 走豆包 Seed 2.0 Lite(全模态理解)读单图 —— OCR / 描述画面 /
读图表表格 / 识别物体。模型自决何时调(用户贴了图问"这写的啥" / "图里是什么")。

模型 ID + 单价全在 `config/media/doubao.yaml` 的 vision 段,本 tool 只装配。
计费:token 计费(同 chat),成功后写一行 usage_events(kind="vision")。
图片路径解析 + base64 复用 tools/image_ref(与 seedream i2i 同一套边界/校验)。
"""
from __future__ import annotations

from pathlib import Path
from typing import Any, Optional
from uuid import UUID

from core.ark_client import ArkConfig
from core.storage.usage import record_vision_usage

from .base import Tool
from .output import compact_tool_output
from .image_ref import load_image_as_data_url
from .media_common import ark_chat_with_retry, extract_chat_answer, record_usage_safe

# question 理应由主模型按用户任务具体填写；这里仅作旧调用兼容兜底。
# 默认只给简洁概览，避免无条件触发「描述 + 全文 OCR + 图表解析」导致长输出和高延迟。
_DEFAULT_QUESTION = (
    "请简洁说明这张图片的主要内容，只提供理解图片所必需的信息。"
    "不要主动全文 OCR、逐项枚举或展开图表数据；无法确定时明确说明。"
)


class LookAtImageTool(Tool):
    name = "look_at_image"
    description = (
        "Read/understand an image using Doubao Seed 2.0 Lite vision (the main model is text-only). "
        "Use to OCR text, describe a picture, read charts/tables/diagrams, or identify objects in an "
        "image the user uploaded (look for a `[用户上传的参考图] <path>` line in their message) or that "
        "was generated/saved in the task. Pass the image path and a task-specific `question`; "
        "only request full-image OCR when the user explicitly needs it. SLOW (tens of seconds per "
        "call) — only call when you genuinely need the image's actual content to proceed, and ask "
        "everything you need in ONE call; never re-read the same image. "
        "Returns the model's textual reading of the image."
    )
    parameters = {
        "type": "object",
        "properties": {
            "image": {
                "type": "string",
                "description": (
                    "要看的图片相对路径(task_dir 内,如 'figures/xxx.png',或用户消息里 "
                    "`[用户上传的参考图]` 行给的路径,或某工具上次返回的 saved 路径)。"
                ),
            },
            "question": {
                "type": "string",
                "description": (
                    "本次任务需要从图里知道什么(强烈建议填写)。如「读出表格中的抗压强度数据」"
                    "「图中仪表读数是多少」「把这页文字完整 OCR 出来」。只问完成用户任务所需内容；"
                    "仅当用户明确要求整图识别时才要求全文 OCR。不传则仅返回简洁画面概览。"
                ),
            },
        },
        "required": ["image"],
    }

    def __init__(
        self,
        *,
        ark_cfg: ArkConfig,
        vision_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,
    ) -> None:
        super().__init__(base_dir, user_root=user_root)
        self.ark_cfg = ark_cfg
        self.cfg = vision_variant_cfg
        self.variant_key = variant_key  # 'seed_2_lite' → usage_events.model_profile = "doubao.seed_2_lite"
        self.working_dir = Path(working_dir)
        self.task_id = task_id
        self.user_id = user_id

    def execute(self, image: str, question: Optional[str] = None) -> str:
        if not (image or "").strip():
            return "[Error] image(图片路径)不能为空"

        cfg = self.cfg
        max_bytes = int(float(cfg.get("max_image_mb", 10)) * 1024 * 1024)
        data_url, disp, err = load_image_as_data_url(
            image.strip(),
            working_dir=self.working_dir,
            user_root=self.user_root,
            display_fn=self._display,
            max_bytes=max_bytes,
        )
        if err:
            return err

        q = (question or "").strip() or _DEFAULT_QUESTION
        model_id = cfg["model_id"]
        timeout_s = float(cfg.get("request_timeout_s", 60))
        endpoint = cfg.get("endpoint", "/chat/completions")

        body: dict[str, Any] = {
            "model": model_id,
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": q},
                        {"type": "image_url", "image_url": {"url": data_url}},
                    ],
                }
            ],
        }

        # 透明重试(超时/网络抖动 tool 内消化,细节见 media_common.ark_chat_with_retry)
        resp, api_err = ark_chat_with_retry(
            self.ark_cfg, endpoint, body,
            timeout_s=timeout_s,
            retries=int(cfg.get("timeout_retries", 1)),
            tool_name="look_at_image",
        )
        if api_err:
            return api_err
        assert resp is not None

        answer, _truncated = extract_chat_answer(resp)
        if not answer:
            return (
                "[Error] vision 响应缺内容(模型未返回文本)。"
                "可能图片格式异常或模型暂不可用,稍后重试。"
            )

        usage = resp.get("usage") or {}
        tin = int(usage.get("prompt_tokens", 0) or 0)
        tout = int(usage.get("completion_tokens", 0) or 0)

        cost = record_usage_safe(
            "look_at_image", record_vision_usage,
            task_id=self.task_id,
            user_id=self.user_id,
            model_profile=f"doubao.{self.variant_key}",
            prompt_tokens=tin,
            completion_tokens=tout,
            input_cny_per_mtoken=float(cfg.get("price_cny_per_mtoken_input", 0)),
            output_cny_per_mtoken=float(cfg.get("price_cny_per_mtoken_output", 0)),
            extra_units={"image": disp},
        )
        cost_cny = float(cost or 0)

        # 第一行 banner(key=value · 分隔,与 seedream/seedance 同协议,便于前端/对账)
        banner = (
            f"[look_at_image] model={model_id} · image={disp} · "
            f"tokens={tin}+{tout} · cost=¥{cost_cny:.4f}"
        )
        # 图片解读正文可能很长(整页 OCR),压一下防爆上下文(保头尾)
        return f"{banner}\n\n{compact_tool_output(answer)}"