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/decoding/mistral_ocr/processor.py

@@ -17,12 +17,10 @@ from mistralai.models import OCRResponse
 from ... schema import Document, TextDocument, Metadata
 from ... schema import document_ingest_queue, text_ingest_queue
 from ... log_level import LogLevel
-from ... base import ConsumerProducer
+from ... base import InputOutputProcessor
 
-module = ".".join(__name__.split(".")[1:-1])
+module = "ocr"
 
-default_input_queue = document_ingest_queue
-default_output_queue = text_ingest_queue
 default_subscriber = module
 default_api_key = os.getenv("MISTRAL_TOKEN")
 
@@ -71,19 +69,17 @@ def get_combined_markdown(ocr_response: OCRResponse) -> str:
 
     return "\n\n".join(markdowns)
 
-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)
+        id = params.get("id")
         subscriber = params.get("subscriber", default_subscriber)
         api_key = params.get("api_key", default_api_key)
 
         super(Processor, self).__init__(
             **params | {
-                "input_queue": input_queue,
-                "output_queue": output_queue,
+                "id": id,
                 "subscriber": subscriber,
                 "input_schema": Document,
                 "output_schema": TextDocument,
@@ -151,7 +147,7 @@ class Processor(ConsumerProducer):
 
             return markdown
 
-    async def handle(self, msg):
+    async def on_message(self, msg, consumer):
 
         print("PDF message received")
 
@@ -166,17 +162,14 @@ class Processor(ConsumerProducer):
             text=markdown.encode("utf-8"),
         )
 
-        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(
             '-k', '--api-key',

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

@@ -9,39 +9,43 @@ import base64
 from langchain_community.document_loaders import PyPDFLoader
 
 from ... schema import Document, TextDocument, Metadata
-from ... schema import document_ingest_queue, text_ingest_queue
 from ... log_level import LogLevel
-from ... base import ConsumerProducer
+from ... base import FlowProcessor, ConsumerSpec, ProducerSpec
 
-module = ".".join(__name__.split(".")[1:-1])
+default_ident = "pdf-decoder"
 
-default_input_queue = document_ingest_queue
-default_output_queue = text_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)
 
         super(Processor, self).__init__(
             **params | {
-                "input_queue": input_queue,
-                "output_queue": output_queue,
-                "subscriber": subscriber,
-                "input_schema": Document,
-                "output_schema": TextDocument,
+                "id": id,
             }
         )
 
-        print("PDF inited")
+        self.register_specification(
+            ConsumerSpec(
+                name = "input",
+                schema = Document,
+                handler = self.on_message,
+            )
+        )
 
-    async def handle(self, msg):
+        self.register_specification(
+            ProducerSpec(
+                name = "output",
+                schema = TextDocument,
+            )
+        )
 
-        print("PDF message received")
+        print("PDF inited", flush=True)
+
+    async def on_message(self, msg, consumer, flow):
+
+        print("PDF message received", flush=True)
 
         v = msg.value()
 
@@ -59,24 +63,22 @@ class Processor(ConsumerProducer):
 
                 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 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)
 
 def run():
 
-    Processor.launch(module, __doc__)
+    Processor.launch(default_ident, __doc__)