From 410e44bec7e2d37babdc19ec6bbd869d00702bf3 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Wed, 4 Mar 2026 16:29:13 +0000 Subject: [PATCH] Use Pulsar comms, not gateway --- .../trustgraph/decoding/pdf/pdf_decoder.py | 164 ++++++++++++------ 1 file changed, 109 insertions(+), 55 deletions(-) diff --git a/trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py b/trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py index d60e9b77..67cf4200 100755 --- a/trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py +++ b/trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py @@ -3,26 +3,31 @@ 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 streaming from librarian API +Supports both inline document data and fetching from librarian via Pulsar for large documents. """ +import asyncio import os import tempfile import base64 import logging +import uuid from langchain_community.document_loaders import PyPDFLoader from ... schema import Document, TextDocument, Metadata +from ... schema import LibrarianRequest, LibrarianResponse +from ... schema import librarian_request_queue, librarian_response_queue from ... base import FlowProcessor, ConsumerSpec, ProducerSpec +from ... base import Consumer, Producer, ConsumerMetrics, ProducerMetrics # Module logger logger = logging.getLogger(__name__) default_ident = "pdf-decoder" -# Default API URL for fetching documents from librarian -default_api_url = os.environ.get("TRUSTGRAPH_API_URL", "http://api:8088") +default_librarian_request_queue = librarian_request_queue +default_librarian_response_queue = librarian_response_queue class Processor(FlowProcessor): @@ -37,8 +42,6 @@ class Processor(FlowProcessor): } ) - self.api_url = params.get("api_url", default_api_url) - self.register_specification( ConsumerSpec( name = "input", @@ -54,60 +57,96 @@ class Processor(FlowProcessor): ) ) + # 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 decoder initialized") - def _fetch_document_to_file(self, document_id, user, file_path, chunk_size=1024*1024): + 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 API and stream to file. - - This avoids loading the entire document into memory at once. + Fetch document content from librarian via Pulsar. """ - import requests + request_id = str(uuid.uuid4()) - logger.info(f"Streaming document {document_id} to temp file...") + request = LibrarianRequest( + operation="get-document-content", + document_id=document_id, + user=user, + ) - # Use chunk-based streaming to minimize memory usage - chunk_index = 0 - total_bytes = 0 + # Create future for response + future = asyncio.get_event_loop().create_future() + self.pending_requests[request_id] = future - with open(file_path, 'wb') as f: - while True: - url = f"{self.api_url}/api/v1/librarian" - payload = { - "operation": "stream-document", - "user": user, - "document-id": document_id, - "chunk-index": chunk_index, - "chunk-size": chunk_size, - } + try: + # Send request + await self.librarian_request_producer.send( + request, properties={"id": request_id} + ) - try: - response = requests.post(url, json=payload, timeout=60) - response.raise_for_status() - data = response.json() - except requests.RequestException as e: - logger.error(f"Failed to fetch chunk {chunk_index}: {e}") - raise + # Wait for response + response = await asyncio.wait_for(future, timeout=timeout) - if "error" in data and data["error"]: - raise RuntimeError(f"API error: {data['error']}") + if response.error: + raise RuntimeError( + f"Librarian error: {response.error.type}: {response.error.message}" + ) - content_b64 = data.get("content", "") - if not content_b64: - break + return response.content - chunk_data = base64.b64decode(content_b64) - f.write(chunk_data) - total_bytes += len(chunk_data) - - total_chunks = data.get("total-chunks", 1) - if chunk_index >= total_chunks - 1: - break - - chunk_index += 1 - - logger.info(f"Downloaded {total_bytes} bytes to temp file") - return total_bytes + except asyncio.TimeoutError: + self.pending_requests.pop(request_id, None) + raise RuntimeError(f"Timeout fetching document {document_id}") async def on_message(self, msg, consumer, flow): @@ -122,14 +161,24 @@ class Processor(FlowProcessor): # Check if we should fetch from librarian or use inline data if v.document_id: - # Stream from librarian API to temp file + # Fetch from librarian via Pulsar logger.info(f"Fetching document {v.document_id} from librarian...") fp.close() - self._fetch_document_to_file( + + content = await self.fetch_document_content( document_id=v.document_id, user=v.metadata.user, - file_path=temp_path, ) + + # Content is base64 encoded + if isinstance(content, str): + content = content.encode('utf-8') + decoded_content = base64.b64decode(content) + + with open(temp_path, 'wb') as f: + f.write(decoded_content) + + logger.info(f"Fetched {len(decoded_content)} bytes from librarian") else: # Use inline data (backward compatibility) fp.write(base64.b64decode(v.data)) @@ -162,12 +211,17 @@ class Processor(FlowProcessor): FlowProcessor.add_args(parser) parser.add_argument( - '--api-url', - default=default_api_url, - help=f'TrustGraph API URL for document streaming (default: {default_api_url})', + '--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__) -