lightrag-cn-git/lightrag/llm/_vision_utils.py

"""Shared image-input normalization for LLM bindings.

All LLM bindings accept a unified ``image_inputs`` keyword parameter. Each
element may be:

- a raw base64 string (the MIME type is inferred via ``imghdr`` / magic bytes,
  defaulting to ``image/png``);
- a data URL of the form ``data:<mime>;base64,<payload>``;
- a dict with keys ``base64`` (required) and optional ``mime_type``,
  ``source_id``, ``source_file``, ``modality``, ``doc_id``.

The provider-specific binding code converts the normalized result to its own
content-block format. The VLM pipeline uses :func:`image_cache_metadata` for
cache-key inputs (deliberately excluding ``source_id`` / ``source_file`` so the
same image at different filenames still hits the same entry) and
:func:`image_audit_metadata` for the human-readable ``original_prompt`` audit
block.
"""

from __future__ import annotations

import base64
import hashlib
import re
import struct
from dataclasses import dataclass
from pathlib import Path
from typing import Any

DATA_URL_RE = re.compile(
    r"^data:(?P<mime>[\w./+-]+);base64,(?P<data>[A-Za-z0-9+/=\s]+)$"
)

_PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"
_JPEG_SIGNATURE = b"\xff\xd8\xff"
_GIF_SIGNATURES = (b"GIF87a", b"GIF89a")
_WEBP_RIFF = b"RIFF"
_WEBP_TAG = b"WEBP"


@dataclass(frozen=True)
class NormalizedImage:
    index: int
    raw_bytes: bytes
    mime_type: str
    sha256: str
    base64_str: str
    source_id: str | None
    source_file: str | None
    modality: str | None
    doc_id: str | None
    # Pixel dimensions parsed from the raster header (None when the format
    # is recognized but dimensions could not be extracted).
    width: int | None = None
    height: int | None = None


def _detect_mime(raw: bytes) -> str:
    if raw.startswith(_PNG_SIGNATURE):
        return "image/png"
    if raw.startswith(_JPEG_SIGNATURE):
        return "image/jpeg"
    if any(raw.startswith(sig) for sig in _GIF_SIGNATURES):
        return "image/gif"
    if len(raw) >= 12 and raw[0:4] == _WEBP_RIFF and raw[8:12] == _WEBP_TAG:
        return "image/webp"
    return "image/png"


def _decode_base64(data: str) -> bytes:
    cleaned = re.sub(r"\s+", "", data)
    try:
        return base64.b64decode(cleaned, validate=True)
    except (base64.binascii.Error, ValueError) as exc:
        raise ValueError(f"invalid base64 image data: {exc}") from exc


def _coerce_item(item: Any) -> dict[str, Any]:
    if isinstance(item, str):
        match = DATA_URL_RE.match(item.strip())
        if match:
            return {"base64": match.group("data"), "mime_type": match.group("mime")}
        return {"base64": item}
    if isinstance(item, dict):
        if "base64" not in item:
            raise ValueError("image_inputs dict element must contain a 'base64' key")
        return item
    raise TypeError(
        f"image_inputs element must be str or dict, got {type(item).__name__}"
    )


def normalize_image_inputs(
    image_inputs: list[Any] | None,
) -> list[NormalizedImage]:
    """Normalize the unified ``image_inputs`` parameter.

    Returns an empty list when ``image_inputs`` is falsy, so callers can do a
    plain ``if normalized:`` check.
    """
    if not image_inputs:
        return []

    result: list[NormalizedImage] = []
    for idx, raw_item in enumerate(image_inputs):
        item = _coerce_item(raw_item)
        raw_bytes = _decode_base64(item["base64"])
        if not raw_bytes:
            raise ValueError(f"image_inputs[{idx}] decoded to empty bytes")
        mime_type = item.get("mime_type") or _detect_mime(raw_bytes)
        sha = hashlib.sha256(raw_bytes).hexdigest()
        clean_b64 = base64.b64encode(raw_bytes).decode("ascii")
        dims = _dimensions_from_bytes(raw_bytes)
        width, height = (dims[0], dims[1]) if dims else (None, None)
        result.append(
            NormalizedImage(
                index=idx,
                raw_bytes=raw_bytes,
                mime_type=mime_type,
                sha256=sha,
                base64_str=clean_b64,
                source_id=item.get("source_id"),
                source_file=item.get("source_file"),
                modality=item.get("modality"),
                doc_id=item.get("doc_id"),
                width=width,
                height=height,
            )
        )
    return result


def image_cache_metadata(images: list[NormalizedImage]) -> list[dict[str, Any]]:
    """Return cache-key-safe image metadata (no source identifiers).

    Includes ``width`` / ``height`` so the cache key reflects the full
    image digest the design contract specifies (mime, sha256, bytes,
    width, height).  The sha256 alone is sufficient for identity, but
    surfacing dimensions matches the documented audit shape and gives
    diagnostics a one-line "what was sent" without re-decoding.
    """
    return [
        {
            "index": img.index,
            "mime_type": img.mime_type,
            "sha256": img.sha256,
            "bytes": len(img.raw_bytes),
            "width": img.width,
            "height": img.height,
        }
        for img in images
    ]


