script-build-and-framework-tmp/agent/agent/core/runner_images.py

"""Image optimization and prompt-cache support for AgentRunner."""

from __future__ import annotations

from typing import Any, Dict, List, Optional


class RunnerImageRuntime:
    """Own image processing while preserving AgentRunner compatibility hooks."""

    def __init__(self, owner: Any) -> None:
        self._owner = owner

    @property
    def _image_opt_cache(self) -> Dict[str, Dict[str, Any]]:
        return self._owner._image_opt_cache

    @property
    def log(self):
        return self._owner.log

    @property
    def llm_call(self):
        return self._owner.llm_call

    async def _optimize_images(self, messages: List[Dict], model: str) -> List[Dict]:
        """
        分级优化已处理的图片，节省 token

        策略（基于图片距离最后一条 assistant 的"轮次"）：
        1. 最近 1-2 轮：保留原图
        2. 3-5 轮：降低分辨率和压缩（节省 token 但保留视觉信息）
        3. 5 轮以上：调用小模型生成文本描述 + 保留 URL

        处理结果会缓存，避免重复的 PIL 解码/编码和 LLM 调用。

        Args:
            messages: 原始消息列表
            model: 当前使用的模型（用于选择描述生成模型）

        Returns:
            优化后的消息列表（深拷贝）
        """
        if not messages:
            return messages

        # 找到最后一条 assistant message 的位置
        last_assistant_idx = -1
        for i in range(len(messages) - 1, -1, -1):
            if messages[i].get("role") == "assistant":
                last_assistant_idx = i
                break

        # 如果没有 assistant message，说明还没开始对话，不优化
        if last_assistant_idx == -1:
            return messages

        # 统计从每个位置到最后一条 assistant 之间的 assistant 数量（作为"轮次"）
        assistant_count_after = [0] * len(messages)
        count = 0
        for i in range(len(messages) - 1, -1, -1):
            assistant_count_after[i] = count
            if messages[i].get("role") == "assistant":
                count += 1

        # 深拷贝避免修改原始数据
        import copy
        import hashlib
        import asyncio
        import base64 as b64mod
        import httpx
        import mimetypes

        messages = copy.deepcopy(messages)

        # 预处理：将所有 HTTP(S) URL 图片下载并转为 base64 data URL
        # Qwen API 无法访问外部签名 URL（如 BFL、火山引擎 TOS），必须在本地转换
        url_download_jobs = []  # [(msg_idx, block_idx, url)]
        for i, msg in enumerate(messages):
            if msg.get("role") != "tool":
                continue
            content = msg.get("content")
            if not isinstance(content, list):
                continue
            for block_idx, block in enumerate(content):
                if isinstance(block, dict) and block.get("type") == "image_url":
                    url = block.get("image_url", {}).get("url", "")
                    if url.startswith(("http://", "https://")):
                        url_download_jobs.append((i, block_idx, url))

        if url_download_jobs:

            async def _download_image_to_data_url(url: str) -> str | None:
                try:
                    async with httpx.AsyncClient(timeout=60, trust_env=False) as client:
                        resp = await client.get(url)
                        resp.raise_for_status()
                        ct = resp.headers.get("content-type", "").split(";")[0].strip()
                        if not ct.startswith("image/"):
                            ct = (
                                mimetypes.guess_type(url.split("?")[0])[0]
                                or "image/png"
                            )
                        b64 = b64mod.b64encode(resp.content).decode()
                        return f"data:{ct};base64,{b64}"
                except Exception:
                    self.log.debug(
                        "Remote image normalization failed: %s", url, exc_info=True
                    )
                    return None

            results = await asyncio.gather(
                *[_download_image_to_data_url(url) for _, _, url in url_download_jobs],
                return_exceptions=True,
            )
            converted = 0
            for (msg_idx, block_idx, original_url), result in zip(
                url_download_jobs, results
            ):
                if isinstance(result, str) and result.startswith("data:"):
                    messages[msg_idx]["content"][block_idx]["image_url"]["url"] = result
                    converted += 1
            if converted:
                self.log.info(
                    f"[Image Optimization] URL→base64 预转换: {converted}/{len(url_download_jobs)} 张"
                )

        # 统计优化情况
        stats = {"kept": 0, "downscaled": 0, "described": 0, "cache_hit": 0}

        # 收集需要降分辨率或尺寸补齐的图片（用于并发处理）
        process_jobs = []  # [(msg_idx, block_idx, image_url, cache_key, max_size, cache_field)]

        # 第一遍：扫描并收集需要处理的图片
        for i in range(last_assistant_idx):
            msg = messages[i]
            if msg.get("role") != "tool":
                continue

            content = msg.get("content")
            if not isinstance(content, list):
                continue

            rounds_ago = assistant_count_after[i]

            for block_idx, block in enumerate(content):
                if isinstance(block, dict) and block.get("type") == "image_url":
                    image_url_obj = block.get("image_url", {})
                    image_url = image_url_obj.get("url", "")

                    if image_url.startswith("data:"):
                        cache_key = hashlib.md5(image_url[:200].encode()).hexdigest()
                    else:
                        cache_key = hashlib.md5(image_url.encode()).hexdigest()

                    # 1-5 轮都需要检查尺寸
                    if rounds_ago <= 5:
                        cached = self._image_opt_cache.get(cache_key, {})
                        cache_field = "pad_only" if rounds_ago <= 2 else "downscaled"

                        if cache_field not in cached and image_url.startswith("data:"):
                            max_size = None if rounds_ago <= 2 else 512
                            process_jobs.append(
                                (
                                    i,
                                    block_idx,
                                    image_url,
                                    cache_key,
                                    max_size,
                                    cache_field,
                                )
                            )

        # 并发处理所有尺寸任务
        if process_jobs:
            process_results = await asyncio.gather(
                *[
                    self._owner._process_image_size(url, max_size=ms)
                    for _, _, url, _, ms, _ in process_jobs
                ],
                return_exceptions=True,
            )
            for (_, _, _, cache_key, _, cache_field), result in zip(
                process_jobs, process_results
            ):
                if not isinstance(result, Exception) and result is not None:
                    self._image_opt_cache.setdefault(cache_key, {})[cache_field] = (
                        result
                    )

        # 第二遍：应用处理结果
        for i in range(last_assistant_idx):
            msg = messages[i]
            if msg.get("role") != "tool":
                continue

            content = msg.get("content")
            if not isinstance(content, list):
                continue

            # 计算这条消息距离最后一条 assistant 的"轮次"
            rounds_ago = assistant_count_after[i]

            # 处理每个 content block
            new_content = []
            for block in content:
                if isinstance(block, dict) and block.get("type") == "image_url":
                    image_url_obj = block.get("image_url", {})
                    image_url = image_url_obj.get("url", "")

                    # 生成缓存 key（URL 图片用 URL 本身，base64 用前 64 字符 hash）
                    if image_url.startswith("data:"):
                        cache_key = hashlib.md5(image_url[:200].encode()).hexdigest()
                    else:
                        cache_key = hashlib.md5(image_url.encode()).hexdigest()

                    # 根据距离决定处理策略
                    if rounds_ago <= 2:
                        # 最近 1-2 轮：只补齐过小图片，保留原分辨率
                        cached = self._image_opt_cache.get(cache_key, {})
                        if "pad_only" in cached:
                            new_content.append(
                                {
                                    "type": "image_url",
                                    "image_url": {"url": cached["pad_only"]},
                                }
                            )
                            stats["kept"] += 1
                            stats["cache_hit"] += 1
                        elif image_url.startswith("data:"):
                            processed = await self._owner._process_image_size(
                                image_url, max_size=None
                            )
                            if processed:
                                self._image_opt_cache.setdefault(cache_key, {})[
                                    "pad_only"
                                ] = processed
                                new_content.append(
                                    {
                                        "type": "image_url",
                                        "image_url": {"url": processed},
                                    }
                                )
                            else:
                                new_content.append(block)
                            stats["kept"] += 1
                        else:
                            new_content.append(block)
                            stats["kept"] += 1

                    elif rounds_ago <= 5:
                        # 3-5 轮：降低分辨率（优先从缓存取）
                        cached = self._image_opt_cache.get(cache_key, {})
                        if "downscaled" in cached:
                            new_content.append(
                                {
                                    "type": "image_url",
                                    "image_url": {"url": cached["downscaled"]},
                                }
                            )
                            stats["downscaled"] += 1
                            stats["cache_hit"] += 1
                        elif image_url.startswith("data:"):
                            processed = await self._owner._process_image_size(
                                image_url, max_size=512
                            )
                            if processed:
                                # 缓存结果
                                self._image_opt_cache.setdefault(cache_key, {})[
                                    "downscaled"
                                ] = processed
                                new_content.append(
                                    {
                                        "type": "image_url",
                                        "image_url": {"url": processed},
                                    }
                                )
                                stats["downscaled"] += 1
                            else:
                                new_content.append(block)
                                stats["kept"] += 1
                        else:
                            # URL 图片：无法直接处理，保留原图
                            new_content.append(block)
                            stats["kept"] += 1

                    else:
                        # 5 轮以上：生成文本描述（优先从缓存取）
                        cached = self._image_opt_cache.get(cache_key, {})
                        if "description" in cached:
                            new_content.append(cached["description"])
                            stats["described"] += 1
                            stats["cache_hit"] += 1
                        else:
                            description = await self._owner._generate_image_description(
                                image_url, model
                            )
                            url_info = (
                                f" (URL: {image_url[:100]}...)"
                                if not image_url.startswith("data:")
                                else ""
                            )
                            desc_block = {
                                "type": "text",
                                "text": f"[Image description: {description}]{url_info}",
                            }
                            # 缓存结果
                            self._image_opt_cache.setdefault(cache_key, {})[
                                "description"
                            ] = desc_block
                            new_content.append(desc_block)
                            stats["described"] += 1
                else:
                    new_content.append(block)

            msg["content"] = new_content
        if stats["downscaled"] > 0 or stats["described"] > 0:
            self.log.info(
                f"[Image Optimization] 保留 {stats['kept']} 张，"
                f"降分辨率 {stats['downscaled']} 张，"
                f"文本描述 {stats['described']} 张，"
                f"缓存命中 {stats['cache_hit']} 次"
            )

        return messages

    async def _process_image_size(
        self, base64_url: str, max_size: Optional[int] = 512, min_size: int = 11
    ) -> Optional[str]:
        """
        处理 base64 图片的尺寸：
        - 若 max_size 不为 None 且大于该值，则等比例缩放
        - 若任意一边小于 min_size，则补充白边 (Padding)
        """
        try:
            from PIL import Image
            import io
            import base64

            # 解析 base64 数据
            if not base64_url.startswith("data:"):
                return None

            _header, data = base64_url.split(",", 1)

            # 解码图片
            img_data = base64.b64decode(data)
            img = Image.open(io.BytesIO(img_data))

            width, height = img.size

            needs_downscale = max_size is not None and (
                width > max_size or height > max_size
            )
            needs_pad = width < min_size or height < min_size

            # 尺寸正常，无需处理
            if not needs_downscale and not needs_pad:
                return base64_url

            new_width, new_height = width, height

            # 1. 降分辨率
            if needs_downscale:
                if width > height:
                    new_width = max_size
                    new_height = int(height * max_size / width)
                else:
                    new_height = max_size
                    new_width = int(width * max_size / height)

            if (new_width, new_height) != (width, height):
                img_resized = img.resize(
                    (new_width, new_height), Image.Resampling.BILINEAR
                )
            else:
                img_resized = img

            # 2. 补齐白边 (Padding)
            pad_width = max(new_width, min_size)
            pad_height = max(new_height, min_size)

            if pad_width > new_width or pad_height > new_height:
                # 创建白色背景
                padded_img = Image.new(
                    "RGBA" if img_resized.mode in ("RGBA", "P") else "RGB",
                    (pad_width, pad_height),
                    (255, 255, 255, 255),
                )
                offset_x = (pad_width - new_width) // 2
                offset_y = (pad_height - new_height) // 2
                padded_img.paste(img_resized, (offset_x, offset_y))
                img_resized = padded_img

            # 转换为 RGB（JPEG不支持 RGBA, P 等具有透明度或索引的模式）
            if img_resized.mode != "RGB":
                if img_resized.mode == "RGBA" or img_resized.mode == "P":
                    # Create a white background for transparent images
                    background = Image.new("RGB", img_resized.size, (255, 255, 255))
                    if img_resized.mode == "P" and "transparency" in img_resized.info:
                        img_resized = img_resized.convert("RGBA")
                    if img_resized.mode == "RGBA":
                        background.paste(img_resized, mask=img_resized.split()[3])
                        img_resized = background
                img_resized = img_resized.convert("RGB")

            # 重新编码为 JPEG（如果只是补齐没有缩放，可以稍微保留高点质量）
            buffer = io.BytesIO()
            quality = 60 if needs_downscale else 85
            img_resized.save(buffer, format="JPEG", quality=quality, optimize=False)
            new_data = base64.b64encode(buffer.getvalue()).decode("utf-8")

            return f"data:image/jpeg;base64,{new_data}"

        except Exception as e:
            self.log.warning(f"[Image Process] 处理图片尺寸失败: {e}")
            return None

    async def _generate_image_description(
        self, image_url: str, current_model: str
    ) -> str:
        """
        使用小模型生成图片的文本描述

        Args:
            image_url: 图片 URL（base64 或 http(s)）
            current_model: 当前使用的模型

        Returns:
            图片描述文本
        """
        # The parameter remains part of AgentRunner's compatibility hook; the
        # description itself always uses the dedicated vision model below.
        del current_model
        try:
            # 使用 qwen-vl-max（通义千问视觉模型）生成描述
            # 注意：qwen-vl 系列专门支持视觉输入
            description_model = "qwen-vl-max"

            # 构建描述请求
            messages = [
                {
                    "role": "user",
                    "content": [
                        {"type": "image_url", "image_url": {"url": image_url}},
                        {
                            "type": "text",
                            "text": "请用 1-2 句话简洁描述这张图片的主要内容。",
                        },
                    ],
                }
            ]

            # 调用 LLM
            result = await self.llm_call(
                messages=messages,
                model=description_model,
                tools=None,
                temperature=0.3,
            )

            description = result.get("content", "").strip()
            return description if description else "图片内容"

        except Exception as e:
            self.log.warning(f"[Image Description] 生成描述失败: {e}")
            return "图片内容"

    def _add_cache_control(
        self, messages: List[Dict], model: str, enable: bool
    ) -> List[Dict]:
        """
        为支持的模型添加 Prompt Caching 标记

        策略：固定位置 + 延迟缓存
        1. 如果有未处理的图片（最后一条 assistant 之后的 tool messages 中有图片），跳过缓存
        2. system message 添加缓存（如果足够长）
        3. 固定位置缓存点（20, 40, 60, 80），确保每个缓存点间隔 >= 1024 tokens
        4. 最多使用 4 个缓存点（含 system）

        Args:
            messages: 原始消息列表
            model: 模型名称
            enable: 是否启用缓存

        Returns:
            添加了 cache_control 的消息列表（深拷贝）
        """
        if not enable:
            return messages

        # 只对 Claude 模型启用
        if "claude" not in model.lower():
            return messages

        # 延迟缓存：检查是否有未处理的图片
        last_assistant_idx = -1
        for i in range(len(messages) - 1, -1, -1):
            if messages[i].get("role") == "assistant":
                last_assistant_idx = i
                break

        # 检查最后一条 assistant 之后是否有包含图片的 tool messages
        has_unprocessed_images = False
        if last_assistant_idx >= 0:
            for i in range(last_assistant_idx + 1, len(messages)):
                msg = messages[i]
                if msg.get("role") == "tool":
                    content = msg.get("content")
                    if isinstance(content, list):
                        has_unprocessed_images = any(
                            isinstance(block, dict) and block.get("type") == "image_url"
                            for block in content
                        )
                        if has_unprocessed_images:
                            break

        if has_unprocessed_images:
            self.log.debug("[Cache] 检测到未处理的图片，延迟缓存建立")
            return messages

        # 深拷贝避免修改原始数据
        import copy

        messages = copy.deepcopy(messages)

        # 策略 1: 为 system message 添加缓存
        system_cached = False
        for msg in messages:
            if msg.get("role") == "system":
                content = msg.get("content", "")
                if isinstance(content, str) and len(content) > 1000:
                    msg["content"] = [
                        {
                            "type": "text",
                            "text": content,
                            "cache_control": {"type": "ephemeral"},
                        }
                    ]
                    system_cached = True
                    self.log.debug(
                        f"[Cache] 为 system message 添加缓存标记 (len={len(content)})"
                    )
                break

        # 策略 2: 固定位置缓存点
        CACHE_INTERVAL = 20
        MAX_POINTS = 3 if system_cached else 4
        MIN_TOKENS = 1024
        AVG_TOKENS_PER_MSG = 70

        total_msgs = len(messages)
        if total_msgs == 0:
            return messages

        cache_positions = []
        last_cache_pos = 0

        for i in range(1, MAX_POINTS + 1):
            target_pos = i * CACHE_INTERVAL - 1  # 19, 39, 59, 79

            if target_pos >= total_msgs:
                break

            # 从目标位置开始查找合适的 user/assistant 消息
            for j in range(target_pos, total_msgs):
                msg = messages[j]

                if msg.get("role") not in ("user", "assistant"):
                    continue

                content = msg.get("content", "")
                if not content:
                    continue

                # 检查 content 是否非空
                is_valid = False
                if isinstance(content, str):
                    is_valid = len(content) > 0
                elif isinstance(content, list):
                    is_valid = any(
                        isinstance(block, dict)
                        and block.get("type") == "text"
                        and len(block.get("text", "")) > 0
                        for block in content
                    )

                if not is_valid:
                    continue

                # 检查 token 距离
                msg_count = j - last_cache_pos
                estimated_tokens = msg_count * AVG_TOKENS_PER_MSG

                if estimated_tokens >= MIN_TOKENS:
                    cache_positions.append(j)
                    last_cache_pos = j
                    self.log.debug(
                        f"[Cache] 在位置 {j} 添加缓存点 (估算 {estimated_tokens} tokens)"
                    )
                    break

        # 应用缓存标记
        for idx in cache_positions:
            msg = messages[idx]
            content = msg.get("content", "")

            if isinstance(content, str):
                msg["content"] = [
                    {
                        "type": "text",
                        "text": content,
                        "cache_control": {"type": "ephemeral"},
                    }
                ]
                self.log.debug(
                    f"[Cache] 为 message[{idx}] ({msg.get('role')}) 添加缓存标记"
                )
            elif isinstance(content, list):
                # 在最后一个 text block 添加 cache_control
                for block in reversed(content):
                    if isinstance(block, dict) and block.get("type") == "text":
                        block["cache_control"] = {"type": "ephemeral"}
                        self.log.debug(
                            f"[Cache] 为 message[{idx}] ({msg.get('role')}) 添加缓存标记"
                        )
                        break

        self.log.debug(
            f"[Cache] 总消息: {total_msgs}, "
            f"缓存点: {len(cache_positions)} at {cache_positions}"
        )
        return messages


__all__ = ["RunnerImageRuntime"]