pdf-quiz-generator/backend/scripts/triage_question_images.py

"""Take the answer figure off the question stem.

PDF extraction attached one image per question to `image_path`, which the quiz
player shows above the stem. For many questions that image is not the vignette's
photograph at all — it is the figure from the *explanation*, the annotated smear
or the labelled radiograph that gives the answer away. Showing it during a test
is a spoiler.

Three groups, decided in this order:

  1. the stem itself refers to a figure  → the image belongs there, leave it;
  2. only the explanation refers to one  → the image belongs to the explanation,
     so it comes off the stem;
     (the source material labels its figures by where they are printed — "Item Q37A" beside the
     question, "Item C37B" beside the critique — which decides most of these;)
  3. neither says anything               → ask a vision model what the picture
     shows and which half of the question it illustrates.

Removal means `image_path = NULL`. The image is not moved to
`explanation_image_path`: an unverified guess about where a figure belongs is
worth less than a clean stem, and the file itself is untouched in object
storage. Every removal snapshots the question first, through the same helper the
edit UI uses, so any of this can be restored from the question's version history.

    docker compose exec backend python -m scripts.triage_question_images
    docker compose exec backend python -m scripts.triage_question_images --apply

    --no-vision     only apply the deterministic rule, ask no model
    --limit N       classify at most N unknowns (for a spot check)
    --model ID      override the model (default: the configured extraction model)
    --report PATH   where the human-review list is written

Anything the model is not confident about is left exactly as it is and listed in
the report, with the question id and the image path, for a person to decide.
"""
import argparse
import base64
import json
import mimetypes
import re
import sys
from concurrent.futures import ThreadPoolExecutor

import litellm

from app.config import settings
from app.database import SessionLocal
from app.models.question import Question
from app.routers.questions import _snapshot_question
from app.services import storage_service
from app.services.ai_service import _proxy_model, get_model_for_task

SCRIPT_EDITOR_ID = None

# Does this text tell the reader to look at a picture? The stem list is the
# wider of the two on purpose: a stem that might need its figure keeps it.
STEM_CUE = re.compile(
    r"figure|photograph|photo\b|image|picture|pictured|shown|depicted|illustrat|"
    r"radiograph|x-ray|xray|roentgenogram|ultrasound|scan\b|smear|tracing|"
    r"electrocardiogram|\becg\b|\bekg\b|\bmri\b|\bct\b|growth chart",
    re.I)

# Evidence that the figure is the explanation's — the same words. An
# explanation saying "the rash shown is typical of…" is describing a figure
# printed with the critique, not with the vignette.
EXPLANATION_CUE = STEM_CUE

# The strongest signal in this bank, and one specific to the source: figures are
# labelled by where they are printed. "Item Q37A" is a figure beside the
# question; "Item C37B" is a figure beside the critique. The stem citing an
# Item Q settles it on its own, and only a critique cites an Item C.
# "ltem" with a lowercase L is the scanner reading a capital I.
STEM_FIGURE_LABEL = re.compile(r"\b[Il]tem\s*Q\s*\d", re.I)
EXPLANATION_FIGURE_LABEL = re.compile(r"\b[Il]tem\s*C\s*\d", re.I)

# How sure the model has to be before an image is removed on its say-so. Set
# from a spot check of the replies: every 0.95 verdict held up against the
# picture, while one of the three 0.85s was wrong — it called an unlabelled
# photograph of the vignette's own rash an explanation figure. Below the bar
# the question is left untouched and goes on the report for a person to read.
MIN_CONFIDENCE = 0.9

PROMPT = """You are auditing a pediatric board-exam question.

The image above was attached to this question during PDF extraction. It may
belong to the QUESTION STEM (a photograph the examinee must look at to answer)
or to the EXPLANATION (a figure printed with the answer critique, which would
give the answer away if shown during the test).

QUESTION STEM:
{stem}

EXPLANATION:
{explanation}

Decide which one the image illustrates. Weigh these signals:
- Does the stem describe a finding that this image shows? Then it is "stem".
- Is the image annotated, labelled, captioned, a diagram, a table, an algorithm,
  a graph, or does it show the diagnosis being named? Then it is "explanation".
- Is it a page header, logo, decorative rule, or blank? Then it is "neither".

Answer with ONLY this JSON, no markdown fence and no preamble:
{{"belongs": "stem" | "explanation" | "neither" | "unclear",
  "confidence": <0.0-1.0>,
  "reason": "<one short sentence>"}}"""


def image_bytes(key: str) -> tuple[bytes, str] | None:
    """The image and the media type a vision model will accept, or None."""
    data = storage_service.s3_object(key) or storage_service.load(key)
    if not data:
        return None
    media_type = mimetypes.guess_type(key)[0] or "image/jpeg"
    if key.lower().endswith((".jpx", ".jp2")) or media_type in ("image/jp2", "image/jpx"):
        # JPEG 2000 is not accepted by the vision APIs; PyMuPDF is already a
        # dependency of the PDF pipeline and re-encodes it without a new one.
        try:
            import fitz

            return fitz.Pixmap(data).tobytes("png"), "image/png"
        except Exception:
            return None
    return data, media_type


