Agent/agent/tools/builtin/file/read_images.py

"""
Read Images Tool - 批量读取图片工具

为"批量读取 + 多图分析"场景设计的工具，与单文件的 read_file 分工：
- read_file: 单个文件（文本 / PDF / 单张图片）
- read_images: 2 张以上图片，支持网格拼图和降采样

核心能力：
1. 并发批量加载本地路径或 URL
2. 自动降采样，防止 token 爆炸
3. 可选拼图（grid 模式），把 N 张图合成一张带索引编号的网格图，
   适合 LLM 横向对比、选图、批量判断场景
4. 自适应布局 + 硬上限，保证拼图即使经过 LLM 内部缩放也能保持可辨
"""

from typing import Any, Dict, List, Literal, Optional, Tuple

from agent.tools import tool, ToolResult, ToolContext
from agent.tools.utils.image import (
    build_image_grid,
    downscale,
    encode_base64,
    load_images,
)


# Grid 模式的硬上限：超过此数量必须分批调用
# 理由：12 张可排成 4x3 网格，每格 ~320px，人物/场景细节清晰可辨。
# 再多格子就太小，分辨不出内容，失去对比价值。
MAX_GRID_IMAGES = 12


def _adaptive_layout(count: int) -> Tuple[int, int]:
    """根据图片数量自动选择 (columns, thumb_size)。

    目标：拼图最终边长不超过 ~1400px，同时每格缩略图保持 >= 320px 以保证可辨认。

    Returns:
        (columns, thumb_size)
    """
    if count <= 2:
        return 2, 500   # 2x1
    if count <= 4:
        return 2, 450   # 2x2
    if count <= 6:
        return 3, 400   # 3x2
    if count <= 9:
        return 3, 380   # 3x3
    # 10-12
    return 4, 320       # 4x3


@tool(
    description="批量读取多张图片，支持自动降采样和网格拼图（用于横向对比/选图场景）",
    hidden_params=["context"],
    display={
        "zh": {
            "name": "批量读取图片",
            "params": {
                "paths": "图片路径列表",
                "layout": "布局模式",
                "max_dimension": "每张图最大边长",
            },
        },
        "en": {
            "name": "Read Images",
            "params": {
                "paths": "Image paths",
                "layout": "Layout mode",
                "max_dimension": "Max dimension per image",
            },
        },
    },
    groups=["core"],
)
async def read_images(
    paths: List[str],
    layout: Literal["grid", "separate"] = "grid",
    max_dimension: int = 1024,
    context: Optional[ToolContext] = None,
) -> ToolResult:
    """批量读取图片并返回给 LLM，支持自动降采样和网格拼图

    为 **2 张以上** 的图片批量分析场景设计。单张图片请用 `read_file`。

    ⚠️ **grid 模式最多 12 张**。超过请分批调用：第一次传前 12 张，第二次传后续，
    以此类推。再多每格就太小，分辨不出内容。

    两种布局模式：

    - **grid**（默认）：把所有图片拼成一张只带索引编号的网格图（1,2,3…）。
      LLM 只看到 1 张拼图，大幅减少结构开销 token。索引对应的原始路径见
      返回文本的对照表，LLM 可以用"第 3 张"来引用具体图片。
      **自适应布局**：根据图片数量自动选择列数和缩略图尺寸，小批量时每张图更清晰：
        * 1-2 张：2 列 × 500px
        * 3-4 张：2 列 × 450px
        * 5-6 张：3 列 × 400px
        * 7-9 张：3 列 × 380px
        * 10-12 张：4 列 × 320px
      适合：从多张候选图中挑选、横向对比质量/风格、批量判断。

    - **separate**：把每张图独立返回（仍然降采样）。无数量限制，但每张图都有
      独立的结构开销 token。适合：
        * 需要逐张做独立的精细分析
        * 每张图之间没有对比关系

    自动降采样：无论哪种模式，每张图都会先降采样到 max_dimension（默认 1024px）
    的最大边长，防止高分辨率图片炸掉 token 预算。

    Args:
        paths: 图片路径列表，支持本地路径和 HTTP(S) URL，可混用。
               grid 模式下不超过 12 张，超过必须分批调用。
        layout: 布局模式，"grid" 拼图（默认）/ "separate" 多张独立
        max_dimension: 每张图的最大边长（等比降采样到不超过此值），默认 1024
        context: 工具上下文（框架注入，无需手动传）

    Returns:
        ToolResult：images 字段包含图片数据（grid 模式 1 张拼图，separate 模式 N 张），
        output 字段包含每张图的索引和来源路径对照表
    """
    if not paths:
        return ToolResult(
            title="批量读图失败",
            output="",
            error="paths 不能为空",
        )

    # 硬上限检查（仅对 grid 模式）
    if layout == "grid" and len(paths) > MAX_GRID_IMAGES:
        return ToolResult(
            title="批量读图失败",
            output="",
            error=(
                f"grid 模式最多支持 {MAX_GRID_IMAGES} 张图片，当前传入 {len(paths)} 张。"
                f"请分批调用：每次最多 {MAX_GRID_IMAGES} 张。"
                f"或者使用 layout='separate' 模式（无数量限制但 token 开销更高）。"
            ),
        )

    if len(paths) == 1:
        hint = "（只有 1 张图片，建议用 read_file 更合适）"
    else:
        hint = ""

    # 1. 并发加载所有图片
    loaded = await load_images(paths)

    # 2. 分离成功和失败
    successes: List[tuple] = []  # [(path, PIL.Image), ...]
    failures: List[str] = []     # [path, ...]
    for source, img in loaded:
        if img is None:
            failures.append(source)
        else:
            successes.append((source, img))

    if not successes:
        return ToolResult(
            title="批量读图失败",
            output="",
            error=f"所有 {len(paths)} 张图片均加载失败",
            metadata={"failed": failures},
        )

    # 3. 每张图降采样
    processed = [(src, downscale(img, max_dimension)) for src, img in successes]

    # 4. 构建索引 → 路径对照表（用完整路径，方便 LLM 后续引用或调用）
    index_lines = [f"{i}. {src}" for i, (src, _) in enumerate(processed, start=1)]
    summary_parts = [f"共加载 {len(processed)}/{len(paths)} 张图片"]
    if hint:
        summary_parts.append(hint)
    if failures:
        summary_parts.append(f"，失败 {len(failures)} 张")
    summary = "".join(summary_parts)

    output_lines = [summary, ""] + index_lines
    if failures:
        output_lines.append("")
        output_lines.append("加载失败的路径：")
        output_lines.extend(f"  - {p}" for p in failures)
    output_text = "\n".join(output_lines)

    # 5. 根据 layout 生成 images 字段
    images_for_llm = []
    if layout == "grid":
        cols, thumb_size = _adaptive_layout(len(processed))
        # 网格只显示序号，不写文件名 —— 索引对应的路径见上方 output 文本
        grid = build_image_grid(
            images=[img for _, img in processed],
            labels=None,
            columns=cols,
            thumb_size=thumb_size,
        )
        # 网格拼图固定用 JPEG 节省 token
        b64, media_type = encode_base64(grid, format="JPEG", quality=80)
        images_for_llm.append({
            "type": "base64",
            "media_type": media_type,
            "data": b64,
        })
    else:  # separate
        for _, img in processed:
            b64, media_type = encode_base64(img, format="JPEG", quality=80)
            images_for_llm.append({
                "type": "base64",
                "media_type": media_type,
                "data": b64,
            })

    return ToolResult(
        title=f"批量读图成功（{layout} 模式，{len(processed)} 张）",
        output=output_text,
        long_term_memory=f"Read {len(processed)} images via {layout} layout",
        images=images_for_llm,
        metadata={
            "count": len(processed),
            "failed_count": len(failures),
            "layout": layout,
        },
    )


