diff --git a/tests/unit/test_provenance/test_uris.py b/tests/unit/test_provenance/test_uris.py index 0e69734c..05bb7a1b 100644 --- a/tests/unit/test_provenance/test_uris.py +++ b/tests/unit/test_provenance/test_uris.py @@ -10,8 +10,7 @@ from trustgraph.provenance.uris import ( _encode_id, document_uri, page_uri, - chunk_uri_from_page, - chunk_uri_from_doc, + chunk_uri, activity_uri, subgraph_uri, agent_uri, @@ -60,31 +59,22 @@ class TestDocumentUris: assert document_uri(iri) == iri def test_page_uri_format(self): - result = page_uri("https://example.com/doc/123", 5) - assert result == "https://example.com/doc/123/p5" + result = page_uri() + assert result.startswith("urn:page:") - def test_page_uri_page_zero(self): - result = page_uri("https://example.com/doc/123", 0) - assert result == "https://example.com/doc/123/p0" + def test_page_uri_unique(self): + r1 = page_uri() + r2 = page_uri() + assert r1 != r2 - def test_chunk_uri_from_page_format(self): - result = chunk_uri_from_page("https://example.com/doc/123", 2, 3) - assert result == "https://example.com/doc/123/p2/c3" + def test_chunk_uri_format(self): + result = chunk_uri() + assert result.startswith("urn:chunk:") - def test_chunk_uri_from_doc_format(self): - result = chunk_uri_from_doc("https://example.com/doc/123", 7) - assert result == "https://example.com/doc/123/c7" - - def test_page_uri_preserves_doc_iri(self): - doc = "urn:isbn:978-3-16-148410-0" - result = page_uri(doc, 1) - assert result.startswith(doc) - - def test_chunk_from_page_hierarchy(self): - """Chunk URI should contain both page and chunk identifiers.""" - result = chunk_uri_from_page("https://example.com/doc", 3, 5) - assert "/p3/" in result - assert result.endswith("/c5") + def test_chunk_uri_unique(self): + r1 = chunk_uri() + r2 = chunk_uri() + assert r1 != r2 class TestActivityAndSubgraphUris: diff --git a/trustgraph-base/trustgraph/provenance/__init__.py b/trustgraph-base/trustgraph/provenance/__init__.py index 18ecb0e8..5b9d2129 100644 --- a/trustgraph-base/trustgraph/provenance/__init__.py +++ b/trustgraph-base/trustgraph/provenance/__init__.py @@ -9,14 +9,14 @@ Provides helpers for: Usage example: from trustgraph.provenance import ( - document_uri, page_uri, chunk_uri_from_page, + document_uri, page_uri, chunk_uri, document_triples, derived_entity_triples, get_vocabulary_triples, ) # Generate URIs doc_uri = document_uri("my-doc-123") - page_uri = page_uri("my-doc-123", page_number=1) + pg_uri = page_uri() # Build provenance triples triples = document_triples( @@ -35,8 +35,7 @@ from . uris import ( TRUSTGRAPH_BASE, document_uri, page_uri, - chunk_uri_from_page, - chunk_uri_from_doc, + chunk_uri, activity_uri, subgraph_uri, agent_uri, @@ -138,8 +137,7 @@ __all__ = [ "TRUSTGRAPH_BASE", "document_uri", "page_uri", - "chunk_uri_from_page", - "chunk_uri_from_doc", + "chunk_uri", "activity_uri", "subgraph_uri", "agent_uri", diff --git a/trustgraph-base/trustgraph/provenance/uris.py b/trustgraph-base/trustgraph/provenance/uris.py index 670143df..d851fa0b 100644 --- a/trustgraph-base/trustgraph/provenance/uris.py +++ b/trustgraph-base/trustgraph/provenance/uris.py @@ -1,12 +1,11 @@ """ URI generation for provenance entities. -Document IDs are already IRIs (e.g., https://trustgraph.ai/doc/abc123). -Child entities (pages, chunks) append path segments to the parent IRI: -- Document: {doc_iri} (as provided) -- Page: {doc_iri}/p{page_number} -- Chunk: {page_iri}/c{chunk_index} (from page) - {doc_iri}/c{chunk_index} (from text doc) +Document IDs are externally provided (e.g., https://trustgraph.ai/doc/abc123). +Child entities (pages, chunks) use UUID-based URNs: +- Document: {doc_iri} (as provided, not generated here) +- Page: urn:page:{uuid} +- Chunk: urn:chunk:{uuid} - Activity: https://trustgraph.ai/activity/{uuid} - Subgraph: https://trustgraph.ai/subgraph/{uuid} """ @@ -28,19 +27,14 @@ def document_uri(doc_iri: str) -> str: return doc_iri -def page_uri(doc_iri: str, page_number: int) -> str: - """Generate URI for a page by appending to document IRI.""" - return f"{doc_iri}/p{page_number}" +def page_uri() -> str: + """Generate a unique URI for a page.""" + return f"urn:page:{uuid.uuid4()}" -def chunk_uri_from_page(doc_iri: str, page_number: int, chunk_index: int) -> str: - """Generate URI for a chunk extracted from a page.""" - return f"{doc_iri}/p{page_number}/c{chunk_index}" - - -def chunk_uri_from_doc(doc_iri: str, chunk_index: int) -> str: - """Generate URI for a chunk extracted directly from a text document.""" - return f"{doc_iri}/c{chunk_index}" +def chunk_uri() -> str: + """Generate a unique URI for a chunk.""" + return f"urn:chunk:{uuid.uuid4()}" def activity_uri(activity_id: str = None) -> str: diff --git a/trustgraph-flow/trustgraph/chunking/recursive/chunker.py b/trustgraph-flow/trustgraph/chunking/recursive/chunker.py index fb84c356..64d58457 100755 --- a/trustgraph-flow/trustgraph/chunking/recursive/chunker.py +++ b/trustgraph-flow/trustgraph/chunking/recursive/chunker.py @@ -12,7 +12,7 @@ from ... schema import TextDocument, Chunk, Metadata, Triples from ... base import ChunkingService, ConsumerSpec, ProducerSpec from ... provenance import ( - derived_entity_triples, + chunk_uri as make_chunk_uri, derived_entity_triples, set_graph, GRAPH_SOURCE, ) @@ -124,10 +124,9 @@ class Processor(ChunkingService): logger.debug(f"Created chunk of size {len(chunk.page_content)}") - # Generate chunk document ID by appending /c{index} to parent - # Works for both page URIs (doc/p3 -> doc/p3/c1) and doc URIs (doc -> doc/c1) - chunk_doc_id = f"{parent_doc_id}/c{chunk_index}" - chunk_uri = chunk_doc_id # URI is same as document ID + # Generate unique chunk ID + c_uri = make_chunk_uri() + chunk_doc_id = c_uri parent_uri = parent_doc_id chunk_content = chunk.page_content.encode("utf-8") @@ -145,7 +144,7 @@ class Processor(ChunkingService): # Emit provenance triples (stored in source graph for separation from core knowledge) prov_triples = derived_entity_triples( - entity_uri=chunk_uri, + entity_uri=c_uri, parent_uri=parent_uri, component_name=COMPONENT_NAME, component_version=COMPONENT_VERSION, @@ -159,7 +158,7 @@ class Processor(ChunkingService): await flow("triples").send(Triples( metadata=Metadata( - id=chunk_uri, + id=c_uri, root=v.metadata.root, user=v.metadata.user, collection=v.metadata.collection, @@ -170,7 +169,7 @@ class Processor(ChunkingService): # Forward chunk ID + content (post-chunker optimization) r = Chunk( metadata=Metadata( - id=chunk_uri, + id=c_uri, root=v.metadata.root, user=v.metadata.user, collection=v.metadata.collection, diff --git a/trustgraph-flow/trustgraph/chunking/token/chunker.py b/trustgraph-flow/trustgraph/chunking/token/chunker.py index 909396c6..4302250e 100755 --- a/trustgraph-flow/trustgraph/chunking/token/chunker.py +++ b/trustgraph-flow/trustgraph/chunking/token/chunker.py @@ -12,7 +12,7 @@ from ... schema import TextDocument, Chunk, Metadata, Triples from ... base import ChunkingService, ConsumerSpec, ProducerSpec from ... provenance import ( - derived_entity_triples, + chunk_uri as make_chunk_uri, derived_entity_triples, set_graph, GRAPH_SOURCE, ) @@ -122,10 +122,9 @@ class Processor(ChunkingService): logger.debug(f"Created chunk of size {len(chunk.page_content)}") - # Generate chunk document ID by appending /c{index} to parent - # Works for both page URIs (doc/p3 -> doc/p3/c1) and doc URIs (doc -> doc/c1) - chunk_doc_id = f"{parent_doc_id}/c{chunk_index}" - chunk_uri = chunk_doc_id # URI is same as document ID + # Generate unique chunk ID + c_uri = make_chunk_uri() + chunk_doc_id = c_uri parent_uri = parent_doc_id chunk_content = chunk.page_content.encode("utf-8") @@ -143,7 +142,7 @@ class Processor(ChunkingService): # Emit provenance triples (stored in source graph for separation from core knowledge) prov_triples = derived_entity_triples( - entity_uri=chunk_uri, + entity_uri=c_uri, parent_uri=parent_uri, component_name=COMPONENT_NAME, component_version=COMPONENT_VERSION, @@ -157,7 +156,7 @@ class Processor(ChunkingService): await flow("triples").send(Triples( metadata=Metadata( - id=chunk_uri, + id=c_uri, root=v.metadata.root, user=v.metadata.user, collection=v.metadata.collection, @@ -168,7 +167,7 @@ class Processor(ChunkingService): # Forward chunk ID + content (post-chunker optimization) r = Chunk( metadata=Metadata( - id=chunk_uri, + id=c_uri, root=v.metadata.root, user=v.metadata.user, collection=v.metadata.collection, diff --git a/trustgraph-flow/trustgraph/decoding/mistral_ocr/processor.py b/trustgraph-flow/trustgraph/decoding/mistral_ocr/processor.py index 3cacb16c..6207d659 100755 --- a/trustgraph-flow/trustgraph/decoding/mistral_ocr/processor.py +++ b/trustgraph-flow/trustgraph/decoding/mistral_ocr/processor.py @@ -1,29 +1,48 @@ """ -Simple decoder, accepts PDF documents on input, outputs pages from the -PDF document as text as separate output objects. +Mistral OCR decoder, accepts PDF documents on input, outputs pages from the +PDF document as markdown text as separate output objects. + +Supports both inline document data and fetching from librarian via Pulsar +for large documents. """ from pypdf import PdfWriter, PdfReader from io import BytesIO +import asyncio import base64 import uuid import os from mistralai import Mistral -from mistralai import DocumentURLChunk, ImageURLChunk, TextChunk from mistralai.models import OCRResponse from ... schema import Document, TextDocument, Metadata +from ... schema import LibrarianRequest, LibrarianResponse, DocumentMetadata +from ... schema import librarian_request_queue, librarian_response_queue +from ... schema import Triples from ... base import FlowProcessor, ConsumerSpec, ProducerSpec +from ... base import Consumer, Producer, ConsumerMetrics, ProducerMetrics + +from ... provenance import ( + document_uri, page_uri as make_page_uri, derived_entity_triples, + set_graph, GRAPH_SOURCE, +) import logging logger = logging.getLogger(__name__) +# Component identification for provenance +COMPONENT_NAME = "mistral-ocr-decoder" +COMPONENT_VERSION = "1.0.0" + default_ident = "pdf-decoder" default_api_key = os.getenv("MISTRAL_TOKEN") +default_librarian_request_queue = librarian_request_queue +default_librarian_response_queue = librarian_response_queue + pages_per_chunk = 5 def chunks(lst, n): @@ -48,27 +67,6 @@ def replace_images_in_markdown(markdown_str: str, images_dict: dict) -> str: ) return markdown_str -def get_combined_markdown(ocr_response: OCRResponse) -> str: - """ - Combine OCR text and images into a single markdown document. - - Args: - ocr_response: Response from OCR processing containing text and images - - Returns: - Combined markdown string with embedded images - """ - markdowns: list[str] = [] - # Extract images from page - for page in ocr_response.pages: - image_data = {} - for img in page.images: - image_data[img.id] = img.image_base64 - # Replace image placeholders with actual images - markdowns.append(replace_images_in_markdown(page.markdown, image_data)) - - return "\n\n".join(markdowns) - class Processor(FlowProcessor): def __init__(self, **params): @@ -97,6 +95,50 @@ class Processor(FlowProcessor): ) ) + self.register_specification( + ProducerSpec( + name = "triples", + schema = Triples, + ) + ) + + # Librarian client for fetching document content + librarian_request_q = params.get( + "librarian_request_queue", default_librarian_request_queue + ) + librarian_response_q = params.get( + "librarian_response_queue", default_librarian_response_queue + ) + + librarian_request_metrics = ProducerMetrics( + processor = id, flow = None, name = "librarian-request" + ) + + self.librarian_request_producer = Producer( + backend = self.pubsub, + topic = librarian_request_q, + schema = LibrarianRequest, + metrics = librarian_request_metrics, + ) + + librarian_response_metrics = ConsumerMetrics( + processor = id, flow = None, name = "librarian-response" + ) + + self.librarian_response_consumer = Consumer( + taskgroup = self.taskgroup, + backend = self.pubsub, + flow = None, + topic = librarian_response_q, + subscriber = f"{id}-librarian", + schema = LibrarianResponse, + handler = self.on_librarian_response, + metrics = librarian_response_metrics, + ) + + # Pending librarian requests: request_id -> asyncio.Future + self.pending_requests = {} + if api_key is None: raise RuntimeError("Mistral API key not specified") @@ -107,15 +149,125 @@ class Processor(FlowProcessor): logger.info("Mistral OCR processor initialized") + async def start(self): + await super(Processor, self).start() + await self.librarian_request_producer.start() + await self.librarian_response_consumer.start() + + async def on_librarian_response(self, msg, consumer, flow): + """Handle responses from the librarian service.""" + response = msg.value() + request_id = msg.properties().get("id") + + if request_id and request_id in self.pending_requests: + future = self.pending_requests.pop(request_id) + future.set_result(response) + else: + logger.warning(f"Received unexpected librarian response: {request_id}") + + async def fetch_document_content(self, document_id, user, timeout=120): + """ + Fetch document content from librarian via Pulsar. + """ + request_id = str(uuid.uuid4()) + + request = LibrarianRequest( + operation="get-document-content", + document_id=document_id, + user=user, + ) + + # Create future for response + future = asyncio.get_event_loop().create_future() + self.pending_requests[request_id] = future + + try: + # Send request + await self.librarian_request_producer.send( + request, properties={"id": request_id} + ) + + # Wait for response + response = await asyncio.wait_for(future, timeout=timeout) + + if response.error: + raise RuntimeError( + f"Librarian error: {response.error.type}: {response.error.message}" + ) + + return response.content + + except asyncio.TimeoutError: + self.pending_requests.pop(request_id, None) + raise RuntimeError(f"Timeout fetching document {document_id}") + + async def save_child_document(self, doc_id, parent_id, user, content, + document_type="page", title=None, timeout=120): + """ + Save a child document to the librarian. + """ + request_id = str(uuid.uuid4()) + + doc_metadata = DocumentMetadata( + id=doc_id, + user=user, + kind="text/plain", + title=title or doc_id, + parent_id=parent_id, + document_type=document_type, + ) + + request = LibrarianRequest( + operation="add-child-document", + document_metadata=doc_metadata, + content=base64.b64encode(content).decode("utf-8"), + ) + + # Create future for response + future = asyncio.get_event_loop().create_future() + self.pending_requests[request_id] = future + + try: + # Send request + await self.librarian_request_producer.send( + request, properties={"id": request_id} + ) + + # Wait for response + response = await asyncio.wait_for(future, timeout=timeout) + + if response.error: + raise RuntimeError( + f"Librarian error saving child document: {response.error.type}: {response.error.message}" + ) + + return doc_id + + except asyncio.TimeoutError: + self.pending_requests.pop(request_id, None) + raise RuntimeError(f"Timeout saving child document {doc_id}") + def ocr(self, blob): + """ + Run Mistral OCR on a PDF blob, returning per-page markdown strings. + + Args: + blob: Raw PDF bytes + + Returns: + List of (page_markdown, page_number) tuples, 1-indexed + """ logger.debug("Parse PDF...") pdfbuf = BytesIO(blob) pdf = PdfReader(pdfbuf) + pages = [] + global_page_num = 0 + for chunk in chunks(pdf.pages, pages_per_chunk): - + logger.debug("Get next pages...") part = PdfWriter() @@ -152,11 +304,19 @@ class Processor(FlowProcessor): logger.debug("Extract markdown...") - markdown = get_combined_markdown(processed) + for page in processed.pages: + global_page_num += 1 + image_data = {} + for img in page.images: + image_data[img.id] = img.image_base64 + markdown = replace_images_in_markdown( + page.markdown, image_data + ) + pages.append((markdown, global_page_num)) - logger.info("OCR complete.") + logger.info(f"OCR complete, {len(pages)} pages.") - return markdown + return pages async def on_message(self, msg, consumer, flow): @@ -166,16 +326,83 @@ class Processor(FlowProcessor): logger.info(f"Decoding {v.metadata.id}...") - markdown = self.ocr(base64.b64decode(v.data)) + # Get PDF content - fetch from librarian or use inline data + if v.document_id: + logger.info(f"Fetching document {v.document_id} from librarian...") + content = await self.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 from librarian") + else: + blob = base64.b64decode(v.data) - r = TextDocument( - metadata=v.metadata, - text=markdown.encode("utf-8"), - ) + # Get the source document ID + source_doc_id = v.document_id or v.metadata.id - await flow("output").send(r) + # Run OCR, get per-page markdown + pages = self.ocr(blob) - logger.info("Done.") + for markdown, page_num in pages: + + logger.debug(f"Processing page {page_num}") + + # Generate unique page ID + pg_uri = make_page_uri() + page_doc_id = pg_uri + page_content = markdown.encode("utf-8") + + # Save page as child document in librarian + await self.save_child_document( + doc_id=page_doc_id, + parent_id=source_doc_id, + user=v.metadata.user, + content=page_content, + document_type="page", + title=f"Page {page_num}", + ) + + # Emit provenance triples + doc_uri = document_uri(source_doc_id) + + prov_triples = derived_entity_triples( + entity_uri=pg_uri, + parent_uri=doc_uri, + component_name=COMPONENT_NAME, + component_version=COMPONENT_VERSION, + label=f"Page {page_num}", + page_number=page_num, + ) + + await flow("triples").send(Triples( + metadata=Metadata( + id=pg_uri, + root=v.metadata.root, + user=v.metadata.user, + collection=v.metadata.collection, + ), + triples=set_graph(prov_triples, GRAPH_SOURCE), + )) + + # Forward page document ID to chunker + # Chunker will fetch content from librarian + r = TextDocument( + metadata=Metadata( + id=pg_uri, + root=v.metadata.root, + user=v.metadata.user, + collection=v.metadata.collection, + ), + document_id=page_doc_id, + text=b"", # Empty, chunker will fetch from librarian + ) + + await flow("output").send(r) + + logger.debug("PDF decoding complete") @staticmethod def add_args(parser): @@ -188,7 +415,18 @@ class Processor(FlowProcessor): help=f'Mistral API Key' ) + parser.add_argument( + '--librarian-request-queue', + default=default_librarian_request_queue, + help=f'Librarian request queue (default: {default_librarian_request_queue})', + ) + + parser.add_argument( + '--librarian-response-queue', + default=default_librarian_response_queue, + help=f'Librarian response queue (default: {default_librarian_response_queue})', + ) + def run(): Processor.launch(default_ident, __doc__) - diff --git a/trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py b/trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py index 550948fe..865b984e 100755 --- a/trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py +++ b/trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py @@ -23,7 +23,7 @@ from ... base import FlowProcessor, ConsumerSpec, ProducerSpec from ... base import Consumer, Producer, ConsumerMetrics, ProducerMetrics from ... provenance import ( - document_uri, page_uri, derived_entity_triples, + document_uri, page_uri as make_page_uri, derived_entity_triples, set_graph, GRAPH_SOURCE, ) @@ -272,8 +272,9 @@ class Processor(FlowProcessor): logger.debug(f"Processing page {page_num}") - # Generate page document ID - page_doc_id = f"{source_doc_id}/p{page_num}" + # Generate unique page ID + pg_uri = make_page_uri() + page_doc_id = pg_uri page_content = page.page_content.encode("utf-8") # Save page as child document in librarian @@ -288,7 +289,6 @@ class Processor(FlowProcessor): # Emit provenance triples (stored in source graph for separation from core knowledge) doc_uri = document_uri(source_doc_id) - pg_uri = page_uri(source_doc_id, page_num) prov_triples = derived_entity_triples( entity_uri=pg_uri, diff --git a/trustgraph-ocr/trustgraph/decoding/ocr/pdf_decoder.py b/trustgraph-ocr/trustgraph/decoding/ocr/pdf_decoder.py index b5aac3c2..0c94039d 100755 --- a/trustgraph-ocr/trustgraph/decoding/ocr/pdf_decoder.py +++ b/trustgraph-ocr/trustgraph/decoding/ocr/pdf_decoder.py @@ -2,22 +2,42 @@ """ Simple decoder, accepts PDF documents on input, outputs pages from the PDF document as text as separate output objects. + +Supports both inline document data and fetching from librarian via Pulsar +for large documents. """ -import tempfile +import asyncio import base64 import logging +import uuid import pytesseract from pdf2image import convert_from_bytes from ... schema import Document, TextDocument, Metadata +from ... schema import LibrarianRequest, LibrarianResponse, DocumentMetadata +from ... schema import librarian_request_queue, librarian_response_queue +from ... schema import Triples from ... base import FlowProcessor, ConsumerSpec, ProducerSpec +from ... base import Consumer, Producer, ConsumerMetrics, ProducerMetrics + +from ... provenance import ( + document_uri, page_uri as make_page_uri, derived_entity_triples, + set_graph, GRAPH_SOURCE, +) + +# Component identification for provenance +COMPONENT_NAME = "tesseract-ocr-decoder" +COMPONENT_VERSION = "1.0.0" # Module logger logger = logging.getLogger(__name__) default_ident = "pdf-decoder" +default_librarian_request_queue = librarian_request_queue +default_librarian_response_queue = librarian_response_queue + class Processor(FlowProcessor): def __init__(self, **params): @@ -45,8 +65,150 @@ class Processor(FlowProcessor): ) ) + self.register_specification( + ProducerSpec( + name = "triples", + schema = Triples, + ) + ) + + # Librarian client for fetching document content + librarian_request_q = params.get( + "librarian_request_queue", default_librarian_request_queue + ) + librarian_response_q = params.get( + "librarian_response_queue", default_librarian_response_queue + ) + + librarian_request_metrics = ProducerMetrics( + processor = id, flow = None, name = "librarian-request" + ) + + self.librarian_request_producer = Producer( + backend = self.pubsub, + topic = librarian_request_q, + schema = LibrarianRequest, + metrics = librarian_request_metrics, + ) + + librarian_response_metrics = ConsumerMetrics( + processor = id, flow = None, name = "librarian-response" + ) + + self.librarian_response_consumer = Consumer( + taskgroup = self.taskgroup, + backend = self.pubsub, + flow = None, + topic = librarian_response_q, + subscriber = f"{id}-librarian", + schema = LibrarianResponse, + handler = self.on_librarian_response, + metrics = librarian_response_metrics, + ) + + # Pending librarian requests: request_id -> asyncio.Future + self.pending_requests = {} + logger.info("PDF OCR processor initialized") + async def start(self): + await super(Processor, self).start() + await self.librarian_request_producer.start() + await self.librarian_response_consumer.start() + + async def on_librarian_response(self, msg, consumer, flow): + """Handle responses from the librarian service.""" + response = msg.value() + request_id = msg.properties().get("id") + + if request_id and request_id in self.pending_requests: + future = self.pending_requests.pop(request_id) + future.set_result(response) + else: + logger.warning(f"Received unexpected librarian response: {request_id}") + + async def fetch_document_content(self, document_id, user, timeout=120): + """ + Fetch document content from librarian via Pulsar. + """ + request_id = str(uuid.uuid4()) + + request = LibrarianRequest( + operation="get-document-content", + document_id=document_id, + user=user, + ) + + # Create future for response + future = asyncio.get_event_loop().create_future() + self.pending_requests[request_id] = future + + try: + # Send request + await self.librarian_request_producer.send( + request, properties={"id": request_id} + ) + + # Wait for response + response = await asyncio.wait_for(future, timeout=timeout) + + if response.error: + raise RuntimeError( + f"Librarian error: {response.error.type}: {response.error.message}" + ) + + return response.content + + except asyncio.TimeoutError: + self.pending_requests.pop(request_id, None) + raise RuntimeError(f"Timeout fetching document {document_id}") + + async def save_child_document(self, doc_id, parent_id, user, content, + document_type="page", title=None, timeout=120): + """ + Save a child document to the librarian. + """ + request_id = str(uuid.uuid4()) + + doc_metadata = DocumentMetadata( + id=doc_id, + user=user, + kind="text/plain", + title=title or doc_id, + parent_id=parent_id, + document_type=document_type, + ) + + request = LibrarianRequest( + operation="add-child-document", + document_metadata=doc_metadata, + content=base64.b64encode(content).decode("utf-8"), + ) + + # Create future for response + future = asyncio.get_event_loop().create_future() + self.pending_requests[request_id] = future + + try: + # Send request + await self.librarian_request_producer.send( + request, properties={"id": request_id} + ) + + # Wait for response + response = await asyncio.wait_for(future, timeout=timeout) + + if response.error: + raise RuntimeError( + f"Librarian error saving child document: {response.error.type}: {response.error.message}" + ) + + return doc_id + + except asyncio.TimeoutError: + self.pending_requests.pop(request_id, None) + raise RuntimeError(f"Timeout saving child document {doc_id}") + async def on_message(self, msg, consumer, flow): logger.info("PDF message received") @@ -55,21 +217,85 @@ class Processor(FlowProcessor): logger.info(f"Decoding {v.metadata.id}...") - blob = base64.b64decode(v.data) + # Get PDF content - fetch from librarian or use inline data + if v.document_id: + logger.info(f"Fetching document {v.document_id} from librarian...") + content = await self.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 from librarian") + else: + blob = base64.b64decode(v.data) + + # Get the source document ID + source_doc_id = v.document_id or v.metadata.id pages = convert_from_bytes(blob) for ix, page in enumerate(pages): + page_num = ix + 1 # 1-indexed + try: text = pytesseract.image_to_string(page, lang='eng') except Exception as e: - logger.warning(f"Page did not OCR: {e}") + logger.warning(f"Page {page_num} did not OCR: {e}") continue + logger.debug(f"Processing page {page_num}") + + # Generate unique page ID + pg_uri = make_page_uri() + page_doc_id = pg_uri + page_content = text.encode("utf-8") + + # Save page as child document in librarian + await self.save_child_document( + doc_id=page_doc_id, + parent_id=source_doc_id, + user=v.metadata.user, + content=page_content, + document_type="page", + title=f"Page {page_num}", + ) + + # Emit provenance triples + doc_uri = document_uri(source_doc_id) + + prov_triples = derived_entity_triples( + entity_uri=pg_uri, + parent_uri=doc_uri, + component_name=COMPONENT_NAME, + component_version=COMPONENT_VERSION, + label=f"Page {page_num}", + page_number=page_num, + ) + + await flow("triples").send(Triples( + metadata=Metadata( + id=pg_uri, + root=v.metadata.root, + user=v.metadata.user, + collection=v.metadata.collection, + ), + triples=set_graph(prov_triples, GRAPH_SOURCE), + )) + + # Forward page document ID to chunker + # Chunker will fetch content from librarian r = TextDocument( - metadata=v.metadata, - text=text.encode("utf-8"), + metadata=Metadata( + id=pg_uri, + root=v.metadata.root, + user=v.metadata.user, + collection=v.metadata.collection, + ), + document_id=page_doc_id, + text=b"", # Empty, chunker will fetch from librarian ) await flow("output").send(r) @@ -78,9 +304,21 @@ class Processor(FlowProcessor): @staticmethod def add_args(parser): + FlowProcessor.add_args(parser) + parser.add_argument( + '--librarian-request-queue', + default=default_librarian_request_queue, + help=f'Librarian request queue (default: {default_librarian_request_queue})', + ) + + parser.add_argument( + '--librarian-response-queue', + default=default_librarian_response_queue, + help=f'Librarian response queue (default: {default_librarian_response_queue})', + ) + def run(): Processor.launch(default_ident, __doc__) -