From 4b1ec364f4763186c8d311f770b7e607e274e77f Mon Sep 17 00:00:00 2001 From: Pier-Jean Malandrino Date: Thu, 2 Apr 2026 12:05:20 +0200 Subject: [PATCH] Add LocalChunker adapter and DoclingDocument serialization LocalChunker implements DocumentChunker port using docling-core chunkers. LocalConverter now serializes DoclingDocument to JSON for re-chunking support. --- document-parser/infra/local_chunker.py | 84 ++++++++++++++++++++++++ document-parser/infra/local_converter.py | 26 ++++++-- 2 files changed, 105 insertions(+), 5 deletions(-) create mode 100644 document-parser/infra/local_chunker.py diff --git a/document-parser/infra/local_chunker.py b/document-parser/infra/local_chunker.py new file mode 100644 index 0000000..7024496 --- /dev/null +++ b/document-parser/infra/local_chunker.py @@ -0,0 +1,84 @@ +"""Local Docling chunker — runs chunking in-process using docling-core. + +This adapter implements the DocumentChunker port. It deserializes a +DoclingDocument from JSON, applies the requested chunker, and returns +domain ChunkResult objects. +""" + +from __future__ import annotations + +import asyncio +import json +import logging + +from docling_core.transforms.chunker import HierarchicalChunker +from docling_core.transforms.chunker.hybrid_chunker import HybridChunker +from docling_core.types.doc.document import DoclingDocument + +from domain.value_objects import ChunkingOptions, ChunkResult + +logger = logging.getLogger(__name__) + + +def _chunk_sync(document_json: str, options: ChunkingOptions) -> list[ChunkResult]: + if not document_json or not document_json.strip(): + raise ValueError("Empty document JSON — nothing to chunk") + + try: + doc_data = json.loads(document_json) + except json.JSONDecodeError as e: + raise ValueError(f"Malformed document JSON: {e}") from e + + doc = DoclingDocument.model_validate(doc_data) + + chunker = _build_chunker(options) + results: list[ChunkResult] = [] + + for chunk in chunker.chunk(doc): + source_page = None + token_count = 0 + + if hasattr(chunk, "meta") and chunk.meta and chunk.meta.doc_items: + for doc_item in chunk.meta.doc_items: + if hasattr(doc_item, "prov") and doc_item.prov: + source_page = doc_item.prov[0].page_no + break + + if hasattr(chunker, "tokenizer") and chunker.tokenizer: + token_count = chunker.tokenizer.count_tokens(chunk.text) + + headings = list(chunk.meta.headings) if chunk.meta and chunk.meta.headings else [] + + results.append( + ChunkResult( + text=chunk.text, + headings=headings, + source_page=source_page, + token_count=token_count, + ) + ) + + logger.info("Chunked document into %d chunks (chunker=%s)", len(results), options.chunker_type) + return results + + +def _build_chunker(options: ChunkingOptions) -> HierarchicalChunker | HybridChunker: + if options.chunker_type == "hierarchical": + return HierarchicalChunker() + + return HybridChunker( + max_tokens=options.max_tokens, + merge_peers=options.merge_peers, + repeat_table_header=options.repeat_table_header, + ) + + +class LocalChunker: + """Adapter that runs docling-core chunking locally.""" + + async def chunk( + self, + document_json: str, + options: ChunkingOptions, + ) -> list[ChunkResult]: + return await asyncio.to_thread(_chunk_sync, document_json, options) diff --git a/document-parser/infra/local_converter.py b/document-parser/infra/local_converter.py index efd4a5a..9e0b23a 100644 --- a/document-parser/infra/local_converter.py +++ b/document-parser/infra/local_converter.py @@ -9,6 +9,7 @@ from __future__ import annotations import asyncio import contextlib +import json import logging import threading @@ -83,6 +84,7 @@ def _get_element_type(item: DocItem) -> str: # Pipeline factory # --------------------------------------------------------------------------- + def _build_docling_converter(options: ConversionOptions) -> DoclingConverter: table_options = TableStructureOptions( do_cell_matching=True, @@ -126,6 +128,7 @@ def _select_converter(options: ConversionOptions) -> DoclingConverter: # Page extraction # --------------------------------------------------------------------------- + def _extract_pages_detail(doc_result) -> tuple[list[PageDetail], int]: pages: dict[int, PageDetail] = {} document = doc_result.document @@ -149,7 +152,9 @@ def _extract_pages_detail(doc_result) -> tuple[list[PageDetail], int]: def _process_content_item( - item: DocItem | GroupItem, level: int, pages: dict[int, PageDetail], + item: DocItem | GroupItem, + level: int, + pages: dict[int, PageDetail], ) -> bool: if isinstance(item, GroupItem): return True @@ -163,9 +168,13 @@ def _process_content_item( if page_no not in pages: logger.warning( "Page %d not found in document metadata — using US Letter fallback (%sx%s pt)", - page_no, _DEFAULT_PAGE_WIDTH, _DEFAULT_PAGE_HEIGHT, + page_no, + _DEFAULT_PAGE_WIDTH, + _DEFAULT_PAGE_HEIGHT, + ) + pages[page_no] = PageDetail( + page_number=page_no, width=_DEFAULT_PAGE_WIDTH, height=_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 @@ -199,6 +208,7 @@ def _process_content_item( # Synchronous conversion (called via asyncio.to_thread) # --------------------------------------------------------------------------- + def _convert_sync(file_path: str, options: ConversionOptions) -> ConversionResult: with _converter_lock: conv = _select_converter(options) @@ -213,7 +223,9 @@ def _convert_sync(file_path: str, options: ConversionOptions) -> ConversionResul 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, + height=doc.pages[i + 1].size.height + if (i + 1) in doc.pages + else _DEFAULT_PAGE_HEIGHT, ) for i in range(page_count) ] @@ -227,6 +239,7 @@ def _convert_sync(file_path: str, options: ConversionOptions) -> ConversionResul content_html=doc.export_to_html(), pages=pages_detail, skipped_items=skipped, + document_json=json.dumps(doc.export_to_dict()), ) @@ -234,10 +247,13 @@ def _convert_sync(file_path: str, options: ConversionOptions) -> ConversionResul # Public adapter class # --------------------------------------------------------------------------- + class LocalConverter: """Adapter that runs Docling locally as a Python library.""" async def convert( - self, file_path: str, options: ConversionOptions, + self, + file_path: str, + options: ConversionOptions, ) -> ConversionResult: return await asyncio.to_thread(_convert_sync, file_path, options)