inquiry_robot/inquiry-agent/agent/handlers/image_inquiry.py

"""
意图：私聊图片询价。

本文件职责：收图 → 回「正在识别」→ 千问抄可见原文 → 拼口播 → 新开后复用文字询价。
禁止：千问填最终缺项；另写海运/空运/陆运图片专线；回调线程同步等千问；群图走本 Handler。
"""

from __future__ import annotations

import logging
import time
from typing import Any, Callable, Optional

from agent.channel.queue import get_message_store
from agent.channel.wecom.models import InboundMessage
from agent.policy import inquiry_copy as copy
from agent.policy.image_wave import get_image_wave_book

logger = logging.getLogger(__name__)

VisionFn = Callable[..., dict[str, Any]]
DownloadFn = Callable[[str], tuple[bytes, str]]
PutFn = Callable[..., dict[str, Any]]


def drop_all_private_bookmarks(sender_id: str) -> None:
    """三线当前书签都丢掉；已出号仍能按工单号找回。"""
    from agent.policy.air_text_flow import get_air_text_flow
    from agent.policy.land_text_flow import get_land_text_flow
    from agent.policy.sea_text_flow import get_sea_text_flow

    get_air_text_flow().drop_current_bookmark(sender_id)
    get_sea_text_flow().drop_current_bookmark(sender_id)
    get_land_text_flow().drop_current_bookmark(sender_id)


def is_active_image_wave(sender_id: str) -> bool:
    """业务结果尚未发出的图片波次。"""
    return get_image_wave_book().active(sender_id) is not None


def append_wave_text(sender_id: str, text: str) -> bool:
    """认图完成前的字并进本波，不走文字询价。"""
    return get_image_wave_book().append_follow_text(sender_id, text)


def note_private_text(sender_id: str, text: str, *, now: Optional[float] = None) -> None:
    """文字询价记下最近一句，供 10 秒内发图带走。"""
    get_image_wave_book().note_text(sender_id, text, now=now if now is not None else time.time())


def handle_image_inquiry(
    message: InboundMessage,
    *,
    store=None,
    vision_fn: Optional[VisionFn] = None,
    download_fn: Optional[DownloadFn] = None,
    put_fn: Optional[PutFn] = None,
    process_now: bool = True,
) -> str:
    """
    私聊一张图。inbox 线程只开波、回识别中；认图在 Worker 或单测 vision_fn。

    须已有 sender_id。群消息不要调用本函数。
    """
    if not message.has_sender():
        logger.warning("image_inquiry 拒绝：无 sender_id")
        return "reject_no_sender"
    if (message.chat_type or "single") == "group":
        logger.info("image_inquiry 忽略群图 chat=%s", message.chat_id)
        return "ignore_group_image"

    book = get_image_wave_book()
    msg_store = store or get_message_store()
    media_id = str((message.media or {}).get("media_id") or "")
    filename = str((message.media or {}).get("filename") or "image.jpg")
    wave, how = book.open_or_join(message.sender_id, media_id, now=time.time())
    if how in {"opened", "new_after_done"}:
        drop_all_private_bookmarks(message.sender_id)
        msg_store.enqueue_outbound(
            touser=message.sender_id,
            content=copy.image_recognizing(),
            dedupe_key=f"img-ack:{wave.wave_id}",
        )
        book.mark_ack(message.sender_id)
        logger.info("image_inquiry 开波 how=%s wave=%s", how, wave.wave_id)
    else:
        logger.info(
            "image_inquiry 并入 how=%s wave=%s pending=%s",
            how,
            wave.wave_id,
            wave.pending,
        )

    if process_now and vision_fn is not None:
        _recognize_one(
            sender_id=message.sender_id,
            media_id=media_id,
            filename=filename,
            vision_fn=vision_fn,
            download_fn=download_fn,
            put_fn=put_fn,
        )
        wave = book.active(message.sender_id) or wave
        if wave.pending <= 0:
            return finish_wave(message.sender_id, store=msg_store, seed_message=message)
        return "recognizing"

    if vision_fn is None:
        try:
            from agent.jobs.recognition import enqueue_recognition

            enqueue_recognition(
                kind="vision",
                payload={
                    "sender_id": message.sender_id,
                    "media_id": media_id,
                    "filename": filename,
                    "message_id": message.message_id,
                    "wave_id": wave.wave_id,
                },
                idempotency_key=message.message_id or media_id,
            )
        except Exception:
            logger.exception("image_inquiry 入识别队失败 sender=%s", message.sender_id)
            _recognize_one(
                sender_id=message.sender_id,
                media_id=media_id,
                filename=filename,
                vision_fn=None,
                download_fn=download_fn,
                put_fn=put_fn,
            )
            return finish_wave(message.sender_id, store=msg_store, seed_message=message)
        return "recognizing"

    return "recognizing"


