feat(copy-check): grade handwritten answer sheets with a vision LLM

ai_service/app/config.py

@@ -204,6 +204,24 @@ class Settings(BaseSettings):
     jwt_algorithm: str = "HS256"
     jwt_token_expiry_minutes: int = 43200  # 30 days in minutes (matching Java 2592000000ms)
 
+    # ---- Copy-check VISION grading ----------------------------------------
+    # When true (product default), the copy-check pipeline sends the actual
+    # answer-sheet page image(s) to a vision LLM and grades the handwriting from
+    # the image, using PaddleOCR only as an assistive text hint + annotation
+    # anchor. When false, it falls back to the legacy text-only OCR grader
+    # unchanged. Toggle with COPY_CHECK_VISION_GRADING=false.
+    copy_check_vision_grading: bool = True
+    # DPI to rasterize pages at for the LLM (150–200 is plenty for handwriting;
+    # higher just inflates the payload). Longest side is then downscaled to
+    # copy_check_vision_max_image_px before JPEG encoding.
+    copy_check_vision_dpi: int = 150
+    copy_check_vision_max_image_px: int = 2000
+    # Cost guards: pages sent per question, and total image sends per copy.
+    copy_check_vision_max_pages_per_question: int = 3
+    copy_check_vision_max_images_per_copy: int = 40
+    # Hard ceiling on pages rasterized from one PDF (pathological-document guard).
+    copy_check_vision_max_pages_per_copy: int = 50
+
     # Internal service-to-service auth.
     # Used by admin_core_service when calling /credits/v1/internal/* endpoints
     # (credit-pack purchase fulfillment from the Razorpay webhook handler).

ai_service/app/services/copy_check/grader.py

@@ -14,16 +14,26 @@
 
 import json
 import logging
+import os
 import re
 from typing import Any, Optional
 
+from ...config import get_settings
 from ..chat_llm_client import ChatLLMClient
-from .prompt_builder import GRADING_SYSTEM, build_grading_prompt
+from .prompt_builder import GRADING_SYSTEM, GRADING_SYSTEM_VISION, build_grading_prompt
 
 logger = logging.getLogger(__name__)
 
+# Text-only path (legacy: grades the printed-text OCR of the handwriting).
 DEFAULT_MODEL = "google/gemini-2.5-flash-lite"
 ESCALATION_MODEL = "google/gemini-2.5-flash"
+# Vision path: the LLM reads the actual page image. Defaults follow the repo's
+# "google/..." OpenRouter model-id convention (see config.py's validated list);
+# override via env without a code change.
+VISION_MODEL = os.getenv("COPY_CHECK_VISION_MODEL", "google/gemini-2.5-flash")
+VISION_ESCALATION_MODEL = os.getenv(
+    "COPY_CHECK_VISION_ESCALATION_MODEL", "google/gemini-2.5-pro"
+)
 ESCALATION_CONF_THRESHOLD = 0.60
 MAX_ESCALATIONS_PER_COPY = 2
 # Budget tuned for typical 8-question copies. Each grading call re-sends the
@@ -73,6 +83,10 @@
         self._prompt_tokens = 0
         self._completion_tokens = 0
         self._escalations_used = 0
