API finalised

trustgraph-base/trustgraph/base/async_processor.py

@@ -48,9 +48,10 @@ class AsyncProcessor:
         self.config_handlers = []
 
         self.config_sub_task = self.subscribe(
-            self.config_push_queue, config_subscriber_id,
+            flow = None,
-            schema=ConfigPush, handler=self.on_config_change,
+            queue = self.config_push_queue, subscriber = config_subscriber_id,
-            start_of_messages=True
+            schema = ConfigPush, handler = self.on_config_change,
+            start_of_messages = True
         )
 
         self.running = True
@@ -84,11 +85,12 @@ class AsyncProcessor:
             await asyncio.sleep(2)
 
     def subscribe(
-            self, queue, subscriber, schema, handler, metrics=None,
+            self, flow, queue, subscriber, schema, handler, metrics=None,
             start_of_messages=False
     ):
 
         return Consumer(
+            flow = flow,
             taskgroup = self.taskgroup,
             client = self.client,
             queue = queue,

trustgraph-base/trustgraph/base/consumer.py

@@ -1,4 +1,5 @@
 
+import _pulsar
 import asyncio
 import time
 
@@ -7,13 +8,14 @@ from .. exceptions import TooManyRequests
 class Consumer:
 
     def __init__(
-            self, taskgroup, client, queue, subscriber, schema,
+            self, taskgroup, flow, client, queue, subscriber, schema,
             handler, rate_limit_retry_time = 10,
             rate_limit_timeout = 7200, metrics = None,
             start_of_messages=False,
     ):
 
         self.taskgroup = taskgroup
+        self.flow = flow
         self.client = client
         self.queue = queue
         self.subscriber = subscriber
@@ -28,6 +30,12 @@ class Consumer:
 
         self.metrics = metrics
 
+    def __del__(self):
+        self.running = False
+
+        if hasattr(self, "consumer"):
+            self.consumer.close()
+
     async def stop(self):
 
         self.running = False
@@ -58,7 +66,15 @@ class Consumer:
 
         while self.running:
 
-            msg = await asyncio.to_thread(self.consumer.receive)
+            try:
+                msg = await asyncio.to_thread(
+                    self.consumer.receive,
+                    timeout_millis=2000
+                )
+            except _pulsar.Timeout:
+                continue
+            except Exception as e:
+                raise e
 
             expiry = time.time() + self.rate_limit_timeout
 
@@ -86,7 +102,7 @@ class Consumer:
                     if self.metrics:
 
                         with self.metrics.record_time():
-                            await self.handler(msg, self)
+                            await self.handler(msg, self, self.flow)
 
                     else:
                         await self.handler(msg, self.consumer)
@@ -132,3 +148,4 @@ class Consumer:
 
         if self.metrics:
             self.metrics.state("stopped")
+

trustgraph-base/trustgraph/base/flow_processor.py

@@ -10,6 +10,9 @@ from .. base import AsyncProcessor, Consumer, Producer
 
 from .. base import ProcessorMetrics, ConsumerMetrics, ProducerMetrics
 
+class Flow:
+    pass
+
 class FlowProcessor(AsyncProcessor):
 
     def __init__(self, **params):
@@ -31,16 +34,12 @@ class FlowProcessor(AsyncProcessor):
 
         self.on_config(self.on_configuration)
 
-        self.consumers = {}
+        self.flows = {}
-        self.producers = {}
 
         # These can be overriden by a derived class
         self.consumer_spec = []
         self.producer_spec = []
-
+        self.config_spec = []
-# "input", self.input_schema)
-#                 "input_schema": Document,
-#                 "output_schema": TextDocument,
 
         print("Service initialised.")
 
@@ -50,9 +49,21 @@ class FlowProcessor(AsyncProcessor):
     def register_producer(self, name, schema):
         self.producer_spec.append((name, schema))
 
+    def register_config(self, name):
+        self.config_spec.append((name,))
+
     async def start_flow(self, flow, defn):
 
-        producers = {}
+        flow_obj = Flow()
+        flow_obj.producer = {}
+        flow_obj.consumer = {}
+        flow_obj.config = {}
+
+        for spec in self.config_spec:
+            name = spec[0]
+            flow_obj.config[name] = defn[name]
+
+            setattr(flow_obj, name, defn[name])
 
         for spec in self.producer_spec:
 
@@ -68,9 +79,9 @@ class FlowProcessor(AsyncProcessor):
                 metrics = producer_metrics,
             )
 
