mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-08 21:02:12 +02:00
24f0190ce7
Adds a RabbitMQ backend as an alternative to Pulsar, selectable via PUBSUB_BACKEND=rabbitmq. Both backends implement the same PubSubBackend protocol — no application code changes needed to switch. RabbitMQ topology: - Single topic exchange per topicspace (e.g. 'tg') - Routing key derived from queue class and topic name - Shared consumers: named queue bound to exchange (competing, round-robin) - Exclusive consumers: anonymous auto-delete queue (broadcast, each gets every message). Used by Subscriber and config push consumer. - Thread-local producer connections (pika is not thread-safe) - Push-based consumption via basic_consume with process_data_events for heartbeat processing Consumer model changes: - Consumer class creates one backend consumer per concurrent task (required for pika thread safety, harmless for Pulsar) - Consumer class accepts consumer_type parameter - Subscriber passes consumer_type='exclusive' for broadcast semantics - Config push consumer uses consumer_type='exclusive' so every processor instance receives config updates - handle_one_from_queue receives consumer as parameter for correct per-connection ack/nack LibrarianClient: - New shared client class replacing duplicated librarian request-response code across 6+ services (chunking, decoders, RAG, etc.) - Uses stream-document instead of get-document-content for fetching document content in 1MB chunks (avoids broker message size limits) - Standalone object (self.librarian = LibrarianClient(...)) not a mixin - get-document-content marked deprecated in schema and OpenAPI spec Serialisation: - Extracted dataclass_to_dict/dict_to_dataclass to shared serialization.py (used by both Pulsar and RabbitMQ backends) Librarian queues: - Changed from flow class (persistent) back to request/response class now that stream-document eliminates large single messages - API upload chunk size reduced from 5MB to 3MB to stay under broker limits after base64 encoding Factory and CLI: - get_pubsub() handles 'rabbitmq' backend with RabbitMQ connection params - add_pubsub_args() includes RabbitMQ options (host, port, credentials) - add_pubsub_args(standalone=True) defaults to localhost for CLI tools - init_trustgraph skips Pulsar admin setup for non-Pulsar backends - tg-dump-queues and tg-monitor-prompts use backend abstraction - BaseClient and ConfigClient accept generic pubsub config
594 lines
19 KiB
Python
594 lines
19 KiB
Python
|
|
"""
|
|
Universal document decoder powered by the unstructured library.
|
|
|
|
Accepts documents in any common format (PDF, DOCX, XLSX, HTML, Markdown,
|
|
plain text, PPTX, etc.) on input, outputs pages or sections as text
|
|
as separate output objects.
|
|
|
|
Supports both inline document data and fetching from librarian via Pulsar
|
|
for large documents. Fetches document metadata from the librarian to
|
|
determine mime type for format detection.
|
|
|
|
Tables are preserved as HTML markup for better downstream extraction.
|
|
Images are stored in the librarian but not sent to the text pipeline.
|
|
"""
|
|
|
|
import base64
|
|
import logging
|
|
import magic
|
|
import tempfile
|
|
import os
|
|
|
|
from unstructured.partition.auto import partition
|
|
|
|
from ... schema import Document, TextDocument, Metadata
|
|
from ... schema import librarian_request_queue, librarian_response_queue
|
|
from ... schema import Triples
|
|
from ... base import FlowProcessor, ConsumerSpec, ProducerSpec, LibrarianClient
|
|
|
|
from ... provenance import (
|
|
document_uri, page_uri as make_page_uri,
|
|
section_uri as make_section_uri, image_uri as make_image_uri,
|
|
derived_entity_triples, set_graph, GRAPH_SOURCE,
|
|
)
|
|
|
|
from . strategies import get_strategy
|
|
|
|
# Component identification for provenance
|
|
COMPONENT_NAME = "universal-decoder"
|
|
COMPONENT_VERSION = "1.0.0"
|
|
|
|
# Module logger
|
|
logger = logging.getLogger(__name__)
|
|
|
|
default_ident = "document-decoder"
|
|
|
|
default_librarian_request_queue = librarian_request_queue
|
|
default_librarian_response_queue = librarian_response_queue
|
|
|
|
# Mime type to unstructured content_type mapping
|
|
# unstructured auto-detects most formats, but we pass the hint when available
|
|
MIME_EXTENSIONS = {
|
|
"application/pdf": ".pdf",
|
|
"application/vnd.openxmlformats-officedocument.wordprocessingml.document": ".docx",
|
|
"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": ".xlsx",
|
|
"application/vnd.ms-excel": ".xls",
|
|
"application/vnd.openxmlformats-officedocument.presentationml.presentation": ".pptx",
|
|
"text/html": ".html",
|
|
"text/markdown": ".md",
|
|
"text/plain": ".txt",
|
|
"text/csv": ".csv",
|
|
"text/tab-separated-values": ".tsv",
|
|
"application/rtf": ".rtf",
|
|
"text/x-rst": ".rst",
|
|
"application/vnd.oasis.opendocument.text": ".odt",
|
|
}
|
|
|
|
# Formats that have natural page boundaries
|
|
PAGE_BASED_FORMATS = {
|
|
"application/pdf",
|
|
"application/vnd.openxmlformats-officedocument.presentationml.presentation",
|
|
"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
|
|
"application/vnd.ms-excel",
|
|
}
|
|
|
|
|
|
def assemble_section_text(elements):
|
|
"""
|
|
Assemble text from a list of unstructured elements.
|
|
|
|
- Text elements: plain text, joined with double newlines
|
|
- Table elements: HTML table markup from text_as_html
|
|
- Image elements: skipped (stored separately, not in text output)
|
|
|
|
Returns:
|
|
tuple: (assembled_text, element_types_set, table_count, image_count)
|
|
"""
|
|
parts = []
|
|
element_types = set()
|
|
table_count = 0
|
|
image_count = 0
|
|
|
|
for el in elements:
|
|
category = getattr(el, 'category', 'UncategorizedText')
|
|
element_types.add(category)
|
|
|
|
if category == 'Image':
|
|
image_count += 1
|
|
continue # Images are NOT included in text output
|
|
|
|
if category == 'Table':
|
|
table_count += 1
|
|
# Prefer HTML representation for tables
|
|
html = getattr(el.metadata, 'text_as_html', None) if hasattr(el, 'metadata') else None
|
|
if html:
|
|
parts.append(html)
|
|
else:
|
|
# Fallback to plain text
|
|
text = getattr(el, 'text', '') or ''
|
|
if text:
|
|
parts.append(text)
|
|
else:
|
|
text = getattr(el, 'text', '') or ''
|
|
if text:
|
|
parts.append(text)
|
|
|
|
return '\n\n'.join(parts), element_types, table_count, image_count
|
|
|
|
|
|
class Processor(FlowProcessor):
|
|
|
|
def __init__(self, **params):
|
|
|
|
id = params.get("id", default_ident)
|
|
|
|
self.partition_strategy = params.get("strategy", "auto")
|
|
self.languages = params.get("languages", "eng").split(",")
|
|
self.section_strategy_name = params.get(
|
|
"section_strategy", "whole-document"
|
|
)
|
|
self.section_element_count = params.get("section_element_count", 20)
|
|
self.section_max_size = params.get("section_max_size", 4000)
|
|
self.section_within_pages = params.get("section_within_pages", False)
|
|
|
|
self.section_strategy = get_strategy(self.section_strategy_name)
|
|
|
|
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,
|
|
)
|
|
)
|
|
|
|
self.register_specification(
|
|
ProducerSpec(
|
|
name="triples",
|
|
schema=Triples,
|
|
)
|
|
)
|
|
|
|
# Librarian client
|
|
self.librarian = LibrarianClient(
|
|
id=id, backend=self.pubsub, taskgroup=self.taskgroup,
|
|
)
|
|
|
|
logger.info("Universal decoder initialized")
|
|
|
|
async def start(self):
|
|
await super(Processor, self).start()
|
|
await self.librarian.start()
|
|
|
|
def extract_elements(self, blob, mime_type=None):
|
|
"""
|
|
Extract elements from a document using unstructured.
|
|
|
|
Args:
|
|
blob: Raw document bytes
|
|
mime_type: Optional mime type hint
|
|
|
|
Returns:
|
|
List of unstructured Element objects
|
|
"""
|
|
# Determine file extension for unstructured
|
|
suffix = MIME_EXTENSIONS.get(mime_type, "") if mime_type else ""
|
|
if not suffix:
|
|
suffix = ".bin"
|
|
|
|
with tempfile.NamedTemporaryFile(
|
|
delete=False, suffix=suffix
|
|
) as fp:
|
|
fp.write(blob)
|
|
temp_path = fp.name
|
|
|
|
try:
|
|
kwargs = {
|
|
"filename": temp_path,
|
|
"strategy": self.partition_strategy,
|
|
"languages": self.languages,
|
|
}
|
|
|
|
# For hi_res strategy, request image extraction
|
|
if self.partition_strategy == "hi_res":
|
|
kwargs["extract_image_block_to_payload"] = True
|
|
|
|
elements = partition(**kwargs)
|
|
|
|
logger.info(
|
|
f"Extracted {len(elements)} elements "
|
|
f"(strategy: {self.partition_strategy})"
|
|
)
|
|
|
|
return elements
|
|
|
|
finally:
|
|
try:
|
|
os.unlink(temp_path)
|
|
except OSError:
|
|
pass
|
|
|
|
def group_by_page(self, elements):
|
|
"""
|
|
Group elements by page number.
|
|
|
|
Returns list of (page_number, elements) tuples.
|
|
"""
|
|
pages = {}
|
|
|
|
for el in elements:
|
|
page_num = getattr(
|
|
el.metadata, 'page_number', None
|
|
) if hasattr(el, 'metadata') else None
|
|
if page_num is None:
|
|
page_num = 1
|
|
if page_num not in pages:
|
|
pages[page_num] = []
|
|
pages[page_num].append(el)
|
|
|
|
return sorted(pages.items())
|
|
|
|
async def emit_section(self, elements, parent_doc_id, doc_uri_str,
|
|
metadata, flow, mime_type=None,
|
|
page_number=None, section_index=None):
|
|
"""
|
|
Process a group of elements as a page or section.
|
|
|
|
Assembles text, saves to librarian, emits provenance, sends
|
|
TextDocument downstream. Returns the entity URI.
|
|
"""
|
|
text, element_types, table_count, image_count = (
|
|
assemble_section_text(elements)
|
|
)
|
|
|
|
if not text.strip():
|
|
logger.debug("Skipping empty section")
|
|
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
|
|
page_content = text.encode("utf-8")
|
|
|
|
# Save to librarian
|
|
await self.librarian.save_child_document(
|
|
doc_id=doc_id,
|
|
parent_id=parent_doc_id,
|
|
user=metadata.user,
|
|
content=page_content,
|
|
document_type="page" if is_page else "section",
|
|
title=label,
|
|
)
|
|
|
|
# Emit provenance triples
|
|
element_types_str = ",".join(sorted(element_types)) if element_types else None
|
|
|
|
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,
|
|
element_types=element_types_str,
|
|
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,
|
|
user=metadata.user,
|
|
collection=metadata.collection,
|
|
),
|
|
triples=set_graph(prov_triples, GRAPH_SOURCE),
|
|
))
|
|
|
|
# Send TextDocument downstream (chunker will fetch from librarian)
|
|
r = TextDocument(
|
|
metadata=Metadata(
|
|
id=entity_uri,
|
|
root=metadata.root,
|
|
user=metadata.user,
|
|
collection=metadata.collection,
|
|
),
|
|
document_id=doc_id,
|
|
text=b"",
|
|
)
|
|
|
|
await flow("output").send(r)
|
|
|
|
return entity_uri
|
|
|
|
async def emit_image(self, element, parent_uri, parent_doc_id,
|
|
metadata, flow, mime_type=None, page_number=None):
|
|
"""
|
|
Store an image element in the librarian with provenance.
|
|
|
|
Images are stored but NOT sent downstream to the text pipeline.
|
|
"""
|
|
img_uri = make_image_uri()
|
|
|
|
# Get image data
|
|
img_data = None
|
|
if hasattr(element, 'metadata'):
|
|
img_data = getattr(element.metadata, 'image_base64', None)
|
|
|
|
if not img_data:
|
|
# No image payload available, just record provenance
|
|
logger.debug("Image element without payload, recording provenance only")
|
|
img_content = b""
|
|
img_kind = "image/unknown"
|
|
else:
|
|
if isinstance(img_data, str):
|
|
img_content = base64.b64decode(img_data)
|
|
else:
|
|
img_content = img_data
|
|
img_kind = "image/png" # unstructured typically extracts as PNG
|
|
|
|
# Save to librarian
|
|
if img_content:
|
|
await self.librarian.save_child_document(
|
|
doc_id=img_uri,
|
|
parent_id=parent_doc_id,
|
|
user=metadata.user,
|
|
content=img_content,
|
|
document_type="image",
|
|
title=f"Image from page {page_number}" if page_number else "Image",
|
|
kind=img_kind,
|
|
)
|
|
|
|
# Emit provenance triples
|
|
prov_triples = derived_entity_triples(
|
|
entity_uri=img_uri,
|
|
parent_uri=parent_uri,
|
|
component_name=COMPONENT_NAME,
|
|
component_version=COMPONENT_VERSION,
|
|
label=f"Image from page {page_number}" if page_number else "Image",
|
|
image=True,
|
|
page_number=page_number,
|
|
mime_type=mime_type,
|
|
)
|
|
|
|
await flow("triples").send(Triples(
|
|
metadata=Metadata(
|
|
id=img_uri,
|
|
root=metadata.root,
|
|
user=metadata.user,
|
|
collection=metadata.collection,
|
|
),
|
|
triples=set_graph(prov_triples, GRAPH_SOURCE),
|
|
))
|
|
|
|
async def on_message(self, msg, consumer, flow):
|
|
|
|
logger.debug("Document message received")
|
|
|
|
v = msg.value()
|
|
|
|
logger.info(f"Decoding {v.metadata.id}...")
|
|
|
|
# Determine content and mime type
|
|
mime_type = None
|
|
|
|
if v.document_id:
|
|
# Librarian path: fetch metadata then content
|
|
logger.info(
|
|
f"Fetching document {v.document_id} from librarian..."
|
|
)
|
|
|
|
doc_meta = await self.librarian.fetch_document_metadata(
|
|
document_id=v.document_id,
|
|
user=v.metadata.user,
|
|
)
|
|
mime_type = doc_meta.kind if doc_meta else None
|
|
|
|
content = await self.librarian.fetch_document_content(
|
|
document_id=v.document_id,
|
|
user=v.metadata.user,
|
|
)
|
|
|
|
if isinstance(content, str):
|
|
content = content.encode('utf-8')
|
|
blob = base64.b64decode(content)
|
|
|
|
logger.info(
|
|
f"Fetched {len(blob)} bytes, mime: {mime_type}"
|
|
)
|
|
else:
|
|
# Inline path: detect format from content
|
|
blob = base64.b64decode(v.data)
|
|
try:
|
|
mime_type = magic.from_buffer(blob, mime=True)
|
|
logger.info(f"Detected mime type: {mime_type}")
|
|
except Exception as e:
|
|
logger.warning(f"Could not detect mime type: {e}")
|
|
|
|
# Get the source document ID
|
|
source_doc_id = v.document_id or v.metadata.id
|
|
doc_uri_str = document_uri(source_doc_id)
|
|
|
|
# Extract elements using unstructured
|
|
elements = self.extract_elements(blob, mime_type)
|
|
|
|
if not elements:
|
|
logger.warning("No elements extracted from document")
|
|
return
|
|
|
|
# Determine if this is a page-based format
|
|
is_page_based = mime_type in PAGE_BASED_FORMATS if mime_type else False
|
|
|
|
# Also check if elements actually have page numbers
|
|
if not is_page_based:
|
|
has_pages = any(
|
|
getattr(el.metadata, 'page_number', None) is not None
|
|
for el in elements
|
|
if hasattr(el, 'metadata')
|
|
)
|
|
if has_pages:
|
|
is_page_based = True
|
|
|
|
if is_page_based:
|
|
# Group by page
|
|
page_groups = self.group_by_page(elements)
|
|
|
|
for page_num, page_elements in page_groups:
|
|
|
|
# Extract and store images separately
|
|
image_elements = [
|
|
el for el in page_elements
|
|
if getattr(el, 'category', '') == 'Image'
|
|
]
|
|
text_elements = [
|
|
el for el in page_elements
|
|
if getattr(el, 'category', '') != 'Image'
|
|
]
|
|
|
|
# Emit the page as a text section
|
|
page_uri_str = await self.emit_section(
|
|
text_elements, source_doc_id, doc_uri_str,
|
|
v.metadata, flow,
|
|
mime_type=mime_type, page_number=page_num,
|
|
)
|
|
|
|
# Store images (not sent to text pipeline)
|
|
for img_el in image_elements:
|
|
await self.emit_image(
|
|
img_el,
|
|
page_uri_str or doc_uri_str,
|
|
source_doc_id,
|
|
v.metadata, flow,
|
|
mime_type=mime_type, page_number=page_num,
|
|
)
|
|
|
|
else:
|
|
# Non-page format: use section strategy
|
|
|
|
# Separate images from text elements
|
|
image_elements = [
|
|
el for el in elements
|
|
if getattr(el, 'category', '') == 'Image'
|
|
]
|
|
text_elements = [
|
|
el for el in elements
|
|
if getattr(el, 'category', '') != 'Image'
|
|
]
|
|
|
|
# Apply section strategy to text elements
|
|
strategy_kwargs = {
|
|
'element_count': self.section_element_count,
|
|
'max_size': self.section_max_size,
|
|
}
|
|
groups = self.section_strategy(text_elements, **strategy_kwargs)
|
|
|
|
for idx, group in enumerate(groups):
|
|
section_idx = idx + 1
|
|
|
|
await self.emit_section(
|
|
group, source_doc_id, doc_uri_str,
|
|
v.metadata, flow,
|
|
mime_type=mime_type, section_index=section_idx,
|
|
)
|
|
|
|
# Store images (not sent to text pipeline)
|
|
for img_el in image_elements:
|
|
await self.emit_image(
|
|
img_el, doc_uri_str, source_doc_id,
|
|
v.metadata, flow,
|
|
mime_type=mime_type,
|
|
)
|
|
|
|
logger.info("Document decoding complete")
|
|
|
|
@staticmethod
|
|
def add_args(parser):
|
|
|
|
FlowProcessor.add_args(parser)
|
|
|
|
parser.add_argument(
|
|
'--strategy',
|
|
default='auto',
|
|
choices=['auto', 'hi_res', 'fast'],
|
|
help='Partitioning strategy (default: auto)',
|
|
)
|
|
|
|
parser.add_argument(
|
|
'--languages',
|
|
default='eng',
|
|
help='Comma-separated OCR language codes (default: eng)',
|
|
)
|
|
|
|
parser.add_argument(
|
|
'--section-strategy',
|
|
default='whole-document',
|
|
choices=[
|
|
'whole-document', 'heading', 'element-type', 'count', 'size'
|
|
],
|
|
help='Section grouping strategy for non-page formats '
|
|
'(default: whole-document)',
|
|
)
|
|
|
|
parser.add_argument(
|
|
'--section-element-count',
|
|
type=int,
|
|
default=20,
|
|
help='Elements per section for count strategy (default: 20)',
|
|
)
|
|
|
|
parser.add_argument(
|
|
'--section-max-size',
|
|
type=int,
|
|
default=4000,
|
|
help='Max chars per section for size strategy (default: 4000)',
|
|
)
|
|
|
|
parser.add_argument(
|
|
'--section-within-pages',
|
|
action='store_true',
|
|
default=False,
|
|
help='Apply section strategy within pages too (default: false)',
|
|
)
|
|
|
|
parser.add_argument(
|
|
'--librarian-request-queue',
|
|
default=default_librarian_request_queue,
|
|
help=f'Librarian request queue '
|
|
f'(default: {default_librarian_request_queue})',
|
|
)
|
|
|
|
parser.add_argument(
|
|
'--librarian-response-queue',
|
|
default=default_librarian_response_queue,
|
|
help=f'Librarian response queue '
|
|
f'(default: {default_librarian_response_queue})',
|
|
)
|
|
|
|
|
|
def run():
|
|
|
|
Processor.launch(default_ident, __doc__)
|