diff --git a/extralit-server/src/extralit_server/contexts/ocr/parsers/service.py b/extralit-server/src/extralit_server/contexts/ocr/parsers/service.py new file mode 100644 index 000000000..6ea1014f6 --- /dev/null +++ b/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"] diff --git a/extralit-server/tests/unit/contexts/ocr/test_parser_service.py b/extralit-server/tests/unit/contexts/ocr/test_parser_service.py new file mode 100644 index 000000000..c8b265898 --- /dev/null +++ b/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 == []