def flush_active_wave(
    sender_id: str,
    *,
    store=None,
    vision_fn: Optional[VisionFn] = None,
    download_fn: Optional[DownloadFn] = None,
    put_fn: Optional[PutFn] = None,
    seed_message: Optional[InboundMessage] = None,
) -> str:
    """单测：两张图都 join 后再认，只出一次业务结果。"""
    book = get_image_wave_book()
    wave = book.active(sender_id)
    if wave is None:
        return "no_wave"
    pending_ids = list(wave.media_ids[-wave.pending :]) if wave.pending else []
    for mid in pending_ids:
        _recognize_one(
            sender_id=sender_id,
            media_id=mid,
            filename="",
            vision_fn=vision_fn,
            download_fn=download_fn,
            put_fn=put_fn,
        )
    return finish_wave(sender_id, store=store, seed_message=seed_message)


def finish_queued_image(payload: dict[str, Any]) -> dict[str, Any]:
    """
    Worker 识别槽：认完一张，pending 为 0 则出业务结果。

    payload.__vision / __bytes 仅单测注入，生产不得依赖。
    """
    sender_id = str(payload.get("sender_id") or "")
    media_id = str(payload.get("media_id") or "")
    filename = str(payload.get("filename") or "image.jpg")
    vision_hook = payload.get("__vision")
    raw = payload.get("__bytes")

    def _vision(**_kwargs: Any) -> dict[str, Any]:
        if isinstance(vision_hook, dict):
            return dict(vision_hook)
        from agent.config import get_settings
        from agent.llm.mode_vision import invoke_vision

        settings = get_settings()
        return invoke_vision(
            payload=_kwargs.get("payload") or {},
            allow_network=bool(settings.llm_allow_network and settings.llm_data_usage_confirmed),
        )

    def _download(mid: str) -> tuple[bytes, str]:
        if isinstance(raw, (bytes, bytearray)):
            return bytes(raw), ""
        from agent.channel.outbox.sender import WeComAppClient
        from agent.config import get_settings

        settings = get_settings()
        client = WeComAppClient(
            corp_id=settings.wecom_corp_id,
            secret=settings.wecom_secret,
            agent_id=int(settings.wecom_agent_id or "0") or 0,
        )
        return client.download_media(mid)

    _recognize_one(
        sender_id=sender_id,
        media_id=media_id,
        filename=filename,
        vision_fn=_vision,
        download_fn=_download,
        put_fn=None,
    )
    book = get_image_wave_book()
    wave = book.active(sender_id)
    if wave is None:
        logger.warning(
            "image_inquiry Worker 认完但无活跃波次 sender=%s media=%s（多半 HTTP/Worker 账不一致）",
            sender_id,
            media_id,
        )
        return {"ok": False, "phase": "no_active_wave", "pending": 0}
    if wave.pending > 0:
        return {"ok": True, "phase": "recognizing", "pending": wave.pending}
    phase = finish_wave(sender_id, seed_message=InboundMessage(
        sender_id=sender_id,
        message_id=str(payload.get("message_id") or media_id),
        content="",
        msg_type="image",
        media={"media_id": media_id, "filename": filename},
    ))
    return {"ok": True, "phase": phase}


def _recognize_one(
    *,
    sender_id: str,
    media_id: str,
    filename: str,
    vision_fn: Optional[VisionFn],
    download_fn: Optional[DownloadFn],
    put_fn: Optional[PutFn],
) -> None:
    """下载、可选落对象存储、读图，写入波次。"""
    import base64

    from agent.llm.mode_vision import sniff_image_mime

    book = get_image_wave_book()
    err = ""
    raw = b""
    if download_fn is None:
        raw, err = b"", "no_download"
    else:
        raw, err = download_fn(media_id)
    if not raw:
        logger.info(
            "image_inquiry 无字节 sender=%s err=%s",
            sender_id,
            err or "media_empty",
        )
        book.add_vision(
            sender_id,
            "",
            object_key="",
            filename=filename,
            error=err or "media_empty",
        )
        return
    mime = sniff_image_mime(raw) or "image/jpeg"
    object_key = ""
    if put_fn is not None:
        key = f"inquiry-image/{sender_id}/{media_id or filename}"
        try:
            stored = put_fn(key=key, data=raw, content_type=mime)
            if stored.get("ok"):
                object_key = str(stored.get("key") or key)
        except Exception:
            logger.exception("原图落存储失败 sender=%s", sender_id)
    else:
        object_key = f"inquiry-image/{sender_id}/{media_id or filename}"
        try:
            from agent.config import get_settings
            from agent.storage.s3_client import S3ObjectStorage

            S3ObjectStorage.from_settings(get_settings(), allow_network=False).put_bytes(
                key=object_key, data=raw, content_type=mime
            )
        except Exception:
            logger.info("原图 dry_run 存储跳过 sender=%s", sender_id)

    vision: dict[str, Any] = {"ok": False, "text": "", "empty": True, "error": err or "no_vision"}
    if vision_fn is not None:
        b64 = base64.b64encode(raw).decode("ascii")
        vision = vision_fn(payload={"image_b64": b64, "mime": mime, "media_id": media_id}) or vision
    else:
        from agent.config import get_settings
        from agent.llm.mode_vision import invoke_vision

        settings = get_settings()
        vision = invoke_vision(
            payload={"image_b64": base64.b64encode(raw).decode("ascii"), "mime": mime},
            allow_network=bool(settings.llm_allow_network and settings.llm_data_usage_confirmed),
        )
    text = str(vision.get("text") or "")
    if not vision.get("ok"):
        err = str(vision.get("error") or err or "vision_failed")
    logger.info(
        "image_inquiry 千问 bytes=%s mime=%s ok=%s empty=%s text_len=%s err=%s",
        len(raw),
        mime,
        bool(vision.get("ok")),
        bool(vision.get("empty")),
        len(text.strip()),
        err or "-",
    )
    book.add_vision(
        sender_id,
        text,
        object_key=object_key,
        filename=filename,
        error=err,
    )


