Incremental / large document loading (#659)

Tech spec

BlobStore (trustgraph-flow/trustgraph/librarian/blob_store.py):
- get_stream() - yields document content in chunks for streaming retrieval
- create_multipart_upload() - initializes S3 multipart upload, returns
  upload_id
- upload_part() - uploads a single part, returns etag
- complete_multipart_upload() - finalizes upload with part etags
- abort_multipart_upload() - cancels and cleans up

Cassandra schema (trustgraph-flow/trustgraph/tables/library.py):
- New upload_session table with 24-hour TTL
- Index on user for listing sessions
- Prepared statements for all operations
- Methods: create_upload_session(), get_upload_session(),
  update_upload_session_chunk(), delete_upload_session(),
  list_upload_sessions()

- Schema extended with UploadSession, UploadProgress, and new
  request/response fields
- Librarian methods: begin_upload, upload_chunk, complete_upload,
  abort_upload, get_upload_status, list_uploads
- Service routing for all new operations
- Python SDK with transparent chunked upload:
  - add_document() auto-switches to chunked for files > 10MB
  - Progress callback support (on_progress)
  - get_pending_uploads(), get_upload_status(), abort_upload(),
    resume_upload()

- Document table: Added parent_id and document_type columns with index
- Document schema (knowledge/document.py): Added document_id field for
  streaming retrieval
- Librarian operations:
  - add-child-document for extracted PDF pages
  - list-children to get child documents
  - stream-document for chunked content retrieval
  - Cascade delete removes children when parent is deleted
  - list-documents filters children by default
- PDF decoder (decoding/pdf/pdf_decoder.py): Updated to stream large
  documents from librarian API to temp file
- Librarian service (librarian/service.py): Sends document_id instead of
  content for large PDFs (>2MB)
- Deprecated tools (load_pdf.py, load_text.py): Added deprecation
  warnings directing users to tg-add-library-document +
  tg-start-library-processing

Remove load_pdf and load_text utils

Move chunker/librarian comms to base class

Updating tests

trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py

@@ -2,21 +2,34 @@
 """
 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 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_librarian_request_queue = librarian_request_queue
+default_librarian_response_queue = librarian_response_queue
+
+
 class Processor(FlowProcessor):
 
     def __init__(self, **params):
@@ -44,8 +57,97 @@ 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")
 
+    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 on_message(self, msg, consumer, flow):
 
         logger.debug("PDF message received")
@@ -54,26 +156,53 @@ class Processor(FlowProcessor):
 
         logger.info(f"Decoding PDF {v.metadata.id}...")
 
-        with tempfile.NamedTemporaryFile(delete_on_close=False) as fp:
+        with tempfile.NamedTemporaryFile(delete_on_close=False, suffix='.pdf') as fp:
+            temp_path = fp.name
 
-            fp.write(base64.b64decode(v.data))
-            fp.close()
+            # Check if we should fetch from librarian or use inline data
+            if v.document_id:
+                # Fetch from librarian via Pulsar
+                logger.info(f"Fetching document {v.document_id} from librarian...")
+                fp.close()
 
-            with open(fp.name, mode='rb') as f:
+                content = await self.fetch_document_content(
+                    document_id=v.document_id,
+                    user=v.metadata.user,
+                )
 
-                loader = PyPDFLoader(fp.name)
-                pages = loader.load()
+                # Content is base64 encoded
+                if isinstance(content, str):
+                    content = content.encode('utf-8')
+                decoded_content = base64.b64decode(content)
 
-                for ix, page in enumerate(pages):
+                with open(temp_path, 'wb') as f:
+                    f.write(decoded_content)
 
-                    logger.debug(f"Processing page {ix}")
+                logger.info(f"Fetched {len(decoded_content)} bytes from librarian")
+            else:
+                # Use inline data (backward compatibility)
+                fp.write(base64.b64decode(v.data))
+                fp.close()
 
-                    r = TextDocument(
-                        metadata=v.metadata,
-                        text=page.page_content.encode("utf-8"),
-                    )
+            loader = PyPDFLoader(temp_path)
+            pages = loader.load()
 
-                    await flow("output").send(r)
+            for ix, page in enumerate(pages):
+
+                logger.debug(f"Processing page {ix}")
+
+                r = TextDocument(
+                    metadata=v.metadata,
+                    text=page.page_content.encode("utf-8"),
+                )
+
+                await flow("output").send(r)
+
+        # Clean up temp file
+        try:
+            os.unlink(temp_path)
+        except OSError:
+            pass
 
         logger.debug("PDF decoding complete")
 
@@ -81,7 +210,18 @@ class Processor(FlowProcessor):
     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__)
-