Prompt template working

Makefile

@@ -65,8 +65,8 @@ some-containers:
 	    -t ${CONTAINER_BASE}/trustgraph-base:${VERSION} .
 	${DOCKER} build -f containers/Containerfile.flow \
 	    -t ${CONTAINER_BASE}/trustgraph-flow:${VERSION} .
-	${DOCKER} build -f containers/Containerfile.vertexai \
+#	${DOCKER} build -f containers/Containerfile.vertexai \
-	    -t ${CONTAINER_BASE}/trustgraph-vertexai:${VERSION} .
+#	    -t ${CONTAINER_BASE}/trustgraph-vertexai:${VERSION} .
 
 basic-containers: update-package-versions
 	${DOCKER} build -f containers/Containerfile.base \

tests/test-llm

@@ -11,11 +11,9 @@ llm = LlmClient(
 )
 
 system = "You are a lovely assistant."
-prompt="Write a funny limerick about a llama"
+prompt="what is 2 + 2 == 5"
 
 resp = llm.request(system, prompt)
 
 print(resp)
 
-
-

tests/test-prompt-question

@@ -3,13 +3,18 @@
 import pulsar
 from trustgraph.clients.prompt_client import PromptClient
 
-p = PromptClient(pulsar_host="pulsar://localhost:6650")
+p = PromptClient(
+    pulsar_host="pulsar://localhost:6650",
+    input_queue="non-persistent://tg/request/prompt:default",
+    output_queue="non-persistent://tg/response/prompt:default",
+    subscriber="test1",
+)
 
 question = """What is the square root of 16?"""
 
 resp = p.request(
     id="question",
-    terms = {
+    variables = {
         "question": question
     }
 )

trustgraph-base/trustgraph/base/__init__.py

@@ -6,7 +6,8 @@ from . producer import Producer
 from . publisher import Publisher
 from . subscriber import Subscriber
 from . metrics import ProcessorMetrics, ConsumerMetrics, ProducerMetrics
-from . flow_processor import FlowProcessor
+from . flow_processor import FlowProcessor, SubscriberSpec, ProducerSpec
+from . flow_processor import ConsumerSpec, SettingSpec
 from . request_response import RequestResponseService
 
 

trustgraph-base/trustgraph/base/consumer.py

@@ -30,6 +30,8 @@ class Consumer:
 
         self.metrics = metrics
 
+        self.consumer = None
+
     def __del__(self):
         self.running = False
 
@@ -39,30 +41,63 @@ class Consumer:
     async def stop(self):
 
         self.running = False
+        await self.task
 
     async def start(self):
 
         self.running = True
 
-        self.consumer = self.client.subscribe(
-            self.queue, self.subscriber, self.schema,
-            start_of_messages = self.start_of_messages,
-        )
-
-        print("Subscribed.", flush=True)
-
         # Puts it in the stopped state, the run thread should set running
         if self.metrics:
             self.metrics.state("stopped")
 
         self.task = self.taskgroup.create_task(self.run())
 
-        print("Subscriber started", flush=True)
-
     async def run(self):
 
         if self.metrics:
-            self.metrics.state("running")
+            self.metrics.state("stopped")
+
+        while self.running:
+
+            try:
+
+                print(self.queue, "subscribing...", flush=True)
+
+                self.consumer = await asyncio.to_thread(
+                    self.client.subscribe,
+                    queue = self.queue, subscriber = self.subscriber,
+                    schema = self.schema,
+                    start_of_messages = self.start_of_messages
+                )
+
+            except Exception as e:
+
+                print("Exception:", e, flush=True)
+                await asyncio.sleep(2)
+                continue
+
+            print(self.queue, "subscribed", flush=True)
+
+            if self.metrics:
+                self.metrics.state("running")
+
+            try:
+
+                await self.consume()
+
+                if self.metrics:
+                    self.metrics.state("stopped")
+
+            except Exception as e:
+
+                print("Exception:", e, flush=True)
+                self.consumer.close()
+                self.consumer = None
+                await asyncio.sleep(2)
+                continue
+
+    async def consume(self):
 
         while self.running:
 
@@ -79,7 +114,7 @@ class Consumer:
             expiry = time.time() + self.rate_limit_timeout
 
             # This loop is for retry on rate-limit / resource limits
-            while True:
+            while self.running:
 
                 if time.time() > expiry:
 
@@ -122,7 +157,6 @@ class Consumer:
 
                     print("TooManyRequests: will retry...", flush=True)
 
-                    # FIXME
                     if self.metrics:
                         self.metrics.rate_limit()
 
@@ -145,7 +179,3 @@ class Consumer:
 
                     # Break out of retry loop, processes next message
                     break
-
-        if self.metrics:
-            self.metrics.state("stopped")
-

trustgraph-base/trustgraph/base/consumer.py.OLD

@@ -1,187 +0,0 @@
-
-import asyncio
-from pulsar.schema import JsonSchema
-import pulsar
-from prometheus_client import Histogram, Info, Counter, Enum
-import time
-
-from . base_processor import BaseProcessor
-from .. exceptions import TooManyRequests
-
-default_rate_limit_retry = 10
-default_rate_limit_timeout = 7200
-
-class Consumer(BaseProcessor):
-
-    def __init__(self, **params):
-
-        if not hasattr(__class__, "state_metric"):
-            __class__.state_metric = Enum(
-                'processor_state', 'Processor state',
-                states=['starting', 'running', 'stopped']
-            )
-            __class__.state_metric.state('starting')
-
-        __class__.state_metric.state('starting')
-
-        super(Consumer, self).__init__(**params)
-
-        self.subscriber = params.get("subscriber")
-        self.input_schema = params.get("input_schema")
-
-        self.rate_limit_retry = params.get(
-            "rate_limit_retry", default_rate_limit_retry
-        )
-        self.rate_limit_timeout = params.get(
-            "rate_limit_timeout", default_rate_limit_timeout
-        )
-
-        if self.input_schema == None:
-            raise RuntimeError("input_schema must be specified")
-
-        if not hasattr(__class__, "request_metric"):
-            __class__.request_metric = Histogram(
-                'request_latency', 'Request latency (seconds)',
-                ["flow"]
-            )
-
-        if not hasattr(__class__, "pubsub_metric"):
-            __class__.pubsub_metric = Info(
-                'pubsub', 'Pub/sub configuration',
-                ["flow"]
-            )
-
-        if not hasattr(__class__, "processing_metric"):
-            __class__.processing_metric = Counter(
-                'processing_count', 'Processing count',
-                ["flow", "status"]
-            )
-
-        if not hasattr(__class__, "rate_limit_metric"):
-            __class__.rate_limit_metric = Counter(
-                'rate_limit_count', 'Rate limit event count',
-                ["flow"]
-            )
-
-        # __class__.pubsub_metric.info({
-        #     "input_queue": self.input_queue,
-        #     "subscriber": self.subscriber,
-        #     "input_schema": self.input_schema.__name__,
-        #     "rate_limit_retry": str(self.rate_limit_retry),
-        #     "rate_limit_timeout": str(self.rate_limit_timeout),
-        # })
-
-        # self.consumer = self.client.subscribe(
-        #     self.input_queue, self.subscriber,
-        #     consumer_type=pulsar.ConsumerType.Shared,
-        #     schema=JsonSchema(self.input_schema),
-        # )
-
-        self.config_queues = {}
-        self.config_handlers = [self.on_config_pubsub]
-
-        print("Initialised consumer.", flush=True)
-
-    async def on_config_pubsub(self, version, config):
-
-        print("Configuring using config", version)
-
-        flows = config.get("flows", {}).get(self.ident, {})
-
-        print(flows)
-
-        print("Configuration completed OK")
-
-    async def run(self):
-
-        __class__.state_metric.state('running')
-
-        while True:
-
-            await asyncio.sleep(1)
-
-            continue
-
-            msg = await asyncio.to_thread(self.consumer.receive)
-
-            expiry = time.time() + self.rate_limit_timeout
-
-            # This loop is for retry on rate-limit / resource limits
-            while True:
-
-                if time.time() > expiry:
-
-                    print("Gave up waiting for rate-limit retry", flush=True)
-
-                    # Message failed to be processed, this causes it to
-                    # be retried
-                    self.consumer.negative_acknowledge(msg)
-
-                    __class__.processing_metric.labels(status="error").inc()
-
-                    # Break out of retry loop, processes next message
-                    break
-
-                try:
-
-                    with __class__.request_metric.time():
-                        await self.handle(msg)
-
-                    # Acknowledge successful processing of the message
-                    self.consumer.acknowledge(msg)
-
-                    __class__.processing_metric.labels(status="success").inc()
-
-                    # Break out of retry loop
-                    break
-
-                except TooManyRequests:
-
-                    print("TooManyRequests: will retry...", flush=True)
-
-                    __class__.rate_limit_metric.inc()
-
-                    # Sleep
-                    time.sleep(self.rate_limit_retry)
-
-                    # Contine from retry loop, just causes a reprocessing
-                    continue
-                
-                except Exception as e:
-
-                    print("Exception:", e, flush=True)
-
-                    # Message failed to be processed, this causes it to
-                    # be retried
-                    self.consumer.negative_acknowledge(msg)
-
-                    __class__.processing_metric.labels(status="error").inc()
-
-                    # Break out of retry loop, processes next message
-                    break
-
-    @staticmethod
-    def add_args(parser, default_input_queue, default_subscriber):
-
-        BaseProcessor.add_args(parser)
-
-        parser.add_argument(
-            '-s', '--subscriber',
-            default=default_subscriber,
-            help=f'Queue subscriber name (default: {default_subscriber})'
-        )
-
-        parser.add_argument(
-            '--rate-limit-retry',
-            type=int,
-            default=default_rate_limit_retry,
-            help=f'Rate limit retry (default: {default_rate_limit_retry})'
-        )
-
-        parser.add_argument(
-            '--rate-limit-timeout',
-            type=int,
-            default=default_rate_limit_timeout,
-            help=f'Rate limit timeout (default: {default_rate_limit_timeout})'
-        )
-

trustgraph-base/trustgraph/base/consumer_producer.py

@@ -1,62 +0,0 @@
-
-from pulsar.schema import JsonSchema
-import pulsar
-from prometheus_client import Histogram, Info, Counter, Enum
-import time
-
-from . consumer import Consumer
-from .. exceptions import TooManyRequests
-
-class ConsumerProducer(Consumer):
-
-    def __init__(self, **params):
-
-        super(ConsumerProducer, self).__init__(**params)
-
-        self.output_queue = params.get("output_queue")
-        self.output_schema = params.get("output_schema")
-
-        if not hasattr(__class__, "output_metric"):
-            __class__.output_metric = Counter(
-                'output_count', 'Output items created'
-            )
-
-        __class__.pubsub_metric.info({
-            "input_queue": self.input_queue,
-            "output_queue": self.output_queue,
-            "subscriber": self.subscriber,
-            "input_schema": self.input_schema.__name__,
-            "output_schema": self.output_schema.__name__,
-            "rate_limit_retry": str(self.rate_limit_retry),
-            "rate_limit_timeout": str(self.rate_limit_timeout),
-        })
-
-        if self.output_schema == None:
-            raise RuntimeError("output_schema must be specified")
-
-        self.producer = self.client.create_producer(
-            topic=self.output_queue,
-            schema=JsonSchema(self.output_schema),
-            chunking_enabled=True,
-        )
-
-        print("Initialised consumer/producer.")
-
-    async def send(self, msg, properties={}):
-        self.producer.send(msg, properties)
-        __class__.output_metric.inc()
-
-    @staticmethod
-    def add_args(
-            parser, default_input_queue, default_subscriber,
-            default_output_queue,
-    ):
-
-        Consumer.add_args(parser, default_input_queue, default_subscriber)
-
-        parser.add_argument(
-            '-o', '--output-queue',
-            default=default_output_queue,
-            help=f'Output queue (default: {default_output_queue})'
-        )
-

trustgraph-base/trustgraph/base/flow_processor.py

@@ -11,70 +11,115 @@ from .. schema import Error
 from .. schema import config_request_queue, config_response_queue
 from .. schema import config_push_queue
 from .. log_level import LogLevel
-from .. base import AsyncProcessor, Consumer, Producer
+from .. base import AsyncProcessor, Consumer, Producer, Subscriber
 from . metrics import ConsumerMetrics, ProducerMetrics
 
+class Spec:
+    pass
+
+class SubscriberSpec(Spec):
+    def __init__(self, name, schema):
+        self.name = name
+        self.schema = schema
+
+    def add(self, flow, processor, definition):
+
+        subscriber_metrics = ConsumerMetrics(
+            flow.id, f"{flow.name}-{self.name}"
+        )
+
+        subscriber = Subscriber(
+            pulsar_client = processor.client.client,
+            topic = definition[self.name],
+            subscription = flow.id,
+            consumer_name = flow.id,
+            schema = self.schema,
+        )
+
+        # Put it in the consumer map, does that work?
+        # It means it gets start/stop call.
+        flow.consumer[self.name] = subscriber
+
+        if not hasattr(flow, self.name):
+            setattr(flow, self.name, subscriber)
+
+class ConsumerSpec(Spec):
+    def __init__(self, name, schema, handler):
+        self.name = name
+        self.schema = schema
+        self.handler = handler
+
+    def add(self, flow, processor, definition):
+
+        consumer_metrics = ConsumerMetrics(
+            flow.id, f"{flow.name}-{self.name}"
+        )
+
+        consumer = processor.subscribe(
+            flow = flow,
+            queue = definition[self.name],
+            subscriber = flow.id,
+            schema = self.schema,
+            handler = self.handler,
+            metrics = consumer_metrics,
+        )
+
+        # Consumer handle gets access to producers and other
+        # metadata
+        consumer.id = flow.id
+        consumer.name = self.name
+        consumer.flow = flow
+
+        flow.consumer[self.name] = consumer
+
+        if not hasattr(flow, self.name):
+            setattr(flow, self.name, consumer)
+
+class ProducerSpec(Spec):
+    def __init__(self, name, schema):
+        self.name = name
+        self.schema = schema
+
+    def add(self, flow, processor, definition):
+
+        producer_metrics = ProducerMetrics(
+            flow.id, f"{flow.name}-{self.name}"
+        )
+
+        producer = processor.publish(
+            queue = definition[self.name],
+            schema = self.schema,
+            metrics = producer_metrics,
+        )
+
+        flow.producer[self.name] = producer
+
+        if not hasattr(self, self.name):
+            setattr(flow, self.name, producer)
+
+class SettingSpec(Spec):
+    def __init__(self, name):
+        self.name = name
+
+    def add(self, flow, processor, definition):
+
+        flow.config[self.name] = definition[self.name]
+
+        if not hasattr(flow, self.name):
+            setattr(flow, self.name, definition[self.name])
+
 class Flow:
     def __init__(self, id, flow, processor, defn):
 
-        self.producer = {}
+        self.id = id
-        self.consumer = {}
-        self.config = {}
         self.name = flow
 
-        for spec in processor.config_spec:
+        self.producer = {}
-            name = spec
+        self.consumer = {}
-            self.config[name] = defn[name]
+        self.setting = {}
 
-            if not hasattr(self, name):
+        for spec in processor.specifications:
-                setattr(self, name, defn[name])
+            spec.add(self, processor, defn)
-
-        for spec in processor.producer_spec:
-
-            name, schema = spec
-
-            producer_metrics = ProducerMetrics(
-                id, f"{flow}-{name}"
-            )
-
-            producer = processor.publish(
-                queue = defn[name],
-                schema = schema,
-                metrics = producer_metrics,
-            )
-
-            self.producer[name] = producer
-
-            if not hasattr(self, name):
-                setattr(self, name, producer)
-
-        for spec in processor.consumer_spec:
-
-            name, schema, handler = spec
-
-            consumer_metrics = ConsumerMetrics(
-                id, f"{flow}-{name}"
-            )
-
-            consumer = processor.subscribe(
-                flow = self,
-                queue = defn[name],
-                subscriber = id,
-                schema = schema,
-                handler = handler,
-                metrics = consumer_metrics,
-            )
-
-            # Consumer handle gets access to producers and other
-            # metadata
-            consumer.id = id
-            consumer.name = name
-            consumer.flow = self
-
-            self.consumer[name] = consumer
-
-            if not hasattr(self, name):
-                setattr(self, name, consumer)
 
     async def start(self):
         for c in self.consumer.values():
@@ -101,29 +146,26 @@ class FlowProcessor(AsyncProcessor):
 
         # These can be overriden by a derived class:
 
-        # Consumer specification, array of ("name", SchemaType, handler)
+        # Array of specifications: ConsumerSpec, ProducerSpec, SettingSpec
-        self.consumer_spec = []
+        self.specifications = []
-
-        # Producer specification, array of ("name", SchemaType)
-        self.producer_spec = []
-
-        # Configuration specification, collects some flow variables from
-        # config, array of "name"
-        self.config_spec = []
 
         print("Service initialised.")
 
     # Register a new consumer name
     def register_consumer(self, name, schema, handler):
-        self.consumer_spec.append((name, schema, handler))
+        self.specifications.append(ConsumerSpec(name, schema, handler))
 
     # Register a producer name
     def register_producer(self, name, schema):
-        self.producer_spec.append((name, schema))
+        self.specifications.append(ProducerSpec(name, schema))
 
     # Register a configuration variable
     def register_config(self, name):
-        self.config_spec.append(name)
+        self.specifications.append(SettingSpec(name))
+
+    # Register a configuration variable
+    def register_specification(self, spec):
+        self.specifications.append(spec)
 
     # Start processing for a new flow
     async def start_flow(self, flow, defn):

trustgraph-base/trustgraph/base/producer.py

@@ -1,4 +1,6 @@
 
+import asyncio
+
 class Producer:
 
     def __init__(self, client, queue, schema, metrics=None):
@@ -6,19 +8,58 @@ class Producer:
         self.queue = queue
         self.schema = schema
 
-        self.producer = self.client.publish(self.queue, self.schema)
-
         self.metrics = metrics
 
+        self.running = True
+        self.producer = None
+
     def __del__(self):
 
+        self.running = False
+
         if hasattr(self, "producer"):
-            self.producer.close()
+            if self.producer:
+                self.producer.close()
+
+    async def start(self):
+        self.running = True
+
+    async def stop(self):
+        self.running = False
 
     async def send(self, msg, properties={}):
 
-        if self.metrics:
+        if not self.running: return
-            self.metrics.inc()
 
-        self.producer.send(msg, properties)
+        while self.running and self.producer is None:
+
+            try:
+                print("Connect publisher to", self.queue, "...", flush=True)
+                self.producer = self.client.publish(self.queue, self.schema)
+                print("Connected to", self.queue, flush=True)
+            except Exception as e:
+                print("Exception:", e, flush=True)
+                await asyncio.sleep(2)
+
+            if not self.running: break
+
+        while self.running:
+
+            try:
+
+                await asyncio.to_thread(
+                    self.producer.send,
+                    msg, properties
+                )
+
+                if self.metrics:
+                    self.metrics.inc()
+
+                # Delivery success, break out of loop
+                break
+
+            except Exception as e:
+                print("Exception:", e, flush=True)
+                self.producer.close()
+                self.producer = None
 

trustgraph-base/trustgraph/base/producer.py.OLD

@@ -1,57 +0,0 @@
-
-from pulsar.schema import JsonSchema
-from prometheus_client import Info, Counter
-
-from . base_processor import BaseProcessor
-
-class Producer(BaseProcessor):
-
-    def __init__(self, **params):
-
-        output_queue = params.get("output_queue")
-        output_schema = params.get("output_schema")
-
-        if not hasattr(__class__, "output_metric"):
-            __class__.output_metric = Counter(
-                'output_count', 'Output items created'
-            )
-
-        if not hasattr(__class__, "pubsub_metric"):
-            __class__.pubsub_metric = Info(
-                'pubsub', 'Pub/sub configuration'
-            )
-
-        __class__.pubsub_metric.info({
-            "output_queue": output_queue,
-            "output_schema": output_schema.__name__,
-        })
-
-        super(Producer, self).__init__(**params)
-
-        if output_schema == None:
-            raise RuntimeError("output_schema must be specified")
-
-        self.producer = self.client.create_producer(
-            topic=output_queue,
-            schema=JsonSchema(output_schema),
-            chunking_enabled=True,
-        )
-
-    async def send(self, msg, properties={}):
-        self.producer.send(msg, properties)
-        __class__.output_metric.inc()
-
-    @staticmethod
-    def add_args(
-            parser, default_input_queue, default_subscriber,
-            default_output_queue,
-    ):
-
-        BaseProcessor.add_args(parser)
-
-        parser.add_argument(
-            '-o', '--output-queue',
-            default=default_output_queue,
-            help=f'Output queue (default: {default_output_queue})'
-        )
-

trustgraph-base/trustgraph/base/publisher.py

@@ -1,8 +1,7 @@
 
-import queue
+import asyncio
 import time
 import pulsar
-import threading
 
 class Publisher:
 
@@ -11,22 +10,22 @@ class Publisher:
         self.client = pulsar_client
         self.topic = topic
         self.schema = schema
-        self.q = queue.Queue(maxsize=max_size)
+        self.q = asyncio.Queue(maxsize=max_size)
         self.chunking_enabled = chunking_enabled
         self.running = True
 
-    def start(self):
+    async def start(self):
-        self.task = threading.Thread(target=self.run)
+        self.task = asyncio.create_task(self.run())
-        self.task.start()
+        await self.task.start()
 
-    def stop(self):
+    async def stop(self):
         self.running = False
 
-    def join(self):
+    async def join(self):
-        self.stop()
+        await self.stop()
-        self.task.join()
+        await self.task
 
-    def run(self):
+    async def run(self):
 
         while self.running:
 
@@ -40,8 +39,13 @@ class Publisher:
                 while self.running:
 
                     try:
-                        id, item = self.q.get(timeout=0.5)
+                        id, item = asyncio.wait_for(
-                    except queue.Empty:
+                            self.q.get(),
+                            timeout=0.5
+                        )
+                    except asyncio.TimeoutError:
+                        continue
+                    except asyncio.QueueEmpty:
                         continue
 
                     if id:
@@ -55,7 +59,6 @@ class Publisher:
             # If handler drops out, sleep a retry
             time.sleep(2)
 
-    def send(self, id, msg):
+    async def send(self, id, item):
-        self.q.put((id, msg))
+        await self.q.put((id, item))
 
-        

trustgraph-base/trustgraph/base/subscriber.py

@@ -1,7 +1,7 @@
 
-import queue
+from pulsar.schema import JsonSchema
-import pulsar
+import asyncio
-import threading
+import _pulsar
 import time
 
 class Subscriber:
@@ -16,41 +16,50 @@ class Subscriber:
         self.q = {}
         self.full = {}
         self.max_size = max_size
-        self.lock = threading.Lock()
+        self.lock = asyncio.Lock()
         self.running = True
 
-    def start(self):
+    async def __del__(self):
-        self.task = threading.Thread(target=self.run)
-        self.task.start()
-
-    def stop(self):
         self.running = False
 
-    def join(self):
+    async def start(self):
-        self.task.join()
+        self.task = asyncio.create_task(self.run())
 
-    def run(self):
+    async def stop(self):
+        self.running = False
 
-        print("SUBSCRIBER LOOP KICKS IN", flush=True)
+    async def join(self):
+        await self.stop()
+        await self.task
+
+    async def run(self):
 
         while self.running:
 
             try:
 
                 consumer = self.client.subscribe(
-                    topic=self.topic,
+                    topic = self.topic,
-                    subscription_name=self.subscription,
+                    subscription_name = self.subscription,
-                    consumer_name=self.consumer_name,
+                    consumer_name = self.consumer_name,
-                    schema=self.schema,
+                    schema = JsonSchema(self.schema),
                 )
 
-                print("SUBSCRIBER RUNNING...", flush=True)
+                print("Subscriber running...", flush=True)
 
                 while self.running:
 
-                    msg = consumer.receive()
+                    try:
-
+                        msg = await asyncio.to_thread(
-                    print("GOT MESSAGE...", flush=True)
+                            consumer.receive,
+                            timeout_millis=2000
+                        )
+                    except _pulsar.Timeout:
+                        continue
+                    except Exception as e:
+                        print("Exception:", e, flush=True)
+                        print(type(e))
+                        raise e
 
                     # Acknowledge successful reception of the message
                     consumer.acknowledge(msg)
@@ -62,57 +71,68 @@ class Subscriber:
 
                     value = msg.value()
 
-                    with self.lock:
+                    async with self.lock:
+
+                        # FIXME: Hard-coded timeouts
 
                         if id in self.q:
+
                             try:
                                 # FIXME: Timeout means data goes missing
-                                self.q[id].put(value, timeout=0.5)
+                                await asyncio.wait_for(
-                            except:
+                                    self.q[id].put(value),
-                                pass
+                                    timeout=2
+                                )
+                            except Exception as e:
+                                print("Q Put:", e, flush=True)
 
                         for q in self.full.values():
                             try:
                                 # FIXME: Timeout means data goes missing
-                                q.put(value, timeout=0.5)
+                                await asyncio.wait_for(
-                            except:
+                                    q.put(value),
-                                pass
+                                    timeout=2
+                                )
+                            except Exception as e:
+                                print("Q Put:", e, flush=True)
 
             except Exception as e:
                 print("Subscriber exception:", e, flush=True)
+
+            consumer.close()
          
             # If handler drops out, sleep a retry
             time.sleep(2)
 
-    def subscribe(self, id):
+    async def subscribe(self, id):
 
-        with self.lock:
+        async with self.lock:
 
-            q = queue.Queue(maxsize=self.max_size)
+            q = asyncio.Queue(maxsize=self.max_size)
             self.q[id] = q
 
         return q
 
-    def unsubscribe(self, id):
+    async def unsubscribe(self, id):
         
-        with self.lock:
+        async with self.lock:
 
             if id in self.q:
 #                self.q[id].shutdown(immediate=True)
                 del self.q[id]
     
-    def subscribe_all(self, id):
+    async def subscribe_all(self, id):
 
-        with self.lock:
+        async with self.lock:
 
-            q = queue.Queue(maxsize=self.max_size)
+            q = asyncio.Queue(maxsize=self.max_size)
             self.full[id] = q
 
         return q
 
-    def unsubscribe_all(self, id):
+    async def unsubscribe_all(self, id):
 
-        with self.lock:
+        async with self.lock:
 
             if id in self.full:
 #                self.full[id].shutdown(immediate=True)

trustgraph-flow/trustgraph/model/prompt/template/prompt_manager.py

@@ -4,8 +4,6 @@ import json
 from jsonschema import validate
 import re
 
-from trustgraph.clients.llm_client import LlmClient
-
 class PromptConfiguration:
     def __init__(self, system_template, global_terms={}, prompts={}):
         self.system_template = system_template
@@ -21,8 +19,7 @@ class Prompt:
 
 class PromptManager:
 
-    def __init__(self, llm, config):
+    def __init__(self, config):
-        self.llm = llm
         self.config = config
         self.terms = config.global_terms
 
@@ -54,7 +51,7 @@ class PromptManager:
 
         return json.loads(json_str)
 
-    def invoke(self, id, input):
+    async def invoke(self, id, input, llm):
 
         if id not in self.prompts:
             raise RuntimeError("ID invalid")
@@ -68,9 +65,7 @@ class PromptManager:
             "prompt": self.templates[id].render(terms)
         }
 
