src/knowledge_parsing/assets.py

"""Asset identity and persistence for figures pulled out of a parse. Added 0.4.0.

Two jobs, deliberately small:

1. **Identity.** `asset_id()` content-addresses the bytes. The parser's own id is
   not usable: Mistral names images per PAGE (`img-0.jpeg`), so the same name
   recurs across pages — the Open Pit textbook produces **7 distinct names for 8
   images**, and a flat store keyed on the parser's name loses one. Content
   addressing also makes a re-parse idempotent and de-duplicates a figure that
   appears in two documents.

2. **Persistence.** An `AssetSink` protocol plus a local-filesystem default.

**Why a protocol and not a blob client.** `src/knowledge_parsing/` imports nothing
from the rest of `src` and that is load-bearing — the package is handed to the
team as a standalone unit, so reaching into `src.storage` (which pulls app
settings, which needs `.env`) would break the export and make the parser
un-runnable outside this repo. The app injects a blob-backed sink instead; see
`knowledge_ingest.asset_sink`.

Persistence is a **never-throw seam**. A figure whose bytes cannot be stored
yields `storage_key=None` and a logged warning, never a failed parse: the parse is
the expensive, cached, hard-to-repeat half, and losing a whole document because
one PNG could not be written is the wrong trade.
"""

from __future__ import annotations

import base64
import binascii
import hashlib
import logging
from pathlib import Path
from typing import Protocol, runtime_checkable

logger = logging.getLogger("knowledge_parsing.assets")

ASSET_ID_CHARS = 16

# Mistral reports ids like "img-0.jpeg"; the suffix is the only extension signal
# available. Anything unrecognised falls back to .bin rather than guessing, so a
# stored object is never mislabelled.
_KNOWN_EXT = {"jpg", "jpeg", "png", "gif", "webp", "bmp", "tiff"}


def asset_id(data: bytes) -> str:
    """Content address for asset bytes: sha256, first 16 hex chars.

    16 chars is ~64 bits. At corpus scale (thousands of figures) the collision
    probability is negligible, and it keeps the id short enough to read in a
    `![](asset://…)` reference inside a prose chunk.
    """
    return hashlib.sha256(data).hexdigest()[:ASSET_ID_CHARS]


def extension_for(parser_id: str | None) -> str:
    """File extension from the parser's own id, defaulting to `bin`."""
    if not parser_id or "." not in parser_id:
        return "bin"
    ext = parser_id.rsplit(".", 1)[-1].strip().lower()
    return ext if ext in _KNOWN_EXT else "bin"


def decode_base64(payload: str | None) -> bytes | None:
    """Decode a base64 image payload, tolerating a `data:` URI prefix.

    Returns None rather than raising on anything malformed — the caller treats a
    missing payload as "no bytes to store", which is the same degraded path as a
    failed write.
    """
    if not payload:
        return None
    text = payload.strip()
    if text.startswith("data:"):
        _, _, text = text.partition(",")
    try:
        return base64.b64decode(text, validate=False)
    except (binascii.Error, ValueError) as exc:
        logger.warning("asset base64 decode failed: %r", exc)
        return None


@runtime_checkable
class AssetSink(Protocol):
    """Somewhere to put asset bytes. Returns a storage key, or None on failure."""

    def put(self, identifier: str, extension: str, data: bytes) -> str | None:
        ...


class LocalAssetSink:
    """Write assets to a local directory, named by content address.

    The default, and what the CLI uses. Idempotent by construction: an asset whose
    id already exists on disk is not rewritten, because the id IS the content.
    """

    def __init__(self, root: Path, prefix: str = "assets") -> None:
        self.root = Path(root)
        self.prefix = prefix

    def put(self, identifier: str, extension: str, data: bytes) -> str | None:
        try:
            target_dir = self.root / self.prefix
            target_dir.mkdir(parents=True, exist_ok=True)
            target = target_dir / f"{identifier}.{extension}"
            if not target.exists():
                target.write_bytes(data)
            return f"{self.prefix}/{identifier}.{extension}"
        except OSError as exc:
            logger.warning("asset write failed for %s: %r", identifier, exc)
            return None


class NullAssetSink:
    """Discard asset bytes, keeping the asset metadata. For `--no-assets`.

    Not the same as having no assets: the figure is still folded into its prose
    chunk with its caption and description, and still carries an `asset_id`. Only
    the bytes are not retained, so `storage_key` is None.
    """

    def put(self, identifier: str, extension: str, data: bytes) -> str | None:
        return None


def store_asset(
    sink: AssetSink | None,
    parser_id: str | None,
    base64_payload: str | None,
) -> tuple[str | None, str | None]:
    """Content-address and persist one asset. Returns `(asset_id, storage_key)`.

    `asset_id` is None only when there are no bytes to address — a table, or a
    figure whose payload the parser did not return. `storage_key` is None whenever
    the bytes were not retained, which is a normal outcome (`NullAssetSink`) as
    well as a failure one.
    """
    data = decode_base64(base64_payload)
    if data is None:
        return None, None
    identifier = asset_id(data)
    if sink is None:
        return identifier, None
    return identifier, sink.put(identifier, extension_for(parser_id), data)


def synthetic_asset_id(*parts: str) -> str:
    """Id for an asset with no bytes of its own — a table.

    Derived from the document and chunk coordinates rather than content, because
    there is no content to hash. Deterministic for the same document, which is
    what matters: a re-parse produces the same id, so a pointer to a table
    survives re-extraction the way a figure's does.
    """
    return hashlib.sha256("|".join(parts).encode("utf-8")).hexdigest()[:ASSET_ID_CHARS]