# ── CLI 入口：图片拼图工具 ──
#
# 这个 CLI 的语义是**拼图工具**，不是"读图工具"——Claude Code 这样的调用方
# 本身就能读单张图（用 Read 工具），真正稀缺的能力是把 N 张图合成一张
# 带索引编号的网格图，让一次 Read 就能横向对比多张。
#
# 因此 CLI 只支持 grid 模式；如果你需要单张图，直接用 Read 工具即可。
#
# 用法：
#   python agent/tools/builtin/file/read_images.py --out=<path> <img1> <img2> ...
#
# 必填参数：
#   --out=/path/grid.jpg     拼图保存路径（必须显式指定，避免污染 /tmp）
#
# 可选参数：
#   --max_dimension=1024     每张图预先降采样的最大边长（默认 1024）
#
# 示例：
#   python agent/tools/builtin/file/read_images.py \
#     --out=/tmp/compare.jpg \
#     ~/Downloads/a.jpg ~/Downloads/b.jpg ~/Downloads/c.jpg
#
# 输出：一行 JSON，包含 out_path、index_map（索引→原始路径对照表）、
# text（文字摘要）。调用方拿到 out_path 后用 Read 工具查看拼图即可。

if __name__ == "__main__":
    import base64
    import json
    import sys
    from pathlib import Path as _Path

    def _print_usage():
        print("用法: python read_images.py --out=<path> <img1> <img2> ...")
        print("     --out=/path/grid.jpg   拼图输出路径（必填）")
        print("     --max_dimension=1024   每张图最大边长（可选，默认 1024）")
        print(f"最多 {MAX_GRID_IMAGES} 张图片")

    if len(sys.argv) < 2 or sys.argv[1] in ("-h", "--help"):
        _print_usage()
        sys.exit(0)

    # 解析参数
    cli_paths: List[str] = []
    cli_out: Optional[str] = None
    cli_max_dim: int = 1024
    for arg in sys.argv[1:]:
        if arg.startswith("--") and "=" in arg:
            k, v = arg.split("=", 1)
            k = k.lstrip("-").replace("-", "_")
            if k == "out":
                cli_out = v
            elif k == "max_dimension":
                cli_max_dim = int(v)
            else:
                print(f"警告: 未知参数 {k}", file=sys.stderr)
        else:
            cli_paths.append(arg)

    if not cli_paths:
        print("错误: 至少提供一个图片路径", file=sys.stderr)
        _print_usage()
        sys.exit(1)

    if not cli_out:
        print("错误: 必须显式指定 --out=<path>", file=sys.stderr)
        _print_usage()
        sys.exit(1)

    import asyncio
    result = asyncio.run(read_images(
        paths=cli_paths,
        layout="grid",
        max_dimension=cli_max_dim,
    ))

    if result.error:
        print(json.dumps({"error": result.error}, ensure_ascii=False, indent=2))
        sys.exit(1)

    # 写入拼图文件
    out_p = _Path(cli_out)
    out_p.parent.mkdir(parents=True, exist_ok=True)
    out_p.write_bytes(base64.b64decode(result.images[0]["data"]))

    # 解析索引 → 原始路径对照表
    index_map: List[Dict[str, Any]] = []
    for line in result.output.split("\n"):
        if line and line[0].isdigit() and ". " in line:
            idx_str, src = line.split(". ", 1)
            if idx_str.isdigit():
                index_map.append({"index": int(idx_str), "source": src})

    print(json.dumps({
        "out_path": str(out_p.resolve()),
        "count": result.metadata.get("count", 0) if result.metadata else 0,
        "index_map": index_map,
        "text": result.output,
    }, ensure_ascii=False, indent=2))