feat(ocr): parser service (routing + per-page merge)

extralit-server/src/extralit_server/contexts/ocr/parsers/service.py

@@ -0,0 +1,145 @@
+"""Route a PDF's pages to parsers and merge what they return into one `DoclingDocument`.
+
+Routing rule: pages triage flags as needing OCR go to liteparse with OCR on; every other page
+goes to pymupdf when the extra is installed, else to liteparse with OCR off. pdf_inspector is
+never chosen automatically. A task may carry a `bbox`, in which case it is a region task and
+only blocks mostly inside that region are kept.
+"""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from collections.abc import Callable, Mapping, Sequence
+from dataclasses import dataclass
+from typing import Optional, Protocol
+
+from docling_core.types.doc import BoundingBox, DoclingDocument, Size
+
+from extralit_server.api.schemas.v1.document.metadata import TriageMetadata
+from extralit_server.contexts.ocr.docling_builder import (
+    CONTAINMENT_THRESHOLD,
+    LayoutBlock,
+    PageContext,
+    append_blocks,
+    content_hash,
+    new_document,
+    register_page,
+)
+from extralit_server.contexts.ocr.parsers.pdf_inspector import page_sizes as _page_sizes
+from extralit_server.contexts.ocr.triage import triage_pdf
+
+OCR_PARSER = "liteparse"
+DIGITAL_PARSER = "pymupdf"
+
+
+class PageParser(Protocol):
+    """Parse a 1-indexed page subset into top-left-origin blocks, keyed by page number."""
+
+    def parse_pages(self, pdf_bytes: bytes, pages: Sequence[int], *, ocr: bool) -> dict[int, list[LayoutBlock]]: ...
+
+
+@dataclass(frozen=True)
+class ParseTask:
+    pages: tuple[int, ...]
+    parser: str
+    ocr: bool = False
+    bbox: Optional[BoundingBox] = None
+
+
+@dataclass(frozen=True)
+class ParsePlan:
+    tasks: tuple[ParseTask, ...]
+
+    @property
+    def pages_by_parser(self) -> dict[str, list[int]]:
+        pages: dict[str, set[int]] = defaultdict(set)
+        for task in self.tasks:
+            pages[task.parser].update(task.pages)
+        return {name: sorted(p) for name, p in pages.items()}
+
+    @property
+    def ocr_pages(self) -> list[int]:
+        return sorted({p for task in self.tasks if task.ocr for p in task.pages})
+
+
+def build_document(
+    blocks: Mapping[int, Sequence[LayoutBlock]],
+    *,
+    page_sizes: Mapping[int, Size],
+    name: str,
+    filename: Optional[str] = None,
+    binary_hash: Optional[int] = None,
+) -> DoclingDocument:
+    """Lay per-page blocks into a fresh `DoclingDocument`; every known page is registered."""
+    doc = new_document(name, filename=filename, binary_hash=binary_hash)
+    for page_no in sorted(page_sizes):
+        ctx = PageContext(page_no=page_no, size=page_sizes[page_no])
+        page_blocks = sorted(blocks.get(page_no, ()), key=lambda b: (b.bbox.t, b.bbox.l))
+        if page_blocks:
+            append_blocks(doc, ctx, page_blocks)
+        else:
+            register_page(doc, ctx)
+    return doc
+
+
+class ParserService:
+    def __init__(
+        self,
+        parsers: Mapping[str, PageParser],
+        *,
+        triage: Callable[[bytes], TriageMetadata] = triage_pdf,
+        page_sizes: Callable[[bytes], dict[int, Size]] = _page_sizes,
+    ):
+        self._parsers = dict(parsers)
+        self._triage = triage
+        self._page_sizes = page_sizes
+
+    def _parser(self, name: str) -> PageParser:
+        try:
+            return self._parsers[name]
+        except KeyError:
+            raise ValueError(f"unknown page parser {name!r}; available: {sorted(self._parsers)}") from None
+
+    def plan(self, pdf_bytes: bytes, *, pages: Optional[Sequence[int]] = None) -> ParsePlan:
+        triage = self._triage(pdf_bytes)
+        page_count = triage.page_count or len(self._page_sizes(pdf_bytes))
+        wanted = [p for p in range(1, page_count + 1) if pages is None or p in set(pages)]
+
+        scanned = [p for p in wanted if p in set(triage.pages_needing_ocr)]
+        digital = [p for p in wanted if p not in set(scanned)]
+        digital_parser = DIGITAL_PARSER if DIGITAL_PARSER in self._parsers else OCR_PARSER
+
+        tasks: list[ParseTask] = []
+        if digital:
+            tasks.append(ParseTask(pages=tuple(digital), parser=digital_parser, ocr=False))
+        if scanned:
+            tasks.append(ParseTask(pages=tuple(scanned), parser=OCR_PARSER, ocr=True))
+        return ParsePlan(tasks=tuple(tasks))
+
+    def parse(
+        self,
+        pdf_bytes: bytes,
+        plan: ParsePlan,
+        *,
+        name: str,
+        filename: Optional[str] = None,
+    ) -> DoclingDocument:
+        """Run every task and merge their blocks per page, ordered by `(top, left)`."""
+        parsers = {task.parser: self._parser(task.parser) for task in plan.tasks}
+        merged: dict[int, list[LayoutBlock]] = defaultdict(list)
+        for task in plan.tasks:
+            for page_no, blocks in parsers[task.parser].parse_pages(pdf_bytes, task.pages, ocr=task.ocr).items():
+                if task.bbox is not None:
+                    blocks = [b for b in blocks if b.bbox.intersection_over_self(task.bbox) >= CONTAINMENT_THRESHOLD]
+                merged[page_no].extend(blocks)
+
+        return build_document(
+            merged,
+            page_sizes=self._page_sizes(pdf_bytes),
+            name=name,
+            filename=filename,
+            binary_hash=content_hash(pdf_bytes),
+        )
+
+
+__all__ = ["PageParser", "ParsePlan", "ParseTask", "ParserService", "build_document"]