def finish_wave(
    sender_id: str,
    *,
    store=None,
    seed_message: Optional[InboundMessage] = None,
) -> str:
    """
    本波图都认完：拼口播、当文字询价、必要时把说明和补问合成一条。
    """
    from agent.handlers.text_inquiry import handle_text_inquiry
    from agent.policy.air_text_flow import get_air_text_flow
    from agent.policy.land_text_flow import get_land_text_flow
    from agent.policy.sea_text_flow import get_sea_text_flow

    book = get_image_wave_book()
    wave = book.active(sender_id)
    oral = book.compose_oral(sender_id)
    reason = ""
    if not oral.strip():
        reason = copy.image_tech_fail() if (wave and wave.errors) else copy.IMAGE_UNREADABLE

    msg_store = store or get_message_store()
    prefix_used = {"v": False}

    class _StoreProxy:
        def enqueue_outbound(self, **kwargs):
            body = str(kwargs.get("content") or "")
            extra = dict(kwargs.get("payload") or {})
            if reason and not prefix_used["v"] and body:
                body = copy.image_fail_then_clarify(reason, body)
                kwargs["content"] = body
                prefix_used["v"] = True
            return msg_store.enqueue_outbound(
                touser=kwargs.get("touser") or sender_id,
                content=body,
                dedupe_key=str(kwargs.get("dedupe_key") or ""),
                payload=extra,
            )

    inbound = InboundMessage(
        sender_id=sender_id,
        message_id=(seed_message.message_id if seed_message else f"img:{sender_id}:{int(time.time())}"),
        content=oral,
        msg_type="text",
        media=dict(seed_message.media) if seed_message else {},
        raw=dict(seed_message.raw) if seed_message else {},
    )
    phase = handle_text_inquiry(inbound, store=_StoreProxy())
    if reason and not prefix_used["v"]:
        msg_store.enqueue_outbound(
            touser=sender_id,
            content=copy.image_fail_then_clarify(reason, copy.ask_transport()),
            dedupe_key=f"img-empty:{inbound.message_id}",
        )
        phase = "need_mode"
    book.mark_business_done(sender_id)
    book.forget_text(sender_id)
    sess = (
        get_sea_text_flow().session_of(sender_id)
        or get_air_text_flow().session_of(sender_id)
        or get_land_text_flow().session_of(sender_id)
    )
    wo = (sess.work_order_no if sess else "") or ""
    if (
        not wo
        and sess
        and wave
        and (sess.business_line or "").upper() == "SEA"
        and wave.object_keys
    ):
        # 海运先出核对卡，工单要等确定。图片键先留在书签上。
        files = []
        for i, key in enumerate(wave.object_keys):
            name = wave.filenames[i] if i < len(wave.filenames) else "image.jpg"
            if not key:
                continue
            files.append(
                {
                    "file_name": name or "image.jpg",
                    "object_key": key,
                    "content_type": "image/jpeg",
                }
            )
        for engine in (get_sea_text_flow(), get_air_text_flow(), get_land_text_flow()):
            hit = engine.session_of(sender_id)
            if hit is sess:
                engine.remember_pending_files(sender_id, files)
                break
    if wo and wave:
        ledger = None
        for engine in (get_sea_text_flow(), get_air_text_flow(), get_land_text_flow()):
            hit = engine.session_of(sender_id)
            if hit and hit.work_order_no == wo:
                ledger = engine.ledger
                break
        if ledger is not None:
            for i, key in enumerate(wave.object_keys):
                name = wave.filenames[i] if i < len(wave.filenames) else "image.jpg"
                try:
                    ledger.attach_file(
                        work_order_no=wo,
                        file_name=name or "image.jpg",
                        object_key=key,
                        content_type="image/jpeg",
                        sender_id=sender_id,
                    )
                except Exception:
                    logger.exception("挂图片附件失败 wo=%s", wo)
    logger.info("image_inquiry 波次完成 phase=%s oral_len=%s", phase, len(oral))
    return phase