From cfdeadfb34e4480f1ff0f7d563257a73923c6814 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Tue, 15 Apr 2025 23:09:32 +0100 Subject: [PATCH] Chunker ported --- trustgraph-base/trustgraph/base/consumer.py | 2 +- .../trustgraph/base/input_output.py | 6 ++-- .../trustgraph/chunking/recursive/chunker.py | 31 +++++++++---------- .../trustgraph/config/service/service.py | 4 --- .../trustgraph/decoding/pdf/pdf_decoder.py | 4 ++- 5 files changed, 23 insertions(+), 24 deletions(-) diff --git a/trustgraph-base/trustgraph/base/consumer.py b/trustgraph-base/trustgraph/base/consumer.py index 14bd7c5d..3a756a0f 100644 --- a/trustgraph-base/trustgraph/base/consumer.py +++ b/trustgraph-base/trustgraph/base/consumer.py @@ -86,7 +86,7 @@ class Consumer: if self.metrics: with self.metrics.record_time(): - await self.handler(msg, self.consumer) + await self.handler(msg, self) else: await self.handler(msg, self.consumer) diff --git a/trustgraph-base/trustgraph/base/input_output.py b/trustgraph-base/trustgraph/base/input_output.py index 95bef6ea..91af74f7 100644 --- a/trustgraph-base/trustgraph/base/input_output.py +++ b/trustgraph-base/trustgraph/base/input_output.py @@ -64,7 +64,9 @@ class InputOutputProcessor(AsyncProcessor): class Queues: pass - consumer.output = Queues() + consumer.q = Queues() + consumer.id = self.id + consumer.flow = flow producers = {} @@ -82,7 +84,7 @@ class InputOutputProcessor(AsyncProcessor): metrics = producer_metrics, ) - setattr(consumer.output, producer_tag, producer) + setattr(consumer.q, producer_tag, producer) self.consumers[flow] = consumer diff --git a/trustgraph-flow/trustgraph/chunking/recursive/chunker.py b/trustgraph-flow/trustgraph/chunking/recursive/chunker.py index 364fc72c..8d7a3df9 100755 --- a/trustgraph-flow/trustgraph/chunking/recursive/chunker.py +++ b/trustgraph-flow/trustgraph/chunking/recursive/chunker.py @@ -10,28 +10,23 @@ from prometheus_client import Histogram from ... schema import TextDocument, Chunk, Metadata from ... schema import text_ingest_queue, chunk_ingest_queue from ... log_level import LogLevel -from ... base import ConsumerProducer +from ... base import InputOutputProcessor module = "chunker" -default_input_queue = text_ingest_queue -default_output_queue = chunk_ingest_queue default_subscriber = module -class Processor(ConsumerProducer): +class Processor(InputOutputProcessor): def __init__(self, **params): - input_queue = params.get("input_queue", default_input_queue) - output_queue = params.get("output_queue", default_output_queue) - subscriber = params.get("subscriber", default_subscriber) + id = params.get("id") + subscriber = params.get("subscriber") chunk_size = params.get("chunk_size", 2000) chunk_overlap = params.get("chunk_overlap", 100) super(Processor, self).__init__( **params | { - "input_queue": input_queue, - "output_queue": output_queue, "subscriber": subscriber, "input_schema": TextDocument, "output_schema": Chunk, @@ -41,6 +36,7 @@ class Processor(ConsumerProducer): if not hasattr(__class__, "chunk_metric"): __class__.chunk_metric = Histogram( 'chunk_size', 'Chunk size', + ["id", "flow"], buckets=[100, 160, 250, 400, 650, 1000, 1600, 2500, 4000, 6400, 10000, 16000] ) @@ -52,7 +48,9 @@ class Processor(ConsumerProducer): is_separator_regex=False, ) - async def handle(self, msg): + print("Chunker initialised", flush=True) + + async def on_message(self, msg, consumer): v = msg.value() print(f"Chunking {v.metadata.id}...", flush=True) @@ -63,24 +61,25 @@ class Processor(ConsumerProducer): for ix, chunk in enumerate(texts): + print("Chunk", len(chunk.page_content), flush=True) + r = Chunk( metadata=v.metadata, chunk=chunk.page_content.encode("utf-8"), ) - __class__.chunk_metric.observe(len(chunk.page_content)) + __class__.chunk_metric.labels( + id=consumer.id, flow=consumer.flow + ).observe(len(chunk.page_content)) - await self.send(r) + await consumer.q.output.send(r) print("Done.", flush=True) @staticmethod def add_args(parser): - ConsumerProducer.add_args( - parser, default_input_queue, default_subscriber, - default_output_queue, - ) + InputOutputProcessor.add_args(parser, default_subscriber) parser.add_argument( '-z', '--chunk-size', diff --git a/trustgraph-flow/trustgraph/config/service/service.py b/trustgraph-flow/trustgraph/config/service/service.py index ac9ca053..d47d1feb 100644 --- a/trustgraph-flow/trustgraph/config/service/service.py +++ b/trustgraph-flow/trustgraph/config/service/service.py @@ -120,8 +120,6 @@ class Processor(AsyncProcessor): await self.response_pub.send(resp, properties={"id": id}) - consumer.acknowledge(msg) - except Exception as e: resp = ConfigResponse( @@ -134,8 +132,6 @@ class Processor(AsyncProcessor): await self.response_pub.send(resp, properties={"id": id}) - consumer.acknowledge(msg) - @staticmethod def add_args(parser): diff --git a/trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py b/trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py index f03d2f75..689df9c4 100755 --- a/trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py +++ b/trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py @@ -55,12 +55,14 @@ class Processor(InputOutputProcessor): for ix, page in enumerate(pages): + print("page", ix, flush=True) + r = TextDocument( metadata=v.metadata, text=page.page_content.encode("utf-8"), ) - await self.send(r) + await consumer.q.output.send(r) print("Done.", flush=True)