extralit-server/tests/unit/contexts/ocr/test_parser_service.py

@@ -0,0 +1,181 @@
+"""Tests for the parser service: routing (plan) and per-page merging (parse)."""
+
+from collections.abc import Sequence
+from typing import Optional
+
+import pytest
+from docling_core.types.doc import BoundingBox, CoordOrigin, DocItemLabel, DoclingDocument, Size
+
+from extralit_server.api.schemas.v1.document.metadata import TriageMetadata
+from extralit_server.contexts.ocr.docling_builder import LayoutBlock
+from extralit_server.contexts.ocr.parsers.service import (
+    ParsePlan,
+    ParserService,
+    ParseTask,
+    build_document,
+)
+
+PAGE = Size(width=100.0, height=200.0)
+
+
+def _box(top: float, left: float, size: float = 10.0) -> BoundingBox:
+    return BoundingBox(l=left, t=top, r=left + size, b=top + size, coord_origin=CoordOrigin.TOPLEFT)
+
+
+def _text(top: float, left: float, text: str) -> LayoutBlock:
+    return LayoutBlock(label=DocItemLabel.TEXT, bbox=_box(top, left), text=text)
+
+
+class FakeParser:
+    """Records what it was asked for and answers with canned blocks per page."""
+
+    def __init__(self, blocks: Optional[dict[int, list[LayoutBlock]]] = None):
+        self.blocks = blocks or {}
+        self.calls: list[tuple[tuple[int, ...], bool]] = []
+
+    def parse_pages(self, pdf_bytes: bytes, pages: Sequence[int], *, ocr: bool) -> dict[int, list[LayoutBlock]]:
+        self.calls.append((tuple(pages), ocr))
+        return {p: list(self.blocks.get(p, [])) for p in pages}
+
+
+def _service(parsers: dict, triage: TriageMetadata, page_count: int = 3) -> ParserService:
+    return ParserService(
+        parsers,
+        triage=lambda _: triage,
+        page_sizes=lambda _: dict.fromkeys(range(1, page_count + 1), PAGE),
+    )
+
+
+class TestPlan:
+    def test_scanned_pages_go_to_liteparse_with_ocr(self):
+        triage = TriageMetadata(pdf_type="mixed", page_count=3, pages_needing_ocr=[2])
+        service = _service({"pymupdf": FakeParser(), "liteparse": FakeParser()}, triage)
+
+        plan = service.plan(b"%PDF")
+
+        assert plan.tasks == (
+            ParseTask(pages=(1, 3), parser="pymupdf", ocr=False),
+            ParseTask(pages=(2,), parser="liteparse", ocr=True),
+        )
+
+    def test_digital_pages_go_to_pymupdf_when_installed(self):
+        triage = TriageMetadata(pdf_type="text_based", page_count=2)
+        service = _service({"pymupdf": FakeParser(), "liteparse": FakeParser()}, triage, page_count=2)
+
+        plan = service.plan(b"%PDF")
+
+        assert plan.tasks == (ParseTask(pages=(1, 2), parser="pymupdf", ocr=False),)
+
+    def test_digital_pages_fall_back_to_liteparse_without_ocr(self):
+        triage = TriageMetadata(pdf_type="text_based", page_count=2)
+        service = _service({"liteparse": FakeParser()}, triage, page_count=2)
+
+        plan = service.plan(b"%PDF")
+
+        assert plan.tasks == (ParseTask(pages=(1, 2), parser="liteparse", ocr=False),)
+
+    def test_fully_scanned_pdf_is_one_ocr_task(self):
+        triage = TriageMetadata(pdf_type="image_based", page_count=2, pages_needing_ocr=[1, 2])
+        service = _service({"pymupdf": FakeParser(), "liteparse": FakeParser()}, triage, page_count=2)
+
+        plan = service.plan(b"%PDF")
+
+        assert plan.tasks == (ParseTask(pages=(1, 2), parser="liteparse", ocr=True),)
+
+    def test_page_count_comes_from_the_pdf_when_triage_is_unknown(self):
+        triage = TriageMetadata(pdf_type="unknown")
+        service = _service({"liteparse": FakeParser()}, triage, page_count=4)
+
+        plan = service.plan(b"%PDF")
+
+        assert plan.tasks == (ParseTask(pages=(1, 2, 3, 4), parser="liteparse", ocr=False),)
+
+    def test_pages_allowlist_restricts_the_plan(self):
+        triage = TriageMetadata(pdf_type="mixed", page_count=3, pages_needing_ocr=[2])
+        service = _service({"pymupdf": FakeParser(), "liteparse": FakeParser()}, triage)
+
+        plan = service.plan(b"%PDF", pages=[2, 3])
+
+        assert plan.tasks == (
+            ParseTask(pages=(3,), parser="pymupdf", ocr=False),
+            ParseTask(pages=(2,), parser="liteparse", ocr=True),
+        )
+
+    def test_plan_summaries(self):
+        plan = ParsePlan(
+            tasks=(
+                ParseTask(pages=(1, 3), parser="pymupdf", ocr=False),
+                ParseTask(pages=(2,), parser="liteparse", ocr=True),
+            )
+        )
+
+        assert plan.pages_by_parser == {"pymupdf": [1, 3], "liteparse": [2]}
+        assert plan.ocr_pages == [2]
+
+
+class TestParse:
+    def test_each_task_runs_on_its_parser_with_its_ocr_flag(self):
+        pymupdf, liteparse = FakeParser(), FakeParser()
+        service = _service({"pymupdf": pymupdf, "liteparse": liteparse}, TriageMetadata(pdf_type="mixed"))
+        plan = ParsePlan(
+            tasks=(
+                ParseTask(pages=(1, 3), parser="pymupdf", ocr=False),
+                ParseTask(pages=(2,), parser="liteparse", ocr=True),
+            )
+        )
+
+        service.parse(b"%PDF", plan, name="doc")
+
+        assert pymupdf.calls == [((1, 3), False)]
+        assert liteparse.calls == [((2,), True)]
+
+    def test_blocks_from_different_tasks_merge_per_page_in_reading_order(self):
+        first = FakeParser({1: [_text(50, 0, "middle"), _text(10, 0, "top")]})
+        second = FakeParser({1: [_text(50, 20, "middle-right"), _text(90, 0, "bottom")]})
+        service = _service({"a": first, "b": second}, TriageMetadata(pdf_type="mixed"))
+        plan = ParsePlan(
+            tasks=(
+                ParseTask(pages=(1,), parser="a", ocr=False),
+                ParseTask(pages=(1,), parser="b", ocr=False, bbox=_box(0, 0, size=200)),
+            )
+        )
+
+        doc = service.parse(b"%PDF", plan, name="doc")
+
+        assert [t.text for t in doc.texts] == ["top", "middle", "middle-right", "bottom"]
+
+    def test_region_task_keeps_only_blocks_inside_its_bbox(self):
+        parser = FakeParser({1: [_text(10, 10, "inside"), _text(150, 10, "outside"), _text(45, 10, "straddling")]})
+        service = _service({"a": parser}, TriageMetadata(pdf_type="mixed"))
+        region = BoundingBox(l=0, t=0, r=100, b=50, coord_origin=CoordOrigin.TOPLEFT)
+        plan = ParsePlan(tasks=(ParseTask(pages=(1,), parser="a", ocr=False, bbox=region),))
+
+        doc = service.parse(b"%PDF", plan, name="doc")
+
+        assert [t.text for t in doc.texts] == ["inside"]
+
+    def test_unknown_parser_in_plan_raises(self):
+        service = _service({"a": FakeParser()}, TriageMetadata(pdf_type="mixed"))
+        plan = ParsePlan(tasks=(ParseTask(pages=(1,), parser="nope", ocr=False),))
+
+        with pytest.raises(ValueError, match="nope"):
+            service.parse(b"%PDF", plan, name="doc")
+
+
+class TestBuildDocument:
+    def test_pages_are_registered_with_their_size_and_blocks_carry_provenance(self):
+        blocks = {1: [_text(10, 0, "one")], 2: [_text(20, 0, "two")]}
+
+        doc = build_document(blocks, page_sizes={1: PAGE, 2: PAGE}, name="doc", filename="doc.pdf")
+
+        assert isinstance(doc, DoclingDocument)
+        assert doc.name == "doc"
+        assert doc.origin.filename == "doc.pdf"
+        assert doc.pages[2].size == PAGE
+        assert [(t.text, t.prov[0].page_no) for t in doc.texts] == [("one", 1), ("two", 2)]
+
+    def test_empty_pages_are_still_registered(self):
+        doc = build_document({}, page_sizes={1: PAGE}, name="doc")
+
+        assert set(doc.pages) == {1}
+        assert doc.texts == []