From fd6196b47c791bfc999683a1960b8e1e1fd84612 Mon Sep 17 00:00:00 2001 From: gavinbee Date: Wed, 27 May 2026 23:57:55 -0400 Subject: [PATCH] Add vision extraction for scanned PDFs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - src/vision_extract.py: per-page extraction via a local vision model (Qwen2.5-VL through Ollama). Rasterizes each page, prompts the model for a JSON object, parses it into the same PageExtraction shape the form-field path emits — so merge / output stay path-agnostic. - Prompt construction: a canonical-schema JSON skeleton built from src.schema constants (can't drift), plus the template's vision_prompt_addendum, plus the source filename so the model can resolve session_number from the form text, the filename, or both (and report which via the "source" field). - Model call: client.generate(format="json", options temperature=0) with the page PNG. One retry on a structural parse failure with a clarifying nudge appended; raises VisionExtractionError if the second attempt also fails. Tolerates markdown code fences some models wrap JSON in despite instructions. - Parsing: each field becomes a typed FieldValue with clamped [0,1] confidence; null values are dropped; bare scalars (model ignored the {value, confidence} shape) get a default 0.5 confidence. successful is coerced to true/false/null defensively — a model emitting the string "false" must NOT become a truthy non-empty string downstream (regression-tested). session_number preserves its source provenance. - Caching: raw model responses are written to a .raw.json sidecar keyed by page number and tagged with the model. Re-runs reuse the cache (fast iteration on parsing logic without re-invoking the slow model); use_cache=False (the CLI's --no-cache) forces a fresh call; a cache built with a different model is ignored. - VisionClient is a Protocol so the module doesn't hard-depend on the ollama package surface; tests inject a fake client returning canned JSON. - tests/test_vision_extract.py: 41 tests over a mocked client. Prompt content, good-response parsing, successful/confidence coercion (incl. the stringy-false regression), retry-then-succeed and retry-then-fail, markdown-fence stripping, missing-rows retry, and the cache round-trip (write → reuse → --no-cache → model-mismatch). - docs/architecture.md: marked implemented; noted it's not yet wired into the CLI (that's #10). Closes #8. Co-Authored-By: Claude Opus 4.7 --- docs/architecture.md | 3 +- src/vision_extract.py | 522 +++++++++++++++++++++++++++++++++++ tests/test_vision_extract.py | 345 +++++++++++++++++++++++ 3 files changed, 869 insertions(+), 1 deletion(-) create mode 100644 src/vision_extract.py create mode 100644 tests/test_vision_extract.py diff --git a/docs/architecture.md b/docs/architecture.md index cbf3e8f..abd4d11 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -17,7 +17,8 @@ | End-to-end CLI (form-field path) | [`main.py`](../main.py) | implemented | | 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, template detection, interactive review | — | pending — see [open issues](https://github.com/gavinbee/canswim-deck-eval-parser/issues) | +| 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) | ## How a parse runs (form-field path) diff --git a/src/vision_extract.py b/src/vision_extract.py new file mode 100644 index 0000000..9895eb2 --- /dev/null +++ b/src/vision_extract.py @@ -0,0 +1,522 @@ +"""Vision-LLM extraction for scanned / flat PDFs. + +For PDFs without fillable widgets, each page is rasterized to PNG and +handed to a local vision model (Qwen2.5-VL via Ollama). The model returns +a JSON object per page that we parse into the same ``PageExtraction`` +shape the form-field path emits, so ``merge`` and ``output`` don't care +which path produced a page. + +What this module owns: + +* **Prompt construction** — a canonical-schema JSON skeleton plus the + detected template's ``vision_prompt_addendum`` plus the source + filename (so the model can resolve ``session_number`` from the form + text, the filename, or both). +* **The model call** — ``client.generate(..., format="json")`` at + temperature 0, with the page PNG as the image. +* **Parsing + validation** — turn the model's JSON into typed + ``FieldValue`` objects, coercing ``successful`` to bool/null and + capturing per-field confidence. One retry on a structural parse + failure with a clarifying nudge appended. +* **The ``.raw.json`` cache** — raw model responses are written to a + sidecar so re-runs (and parsing-logic iteration) don't re-invoke the + slow model. ``use_cache=False`` (the CLI's ``--no-cache``) forces a + fresh call. +""" +from __future__ import annotations + +import json +import logging +from pathlib import Path +from typing import Any, Optional, Protocol + +from . import pdf_io, schema as s +from .form_extract import PageExtraction +from .templates.base import Template + +log = logging.getLogger(__name__) + + +DEFAULT_VISION_MODEL = "qwen2.5vl:7b" + +# Default confidence when the model gives a value but omits a confidence. +# Mid-scale so it neither hides (too high) nor screams (too low) — the +# missing confidence itself is a mild signal something's off. +_DEFAULT_CONFIDENCE = 0.5 + +# Truthy / falsy spellings the model might emit for ``successful`` before +# we've forced it to a real boolean. +_TRUE_WORDS = {"true", "yes", "y", "1", "signed", "signed off", "pass", "passed"} +_FALSE_WORDS = {"false", "no", "n", "0", "fail", "failed", "unsuccessful"} + + +# --------------------------------------------------------------------------- +# Minimal client protocol (so tests can inject a stub) +# --------------------------------------------------------------------------- + + +class VisionClient(Protocol): + """The slice of ``ollama.Client`` this module uses.""" + + def generate(self, *args: Any, **kwargs: Any) -> Any: ... + + +class VisionExtractionError(RuntimeError): + """Raised when the model output can't be parsed even after a retry.""" + + +# --------------------------------------------------------------------------- +# Prompt construction +# --------------------------------------------------------------------------- + + +# The JSON skeleton the model is asked to fill. Built from canonical +# schema names so it can never drift from src.schema. +def _json_skeleton() -> str: + def field(extra: str = "") -> str: + return '{"value": , "confidence": <0.0-1.0>' + extra + "}" + + meet = ",\n ".join(f'"{k}": {field()}' for k in s.MEET_FIELDS) + session_keys = [ + s.COMPETITION_COORDINATOR, s.CC_LEVEL, s.DATE_SESSION, s.SESSION_NUMBER, + ] + session_lines = [] + for k in session_keys: + if k == s.SESSION_NUMBER: + session_lines.append( + f'"{k}": {{"value": , "confidence": <0.0-1.0>, ' + '"source": "form" | "filename" | "form+filename" | "unknown"}' + ) + else: + session_lines.append(f'"{k}": {field()}') + session = ",\n ".join(session_lines) + + row_keys = [ + s.OFFICIAL_NAME, s.CLUB, s.POSITION, s.LANE_NUMBER, + s.TIMES_WORKED_POSITION, s.MENTOR, s.LEVEL, + ] + row = ",\n ".join(f'"{k}": {field()}' for k in row_keys) + successful = ( + f'"{s.SUCCESSFUL}": {{"value": true | false | null, ' + '"confidence": <0.0-1.0>, "rationale": }' + ) + + return f"""{{ + "meet": {{ + {meet} + }}, + "session": {{ + {session} + }}, + "rows": [ + {{ + {row}, + {successful} + }} + ] +}}""" + + +_BASE_INSTRUCTIONS = """\ +You are extracting structured data from one page of a swimming On-Deck \ +Evaluation form. The page image is attached. + +Return a SINGLE JSON object, nothing else, matching exactly this shape \ +(one entry in "rows" per official visible on the page; omit blank rows): + +{skeleton} + +Rules: +- Every field is an object with a "value" and a "confidence" in [0.0, 1.0]. + Use null for "value" when a field is blank or unreadable. Set confidence + to how sure you are of the value you read — be honest; low confidence on + messy handwriting is expected and useful. +- "meet" fields are the same for the whole page (competition name, host + club, COC). +- "session" fields describe this session. For "session_number", read it + from the form's date/session text if present; otherwise infer it from + the source filename "{filename}"; set "source" to where you got it + ("form", "filename", "form+filename", or "unknown" if you truly can't + tell). +- "rows" is one object per official. Read names, clubs, positions, lane + numbers, etc. as written. +- "successful": judge the WHOLE ROW, not just the initials cell. Emit true + if the official was clearly signed off, false if clearly not, null if + genuinely ambiguous. Put your reasoning in "rationale". +{addendum} + +Output ONLY the JSON object.""" + + +def build_prompt(template: Template, filename: str) -> str: + """Assemble the page-extraction prompt for a given template + file.""" + addendum = template.vision_prompt_addendum.strip() + addendum_block = ( + f"\nTemplate-specific notes:\n{addendum}" if addendum else "" + ) + return _BASE_INSTRUCTIONS.format( + skeleton=_json_skeleton(), + filename=filename, + addendum=addendum_block, + ) + + +# --------------------------------------------------------------------------- +# Public extraction entry points +# --------------------------------------------------------------------------- + + +def extract_pdf( + pdf_path: str, + template: Template, + *, + client: VisionClient, + model: str = DEFAULT_VISION_MODEL, + cache_path: Optional[str | Path] = None, + use_cache: bool = True, + dpi: int = pdf_io.DEFAULT_DPI, +) -> list[PageExtraction]: + """Vision-extract every page of ``pdf_path``. + + Args: + pdf_path: Path to the (scanned / flat) PDF. + template: Detected template — supplies the prompt addendum. + client: Anything with a ``generate`` method (real ``ollama.Client`` + or a test stub). + model: Ollama model tag. + cache_path: Where to read/write the ``.raw.json`` sidecar. If + ``None``, no caching. + use_cache: If True and a matching cache exists, re-parse cached + raw responses instead of calling the model. + dpi: Rasterization DPI. + + Returns: + One ``PageExtraction`` per page. + """ + filename = Path(pdf_path).name + n_pages = pdf_io.page_count(pdf_path) + + cached_raw = _load_cache(cache_path, model) if (use_cache and cache_path) else None + + pages: list[PageExtraction] = [] + raw_by_page: dict[str, Any] = {} + + for i in range(n_pages): + page_no = i + 1 + cached = cached_raw.get(str(page_no)) if cached_raw else None + if cached is not None: + log.debug("Using cached vision response for page %d", page_no) + raw = cached + else: + png = pdf_io.rasterize_page(pdf_path, i, dpi=dpi) + raw = _call_model(client, model, template, filename, png) + raw_by_page[str(page_no)] = raw + pages.append(_parse_response(raw, page_number=page_no)) + + if cache_path: + _write_cache(cache_path, model, filename, raw_by_page) + + return pages + + +def extract_page( + png_bytes: bytes, + template: Template, + filename: str, + page_number: int, + *, + client: VisionClient, + model: str = DEFAULT_VISION_MODEL, +) -> tuple[PageExtraction, dict]: + """Vision-extract a single rasterized page. + + Returns the parsed ``PageExtraction`` and the raw model JSON (so + callers can cache it). + """ + raw = _call_model(client, model, template, filename, png_bytes) + return _parse_response(raw, page_number=page_number), raw + + +# --------------------------------------------------------------------------- +# Model invocation +# --------------------------------------------------------------------------- + + +_RETRY_NUDGE = ( + "\n\nYour previous reply was not valid JSON matching the requested " + "shape. Reply with ONLY the JSON object, no prose, no markdown fences." +) + + +def _call_model( + client: VisionClient, + model: str, + template: Template, + filename: str, + png_bytes: bytes, +) -> dict: + """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) + 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) + if _is_structural(parsed): + return parsed + + raise VisionExtractionError( + "Vision model did not return parseable JSON after a retry. " + f"Last response (truncated): {text[:300]!r}" + ) + + +def _generate( + client: VisionClient, + model: str, + prompt: str, + png_bytes: bytes, +) -> str: + """One ``generate`` call. Returns the response text.""" + response = client.generate( + model=model, + prompt=prompt, + images=[png_bytes], + format="json", + options={"temperature": 0}, + ) + # ollama-python returns a GenerateResponse with a ``.response`` attr; + # also dict-accessible. Support both for forward/backward compat. + text = getattr(response, "response", None) + if text is None and isinstance(response, dict): + text = response.get("response") + return text or "" + + +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. + """ + if not text: + return None + stripped = text.strip() + # Strip ```json ... ``` fences some models add despite instructions. + if stripped.startswith("```"): + stripped = stripped.split("```", 2) + stripped = stripped[1] if len(stripped) > 1 else "" + if stripped.startswith("json"): + stripped = stripped[len("json"):] + stripped = stripped.strip().rstrip("`").strip() + try: + return json.loads(stripped) + except (json.JSONDecodeError, ValueError): + return None + + +def _is_structural(parsed: Any) -> bool: + """True if ``parsed`` is at least shaped like our expected object. + + We require a dict with a "rows" list — that's the structural floor + worth retrying for. Individual fields are allowed to be missing/null. + """ + return isinstance(parsed, dict) and isinstance(parsed.get("rows"), list) + + +# --------------------------------------------------------------------------- +# Response parsing +# --------------------------------------------------------------------------- + + +def _parse_response(raw: dict, page_number: int) -> PageExtraction: + """Turn one page's raw model JSON into a ``PageExtraction``.""" + result = PageExtraction(page_number=page_number) + + meet_obj = raw.get("meet") or {} + for key in s.MEET_FIELDS: + fv = _parse_field(meet_obj.get(key)) + if fv is not None: + result.meet[key] = fv + + session_obj = raw.get("session") or {} + for key in (s.COMPETITION_COORDINATOR, s.CC_LEVEL, s.DATE_SESSION): + fv = _parse_field(session_obj.get(key)) + if fv is not None: + result.session[key] = fv + sn = _parse_session_number(session_obj.get(s.SESSION_NUMBER)) + if sn is not None: + result.session[s.SESSION_NUMBER] = sn + + for row_obj in raw.get("rows") or []: + if not isinstance(row_obj, dict): + continue + row: dict[str, s.FieldValue] = {} + for key in (s.OFFICIAL_NAME, s.CLUB, s.POSITION, s.LANE_NUMBER, + s.TIMES_WORKED_POSITION, s.MENTOR, s.LEVEL): + fv = _parse_field(row_obj.get(key)) + if fv is not None: + row[key] = fv + suc = _parse_successful(row_obj.get(s.SUCCESSFUL)) + if suc is not None: + row[s.SUCCESSFUL] = suc + if row: + result.rows.append(row) + + return result + + +def _parse_field(obj: Any) -> Optional[s.FieldValue]: + """Parse a ``{"value":..,"confidence":..}`` object into a FieldValue. + + Tolerates the model emitting a bare scalar instead of the object + form. Returns None when there's no usable value. + """ + if obj is None: + return None + if isinstance(obj, dict): + value = obj.get("value") + if value is None: + return None + return s.FieldValue( + value=value, + confidence=_clamp_confidence(obj.get("confidence")), + ) + # Bare scalar — model ignored the object shape for this field. + if isinstance(obj, (str, int, float, bool)): + if isinstance(obj, str) and not obj.strip(): + return None + return s.FieldValue(value=obj, confidence=_DEFAULT_CONFIDENCE) + return None + + +def _parse_session_number(obj: Any) -> Optional[s.FieldValue]: + """Parse the session_number field, preserving its ``source``.""" + if obj is None: + return None + if isinstance(obj, dict): + value = obj.get("value") + if value is None: + return None + source = obj.get("source") + if source not in {"form", "filename", "form+filename", "unknown"}: + source = "unknown" + return s.FieldValue( + value=value, + confidence=_clamp_confidence(obj.get("confidence")), + source=source, + ) + if isinstance(obj, (int, str)): + return s.FieldValue( + value=obj, confidence=_DEFAULT_CONFIDENCE, source="unknown", + ) + return None + + +def _parse_successful(obj: Any) -> Optional[s.FieldValue]: + """Parse ``successful`` into a bool/null FieldValue with rationale. + + The model is asked for true/false/null. We also defend against + stringy spellings ("yes"/"no") so a slightly-off model response + doesn't silently become a wrong boolean downstream. + """ + if obj is None: + return None + if isinstance(obj, dict): + value = _coerce_bool(obj.get("value")) + return s.FieldValue( + value=value, + confidence=_clamp_confidence(obj.get("confidence")), + rationale=_as_str_or_none(obj.get("rationale")), + ) + # Bare scalar fallback. + return s.FieldValue( + value=_coerce_bool(obj), + confidence=_DEFAULT_CONFIDENCE, + ) + + +# --------------------------------------------------------------------------- +# Coercion helpers +# --------------------------------------------------------------------------- + + +def _coerce_bool(value: Any) -> Optional[bool]: + """Map the model's successful value to True / False / None.""" + if value is None: + return None + if isinstance(value, bool): + return value + if isinstance(value, (int, float)): + return bool(value) + if isinstance(value, str): + v = value.strip().lower() + if not v or v in {"null", "none", "n/a", "unknown", "?"}: + return None + if v in _TRUE_WORDS: + return True + if v in _FALSE_WORDS: + return False + # Unrecognised — treat as ambiguous rather than guessing. + return None + + +def _clamp_confidence(value: Any) -> float: + """Coerce a model-reported confidence into [0.0, 1.0].""" + try: + c = float(value) + except (TypeError, ValueError): + return _DEFAULT_CONFIDENCE + return max(0.0, min(1.0, c)) + + +def _as_str_or_none(value: Any) -> Optional[str]: + if value is None: + return None + text = str(value).strip() + return text or None + + +# --------------------------------------------------------------------------- +# Raw-response cache (.raw.json sidecar) +# --------------------------------------------------------------------------- + + +def _load_cache(cache_path: str | Path, model: str) -> Optional[dict]: + """Load the per-page raw responses if the cache matches ``model``.""" + path = Path(cache_path) + if not path.is_file(): + return None + try: + data = json.loads(path.read_text(encoding="utf-8")) + except (json.JSONDecodeError, OSError) as exc: + log.warning("Ignoring unreadable vision cache %s: %s", path, exc) + return None + if data.get("model") != model: + log.info( + "Vision cache %s was built with model %r, not %r — ignoring.", + path, data.get("model"), model, + ) + return None + pages = data.get("pages") + return pages if isinstance(pages, dict) else None + + +def _write_cache( + cache_path: str | Path, + model: str, + source_pdf: str, + raw_by_page: dict[str, Any], +) -> None: + path = Path(cache_path) + path.parent.mkdir(parents=True, exist_ok=True) + payload = { + "model": model, + "source_pdf": source_pdf, + "pages": raw_by_page, + } + with path.open("w", encoding="utf-8") as f: + json.dump(payload, f, indent=2, ensure_ascii=False) + f.write("\n") diff --git a/tests/test_vision_extract.py b/tests/test_vision_extract.py new file mode 100644 index 0000000..6cf4a11 --- /dev/null +++ b/tests/test_vision_extract.py @@ -0,0 +1,345 @@ +"""Tests for src/vision_extract.py. + +The vision model is always mocked — a fake client whose ``generate`` +returns canned JSON. Tests pin: prompt construction (template addendum + +filename + schema field names present), JSON parsing into the +PageExtraction shape, the retry-on-bad-JSON path, successful/session +coercion, confidence clamping, and the raw.json cache round-trip. +""" +from __future__ import annotations + +import json +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +from src import schema as s, vision_extract +from src.templates import TEMPLATES +from src.vision_extract import ( + VisionExtractionError, + build_prompt, + extract_page, + extract_pdf, +) + + +ONTARIO = TEMPLATES["swim_ontario_v1"] + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _good_response( + *, + competition_name="Aurora Open 2026", + n_rows=2, + session_source="form", +) -> dict: + rows = [] + for i in range(n_rows): + rows.append({ + "official_name": {"value": f"Official {i+1}", "confidence": 0.9}, + "club": {"value": "AAC", "confidence": 0.95}, + "position": {"value": "Starter", "confidence": 0.88}, + "lane_number": {"value": None, "confidence": 1.0}, + "times_worked_position": {"value": "3", "confidence": 0.7}, + "mentor": {"value": "Beth Brown", "confidence": 0.8}, + "level": {"value": "Level 4", "confidence": 0.75}, + "successful": {"value": True, "confidence": 0.92, + "rationale": "initials present"}, + }) + return { + "meet": { + "competition_name": {"value": competition_name, "confidence": 0.97}, + "host_club": {"value": "AAC", "confidence": 0.9}, + "coc": {"value": "Carol Smith", "confidence": 0.85}, + }, + "session": { + "competition_coordinator": {"value": "Alex A", "confidence": 0.9}, + "cc_level": {"value": "Level 5", "confidence": 0.8}, + "date_session": {"value": "2026-01-09", "confidence": 0.93}, + "session_number": {"value": 1, "confidence": 0.95, + "source": session_source}, + }, + "rows": rows, + } + + +def _client_returning(*json_objects_or_text) -> MagicMock: + """Build a fake client whose generate() returns the given payloads in + order. Each item may be a dict (JSON-dumped) or a raw string.""" + client = MagicMock() + responses = [] + for item in json_objects_or_text: + text = item if isinstance(item, str) else json.dumps(item) + responses.append(SimpleNamespace(response=text)) + client.generate.side_effect = responses + return client + + +# --------------------------------------------------------------------------- +# Prompt construction +# --------------------------------------------------------------------------- + + +class TestBuildPrompt: + def test_includes_filename(self): + prompt = build_prompt(ONTARIO, "meet-session_3.pdf") + assert "meet-session_3.pdf" in prompt + + def test_includes_template_addendum(self): + prompt = build_prompt(ONTARIO, "x.pdf") + # Swim Ontario's addendum mentions ditto marks and the + # whole-row successful judgement. + assert "Swim Ontario" in prompt + assert "ditto" in prompt.lower() + + def test_includes_canonical_field_names(self): + prompt = build_prompt(ONTARIO, "x.pdf") + for field in (s.OFFICIAL_NAME, s.SUCCESSFUL, s.SESSION_NUMBER, + s.COMPETITION_NAME, s.MENTOR): + assert field in prompt + + def test_explains_session_number_source(self): + prompt = build_prompt(ONTARIO, "x.pdf") + assert "form+filename" in prompt + assert "source" in prompt + + +# --------------------------------------------------------------------------- +# Parsing a good response +# --------------------------------------------------------------------------- + + +class TestParseGoodResponse: + def setup_method(self): + self.page = vision_extract._parse_response(_good_response(), page_number=1) + + def test_page_number(self): + assert self.page.page_number == 1 + + def test_meet_fields(self): + assert self.page.meet[s.COMPETITION_NAME].value == "Aurora Open 2026" + assert self.page.meet[s.COMPETITION_NAME].confidence == 0.97 + assert self.page.meet[s.HOST_CLUB].value == "AAC" + + def test_session_number_with_source(self): + sn = self.page.session[s.SESSION_NUMBER] + assert sn.value == 1 + assert sn.source == "form" + assert sn.confidence == 0.95 + + def test_rows_parsed(self): + assert len(self.page.rows) == 2 + assert self.page.rows[0][s.OFFICIAL_NAME].value == "Official 1" + assert self.page.rows[0][s.POSITION].value == "Starter" + + def test_successful_is_typed_bool(self): + suc = self.page.rows[0][s.SUCCESSFUL] + assert suc.value is True + assert suc.rationale == "initials present" + + def test_null_value_field_dropped(self): + # lane_number had value=None — it should not appear in the row. + assert s.LANE_NUMBER not in self.page.rows[0] + + +# --------------------------------------------------------------------------- +# Coercion: successful +# --------------------------------------------------------------------------- + + +class TestSuccessfulCoercion: + @pytest.mark.parametrize("raw, expected", [ + (True, True), + (False, False), + (None, None), + ("true", True), + ("false", False), + ("yes", True), + ("no", False), + ("", None), + ("unknown", None), + ("?", None), + ("Pass", True), + ("FAILED", False), + ]) + def test_coerce_bool(self, raw, expected): + assert vision_extract._coerce_bool(raw) == expected + + def test_stringy_false_does_not_become_true(self): + # The important regression: a model emitting "false" must not be + # read as a truthy non-empty string downstream. + page = vision_extract._parse_response({ + "rows": [{ + "official_name": {"value": "X", "confidence": 0.9}, + "successful": {"value": "false", "confidence": 0.8}, + }], + }, page_number=1) + assert page.rows[0][s.SUCCESSFUL].value is False + + +# --------------------------------------------------------------------------- +# Confidence clamping +# --------------------------------------------------------------------------- + + +class TestConfidenceClamping: + @pytest.mark.parametrize("raw, expected", [ + (0.9, 0.9), + (1.5, 1.0), # over-range clamps down + (-0.3, 0.0), # under-range clamps up + ("0.7", 0.7), # stringy float + (None, 0.5), # missing → default + ("abc", 0.5), # garbage → default + ]) + def test_clamp(self, raw, expected): + assert vision_extract._clamp_confidence(raw) == expected + + def test_bare_scalar_field_gets_default_confidence(self): + # Model emitted a bare string instead of {value, confidence}. + page = vision_extract._parse_response({ + "meet": {"competition_name": "Stray Meet"}, + "rows": [], + }, page_number=1) + fv = page.meet[s.COMPETITION_NAME] + assert fv.value == "Stray Meet" + assert fv.confidence == 0.5 + + +# --------------------------------------------------------------------------- +# Model call + retry +# --------------------------------------------------------------------------- + + +class TestModelCallAndRetry: + def test_single_call_when_valid(self): + client = _client_returning(_good_response(n_rows=1)) + page, raw = extract_page( + b"PNGBYTES", ONTARIO, "x.pdf", 1, client=client, model="m", + ) + assert client.generate.call_count == 1 + assert len(page.rows) == 1 + + def test_passes_image_and_json_format(self): + client = _client_returning(_good_response(n_rows=1)) + extract_page(b"PNGBYTES", ONTARIO, "x.pdf", 1, client=client, model="qwen2.5vl:7b") + _, kwargs = client.generate.call_args + assert kwargs["model"] == "qwen2.5vl:7b" + assert kwargs["images"] == [b"PNGBYTES"] + assert kwargs["format"] == "json" + assert kwargs["options"]["temperature"] == 0 + + def test_retries_once_on_bad_json_then_succeeds(self): + client = _client_returning("not json at all", _good_response(n_rows=1)) + page, _ = extract_page(b"x", ONTARIO, "x.pdf", 1, client=client, model="m") + assert client.generate.call_count == 2 + assert len(page.rows) == 1 + + def test_retry_nudge_appended_on_second_call(self): + client = _client_returning("garbage", _good_response(n_rows=1)) + extract_page(b"x", ONTARIO, "x.pdf", 1, client=client, model="m") + first_prompt = client.generate.call_args_list[0].kwargs["prompt"] + second_prompt = client.generate.call_args_list[1].kwargs["prompt"] + assert len(second_prompt) > len(first_prompt) + assert "not valid JSON" in second_prompt + + def test_raises_after_two_failures(self): + client = _client_returning("garbage", "still garbage") + with pytest.raises(VisionExtractionError): + extract_page(b"x", ONTARIO, "x.pdf", 1, client=client, model="m") + + def test_strips_markdown_fences(self): + fenced = "```json\n" + json.dumps(_good_response(n_rows=1)) + "\n```" + client = _client_returning(fenced) + page, _ = extract_page(b"x", ONTARIO, "x.pdf", 1, client=client, model="m") + # Parsed fine despite the fence → only one call, rows present. + assert client.generate.call_count == 1 + assert len(page.rows) == 1 + + def test_missing_rows_key_triggers_retry(self): + # A dict without "rows" is structurally invalid → retry. + client = _client_returning({"meet": {}}, _good_response(n_rows=1)) + extract_page(b"x", ONTARIO, "x.pdf", 1, client=client, model="m") + assert client.generate.call_count == 2 + + +# --------------------------------------------------------------------------- +# extract_pdf orchestration + cache +# --------------------------------------------------------------------------- + + +class TestExtractPdfOrchestration: + def test_calls_model_once_per_page(self, tmp_path): + client = _client_returning(_good_response(n_rows=2), _good_response(n_rows=1)) + with patch("src.vision_extract.pdf_io.page_count", return_value=2), \ + patch("src.vision_extract.pdf_io.rasterize_page", return_value=b"PNG"): + pages = extract_pdf( + "scan.pdf", ONTARIO, client=client, model="m", + cache_path=None, + ) + assert client.generate.call_count == 2 + assert len(pages) == 2 + assert len(pages[0].rows) == 2 + assert len(pages[1].rows) == 1 + + +class TestCache: + def test_writes_cache_then_reuses_without_calling_model(self, tmp_path): + cache = tmp_path / "scan.raw.json" + + # First run: model called once per page, cache written. + client1 = _client_returning(_good_response(n_rows=1)) + with patch("src.vision_extract.pdf_io.page_count", return_value=1), \ + patch("src.vision_extract.pdf_io.rasterize_page", return_value=b"PNG"): + extract_pdf("scan.pdf", ONTARIO, client=client1, model="m", + cache_path=cache) + assert client1.generate.call_count == 1 + assert cache.is_file() + + # Second run: cache hit → model NOT called. + client2 = _client_returning(_good_response(n_rows=1)) + with patch("src.vision_extract.pdf_io.page_count", return_value=1), \ + patch("src.vision_extract.pdf_io.rasterize_page", return_value=b"PNG"): + pages = extract_pdf("scan.pdf", ONTARIO, client=client2, model="m", + cache_path=cache) + client2.generate.assert_not_called() + assert len(pages[0].rows) == 1 + + def test_no_cache_flag_forces_fresh_call(self, tmp_path): + cache = tmp_path / "scan.raw.json" + # Pre-seed a cache. + cache.write_text(json.dumps({ + "model": "m", "source_pdf": "scan.pdf", + "pages": {"1": _good_response(n_rows=5)}, + }), encoding="utf-8") + + client = _client_returning(_good_response(n_rows=1)) + with patch("src.vision_extract.pdf_io.page_count", return_value=1), \ + patch("src.vision_extract.pdf_io.rasterize_page", return_value=b"PNG"): + pages = extract_pdf("scan.pdf", ONTARIO, client=client, model="m", + cache_path=cache, use_cache=False) + # use_cache=False → model called, and we get 1 row not the + # cached 5. + assert client.generate.call_count == 1 + assert len(pages[0].rows) == 1 + + def test_cache_ignored_when_model_differs(self, tmp_path): + cache = tmp_path / "scan.raw.json" + cache.write_text(json.dumps({ + "model": "OLD-MODEL", "source_pdf": "scan.pdf", + "pages": {"1": _good_response(n_rows=5)}, + }), encoding="utf-8") + + client = _client_returning(_good_response(n_rows=1)) + with patch("src.vision_extract.pdf_io.page_count", return_value=1), \ + patch("src.vision_extract.pdf_io.rasterize_page", return_value=b"PNG"): + pages = extract_pdf("scan.pdf", ONTARIO, client=client, model="NEW-MODEL", + cache_path=cache) + # Different model → cache ignored → fresh call. + assert client.generate.call_count == 1 + assert len(pages[0].rows) == 1