morpheus-docs/scrape/vision.py

"""Local-Ollama vision OCR for corpus images (opt-in, off by default).

Some HPE doc pages carry information *only* inside an image. The Morpheus
release schedule / support matrix (sf000111242en_us) is the canonical case:
the entire lifecycle table — which version ships when, when a stream hits
Maintenance/EOL — is a JPEG, so html_to_md would drop every value and
retrieval could never answer "when does Morpheus 8.x.x reach end of life".

When VISION_OCR=1, `transcribe()` sends qualifying images to a local vision
model (qwen2.5vl:7b on the host Ollama by default) and returns a labeled
markdown block that html_to_md appends to the page, so the text gets chunked,
embedded and retrieved.

Model choice (measured on the release-schedule matrix, known ground truth):
qwen2.5vl:7b scored 55/56 cells with a bare prompt, ~18s/image, fully on GPU,
and — unlike gemma3:12b — needs no hand-tuned prompt (gemma is accurate only
with an exact "read row by row, keep cells aligned" instruction). qwen2.5vl:32b
was no more accurate (same single miss) and 5x slower (CPU spill). See the
_PROMPT comment for the per-model wording lore.

Reliability (the whole point of this module):
  - temperature 0 + a model/prompt pairing chosen against ground truth.
  - Two samples with different seeds must AGREE after normalization
    (self-consistency). On disagreement we take a third sample and keep the
    majority, flagged `certain=False` so the reader knows to double-check.
    Self-consistency catches RANDOM flakiness only — a model that reads a
    table wrong does so identically every run — so every transcription also
    ships with a "verify against the source image" caveat.

Cost control (a full --force re-scrape re-touches every page weekly):
  - Content-hash cache in corpus/.vision-cache/ — each unique image is
    OCR'd once, ever, and the result is committed so CI reuses it. The same
    diagram shared across version bundles collapses to one OCR.
  - VISION_MAX_NEW bounds NEW OCRs per run so the first run can't balloon
    into hours; the cache fills incrementally over subsequent refreshes.
    Deferred images are logged (never silently dropped).

Every failure path degrades to None — a down endpoint, a timeout, a decode
error — so the scrape never blocks on vision. Defaults target the git.jpaul.io
Ollama host; override via env for other deployments.

Config (env):
  VISION_OCR      "1" to enable (default off)
  VISION_URL      Ollama base URL           (default http://192.168.0.2:11434)
  VISION_MODEL    vision model tag          (default qwen2.5vl:7b)
  VISION_PROMPT   override the OCR prompt    (default is tuned for qwen2.5vl)
  VISION_MIN_W/H  min image dims to OCR     (default 600 x 300 — skips icons)
  VISION_MAX_NEW  new OCRs per run (budget) (default 60)
  VISION_NUM_CTX  Ollama context window     (default 8192)
  VISION_TIMEOUT  per-call seconds          (default 300)
"""
from __future__ import annotations

import base64
import hashlib
import io
import json
import os
import re
import sys
import threading
from pathlib import Path

import requests

ROOT = Path(__file__).resolve().parent.parent
CACHE_DIR = ROOT / "corpus" / ".vision-cache"

VISION_OCR = os.environ.get("VISION_OCR", "") == "1"
VISION_URL = os.environ.get("VISION_URL", "http://192.168.0.2:11434").rstrip("/")
VISION_MODEL = os.environ.get("VISION_MODEL", "qwen2.5vl:7b")
VISION_MIN_W = int(os.environ.get("VISION_MIN_W", "600"))
VISION_MIN_H = int(os.environ.get("VISION_MIN_H", "300"))
VISION_MAX_NEW = int(os.environ.get("VISION_MAX_NEW", "60"))
VISION_TIMEOUT = int(os.environ.get("VISION_TIMEOUT", "300"))
VISION_NUM_CTX = int(os.environ.get("VISION_NUM_CTX", "8192"))

# Prompt wording is load-bearing, and the right wording is MODEL-SPECIFIC —
# measured on the release-schedule matrix (14 rows x 4 version columns, known
# ground truth). Errors were stable across seeds, so self-consistency does NOT
# catch them; only prompt/model choice does.
#   - qwen2.5vl:7b (default): a bare "transcribe the table" instruction scores
#     55/56 cells; adding ANY extra clause (a non-table fallback, "preserve
#     blank cells") regressed it to 51/56. So keep it minimal. qwen also reads
#     it correctly WITHOUT hand-tuning — the prompt-robustness we want.
#   - gemma3:12b: the opposite — the bare prompt shifts a column by one row;
#     it needs an explicit "read row by row, keep cells aligned to their
#     column" instruction to hit 55/56.
# The single cell every local model (incl. qwen2.5vl:32b) misses is the merged
# "9.1.2 - 9.3.2" range row. Set VISION_PROMPT to override for other models.
_DEFAULT_PROMPT = "Transcribe the table in this image to a GitHub markdown table."
_PROMPT = os.environ.get("VISION_PROMPT", _DEFAULT_PROMPT)

_lock = threading.Lock()
_new_ocr = 0
_stats = {
    "cached": 0, "ocr": 0, "uncertain": 0,
    "skipped_small": 0, "deferred": 0, "errors": 0, "empty": 0,
}


def enabled() -> bool:
    return VISION_OCR


def _bump(key: str, n: int = 1) -> None:
    with _lock:
        _stats[key] += n


