llm.py

"""
LLM verification — Gemma 4 12B on transformers + ZeroGPU (GPU, no external API).

Runs Google's Gemma 4 12B (a <=32B open model) on the Space's GPU via Hugging
Face transformers, decorated with @spaces.GPU for ZeroGPU. It:
  1. OCRs each uploaded file (images, or PDF pages rendered to images), and
  2. Checks the document set against the required-documents checklist and the STG
     content rules for the selected package + stage.

The model is loaded once at module level on CUDA (per ZeroGPU guidance); the GPU
is attached only inside the @spaces.GPU-decorated generate function. PDFs are
rendered to images with PyMuPDF (no poppler). The model returns strict JSON,
which we render into a readable verdict.

Model is configurable via CLAIMREADY_MODEL_ID (default google/gemma-4-12B-it).
"""

import io
import json
import os

import fitz  # PyMuPDF
import spaces
import transformers
from PIL import Image
from transformers import AutoProcessor

MODEL_ID = os.environ.get("CLAIMREADY_MODEL_ID", "google/gemma-3-12b-it")

_MAX_IMAGE_SIDE = 896      # one vision tile per image
_MAX_IMAGES = 12           # safety cap (PDF pages + image files)
_PDF_DPI = 150
_IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp", ".tiff", ".tif"}


class LLMConfigError(RuntimeError):
    """Raised when the model cannot be loaded."""


def _model_class():
    """Pick a vision-language Auto-model class available in this transformers version.

    AutoModelForImageTextToText is the correct class for Gemma 3 vision.
    """
    for name in ("AutoModelForImageTextToText", "AutoModelForMultimodalLM", "AutoModelForCausalLM"):
        cls = getattr(transformers, name, None)
        if cls is not None:
            return cls
    raise LLMConfigError("No suitable Auto model class found in transformers.")


# --- Load once at module level (ZeroGPU: model goes on CUDA here) -------------
try:
    processor = AutoProcessor.from_pretrained(MODEL_ID)
    model = _model_class().from_pretrained(MODEL_ID, dtype="auto", device_map="cuda")
except Exception as e:  # surfaced on first verify() call
    processor = None
    model = None
    _LOAD_ERROR = e
else:
    _LOAD_ERROR = None


# --- File -> PIL images -------------------------------------------------------

def _downscale(img):
    img = img.convert("RGB")
    w, h = img.size
    scale = min(1.0, _MAX_IMAGE_SIDE / max(w, h))
    if scale < 1.0:
        img = img.resize((max(1, int(w * scale)), max(1, int(h * scale))))
    return img


def _file_to_images(path):
    ext = os.path.splitext(path)[1].lower()
    if ext == ".pdf":
        imgs = []
        with fitz.open(path) as doc:
            for page in doc:
                pix = page.get_pixmap(dpi=_PDF_DPI)
                imgs.append(_downscale(Image.open(io.BytesIO(pix.tobytes("png")))))
        return imgs
    try:
        return [_downscale(Image.open(path))]
    except Exception:
        return []


# --- Prompt -------------------------------------------------------------------

def _checklist_text(docs, rules):
    lines = ["REQUIRED DOCUMENTS:"]
    for d in docs:
        lines.append(f"- id={d['id']} | {d['label']} — verify: {d['verify']}")
    if rules:
        lines.append("")
        lines.append("CONTENT RULES (check against the text/values in the documents):")
        for r in rules:
            lines.append(f"- id={r['id']} | {r['check']}")
    return "\n".join(lines)


