Address review: tighten return type, cache label, document key asymmetry

- _create_pdf_annotation return type Annotation (not Annotation | None) and remove dead 'if annot is not None' guard in caller — function raises on failure
- Cache ensure_label_and_labelset across loop iterations; reset on savepoint rollback so the happy path is one DB lookup, not N
- Document why json is in defaults for PDFs but a lookup key for spans
- Simplify test stub so it doesn't depend on SpanAnnotation internals
- Fix test page assertion: PlasmaPDF returns 0-indexed pages

Author: JSv4

opencontractserver/tests/test_extraction_grounding.py

@@ -514,17 +514,18 @@ def test_ground_pdf_creates_token_label_annotations(self):
             self.assertIsNotNone(annot.annotation_label)
             assert annot.annotation_label is not None  # narrow for mypy
             self.assertEqual(annot.annotation_label.text, OC_EXTRACT_SOURCE_LABEL)
-            # Page must be a positive integer; never the silent default of 1
-            # for a span that actually lives on page 2.
+            # PlasmaPDF returns 0-indexed pages; valid range is
+            # [0, len(pages) - 1]. The bug we're guarding against is the
+            # silent fallback that would have saved everything on a single
+            # default page even when the span lives on a different one.
             self.assertIsInstance(annot.page, int)
-            self.assertGreaterEqual(annot.page, 1)
-            self.assertLessEqual(annot.page, len(self.pages_text))
+            self.assertGreaterEqual(annot.page, 0)
+            self.assertLess(annot.page, len(self.pages_text))
             self.assertTrue(annot.raw_text)
 
-        # "Acme Holdings" is on page 1 (1-indexed) and
-        # "Global Acquisitions" on page 2 — confirm the per-page mapping
-        # actually works by checking we got at least one annotation off
-        # page 1.
+        # "Acme Holdings" is on page 0 (0-indexed) and "Global Acquisitions"
+        # on page 1 — confirm the per-page mapping actually works by
+        # checking we got annotations on more than one page.
         pages_seen = {a.page for a in annotations}
         self.assertGreater(
             len(pages_seen),
@@ -587,11 +588,7 @@ def test_ground_pdf_skips_when_page_is_none(self):
         def stub_create(self, span_annotation):
             # Mimic PlasmaPDF's payload but force page=None so the grounding
             # pipeline must take the skip-rather-than-fallback path.
-            return {
-                "page": None,
-                "rawText": span_annotation.span["text"],
-                "annotation_json": {},
-            }
+            return {"page": None, "rawText": "stub", "annotation_json": {}}
 
         with patch(
             "plasmapdf.models.PdfDataLayer.PdfDataLayer."

opencontractserver/utils/extraction_grounding.py

@@ -207,16 +207,21 @@ def _create_grounding_annotations(
 
     annotations: list[Annotation] = []
 
+    # Cache the label across iterations so the happy path is one DB lookup,
+    # not N. Reset to None on savepoint rollback so a transient labelset
+    # failure can be retried on the next annotation.
+    cached_label: AnnotationLabel | None = None
+
     for result in alignment_results:
         try:
             with transaction.atomic():
-                # Resolve label inside the savepoint so a labelset failure
-                # only skips this annotation, not the whole batch.
-                label_obj = corpus.ensure_label_and_labelset(
-                    label_text=OC_EXTRACT_SOURCE_LABEL,
-                    creator_id=creator_id,
-                    label_type=annotation_type,
-                )
+                if cached_label is None:
+                    cached_label = corpus.ensure_label_and_labelset(
+                        label_text=OC_EXTRACT_SOURCE_LABEL,
+                        creator_id=creator_id,
+                        label_type=annotation_type,
+                    )
+                label_obj = cached_label
 
                 if annotation_type == TOKEN_LABEL and pdf_layer is not None:
                     annot = _create_pdf_annotation(
@@ -229,10 +234,10 @@ def _create_grounding_annotations(
                 else:
                     continue
 
-                if annot is not None:
-                    annotations.append(annot)
+                annotations.append(annot)
 
         except Exception:
+            cached_label = None
             logger.warning(
                 "Failed to create grounding annotation for %r "
                 "at [%d:%d] in document %d",
@@ -253,15 +258,16 @@ def _create_pdf_annotation(
     creator_id: int,
     pdf_layer: Any,
     label_obj: AnnotationLabel,
-) -> Annotation | None:
+) -> Annotation:
     """Create (or fetch) a TOKEN_LABEL annotation for a PDF document via PlasmaPDF.
 
     Idempotent: returns the existing annotation if one with the same
-    document, label, and span coordinates already exists.
+    document, label, page, and raw text already exists.
 
-    Returns ``None`` if the annotation cannot be created safely (e.g.
-    PlasmaPDF could not determine the page) so the caller can skip it.
-    Raises inside the savepoint to roll back on hard failures.
+    Raises ``ValueError`` inside the savepoint when the annotation cannot
+    be created safely (e.g. PlasmaPDF could not determine the page); the
+    caller's per-annotation ``try/except`` rolls back the savepoint and
+    logs the skip.
     """
     from plasmapdf.models.types import SpanAnnotation, TextSpan
 
@@ -289,6 +295,13 @@ def _create_pdf_annotation(
             f"{document.id}; skipping grounding annotation."
         )
 
+    # Note: ``json`` (bounding boxes) is in ``defaults``, NOT a lookup key.
+    # For PDFs the (document, label, page, raw_text) tuple already uniquely
+    # identifies the span; PlasmaPDF's bounding-box layout is deterministic
+    # for stable input, so on Celery retry we want to reuse the existing
+    # annotation rather than create a near-duplicate that differs only by
+    # bounding-box reformatting. Span annotations key on ``json`` because
+    # the char offsets ARE the identity for a text/DOCX document.
     annot, _ = Annotation.objects.get_or_create(
         document=document,
         annotation_label=label_obj,