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 (%s×%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,
|
||
)
|