smoke_signal/scripts/05_llm_normalise.py

#!/usr/bin/env python3
"""
Smoke Signal — Stage 7: LLM Normalisation Layer
=================================================
Takes low-confidence or visually complex regions from Stage 5
and sends them to an LLM (via HF Inference API) for cleanup.

Rules (non-negotiable):
  - LLM receives: page image crop + raw OCR candidates
  - LLM must return strict JSON only — no free text, no preamble
  - Uncertain words must be marked uncertain, NEVER silently guessed
  - LLM may NOT invent text that isn't visually present
  - Only low-confidence / flagged pages are sent (cost + drift control)
  - Every call logs: prompt version, model, input hash, output hash, cost

Output per region:
  {
    "text_raw":         "original OCR text",
    "text_clean":       "LLM corrected text",
    "uncertain_words":  ["word1", "word2"],
    "confidence_notes": "why confidence is low",
    "changed_from_ocr": true/false,
    "visible_context":  "brief note on what the LLM can see"
  }

Usage:
    python scripts/05_llm_normalise.py
    python scripts/05_llm_normalise.py --book-id SS-BOOK-0001
    python scripts/05_llm_normalise.py --dry-run
    python scripts/05_llm_normalise.py --confidence-below 0.80
"""

import argparse
import base64
import hashlib
import json
import sys
import time
from datetime import datetime
from pathlib import Path
from typing import Optional

# ── Paths ──────────────────────────────────────────────────────────────────────
ROOT          = Path(__file__).resolve().parents[1]
MANIFEST_CSV  = ROOT / "manifest" / "source_manifest.csv"
REGIONS_DIR   = ROOT / "regions"
OCR_RAW_DIR   = ROOT / "ocr_raw"
RENDERS_DIR   = ROOT / "renders"
LOGS_DIR      = ROOT / "logs"
REVIEW_DIR    = ROOT / "review"
CLEANED_DIR   = ROOT / "regions" / "cleaned"

CLEANED_DIR.mkdir(parents=True, exist_ok=True)
LOGS_DIR.mkdir(parents=True, exist_ok=True)

# ── Config (freeze before running — do not change mid-batch) ──────────────────
CONFIG = {
    "config_version":       "ss_llm_v0.1",
    "prompt_version":       "ss_prompt_v0.1",
    "model":                "meta-llama/Llama-3.2-11B-Vision-Instruct",  # HF Inference API
    "confidence_threshold": 0.75,   # only send pages below this
    "max_pages_per_run":    50,     # cost control — cap per batch
    "temperature":          0.1,    # low = deterministic, no hallucination
    "max_new_tokens":       512,
    "eligible_statuses":    ["ocred"],
    "story_classes":        ["narration", "dialogue-speech-bubble", "caption"],
}

# ── Prompt (versioned — never change without bumping prompt_version) ───────────
SYSTEM_PROMPT = """You are a precise OCR correction assistant for children's picture books.

RULES — follow exactly:
1. Return ONLY valid JSON. No preamble, no explanation, no markdown fences.
2. Correct OCR errors you can see in the image. Do NOT invent text.
3. If a word is unclear or unreadable, add it to uncertain_words — do NOT guess.
4. Keep the author's exact words, punctuation, and line breaks.
5. Do not add, remove, or reorder words unless fixing a clear OCR error.
6. changed_from_ocr must be true only if you changed something.

Return this exact JSON structure:
{
  "text_clean": "corrected text here",
  "uncertain_words": ["list", "of", "unclear", "words"],
  "confidence_notes": "brief note on what made this hard to read",
  "changed_from_ocr": false,
  "visible_context_notes": "brief note on what you can see in the image"
}"""

USER_PROMPT_TEMPLATE = """Here is the raw OCR output for a picture book page region:

RAW OCR: {raw_text}
REGION CLASS: {region_class}
OCR CONFIDENCE: {confidence}

Please examine the image crop and return the corrected JSON."""


# ── HF Inference API ───────────────────────────────────────────────────────────
def _get_hf_token() -> Optional[str]:
    """Get HF token from environment or huggingface_hub cache."""
    import os
    token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN")
    if token:
        return token
    try:
        from huggingface_hub import get_token
        return get_token()
    except Exception:
        return None


