Blame for src/knowledge_parsing/parse.py

"""Stage 2 — parse a document with MinerU, cached by content.

Why the cache is keyed to document CONTENT rather than to a run id: parsing is
by far the most expensive stage (19 s/page on CPU). A per-run cache would
re-parse everything each time the normaliser changes and the pipeline is run
again, for nothing. Keyed by content + settings + version, the same document is
parsed once — however many times the normaliser is revised.

Cache key = sha256(file bytes) + settings fingerprint + MinerU version. The
version is in there so figures from an older MinerU never quietly mix with
figures from a newer one — which matters, because those figures are used to
justify spend.
"""

from __future__ import annotations

import hashlib
import json
import shutil
import time
from dataclasses import dataclass
from pathlib import Path

from .config import PipelineConfig


@dataclass
class ParseResult:
    doc_id: str
    source: Path
    cache_dir: Path          # folder holding MinerU's output untouched
    content_list: Path       # *_content_list.json
    middle_json: Path | None
    pages: int
    seconds: float
    from_cache: bool
    backend_recorded: str | None    # read from _middle.json, NOT from the config
    mineru_version: str | None


def mineru_version() -> str:
    try:
        from mineru.version import __version__  # type: ignore
        return str(__version__)
    except Exception:
        try:
            from importlib.metadata import version
            return version("mineru")
        except Exception:
            return "unknown"


def file_hash(path: Path, length: int = 16) -> str:
    h = hashlib.sha256()
    with path.open("rb") as f:
        for block in iter(lambda: f.read(1 << 20), b""):
            h.update(block)
    return h.hexdigest()[:length]


def cache_key(source: Path, cfg: PipelineConfig) -> str:
    return f"{file_hash(source)}-{cfg.fingerprint()}-{mineru_version()}"


def _find_output(root: Path) -> tuple[Path | None, Path | None]:
    """Locate content_list & middle json anywhere under root.

    Deliberately a glob rather than a fixed path: MinerU's subfolder layout
    differs per backend ('auto' for pipeline, something else for vlm). Globbing
    means this module needs no change when the backend does.
    """
    content = next(iter(sorted(root.rglob("*_content_list.json"))), None)
    middle = next(iter(sorted(root.rglob("*_middle.json"))), None)
    return content, middle


def _read_middle(middle: Path | None) -> tuple[str | None, str | None]:
    """The backend & version that ACTUALLY ran, taken from MinerU's own output."""
    if not middle or not middle.exists():
        return None, None
    try:
        d = json.loads(middle.read_text(encoding="utf-8"))
        return d.get("_backend"), d.get("_version_name")
    except Exception:
        return None, None


def parse_document(source: Path, cfg: PipelineConfig) -> ParseResult:
    """Parse one document. A cache hit is not re-run."""
    doc_id = source.stem
    target = cfg.cache_dir / "parse" / cache_key(source, cfg)
    marker = target / ".done"

    if marker.exists():
        content, middle = _find_output(target)
        if content:
            meta = json.loads(marker.read_text(encoding="utf-8"))
            backend_recorded, version = _read_middle(middle)
            return ParseResult(
                doc_id=doc_id, source=source, cache_dir=target,
                content_list=content, middle_json=middle,
                pages=meta.get("pages", 0), seconds=meta.get("seconds", 0.0),
                from_cache=True,
                backend_recorded=backend_recorded, mineru_version=version,
            )
        shutil.rmtree(target, ignore_errors=True)   # cache corrupt, redo it

    from mineru.cli.common import do_parse, read_fn

    # MinerU dispatches on prefixes: "pipeline", then `startswith("vlm-")` and
    # `startswith("hybrid-")`. The bare names "vlm"/"hybrid" match NONE of those
    # branches — MinerU simply does nothing, without error, and the failure only
    # surfaces here as "content_list not found", which points at the wrong place
    # entirely. The short names stay in the config because they are part of the
    # cache key and read better; the translation happens here, right before
    # MinerU is touched.
    mineru_backend = {
        "pipeline": "pipeline",
        "vlm": "vlm-engine",
        "hybrid": "hybrid-engine",
    }[cfg.backend]

    # The playground does this before parsing; aliases like "latin"/"en" are
    # resolved to a canonical code. Matched here so results are identical.
    lang = cfg.canonical_lang()

    # with_suffix() REPLACES the last suffix, and the folder name ends in the
    # MinerU version ("...-3.4.4") — so ".4" would be thrown away and two
    # different versions would collide on one temp directory.
    staging = target.with_name(target.name + ".partial")
    shutil.rmtree(staging, ignore_errors=True)
    staging.mkdir(parents=True, exist_ok=True)

    pdf_bytes = read_fn(source)
    started = time.perf_counter()
    do_parse(
        output_dir=str(staging),
        pdf_file_names=[doc_id],
        pdf_bytes_list=[pdf_bytes],
        p_lang_list=[lang],
        backend=mineru_backend,
        formula_enable=cfg.formula_enable,
        table_enable=cfg.table_enable,
        f_draw_layout_bbox=cfg.write_debug_pdf,
        f_draw_span_bbox=cfg.write_debug_pdf,
        start_page_id=cfg.start_page,
        end_page_id=cfg.end_page,
        **({"effort": cfg.effort} if cfg.backend == "hybrid" else {}),
    )
    seconds = time.perf_counter() - started

    content, middle = _find_output(staging)
    if content is None:
        raise RuntimeError(
            f"MinerU finished but no *_content_list.json was found in {staging}"
        )

    pages = _count_pages(source, cfg)
    staging.rename(target)
    content = target / content.relative_to(staging)
    middle = target / middle.relative_to(staging) if middle else None

    marker.write_text(
        json.dumps({"pages": pages, "seconds": seconds, "doc_id": doc_id}),
        encoding="utf-8",
    )
    backend_recorded, version = _read_middle(middle)

    return ParseResult(
        doc_id=doc_id, source=source, cache_dir=target,
        content_list=content, middle_json=middle,
        pages=pages, seconds=seconds, from_cache=False,
        backend_recorded=backend_recorded, mineru_version=version,
    )


def _count_pages(source: Path, cfg: PipelineConfig) -> int:
    if cfg.end_page is not None:
        return cfg.end_page - cfg.start_page + 1
    try:
        from pypdf import PdfReader
        return len(PdfReader(str(source)).pages)
    except Exception:
        return 0