def _system_prompt():
    return (
        "You are a PMJAY (Ayushman Bharat) claims processing doctor reviewing a "
        "hospital's pre-authorization / claim document set against the Standard "
        "Treatment Guidelines BEFORE it is submitted to the government portal.\n\n"
        "You will be given: (a) the package + stage, (b) a checklist of required "
        "documents, (c) content rules to confirm, and (d) the uploaded files as "
        "images (PDF pages are also given as images). Read every image carefully "
        "(OCR the text).\n\n"
        "For EACH required document, decide whether it is present among the "
        "uploaded files and whether its content matches what must be verified. "
        "For EACH content rule, decide pass / fail / unclear based on the text or "
        "values you can read. Do not invent evidence — if you cannot read it, say "
        "so and mark it unclear.\n\n"
        "STRICT RULE EVALUATION — apply these without exception:\n"
        "1. NUMERIC THRESHOLDS: When a rule states a numeric limit (e.g. '>= 101°F', "
        "'>= 38.3°C', '< 7 g/dl', '> 55 years'), find the actual value(s) and compare "
        "EXACTLY. For EACH value, FIRST write the comparison out in the evidence in "
        "the form '<value> vs <threshold>: <value> is greater/less than <threshold>, "
        "so it MEETS/does NOT meet the rule', THEN decide. Be careful with direction: "
        "a LARGER number is NEVER 'below' a smaller threshold.\n"
        "   Worked examples (follow this logic exactly):\n"
        "   • 102.4°F, 102.0°F, 101.6°F each satisfy '>= 101°F' (each is greater than "
        "101) -> PASS.\n"
        "   • 100.8°F does NOT satisfy '>= 101°F' (100.8 is less than 101) -> FAIL.\n"
        "   • Hb 6.2 satisfies '< 7 g/dl' (6.2 is less than 7) -> PASS;  Hb 7.2 does "
        "NOT (7.2 is not less than 7) -> FAIL.\n"
        "   Do NOT round, approximate, or give the benefit of the doubt.\n"
        "2. DURATION / FREQUENCY: When a rule includes a time span (e.g. 'for more "
        "than 2 days'), you must find explicit dated evidence covering that span. A "
        "single reading never satisfies a multi-day requirement — mark it 'unclear' "
        "if duration is not documented.\n"
        "3. PASS ONLY ON FULL MATCH: Mark 'pass' only when you can cite a specific, "
        "readable value that meets EVERY part of the criterion. If a value is present "
        "but below/outside the bar, mark 'fail'. If you cannot read or locate the "
        "value, mark 'unclear'. Never 'pass' by assumption.\n"
        "4. UNITS: Convert units when needed to compare (°F<->°C, etc.), but always "
        "state the actual value you read.\n"
        "5. EVIDENCE: In every rule's 'evidence', quote the exact value you read and "
        "state explicitly whether it meets the threshold (e.g. \"Temp 100.8°F read; "
        "below the 101°F threshold -> fail\").\n"
        "6. STATUS MUST MATCH EVIDENCE: The 'status' you assign to a rule (and the "
        "'present' boolean for a document) MUST agree with the conclusion in its own "
        "'evidence'. If the evidence says the value is below/outside the bar or ends "
        "in '-> fail', the status is 'fail' — NEVER 'pass'. If the value cannot be "
        "read/located, the status is 'unclear'. A 'pass' whose evidence concludes "
        "'not met' is forbidden.\n\n"
        "Respond with ONLY a JSON object (no markdown fences, no prose) matching:\n"
        "{\n"
        '  "documents": [{"id": str, "present": bool, "confidence": "high|medium|low", "evidence": str}],\n'
        '  "rules": [{"id": str, "status": "pass|fail|unclear", "evidence": str}],\n'
        '  "overall": "ready|not_ready|needs_review",\n'
        '  "summary": str\n'
        "}\n"
        "Use the exact id values given. 'evidence' is a short phrase citing which "
        "file and what you saw. Keep 'summary' to 1-3 sentences."
    )


def _build_messages(pkg, stage_label, docs, rules, files):
    intro = (
        f"{_system_prompt()}\n\n"
        f"PACKAGE: {pkg['code']} — {pkg['name']} ({pkg.get('procedure', '')})\n"
        f"STAGE: {stage_label}\n\n"
        f"{_checklist_text(docs, rules)}\n\n"
        f"UPLOADED FILES: the documents follow as images."
    )
    content = [{"type": "text", "text": intro}]
    n = 0
    for path in files:
        if n >= _MAX_IMAGES:
            break
        for img in _file_to_images(path):
            if n >= _MAX_IMAGES:
                break
            content.append({"type": "image", "image": img})
            n += 1
    return [{"role": "user", "content": content}]


@spaces.GPU(duration=120)
def _generate(messages, max_new_tokens):
    inputs = processor.apply_chat_template(
        messages,
        tokenize=True,
        return_dict=True,
        return_tensors="pt",
        add_generation_prompt=True,
    ).to(model.device)
    out = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False)
    gen = out[0][inputs["input_ids"].shape[-1]:]
    return processor.decode(gen, skip_special_tokens=True)


