mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-27 22:21:04 +02:00
bdfdbb3c1f
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.
852 lines
30 KiB
Python
852 lines
30 KiB
Python
"""
|
|
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__])
|