def _dims(data: bytes) -> tuple[int, int] | None:
    try:
        from PIL import Image
        with Image.open(io.BytesIO(data)) as im:
            return im.size
    except Exception:
        return None


def _generate(img_b64: str, seed: int) -> str | None:
    body = {
        "model": VISION_MODEL,
        "prompt": _PROMPT,
        "images": [img_b64],
        "stream": False,
        "keep_alive": "10m",
        "options": {"temperature": 0, "seed": seed, "num_ctx": VISION_NUM_CTX},
    }
    r = requests.post(f"{VISION_URL}/api/generate", json=body, timeout=VISION_TIMEOUT)
    r.raise_for_status()
    return (r.json() or {}).get("response")


def _clean(s: str) -> str:
    """Extract just the transcription from a model response, dropping chatty
    preamble ("Here's the transcription:") and code fences. gemma3 in
    particular wraps output in ```markdown fences and a lead-in sentence
    despite being told not to."""
    s = s.strip()
    # Prefer the contents of a fenced block if the model wrapped its answer.
    m = re.search(r"```[a-zA-Z]*\n(.*?)```", s, re.S)
    if m:
        s = m.group(1).strip()
    # For a table, drop any preamble before the first table row.
    lines = s.splitlines()
    for i, ln in enumerate(lines):
        if ln.lstrip().startswith("|"):
            return "\n".join(lines[i:]).strip()
    return s.strip()


def _normalize(s: str) -> str:
    """Collapse to a comparable form so two OCR samples that differ only in
    whitespace/case count as agreeing (compared on cleaned text)."""
    return re.sub(r"\s+", " ", s).strip().lower()


def _ocr_bytes(data: bytes) -> tuple[str | None, bool]:
    """Two-sample self-consistency → (cleaned_markdown, certain). A third
    sample breaks a tie; `certain` is True only when >=2 samples agree.

    NOTE: self-consistency catches RANDOM flakiness, not SYSTEMATIC error — a
    prompt/model that reads a table wrong tends to do so identically every
    time (we saw exactly this). It is a stability check, not a correctness
    proof; the transcription is always emitted with a verify-against-source
    caveat. The real correctness lever is the prompt/model choice."""
    b64 = base64.b64encode(data).decode()
    a = _clean(_generate(b64, 7) or "")
    b = _clean(_generate(b64, 99) or "")
    if a and b and _normalize(a) == _normalize(b):
        return a, True
    cands = [x for x in (a, b) if x]
    if not cands:
        return None, False
    c = _clean(_generate(b64, 1234) or "")
    if c:
        cands.append(c)
    groups: dict[str, list[str]] = {}
    for x in cands:
        groups.setdefault(_normalize(x), []).append(x)
    best = max(groups.values(), key=len)
    return best[0], len(best) >= 2


def _cache_path(sha: str) -> Path:
    return CACHE_DIR / f"{sha}.json"


def _write_cache(sha: str, url: str, markdown: str | None, certain: bool) -> None:
    try:
        CACHE_DIR.mkdir(parents=True, exist_ok=True)
        _cache_path(sha).write_text(json.dumps({
            "sha": sha, "src": url, "model": VISION_MODEL,
            "certain": certain, "markdown": markdown,
        }, indent=2) + "\n")
    except Exception:
        pass


def _format(markdown: str, certain: bool) -> str:
    note = (f"Transcribed from the image above by local vision OCR "
            f"({VISION_MODEL}); verify against the source image.")
    if not certain:
        note += " NOTE: OCR samples disagreed — treat values as approximate."
    return f"\n> {note}\n\n{markdown}\n"


def transcribe(img_url: str, session: requests.Session) -> str | None:
    """Return a labeled markdown transcription block for `img_url`, or None if
    vision is disabled, the image is too small, the budget is spent, the image
    has no readable text, or anything errors. Never raises."""
    global _new_ocr
    if not VISION_OCR:
        return None

    try:
        resp = session.get(img_url, timeout=60)
        if resp.status_code != 200 or not resp.content:
            return None
        data = resp.content
    except Exception:
        _bump("errors")
        return None

    sha = hashlib.sha256(data).hexdigest()
    cp = _cache_path(sha)
    if cp.exists():
        try:
            entry = json.loads(cp.read_text())
        except Exception:
            entry = None
        if entry is not None:
            _bump("cached")
            md = entry.get("markdown")
            return _format(md, entry.get("certain", True)) if md else None

    dims = _dims(data)
    if not dims or dims[0] < VISION_MIN_W or dims[1] < VISION_MIN_H:
        _bump("skipped_small")
        return None

    with _lock:
        if _new_ocr >= VISION_MAX_NEW:
            _stats["deferred"] += 1
            return None
        _new_ocr += 1

    try:
        markdown, certain = _ocr_bytes(data)
    except Exception:
        _bump("errors")
        return None

    if not markdown or len(markdown.strip()) < 20:
        # Negative-cache tiny/no-text images so we never re-spend budget on them.
        _write_cache(sha, img_url, None, True)
        _bump("empty")
        return None

    _write_cache(sha, img_url, markdown, certain)
    _bump("ocr")
    if not certain:
        _bump("uncertain")
    return _format(markdown, certain)


def summary() -> str:
    parts = ", ".join(f"{k}={v}" for k, v in _stats.items())
    return f"vision[{VISION_MODEL} @ {VISION_URL}]: {parts}"