diff --git a/.github/workflows/pull-request.yaml b/.github/workflows/pull-request.yaml index f4bbaaf2..b8f74c09 100644 --- a/.github/workflows/pull-request.yaml +++ b/.github/workflows/pull-request.yaml @@ -42,6 +42,9 @@ jobs: - name: Install trustgraph-unstructured run: (cd trustgraph-unstructured; pip install .) + - name: Install trustgraph-docling + run: (cd trustgraph-docling; pip install .) + - name: Install trustgraph-vertexai run: (cd trustgraph-vertexai; pip install .) diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml index 02c546df..be51e947 100644 --- a/.github/workflows/release.yaml +++ b/.github/workflows/release.yaml @@ -58,6 +58,7 @@ jobs: - hf - ocr - unstructured + - docling - mcp platform: - amd64 @@ -118,6 +119,7 @@ jobs: - hf - ocr - unstructured + - docling - mcp steps: diff --git a/.gitignore b/.gitignore index 366edb4a..354be3e0 100644 --- a/.gitignore +++ b/.gitignore @@ -15,6 +15,7 @@ trustgraph-parquet/trustgraph/parquet_version.py trustgraph-vertexai/trustgraph/vertexai_version.py trustgraph-unstructured/trustgraph/unstructured_version.py trustgraph-mcp/trustgraph/mcp_version.py +trustgraph-docling/trustgraph/docling_version.py trustgraph/trustgraph/trustgraph_version.py vertexai/ venv/ diff --git a/Makefile b/Makefile index 6c5f5ce5..e14c85ad 100644 --- a/Makefile +++ b/Makefile @@ -18,6 +18,7 @@ wheels: pip3 wheel --no-deps --wheel-dir dist trustgraph-cli/ pip3 wheel --no-deps --wheel-dir dist trustgraph-ocr/ pip3 wheel --no-deps --wheel-dir dist trustgraph-unstructured/ + pip3 wheel --no-deps --wheel-dir dist trustgraph-docling/ pip3 wheel --no-deps --wheel-dir dist trustgraph-mcp/ packages: update-package-versions @@ -31,6 +32,7 @@ packages: update-package-versions cd trustgraph-cli && python -m build --sdist --outdir ../dist/ cd trustgraph-ocr && python -m build --sdist --outdir ../dist/ cd trustgraph-unstructured && python -m build --sdist --outdir ../dist/ + cd trustgraph-docling && python -m build --sdist --outdir ../dist/ cd trustgraph-mcp && python -m build --sdist --outdir ../dist/ pypi-upload: @@ -49,13 +51,14 @@ update-package-versions: echo __version__ = \"${VERSION}\" > trustgraph-cli/trustgraph/cli_version.py echo __version__ = \"${VERSION}\" > trustgraph-ocr/trustgraph/ocr_version.py echo __version__ = \"${VERSION}\" > trustgraph-unstructured/trustgraph/unstructured_version.py + echo __version__ = \"${VERSION}\" > trustgraph-docling/trustgraph/docling_version.py echo __version__ = \"${VERSION}\" > trustgraph/trustgraph/trustgraph_version.py echo __version__ = \"${VERSION}\" > trustgraph-mcp/trustgraph/mcp_version.py containers: container-base container-flow \ container-bedrock container-vertexai \ container-hf container-ocr \ -container-unstructured container-mcp +container-unstructured container-docling container-mcp some-containers: container-base container-flow # container-unstructured @@ -68,6 +71,7 @@ push: ${DOCKER} push ${CONTAINER_BASE}/trustgraph-hf:${VERSION} ${DOCKER} push ${CONTAINER_BASE}/trustgraph-ocr:${VERSION} ${DOCKER} push ${CONTAINER_BASE}/trustgraph-unstructured:${VERSION} + ${DOCKER} push ${CONTAINER_BASE}/trustgraph-docling:${VERSION} ${DOCKER} push ${CONTAINER_BASE}/trustgraph-mcp:${VERSION} # Individual container build targets diff --git a/containers/Containerfile.docling b/containers/Containerfile.docling new file mode 100644 index 00000000..54e5163f --- /dev/null +++ b/containers/Containerfile.docling @@ -0,0 +1,57 @@ + +# ---------------------------------------------------------------------------- +# Base container with system dependencies +# ---------------------------------------------------------------------------- + +FROM docker.io/fedora:42 AS base + +ENV PIP_BREAK_SYSTEM_PACKAGES=1 + +# This is cached, same as other containers +RUN dnf install -y python3.11 && \ + alternatives --install /usr/bin/python python /usr/bin/python3.11 1 && \ + python -m ensurepip --upgrade && \ + pip3 install --no-cache-dir --upgrade 'pip>=26.0' 'setuptools>=78.1.1' && \ + pip3 install --no-cache-dir build wheel aiohttp && \ + pip3 install --no-cache-dir pulsar-client==3.11.0 && \ + dnf clean all + +# Installs local to me +RUN pip install --no-cache-dir \ + torch torchvision --index-url https://download.pytorch.org/whl/cpu + +RUN dnf install -y libxcb mesa-libGL && \ + dnf clean all + +# ---------------------------------------------------------------------------- +# Build a container which contains the built Python packages. The build +# creates a bunch of left-over cruft, a separate phase means this is only +# needed to support package build +# ---------------------------------------------------------------------------- + +FROM base AS build + +COPY trustgraph-base/ /root/build/trustgraph-base/ +COPY trustgraph-docling/ /root/build/trustgraph-docling/ + +WORKDIR /root/build/ + +RUN pip3 wheel -w /root/wheels/ --no-deps ./trustgraph-base/ +RUN pip3 wheel -w /root/wheels/ --no-deps ./trustgraph-docling/ + +RUN ls /root/wheels + +# ---------------------------------------------------------------------------- +# Finally, the target container. Start with base and add the package. +# ---------------------------------------------------------------------------- + +FROM base + +COPY --from=build /root/wheels /root/wheels + +RUN \ + pip3 install --no-cache-dir /root/wheels/trustgraph_base-* && \ + pip3 install --no-cache-dir /root/wheels/trustgraph_docling-* && \ + rm -rf /root/wheels + +WORKDIR / diff --git a/docs/tech-specs/docling-decoder.md b/docs/tech-specs/docling-decoder.md new file mode 100644 index 00000000..489c6466 --- /dev/null +++ b/docs/tech-specs/docling-decoder.md @@ -0,0 +1,206 @@ +--- +layout: default +title: "Docling Document Decoder" +parent: "Tech Specs" +--- + +# Docling Document Decoder + +## Problem Statement + +The current universal document decoder is built on the `unstructured` +library. While functional, it has significant operational problems: + +- **Heavyweight.** `unstructured` pulls in PyTorch and a large + dependency tree. Container images are bloated (~2 GB+) and slow + to build. + +- **Fragile processing strategies.** `unstructured`'s internal + strategies (e.g. `is_pdf_too_complex`) change between releases in + ways that alter output for the same input document. This makes + decoder behaviour unpredictable across upgrades. + +- **Flat element model.** `unstructured` returns a flat list of + elements with no structural hierarchy. The decoder must manually + reconstruct page/section boundaries using custom grouping + strategies (`strategies.py`), which is brittle and loses the + document's semantic structure. + +- **Temp file dance.** `unstructured` requires files on disk, + forcing the decoder to write blobs to temp files, parse, then + clean up. + +- **Table reconstruction.** Table extraction relies on a hidden + `text_as_html` metadata attribute that is inconsistently + populated, requiring fallback logic. + +## Approach + +Replace `unstructured` with [Docling](https://github.com/DS4SD/docling), +a document conversion library from IBM Research that provides: + +- A unified `DoclingDocument` schema across all input formats +- Preserved semantic hierarchy (headings, sections, lists, tables) +- Native table reconstruction to Markdown or structured output +- In-memory stream processing (no temp files) +- Built-in `HybridChunker` for semantic chunking +- Lighter dependency footprint + +## Requirements + +The decoder is a critical part of the provenance chain. Every +piece of text that enters the knowledge graph must be traceable +back to its source document, and the decoder is where that chain +begins. Docling must support the following: + +### Source provenance + +Each output (page, section, or chunk) must carry enough metadata +to trace it back to the source document: + +- **Parent document ID.** Every emitted page/section/chunk links + to its source document via `prov:wasDerivedFrom`. The + explainability UI walks this chain to show users exactly which + document produced a given piece of knowledge. + +- **Page numbers.** For page-based formats (PDF, PPTX), each + output must record which page it came from. Docling must + expose page numbers in its document model so the decoder can + propagate them into provenance triples. + +- **Character offsets and lengths.** The downstream chunker + records character offset and length for each chunk relative + to its parent section/page. In hybrid mode, the decoder + must provide equivalent positional metadata so chunks can be + mapped back to their location within the source text. + +- **Structural context.** Docling's hierarchical model + (heading path, section nesting) should be preserved where + possible. This supports future explainability features + that show users not just which page, but which section of + the document a fact came from. + +### Content fidelity + +- **Tables as markup.** Tables must be preserved as structured + content (HTML or Markdown), not flattened to plain text. + Downstream extractors rely on table structure to correctly + identify relationships between cells. + +- **Image separation.** Images must be extractable and stored + separately in the librarian with their own provenance, but + must not be injected into the text pipeline. The provenance + chain links each image to its source page/document. + +- **Text ordering.** Elements must be emitted in reading order. + Docling's layout analysis handles multi-column and complex + layouts — the decoder must preserve that ordering. + +### Operational + +- **Format coverage.** Must handle at minimum: PDF, DOCX, XLSX, + PPTX, HTML, Markdown, plain text. These cover the majority + of enterprise document ingestion. + +- **In-memory processing.** No temp files. The decoder runs in + a container processing documents from Pulsar — temp file + management adds failure modes and cleanup complexity. + +- **Deterministic output.** The same document must produce the + same output across library upgrades. This is a key + motivation for moving away from unstructured, whose + processing strategies change between releases. + +## Architecture + +### Separate package + +The Docling decoder ships as `trustgraph-docling`, a new package +alongside the existing `trustgraph-unstructured`. Users choose +which decoder to deploy — both produce the same downstream +interface (`TextDocument` or `Chunk` messages over Pulsar). + +Package structure mirrors `trustgraph-unstructured`: + +``` +trustgraph-docling/ + pyproject.toml + trustgraph/ + docling_version.py + decoding/ + docling/ + __init__.py + __main__.py + processor.py +``` + +### Two chunking modes + +The decoder supports two modes via `--chunking-mode`: + +**`page` (default)** — drop-in replacement for the unstructured +decoder. Docling converts the document, the decoder groups +output by page (for page-based formats) or emits as a single +section (for non-page formats). Emits `TextDocument` messages +for the existing downstream chunker to split further. + +**`hybrid`** — Docling's `HybridChunker` performs semantic +chunking natively, respecting document hierarchy (headings, +tables, lists). Emits `Chunk` messages directly, bypassing +the downstream chunker. This mode produces higher-quality +chunks because the chunker understands document structure +rather than splitting on character count. + +### Pipeline interface + +Both modes use the same `FlowProcessor` base class and the same +input schema (`Document`). The output schema depends on mode: + +- `page` mode: `TextDocument` (same as unstructured decoder) +- `hybrid` mode: `Chunk` (same as downstream chunker output) + +Provenance and librarian integration work identically to the +unstructured decoder — each page/section/chunk gets a URI, +provenance triples linking it to the source document, and +content saved to the librarian. + +### Image handling + +Same as the unstructured decoder: images are stored in the +librarian with provenance but are not sent to the text pipeline. + +### Document conversion + +Docling's `DocumentConverter` handles format detection and +conversion. The decoder accepts the document as bytes plus an +optional mime type hint, wraps it in a `DocumentStream`, and +calls `converter.convert()`. No temp files. + +### What we don't need + +- `strategies.py` — Docling's hierarchical document model and + `HybridChunker` replace all five custom grouping strategies. +- Mime-to-extension mapping — Docling uses `InputFormat` enums, + not file extensions. +- Manual page grouping logic — Docling tracks page numbers + natively in its document model. + +## Configuration + +| Flag | Default | Description | +|---|---|---| +| `--chunking-mode` | `page` | `page` or `hybrid` | +| `--languages` | `eng` | OCR language codes (comma-separated) | +| `--chunk-max-tokens` | `512` | Max tokens per chunk, hybrid mode only. Controls chunk sizing internally via tokenizer — output is still plain text | + +## Open Questions + +- **OCR configuration.** Docling has its own OCR pipeline + (EasyOCR, Tesseract). Need to verify language configuration + maps cleanly. +- **Format coverage.** Verify Docling handles all formats + currently supported by the unstructured decoder (CSV, TSV, + RST, RTF, ODT). Some niche formats may need fallback. +- **Performance.** Docling uses AI models for layout analysis. + Need to benchmark against unstructured on representative + documents to confirm comparable or better throughput. diff --git a/tests/unit/test_decoding/test_docling_processor.py b/tests/unit/test_decoding/test_docling_processor.py new file mode 100644 index 00000000..7149bf11 --- /dev/null +++ b/tests/unit/test_decoding/test_docling_processor.py @@ -0,0 +1,852 @@ +""" +Unit tests for trustgraph.decoding.docling.processor +""" + +import pytest +import base64 +from unittest.mock import AsyncMock, MagicMock, patch +from unittest import IsolatedAsyncioTestCase + +from trustgraph.decoding.docling.processor import ( + Processor, MIME_EXTENSIONS, PAGE_BASED_FORMATS, + COMPONENT_NAME, COMPONENT_VERSION, +) +from trustgraph.schema import ( + Document, TextDocument, Chunk, Metadata, Triples, + Triple, Term, IRI, LITERAL, +) +from trustgraph.provenance import TG_PAGE_NUMBER + + +class MockAsyncProcessor: + def __init__(self, **params): + self.config_handlers = [] + self.id = params.get('id', 'test-service') + self.specifications = [] + self.pubsub = MagicMock() + self.taskgroup = params.get('taskgroup', MagicMock()) + + +def make_text_item(text="Some text", page_no=None): + """Create a mock Docling TextItem.""" + item = MagicMock() + type(item).__name__ = "TextItem" + item.text = text + if page_no is not None: + prov_entry = MagicMock() + prov_entry.page_no = page_no + item.prov = [prov_entry] + else: + item.prov = [] + return item + + +def make_table_item(html="
data
", + page_no=None): + """Create a mock Docling TableItem.""" + item = MagicMock() + type(item).__name__ = "TableItem" + item.text = "table text" + item.export_to_html = MagicMock(return_value=html) + if page_no is not None: + prov_entry = MagicMock() + prov_entry.page_no = page_no + item.prov = [prov_entry] + else: + item.prov = [] + return item + + +def make_picture_item(page_no=None): + """Create a mock Docling PictureItem.""" + item = MagicMock() + type(item).__name__ = "PictureItem" + item.text = "" + if page_no is not None: + prov_entry = MagicMock() + prov_entry.page_no = page_no + item.prov = [prov_entry] + else: + item.prov = [] + return item + + +class TestMimeExtensions: + """Test the MIME type to extension mapping.""" + + def test_pdf_extension(self): + assert MIME_EXTENSIONS["application/pdf"] == ".pdf" + + def test_docx_extension(self): + key = "application/vnd.openxmlformats-officedocument.wordprocessingml.document" + assert MIME_EXTENSIONS[key] == ".docx" + + def test_xlsx_extension(self): + key = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" + assert MIME_EXTENSIONS[key] == ".xlsx" + + def test_pptx_extension(self): + key = "application/vnd.openxmlformats-officedocument.presentationml.presentation" + assert MIME_EXTENSIONS[key] == ".pptx" + + def test_html_extension(self): + assert MIME_EXTENSIONS["text/html"] == ".html" + + def test_markdown_extension(self): + assert MIME_EXTENSIONS["text/markdown"] == ".md" + + def test_plain_text_extension(self): + assert MIME_EXTENSIONS["text/plain"] == ".md" + + def test_csv_extension(self): + assert MIME_EXTENSIONS["text/csv"] == ".csv" + + def test_unknown_mime_falls_back(self): + assert MIME_EXTENSIONS.get("application/octet-stream", ".bin") == ".bin" + + +class TestPageBasedFormats: + """Test page-based format detection.""" + + def test_pdf_is_page_based(self): + assert "application/pdf" in PAGE_BASED_FORMATS + + def test_pptx_is_page_based(self): + pptx = "application/vnd.openxmlformats-officedocument.presentationml.presentation" + assert pptx in PAGE_BASED_FORMATS + + def test_html_is_not_page_based(self): + assert "text/html" not in PAGE_BASED_FORMATS + + def test_docx_is_not_page_based(self): + docx = "application/vnd.openxmlformats-officedocument.wordprocessingml.document" + assert docx not in PAGE_BASED_FORMATS + + def test_markdown_is_not_page_based(self): + assert "text/markdown" not in PAGE_BASED_FORMATS + + +class TestGetPageTexts: + """Test page text extraction from DoclingDocument.""" + + @patch('trustgraph.decoding.docling.processor.DocumentConverter') + @patch('trustgraph.base.librarian_client.Consumer') + @patch('trustgraph.base.librarian_client.Producer') + @patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor) + def test_single_page_text(self, mock_prod, mock_cons, mock_conv): + processor = Processor( + id='test', taskgroup=MagicMock(), + ) + + doc = MagicMock() + doc.iterate_items.return_value = [ + (make_text_item("Hello world", page_no=1), 0), + (make_text_item("More text", page_no=1), 0), + ] + + pages = processor.get_page_texts(doc) + + assert len(pages) == 1 + page_no, text, table_count, image_count = pages[0] + assert page_no == 1 + assert "Hello world" in text + assert "More text" in text + assert table_count == 0 + assert image_count == 0 + + @patch('trustgraph.decoding.docling.processor.DocumentConverter') + @patch('trustgraph.base.librarian_client.Consumer') + @patch('trustgraph.base.librarian_client.Producer') + @patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor) + def test_multiple_pages(self, mock_prod, mock_cons, mock_conv): + processor = Processor( + id='test', taskgroup=MagicMock(), + ) + + doc = MagicMock() + doc.iterate_items.return_value = [ + (make_text_item("Page 1 text", page_no=1), 0), + (make_text_item("Page 2 text", page_no=2), 0), + (make_text_item("Page 3 text", page_no=3), 0), + ] + + pages = processor.get_page_texts(doc) + + assert len(pages) == 3 + assert pages[0][0] == 1 + assert pages[1][0] == 2 + assert pages[2][0] == 3 + + @patch('trustgraph.decoding.docling.processor.DocumentConverter') + @patch('trustgraph.base.librarian_client.Consumer') + @patch('trustgraph.base.librarian_client.Producer') + @patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor) + def test_table_counted_and_html_included(self, mock_prod, mock_cons, + mock_conv): + processor = Processor( + id='test', taskgroup=MagicMock(), + ) + + html = "
A
" + doc = MagicMock() + doc.iterate_items.return_value = [ + (make_text_item("Intro", page_no=1), 0), + (make_table_item(html=html, page_no=1), 0), + ] + + pages = processor.get_page_texts(doc) + + assert len(pages) == 1 + _, text, table_count, image_count = pages[0] + assert "" in text + assert "Intro" in text + assert table_count == 1 + assert image_count == 0 + + @patch('trustgraph.decoding.docling.processor.DocumentConverter') + @patch('trustgraph.base.librarian_client.Consumer') + @patch('trustgraph.base.librarian_client.Producer') + @patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor) + def test_picture_counted_not_in_text(self, mock_prod, mock_cons, + mock_conv): + processor = Processor( + id='test', taskgroup=MagicMock(), + ) + + doc = MagicMock() + doc.iterate_items.return_value = [ + (make_text_item("Text content", page_no=1), 0), + (make_picture_item(page_no=1), 0), + ] + + pages = processor.get_page_texts(doc) + + assert len(pages) == 1 + _, text, table_count, image_count = pages[0] + assert "Text content" in text + assert image_count == 1 + assert table_count == 0 + + @patch('trustgraph.decoding.docling.processor.DocumentConverter') + @patch('trustgraph.base.librarian_client.Consumer') + @patch('trustgraph.base.librarian_client.Producer') + @patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor) + def test_no_prov_defaults_to_page_1(self, mock_prod, mock_cons, mock_conv): + processor = Processor( + id='test', taskgroup=MagicMock(), + ) + + doc = MagicMock() + doc.iterate_items.return_value = [ + (make_text_item("No page info"), 0), + ] + + pages = processor.get_page_texts(doc) + + assert len(pages) == 1 + assert pages[0][0] == 1 + + @patch('trustgraph.decoding.docling.processor.DocumentConverter') + @patch('trustgraph.base.librarian_client.Consumer') + @patch('trustgraph.base.librarian_client.Producer') + @patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor) + def test_empty_text_pages_skipped(self, mock_prod, mock_cons, mock_conv): + processor = Processor( + id='test', taskgroup=MagicMock(), + ) + + item = make_text_item(" ", page_no=1) + doc = MagicMock() + doc.iterate_items.return_value = [(item, 0)] + + pages = processor.get_page_texts(doc) + + assert len(pages) == 0 + + +class TestGetFullText: + """Test full document text extraction.""" + + @patch('trustgraph.decoding.docling.processor.DocumentConverter') + @patch('trustgraph.base.librarian_client.Consumer') + @patch('trustgraph.base.librarian_client.Producer') + @patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor) + def test_concatenates_text_items(self, mock_prod, mock_cons, mock_conv): + processor = Processor( + id='test', taskgroup=MagicMock(), + ) + + doc = MagicMock() + doc.iterate_items.return_value = [ + (make_text_item("First paragraph"), 0), + (make_text_item("Second paragraph"), 0), + ] + + text, table_count, image_count = processor.get_full_text(doc) + + assert "First paragraph" in text + assert "Second paragraph" in text + assert "\n\n" in text + assert table_count == 0 + assert image_count == 0 + + @patch('trustgraph.decoding.docling.processor.DocumentConverter') + @patch('trustgraph.base.librarian_client.Consumer') + @patch('trustgraph.base.librarian_client.Producer') + @patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor) + def test_tables_as_html(self, mock_prod, mock_cons, mock_conv): + processor = Processor( + id='test', taskgroup=MagicMock(), + ) + + html = "
cell
" + doc = MagicMock() + doc.iterate_items.return_value = [ + (make_text_item("Before"), 0), + (make_table_item(html=html), 0), + (make_text_item("After"), 0), + ] + + text, table_count, image_count = processor.get_full_text(doc) + + assert "" in text + assert "Before" in text + assert "After" in text + assert table_count == 1 + assert image_count == 0 + + @patch('trustgraph.decoding.docling.processor.DocumentConverter') + @patch('trustgraph.base.librarian_client.Consumer') + @patch('trustgraph.base.librarian_client.Producer') + @patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor) + def test_pictures_counted_not_included(self, mock_prod, mock_cons, + mock_conv): + processor = Processor( + id='test', taskgroup=MagicMock(), + ) + + doc = MagicMock() + doc.iterate_items.return_value = [ + (make_text_item("Text"), 0), + (make_picture_item(), 0), + (make_picture_item(), 0), + ] + + text, table_count, image_count = processor.get_full_text(doc) + + assert text == "Text" + assert image_count == 2 + assert table_count == 0 + + +class TestEmitSection(IsolatedAsyncioTestCase): + """Test emit_section output shape and provenance.""" + + @patch('trustgraph.decoding.docling.processor.DocumentConverter') + @patch('trustgraph.base.librarian_client.Consumer') + @patch('trustgraph.base.librarian_client.Producer') + @patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor) + async def test_emit_page_section(self, mock_prod, mock_cons, mock_conv): + processor = Processor( + id='test', taskgroup=MagicMock(), + ) + + mock_output_flow = AsyncMock() + mock_triples_flow = AsyncMock() + mock_flow = MagicMock(side_effect=lambda name: { + "output": mock_output_flow, + "triples": mock_triples_flow, + }.get(name)) + mock_flow.librarian.save_child_document = AsyncMock() + + metadata = Metadata(id="doc-1", root="root-1", collection="default") + + uri = await processor.emit_section( + text="Page content here", + parent_doc_id="parent-doc", + doc_uri_str="urn:doc:parent", + metadata=metadata, + flow=mock_flow, + mime_type="application/pdf", + page_number=3, + ) + + assert uri is not None + assert uri.startswith("urn:page:") + + mock_flow.librarian.save_child_document.assert_called_once() + save_kwargs = mock_flow.librarian.save_child_document.call_args + assert save_kwargs[1]["document_type"] == "page" + assert save_kwargs[1]["title"] == "Page 3" + + mock_triples_flow.send.assert_called_once() + triples_msg = mock_triples_flow.send.call_args[0][0] + assert isinstance(triples_msg, Triples) + assert triples_msg.metadata.root == "root-1" + + mock_output_flow.send.assert_called_once() + text_doc = mock_output_flow.send.call_args[0][0] + assert isinstance(text_doc, TextDocument) + assert text_doc.text == b"Page content here" + assert text_doc.metadata.id == uri + + @patch('trustgraph.decoding.docling.processor.DocumentConverter') + @patch('trustgraph.base.librarian_client.Consumer') + @patch('trustgraph.base.librarian_client.Producer') + @patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor) + async def test_emit_non_page_section(self, mock_prod, mock_cons, + mock_conv): + processor = Processor( + id='test', taskgroup=MagicMock(), + ) + + mock_output_flow = AsyncMock() + mock_triples_flow = AsyncMock() + mock_flow = MagicMock(side_effect=lambda name: { + "output": mock_output_flow, + "triples": mock_triples_flow, + }.get(name)) + mock_flow.librarian.save_child_document = AsyncMock() + + metadata = Metadata(id="doc-1", root="root-1", collection="default") + + uri = await processor.emit_section( + text="Section content", + parent_doc_id="parent-doc", + doc_uri_str="urn:doc:parent", + metadata=metadata, + flow=mock_flow, + section_index=1, + ) + + assert uri is not None + assert uri.startswith("urn:section:") + + save_kwargs = mock_flow.librarian.save_child_document.call_args + assert save_kwargs[1]["document_type"] == "section" + + @patch('trustgraph.decoding.docling.processor.DocumentConverter') + @patch('trustgraph.base.librarian_client.Consumer') + @patch('trustgraph.base.librarian_client.Producer') + @patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor) + async def test_emit_section_skips_blank_text(self, mock_prod, mock_cons, + mock_conv): + processor = Processor( + id='test', taskgroup=MagicMock(), + ) + + mock_flow = MagicMock() + metadata = Metadata(id="doc-1", collection="default") + + uri = await processor.emit_section( + text=" ", + parent_doc_id="parent-doc", + doc_uri_str="urn:doc:parent", + metadata=metadata, + flow=mock_flow, + ) + + assert uri is None + + +class TestEmitChunk(IsolatedAsyncioTestCase): + """Test emit_chunk output shape and multi-page provenance.""" + + @patch('trustgraph.decoding.docling.processor.DocumentConverter') + @patch('trustgraph.base.librarian_client.Consumer') + @patch('trustgraph.base.librarian_client.Producer') + @patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor) + async def test_emit_chunk_single_page(self, mock_prod, mock_cons, + mock_conv): + processor = Processor( + id='test', taskgroup=MagicMock(), + chunking_mode='hybrid', + ) + + mock_chunks_flow = AsyncMock() + mock_triples_flow = AsyncMock() + mock_flow = MagicMock(side_effect=lambda name: { + "chunks": mock_chunks_flow, + "triples": mock_triples_flow, + }.get(name)) + mock_flow.librarian.save_child_document = AsyncMock() + + metadata = Metadata(id="doc-1", root="root-1", collection="default") + + uri = await processor.emit_chunk( + text="Chunk text content", + chunk_index=1, + parent_doc_id="parent-doc", + doc_uri_str="urn:doc:parent", + metadata=metadata, + flow=mock_flow, + page_numbers=[5], + ) + + assert uri is not None + assert uri.startswith("urn:chunk:") + + mock_chunks_flow.send.assert_called_once() + chunk_msg = mock_chunks_flow.send.call_args[0][0] + assert isinstance(chunk_msg, Chunk) + assert chunk_msg.chunk == b"Chunk text content" + + mock_triples_flow.send.assert_called_once() + triples_msg = mock_triples_flow.send.call_args[0][0] + assert isinstance(triples_msg, Triples) + + @patch('trustgraph.decoding.docling.processor.DocumentConverter') + @patch('trustgraph.base.librarian_client.Consumer') + @patch('trustgraph.base.librarian_client.Producer') + @patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor) + async def test_emit_chunk_multi_page_provenance(self, mock_prod, mock_cons, + mock_conv): + """Chunks spanning multiple pages emit extra TG_PAGE_NUMBER triples.""" + processor = Processor( + id='test', taskgroup=MagicMock(), + chunking_mode='hybrid', + ) + + mock_chunks_flow = AsyncMock() + mock_triples_flow = AsyncMock() + mock_flow = MagicMock(side_effect=lambda name: { + "chunks": mock_chunks_flow, + "triples": mock_triples_flow, + }.get(name)) + mock_flow.librarian.save_child_document = AsyncMock() + + metadata = Metadata(id="doc-1", root="root-1", collection="default") + + uri = await processor.emit_chunk( + text="Chunk spanning pages", + chunk_index=2, + parent_doc_id="parent-doc", + doc_uri_str="urn:doc:parent", + metadata=metadata, + flow=mock_flow, + page_numbers=[3, 4, 5], + ) + + assert uri is not None + + triples_msg = mock_triples_flow.send.call_args[0][0] + all_triples = triples_msg.triples + + page_number_triples = [ + t for t in all_triples + if t.p.iri == TG_PAGE_NUMBER + ] + + extra_pages = {t.o.value for t in page_number_triples} + assert "4" in extra_pages + assert "5" in extra_pages + + @patch('trustgraph.decoding.docling.processor.DocumentConverter') + @patch('trustgraph.base.librarian_client.Consumer') + @patch('trustgraph.base.librarian_client.Producer') + @patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor) + async def test_emit_chunk_skips_blank(self, mock_prod, mock_cons, + mock_conv): + processor = Processor( + id='test', taskgroup=MagicMock(), + chunking_mode='hybrid', + ) + + mock_flow = MagicMock() + metadata = Metadata(id="doc-1", collection="default") + + uri = await processor.emit_chunk( + text=" \n ", + chunk_index=1, + parent_doc_id="parent-doc", + doc_uri_str="urn:doc:parent", + metadata=metadata, + flow=mock_flow, + ) + + assert uri is None + + +class TestProcessPageMode(IsolatedAsyncioTestCase): + """Test page mode processing.""" + + @patch('trustgraph.decoding.docling.processor.DocumentConverter') + @patch('trustgraph.base.librarian_client.Consumer') + @patch('trustgraph.base.librarian_client.Producer') + @patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor) + async def test_page_based_format_emits_per_page(self, mock_prod, mock_cons, + mock_conv): + processor = Processor( + id='test', taskgroup=MagicMock(), + ) + + doc = MagicMock() + doc.iterate_items.return_value = [ + (make_text_item("Page 1", page_no=1), 0), + (make_text_item("Page 2", page_no=2), 0), + ] + + mock_output_flow = AsyncMock() + mock_triples_flow = AsyncMock() + mock_flow = MagicMock(side_effect=lambda name: { + "output": mock_output_flow, + "triples": mock_triples_flow, + }.get(name)) + mock_flow.librarian.save_child_document = AsyncMock() + + metadata = Metadata(id="doc-1", root="root-1", collection="default") + + await processor.process_page_mode( + doc, "parent-doc", "urn:doc:parent", + metadata, mock_flow, "application/pdf", + ) + + assert mock_output_flow.send.call_count == 2 + assert mock_triples_flow.send.call_count == 2 + + @patch('trustgraph.decoding.docling.processor.DocumentConverter') + @patch('trustgraph.base.librarian_client.Consumer') + @patch('trustgraph.base.librarian_client.Producer') + @patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor) + async def test_non_page_format_emits_single_section(self, mock_prod, + mock_cons, mock_conv): + processor = Processor( + id='test', taskgroup=MagicMock(), + ) + + doc = MagicMock() + doc.iterate_items.return_value = [ + (make_text_item("Full document text"), 0), + ] + + mock_output_flow = AsyncMock() + mock_triples_flow = AsyncMock() + mock_flow = MagicMock(side_effect=lambda name: { + "output": mock_output_flow, + "triples": mock_triples_flow, + }.get(name)) + mock_flow.librarian.save_child_document = AsyncMock() + + metadata = Metadata(id="doc-1", root="root-1", collection="default") + + await processor.process_page_mode( + doc, "parent-doc", "urn:doc:parent", + metadata, mock_flow, "text/html", + ) + + assert mock_output_flow.send.call_count == 1 + + text_doc = mock_output_flow.send.call_args[0][0] + assert isinstance(text_doc, TextDocument) + assert text_doc.document_id.startswith("urn:section:") + + +class TestProcessorInitialization(IsolatedAsyncioTestCase): + """Test processor initialization and configuration.""" + + @patch('trustgraph.decoding.docling.processor.DocumentConverter') + @patch('trustgraph.base.librarian_client.Consumer') + @patch('trustgraph.base.librarian_client.Producer') + @patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor) + async def test_default_page_mode(self, mock_prod, mock_cons, mock_conv): + processor = Processor( + id='test', taskgroup=MagicMock(), + ) + + assert processor.chunking_mode == "page" + assert processor.chunk_max_tokens == 512 + + consumer_specs = [ + s for s in processor.specifications if hasattr(s, 'handler') + ] + assert len(consumer_specs) >= 1 + assert consumer_specs[0].name == "input" + assert consumer_specs[0].schema == Document + + @patch('trustgraph.decoding.docling.processor.DocumentConverter') + @patch('trustgraph.base.librarian_client.Consumer') + @patch('trustgraph.base.librarian_client.Producer') + @patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor) + async def test_hybrid_mode_registers_chunks_producer(self, mock_prod, + mock_cons, mock_conv): + processor = Processor( + id='test', taskgroup=MagicMock(), + chunking_mode='hybrid', + ) + + assert processor.chunking_mode == "hybrid" + + producer_specs = [ + s for s in processor.specifications + if not hasattr(s, 'handler') and hasattr(s, 'schema') + ] + producer_names = [s.name for s in producer_specs] + assert "chunks" in producer_names + assert "output" in producer_names + assert "triples" in producer_names + + @patch('trustgraph.decoding.docling.processor.DocumentConverter') + @patch('trustgraph.base.librarian_client.Consumer') + @patch('trustgraph.base.librarian_client.Producer') + @patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor) + async def test_custom_languages(self, mock_prod, mock_cons, mock_conv): + processor = Processor( + id='test', taskgroup=MagicMock(), + languages='eng,deu,fra', + ) + + assert processor.languages == ['eng', 'deu', 'fra'] + + +class TestAddArgs: + """Test CLI argument registration.""" + + @patch('trustgraph.base.flow_processor.FlowProcessor.add_args') + def test_add_args_registers_expected_args(self, mock_parent): + mock_parser = MagicMock() + + Processor.add_args(mock_parser) + + mock_parent.assert_called_once_with(mock_parser) + + arg_names = [c[0] for c in mock_parser.add_argument.call_args_list] + assert ('--chunking-mode',) in arg_names + assert ('--languages',) in arg_names + assert ('--chunk-max-tokens',) in arg_names + + +class TestRun: + """Test the run entry point.""" + + @patch('trustgraph.decoding.docling.processor.Processor.launch') + def test_run_calls_launch(self, mock_launch): + from trustgraph.decoding.docling.processor import run + run() + + mock_launch.assert_called_once() + args = mock_launch.call_args[0] + assert args[0] == "document-decoder" + + +class TestOnMessage(IsolatedAsyncioTestCase): + """Test the on_message entry point.""" + + @patch('trustgraph.decoding.docling.processor.DocumentConverter') + @patch('trustgraph.base.librarian_client.Consumer') + @patch('trustgraph.base.librarian_client.Producer') + @patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor) + async def test_inline_document_page_mode(self, mock_prod, mock_cons, + mock_conv): + processor = Processor( + id='test', taskgroup=MagicMock(), + ) + + content = b"fake document bytes" + b64 = base64.b64encode(content).decode('utf-8') + + mock_metadata = Metadata(id="test-doc", root="root", collection="col") + mock_document = Document(metadata=mock_metadata, data=b64) + mock_msg = MagicMock() + mock_msg.value.return_value = mock_document + + mock_doc = MagicMock() + mock_doc.iterate_items.return_value = [ + (make_text_item("Converted text"), 0), + ] + processor.convert_document = MagicMock(return_value=mock_doc) + + mock_output_flow = AsyncMock() + mock_triples_flow = AsyncMock() + mock_flow = MagicMock(side_effect=lambda name: { + "output": mock_output_flow, + "triples": mock_triples_flow, + }.get(name)) + mock_flow.librarian.save_child_document = AsyncMock() + + await processor.on_message(mock_msg, None, mock_flow) + + assert mock_output_flow.send.call_count == 1 + + @patch('trustgraph.decoding.docling.processor.DocumentConverter') + @patch('trustgraph.base.librarian_client.Consumer') + @patch('trustgraph.base.librarian_client.Producer') + @patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor) + async def test_librarian_document_fetch(self, mock_prod, mock_cons, + mock_conv): + processor = Processor( + id='test', taskgroup=MagicMock(), + ) + + content = b"fetched document bytes" + b64_content = base64.b64encode(content) + + mock_metadata = Metadata(id="test-doc", root="root", collection="col") + mock_document = Document( + metadata=mock_metadata, + document_id="lib-doc-123", + ) + mock_msg = MagicMock() + mock_msg.value.return_value = mock_document + + mock_doc = MagicMock() + mock_doc.iterate_items.return_value = [ + (make_text_item("Fetched text"), 0), + ] + processor.convert_document = MagicMock(return_value=mock_doc) + + mock_output_flow = AsyncMock() + mock_triples_flow = AsyncMock() + mock_flow = MagicMock(side_effect=lambda name: { + "output": mock_output_flow, + "triples": mock_triples_flow, + }.get(name)) + mock_flow.librarian.fetch_document_metadata = AsyncMock( + return_value=MagicMock(kind="application/pdf") + ) + mock_flow.librarian.fetch_document_content = AsyncMock( + return_value=b64_content + ) + mock_flow.librarian.save_child_document = AsyncMock() + + await processor.on_message(mock_msg, None, mock_flow) + + mock_flow.librarian.fetch_document_metadata.assert_called_once_with( + document_id="lib-doc-123", + ) + mock_flow.librarian.fetch_document_content.assert_called_once_with( + document_id="lib-doc-123", + ) + + processor.convert_document.assert_called_once() + call_args = processor.convert_document.call_args + assert call_args[0][1] == "application/pdf" + + @patch('trustgraph.decoding.docling.processor.DocumentConverter') + @patch('trustgraph.base.librarian_client.Consumer') + @patch('trustgraph.base.librarian_client.Producer') + @patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor) + async def test_conversion_failure_logged_not_raised(self, mock_prod, + mock_cons, mock_conv): + processor = Processor( + id='test', taskgroup=MagicMock(), + ) + + content = b"bad document" + b64 = base64.b64encode(content).decode('utf-8') + + mock_metadata = Metadata(id="test-doc", root="root", collection="col") + mock_document = Document(metadata=mock_metadata, data=b64) + mock_msg = MagicMock() + mock_msg.value.return_value = mock_document + + processor.convert_document = MagicMock( + side_effect=ValueError("Conversion failed") + ) + + mock_flow = MagicMock() + + await processor.on_message(mock_msg, None, mock_flow) + + +if __name__ == '__main__': + pytest.main([__file__]) diff --git a/trustgraph-docling/pyproject.toml b/trustgraph-docling/pyproject.toml new file mode 100644 index 00000000..17b35c86 --- /dev/null +++ b/trustgraph-docling/pyproject.toml @@ -0,0 +1,34 @@ +[build-system] +requires = ["setuptools>=61.0", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "trustgraph-docling" +dynamic = ["version"] +authors = [{name = "trustgraph.ai", email = "security@trustgraph.ai"}] +description = "TrustGraph document decoder powered by Docling — lightweight alternative to the unstructured-based decoder." +readme = "README.md" +requires-python = ">=3.8" +dependencies = [ + "trustgraph-base>=2.7,<2.8", + "pulsar-client", + "prometheus-client", + "docling", + "onnxruntime", +] +classifiers = [ + "Programming Language :: Python :: 3", + "Operating System :: OS Independent", +] + +[project.urls] +Homepage = "https://github.com/trustgraph-ai/trustgraph" + +[project.scripts] +docling-decoder = "trustgraph.decoding.docling:run" + +[tool.setuptools.packages.find] +include = ["trustgraph*"] + +[tool.setuptools.dynamic] +version = {attr = "trustgraph.docling_version.__version__"} diff --git a/trustgraph-docling/trustgraph/decoding/docling/__init__.py b/trustgraph-docling/trustgraph/decoding/docling/__init__.py new file mode 100644 index 00000000..bd3b0e96 --- /dev/null +++ b/trustgraph-docling/trustgraph/decoding/docling/__init__.py @@ -0,0 +1,2 @@ + +from . processor import * diff --git a/trustgraph-docling/trustgraph/decoding/docling/__main__.py b/trustgraph-docling/trustgraph/decoding/docling/__main__.py new file mode 100644 index 00000000..1ebce4d4 --- /dev/null +++ b/trustgraph-docling/trustgraph/decoding/docling/__main__.py @@ -0,0 +1,6 @@ +#!/usr/bin/env python3 + +from . processor import run + +if __name__ == '__main__': + run() diff --git a/trustgraph-docling/trustgraph/decoding/docling/processor.py b/trustgraph-docling/trustgraph/decoding/docling/processor.py new file mode 100644 index 00000000..9274d3b6 --- /dev/null +++ b/trustgraph-docling/trustgraph/decoding/docling/processor.py @@ -0,0 +1,565 @@ + +""" +Document decoder powered by Docling. + +Accepts documents in common formats (PDF, DOCX, XLSX, HTML, Markdown, +plain text, PPTX, etc.) on input. Two output modes: + +- page mode: emits pages/sections as TextDocument objects for + downstream chunking (drop-in replacement for the unstructured decoder) +- hybrid mode: uses Docling's HybridChunker to emit Chunk objects + directly, bypassing the downstream chunker + +Supports both inline document data and fetching from librarian via +Pulsar for large documents. + +Tables are preserved as HTML markup for better downstream extraction. +Images are stored in the librarian but not sent to the text pipeline. +""" + +import os +os.environ.setdefault("DOCLING_NUM_THREADS", "1") +os.environ.setdefault("OMP_NUM_THREADS", "1") + +import base64 +import io +import logging + +from docling.datamodel.base_models import InputFormat +from docling.datamodel.pipeline_options import PdfPipelineOptions +from docling.document_converter import DocumentConverter, DocumentStream, PdfFormatOption +from docling.backend.pypdfium2_backend import PyPdfiumDocumentBackend +from docling.chunking import HybridChunker + +from ... schema import Document, TextDocument, Chunk, Metadata +from ... schema import Triples, Triple, Term, IRI, LITERAL +from ... base import FlowProcessor, ConsumerSpec, ProducerSpec, LibrarianSpec + +from ... provenance import ( + document_uri, page_uri as make_page_uri, + section_uri as make_section_uri, chunk_uri as make_chunk_uri, + image_uri as make_image_uri, + derived_entity_triples, set_graph, GRAPH_SOURCE, + TG_PAGE_NUMBER, +) + +COMPONENT_NAME = "docling-decoder" +COMPONENT_VERSION = "1.0.0" + +logger = logging.getLogger(__name__) + +default_ident = "document-decoder" + +MIME_EXTENSIONS = { + "application/pdf": ".pdf", + "application/vnd.openxmlformats-officedocument.wordprocessingml.document": ".docx", + "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": ".xlsx", + "application/vnd.openxmlformats-officedocument.presentationml.presentation": ".pptx", + "text/html": ".html", + "text/markdown": ".md", + "text/plain": ".md", + "text/csv": ".csv", +} + +PAGE_BASED_FORMATS = { + "application/pdf", + "application/vnd.openxmlformats-officedocument.presentationml.presentation", +} + + +class Processor(FlowProcessor): + + def __init__(self, **params): + + id = params.get("id", default_ident) + + self.chunking_mode = params.get("chunking_mode", "page") + self.chunk_max_tokens = params.get("chunk_max_tokens", 512) + self.languages = params.get("languages", "eng").split(",") + + pipeline_options = PdfPipelineOptions() + pipeline_options.do_ocr = True + pipeline_options.do_table_structure = True + pipeline_options.generate_page_images = True + + self.converter = DocumentConverter( + allowed_formats=[ + InputFormat.PDF, InputFormat.DOCX, InputFormat.XLSX, + InputFormat.PPTX, InputFormat.HTML, InputFormat.MD, + InputFormat.CSV, + ], + format_options={ + InputFormat.PDF: PdfFormatOption( + pipeline_options=pipeline_options, + backend=PyPdfiumDocumentBackend, + ), + }, + ) + + super(Processor, self).__init__( + **params | { + "id": id, + } + ) + + self.register_specification( + ConsumerSpec( + name="input", + schema=Document, + handler=self.on_message, + ) + ) + + self.register_specification( + ProducerSpec( + name="output", + schema=TextDocument, + ) + ) + + if self.chunking_mode == "hybrid": + self.register_specification( + ProducerSpec( + name="chunks", + schema=Chunk, + ) + ) + + self.register_specification( + ProducerSpec( + name="triples", + schema=Triples, + ) + ) + + self.register_specification( + LibrarianSpec() + ) + + logger.info( + f"Docling decoder initialized (mode: {self.chunking_mode})" + ) + + def convert_document(self, blob, mime_type=None, name_hint="document"): + """ + Convert raw bytes to a DoclingDocument. + + Uses single-document convert with raises_on_error=False + to capture page-level errors without batch panic. + """ + suffix = MIME_EXTENSIONS.get(mime_type, ".bin") + routing_filename = f"source_file{suffix}" + + buf = io.BytesIO(blob) + stream = DocumentStream(name=routing_filename, stream=buf) + + result = self.converter.convert(stream, raises_on_error=True) + + if result.errors: + logger.error( + f"Docling backend failed processing pages for {name_hint}" + ) + for err in result.errors: + logger.error( + f"-> Conversion error: {err!r} " + f"attrs={vars(err) if hasattr(err, '__dict__') else 'N/A'}" + ) + + if result.document: + logger.warning( + "Returning partially recovered layout context." + ) + return result.document + + raise ValueError( + f"Conversion completely failed for {name_hint}. " + f"First error: {repr(result.errors[0])}" + ) + + logger.info(f"Docling conversion complete for {routing_filename}") + return result.document + + def get_page_texts(self, doc): + """ + Extract text grouped by page from a DoclingDocument. + + Returns list of (page_number, text, table_count, image_count). + An item that spans multiple pages is included in each page. + """ + page_map = {} + + for item, _level in doc.iterate_items(): + prov = getattr(item, 'prov', None) + + page_nos = set() + if prov: + for p in prov: + pg = getattr(p, 'page_no', None) + if pg is not None: + page_nos.add(pg) + + if not page_nos: + page_nos.add(1) + + item_type = type(item).__name__ + + for page_no in page_nos: + if page_no not in page_map: + page_map[page_no] = { + "texts": [], + "table_count": 0, + "image_count": 0, + } + + if item_type == "PictureItem": + page_map[page_no]["image_count"] += 1 + continue + + if item_type == "TableItem": + page_map[page_no]["table_count"] += 1 + html = item.export_to_html(doc=doc) + if html: + page_map[page_no]["texts"].append(html) + continue + + text = getattr(item, 'text', None) + if text and text.strip(): + page_map[page_no]["texts"].append(text) + + pages = [] + for page_no in sorted(page_map.keys()): + info = page_map[page_no] + text = "\n\n".join(info["texts"]) + if text.strip(): + pages.append(( + page_no, text, + info["table_count"], info["image_count"], + )) + + return pages + + def get_full_text(self, doc): + """ + Export the full document as markdown text. + + Tables are exported as HTML for structural fidelity. + """ + parts = [] + table_count = 0 + image_count = 0 + + for item, _level in doc.iterate_items(): + item_type = type(item).__name__ + + if item_type == "PictureItem": + image_count += 1 + continue + + if item_type == "TableItem": + table_count += 1 + html = item.export_to_html(doc=doc) + if html: + parts.append(html) + continue + + text = getattr(item, 'text', None) + if text and text.strip(): + parts.append(text) + + return "\n\n".join(parts), table_count, image_count + + async def emit_section(self, text, parent_doc_id, doc_uri_str, + metadata, flow, mime_type=None, + page_number=None, section_index=None, + table_count=0, image_count=0): + """ + Save a page or section to the librarian, emit provenance, + and send a TextDocument downstream. + """ + if not text.strip(): + return None + + is_page = page_number is not None + char_length = len(text) + + if is_page: + entity_uri = make_page_uri() + label = f"Page {page_number}" + else: + entity_uri = make_section_uri() + label = f"Section {section_index}" if section_index else "Section" + + doc_id = entity_uri + content = text.encode("utf-8") + + await flow.librarian.save_child_document( + doc_id=doc_id, + parent_id=parent_doc_id, + content=content, + document_type="page" if is_page else "section", + title=label, + ) + + prov_triples = derived_entity_triples( + entity_uri=entity_uri, + parent_uri=doc_uri_str, + component_name=COMPONENT_NAME, + component_version=COMPONENT_VERSION, + label=label, + page_number=page_number, + section=not is_page, + char_length=char_length, + mime_type=mime_type, + table_count=table_count if table_count > 0 else None, + image_count=image_count if image_count > 0 else None, + ) + + await flow("triples").send(Triples( + metadata=Metadata( + id=entity_uri, + root=metadata.root, + collection=metadata.collection, + ), + triples=set_graph(prov_triples, GRAPH_SOURCE), + )) + + r = TextDocument( + metadata=Metadata( + id=entity_uri, + root=metadata.root, + collection=metadata.collection, + ), + document_id=doc_id, + text=content, + ) + + await flow("output").send(r) + + return entity_uri + + async def emit_chunk(self, text, chunk_index, parent_doc_id, + doc_uri_str, metadata, flow, + page_numbers=None, char_offset=None): + """ + Save a chunk to the librarian, emit provenance, and send + a Chunk message downstream. + """ + if not text.strip(): + return None + + c_uri = make_chunk_uri() + chunk_content = text.encode("utf-8") + char_length = len(text) + + primary_page = page_numbers[0] if page_numbers else None + + await flow.librarian.save_child_document( + doc_id=c_uri, + parent_id=parent_doc_id, + content=chunk_content, + document_type="chunk", + title=f"Chunk {chunk_index}", + ) + + prov_triples = derived_entity_triples( + entity_uri=c_uri, + parent_uri=doc_uri_str, + component_name=COMPONENT_NAME, + component_version=COMPONENT_VERSION, + label=f"Chunk {chunk_index}", + chunk_index=chunk_index, + char_offset=char_offset, + char_length=char_length, + page_number=primary_page, + ) + + if page_numbers and len(page_numbers) > 1: + for extra_page in page_numbers[1:]: + prov_triples.append(Triple( + s=Term(type=IRI, iri=c_uri), + p=Term(type=IRI, iri=TG_PAGE_NUMBER), + o=Term(type=LITERAL, value=str(extra_page)), + )) + + await flow("triples").send(Triples( + metadata=Metadata( + id=c_uri, + root=metadata.root, + collection=metadata.collection, + ), + triples=set_graph(prov_triples, GRAPH_SOURCE), + )) + + r = Chunk( + metadata=Metadata( + id=c_uri, + root=metadata.root, + collection=metadata.collection, + ), + chunk=chunk_content, + document_id=c_uri, + ) + + await flow("chunks").send(r) + + return c_uri + + async def process_page_mode(self, doc, source_doc_id, doc_uri_str, + metadata, flow, mime_type): + """Page mode: emit pages or sections as TextDocument.""" + is_page_based = mime_type in PAGE_BASED_FORMATS if mime_type else False + + logger.info( + f"Page mode: mime={mime_type} is_page_based={is_page_based}" + ) + + if is_page_based: + pages = self.get_page_texts(doc) + logger.info(f"Extracted {len(pages)} pages") + for page_num, text, table_count, image_count in pages: + await self.emit_section( + text, source_doc_id, doc_uri_str, + metadata, flow, + mime_type=mime_type, page_number=page_num, + table_count=table_count, image_count=image_count, + ) + else: + text, table_count, image_count = self.get_full_text(doc) + if text.strip(): + await self.emit_section( + text, source_doc_id, doc_uri_str, + metadata, flow, + mime_type=mime_type, section_index=1, + table_count=table_count, image_count=image_count, + ) + + async def process_hybrid_mode(self, doc, source_doc_id, doc_uri_str, + metadata, flow, mime_type): + """Hybrid mode: use Docling HybridChunker, emit Chunk directly.""" + chunker = HybridChunker( + max_tokens=self.chunk_max_tokens, + ) + + chunks = chunker.chunk(doc) + char_offset = 0 + + for idx, chunk in enumerate(chunks): + chunk_index = idx + 1 + text = chunker.serialize(chunk) + + page_numbers = None + meta = getattr(chunk, 'meta', None) + if meta: + pn = getattr(meta, 'doc_items', None) + if pn: + pages = set() + for di in pn: + for p in getattr(di, 'prov', []): + pg = getattr(p, 'page_no', None) + if pg is not None: + pages.add(pg) + if pages: + page_numbers = sorted(pages) + + await self.emit_chunk( + text, chunk_index, source_doc_id, doc_uri_str, + metadata, flow, + page_numbers=page_numbers, + char_offset=char_offset, + ) + + char_offset += len(text) + + async def on_message(self, msg, consumer, flow): + + logger.debug("Document message received") + + v = msg.value() + + logger.info(f"Decoding {v.metadata.id}...") + + mime_type = None + + if v.document_id: + logger.info( + f"Fetching document {v.document_id} from librarian..." + ) + + doc_meta = await flow.librarian.fetch_document_metadata( + document_id=v.document_id, + ) + mime_type = doc_meta.kind if doc_meta else None + + content = await flow.librarian.fetch_document_content( + document_id=v.document_id, + ) + + if isinstance(content, str): + content = content.encode('utf-8') + blob = base64.b64decode(content) + + logger.info( + f"Fetched {len(blob)} bytes, mime: {mime_type}" + ) + else: + blob = base64.b64decode(v.data) + + source_doc_id = v.document_id or v.metadata.id + doc_uri_str = document_uri(source_doc_id) + + try: + doc = self.convert_document( + blob, mime_type, name_hint=source_doc_id, + ) + except Exception as e: + logger.error( + f"Failed to convert {source_doc_id}: " + f"{type(e).__name__}: {e}", + exc_info=True, + ) + return + + if self.chunking_mode == "hybrid": + await self.process_hybrid_mode( + doc, source_doc_id, doc_uri_str, + v.metadata, flow, mime_type, + ) + else: + await self.process_page_mode( + doc, source_doc_id, doc_uri_str, + v.metadata, flow, mime_type, + ) + + logger.info("Document decoding complete") + + @staticmethod + def add_args(parser): + + FlowProcessor.add_args(parser) + + parser.add_argument( + '--chunking-mode', + default='page', + choices=['page', 'hybrid'], + help='Chunking mode: page emits TextDocument for downstream ' + 'chunker, hybrid uses Docling HybridChunker directly ' + '(default: page)', + ) + + parser.add_argument( + '--languages', + default='eng', + help='Comma-separated OCR language codes (default: eng)', + ) + + parser.add_argument( + '--chunk-max-tokens', + type=int, + default=512, + help='Max tokens per chunk for hybrid mode (default: 512)', + ) + + +def run(): + + Processor.launch(default_ident, __doc__)