def image_audit_metadata(images: list[NormalizedImage]) -> list[dict[str, Any]]:
    """Return audit metadata suitable for the ``original_prompt`` block.

    Never includes the raw base64 payload — only digests and source pointers.
    """
    return [
        {
            "index": img.index,
            "mime_type": img.mime_type,
            "sha256": img.sha256,
            "bytes": len(img.raw_bytes),
            "width": img.width,
            "height": img.height,
            "source_id": img.source_id,
            "source_file": img.source_file,
            "modality": img.modality,
            "doc_id": img.doc_id,
        }
        for img in images
    ]


def _read_png_dimensions(data: bytes) -> tuple[int, int] | None:
    # IHDR is the first chunk; width/height are big-endian uint32 at offsets
    # 16/20 (8-byte signature + 4 length + 4 "IHDR" + 4 width + 4 height).
    if len(data) < 24 or not data.startswith(_PNG_SIGNATURE):
        return None
    width, height = struct.unpack(">II", data[16:24])
    return width, height


def _read_gif_dimensions(data: bytes) -> tuple[int, int] | None:
    # Logical screen descriptor: width/height are little-endian uint16 at
    # offsets 6/8.
    if len(data) < 10 or not any(data.startswith(sig) for sig in _GIF_SIGNATURES):
        return None
    width, height = struct.unpack("<HH", data[6:10])
    return width, height


def _read_jpeg_dimensions(data: bytes) -> tuple[int, int] | None:
    # Scan for a Start-Of-Frame marker (SOF0 / SOF2 / etc.). Skip segments by
    # their length field. We deliberately accept any SOF variant the codec
    # might emit rather than enumerating each one.
    if len(data) < 4 or not data.startswith(_JPEG_SIGNATURE):
        return None
    i = 2
    n = len(data)
    while i < n:
        if data[i] != 0xFF:
            return None
        # Skip fill bytes.
        while i < n and data[i] == 0xFF:
            i += 1
        if i >= n:
            return None
        marker = data[i]
        i += 1
        # Standalone markers without a length field.
        if marker in (0xD8, 0xD9) or 0xD0 <= marker <= 0xD7:
            continue
        if i + 2 > n:
            return None
        segment_len = struct.unpack(">H", data[i : i + 2])[0]
        if segment_len < 2 or i + segment_len > n:
            return None
        # SOF0..SOF15 except 0xC4 (DHT), 0xC8 (JPG reserved), 0xCC (DAC).
        if 0xC0 <= marker <= 0xCF and marker not in (0xC4, 0xC8, 0xCC):
            # SOF payload: precision(1) + height(2) + width(2) + …
            if i + 7 > n:
                return None
            height, width = struct.unpack(">HH", data[i + 3 : i + 7])
            return width, height
        i += segment_len
    return None


def _read_webp_dimensions(data: bytes) -> tuple[int, int] | None:
    if len(data) < 30 or data[0:4] != _WEBP_RIFF or data[8:12] != _WEBP_TAG:
        return None
    chunk_type = data[12:16]
    if chunk_type == b"VP8 ":
        # Lossy: 3-byte tag + 3-byte sync code at offset 23, then 4 bytes
        # holding 14-bit width / 14-bit height in little-endian halves.
        if len(data) < 30:
            return None
        width = struct.unpack("<H", data[26:28])[0] & 0x3FFF
        height = struct.unpack("<H", data[28:30])[0] & 0x3FFF
        return width, height
    if chunk_type == b"VP8L":
        # Lossless: signature(0x2F) + 4 bytes encoding 14-bit width-1 / 14-bit
        # height-1 starting at offset 21.
        if len(data) < 25 or data[20] != 0x2F:
            return None
        b0, b1, b2, b3 = data[21], data[22], data[23], data[24]
        width = ((b1 & 0x3F) << 8 | b0) + 1
        height = ((b3 & 0x0F) << 10 | b2 << 2 | (b1 & 0xC0) >> 6) + 1
        return width, height
    if chunk_type == b"VP8X":
        # Extended: 3 bytes width-1 / 3 bytes height-1, little-endian, at
        # offsets 24/27.
        if len(data) < 30:
            return None
        width = (data[24] | data[25] << 8 | data[26] << 16) + 1
        height = (data[27] | data[28] << 8 | data[29] << 16) + 1
        return width, height
    return None


def read_image_dimensions(path: Path) -> tuple[int, int] | None:
    """Return ``(width, height)`` for a raster image, or ``None`` if unknown.

    Reads only the file header — no Pillow dependency. Supports PNG, JPEG,
    GIF and WebP (VP8 / VP8L / VP8X). Returns ``None`` for unsupported
    formats and on any I/O or parse error so callers can fall back to a
    skipped/failure decision without raising.
    """
    try:
        with open(path, "rb") as fh:
            header = fh.read(64 * 1024)
    except OSError:
        return None
    return _dimensions_from_bytes(header)


def _dimensions_from_bytes(data: bytes) -> tuple[int, int] | None:
    """Run the four header readers against a byte buffer.

    Shared between the file-path entry point (:func:`read_image_dimensions`)
    and :func:`normalize_image_inputs`, which receives raster payloads
    decoded from the unified ``image_inputs`` parameter.
    """
    if not data:
        return None
    for reader in (
        _read_png_dimensions,
        _read_gif_dimensions,
        _read_jpeg_dimensions,
        _read_webp_dimensions,
    ):
        try:
            dims = reader(data)
        except (struct.error, IndexError, ValueError):
            continue
        if dims:
            return dims
    return None