-            producers[name] = producer
+            flow_obj.producer[name] = producer
 
-        consumers = {}
+            setattr(flow_obj, name, producer)
 
         for spec in self.consumer_spec:
 
@@ -81,6 +92,7 @@ class FlowProcessor(AsyncProcessor):
             )
 
             consumer = self.subscribe(
+                flow = flow_obj,
                 queue = defn[name],
                 subscriber = self.subscriber,
                 schema = schema,
@@ -92,25 +104,25 @@ class FlowProcessor(AsyncProcessor):
             # metadata
             consumer.id = self.id
             consumer.name = name
-            consumer.flow = flow
+            consumer.flow = flow_obj
-            consumer.q = producers
+            consumer.flow.name = flow
-
-            consumers[name] = consumer
 
             await consumer.start()
 
-        self.consumers[flow] = consumers
+            flow_obj.consumer[name] = consumer
-        self.producers[flow] = producers
+
+            setattr(flow_obj, name, consumer)
+
+        self.flows[flow] = flow_obj
             
         print("Started flow: ", flow)
 
     async def stop_flow(self, flow):
 
-        for c in self.consumers[flow]:
+        for c in self.flows[flow].consumer:
             await c.stop()
 
-        del self.consumers[flow]
+        del self.flows[flow]
-        del self.producers[flow]
 
         print("Stopped flow: ", flow, flush=True)
 
@@ -125,7 +137,7 @@ class FlowProcessor(AsyncProcessor):
             flow_config = json.loads(config["flows"][self.id])
 
             wanted_keys = flow_config.keys()
-            current_keys = self.consumers.keys()
+            current_keys = self.flows.keys()
 
             for key in wanted_keys:
                 if key not in current_keys:

trustgraph-base/trustgraph/base/producer.py

@@ -12,7 +12,8 @@ class Producer:
 
     def __del__(self):
 
-        self.producer.close()
+        if hasattr(self, "producer"):
+            self.producer.close()
 
     async def send(self, msg, properties={}):
 
@@ -21,4 +22,3 @@ class Producer:
 
         self.producer.send(msg, properties)
 
-

trustgraph-base/trustgraph/base/pubsub.py

@@ -51,7 +51,8 @@ class PulsarClient:
             if self.client:
                 self.client.close()
 
-    def subscribe(self, queue, subscriber, schema, start_of_messages=False):
+    def subscribe(self, queue, subscriber, schema,
+                  start_of_messages=False):
 
         if start_of_messages:
             pos = pulsar.InitialPosition.Earliest
@@ -59,10 +60,11 @@ class PulsarClient:
             pos =  pulsar.InitialPosition.Latest
 
         return self.client.subscribe(
-            queue, subscriber,
+            topic = queue,
-            consumer_type=pulsar.ConsumerType.Shared,
+            subscription_name = subscriber,
-            schema=JsonSchema(schema),
+            consumer_type = pulsar.ConsumerType.Shared,
-            initial_position=pos,
+            schema = JsonSchema(schema),
+            initial_position = pos,
         )
 
     def publish(self, queue, schema):

trustgraph-base/trustgraph/base/request_response.py

@@ -1,4 +1,5 @@
 
+import json
 from pulsar.schema import JsonSchema
 
 from .. schema import Error
@@ -8,109 +9,28 @@ from .. log_level import LogLevel
 from .. base import AsyncProcessor, Consumer, Producer
 
 from .. base import ProcessorMetrics, ConsumerMetrics, ProducerMetrics
+from . flow_processor import FlowProcessor
 
-class RequestResponseProcessor(AsyncProcessor):
+class RequestResponseService(FlowProcessor):
 
     def __init__(self, **params):
-        
-        id = params.get("id")
-        self.subscriber = params.get("subscriber")
-        self.request_schema = params.get("request_schema")
-        self.response_schema = params.get("response_schema")
 
-        ProcessorMetrics(id=id).info(
+        super(RequestResponseService, self).__init__(**params)
-            {
-                "subscriber": self.subscriber,
-                "request_schema": self.request_schema.__name__,
-                "response_schema": self.response_schema.__name__,
-            }
-        )
 
-        super(Processor, self).__init__(
+        # These can be overriden by a derived class
-            **params | {
+        self.consumer_spec = [
-                "request_schema": self.request_schema.__name__,
+            ("request", params.get("request_schema"), self.on_request)
-                "response_schema": self.response_schema.__name__,
+        ]
-            }
+        self.producer_spec = [
-        )
+            ("response", params.get("response_schema"))
-
+        ]
-        self.on_config(self.on_configuration)
-
-        self.subs = {}
-        self.pubs = {}
 
         print("Service initialised.")
 
-    async def start_handler(self, flow, defn):
-
-        request_metrics = ConsumerMetrics(self.id, flow)
-        response_metrics = ProducerMetrics(self.id, flow)
-
-        self.subs[flow] = self.subscribe(
-            queue = defn["input"]
-            subscriber = subscriber,
-            schema = self.request_schema,
-            handler = self.on_message,
-            metrics = request_metrics,
-        )
-        
-        self.pubs[flow] = self.publish(
-            queue = defn["output"],
-            schema = self.response_schema,
-            metrics = response_metrics,
-        )
-
-        self.subs[flow].start()
-
-        print("Started flow for", flow)
-
-    async def on_configuration(self, config, version):
-
-        if "flows" not in config: return
-
-        if self.id in config["flows"]:
-
-            wanted_keys = config["flows"][self.id].keys()
-            current_keys = self.subs.keys()
-
-            for key in wanted_keys:
-                if key not in current_keys:
-                    self.start_handler(key, config["flows"][key])
-
-            for key in current_keys:
-                if key not in wanted_keys:
-                    self.stop_handler(key, config["flows"][key])
-
-            print("Handled config update")
-
-    async def start(self):
-
-        await self.request_sub.start()
-        
-    async def on_message(self, msg, consumer):
-
-        try:
-
-            v = msg.value()
-
-            # Sender-produced ID
-            id = msg.properties()["id"]
-
-            print(f"Handling {id}...", flush=True)
-
-            consumer.acknowledge(msg)
-
-        except Exception as e:
-            
-            resp = "error"
-
-            await self.response_pub.send(resp, properties={"id": id})
-
-            consumer.acknowledge(msg)
-
     @staticmethod
-    def add_args(parser):
+    def add_args(parser, default_subscriber):
 
-        AsyncProcessor.add_args(parser)
+        FlowProcessor.add_args(parser)
 
 def run():
 

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

@@ -59,7 +59,7 @@ class Processor(FlowProcessor):
 
         print("Chunker initialised", flush=True)
 
-    async def on_message(self, msg, consumer):
+    async def on_message(self, msg, consumer, flow):
 
         v = msg.value()
         print(f"Chunking {v.metadata.id}...", flush=True)
@@ -81,7 +81,7 @@ class Processor(FlowProcessor):
                 id=consumer.id, flow=consumer.flow
             ).observe(len(chunk.page_content))
 