def call_llm_with_image(
    image_path: Path,
    raw_text: str,
    region_class: str,
    confidence: float,
) -> dict:
    """
    Call HF Inference API with image + OCR text.
    Returns parsed JSON result or error dict.
    """
    import urllib.request
    import urllib.error

    token = _get_hf_token()
    if not token:
        return {
            "error": "no_hf_token",
            "text_clean": raw_text,
            "uncertain_words": [],
            "confidence_notes": "HF token not found — set HF_TOKEN env var or run huggingface-cli login",
            "changed_from_ocr": False,
        }

    # Encode image as base64
    try:
        with open(image_path, "rb") as f:
            image_b64 = base64.b64encode(f.read()).decode("utf-8")
        image_ext = image_path.suffix.lower().replace(".", "")
        media_type = f"image/{image_ext if image_ext in ('png','jpg','jpeg','webp') else 'png'}"
    except Exception as e:
        return {"error": f"image_load_failed: {e}", "text_clean": raw_text,
                "uncertain_words": [], "changed_from_ocr": False}

    user_message = USER_PROMPT_TEMPLATE.format(
        raw_text=raw_text[:800],  # truncate for token budget
        region_class=region_class,
        confidence=round(confidence, 3),
    )

    payload = json.dumps({
        "model": CONFIG["model"],
        "messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {
                "role": "user",
                "content": [
                    {
                        "type": "image_url",
                        "image_url": {"url": f"data:{media_type};base64,{image_b64}"}
                    },
                    {"type": "text", "text": user_message}
                ]
            }
        ],
        "max_tokens": CONFIG["max_new_tokens"],
        "temperature": CONFIG["temperature"],
    }).encode("utf-8")

    api_url = f"https://api-inference.huggingface.co/models/{CONFIG['model']}/v1/chat/completions"

    req = urllib.request.Request(
        api_url,
        data=payload,
        headers={
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json",
        },
        method="POST",
    )

    try:
        with urllib.request.urlopen(req, timeout=30) as resp:
            response_data = json.loads(resp.read().decode("utf-8"))
        raw_response = response_data["choices"][0]["message"]["content"].strip()
    except urllib.error.HTTPError as e:
        return {"error": f"http_{e.code}: {e.reason}", "text_clean": raw_text,
                "uncertain_words": [], "changed_from_ocr": False}
    except Exception as e:
        return {"error": str(e), "text_clean": raw_text,
                "uncertain_words": [], "changed_from_ocr": False}

    # Parse JSON response — strip markdown fences if model added them
    try:
        clean = raw_response.strip()
        if clean.startswith("```"):
            clean = clean.split("```")[1]
            if clean.startswith("json"):
                clean = clean[4:]
        result = json.loads(clean.strip())
    except json.JSONDecodeError:
        return {
            "error": "invalid_json_response",
            "raw_response": raw_response[:500],
            "text_clean": raw_text,
            "uncertain_words": [],
            "changed_from_ocr": False,
        }

    return result


# ── Input hash (for audit trail) ──────────────────────────────────────────────
def _input_hash(text: str, image_path: Path) -> str:
    h = hashlib.sha256()
    h.update(text.encode("utf-8"))
    if image_path.exists():
        with open(image_path, "rb") as f:
            h.update(f.read(4096))  # first 4kb sufficient for fingerprint
    return h.hexdigest()[:16]


def _output_hash(result: dict) -> str:
    return hashlib.sha256(
        json.dumps(result, sort_keys=True).encode("utf-8")
    ).hexdigest()[:16]


# ── Load region data ───────────────────────────────────────────────────────────
def load_ordered_regions(book_id: str) -> Optional[dict]:
    path = REGIONS_DIR / f"{book_id}_ordered_regions.json"
    return json.load(open(path)) if path.exists() else None