def classify(stem: str, explanation: str, key: str, model: str,
             api_key: str | None) -> dict:
    """Ask the vision model where the image belongs. Never raises."""
    loaded = image_bytes(key)
    if not loaded:
        return {"belongs": "unclear", "confidence": 0.0, "reason": "image could not be read"}
    data, media_type = loaded

    kwargs = {
        "model": _proxy_model(model),
        "temperature": 0,
        "messages": [{"role": "user", "content": [
            {"type": "image_url", "image_url": {
                "url": f"data:{media_type};base64,{base64.b64encode(data).decode()}"}},
            {"type": "text", "text": PROMPT.format(
                stem=(stem or "")[:4000] or "(empty)",
                explanation=(explanation or "")[:4000] or "(empty)")},
        ]}],
    }
    if api_key or settings.LITELLM_API_KEY:
        kwargs["api_key"] = api_key or settings.LITELLM_API_KEY
    if settings.LITELLM_API_BASE:
        kwargs["api_base"] = settings.LITELLM_API_BASE

    try:
        text = litellm.completion(**kwargs).choices[0].message.content or ""
    except Exception as error:  # a failed call must not decide anything
        return {"belongs": "unclear", "confidence": 0.0, "reason": f"model error: {error}"[:160]}

    match = re.search(r"\{.*\}", text, re.S)
    if not match:
        return {"belongs": "unclear", "confidence": 0.0, "reason": "unparseable reply"}
    try:
        answer = json.loads(match.group(0))
    except json.JSONDecodeError:
        return {"belongs": "unclear", "confidence": 0.0, "reason": "unparseable reply"}
    return {
        "belongs": str(answer.get("belongs", "unclear")).lower(),
        "confidence": float(answer.get("confidence") or 0.0),
        "reason": str(answer.get("reason", ""))[:200],
    }


def deterministic_verdict(stem: str, explanation: str) -> str:
    """"keep" when the stem needs the figure, "remove" when only the explanation
    mentions one, "unknown" when neither says anything."""
    stem, explanation = stem or "", explanation or ""
    if STEM_FIGURE_LABEL.search(stem) or STEM_CUE.search(stem):
        return "keep"
    if EXPLANATION_FIGURE_LABEL.search(explanation) or EXPLANATION_CUE.search(explanation):
        return "remove"
    return "unknown"


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--apply", action="store_true", help="write the changes")
    parser.add_argument("--no-vision", action="store_true",
                        help="deterministic rule only; ask no model")
    parser.add_argument("--limit", type=int, default=0, help="classify at most N unknowns")
    parser.add_argument("--model", help="vision model id (default: the extraction model)")
    parser.add_argument("--workers", type=int, default=4, help="parallel vision calls")
    parser.add_argument("--report", default="image_triage_report.txt",
                        help="where the human-review list is written")
    args = parser.parse_args()

    db = SessionLocal()
    try:
        questions = db.query(Question).filter(
            Question.image_path.isnot(None), Question.image_path != ""
        ).order_by(Question.id).all()

        keep, remove, unknown = [], [], []
        for question in questions:
            verdict = deterministic_verdict(question.question_text, question.explanation)
            {"keep": keep, "remove": remove, "unknown": unknown}[verdict].append(question)

        print(f"stem images        : {len(questions)}")
        print(f"  stem names a figure     : {len(keep)}  (left alone)")
        print(f"  only the explanation    : {len(remove)}  (remove)")
        print(f"  neither                 : {len(unknown)}")

        judged: dict[int, dict] = {}
        if unknown and not args.no_vision:
            model = args.model or get_model_for_task(db, "extraction")[0]
            api_key = None if args.model else get_model_for_task(db, "extraction")[1]
            batch = unknown[: args.limit] if args.limit else unknown
            print(f"\n  asking {model} about {len(batch)} image(s)…")
            payload = [(q.id, q.question_text, q.explanation, q.image_path) for q in batch]
            with ThreadPoolExecutor(max_workers=args.workers) as pool:
                results = list(pool.map(
                    lambda row: (row[0], classify(row[1], row[2], row[3], model, api_key)),
                    payload))
            judged = dict(results)

        # Only a confident "explanation" removes an image. "neither" is a page
        # header or a logo — junk, but junk that gives nothing away, so it stays
        # for a person to look at rather than being deleted on a guess.
        by_id = {q.id: q for q in unknown}
        vision_removals, uncertain = [], []
        for qid, answer in judged.items():
            if answer["belongs"] == "explanation" and answer["confidence"] >= MIN_CONFIDENCE:
                vision_removals.append((by_id[qid], answer))
            elif answer["belongs"] != "stem":
                uncertain.append((by_id[qid], answer))

        if judged:
            counts: dict[str, int] = {}
            for answer in judged.values():
                counts[answer["belongs"]] = counts.get(answer["belongs"], 0) + 1
            print("  model verdicts          : " +
                  ", ".join(f"{k}={v}" for k, v in sorted(counts.items())))
            print(f"  confident enough to act : {len(vision_removals)}")
            print(f"  for human review        : {len(uncertain)}")

        removals = [(q, {"reason": "only the explanation mentions a figure"}) for q in remove]
        removals += vision_removals

        for question, answer in removals:
            print(f"    remove  q{question.id}  {question.image_path}  — {answer['reason']}")
            if args.apply:
                _snapshot_question(db, question, SCRIPT_EDITOR_ID)
                question.image_path = None
        if args.apply:
            db.commit()

        if uncertain:
            lines = ["Questions whose stem image could not be classified confidently.",
                     "The image is still attached to the stem; decide by hand.",
                     ""]
            lines += [f"question {q.id}\t{q.image_path}\t{a['belongs']} "
                      f"({a['confidence']:.2f})\t{a['reason']}" for q, a in uncertain]
            with open(args.report, "w", encoding="utf-8") as handle:
                handle.write("\n".join(lines) + "\n")
            print(f"\n  {len(uncertain)} left for review -> {args.report}")

        print("\n" + ("APPLIED" if args.apply else "DRY RUN"))
        print(f"  images removed from stems : {len(removals)}")
        print(f"  images kept               : {len(questions) - len(removals)}")
        if not args.apply:
            print("\n  Re-run with --apply to write these changes.")
        return 0
    finally:
        db.close()


if __name__ == "__main__":
    sys.exit(main())