-            await consumer.q["output"].send(r)
+            await flow.producer["output"].send(r)
 
         print("Done.", flush=True)
 

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

@@ -58,7 +58,7 @@ class Processor(FlowProcessor):
 
         print("Chunker initialised", flush=True)
 
-    async def on_message(self, msg, consumer):
+    async def on_message(self, msg, consumer, flow):
 
         v = msg.value()
         print(f"Chunking {v.metadata.id}...", flush=True)
@@ -80,7 +80,7 @@ class Processor(FlowProcessor):
                 id=consumer.id, flow=consumer.flow
             ).observe(len(chunk.page_content))
 
-            await consumer.q["output"].send(r)
+            await flow.producer["output"].send(r)
 
         print("Done.", flush=True)
 

trustgraph-flow/trustgraph/config/service/service.py

@@ -74,6 +74,7 @@ class Processor(AsyncProcessor):
         )
 
         self.subs = self.subscribe(
+            flow = None,
             queue = request_queue,
             subscriber = subscriber,
             schema = request_schema,
@@ -105,7 +106,7 @@ class Processor(AsyncProcessor):
 
         print("Pushed version ", self.config.version)
         
-    async def on_message(self, msg, consumer):
+    async def on_message(self, msg, consumer, flow):
 
         try:
 

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

@@ -44,7 +44,7 @@ class Processor(FlowProcessor):
 
         print("PDF inited", flush=True)
 
-    async def on_message(self, msg, consumer):
+    async def on_message(self, msg, consumer, flow):
 
         print("PDF message received", flush=True)
 
@@ -71,7 +71,7 @@ class Processor(FlowProcessor):
                         text=page.page_content.encode("utf-8"),
                     )
 
-                    await consumer.q["output"].send(r)
+                    await flow.producer["output"].send(r)
 
         print("Done.", flush=True)