def _parse_json(text):
    text = (text or "").strip()
    if text.startswith("```"):
        text = text.strip("`")
        if text.lower().startswith("json"):
            text = text[4:]
        text = text.strip()
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        start, end = text.find("{"), text.rfind("}")
        if start != -1 and end != -1 and end > start:
            try:
                return json.loads(text[start : end + 1])
            except json.JSONDecodeError:
                return None
    return None


def verify(pkg, stage_label, docs, rules, files):
    """Run GPU verification. Returns a markdown report string."""
    if model is None:
        raise LLMConfigError(
            f"Model failed to load: {type(_LOAD_ERROR).__name__}: {_LOAD_ERROR}"
        )
    messages = _build_messages(pkg, stage_label, docs, rules, files)
    raw = _generate(messages, int(os.environ.get("CLAIMREADY_MAX_TOKENS", "1024")))
    parsed = _parse_json(raw)
    return _render(pkg, stage_label, docs, rules, parsed, raw)


# --- Rendering -----------------------------------------------------------------

_DOC_ICON = {True: "✅", False: "❌"}
_RULE_ICON = {"pass": "✅", "fail": "❌", "unclear": "⚠️"}
_OVERALL = {
    "ready": "✅ **Looks complete** — no blocking gaps spotted. _A reviewer should confirm before submission._",
    "not_ready": "⚠️ **Possible gaps found** — review the items below before submitting.",
    "needs_review": "⚠️ **Needs manual review** — some items could not be confirmed.",
}


def _doc_present(r):
    """A low-confidence 'present' is treated as NOT present — safer for a
    compliance pre-check (the reviewer should re-check it)."""
    return bool(r.get("present")) and (r.get("confidence") or "").lower() != "low"


def _render(pkg, stage_label, docs, rules, parsed, raw):
    if parsed is None:
        return (
            "## ⚠️ Could not parse the model's response\n\n"
            "The verification model replied, but not in the expected format. "
            "Raw response below:\n\n"
            f"```\n{(raw or '').strip()[:2000]}\n```"
        )

    rule_by_id = {r["id"]: r for r in rules}
    res_docs = {d.get("id"): d for d in parsed.get("documents", []) if isinstance(d, dict)}
    res_rules = {r.get("id"): r for r in parsed.get("rules", []) if isinstance(r, dict)}

    lines = [
        "## Assistive review",
        f"**{pkg['code']} — {pkg['name']}** · {stage_label}",
        "",
        _OVERALL.get(parsed.get("overall"), "⚠️ **Review the findings below.**"),
    ]
    if parsed.get("summary"):
        lines += ["", f"> {parsed['summary']}"]

    lines += ["", "### Documents"]
    for d in docs:
        r = res_docs.get(d["id"], {})
        present = _doc_present(r)
        icon = _DOC_ICON[present]
        conf = (r.get("confidence") or "").lower()
        conf_txt = f" _(confidence: {conf})_" if conf else ""
        lines.append(f"- {icon} **{d['label']}**{conf_txt}")
        if r.get("evidence"):
            ev = r["evidence"]
            if bool(r.get("present")) and conf == "low":
                ev += "  — _low confidence; flagged for manual verification_"
            lines.append(f"   - {ev}")
        elif not r:
            lines.append("   - _Not assessed by the model._")

    if rules:
        lines += ["", "### Content checks"]
        for rule in rules:
            r = res_rules.get(rule["id"], {})
            status = r.get("status", "unclear")
            icon = _RULE_ICON.get(status, "⚠️")
            lines.append(f"- {icon} {rule['check']} — _{status}_")
            if r.get("evidence"):
                lines.append(f"   - {r['evidence']}")

    missing = [d["label"] for d in docs if not _doc_present(res_docs.get(d["id"], {}))]
    failed = [
        rule_by_id[rid]["check"]
        for rid, r in res_rules.items()
        if r.get("status") == "fail" and rid in rule_by_id
    ]
    if missing or failed:
        lines += ["", "### ⚠️ Action needed before submission"]
        for m in missing:
            lines.append(f"- Missing / unreadable document: **{m}**")
        for f in failed:
            lines.append(f"- Failed content check: {f}")

    return "\n".join(lines)