# ── Per-region normalisation ───────────────────────────────────────────────────
def normalise_region(
    book_id: str,
    page_num: int,
    region: dict,
    render_path: Optional[str],
    dry_run: bool,
) -> dict:
    """Normalise a single region via LLM. Returns enriched region dict."""

    raw_text     = region.get("text", "").strip()
    region_class = region.get("region_class", "narration")
    confidence   = region.get("confidence", 1.0)
    region_id    = region.get("region_id", f"{book_id}_p{page_num:04d}")

    if dry_run:
        return {
            **region,
            "text_clean":           raw_text,
            "uncertain_words":      [],
            "confidence_notes":     "dry-run",
            "changed_from_ocr":     False,
            "visible_context_notes": "dry-run",
            "llm_model":            CONFIG["model"],
            "prompt_version":       CONFIG["prompt_version"],
            "normalised_at":        datetime.utcnow().isoformat() + "Z",
            "dry_run":              True,
        }

    # Find render image
    img_path = None
    if render_path:
        candidate = ROOT / render_path if not Path(render_path).is_absolute() else Path(render_path)
        if candidate.exists():
            img_path = candidate

    if img_path is None:
        # Try to find any render for this book/page
        book_renders = RENDERS_DIR / book_id
        if book_renders.exists():
            candidates = sorted(book_renders.glob(f"{book_id}_page_{page_num:04d}_*.png"))
            if candidates:
                img_path = candidates[0]

    if img_path is None:
        return {
            **region,
            "text_clean":       raw_text,
            "uncertain_words":  [],
            "confidence_notes": "no_render_available",
            "changed_from_ocr": False,
            "error":            "no_render_found",
        }

    in_hash  = _input_hash(raw_text, img_path)

    llm_result = call_llm_with_image(img_path, raw_text, region_class, confidence)

    out_hash = _output_hash(llm_result)

    return {
        **region,
        "text_clean":            llm_result.get("text_clean", raw_text),
        "uncertain_words":       llm_result.get("uncertain_words", []),
        "confidence_notes":      llm_result.get("confidence_notes", ""),
        "changed_from_ocr":      llm_result.get("changed_from_ocr", False),
        "visible_context_notes": llm_result.get("visible_context_notes", ""),
        "llm_model":             CONFIG["model"],
        "prompt_version":        CONFIG["prompt_version"],
        "input_hash":            in_hash,
        "output_hash":           out_hash,
        "llm_error":             llm_result.get("error"),
        "normalised_at":         datetime.utcnow().isoformat() + "Z",
    }


# ── Per-book normalisation ────────────────────────────────────────────────────
def normalise_book(
    book_id: str,
    filename: str,
    confidence_threshold: float,
    max_pages: int,
    dry_run: bool,
) -> tuple:
    print(f"\n  [{book_id}] {filename}")

    regions_data = load_ordered_regions(book_id)
    if regions_data is None:
        print(f"    ✗ No region data. Run 04_region_detector.py first.")
        return None, {"error": "no_region_data"}

    pages_to_process = []
    for page in regions_data.get("pages", []):
        page_conf = page.get("page_confidence", 1.0)
        route     = page.get("route", "embedded_text")
        if route == "embedded_text":
            continue
        if page_conf < confidence_threshold:
            pages_to_process.append(page)

    if not pages_to_process:
        print(f"    ✓ No pages below confidence threshold {confidence_threshold} — skipping.")
        return regions_data, {}

    # Apply page cap
    if len(pages_to_process) > max_pages:
        print(f"    ⚠️  {len(pages_to_process)} pages need normalisation — capping at {max_pages}")
        pages_to_process = pages_to_process[:max_pages]

    print(f"    Pages to normalise: {len(pages_to_process)}")

    total_changed  = 0
    total_uncertain = 0
    total_errors   = 0
    call_log       = []

    # Process each page
    page_index = {p["page_number"]: i for i, p in enumerate(regions_data["pages"])}

    for page in pages_to_process:
        page_num    = page["page_number"]
        page_conf   = page.get("page_confidence", 0)
        render_path = None

        # Find render path from OCR raw data
        ocr_path = OCR_RAW_DIR / book_id / f"{book_id}_ocr_raw.json"
        if ocr_path.exists():
            ocr_data = json.load(open(ocr_path))
            for ocr_page in ocr_data.get("pages", []):
                if ocr_page.get("page_number") == page_num:
                    render_path = ocr_page.get("render_path")
                    break

        print(f"    Page {page_num:3d} (conf={page_conf:.2f}): ", end="", flush=True)

        normalised_regions = []
        for region in page.get("regions", []):
            region_class = region.get("region_class", "narration")

            # Only normalise story-relevant regions
            if region_class not in CONFIG["story_classes"]:
                normalised_regions.append(region)
                continue

            region_conf = region.get("confidence", 1.0)
            if region_conf >= confidence_threshold:
                normalised_regions.append(region)
                continue

            result = normalise_region(book_id, page_num, region, render_path, dry_run)

            if result.get("changed_from_ocr"):
                total_changed += 1
            if result.get("uncertain_words"):
                total_uncertain += len(result["uncertain_words"])
            if result.get("llm_error"):
                total_errors += 1

            call_log.append({
                "book_id":        book_id,
                "page_number":    page_num,
                "region_id":      region.get("region_id"),
                "input_hash":     result.get("input_hash"),
                "output_hash":    result.get("output_hash"),
                "changed":        result.get("changed_from_ocr", False),
                "uncertain_count": len(result.get("uncertain_words", [])),
                "error":          result.get("llm_error"),
                "model":          CONFIG["model"],
                "prompt_version": CONFIG["prompt_version"],
            })

            normalised_regions.append(result)

        # Update page in regions data
        idx = page_index.get(page_num)
        if idx is not None:
            regions_data["pages"][idx]["regions"] = normalised_regions
            regions_data["pages"][idx]["normalised"] = True

        changed_count   = sum(1 for r in normalised_regions if r.get("changed_from_ocr"))
        uncertain_count = sum(len(r.get("uncertain_words", [])) for r in normalised_regions)
        print(f"changed={changed_count} uncertain_words={uncertain_count}")

    # ── Save cleaned regions ──────────────────────────────────────────────────
    if not dry_run:
        regions_data["normalised_at"]    = datetime.utcnow().isoformat() + "Z"
        regions_data["prompt_version"]   = CONFIG["prompt_version"]
        regions_data["llm_model"]        = CONFIG["model"]
        regions_data["normalisation_log"] = call_log

        cleaned_path = CLEANED_DIR / f"{book_id}_cleaned_regions.json"
        with open(cleaned_path, "w", encoding="utf-8") as f:
            json.dump(regions_data, f, indent=2)
        print(f"    Cleaned regions → {cleaned_path.relative_to(ROOT)}")

        # Log file
        log_path = LOGS_DIR / f"llm_calls_{book_id}_{datetime.utcnow().strftime('%Y%m%d')}.json"
        with open(log_path, "w") as f:
            json.dump(call_log, f, indent=2)

    print(f"    Total changed={total_changed} | uncertain_words={total_uncertain} | errors={total_errors}")
    return regions_data, {}


