Chunker ported

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)

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
 

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)
+        id = params.get("id")
-        output_queue = params.get("output_queue", default_output_queue)
+        subscriber = params.get("subscriber")
-        subscriber = params.get("subscriber", default_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(
+        InputOutputProcessor.add_args(parser, default_subscriber)
-            parser, default_input_queue, default_subscriber,
-            default_output_queue,
-        )
 
         parser.add_argument(
             '-z', '--chunk-size',

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):
 

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)