From b6ab16d6df85249370edf5188dcd0554b7551f59 Mon Sep 17 00:00:00 2001 From: gavinbee <29419542+gavinbee@users.noreply.github.com> Date: Thu, 28 May 2026 22:57:16 -0400 Subject: [PATCH] Add template detection from page 1 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - src/template_detect.py: classify page 1 of a PDF against the template registry via the vision model. Rasterizes page 1, lists every known template ID (implemented AND stubs) in the prompt, asks for {template_id, confidence}, and validates the result. - Recognising a stub is deliberate: if the model confidently says "this is a Natation Québec form", detect_template returns that id with is_implemented=False so the caller can surface the specific "not implemented yet, see issue X" message (via templates.get_template) rather than a vague "couldn't identify". TemplateDetection carries template_id, confidence, and is_implemented. - Error paths raise TemplateDetectionError with a --template override hint: model said "unknown", model returned an unregistered id (a hallucinated province), confidence below threshold (default 0.7), or a non-dict response. The unknown/unregistered case also links the new-issue page so users can contribute a sample PDF. - Refactor: promoted vision_extract._try_parse_json to public try_parse_json() (underscore dropped, name otherwise unchanged) so both modules strip markdown ```json fences the same way. Updated the two call sites in vision_extract. - tests/test_template_detect.py: 19 tests over a mocked client — prompt lists impl + stub ids with display names, happy path for an implemented template, stub-detected-but-unimplemented, every error path (unknown / unregistered / below-threshold / non-dict / missing confidence), custom threshold, and confidence clamping. - docs/architecture.md: marked implemented; noted CLI wiring is #10. Closes #9. Co-Authored-By: Claude Opus 4.7 --- docs/architecture.md | 3 +- src/template_detect.py | 184 ++++++++++++++++++++++++++++++++++ src/vision_extract.py | 8 +- tests/test_template_detect.py | 156 ++++++++++++++++++++++++++++ 4 files changed, 347 insertions(+), 4 deletions(-) create mode 100644 src/template_detect.py create mode 100644 tests/test_template_detect.py diff --git a/docs/architecture.md b/docs/architecture.md index abd4d11..d4f5cdc 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -18,7 +18,8 @@ | Ollama runtime + lifecycle | [`src/ollama_runtime.py`](../src/ollama_runtime.py) | implemented | | GPU detection + tier picker | [`src/gpu_detect.py`](../src/gpu_detect.py) | implemented | | Vision extraction | [`src/vision_extract.py`](../src/vision_extract.py) | implemented (not yet wired into the CLI — see #10) | -| Template detection, interactive review | — | pending — see [open issues](https://github.com/gavinbee/canswim-deck-eval-parser/issues) | +| Template detection | [`src/template_detect.py`](../src/template_detect.py) | implemented (not yet wired into the CLI — see #10) | +| Interactive review | — | pending — see [open issues](https://github.com/gavinbee/canswim-deck-eval-parser/issues) | ## How a parse runs (form-field path) diff --git a/src/template_detect.py b/src/template_detect.py new file mode 100644 index 0000000..9c99534 --- /dev/null +++ b/src/template_detect.py @@ -0,0 +1,184 @@ +"""Detect which provincial template a PDF uses, from page 1. + +The first thing the pipeline does with a scanned PDF is figure out which +provincial On-Deck Evaluation template it is — that choice drives the +field labels, the vision-prompt addendum, the date formats, and the +form-field widget map for everything downstream. + +We rasterize page 1, show it to the vision model, and ask it to classify +against the registry's known template IDs (implemented *and* stubs). +Recognising a stub is useful: if the model confidently says "this is a +Natation Québec form", the caller can surface the specific +"not implemented yet, see issue X" message via +``templates.get_template`` rather than a vague "couldn't identify". + +If the model isn't confident (below threshold) or says ``unknown``, we +raise :class:`TemplateDetectionError` telling the user to re-run with an +explicit ``--template`` override. +""" +from __future__ import annotations + +import logging +from dataclasses import dataclass +from pathlib import Path +from typing import Optional + +from . import pdf_io +from .templates import TEMPLATES, TEMPLATE_STUBS, known_template_ids +from .vision_extract import DEFAULT_VISION_MODEL, VisionClient, try_parse_json + +log = logging.getLogger(__name__) + + +# Below this, we refuse to guess and ask the user to pass --template. +DEFAULT_CONFIDENCE_THRESHOLD = 0.7 + + +@dataclass(frozen=True) +class TemplateDetection: + """Result of classifying page 1. + + ``is_implemented`` distinguishes a detected-and-parseable template + from a detected-but-stubbed one (e.g. Quebec), so the caller can + choose the right message. + """ + + template_id: str + confidence: float + is_implemented: bool + + +class TemplateDetectionError(RuntimeError): + """Raised when the template can't be identified confidently.""" + + +def _template_descriptions() -> dict[str, str]: + """``{id: display_name}`` for every known template (impl + stubs).""" + out: dict[str, str] = {tid: t.display_name for tid, t in TEMPLATES.items()} + out.update(TEMPLATE_STUBS) + return dict(sorted(out.items())) + + +def build_prompt() -> str: + """Classification prompt listing every known template ID.""" + lines = [ + f' - "{tid}": {name}' + for tid, name in _template_descriptions().items() + ] + catalog = "\n".join(lines) + return f"""\ +You are looking at page 1 of a Canadian swimming On-Deck Evaluation form. +Identify which provincial template it is. The known templates are: + +{catalog} + +Respond with a SINGLE JSON object, nothing else: + +{{"template_id": "", "confidence": <0.0-1.0>}} + +Base your answer on the form's title, the issuing organisation's name or +logo, the language, and the field layout. If you genuinely can't tell, +use "unknown" with a low confidence. Output ONLY the JSON object.""" + + +def detect_template( + pdf_path: str, + *, + client: VisionClient, + model: str = DEFAULT_VISION_MODEL, + threshold: float = DEFAULT_CONFIDENCE_THRESHOLD, + dpi: int = pdf_io.DEFAULT_DPI, +) -> TemplateDetection: + """Classify page 1 of ``pdf_path`` against the template registry. + + Args: + pdf_path: Path to the PDF. + client: Vision client (real ``ollama.Client`` or a test stub). + model: Vision model tag. + threshold: Minimum confidence to accept a classification. + dpi: Rasterization DPI for page 1. + + Returns: + A :class:`TemplateDetection`. The id may be a not-yet-implemented + stub — the caller decides how to handle that (typically by + calling ``templates.get_template`` which raises a helpful + ``NotImplementedError``). + + Raises: + TemplateDetectionError: if the model returns ``unknown``, an + unrecognised id, or a confidence below ``threshold``. + """ + png = pdf_io.rasterize_page(pdf_path, 0, dpi=dpi) + prompt = build_prompt() + + response = client.generate( + model=model, + prompt=prompt, + images=[png], + format="json", + options={"temperature": 0}, + ) + text = getattr(response, "response", None) + if text is None and isinstance(response, dict): + text = response.get("response") + parsed = try_parse_json(text or "") + + detection = _interpret(parsed, threshold=threshold, pdf_path=pdf_path) + log.info( + "Detected template %s (confidence %.2f) for %s", + detection.template_id, detection.confidence, Path(pdf_path).name, + ) + return detection + + +def _interpret( + parsed: object, + *, + threshold: float, + pdf_path: str, +) -> TemplateDetection: + """Validate the model's JSON and turn it into a TemplateDetection.""" + name = Path(pdf_path).name + override_hint = ( + "Re-run with an explicit template, e.g. " + "`--template swim_ontario_v1`." + ) + + if not isinstance(parsed, dict): + raise TemplateDetectionError( + f"Could not identify the template for {name}: the model did " + f"not return a usable classification. {override_hint}" + ) + + template_id = parsed.get("template_id") + confidence = _clamp(parsed.get("confidence")) + + if template_id == "unknown" or template_id not in known_template_ids(): + raise TemplateDetectionError( + f"Could not confidently identify the template for {name} " + f"(model said {template_id!r}). {override_hint} " + "If this is a province we don't support yet, please file an " + "issue with a sample PDF: " + "https://github.com/gavinbee/canswim-deck-eval-parser/issues/new" + ) + + if confidence < threshold: + raise TemplateDetectionError( + f"Low-confidence template detection for {name}: best guess " + f"was {template_id!r} at {confidence:.2f} (threshold " + f"{threshold:.2f}). {override_hint}" + ) + + return TemplateDetection( + template_id=template_id, + confidence=confidence, + is_implemented=template_id in TEMPLATES, + ) + + +def _clamp(value: object) -> float: + try: + c = float(value) # type: ignore[arg-type] + except (TypeError, ValueError): + return 0.0 + return max(0.0, min(1.0, c)) diff --git a/src/vision_extract.py b/src/vision_extract.py index 9895eb2..7b81692 100644 --- a/src/vision_extract.py +++ b/src/vision_extract.py @@ -258,13 +258,13 @@ def _call_model( """Call the model and return parsed JSON, retrying once on failure.""" prompt = build_prompt(template, filename) text = _generate(client, model, prompt, png_bytes) - parsed = _try_parse_json(text) + parsed = try_parse_json(text) if _is_structural(parsed): return parsed log.warning("Vision model returned unparseable output; retrying once.") text = _generate(client, model, prompt + _RETRY_NUDGE, png_bytes) - parsed = _try_parse_json(text) + parsed = try_parse_json(text) if _is_structural(parsed): return parsed @@ -296,10 +296,12 @@ def _generate( return text or "" -def _try_parse_json(text: str) -> Any: +def try_parse_json(text: str) -> Any: """Parse model text to JSON, tolerating markdown code fences. Returns the parsed object, or ``None`` if it isn't valid JSON. + Shared with ``src.template_detect`` so both modules handle the + ```json ... ``` fences some models emit the same way. """ if not text: return None diff --git a/tests/test_template_detect.py b/tests/test_template_detect.py new file mode 100644 index 0000000..2fc6bcc --- /dev/null +++ b/tests/test_template_detect.py @@ -0,0 +1,156 @@ +"""Tests for src/template_detect.py. + +The vision client is mocked. We patch pdf_io.rasterize_page so no real +PDF is needed for the classification-logic tests. +""" +from __future__ import annotations + +import json +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +from src import template_detect +from src.template_detect import ( + DEFAULT_CONFIDENCE_THRESHOLD, + TemplateDetection, + TemplateDetectionError, + build_prompt, + detect_template, +) + + +def _client_returning(payload) -> MagicMock: + client = MagicMock() + text = payload if isinstance(payload, str) else json.dumps(payload) + client.generate.return_value = SimpleNamespace(response=text) + return client + + +def _detect(payload, **kwargs): + client = _client_returning(payload) + with patch("src.template_detect.pdf_io.rasterize_page", return_value=b"PNG"): + return detect_template("scan.pdf", client=client, model="m", **kwargs), client + + +# --------------------------------------------------------------------------- +# Prompt +# --------------------------------------------------------------------------- + + +class TestBuildPrompt: + def test_lists_implemented_and_stub_ids(self): + prompt = build_prompt() + assert "swim_ontario_v1" in prompt + assert "swim_quebec_v1" in prompt + assert "swim_alberta_v1" in prompt + assert "swim_bc_v1" in prompt + + def test_includes_display_names(self): + prompt = build_prompt() + assert "Swim Ontario On-Deck Evaluation" in prompt + assert "Natation Québec" in prompt + + def test_mentions_unknown_escape_hatch(self): + prompt = build_prompt() + assert "unknown" in prompt + + +# --------------------------------------------------------------------------- +# Happy path +# --------------------------------------------------------------------------- + + +class TestDetectImplementedTemplate: + def test_returns_detection_for_ontario(self): + detection, client = _detect( + {"template_id": "swim_ontario_v1", "confidence": 0.96} + ) + assert detection.template_id == "swim_ontario_v1" + assert detection.confidence == 0.96 + assert detection.is_implemented is True + client.generate.assert_called_once() + + def test_passes_image_and_json_format(self): + _, client = _detect({"template_id": "swim_ontario_v1", "confidence": 0.9}) + kwargs = client.generate.call_args.kwargs + assert kwargs["images"] == [b"PNG"] + assert kwargs["format"] == "json" + assert kwargs["options"]["temperature"] == 0 + + +class TestDetectStubTemplate: + def test_stub_detected_but_marked_unimplemented(self): + # The model confidently recognises a Quebec form. We return it + # (so the caller can show the specific "not yet implemented" + # message) but flag is_implemented=False. + detection, _ = _detect( + {"template_id": "swim_quebec_v1", "confidence": 0.93} + ) + assert detection.template_id == "swim_quebec_v1" + assert detection.is_implemented is False + + +# --------------------------------------------------------------------------- +# Error paths +# --------------------------------------------------------------------------- + + +class TestDetectionErrors: + def test_unknown_raises(self): + with pytest.raises(TemplateDetectionError) as exc: + _detect({"template_id": "unknown", "confidence": 0.1}) + assert "--template" in str(exc.value) + assert "file an issue" in str(exc.value) + + def test_unrecognised_id_raises(self): + # Model hallucinates a province we don't have registered. + with pytest.raises(TemplateDetectionError): + _detect({"template_id": "swim_yukon_v1", "confidence": 0.99}) + + def test_below_threshold_raises(self): + with pytest.raises(TemplateDetectionError) as exc: + _detect({"template_id": "swim_ontario_v1", "confidence": 0.5}) + assert "Low-confidence" in str(exc.value) + assert "0.50" in str(exc.value) + + def test_custom_threshold_respected(self): + # 0.5 passes if the caller lowers the bar. + detection, _ = _detect( + {"template_id": "swim_ontario_v1", "confidence": 0.55}, + threshold=0.5, + ) + assert detection.template_id == "swim_ontario_v1" + + def test_default_threshold_value(self): + # Guard against an accidental change to the default. + assert DEFAULT_CONFIDENCE_THRESHOLD == 0.7 + + def test_non_dict_response_raises(self): + with pytest.raises(TemplateDetectionError) as exc: + _detect("not json at all") + assert "did not return a usable classification" in str(exc.value) + + def test_missing_confidence_treated_as_zero(self): + # No confidence key → 0.0 → below threshold → raises. + with pytest.raises(TemplateDetectionError): + _detect({"template_id": "swim_ontario_v1"}) + + +# --------------------------------------------------------------------------- +# Confidence clamping +# --------------------------------------------------------------------------- + + +class TestClamp: + @pytest.mark.parametrize("raw, expected", [ + (0.8, 0.8), + (1.4, 1.0), + (-0.2, 0.0), + ("0.9", 0.9), + (None, 0.0), + ("garbage", 0.0), + ]) + def test_clamp(self, raw, expected): + assert template_detect._clamp(raw) == expected