# ── Main ───────────────────────────────────────────────────────────────────────
def main():
    parser = argparse.ArgumentParser(description="Smoke Signal — Stage 7: LLM Normalisation")
    parser.add_argument("--book-id",           help="Process a single book by ID")
    parser.add_argument("--batch-id",          help="Tag this run with a batch ID")
    parser.add_argument("--confidence-below",  type=float, default=CONFIG["confidence_threshold"],
                        help=f"Only process pages below this confidence (default {CONFIG['confidence_threshold']})")
    parser.add_argument("--max-pages",         type=int, default=CONFIG["max_pages_per_run"],
                        help=f"Max pages per book per run (default {CONFIG['max_pages_per_run']})")
    parser.add_argument("--dry-run",           action="store_true")
    args = parser.parse_args()

    run_id   = args.batch_id or f"SS-RUN-{datetime.utcnow().strftime('%Y%m%d-%H%M%S')}"
    dry_run  = args.dry_run

    print(f"\n{'='*60}")
    print(f"  Smoke Signal — Stage 7: LLM Normalisation")
    print(f"  Run ID    : {run_id}")
    print(f"  Model     : {CONFIG['model']}")
    print(f"  Prompt    : {CONFIG['prompt_version']}")
    print(f"  Threshold : conf < {args.confidence_below}")
    print(f"  Max pages : {args.max_pages}")
    if dry_run:
        print(f"  Mode      : DRY RUN")
    print(f"{'='*60}")

    # Find books with region data
    if args.book_id:
        region_files = [REGIONS_DIR / f"{args.book_id}_ordered_regions.json"]
    else:
        region_files = sorted(REGIONS_DIR.glob("*_ordered_regions.json"))

    if not region_files:
        print("\n  No region files found. Run 04_region_detector.py first.")
        sys.exit(0)

    print(f"\n  Books to process: {len(region_files)}")

    results  = []
    t_start  = time.time()

    for region_file in region_files:
        book_id  = region_file.stem.replace("_ordered_regions", "")
        filename = book_id  # fallback

        _, error = normalise_book(
            book_id,
            filename,
            args.confidence_below,
            args.max_pages,
            dry_run,
        )

        results.append({"book_id": book_id, "error": error})

    elapsed   = round(time.time() - t_start, 1)
    succeeded = sum(1 for r in results if not r.get("error"))

    print(f"\n{'─'*60}")
    print(f"  Books processed : {len(results)}")
    print(f"  Succeeded       : {succeeded}")
    print(f"  Time            : {elapsed}s")
    print(f"{'─'*60}")
    print(f"\n  Cleaned regions → {CLEANED_DIR.relative_to(ROOT)}")
    print(f"  Next: Run 06_review_workbench.py (Stage 8)\n")


if __name__ == "__main__":
    main()