From 57862fba8ec5e525a39e69d0f5a67d177715e408 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Wed, 4 Mar 2026 16:32:13 +0000 Subject: [PATCH] Fixed text loader to use librarian as a source --- .../trustgraph/schema/knowledge/document.py | 3 + .../trustgraph/chunking/recursive/chunker.py | 133 +++++++++++++++++- .../trustgraph/librarian/service.py | 33 +++-- 3 files changed, 155 insertions(+), 14 deletions(-) diff --git a/trustgraph-base/trustgraph/schema/knowledge/document.py b/trustgraph-base/trustgraph/schema/knowledge/document.py index f3996136..5d18b265 100644 --- a/trustgraph-base/trustgraph/schema/knowledge/document.py +++ b/trustgraph-base/trustgraph/schema/knowledge/document.py @@ -22,6 +22,9 @@ class Document: class TextDocument: metadata: Metadata | None = None text: bytes = b"" + # For large document streaming: if document_id is set, the receiver should + # fetch content from librarian instead of using inline text + document_id: str = "" ############################################################################ diff --git a/trustgraph-flow/trustgraph/chunking/recursive/chunker.py b/trustgraph-flow/trustgraph/chunking/recursive/chunker.py index bc6d9cb9..c338c3d8 100755 --- a/trustgraph-flow/trustgraph/chunking/recursive/chunker.py +++ b/trustgraph-flow/trustgraph/chunking/recursive/chunker.py @@ -4,18 +4,28 @@ Simple decoder, accepts text documents on input, outputs chunks from the as text as separate output objects. """ +import asyncio +import base64 import logging +import uuid from langchain_text_splitters import RecursiveCharacterTextSplitter from prometheus_client import Histogram from ... schema import TextDocument, Chunk +from ... schema import LibrarianRequest, LibrarianResponse +from ... schema import librarian_request_queue, librarian_response_queue from ... base import ChunkingService, ConsumerSpec, ProducerSpec +from ... base import Consumer, Producer, ConsumerMetrics, ProducerMetrics # Module logger logger = logging.getLogger(__name__) default_ident = "chunker" +default_librarian_request_queue = librarian_request_queue +default_librarian_response_queue = librarian_response_queue + + class Processor(ChunkingService): def __init__(self, **params): @@ -23,7 +33,7 @@ class Processor(ChunkingService): id = params.get("id", default_ident) chunk_size = params.get("chunk_size", 2000) chunk_overlap = params.get("chunk_overlap", 100) - + super(Processor, self).__init__( **params | { "id": id } ) @@ -62,13 +72,117 @@ class Processor(ChunkingService): ) ) + # 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("Recursive chunker 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): v = msg.value() logger.info(f"Chunking document {v.metadata.id}...") + # Check if we need to fetch content from librarian + if v.document_id and not v.text: + 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, + ) + # Content is base64 encoded + if isinstance(content, str): + content = content.encode('utf-8') + text = base64.b64decode(content).decode("utf-8") + logger.info(f"Fetched {len(text)} characters from librarian") + else: + text = v.text.decode("utf-8") + # Extract chunk parameters from flow (allows runtime override) chunk_size, chunk_overlap = await self.chunk_document( msg, consumer, flow, @@ -90,9 +204,7 @@ class Processor(ChunkingService): is_separator_regex=False, ) - texts = text_splitter.create_documents( - [v.text.decode("utf-8")] - ) + texts = text_splitter.create_documents([text]) for ix, chunk in enumerate(texts): @@ -130,7 +242,18 @@ class Processor(ChunkingService): help=f'Chunk overlap (default: 100)' ) + 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/librarian/service.py b/trustgraph-flow/trustgraph/librarian/service.py index 217da4f1..d374adc2 100755 --- a/trustgraph-flow/trustgraph/librarian/service.py +++ b/trustgraph-flow/trustgraph/librarian/service.py @@ -295,15 +295,30 @@ class Processor(AsyncProcessor): q = flow["interfaces"][kind] if kind == "text-load": - doc = TextDocument( - metadata = Metadata( - id = document.id, - metadata = document.metadata, - user = processing.user, - collection = processing.collection - ), - text = content, - ) + # For large text documents, send document_id for streaming retrieval + if len(content) >= self.STREAMING_THRESHOLD: + logger.info(f"Text document {document.id} is large ({len(content)} bytes), " + f"sending document_id for streaming retrieval") + doc = TextDocument( + metadata = Metadata( + id = document.id, + metadata = document.metadata, + user = processing.user, + collection = processing.collection + ), + document_id = document.id, + text = b"", # Empty, receiver will fetch via librarian + ) + else: + doc = TextDocument( + metadata = Metadata( + id = document.id, + metadata = document.metadata, + user = processing.user, + collection = processing.collection + ), + text = content, + ) schema = TextDocument else: # For large PDF documents, send document_id for streaming retrieval