diff --git a/extralit-server/scripts/checkpoint1_dump.py b/extralit-server/scripts/checkpoint1_dump.py new file mode 100644 index 000000000..fd604fdf8 --- /dev/null +++ b/extralit-server/scripts/checkpoint1_dump.py @@ -0,0 +1,166 @@ +"""Checkpoint 1: run `ParserService` on the gold PDFs and check table header rows against `gold.md`. + +Usage: uv run python scripts/checkpoint1_dump.py [--gold-dir DIR] [--out DIR] +Exit code 1 when a gold table header row is missing from the markdown (adageba Table 2 excepted) +or the scanned paper does not route through liteparse with OCR. +""" + +from __future__ import annotations + +import argparse +import itertools +import json +import re +import sys +from dataclasses import dataclass +from pathlib import Path + +from extralit_server.contexts.ocr.parsers.registry import default_service +from extralit_server.contexts.ocr.parsers.service import OCR_PARSER + +PAPERS = ("robert1991-bullwho-scanned", "adageba2022-plosone-digital") +EXCEPTIONS = {("adageba2022-plosone-digital", "Table 2")} +SCANNED = "robert1991-bullwho-scanned" + + +@dataclass +class HeaderCheck: + """Whether one gold table's header row survived into the parsed markdown.""" + + table: str + header: list[str] + status: str # ok | cells-elsewhere | missing | excepted + absent: tuple[str, ...] = () + + +HEADER_ROWS_SCANNED = 3 + + +def _norm(text: str) -> str: + """Letters and digits only, so spacing, hyphens and markdown padding do not matter.""" + return re.sub(r"[^a-z0-9]", "", text.lower()) + + +def markdown_tables(markdown: str) -> list[list[str]]: + """Each markdown table as its non-separator rows.""" + tables: list[list[str]] = [] + current: list[str] = [] + for line in [*markdown.splitlines(), ""]: + if line.lstrip().startswith("|"): + if not re.fullmatch(r"\|[\s:|-]+\|?", line.strip()): + current.append(line) + elif current: + tables.append(current) + current = [] + return tables + + +def gold_headers(gold_md: str) -> list[tuple[str, list[str]]]: + """Header row of each `### Table N` section: the pipe row followed by a `|---|` separator.""" + headers: list[tuple[str, list[str]]] = [] + for section in re.split(r"^### ", gold_md, flags=re.MULTILINE)[1:]: + title = re.match(r"(Table \d+)", section) + lines = section.splitlines() + for line, nxt in itertools.pairwise(lines): + if line.startswith("|") and re.fullmatch(r"\|[\s:|-]+\|?", nxt.strip()): + cells = [c.strip() for c in line.strip().strip("|").split("|")] + headers.append((title.group(1) if title else "?", [c for c in cells if c])) + break + return headers + + +def check_headers(paper: str, gold_md: str, markdown: str) -> list[HeaderCheck]: + """A gold header is `ok` when all its cells sit in the first rows of one markdown table.""" + heads = ["".join(_norm(row) for row in table[:HEADER_ROWS_SCANNED]) for table in markdown_tables(markdown)] + flat = _norm(markdown) + results = [] + for table, header in gold_headers(gold_md): + cells = [_norm(c) for c in header] + if any(all(c in head for c in cells) for head in heads): + status = "ok" + elif all(c in flat for c in cells): + status = "cells-elsewhere" + else: + status = "missing" + if (paper, table) in EXCEPTIONS and status != "ok": + status = "excepted" + absent = tuple(c for c, n in zip(header, cells, strict=True) if n not in flat) + results.append(HeaderCheck(table, header, status, absent)) + return results + + +def dump_paper(paper: str, gold_dir: Path, out_dir: Path) -> tuple[list[str], list[HeaderCheck], bool]: + """Parse one gold paper, write its artifacts, and return its report lines, header checks and routing verdict.""" + pdf_bytes = (gold_dir / paper / "source.pdf").read_bytes() + service = default_service() + plan = service.plan(pdf_bytes) + doc = service.parse(pdf_bytes, plan, name=paper, filename="source.pdf") + markdown = doc.export_to_markdown() + + target = out_dir / paper + target.mkdir(parents=True, exist_ok=True) + (target / "markdown.md").write_text(markdown) + lines = ["ref\tlabel\tpage\tl\tt\tr\tb"] + for item, _ in doc.iterate_items(with_groups=False): + for prov in item.prov or [None]: + box = prov.bbox if prov else None + coords = f"{box.l:.1f}\t{box.t:.1f}\t{box.r:.1f}\t{box.b:.1f}" if box else "\t\t\t" + lines.append(f"{item.self_ref}\t{item.label.value}\t{prov.page_no if prov else ''}\t{coords}") + (target / "items.tsv").write_text("\n".join(lines) + "\n") + plan_json = { + "tasks": [{"pages": list(t.pages), "parser": t.parser, "ocr": t.ocr} for t in plan.tasks], + "parsers": plan.pages_by_parser, + "pages_needing_ocr": list(plan.pages_needing_ocr), + "ocr_pages": plan.ocr_pages, + "ocr_skipped_pages": plan.ocr_skipped_pages, + } + (target / "plan.json").write_text(json.dumps(plan_json, indent=2)) + + checks = check_headers(paper, (gold_dir / paper / "gold.md").read_text(), markdown) + routed = True + if paper == SCANNED: + ocr_covered = {p for t in plan.tasks if t.parser == OCR_PARSER and t.ocr for p in t.pages} + routed = bool(plan.pages_needing_ocr) and set(plan.pages_needing_ocr) <= ocr_covered + + labels = {} + for item, _ in doc.iterate_items(with_groups=False): + labels[item.label.value] = labels.get(item.label.value, 0) + 1 + summary = [ + f"## {paper}", + f"- plan: `{json.dumps(plan_json['parsers'])}`; ocr_pages={plan.ocr_pages}; " + f"ocr_skipped_pages (text layer read, no OCR)={plan.ocr_skipped_pages}; " + f"pages_needing_ocr={list(plan.pages_needing_ocr)}", + f"- items by label: `{json.dumps(labels, sort_keys=True)}`", + ] + if paper == SCANNED: + summary.append(f"- every page needing OCR routed through liteparse+OCR: **{routed}**") + summary.append("- gold table header rows:") + summary += [ + f" - {c.table}: **{c.status}** — {' | '.join(c.header)}" + + (f" (absent cells: {list(c.absent)})" if c.absent else "") + for c in checks + ] + return summary, checks, routed + + +def main() -> int: + """Dump every checkpoint paper and return non-zero if any gate fails.""" + parser = argparse.ArgumentParser() + parser.add_argument("--gold-dir", type=Path, default=Path("~/Projects/Extralit/data/gold").expanduser()) + parser.add_argument("--out", type=Path, default=Path("/tmp/eng85-checkpoint1")) + args = parser.parse_args() + args.out.mkdir(parents=True, exist_ok=True) + + report: list[str] = ["# Checkpoint 1 report", ""] + failed = False + for paper in PAPERS: + summary, checks, routed = dump_paper(paper, args.gold_dir, args.out) + report += [*summary, ""] + failed |= not routed or any(c.status in ("missing", "cells-elsewhere") for c in checks) + (args.out / "report.md").write_text("\n".join(report)) + print("\n".join(report)) + return 1 if failed else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/extralit-server/src/extralit_server/api/schemas/v1/document/metadata.py b/extralit-server/src/extralit_server/api/schemas/v1/document/metadata.py index da6d88383..7462eb5b6 100644 --- a/extralit-server/src/extralit_server/api/schemas/v1/document/metadata.py +++ b/extralit-server/src/extralit_server/api/schemas/v1/document/metadata.py @@ -79,6 +79,10 @@ class LayoutMetadata(BaseModel): Only pointers and counts — the layout itself lives in object storage, because `documents.metadata_` is returned in full by every `GET /documents` listing. + + `parsers` is the executed plan (parser name -> pages). `ocr_pages` are pages where OCR actually + ran; `ocr_skipped_pages` are pages planned for OCR (triage flagged them) whose existing text + layer the parser read instead, so `pages_needing_ocr` = planned, `ocr_pages` = actual. """ layout_url: str = Field(..., description="S3 object path of the canonical DoclingDocument JSON") @@ -93,6 +97,14 @@ class LayoutMetadata(BaseModel): pages_needing_ocr: list[int] = Field( default_factory=list, description="1-indexed pages with no reliable text layer" ) + parsers: dict[str, list[int]] = Field( + default_factory=dict, description="Executed plan: parser name -> 1-indexed pages it parsed" + ) + ocr_pages: list[int] = Field(default_factory=list, description="1-indexed pages where OCR actually ran") + ocr_skipped_pages: list[int] = Field( + default_factory=list, + description="Pages planned for OCR that already had a text layer, which the parser read instead of running OCR", + ) class DocumentProcessingMetadata(BaseModel): diff --git a/extralit-server/src/extralit_server/contexts/ocr/docling_builder.py b/extralit-server/src/extralit_server/contexts/ocr/docling_builder.py index 3de588706..22d6355e3 100644 --- a/extralit-server/src/extralit_server/contexts/ocr/docling_builder.py +++ b/extralit-server/src/extralit_server/contexts/ocr/docling_builder.py @@ -6,6 +6,8 @@ from __future__ import annotations +import re +from collections.abc import Sequence from dataclasses import dataclass from typing import Optional @@ -27,6 +29,9 @@ #: Fraction of a text block that must fall inside a table/picture for it to count as duplicated. CONTAINMENT_THRESHOLD = 0.6 +#: Share of the cell text under a block that the block's own words must account for to count as a duplicate. +CELL_SHARE_THRESHOLD = 0.5 + #: Labels routed to `add_table` / `add_picture` rather than `add_text`. TABLE_LABELS = frozenset({DocItemLabel.TABLE, DocItemLabel.DOCUMENT_INDEX}) PICTURE_LABELS = frozenset({DocItemLabel.PICTURE, DocItemLabel.CHART}) @@ -61,10 +66,12 @@ class PageContext: @property def height(self) -> float: + """Page height in points.""" return self.size.height @property def width(self) -> float: + """Page width in points.""" return self.size.width @@ -121,41 +128,49 @@ def make_prov(ctx: PageContext, bbox: BoundingBox, text: Optional[str]) -> Prove ) -def is_contained( - doc: DoclingDocument, - bbox: BoundingBox, - page_no: int, - threshold: float = CONTAINMENT_THRESHOLD, -) -> bool: - """Whether `bbox` is mostly swallowed by a table or picture already on this page.""" - for item in [*doc.tables, *doc.pictures]: - for prov in item.prov: - if prov.page_no != page_no: - continue - if bbox.intersection_over_self(prov.bbox) >= threshold: - return True - return False +def _tokens(text: str) -> list[str]: + return re.findall(r"\w+", text.lower()) + + +def _carries_text(block_words: list[str], container: LayoutBlock, box: BoundingBox, ctx: PageContext) -> bool: + """Whether the cells under `box` hold this block's words and little else. + + A cell much longer than the block is a false table swallowing prose, so the block is kept. + Cells without a bbox count only if they share a word with the block. + """ + + block_set = set(block_words) + def touches(cell: TableCell) -> bool: + if cell.bbox is None: + return not block_set.isdisjoint(_tokens(cell.text)) + return box.intersection_over_self(clamp_to_page(flip_to_top_left(cell.bbox, ctx.height), ctx)) > 0 -def _sort_key(doc: DoclingDocument, node: NodeItem) -> tuple[float, float, float]: - """Position a body child by the earliest provenance anywhere in its subtree.""" - positions = [] - stack = [node] - while stack: - current = stack.pop() - for prov in getattr(current, "prov", []) or []: - positions.append((float(prov.page_no), prov.bbox.t, prov.bbox.l)) - stack.extend(child.resolve(doc) for child in current.children) - return min(positions) if positions else (float("inf"),) * 3 + under = [cell for cell in container.cells or [] if touches(cell)] + cell_words = [w for cell in under for w in _tokens(cell.text)] + if not cell_words: + return False + known = set(cell_words) + covered = sum(w in known for w in block_words) / len(block_words) + return covered >= CONTAINMENT_THRESHOLD and len(block_words) >= CELL_SHARE_THRESHOLD * len(cell_words) -def sort_body_by_position(doc: DoclingDocument) -> None: - """Restore geometric reading order across text, tables and pictures alike. +def is_contained(block: LayoutBlock, bbox: BoundingBox, containers: Sequence[LayoutBlock], ctx: PageContext) -> bool: + """Whether `block` duplicates a table or picture on its page. - The dedup pass has to add tables and pictures first, which leaves them ahead of every - paragraph in the body; this puts the page back in the order a reader would see it. + The bbox must be mostly inside the container; for a table the block's words must also be + what the cells under it hold, so a false table over real prose or a row it missed drops nothing. """ - doc.body.children.sort(key=lambda ref: _sort_key(doc, ref.resolve(doc))) + words = _tokens(block.text) + for container in containers: + container_box = clamp_to_page(flip_to_top_left(container.bbox, ctx.height), ctx) + if bbox.intersection_over_self(container_box) < CONTAINMENT_THRESHOLD: + continue + if container.label in PICTURE_LABELS or not words: + return True + if _carries_text(words, container, bbox, ctx): + return True + return False def _page_of(item: NodeItem) -> Optional[int]: @@ -184,19 +199,15 @@ def link_captions(doc: DoclingDocument) -> None: def append_blocks( doc: DoclingDocument, ctx: PageContext, - blocks: list[LayoutBlock], + blocks: Sequence[LayoutBlock], ) -> list[NodeItem]: - """Add one page's blocks, then restore reading order. - - Tables and pictures go in first so the text pass can drop anything they already contain — - without that ordering the same words land in the document twice, once as a table cell and - once as a stray paragraph. - """ + """Add one page's blocks in the order given, dropping text a table or picture already carries.""" from extralit_server.contexts.ocr.figures import add_picture_block from extralit_server.contexts.ocr.tables import add_table_block from extralit_server.contexts.ocr.text import add_text_block register_page(doc, ctx) + containers = [b for b in blocks if b.label in TABLE_LABELS or b.label in PICTURE_LABELS] added: list[NodeItem] = [] for block in blocks: @@ -204,14 +215,8 @@ def append_blocks( added.append(add_table_block(doc, block, ctx)) elif block.label in PICTURE_LABELS: added.append(add_picture_block(doc, block, ctx)) - - for block in blocks: - if block.label in TABLE_LABELS or block.label in PICTURE_LABELS: - continue - item = add_text_block(doc, block, ctx) - if item is not None: + elif (item := add_text_block(doc, block, ctx, containers=containers)) is not None: added.append(item) - sort_body_by_position(doc) link_captions(doc) return added diff --git a/extralit-server/src/extralit_server/contexts/ocr/parsers/pdf_inspector.py b/extralit-server/src/extralit_server/contexts/ocr/parsers/pdf_inspector.py index 0bcd94404..ee3d410c4 100644 --- a/extralit-server/src/extralit_server/contexts/ocr/parsers/pdf_inspector.py +++ b/extralit-server/src/extralit_server/contexts/ocr/parsers/pdf_inspector.py @@ -254,6 +254,21 @@ def _blocks_for_page( return blocks +class PdfInspectorParser: + """`PageParser` over pdf-inspector; it reads the text layer, so `ocr` has no effect.""" + + def parse_pages(self, pdf_bytes: bytes, pages: Sequence[int], *, ocr: bool) -> dict[int, list[LayoutBlock]]: + """Parse the given 1-indexed pages into top-left-origin blocks, keyed by page number.""" + if not pages: + return {} + sizes = page_sizes(pdf_bytes) + roles = _structure_roles(pdf_bytes) + by_page: dict[int, list[Any]] = defaultdict(list) + for item in pdf_inspector.extract_text_with_positions_bytes(pdf_bytes): + by_page[item.page].append(item) + return {p: _blocks_for_page(by_page.get(p, []), roles, p) for p in sorted(set(pages)) if p in sizes} + + def parse( pdf_bytes: bytes, *, diff --git a/extralit-server/src/extralit_server/contexts/ocr/parsers/pymupdf.py b/extralit-server/src/extralit_server/contexts/ocr/parsers/pymupdf.py index 113808b28..75b8ead9e 100644 --- a/extralit-server/src/extralit_server/contexts/ocr/parsers/pymupdf.py +++ b/extralit-server/src/extralit_server/contexts/ocr/parsers/pymupdf.py @@ -13,7 +13,7 @@ import pymupdf from docling_core.types.doc import BoundingBox, CoordOrigin, DocItemLabel, DoclingDocument, Size, TableCell -# pymupdf4llm drops its top-level IdentifyHeaders/TocHeaders when pymupdf-layout is installed. +# pymupdf-layout removes these names from the package namespace; the submodule keeps them. from pymupdf4llm.helpers.pymupdf_rag import IdentifyHeaders, TocHeaders from extralit_server.contexts.ocr.docling_builder import ( @@ -154,6 +154,18 @@ def _blocks_for_page(page: pymupdf.Page, headers: Any) -> list[LayoutBlock]: return blocks +class PyMuPDFParser: + """`PageParser` over pymupdf; it reads the text layer, so `ocr` has no effect.""" + + def parse_pages(self, pdf_bytes: bytes, pages: Sequence[int], *, ocr: bool) -> dict[int, list[LayoutBlock]]: + """Parse the given 1-indexed pages into top-left-origin blocks, keyed by page number.""" + if not pages: + return {} + with pymupdf.open(stream=pdf_bytes, filetype="pdf") as pdf: + headers = _header_levels(pdf, pages) + return {p: _blocks_for_page(pdf[p - 1], headers) for p in sorted(set(pages)) if 0 < p <= pdf.page_count} + + def parse( pdf_bytes: bytes, *, diff --git a/extralit-server/src/extralit_server/contexts/ocr/parsers/registry.py b/extralit-server/src/extralit_server/contexts/ocr/parsers/registry.py index c228161f9..9417fb53c 100644 --- a/extralit-server/src/extralit_server/contexts/ocr/parsers/registry.py +++ b/extralit-server/src/extralit_server/contexts/ocr/parsers/registry.py @@ -1,42 +1,22 @@ -"""Swappable PDF→`DoclingDocument` parsers. - -Each parser normalizes its backend into `LayoutBlock`s and hands them to the shared builder, -so the document that comes out is the same shape regardless of which one ran. -""" +"""Named `PageParser`s available in this environment, and the default `ParserService` over them.""" from __future__ import annotations import logging -from collections.abc import Sequence -from typing import Optional, Protocol - -from docling_core.types.doc import DoclingDocument -from extralit_server.contexts.ocr.parsers.pdf_inspector import parse as _parse_pdf_inspector +from extralit_server.contexts.ocr.parsers.liteparse import LiteParseParser +from extralit_server.contexts.ocr.parsers.pdf_inspector import PdfInspectorParser +from extralit_server.contexts.ocr.parsers.service import DIGITAL_PARSER, OCR_PARSER, PageParser, ParserService try: - from extralit_server.contexts.ocr.parsers.pymupdf import parse as _parse_pymupdf + from extralit_server.contexts.ocr.parsers.pymupdf import PyMuPDFParser except ImportError as e: # AGPL extra, deliberately optional - _parse_pymupdf = None + PyMuPDFParser = None logging.getLogger(__name__).debug(f"pymupdf layout parser unavailable: {e}") - -class LayoutParser(Protocol): - """Parse PDF bytes into a `DoclingDocument`.""" - - def __call__( - self, - pdf_bytes: bytes, - *, - name: str, - pages: Optional[Sequence[int]] = None, - filename: Optional[str] = None, - ) -> DoclingDocument: ... - - -_PARSERS: dict[str, LayoutParser] = {"pdf_inspector": _parse_pdf_inspector} -if _parse_pymupdf is not None: - _PARSERS["pymupdf"] = _parse_pymupdf +_PARSERS: dict[str, PageParser] = {"pdf_inspector": PdfInspectorParser(), OCR_PARSER: LiteParseParser()} +if PyMuPDFParser is not None: + _PARSERS[DIGITAL_PARSER] = PyMuPDFParser() def list_parsers() -> list[str]: @@ -44,7 +24,7 @@ def list_parsers() -> list[str]: return sorted(_PARSERS) -def get_parser(name: str) -> LayoutParser: +def get_parser(name: str) -> PageParser: """Look up a parser by name.""" try: return _PARSERS[name] @@ -52,9 +32,9 @@ def get_parser(name: str) -> LayoutParser: raise ValueError(f"unknown layout parser {name!r}; available: {list_parsers()}") from None -def default_parser_name() -> str: - """Prefer pymupdf's higher-fidelity geometry when the extra is installed.""" - return "pymupdf" if "pymupdf" in _PARSERS else "pdf_inspector" +def default_service() -> ParserService: + """Routing service over every installed parser; pdf_inspector is only used when named explicitly.""" + return ParserService(_PARSERS) -__all__ = ["LayoutParser", "default_parser_name", "get_parser", "list_parsers"] +__all__ = ["default_service", "get_parser", "list_parsers"] diff --git a/extralit-server/src/extralit_server/contexts/ocr/parsers/service.py b/extralit-server/src/extralit_server/contexts/ocr/parsers/service.py index 6ea1014f6..69531899d 100644 --- a/extralit-server/src/extralit_server/contexts/ocr/parsers/service.py +++ b/extralit-server/src/extralit_server/contexts/ocr/parsers/service.py @@ -3,7 +3,8 @@ 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. +only blocks mostly inside that region are kept. Each parser's own reading order is kept; blocks +from a later task on the same page are slotted in by position (see `merge_blocks`). """ from __future__ import annotations @@ -22,6 +23,7 @@ PageContext, append_blocks, content_hash, + flip_to_top_left, new_document, register_page, ) @@ -31,15 +33,22 @@ OCR_PARSER = "liteparse" DIGITAL_PARSER = "pymupdf" +#: Triage reason for a page that already has a (hidden) text layer. +TEXT_LAYER_REASON = "invisible_text_layer" + 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]]: ... + def parse_pages(self, pdf_bytes: bytes, pages: Sequence[int], *, ocr: bool) -> dict[int, list[LayoutBlock]]: + """Parse the given pages; `ocr` asks for OCR where the parser supports it.""" + ... @dataclass(frozen=True) class ParseTask: + """One parser run over a page subset, optionally with OCR and clipped to a region.""" + pages: tuple[int, ...] parser: str ocr: bool = False @@ -48,19 +57,58 @@ class ParseTask: @dataclass(frozen=True) class ParsePlan: + """The tasks that together parse a document, plus the triage facts behind them.""" + tasks: tuple[ParseTask, ...] + #: Triage's full OCR-needed list, regardless of the page allowlist. + pages_needing_ocr: tuple[int, ...] = () + #: OCR-planned pages that already carry a text layer, which liteparse reads without running OCR. + text_layer_pages: tuple[int, ...] = () @property def pages_by_parser(self) -> dict[str, list[int]]: + """Sorted pages assigned to each parser.""" 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]: + def planned_ocr_pages(self) -> list[int]: + """Pages in a task with OCR enabled.""" return sorted({p for task in self.tasks if task.ocr for p in task.pages}) + @property + def ocr_pages(self) -> list[int]: + """Pages where OCR actually runs: planned, minus those whose text layer is read instead.""" + return [p for p in self.planned_ocr_pages if p not in set(self.text_layer_pages)] + + @property + def ocr_skipped_pages(self) -> list[int]: + """OCR-planned pages whose text layer is read instead of running OCR.""" + return [p for p in self.planned_ocr_pages if p in set(self.text_layer_pages)] + + +def _top_left_box(block: LayoutBlock, height: float) -> BoundingBox: + return flip_to_top_left(block.bbox, height) + + +def _position(block: LayoutBlock, height: float) -> tuple[float, float]: + box = _top_left_box(block, height) + return (box.t, box.l) + + +def merge_blocks(base: Sequence[LayoutBlock], extra: Sequence[LayoutBlock], height: float) -> list[LayoutBlock]: + """Slot `extra` into `base` by `(top, left)`, keeping each list's own internal order.""" + merged = list(base) + cursor = 0 + for block in extra: + key = _position(block, height) + index = next((i for i in range(cursor, len(merged)) if _position(merged[i], height) > key), len(merged)) + merged.insert(index, block) + cursor = index + 1 + return merged + def build_document( blocks: Mapping[int, Sequence[LayoutBlock]], @@ -70,11 +118,11 @@ def build_document( filename: Optional[str] = None, binary_hash: Optional[int] = None, ) -> DoclingDocument: - """Lay per-page blocks into a fresh `DoclingDocument`; every known page is registered.""" + """Lay per-page blocks into a fresh `DoclingDocument`, in the order given; 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)) + page_blocks = blocks.get(page_no, ()) if page_blocks: append_blocks(doc, ctx, page_blocks) else: @@ -83,6 +131,8 @@ def build_document( class ParserService: + """Plan per-page parser routing for a PDF and merge the parsers' output into one `DoclingDocument`.""" + def __init__( self, parsers: Mapping[str, PageParser], @@ -100,21 +150,36 @@ def _parser(self, name: str) -> PageParser: 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: + def plan( + self, + pdf_bytes: bytes, + *, + pages: Optional[Sequence[int]] = None, + parser: Optional[str] = None, + ) -> ParsePlan: + """Route pages to parsers; `parser` forces one parser for every page (liteparse still OCRs flagged pages).""" 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 + allowed = None if pages is None else set(pages) + wanted = [p for p in range(1, page_count + 1) if allowed is None or p in allowed] + + flagged = set(triage.pages_needing_ocr) + scanned = [p for p in wanted if p in flagged] + digital = [p for p in wanted if p not in flagged] + + if parser is not None: + self._parser(parser) + if parser == OCR_PARSER: + routes = [(digital, parser, False), (scanned, parser, True)] + else: + routes = [(wanted, parser, False)] + else: + digital_parser = DIGITAL_PARSER if DIGITAL_PARSER in self._parsers else OCR_PARSER + routes = [(digital, digital_parser, False), (scanned, OCR_PARSER, True)] - 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)) + tasks = tuple(ParseTask(pages=tuple(p), parser=name, ocr=ocr) for p, name, ocr in routes if p) + text_layer = tuple(p for p in scanned if TEXT_LAYER_REASON in triage.ocr_reasons_by_page.get(str(p), ())) + return ParsePlan(tasks=tasks, pages_needing_ocr=tuple(sorted(flagged)), text_layer_pages=text_layer) def parse( self, @@ -124,22 +189,28 @@ def parse( name: str, filename: Optional[str] = None, ) -> DoclingDocument: - """Run every task and merge their blocks per page, ordered by `(top, left)`.""" + """Run every task and merge their blocks per page; blocking, so async callers use `to_thread`.""" + sizes = self._page_sizes(pdf_bytes) parsers = {task.parser: self._parser(task.parser) for task in plan.tasks} - merged: dict[int, list[LayoutBlock]] = defaultdict(list) + merged: dict[int, list[LayoutBlock]] = {} for task in plan.tasks: for page_no, blocks in parsers[task.parser].parse_pages(pdf_bytes, task.pages, ocr=task.ocr).items(): + height = sizes[page_no].height if page_no in sizes else 0.0 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) + blocks = [ + b + for b in blocks + if _top_left_box(b, height).intersection_over_self(task.bbox) >= CONTAINMENT_THRESHOLD + ] + merged[page_no] = merge_blocks(merged.get(page_no, []), blocks, height) return build_document( merged, - page_sizes=self._page_sizes(pdf_bytes), + page_sizes=sizes, name=name, filename=filename, binary_hash=content_hash(pdf_bytes), ) -__all__ = ["PageParser", "ParsePlan", "ParseTask", "ParserService", "build_document"] +__all__ = ["PageParser", "ParsePlan", "ParseTask", "ParserService", "build_document", "merge_blocks"] diff --git a/extralit-server/src/extralit_server/contexts/ocr/text.py b/extralit-server/src/extralit_server/contexts/ocr/text.py index 00822ef6d..61e24c84b 100644 --- a/extralit-server/src/extralit_server/contexts/ocr/text.py +++ b/extralit-server/src/extralit_server/contexts/ocr/text.py @@ -2,6 +2,7 @@ from __future__ import annotations +from collections.abc import Sequence from typing import Optional from docling_core.types.doc import ContentLayer, DocItemLabel, DoclingDocument @@ -24,6 +25,7 @@ def add_text_block( block: LayoutBlock, ctx: PageContext, parent: Optional[NodeItem] = None, + containers: Sequence[LayoutBlock] = (), ) -> Optional[NodeItem]: """Add one text block, skipping anything empty or already covered by a table/picture.""" text = (block.text or "").strip() @@ -32,7 +34,7 @@ def add_text_block( prov = make_prov(ctx, block.bbox, text) - if block.label not in CONTAINMENT_EXEMPT_LABELS and is_contained(doc, prov.bbox, ctx.page_no): + if block.label not in CONTAINMENT_EXEMPT_LABELS and is_contained(block, prov.bbox, containers, ctx): return None # add_text dispatches SECTION_HEADER to add_heading but drops the level, so route it here. diff --git a/extralit-server/src/extralit_server/jobs/ocr_jobs.py b/extralit-server/src/extralit_server/jobs/ocr_jobs.py index bddc9148f..aeb0580e0 100644 --- a/extralit-server/src/extralit_server/jobs/ocr_jobs.py +++ b/extralit-server/src/extralit_server/jobs/ocr_jobs.py @@ -1,5 +1,6 @@ """Document layout extraction jobs.""" +import asyncio import logging from collections.abc import Sequence from typing import Any, Optional @@ -13,8 +14,7 @@ from extralit_server.contexts.document.metadata import update_processing_metadata from extralit_server.contexts.ocr import storage as layout_storage from extralit_server.contexts.ocr.layout_store import LayoutStore -from extralit_server.contexts.ocr.parsers.pdf_inspector import classify -from extralit_server.contexts.ocr.parsers.registry import default_parser_name, get_parser +from extralit_server.contexts.ocr.parsers.registry import default_service from extralit_server.contexts.workflows import writer_skip_reason from extralit_server.database import AsyncSessionLocal from extralit_server.jobs.queues import OCR_QUEUE, REDIS_CONNECTION @@ -22,20 +22,6 @@ _LOGGER = logging.getLogger(__name__) -def route_parser(pdf_bytes: bytes) -> tuple[str, dict[str, Any]]: - """Pick a parser and record what classification saw. - - `pages_needing_ocr` is surfaced rather than acted on, so scanned pages show up as an - explicit gap instead of silently producing an empty layout. - """ - try: - classification = classify(pdf_bytes) - except Exception as e: - _LOGGER.warning(f"PDF classification failed, falling back to the default parser: {e}") - classification = {"pages_needing_ocr": [], "page_count": 0, "pdf_type": "unknown"} - return default_parser_name(), classification - - @job( queue=OCR_QUEUE, connection=REDIS_CONNECTION, @@ -56,7 +42,7 @@ async def async_document_layout_job( document_id: UUID of the document to process s3_url: proxy URL of the PDF, as stored on the document workspace_name: workspace the artifacts are written to - parser: layout parser name; None routes automatically + parser: force one parser for every page; None routes per page pages: 1-indexed page allowlist; None processes every page Returns: @@ -78,15 +64,14 @@ async def async_document_layout_job( storage = await files.get_storage() pdf_bytes = await files.download_file_content(storage, s3_url) - routed, classification = route_parser(pdf_bytes) - parser_name = parser or routed + # Both calls are CPU-bound and sync (liteparse holds the GIL for long stretches), so keep them off the loop. + service = default_service() + plan = await asyncio.to_thread(service.plan, pdf_bytes, pages=pages, parser=parser) + parser_name = "+".join(sorted(plan.pages_by_parser)) or (parser or "") _LOGGER.info(f"Extracting layout for document {document_id} with parser {parser_name}") - doc = get_parser(parser_name)( - pdf_bytes, - name=str(document_id), - pages=pages, - filename=s3_url.split("/")[-1], + doc = await asyncio.to_thread( + service.parse, pdf_bytes, plan, name=str(document_id), filename=s3_url.split("/")[-1] ) # Ordering closes the delete-vs-running-job race: nothing may write rows for a document @@ -108,7 +93,10 @@ async def async_document_layout_job( docling_version=doc.version, num_items=sum(1 for _ in doc.iterate_items(with_groups=False)), num_pages=len(doc.pages), - pages_needing_ocr=classification.get("pages_needing_ocr", []), + pages_needing_ocr=list(plan.pages_needing_ocr), + parsers=plan.pages_by_parser, + ocr_pages=plan.ocr_pages, + ocr_skipped_pages=plan.ocr_skipped_pages, ) # Outside the workspace lock: the row lock only serializes writers of this JSON column. diff --git a/extralit-server/src/extralit_server/workflows/documents.py b/extralit-server/src/extralit_server/workflows/documents.py index 77794c6d4..2e11af492 100644 --- a/extralit-server/src/extralit_server/workflows/documents.py +++ b/extralit-server/src/extralit_server/workflows/documents.py @@ -5,7 +5,6 @@ from rq.group import Group from rq.job import Dependency -from extralit_server.contexts.ocr.parsers.registry import default_parser_name from extralit_server.database import AsyncSessionLocal from extralit_server.jobs.document_jobs import analysis_and_preprocess_job from extralit_server.jobs.ocr_jobs import async_document_layout_job @@ -38,7 +37,7 @@ async def create_document_workflow( reference: Reference key for tracking workspace_name: Workspace name for job context workspace_id: UUID of the workspace - layout_parser: Layout parser to run; None uses the default parser + layout_parser: Force one layout parser for every page; None routes per page Returns: Dictionary containing workflow_id and group_id for tracking @@ -106,7 +105,7 @@ def meta(step: str) -> dict[str, str]: # classifies those pages but bundles no engine, so today they stay an explicitly surfaced gap. layout_job_data = OCR_QUEUE.prepare_data( async_document_layout_job, - (document_id, s3_url, workspace_name, layout_parser or default_parser_name()), + (document_id, s3_url, workspace_name, layout_parser), timeout=1800, job_id=f"document_layout_{document_id}_{run_suffix}", depends_on=on_analysis, diff --git a/extralit-server/tests/unit/contexts/ocr/test_docling_builder.py b/extralit-server/tests/unit/contexts/ocr/test_docling_builder.py index 64dc30791..9dc7d1f38 100644 --- a/extralit-server/tests/unit/contexts/ocr/test_docling_builder.py +++ b/extralit-server/tests/unit/contexts/ocr/test_docling_builder.py @@ -13,6 +13,7 @@ new_document, register_page, ) +from extralit_server.contexts.ocr.tables import make_cell PAGE_HEIGHT = 792.0 PAGE_WIDTH = 612.0 @@ -105,44 +106,98 @@ def test_bbox_is_clamped_into_the_page(self, ctx): assert prov.bbox.r == PAGE_WIDTH +def _table(box: BoundingBox, *cell_texts: str) -> LayoutBlock: + cells = [make_cell(text, row=0, col=i) for i, text in enumerate(cell_texts)] + return LayoutBlock(label=DocItemLabel.TABLE, bbox=box, cells=cells) + + +def _text(box: BoundingBox, text: str) -> LayoutBlock: + return LayoutBlock(label=DocItemLabel.TEXT, bbox=box, text=text) + + class TestIsContained: - def test_text_inside_an_existing_table_is_contained(self, doc, ctx): - register_page(doc, ctx) - append_blocks(doc, ctx, [LayoutBlock(label=DocItemLabel.TABLE, bbox=bbox(t=100, b=300, left=0, right=500))]) + def test_text_inside_a_table_that_carries_its_words_is_contained(self, ctx): + table = _table(bbox(t=100, b=300, left=0, right=500), "cell", "text") + block = _text(bbox(t=150, b=170, left=10, right=200), "cell text") - assert is_contained(doc, bbox(t=150, b=170, left=10, right=200), page_no=1) + assert is_contained(block, block.bbox, [table], ctx) - def test_text_outside_every_table_is_not_contained(self, doc, ctx): - register_page(doc, ctx) - append_blocks(doc, ctx, [LayoutBlock(label=DocItemLabel.TABLE, bbox=bbox(t=100, b=300, left=0, right=500))]) + def test_text_inside_a_table_that_lacks_its_words_is_kept(self, ctx): + table = _table(bbox(t=100, b=300, left=0, right=500), "Variable", "Number") + block = _text(bbox(t=150, b=170, left=10, right=200), "(192) (217) (212) (621)") - assert not is_contained(doc, bbox(t=400, b=420, left=10, right=200), page_no=1) + assert not is_contained(block, block.bbox, [table], ctx) - def test_overlap_below_threshold_is_not_contained(self, doc, ctx): - register_page(doc, ctx) - append_blocks(doc, ctx, [LayoutBlock(label=DocItemLabel.TABLE, bbox=bbox(t=100, b=200, left=0, right=100))]) + def test_prose_inside_one_giant_cell_is_kept(self, ctx): + # A false table whose single cell holds the whole page's text must not absorb the page's blocks. + prose = "daily biting rate parity rate sporozoite rate inoculation rate " * 5 + giant = make_cell(prose + " extra words", row=0, col=0, bbox=bbox(t=100, b=500, left=0, right=500)) + table = LayoutBlock(label=DocItemLabel.TABLE, bbox=bbox(t=100, b=500, left=0, right=500), cells=[giant]) + block = _text(bbox(t=120, b=140, left=10, right=400), "daily biting rate parity rate") + + assert not is_contained(block, block.bbox, [table], ctx) + + def test_text_matching_a_small_cell_under_it_is_contained(self, ctx): + small = make_cell("Anopheles gambiae", row=0, col=0, bbox=bbox(t=100, b=110, left=0, right=100)) + other = make_cell("14999", row=0, col=1, bbox=bbox(t=300, b=310, left=0, right=100)) + table = LayoutBlock(label=DocItemLabel.TABLE, bbox=bbox(t=90, b=320, left=0, right=200), cells=[small, other]) + block = _text(bbox(t=100, b=110, left=0, right=100), "Anopheles gambiae") + + assert is_contained(block, block.bbox, [table], ctx) + + def test_one_row_of_a_table_whose_cells_lack_bboxes_is_contained(self, ctx): + rows = [("Anopheles", "gambiae", "14999"), ("Culex", "pipiens", "212"), ("Aedes", "aegypti", "37")] + cells = [make_cell(text, row=r, col=c) for r, row in enumerate(rows) for c, text in enumerate(row)] + table = LayoutBlock(label=DocItemLabel.TABLE, bbox=bbox(t=100, b=300, left=0, right=500), cells=cells) + block = _text(bbox(t=150, b=170, left=10, right=200), "Culex pipiens 212") + + assert all(cell.bbox is None for cell in cells) + assert is_contained(block, block.bbox, [table], ctx) + + def test_prose_inside_one_giant_cell_without_a_bbox_is_kept(self, ctx): + prose = "daily biting rate parity rate sporozoite rate inoculation rate " * 5 + table = _table(bbox(t=100, b=500, left=0, right=500), prose) + block = _text(bbox(t=120, b=140, left=10, right=400), "daily biting rate parity rate") + + assert not is_contained(block, block.bbox, [table], ctx) + + def test_a_table_without_cells_absorbs_no_text(self, ctx): + table = LayoutBlock(label=DocItemLabel.TABLE, bbox=bbox(t=100, b=300, left=0, right=500)) + block = _text(bbox(t=150, b=170, left=10, right=200), "real prose") + assert not is_contained(block, block.bbox, [table], ctx) + + def test_text_outside_every_table_is_not_contained(self, ctx): + table = _table(bbox(t=100, b=300, left=0, right=500), "cell") + block = _text(bbox(t=400, b=420, left=10, right=200), "cell") + + assert not is_contained(block, block.bbox, [table], ctx) + + def test_overlap_below_threshold_is_not_contained(self, ctx): + table = _table(bbox(t=100, b=200, left=0, right=100), "cell") # Half of this box lies inside the table -> IoSelf 0.5, under the 0.6 default. - assert not is_contained(doc, bbox(t=150, b=250, left=0, right=100), page_no=1) + block = _text(bbox(t=150, b=250, left=0, right=100), "cell") - def test_overlap_above_threshold_is_contained(self, doc, ctx): - register_page(doc, ctx) - append_blocks(doc, ctx, [LayoutBlock(label=DocItemLabel.TABLE, bbox=bbox(t=100, b=200, left=0, right=100))]) + assert not is_contained(block, block.bbox, [table], ctx) - # 70% of this box lies inside the table. - assert is_contained(doc, bbox(t=130, b=230, left=0, right=100), page_no=1) + def test_overlap_above_threshold_is_contained(self, ctx): + table = _table(bbox(t=100, b=200, left=0, right=100), "cell") + block = _text(bbox(t=130, b=230, left=0, right=100), "cell") - def test_containment_is_scoped_to_the_page(self, doc, ctx): - register_page(doc, ctx) - append_blocks(doc, ctx, [LayoutBlock(label=DocItemLabel.TABLE, bbox=bbox(t=100, b=300, left=0, right=500))]) + assert is_contained(block, block.bbox, [table], ctx) - assert not is_contained(doc, bbox(t=150, b=170, left=10, right=200), page_no=2) + def test_pictures_absorb_text_by_position_alone(self, ctx): + picture = LayoutBlock(label=DocItemLabel.PICTURE, bbox=bbox(t=100, b=300, left=0, right=500)) + block = _text(bbox(t=150, b=170, left=10, right=200), "axis label") - def test_pictures_also_absorb_text(self, doc, ctx): - register_page(doc, ctx) - append_blocks(doc, ctx, [LayoutBlock(label=DocItemLabel.PICTURE, bbox=bbox(t=100, b=300, left=0, right=500))]) + assert is_contained(block, block.bbox, [picture], ctx) + + def test_bottom_left_containers_are_flipped_before_comparing(self, ctx): + box = BoundingBox(l=0, r=500, b=PAGE_HEIGHT - 300, t=PAGE_HEIGHT - 100, coord_origin=CoordOrigin.BOTTOMLEFT) + table = _table(box, "cell") + block = _text(bbox(t=150, b=170, left=10, right=200), "cell") - assert is_contained(doc, bbox(t=150, b=170, left=10, right=200), page_no=1) + assert is_contained(block, block.bbox, [table], ctx) class TestAppendBlocks: @@ -194,7 +249,7 @@ def test_reading_order_of_text_is_preserved(self, doc, ctx): def test_text_inside_a_table_is_dropped(self, doc, ctx): blocks = [ - LayoutBlock(label=DocItemLabel.TABLE, bbox=bbox(t=100, b=300, left=0, right=500)), + _table(bbox(t=100, b=300, left=0, right=500), "cell text"), LayoutBlock(label=DocItemLabel.TEXT, bbox=bbox(t=150, b=170, left=10, right=200), text="cell text"), LayoutBlock(label=DocItemLabel.TEXT, bbox=bbox(t=400, b=420), text="body text"), ] @@ -396,3 +451,19 @@ def test_document_round_trips_through_json(self, doc, ctx): assert restored.texts[0].text == "hello" assert restored.tables[0].prov[0].page_no == 1 + + +class TestReadingOrderIsKept: + def test_blocks_keep_the_order_the_parser_gave(self, doc, ctx): + # Two columns: the parser reads the whole left column, then the right; a (top, left) sort would interleave. + blocks = [ + LayoutBlock(label=DocItemLabel.TEXT, bbox=bbox(t=100, b=120, left=0, right=250), text="left one"), + LayoutBlock(label=DocItemLabel.TEXT, bbox=bbox(t=300, b=320, left=0, right=250), text="left two"), + LayoutBlock(label=DocItemLabel.TEXT, bbox=bbox(t=100, b=120, left=300, right=550), text="right one"), + LayoutBlock(label=DocItemLabel.TEXT, bbox=bbox(t=300, b=320, left=300, right=550), text="right two"), + ] + + append_blocks(doc, ctx, blocks) + + assert [t.text for t in doc.texts] == ["left one", "left two", "right one", "right two"] + assert [i.text for i, _ in doc.iterate_items(with_groups=False)] == [b.text for b in blocks] diff --git a/extralit-server/tests/unit/contexts/ocr/test_parser_service.py b/extralit-server/tests/unit/contexts/ocr/test_parser_service.py index c8b265898..98c0945dc 100644 --- a/extralit-server/tests/unit/contexts/ocr/test_parser_service.py +++ b/extralit-server/tests/unit/contexts/ocr/test_parser_service.py @@ -14,6 +14,7 @@ ParseTask, build_document, ) +from extralit_server.contexts.ocr.tables import make_cell PAGE = Size(width=100.0, height=200.0) @@ -111,9 +112,104 @@ def test_plan_summaries(self): assert plan.pages_by_parser == {"pymupdf": [1, 3], "liteparse": [2]} assert plan.ocr_pages == [2] + assert plan.ocr_skipped_pages == [] + + def test_ocr_pages_exclude_pages_read_from_a_text_layer(self): + plan = ParsePlan( + tasks=(ParseTask(pages=(1, 2), parser="liteparse", ocr=True),), + pages_needing_ocr=(1, 2), + text_layer_pages=(2,), + ) + + assert plan.ocr_pages == [1] + assert plan.ocr_skipped_pages == [2] + + +class TestPlanDetails: + def test_hidden_text_layer_pages_are_flagged_as_not_ocrd(self): + triage = TriageMetadata( + pdf_type="image_based", + page_count=3, + pages_needing_ocr=[1, 2, 3], + ocr_reasons_by_page={"1": ["invisible_text_layer"], "2": ["no_text"], "3": ["invisible_text_layer"]}, + ) + service = _service({"pymupdf": FakeParser(), "liteparse": FakeParser()}, triage) + + plan = service.plan(b"%PDF") + + assert plan.tasks == (ParseTask(pages=(1, 2, 3), parser="liteparse", ocr=True),) + assert plan.ocr_pages == [2] + assert plan.ocr_skipped_pages == [1, 3] + assert plan.pages_needing_ocr == (1, 2, 3) + + def test_a_named_parser_takes_every_page(self): + triage = TriageMetadata(pdf_type="mixed", page_count=3, pages_needing_ocr=[2]) + service = _service({"pymupdf": FakeParser(), "pdf_inspector": FakeParser()}, triage) + + plan = service.plan(b"%PDF", parser="pdf_inspector") + + assert plan.tasks == (ParseTask(pages=(1, 2, 3), parser="pdf_inspector", ocr=False),) + + def test_a_named_liteparse_still_ocrs_flagged_pages(self): + triage = TriageMetadata(pdf_type="mixed", page_count=3, pages_needing_ocr=[2]) + service = _service({"liteparse": FakeParser()}, triage) + + plan = service.plan(b"%PDF", parser="liteparse") + + assert plan.tasks == ( + ParseTask(pages=(1, 3), parser="liteparse", ocr=False), + ParseTask(pages=(2,), parser="liteparse", ocr=True), + ) + + def test_an_unknown_named_parser_raises(self): + service = _service({"pymupdf": FakeParser()}, TriageMetadata(pdf_type="mixed")) + + with pytest.raises(ValueError, match="nope"): + service.plan(b"%PDF", parser="nope") class TestParse: + def test_a_parsers_reading_order_survives_the_merge(self): + # Two columns read left-then-right: a (top, left) sort would interleave them. + parser = FakeParser({1: [_text(10, 0, "L1"), _text(100, 0, "L2"), _text(10, 50, "R1"), _text(100, 50, "R2")]}) + service = _service({"a": parser}, TriageMetadata(pdf_type="mixed"), page_count=1) + plan = ParsePlan(tasks=(ParseTask(pages=(1,), parser="a", ocr=False),)) + + doc = service.parse(b"%PDF", plan, name="doc") + + assert [t.text for t in doc.texts] == ["L1", "L2", "R1", "R2"] + + def test_a_later_task_keeps_its_own_order_when_slotted_in(self): + first = FakeParser({1: [_text(10, 0, "a"), _text(100, 0, "d")]}) + second = FakeParser({1: [_text(50, 0, "b"), _text(60, 0, "c")]}) + service = _service({"a": first, "b": second}, TriageMetadata(pdf_type="mixed"), page_count=1) + plan = ParsePlan( + tasks=( + ParseTask(pages=(1,), parser="a", ocr=False), + ParseTask(pages=(1,), parser="b", ocr=False), + ) + ) + + doc = service.parse(b"%PDF", plan, name="doc") + + assert [t.text for t in doc.texts] == ["a", "b", "c", "d"] + + def test_prose_over_a_false_table_survives_but_its_own_words_do_not(self): + table = LayoutBlock( + label=DocItemLabel.TABLE, + bbox=_box(0, 0, size=150), + cells=[make_cell("Variable", row=0, col=0), make_cell("1246", row=0, col=1)], + ) + parser = FakeParser( + {1: [table, _text(20, 10, "Table 2. Ratio of people per net"), _text(60, 10, "Variable 1246")]} + ) + service = _service({"a": parser}, TriageMetadata(pdf_type="mixed"), page_count=1) + plan = ParsePlan(tasks=(ParseTask(pages=(1,), parser="a", ocr=False),)) + + doc = service.parse(b"%PDF", plan, name="doc") + + assert [t.text for t in doc.texts] == ["Table 2. Ratio of people per net"] + 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")) @@ -129,8 +225,8 @@ def test_each_task_runs_on_its_parser_with_its_ocr_flag(self): 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")]}) + def test_blocks_from_different_tasks_merge_per_page_by_position(self): + first = FakeParser({1: [_text(10, 0, "top"), _text(50, 0, "middle")]}) second = FakeParser({1: [_text(50, 20, "middle-right"), _text(90, 0, "bottom")]}) service = _service({"a": first, "b": second}, TriageMetadata(pdf_type="mixed")) plan = ParsePlan( diff --git a/extralit-server/tests/unit/contexts/ocr/test_pdf_inspector_parser.py b/extralit-server/tests/unit/contexts/ocr/test_pdf_inspector_parser.py index d1919a676..ce4606a18 100644 --- a/extralit-server/tests/unit/contexts/ocr/test_pdf_inspector_parser.py +++ b/extralit-server/tests/unit/contexts/ocr/test_pdf_inspector_parser.py @@ -41,10 +41,11 @@ class TestRegistry: def test_pdf_inspector_is_always_registered(self): assert "pdf_inspector" in list_parsers() - def test_get_parser_returns_a_callable(self, untagged_bytes): - doc = get_parser("pdf_inspector")(untagged_bytes, name="sample") + def test_get_parser_returns_a_page_parser(self, untagged_bytes): + blocks = get_parser("pdf_inspector").parse_pages(untagged_bytes, [1, 99], ocr=False) - assert doc.name == "sample" + assert list(blocks) == [1] + assert blocks[1] def test_unknown_parser_raises(self): with pytest.raises(ValueError, match="unknown"): diff --git a/extralit-server/tests/unit/contexts/ocr/test_pymupdf_parser.py b/extralit-server/tests/unit/contexts/ocr/test_pymupdf_parser.py index ece1d814c..25e11f957 100644 --- a/extralit-server/tests/unit/contexts/ocr/test_pymupdf_parser.py +++ b/extralit-server/tests/unit/contexts/ocr/test_pymupdf_parser.py @@ -28,8 +28,19 @@ class TestRegistry: def test_pymupdf_is_registered_when_the_extra_is_installed(self): assert "pymupdf" in list_parsers() - def test_get_parser_resolves_it(self, pdf_bytes): - assert get_parser("pymupdf")(pdf_bytes, name="sample").name == "sample" + def test_get_parser_resolves_a_page_parser(self, pdf_bytes): + blocks = get_parser("pymupdf").parse_pages(pdf_bytes, [1, 99], ocr=False) + + assert list(blocks) == [1] + assert blocks[1] + + def test_ocr_flag_does_not_change_the_output(self, pdf_bytes): + parser = get_parser("pymupdf") + + assert parser.parse_pages(pdf_bytes, [1], ocr=True) == parser.parse_pages(pdf_bytes, [1], ocr=False) + + def test_no_pages_means_no_output(self, pdf_bytes): + assert get_parser("pymupdf").parse_pages(pdf_bytes, [], ocr=False) == {} class TestGeometry: diff --git a/extralit-server/tests/unit/jobs/test_ocr_jobs.py b/extralit-server/tests/unit/jobs/test_ocr_jobs.py index aadf95854..be012de76 100644 --- a/extralit-server/tests/unit/jobs/test_ocr_jobs.py +++ b/extralit-server/tests/unit/jobs/test_ocr_jobs.py @@ -1,15 +1,28 @@ """Tests for how the layout job orders its writes.""" +import threading +from pathlib import Path from unittest.mock import AsyncMock, MagicMock, patch from uuid import uuid4 import pytest +from extralit_server.api.schemas.v1.document.metadata import TriageMetadata +from extralit_server.contexts.ocr.parsers.registry import default_service +from extralit_server.contexts.ocr.parsers.service import ParserService + +FIXTURES = Path(__file__).parents[2] / "fixtures" / "pdf" + MODULE = "extralit_server.jobs.ocr_jobs" pytestmark = pytest.mark.asyncio +class StubParser: + def parse_pages(self, pdf_bytes, pages, *, ocr): + return {p: [] for p in pages} + + @pytest.fixture def job_context(tmp_path, monkeypatch): """Everything the layout job touches outside its own logic, with the document present.""" @@ -17,8 +30,12 @@ def job_context(tmp_path, monkeypatch): calls = [] store = LayoutStore(str(tmp_path / "ws" / "layout")) - doc = MagicMock(version="1.7.0", pages={1: MagicMock()}) - doc.iterate_items.return_value = [] + pdf_bytes = (FIXTURES / "two_page.pdf").read_bytes() + triage = TriageMetadata( + pdf_type="mixed", page_count=2, pages_needing_ocr=[2], ocr_reasons_by_page={"2": ["invisible_text_layer"]} + ) + service = ParserService({"pymupdf": StubParser(), "liteparse": StubParser()}, triage=lambda _: triage) + layouts = [] async def store_layout(s3_client, workspace_name, document_id, document, store=None): calls.append(("store_layout", store._lock_depth)) @@ -26,15 +43,17 @@ async def store_layout(s3_client, workspace_name, document_id, document, store=N async def update_metadata(db, document_id, mutate): calls.append(("update_metadata", store._lock_depth)) + holder = MagicMock() + mutate(holder) + layouts.append(holder.layout_metadata) return None db = MagicMock() with ( patch(f"{MODULE}.files.get_storage", AsyncMock(return_value=AsyncMock())), - patch(f"{MODULE}.files.download_file_content", AsyncMock(return_value=b"%PDF-1.4")), - patch(f"{MODULE}.route_parser", return_value=("pdf_inspector", {"pages_needing_ocr": [2]})), - patch(f"{MODULE}.get_parser", return_value=lambda *args, **kwargs: doc), + patch(f"{MODULE}.files.download_file_content", AsyncMock(return_value=pdf_bytes)), + patch(f"{MODULE}.default_service", return_value=service), patch(f"{MODULE}.LayoutStore.for_workspace", return_value=store), patch(f"{MODULE}.layout_storage.store_layout", store_layout), patch(f"{MODULE}.update_processing_metadata", update_metadata), @@ -44,17 +63,17 @@ async def update_metadata(db, document_id, mutate): ): session.return_value.__aenter__ = AsyncMock(return_value=db) session.return_value.__aexit__ = AsyncMock(return_value=False) - yield {"calls": calls, "store": store, "skip": skip} + yield {"calls": calls, "store": store, "skip": skip, "layouts": layouts} -async def run_job(document_id=None): +async def run_job(document_id=None, parser=None): from extralit_server.jobs.ocr_jobs import async_document_layout_job return await async_document_layout_job( document_id or uuid4(), "/api/v1/file/ws/pdf/doc.pdf", "ws", - "pdf_inspector", + parser, ) @@ -85,6 +104,51 @@ async def test_pages_needing_ocr_are_surfaced(self, job_context): assert result["pages_needing_ocr"] == [2] +class TestExecutedPlanMetadata: + async def test_the_executed_plan_is_recorded(self, job_context): + result = await run_job() + + layout = job_context["layouts"][0] + assert layout.parsers == {"pymupdf": [1], "liteparse": [2]} + assert layout.parser == "liteparse+pymupdf" + assert result["parsers"] == layout.parsers + + async def test_ocr_pages_exclude_pages_whose_text_layer_was_read(self, job_context): + await run_job() + + layout = job_context["layouts"][0] + assert layout.ocr_pages == [] + assert layout.ocr_skipped_pages == [2] + assert layout.pages_needing_ocr == [2] + + async def test_the_parser_runs_off_the_event_loop(self, job_context): + threads = [] + + class Spy(StubParser): + def parse_pages(self, pdf_bytes, pages, *, ocr): + threads.append(threading.current_thread()) + return super().parse_pages(pdf_bytes, pages, ocr=ocr) + + service = ParserService( + {"pymupdf": Spy(), "liteparse": Spy()}, triage=lambda _: TriageMetadata(pdf_type="mixed") + ) + with patch(f"{MODULE}.default_service", return_value=service): + await run_job() + + assert threads + assert all(t is not threading.main_thread() for t in threads) + + async def test_layout_job_on_the_two_page_fixture_writes_parsers_and_ocr_pages(self, job_context): + pytest.importorskip("pymupdf4llm") + with patch(f"{MODULE}.default_service", default_service): + await run_job() + + layout = job_context["layouts"][0] + assert layout.parsers == {"pymupdf": [1, 2]} + assert layout.ocr_pages == [] + assert layout.num_pages == 2 + + class TestSupersededRuns: async def test_a_superseded_run_does_not_write(self, job_context): # Stopping a started job is only a request, so a forced restart can overlap this run. diff --git a/extralit-server/tests/unit/workflows/test_document_workflow_layout.py b/extralit-server/tests/unit/workflows/test_document_workflow_layout.py index 9417752fa..d9cfb4d3d 100644 --- a/extralit-server/tests/unit/workflows/test_document_workflow_layout.py +++ b/extralit-server/tests/unit/workflows/test_document_workflow_layout.py @@ -62,14 +62,12 @@ def prepared_for(calls, step): @pytest.mark.asyncio class TestLayoutJobSequencing: - async def test_layout_runs_by_default_with_the_default_parser(self, enqueued): - from extralit_server.contexts.ocr.parsers.registry import default_parser_name - + async def test_layout_runs_by_default_with_per_page_routing(self, enqueued): await run_workflow(layout_parser=None) layout = prepared_for(enqueued, "document_layout") assert layout is not None - assert layout["args"][3] == default_parser_name() + assert layout["args"][3] is None async def test_a_named_parser_overrides_the_default(self, enqueued): await run_workflow(layout_parser="pdf_inspector")