+        # Vision-image accounting: how many page images we've attached across all
+        # grading calls for this copy, and the per-copy cap that bounds cost.
+        self._images_used = 0
+        self._max_images_per_copy = get_settings().copy_check_vision_max_images_per_copy
 
     def add_tokens(self, n: int) -> None:
         """External counter for non-grading calls (e.g. criteria generation
@@ -119,6 +133,10 @@
     def tokens_used(self) -> int:
         return self._tokens_used
 
+    @property
+    def images_used(self) -> int:
+        return self._images_used
+
     @property
     def prompt_tokens(self) -> int:
         return self._prompt_tokens
@@ -133,20 +151,34 @@
         rubric: dict[str, Any],
         layout_map: dict[str, Any],
         preferred_model: Optional[str] = None,
+        page_images: Optional[list[str]] = None,
     ) -> dict[str, Any]:
-        model = preferred_model or DEFAULT_MODEL
-        verdict = await self._call(question, rubric, layout_map, model)
+        """Grade one question.
+
+        `page_images`: base64 image data URLs of the page(s) this question's
+        answer occupies. When present, the LLM reads the handwriting from the
+        image and the OCR transcript is demoted to an assistive hint (vision
+        path). When None/empty, the legacy text-only path runs unchanged.
+        """
+        has_images = bool(page_images)
+        # Teacher's explicit model pick always wins; otherwise choose the vision
+        # default when images are attached, else the cheap text default.
+        model = preferred_model or (VISION_MODEL if has_images else DEFAULT_MODEL)
+        escalation_model = VISION_ESCALATION_MODEL if has_images else ESCALATION_MODEL
+        verdict = await self._call(question, rubric, layout_map, model, page_images)
         if (
             float(verdict.get("confidence", 0)) < ESCALATION_CONF_THRESHOLD
             and self._escalations_used < MAX_ESCALATIONS_PER_COPY
         ):
             self._escalations_used += 1
             logger.info(
                 "Escalating Q%s to %s (conf=%.2f)",
-                question["question_id"], ESCALATION_MODEL, verdict.get("confidence", 0),
+                question["question_id"], escalation_model, verdict.get("confidence", 0),
             )
             try:
-                verdict = await self._call(question, rubric, layout_map, ESCALATION_MODEL)
+                verdict = await self._call(
+                    question, rubric, layout_map, escalation_model, page_images,
+                )
             except Exception as e:
                 logger.warning(f"Escalation failed, keeping initial verdict: {e}")
         return verdict
@@ -157,15 +189,42 @@
         rubric: dict[str, Any],
         layout_map: dict[str, Any],
         model: str,
+        page_images: Optional[list[str]] = None,
     ) -> dict[str, Any]:
         if self._tokens_used >= FAIL_TOKENS_PER_COPY:
             raise RuntimeError(
                 f"copy-check token budget exhausted: {self._tokens_used} >= {FAIL_TOKENS_PER_COPY}"
             )
-        prompt = build_grading_prompt(question, rubric, layout_map)
+        # Enforce the per-copy image budget: trim (or drop) this call's images if
+        # we're at/over the cap, so a big copy can't run up unbounded vision cost.
+        imgs = list(page_images or [])
+        if imgs:
+            remaining = self._max_images_per_copy - self._images_used
+            if remaining <= 0:
+                logger.warning(
+                    "copy-check image budget exhausted (%d/%d sent); grading Q%s text-only",
+                    self._images_used, self._max_images_per_copy, question.get("question_id"),
+                )
+                imgs = []
+            elif len(imgs) > remaining:
+                logger.info(
+                    "copy-check image budget: trimming Q%s from %d to %d page image(s)",
+                    question.get("question_id"), len(imgs), remaining,
+                )
+                imgs = imgs[:remaining]
+        vision = bool(imgs)
+        system = GRADING_SYSTEM_VISION if vision else GRADING_SYSTEM
+        prompt = build_grading_prompt(question, rubric, layout_map, vision=vision)
+        user_msg: dict[str, Any] = {"role": "user", "content": prompt}
+        if vision:
+            # Shape expected by ChatLLMClient._convert_to_multimodal_messages:
+            # a user message with an `attachments` list of {"type","url"} dicts;
+            # the client turns each into an OpenAI "image_url" content part.
+            user_msg["attachments"] = [{"type": "image", "url": u} for u in imgs]
+            self._images_used += len(imgs)
         messages = [
-            {"role": "system", "content": GRADING_SYSTEM},
-            {"role": "user", "content": prompt},
+            {"role": "system", "content": system},
+            user_msg,
         ]
         try:
             response = await self.llm.chat_completion(

ai_service/app/services/copy_check/mathpix_fallback.py

@@ -14,9 +14,8 @@
 from pathlib import Path
 from typing import Any
 
-import httpx
-
 from ..mathpix_service import MathpixService
+from .page_images import download_pdf, rasterize_pages
 
 logger = logging.getLogger(__name__)
 
@@ -51,10 +50,13 @@ async def enrich_layout_for_math(self, pdf_url: str, layout_map: dict[str, Any])
 
         with tempfile.TemporaryDirectory(prefix="mathpix-crops-") as tmp:
             pdf_path = Path(tmp) / "input.pdf"
-            await _download(pdf_url, pdf_path)
-            page_imgs = await asyncio.get_event_loop().run_in_executor(
-                None, _rasterize_pages, pdf_path,
-            )
+            await download_pdf(pdf_url, pdf_path)
+            # 200 DPI matches the layout_map's box coordinates (full_res px), so
+            # crops line up. rasterize_pages returns [(page_id, img)]; keyed by
+            # page_id here for the crop lookup below.
+            page_imgs = dict(await asyncio.get_event_loop().run_in_executor(
+                None, rasterize_pages, pdf_path,
+            ))
             for page, line in flagged:
                 if not self.can_run:
                     logger.info("Mathpix budget exhausted (%d crops), skipping rest", self._used)
@@ -76,47 +78,6 @@ async def enrich_layout_for_math(self, pdf_url: str, layout_map: dict[str, Any])
         return layout_map
 
 
-async def _download(url: str, dest: Path) -> None:
-    async with httpx.AsyncClient(timeout=60, follow_redirects=True) as client:
-        async with client.stream("GET", url) as resp:
-            resp.raise_for_status()
-            with dest.open("wb") as f:
-                async for chunk in resp.aiter_bytes(chunk_size=1 << 16):
-                    f.write(chunk)
-
-
-def _rasterize_pages(pdf_path: Path) -> dict[str, Any]:
-    """Return {page_id: PIL.Image} for all pages of the PDF at 200 DPI.
-
-    Handles all PyMuPDF colorspaces (gray/RGB/CMYK) since PDFs in the wild
-    aren't always sRGB. A naive Image.frombytes("RGB", …) corrupts the buffer
-    for pix.n != 3.
-    """
-    import fitz  # PyMuPDF
-    from PIL import Image
-
-    out: dict[str, Any] = {}
-    doc = fitz.open(pdf_path)
-    try:
-        matrix = fitz.Matrix(200 / 72.0, 200 / 72.0)
-        for i, page in enumerate(doc):
-            pix = page.get_pixmap(matrix=matrix, alpha=False)
-            mode_map = {1: "L", 3: "RGB", 4: "CMYK"}
-            mode = mode_map.get(pix.n)
-            if mode is None:
-                # Drop alpha or unknown channels via an intermediate RGB pixmap.
-                pix_rgb = fitz.Pixmap(fitz.csRGB, pix)
-                img = Image.frombytes("RGB", (pix_rgb.width, pix_rgb.height), pix_rgb.samples)
-            else:
-                img = Image.frombytes(mode, (pix.width, pix.height), pix.samples)
-                if mode != "RGB":
-                    img = img.convert("RGB")
-            out[f"p{i + 1}"] = img
-    finally:
-        doc.close()
-    return out
-
-
 def _crop_to_base64(img, box: list[int]) -> str:
     from PIL import Image  # noqa: F401  — typed via duck typing