298 lines
9.3 KiB
Python
298 lines
9.3 KiB
Python
"""Docling document extraction logic — pure domain, no HTTP concerns.
|
|
|
|
Wraps the Docling library to convert documents and extract structured
|
|
per-page elements with bounding boxes and hierarchy levels.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import contextlib
|
|
import logging
|
|
import threading
|
|
from dataclasses import dataclass, field
|
|
|
|
from docling.datamodel.base_models import InputFormat
|
|
from docling.datamodel.pipeline_options import (
|
|
PdfPipelineOptions,
|
|
TableFormerMode,
|
|
TableStructureOptions,
|
|
)
|
|
from docling.document_converter import DocumentConverter, PdfFormatOption
|
|
from docling_core.types.doc import (
|
|
CodeItem,
|
|
DocItem,
|
|
FloatingItem,
|
|
FormulaItem,
|
|
GroupItem,
|
|
ListItem,
|
|
PictureItem,
|
|
SectionHeaderItem,
|
|
TableItem,
|
|
TextItem,
|
|
TitleItem,
|
|
)
|
|
|
|
from domain.bbox import to_topleft_list
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Thread lock — DocumentConverter is not thread-safe
|
|
_converter_lock = threading.Lock()
|
|
|
|
# US Letter page dimensions (points) — fallback when page size is unknown
|
|
_DEFAULT_PAGE_WIDTH = 612.0
|
|
_DEFAULT_PAGE_HEIGHT = 792.0
|
|
|
|
# Default converter (lazy-init on first request)
|
|
_default_converter: DocumentConverter | None = None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Domain value objects
|
|
# ---------------------------------------------------------------------------
|
|
|
|
@dataclass
|
|
class PageElement:
|
|
type: str
|
|
bbox: list[float]
|
|
content: str
|
|
level: int = 0
|
|
|
|
|
|
@dataclass
|
|
class PageDetail:
|
|
page_number: int
|
|
width: float
|
|
height: float
|
|
elements: list[PageElement] = field(default_factory=list)
|
|
|
|
|
|
@dataclass
|
|
class ConversionOptions:
|
|
do_ocr: bool = True
|
|
do_table_structure: bool = True
|
|
table_mode: str = "accurate"
|
|
do_code_enrichment: bool = False
|
|
do_formula_enrichment: bool = False
|
|
do_picture_classification: bool = False
|
|
do_picture_description: bool = False
|
|
generate_picture_images: bool = False
|
|
generate_page_images: bool = False
|
|
images_scale: float = 1.0
|
|
|
|
def is_default(self) -> bool:
|
|
return self == ConversionOptions()
|
|
|
|
|
|
@dataclass
|
|
class ConversionResult:
|
|
page_count: int
|
|
content_markdown: str
|
|
content_html: str
|
|
pages: list[PageDetail]
|
|
skipped_items: int = 0
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Element type detection
|
|
# ---------------------------------------------------------------------------
|
|
|
|
# Mapping from Docling type to element type string.
|
|
# Order matters: most specific types before their parents.
|
|
_ELEMENT_TYPE_MAP: list[tuple[type, str]] = [
|
|
(TableItem, "table"),
|
|
(PictureItem, "picture"),
|
|
(TitleItem, "title"),
|
|
(SectionHeaderItem, "section_header"),
|
|
(ListItem, "list"),
|
|
(FormulaItem, "formula"),
|
|
(CodeItem, "code"),
|
|
(FloatingItem, "floating"),
|
|
(TextItem, "text"),
|
|
]
|
|
|
|
|
|
def _get_element_type(item: DocItem) -> str:
|
|
"""Determine element type via isinstance on Docling's type hierarchy."""
|
|
for cls, type_name in _ELEMENT_TYPE_MAP:
|
|
if isinstance(item, cls):
|
|
return type_name
|
|
return "text"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Pipeline factory
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def build_converter(options: ConversionOptions | None = None) -> DocumentConverter:
|
|
"""Build a DocumentConverter with the given pipeline options."""
|
|
opts = options or ConversionOptions()
|
|
|
|
table_options = TableStructureOptions(
|
|
do_cell_matching=True,
|
|
mode=TableFormerMode.ACCURATE if opts.table_mode == "accurate" else TableFormerMode.FAST,
|
|
)
|
|
|
|
pipeline_options = PdfPipelineOptions(
|
|
do_ocr=opts.do_ocr,
|
|
do_table_structure=opts.do_table_structure,
|
|
table_structure_options=table_options,
|
|
do_code_enrichment=opts.do_code_enrichment,
|
|
do_formula_enrichment=opts.do_formula_enrichment,
|
|
do_picture_classification=opts.do_picture_classification,
|
|
do_picture_description=opts.do_picture_description,
|
|
generate_page_images=opts.generate_page_images,
|
|
generate_picture_images=opts.generate_picture_images,
|
|
images_scale=opts.images_scale,
|
|
)
|
|
|
|
return DocumentConverter(
|
|
format_options={
|
|
InputFormat.PDF: PdfFormatOption(pipeline_options=pipeline_options),
|
|
}
|
|
)
|
|
|
|
|
|
def get_default_converter() -> DocumentConverter:
|
|
global _default_converter
|
|
if _default_converter is None:
|
|
_default_converter = build_converter()
|
|
return _default_converter
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Page extraction
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def extract_pages_detail(doc_result) -> tuple[list[PageDetail], int]:
|
|
"""Extract per-page element details with bounding boxes from Docling result.
|
|
|
|
Returns (pages, skipped_count) for transparent error reporting.
|
|
"""
|
|
pages: dict[int, PageDetail] = {}
|
|
document = doc_result.document
|
|
skipped = 0
|
|
|
|
for page_key, page_obj in document.pages.items():
|
|
page_no = int(page_key) if isinstance(page_key, str) else page_key
|
|
pages[page_no] = PageDetail(
|
|
page_number=page_no,
|
|
width=page_obj.size.width,
|
|
height=page_obj.size.height,
|
|
)
|
|
|
|
for item, level in document.iterate_items():
|
|
ok = _process_content_item(item, level, pages)
|
|
if not ok:
|
|
skipped += 1
|
|
|
|
sorted_pages = sorted(pages.values(), key=lambda p: p.page_number)
|
|
return sorted_pages, skipped
|
|
|
|
|
|
def _process_content_item(
|
|
item: DocItem | GroupItem, level: int, pages: dict[int, PageDetail],
|
|
) -> bool:
|
|
"""Process a single content item and add it to the appropriate page."""
|
|
if isinstance(item, GroupItem):
|
|
return True
|
|
|
|
if not isinstance(item, DocItem) or not item.prov:
|
|
return False
|
|
|
|
for prov in item.prov:
|
|
try:
|
|
page_no = prov.page_no
|
|
if page_no not in pages:
|
|
# Fallback: page was not found in document.pages (corrupted PDF or
|
|
# Docling edge case). US Letter dimensions are used as a safe default.
|
|
# This may cause slight bbox misalignment on non-Letter pages (e.g. A4).
|
|
logger.warning(
|
|
"Page %d not found in document metadata — using US Letter fallback (%sx%s pt)",
|
|
page_no, _DEFAULT_PAGE_WIDTH, _DEFAULT_PAGE_HEIGHT,
|
|
)
|
|
pages[page_no] = PageDetail(page_number=page_no, width=_DEFAULT_PAGE_WIDTH, height=_DEFAULT_PAGE_HEIGHT)
|
|
|
|
page_height = pages[page_no].height
|
|
|
|
bbox = [0.0, 0.0, 0.0, 0.0]
|
|
if prov.bbox:
|
|
bbox = to_topleft_list(prov.bbox, page_height)
|
|
|
|
element_type = _get_element_type(item)
|
|
|
|
content = getattr(item, "text", "") or ""
|
|
if isinstance(item, TableItem):
|
|
with contextlib.suppress(AttributeError, ValueError):
|
|
content = item.export_to_markdown()
|
|
|
|
pages[page_no].elements.append(
|
|
PageElement(type=element_type, bbox=bbox, content=content, level=level)
|
|
)
|
|
except (AttributeError, KeyError, TypeError, ValueError):
|
|
logger.warning(
|
|
"Skipping item %s on page %s",
|
|
type(item).__name__,
|
|
getattr(prov, "page_no", "?"),
|
|
exc_info=True,
|
|
)
|
|
return False
|
|
|
|
return True
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Main conversion entry point
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _select_converter(options: ConversionOptions) -> DocumentConverter:
|
|
"""Return the cached default converter or build a custom one."""
|
|
if options.is_default():
|
|
return get_default_converter()
|
|
return build_converter(options)
|
|
|
|
|
|
def _build_fallback_pages(doc, page_count: int) -> list[PageDetail]:
|
|
"""Create empty PageDetail entries when extraction yields nothing."""
|
|
return [
|
|
PageDetail(
|
|
page_number=i + 1,
|
|
width=doc.pages[i + 1].size.width if (i + 1) in doc.pages else _DEFAULT_PAGE_WIDTH,
|
|
height=doc.pages[i + 1].size.height if (i + 1) in doc.pages else _DEFAULT_PAGE_HEIGHT,
|
|
)
|
|
for i in range(page_count)
|
|
]
|
|
|
|
|
|
def convert_document(
|
|
file_path: str,
|
|
options: ConversionOptions | None = None,
|
|
) -> ConversionResult:
|
|
"""Convert a document and return structured results.
|
|
|
|
This is the main entry point for document parsing. Runs synchronously
|
|
(caller should use asyncio.to_thread for non-blocking execution).
|
|
"""
|
|
opts = options or ConversionOptions()
|
|
|
|
with _converter_lock:
|
|
conv = _select_converter(opts)
|
|
result = conv.convert(file_path)
|
|
|
|
doc = result.document
|
|
page_count = len(doc.pages)
|
|
pages_detail, skipped = extract_pages_detail(result)
|
|
|
|
if not pages_detail and page_count > 0:
|
|
pages_detail = _build_fallback_pages(doc, page_count)
|
|
|
|
if skipped > 0:
|
|
logger.info("Parsed: %d pages, %d items skipped", page_count, skipped)
|
|
|
|
return ConversionResult(
|
|
page_count=page_count or len(pages_detail) or 1,
|
|
content_markdown=doc.export_to_markdown(),
|
|
content_html=doc.export_to_html(),
|
|
pages=pages_detail,
|
|
skipped_items=skipped,
|
|
)
|