zcbot/core/attachments.py

"""Structured user-message attachments and model-context compatibility rendering."""
from __future__ import annotations

import base64
import mimetypes
from copy import deepcopy
from collections.abc import Iterable
from pathlib import Path
from typing import Any

from .artifacts import ArtifactPathError, resolve_artifact_path

ATTACHMENT_REF_VERSION = 1
MAX_ATTACHMENTS_PER_MESSAGE = 10
MAX_NATIVE_IMAGE_BYTES = 10 * 1024 * 1024
_NATIVE_IMAGE_MIME = {
    ".gif": "image/gif",
    ".jpeg": "image/jpeg",
    ".jpg": "image/jpeg",
    ".png": "image/png",
    ".webp": "image/webp",
}
_IMAGE_EXTENSIONS = {
    ".avif", ".bmp", ".gif", ".heic", ".heif", ".jpeg", ".jpg",
    ".png", ".svg", ".tif", ".tiff", ".webp",
}


def _attachment_kind(requested: str, path: Path, media_type: str) -> str:
    detected_image = media_type.startswith("image/") or path.suffix.lower() in _IMAGE_EXTENSIONS
    if requested == "image" and detected_image:
        return "image"
    return "file"


def normalize_attachment_refs(
    refs: Iterable[Any], *, working_dir: Path, user_root: Path,
) -> list[dict]:
    """Validate client refs and return task-relative, display-ready metadata."""
    output: list[dict] = []
    seen: set[str] = set()
    for raw in refs:
        if len(output) >= MAX_ATTACHMENTS_PER_MESSAGE:
            raise ArtifactPathError(
                f"at most {MAX_ATTACHMENTS_PER_MESSAGE} attachments are allowed"
            )
        data = raw.model_dump() if hasattr(raw, "model_dump") else dict(raw or {})
        absolute, rel = resolve_artifact_path(
            str(data.get("path") or ""),
            working_dir=working_dir,
            user_root=user_root,
            require_file=True,
        )
        if rel in seen:
            continue
        seen.add(rel)
        media_type = mimetypes.guess_type(absolute.name)[0] or "application/octet-stream"
        label = str(data.get("label") or absolute.name).strip() or absolute.name
        output.append({
            "version": ATTACHMENT_REF_VERSION,
            "scope": "working_dir",
            "path": rel,
            "label": label,
            "kind": _attachment_kind(str(data.get("kind") or "file"), absolute, media_type),
            "media_type": media_type,
            "size_bytes": absolute.stat().st_size,
        })
    return output


def content_for_model(content: str, refs: Iterable[dict], working_dir_name: str) -> str:
    """Append the legacy-readable file hints only to the provider-bound message."""
    lines: list[str] = []
    wd = str(working_dir_name or "").strip().strip("/\\")
    for ref in refs or ():
        path = str((ref or {}).get("path") or "").strip().replace("\\", "/")
        if not path:
            continue
        full_path = f"{wd}/{path}" if wd else path
        marker = "[用户上传的参考图]" if (ref or {}).get("kind") == "image" else "[用户上传的文件]"
        lines.append(f"{marker} {full_path}")
    text = str(content or "").strip()
    if not lines:
        return text
    suffix = "\n".join(lines)
    return f"{text}\n\n{suffix}" if text else suffix


def materialize_native_images(
    messages: Iterable[dict],
    *,
    enabled: bool,
    working_dir: Path,
    user_root: Path,
) -> tuple[list[dict], int]:
    """把内存附件引用按需物化为 OpenAI-compatible image_url blocks。

    Base64 只存在于本次 provider-bound 副本，不写数据库。历史文件已删除、越界、
    格式不受支持或超过上限时保留原有文字路径提示并跳过图片块，不能让旧附件阻断续聊。
    所有 ``_`` 内部字段在返回前剥离。
    """
    output: list[dict] = []
    image_count = 0
    for message in messages:
        new_msg = deepcopy(message)
        refs = list(new_msg.pop("_attachment_refs", []) or [])
        new_msg.pop("_model_profile", None)
        if not enabled or new_msg.get("role") != "user" or not refs:
            output.append(new_msg)
            continue

        blocks: list[dict] = []
        text = new_msg.get("content")
        if isinstance(text, str) and text:
            blocks.append({"type": "text", "text": text})
        message_image_count = 0
        for ref in refs:
            if not isinstance(ref, dict) or ref.get("kind") != "image":
                continue
            raw_path = str(ref.get("path") or "")
            try:
                path, _ = resolve_artifact_path(
                    raw_path,
                    working_dir=working_dir,
                    user_root=user_root,
                    require_file=True,
                )
                mime = _NATIVE_IMAGE_MIME.get(path.suffix.lower())
                if mime is None or path.stat().st_size > MAX_NATIVE_IMAGE_BYTES:
                    continue
                data = base64.b64encode(path.read_bytes()).decode("ascii")
            except (ArtifactPathError, OSError):
                continue
            blocks.append({
                "type": "image_url",
                "image_url": {"url": f"data:{mime};base64,{data}"},
            })
            image_count += 1
            message_image_count += 1
        if message_image_count:
            new_msg["content"] = blocks
        output.append(new_msg)
    return output, image_count