mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-27 22:21:04 +02:00
feat: add Docling-based document decoder as alternative to unstructured (#1054)
Add trustgraph-docling package providing a document decoder powered by IBM's Docling library. Supports PDF, DOCX, XLSX, PPTX, HTML, Markdown, and CSV with two chunking modes: - page mode (default): groups output by page for page-based formats, emits TextDocument messages for the existing downstream chunker - hybrid mode: uses Docling's HybridChunker for structure-aware semantic chunking, emits Chunk messages directly Tables are preserved as HTML markup. Images are stored in the librarian with provenance but excluded from the text pipeline. Full provenance chain is maintained for explainability tracing. Includes Containerfile, CI/CD integration (PR checks and release workflows), Makefile targets, and tech spec documentation. Added unit tests for the package.
This commit is contained in:
parent
6d51773e00
commit
bdfdbb3c1f
11 changed files with 1733 additions and 1 deletions
3
.github/workflows/pull-request.yaml
vendored
3
.github/workflows/pull-request.yaml
vendored
|
|
@ -42,6 +42,9 @@ jobs:
|
||||||
- name: Install trustgraph-unstructured
|
- name: Install trustgraph-unstructured
|
||||||
run: (cd trustgraph-unstructured; pip install .)
|
run: (cd trustgraph-unstructured; pip install .)
|
||||||
|
|
||||||
|
- name: Install trustgraph-docling
|
||||||
|
run: (cd trustgraph-docling; pip install .)
|
||||||
|
|
||||||
- name: Install trustgraph-vertexai
|
- name: Install trustgraph-vertexai
|
||||||
run: (cd trustgraph-vertexai; pip install .)
|
run: (cd trustgraph-vertexai; pip install .)
|
||||||
|
|
||||||
|
|
|
||||||
2
.github/workflows/release.yaml
vendored
2
.github/workflows/release.yaml
vendored
|
|
@ -58,6 +58,7 @@ jobs:
|
||||||
- hf
|
- hf
|
||||||
- ocr
|
- ocr
|
||||||
- unstructured
|
- unstructured
|
||||||
|
- docling
|
||||||
- mcp
|
- mcp
|
||||||
platform:
|
platform:
|
||||||
- amd64
|
- amd64
|
||||||
|
|
@ -118,6 +119,7 @@ jobs:
|
||||||
- hf
|
- hf
|
||||||
- ocr
|
- ocr
|
||||||
- unstructured
|
- unstructured
|
||||||
|
- docling
|
||||||
- mcp
|
- mcp
|
||||||
|
|
||||||
steps:
|
steps:
|
||||||
|
|
|
||||||
1
.gitignore
vendored
1
.gitignore
vendored
|
|
@ -15,6 +15,7 @@ trustgraph-parquet/trustgraph/parquet_version.py
|
||||||
trustgraph-vertexai/trustgraph/vertexai_version.py
|
trustgraph-vertexai/trustgraph/vertexai_version.py
|
||||||
trustgraph-unstructured/trustgraph/unstructured_version.py
|
trustgraph-unstructured/trustgraph/unstructured_version.py
|
||||||
trustgraph-mcp/trustgraph/mcp_version.py
|
trustgraph-mcp/trustgraph/mcp_version.py
|
||||||
|
trustgraph-docling/trustgraph/docling_version.py
|
||||||
trustgraph/trustgraph/trustgraph_version.py
|
trustgraph/trustgraph/trustgraph_version.py
|
||||||
vertexai/
|
vertexai/
|
||||||
venv/
|
venv/
|
||||||
|
|
|
||||||
6
Makefile
6
Makefile
|
|
@ -18,6 +18,7 @@ wheels:
|
||||||
pip3 wheel --no-deps --wheel-dir dist trustgraph-cli/
|
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-ocr/
|
||||||
pip3 wheel --no-deps --wheel-dir dist trustgraph-unstructured/
|
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/
|
pip3 wheel --no-deps --wheel-dir dist trustgraph-mcp/
|
||||||
|
|
||||||
packages: update-package-versions
|
packages: update-package-versions
|
||||||
|
|
@ -31,6 +32,7 @@ packages: update-package-versions
|
||||||
cd trustgraph-cli && python -m build --sdist --outdir ../dist/
|
cd trustgraph-cli && python -m build --sdist --outdir ../dist/
|
||||||
cd trustgraph-ocr && 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-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/
|
cd trustgraph-mcp && python -m build --sdist --outdir ../dist/
|
||||||
|
|
||||||
pypi-upload:
|
pypi-upload:
|
||||||
|
|
@ -49,13 +51,14 @@ update-package-versions:
|
||||||
echo __version__ = \"${VERSION}\" > trustgraph-cli/trustgraph/cli_version.py
|
echo __version__ = \"${VERSION}\" > trustgraph-cli/trustgraph/cli_version.py
|
||||||
echo __version__ = \"${VERSION}\" > trustgraph-ocr/trustgraph/ocr_version.py
|
echo __version__ = \"${VERSION}\" > trustgraph-ocr/trustgraph/ocr_version.py
|
||||||
echo __version__ = \"${VERSION}\" > trustgraph-unstructured/trustgraph/unstructured_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/trustgraph/trustgraph_version.py
|
||||||
echo __version__ = \"${VERSION}\" > trustgraph-mcp/trustgraph/mcp_version.py
|
echo __version__ = \"${VERSION}\" > trustgraph-mcp/trustgraph/mcp_version.py
|
||||||
|
|
||||||
containers: container-base container-flow \
|
containers: container-base container-flow \
|
||||||
container-bedrock container-vertexai \
|
container-bedrock container-vertexai \
|
||||||
container-hf container-ocr \
|
container-hf container-ocr \
|
||||||
container-unstructured container-mcp
|
container-unstructured container-docling container-mcp
|
||||||
|
|
||||||
some-containers: container-base container-flow
|
some-containers: container-base container-flow
|
||||||
# container-unstructured
|
# container-unstructured
|
||||||
|
|
@ -68,6 +71,7 @@ push:
|
||||||
${DOCKER} push ${CONTAINER_BASE}/trustgraph-hf:${VERSION}
|
${DOCKER} push ${CONTAINER_BASE}/trustgraph-hf:${VERSION}
|
||||||
${DOCKER} push ${CONTAINER_BASE}/trustgraph-ocr:${VERSION}
|
${DOCKER} push ${CONTAINER_BASE}/trustgraph-ocr:${VERSION}
|
||||||
${DOCKER} push ${CONTAINER_BASE}/trustgraph-unstructured:${VERSION}
|
${DOCKER} push ${CONTAINER_BASE}/trustgraph-unstructured:${VERSION}
|
||||||
|
${DOCKER} push ${CONTAINER_BASE}/trustgraph-docling:${VERSION}
|
||||||
${DOCKER} push ${CONTAINER_BASE}/trustgraph-mcp:${VERSION}
|
${DOCKER} push ${CONTAINER_BASE}/trustgraph-mcp:${VERSION}
|
||||||
|
|
||||||
# Individual container build targets
|
# Individual container build targets
|
||||||
|
|
|
||||||
57
containers/Containerfile.docling
Normal file
57
containers/Containerfile.docling
Normal file
|
|
@ -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 /
|
||||||
206
docs/tech-specs/docling-decoder.md
Normal file
206
docs/tech-specs/docling-decoder.md
Normal file
|
|
@ -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.
|
||||||
852
tests/unit/test_decoding/test_docling_processor.py
Normal file
852
tests/unit/test_decoding/test_docling_processor.py
Normal file
|
|
@ -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="<table><tr><td>data</td></tr></table>",
|
||||||
|
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 = "<table><tr><td>A</td></tr></table>"
|
||||||
|
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 "<table>" 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 = "<table><tr><td>cell</td></tr></table>"
|
||||||
|
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 "<table>" 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__])
|
||||||
34
trustgraph-docling/pyproject.toml
Normal file
34
trustgraph-docling/pyproject.toml
Normal file
|
|
@ -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__"}
|
||||||
|
|
@ -0,0 +1,2 @@
|
||||||
|
|
||||||
|
from . processor import *
|
||||||
|
|
@ -0,0 +1,6 @@
|
||||||
|
#!/usr/bin/env python3
|
||||||
|
|
||||||
|
from . processor import run
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
run()
|
||||||
565
trustgraph-docling/trustgraph/decoding/docling/processor.py
Normal file
565
trustgraph-docling/trustgraph/decoding/docling/processor.py
Normal file
|
|
@ -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__)
|
||||||
Loading…
Add table
Add a link
Reference in a new issue