src/lib/processing/ingest_progress.py

"""Stdout-line parser for ingest subprocess output (lifted from stage9_ui).

Why this module exists (BACKEND_BUILD.md §3.4 — allowed additive edit):

Both the Streamlit UI (during its lifetime) and the API job runner need
to parse the same stdout lines that ``src/pipeline/ingest.py`` emits — a
``[stage N] ...`` header for each phase, ``    page N: ok`` lines as
pages are completed, and a final ``=== ingest finished ===`` /
``ingest exited with code`` terminator. Forking two copies of that
parser would diverge over time; lifting the existing one into
``src/lib/processing/`` lets both consumers import from a single
authoritative location.

This is a verbatim copy of the parser that was in
``src/stage9_ui/shared.py``. The Streamlit module imports back from
here as a re-export shim so existing call sites keep working without a
behavior change. The job runner in ``src/api/jobs/types/ingest.py``
imports from here directly.

The parser is **stateless** in the sense that it takes a full log
buffer (list of lines) and returns the current progress snapshot — the
API job runner wraps it in a small streaming adapter that re-runs the
parser as new lines arrive, so the runner can compare snapshots and
decide when to emit a fresh ``stage`` / ``progress`` SSE event.
"""

from __future__ import annotations

import re

# Lines from `transformers` / Streamlit's tf-import are routinely on
# stdout when a subprocess imports them at startup. The UI elided them
# from the visible log; we keep that behavior here so the job runner's
# log_tail doesn't fill up with import-time chatter that says nothing
# about the actual ingest.
_INGEST_LOG_NOISE = (
    "PyTorch was not found",
    "TensorFlow was not found",
    "JAX was not found",
)
_INGEST_STAGE_RE = re.compile(r"^\[stage\s+([\w]+)\]\s*(.*?)(?:\.\.\.|…|\.)?$")
_INGEST_PAGE_OK_RE = re.compile(r"^\s+page\s+\d+:\s+ok\b")


def filter_ingest_log_noise(lines: list[str]) -> list[str]:
    """Drop import-time chatter (transformers / TF / JAX warnings) from the log."""
    return [l for l in lines if not any(n in l for n in _INGEST_LOG_NOISE)]


def parse_ingest_progress(lines: list[str], total_pages: int) -> dict | None:
    """Walk log lines, return the current stage + per-stage page count + pct.

    Resets the page counter on each new ``[stage X]`` so the bar snaps back to
    0% when OCR ends and cleanup starts. Returns ``None`` when the log has no
    useful info yet.

    Output keys (locked — Streamlit's UI and the API runner both read these):
      * ``stage_label`` — human-readable, e.g. ``"Stage 3: OCR"`` (or
        ``"Starting…"`` when no header has been seen yet).
      * ``pages_done`` — count of ``page N: ok`` lines seen since the
        last stage header.
      * ``total_pages`` — passthrough from the caller; the parser doesn't
        infer this (the caller looks it up from the books table).
      * ``pct`` — clamped to ``[0, 1]``; 1.0 when ``finished`` is True.
      * ``finished`` — True if we've seen ``ingest finished cleanly``.
      * ``failed`` — True if we've seen ``FAILED:`` or
        ``ingest exited with code``.
    """
    stage_label = None
    pages_done = 0
    finished = False
    failed = False
    for line in lines:
        m = _INGEST_STAGE_RE.match(line)
        if m:
            stage_label = f"Stage {m.group(1)}: {m.group(2).strip().rstrip('…').rstrip('.')}"
            pages_done = 0
            continue
        if _INGEST_PAGE_OK_RE.match(line):
            pages_done += 1
            continue
        low = line.lower()
        if "ingest finished cleanly" in low or "=== ingest finished" in low:
            finished = True
        elif "ingest exited with code" in low or "FAILED:" in line:
            failed = True
    if stage_label is None and pages_done == 0 and not finished and not failed:
        return None
    pct = (pages_done / total_pages) if total_pages else 0.0
    pct = min(max(pct, 0.0), 1.0)
    if finished:
        pct = 1.0
    return {
        "stage_label": stage_label or "Starting…",
        "pages_done":  pages_done,
        "total_pages": total_pages,
        "pct":         pct,
        "finished":    finished,
        "failed":      failed,
    }


# ---------- Canonical stage-name mapping (CONTRACT.md §5) -----------------
# The stdout headers carry a *number* (``[stage 2]``, ``[stage 3]``, etc.)
# but the SSE event taxonomy uses the canonical string names from
# CONTRACT.md §5. Both sides agreed on these names: ``acquisition``,
# ``ocr``, ``cleanup``, ``chunking``, ``indexing``. We map number → name
# here so the API runner only has to do one lookup. Stages 4b (printed-
# page extraction) is folded into ``cleanup`` because there's no separate
# canonical name for it and the FE doesn't show a distinct phase.
_STAGE_NUMBER_TO_NAME: dict[str, str] = {
    "2":  "acquisition",
    "3":  "ocr",
    "4":  "cleanup",
    "4b": "cleanup",
    "5":  "chunking",
    "6":  "indexing",
}


def stage_number_to_name(stage_number: str | None) -> str | None:
    """Map a stdout stage-header number → canonical SSE stage name.

    Returns ``None`` for unknown numbers so the caller can decide whether
    to emit a generic ``stage`` event or skip it entirely.
    """
    if not stage_number:
        return None
    return _STAGE_NUMBER_TO_NAME.get(stage_number.strip().lower())