-        resp = self.llm.request(**prompt)
+        resp = await llm(**prompt)
-
-        print(resp, flush=True)
 
         if resp_type == "text":
             return resp
@@ -83,10 +78,9 @@ class PromptManager:
         except:
             raise RuntimeError("JSON parse fail")
 
-        print(obj, flush=True)
         if self.prompts[id].schema:
             try:
-                print(self.prompts[id].schema)
+                print(self.prompts[id].schema, flush=True)
                 validate(instance=obj, schema=self.prompts[id].schema)
             except Exception as e:
                 raise RuntimeError(f"Schema validation fail: {e}")

trustgraph-flow/trustgraph/model/prompt/template/service.py

@@ -3,6 +3,7 @@
 Language service abstracts prompt engineering from LLM.
 """
 
+import asyncio
 import json
 import re
 
@@ -10,63 +11,73 @@ from .... schema import Definition, Relationship, Triple
 from .... schema import Topic
 from .... schema import PromptRequest, PromptResponse, Error
 from .... schema import TextCompletionRequest, TextCompletionResponse
-from .... schema import text_completion_request_queue
+from .... base import FlowProcessor, SubscriberSpec, ConsumerSpec
-from .... schema import text_completion_response_queue
+from .... base import ProducerSpec
-from .... schema import prompt_request_queue, prompt_response_queue
-from .... clients.llm_client import LlmClient
-from .... base import RequestResponseService
 
 from . prompt_manager import PromptConfiguration, Prompt, PromptManager
 
-module = "prompt"
+default_ident = "prompt"
-default_subscriber = module
 
-class Processor(RequestResponseService):
+class Processor(FlowProcessor):
 
     def __init__(self, **params):
 
         id = params.get("id")
-        subscriber = params.get("subscriber", default_subscriber)
 
+        # Config key for prompts
         self.config_key = params.get("config_type", "prompt")
 
         super(Processor, self).__init__(
             **params | {
                 "id": id,
-                "subscriber": subscriber,
-                "request_schema": PromptRequest,
-                "response_schema": PromptResponse,
             }
         )
 
-        self.llm = LlmClient(
+        # self.llm = LlmClient(
-            subscriber=subscriber,
+        #     subscriber=subscriber,
-            input_queue=tc_request_queue,
+        #     input_queue=tc_request_queue,
-            output_queue=tc_response_queue,
+        #     output_queue=tc_response_queue,
-            pulsar_host = self.pulsar_host,
+        #     pulsar_host = self.pulsar_host,
-            pulsar_api_key=self.pulsar_api_key,
+        #     pulsar_api_key=self.pulsar_api_key,
+        # )
+
+        self.register_specification(
+            ConsumerSpec(
+                name = "request",
+                schema = PromptRequest,
+                handler = self.on_request
+            )
         )
 
-        # System prompt hack
+        self.register_specification(
-        class Llm:
+            ProducerSpec(
-            def __init__(self, llm):
+                name = "text-completion-request",
-                self.llm = llm
+                schema = TextCompletionRequest
-            def request(self, system, prompt):
+            )
-                print(system)
+        )
-                print(prompt, flush=True)
-                return self.llm.request(system, prompt)
 
-        self.llm = Llm(self.llm)
+        self.register_specification(
+            SubscriberSpec(
+                name = "text-completion-response",
+                schema = TextCompletionResponse,
+            )
+        )
 
-        self.config_handlers.append(self.on_config)
+        self.register_specification(
+            ProducerSpec(
+                name = "response",
+                schema = PromptResponse
+            )
+        )
+
+        self.register_config_handler(self.on_prompt_config)
 
         # Null configuration, should reload quickly
         self.manager = PromptManager(
-            llm = self.llm,
             config = PromptConfiguration("", {}, {})
         )
 
-    async def on_config(self, version, config):
+    async def on_prompt_config(self, config, version):
 
         print("Loading configuration version", version)
 
@@ -100,7 +111,6 @@ class Processor(RequestResponseService):
                 )
 
             self.manager = PromptManager(
-                self.llm,
                 PromptConfiguration(
                     system,
                     {},
@@ -115,7 +125,7 @@ class Processor(RequestResponseService):
             print("Exception:", e, flush=True)
             print("Configuration reload failed", flush=True)
 
-    async def handle(self, msg):
+    async def on_request(self, msg, consumer, flow):
 
         v = msg.value()
 
@@ -127,7 +137,7 @@ class Processor(RequestResponseService):
 
         try:
 
-            print(v.terms)
+            print(v.terms, flush=True)
 
             input = {
                 k: json.loads(v)
@@ -135,9 +145,40 @@ class Processor(RequestResponseService):
             }
             
             print(f"Handling kind {kind}...", flush=True)
-            print(input, flush=True)
 
-            resp = self.manager.invoke(kind, input)
+            q = await flow.consumer["text-completion-response"].subscribe(id)
+
+            async def llm(system, prompt):
+
+                print(system, flush=True)
+                print(prompt, flush=True)
+
+                await flow.producer["text-completion-request"].send(
+                    TextCompletionRequest(
+                        system=system, prompt=prompt
+                    ),
+                    properties={"id": id}
+                )
+
+                # FIXME: hard-coded?
+                resp = await asyncio.wait_for(
+                    q.get(),
+                    timeout=600
+                )
+
+                try:
+                    return resp.response
+                except Exception as e:
+                    print("LLM Exception:", e, flush=True)
+                    return None
+
+            try:
+                resp = await self.manager.invoke(kind, input, llm)
+            except Exception as e:
+                print("Invocation exception:", e, flush=True)
+                raise e
+            finally:
+                await flow.consumer["text-completion-response"].unsubscribe(id)
 
             if isinstance(resp, str):
 
@@ -150,7 +191,7 @@ class Processor(RequestResponseService):
                     error=None,
                 )
 
-                await self.send(r, properties={"id": id})
+                await flow.response.send(r, properties={"id": id})
 
                 return
 
@@ -165,13 +206,13 @@ class Processor(RequestResponseService):
                     error=None,
                 )
 
-                await self.send(r, properties={"id": id})
+                await flow.response.send(r, properties={"id": id})
 
                 return
             
         except Exception as e:
 
-            print(f"Exception: {e}")
+            print(f"Exception: {e}", flush=True)
 
             print("Send error response...", flush=True)
 
@@ -183,11 +224,11 @@ class Processor(RequestResponseService):
                 response=None,
             )
 
-            await self.send(r, properties={"id": id})
+            await flow.response.send(r, properties={"id": id})
 
         except Exception as e:
 
-            print(f"Exception: {e}")
+            print(f"Exception: {e}", flush=True)
 
             print("Send error response...", flush=True)
 
@@ -204,22 +245,7 @@ class Processor(RequestResponseService):
     @staticmethod
     def add_args(parser):
 
-        RequestResponseService.add_args(
+        FlowProcessor.add_args(parser)
-            parser, default_input_queue, default_subscriber,
-            default_output_queue,
-        )
-
-        parser.add_argument(
-            '--text-completion-request-queue',
-            default=text_completion_request_queue,
-            help=f'Text completion request queue (default: {text_completion_request_queue})',
-        )
-
-        parser.add_argument(
-            '--text-completion-response-queue',
-            default=text_completion_response_queue,
-            help=f'Text completion response queue (default: {text_completion_response_queue})',
-        )
 
         parser.add_argument(
             '--config-type',
@@ -229,5 +255,5 @@ class Processor(RequestResponseService):
 
 def run():
 
-    Processor.launch(module, __doc__)
+    Processor.launch(default_ident, __doc__)
 

trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py

@@ -181,7 +181,7 @@ class Processor(RequestResponseService):
     @staticmethod
     def add_args(parser):
 
-        RequestResponseService.add_args(parser, default_subscriber)
+        RequestResponseService.add_args(parser)
 
         parser.add_argument(
             '-m', '--model',