Feature/configure flows (#345)

- Keeps processing in different flows separate so that data can go to different stores / collections etc.
- Potentially supports different processing flows
- Tidies the processing API with common base-classes for e.g. LLMs, and automatic configuration of 'clients' to use the right queue names in a flow

trustgraph-flow/trustgraph/chunking/recursive/chunker.py

@@ -7,40 +7,27 @@ as text as separate output objects.
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 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 ... schema import TextDocument, Chunk
+from ... base import FlowProcessor, ConsumerSpec, ProducerSpec
 
-module = ".".join(__name__.split(".")[1:-1])
+default_ident = "chunker"
 
-default_input_queue = text_ingest_queue
-default_output_queue = chunk_ingest_queue
-default_subscriber = module
-
-class Processor(ConsumerProducer):
+class Processor(FlowProcessor):
 
     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", default_ident)
         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,
-            }
+            **params | { "id": id }
         )
 
         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 +39,24 @@ class Processor(ConsumerProducer):
             is_separator_regex=False,
         )
 
-    async def handle(self, msg):
+        self.register_specification(
+            ConsumerSpec(
+                name = "input",
+                schema = TextDocument,
+                handler = self.on_message,
+            )
+        )
+
+        self.register_specification(
+            ProducerSpec(
+                name = "output",
+                schema = Chunk,
+            )
+        )
+
+        print("Chunker initialised", flush=True)
+
+    async def on_message(self, msg, consumer, flow):
 
         v = msg.value()
         print(f"Chunking {v.metadata.id}...", flush=True)
@@ -63,24 +67,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 flow("output").send(r)
 
         print("Done.", flush=True)
 
     @staticmethod
     def add_args(parser):
 
-        ConsumerProducer.add_args(
-            parser, default_input_queue, default_subscriber,
-            default_output_queue,
-        )
+        FlowProcessor.add_args(parser)
 
         parser.add_argument(
             '-z', '--chunk-size',
@@ -98,5 +103,5 @@ class Processor(ConsumerProducer):
 
 def run():
 
-    Processor.launch(module, __doc__)
+    Processor.launch(default_ident, __doc__)