Request/response spec working, implemented in extract-defs and vertexai

trustgraph-base/trustgraph/base/__init__.py

@@ -7,7 +7,9 @@ from . publisher import Publisher
 from . subscriber import Subscriber
 from . metrics import ProcessorMetrics, ConsumerMetrics, ProducerMetrics
 from . flow_processor import FlowProcessor
-from . flow_processor import ConsumerSpec, SettingSpec
+from . consumer_spec import ConsumerSpec
-from . flow_processor import SubscriberSpec, ProducerSpec
+from . setting_spec import SettingSpec
-from . request_response import RequestResponseService
+from . producer_spec import ProducerSpec
+from . subscriber_spec import SubscriberSpec
+from . request_response_spec import RequestResponseSpec
 

trustgraph-base/trustgraph/base/async_processor.py

@@ -30,7 +30,7 @@ class AsyncProcessor:
         self.id = params.get("id")
 
         # Register a pulsar client
-        self.client = PulsarClient(**params)
+        self.pulsar_client = PulsarClient(**params)
 
         # Initialise metrics, records the parameters
         ProcessorMetrics(id=self.id).info({
@@ -58,10 +58,17 @@ class AsyncProcessor:
         config_subscriber_id = str(uuid.uuid4())
 
         # Subscribe to config queue
-        self.config_sub_task = self.subscribe(
+        self.config_sub_task = Consumer(
+
+            taskgroup = self.taskgroup,
+            client = self.client,
+            subscriber = config_subscriber_id,
             flow = None,
-            queue = self.config_push_queue, subscriber = config_subscriber_id,
+
-            schema = ConfigPush, handler = self.on_config_change,
+            topic = self.config_push_queue,
+            schema = ConfigPush,
+
+            handler = self.on_config_change,
 
             # This causes new subscriptions to view the entire history of
             # configuration
@@ -85,6 +92,10 @@ class AsyncProcessor:
     @property
     def pulsar_host(self): return self.client.pulsar_host
 
+    # Returns the pulsar client
+    @property
+    def client(self): return self.pulsar_client.client
+
     # Register a new event handler for configuration change
     def register_config_handler(self, handler):
         self.config_handlers.append(handler)
@@ -112,33 +123,6 @@ class AsyncProcessor:
         while self.running:
             await asyncio.sleep(2)
 
-    # Subscribe to a topic, returns a Consumer
-    def subscribe(
-            self, flow, queue, subscriber, schema, handler, metrics=None,
-            start_of_messages=False
-    ):
-
-        return Consumer(
-            flow = flow,
-            taskgroup = self.taskgroup,
-            client = self.client,
-            queue = queue,
-            subscriber = subscriber,
-            schema = schema,
-            handler = handler,
-            metrics = metrics,
-            start_of_messages=start_of_messages,
-        )
-
-    # Open a mechanism to publish messages to a topic.  Returns a subscriber
-    def publish(self, queue, schema, metrics=None):
-        return Producer(
-            client = self.client,
-            queue = queue,
-            schema = schema,
-            metrics = metrics,
-        )
-
     # Startup fabric.  This runs in 'async' mode, creates a taskgroup and
     # runs the producer.
     @classmethod

trustgraph-base/trustgraph/base/consumer.py

@@ -1,4 +1,6 @@
 
+from pulsar.schema import JsonSchema
+import pulsar
 import _pulsar
 import asyncio
 import time
@@ -8,19 +10,18 @@ from .. exceptions import TooManyRequests
 class Consumer:
 
     def __init__(
-            self, taskgroup, flow, client, queue, subscriber, schema,
+            self, taskgroup, flow, client, topic, subscriber, schema,
             handler, 
             metrics = None,
             start_of_messages=False,
             rate_limit_retry_time = 10, rate_limit_timeout = 7200,
             reconnect_time = 5,
-            
     ):
 
         self.taskgroup = taskgroup
         self.flow = flow
         self.client = client
-        self.queue = queue
+        self.topic = topic
         self.subscriber = subscriber
         self.schema = schema
         self.handler = handler
@@ -43,7 +44,8 @@ class Consumer:
         self.running = False
 
         if hasattr(self, "consumer"):
-            self.consumer.close()
+            if self.consumer:
+                self.consumer.close()
 
     async def stop(self):
 
@@ -69,13 +71,19 @@ class Consumer:
 
             try:
 
-                print(self.queue, "subscribing...", flush=True)
+                print(self.topic, "subscribing...", flush=True)
+
+                if self.start_of_messages:
+                    pos = pulsar.InitialPosition.Earliest
+                else:
+                    pos = pulsar.InitialPosition.Latest
 
                 self.consumer = await asyncio.to_thread(
                     self.client.subscribe,
-                    queue = self.queue, subscriber = self.subscriber,
+                    topic = self.topic,
-                    schema = self.schema,
+                    subscription_name = self.subscriber,
-                    start_of_messages = self.start_of_messages
+                    schema = JsonSchema(self.schema),
+                    initial_position = pos
                 )
 
             except Exception as e:
@@ -84,7 +92,7 @@ class Consumer:
                 await asyncio.sleep(self.reconnect_time)
                 continue
 
-            print(self.queue, "subscribed", flush=True)
+            print(self.topic, "subscribed", flush=True)
 
             if self.metrics:
                 self.metrics.state("running")

trustgraph-base/trustgraph/base/consumer_spec.py

@@ -0,0 +1,36 @@
+
+from . metrics import ConsumerMetrics
+from . consumer import Consumer
+from . spec import Spec
+
+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 = Consumer(
+            taskgroup = processor.taskgroup,
+            flow = flow,
+            client = processor.client,
+            topic = 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
+

trustgraph-base/trustgraph/base/flow.py

@@ -0,0 +1,32 @@
+
+import asyncio
+
+class Flow:
+    def __init__(self, id, flow, processor, defn):
+
+        self.id = id
+        self.name = flow
+
+        self.producer = {}
+
+        # Consumers and publishers.  Is this a bit untidy?
+        self.consumer = {}
+
+        self.setting = {}
+
+        for spec in processor.specifications:
+            spec.add(self, processor, defn)
+
+    async def start(self):
+        for c in self.consumer.values():
+            await c.start()
+
+    async def stop(self):
+        for c in self.consumer.values():
+            await c.stop()
+
+    def __call__(self, key):
+        if key in self.producer: return self.producer[key]
+        if key in self.consumer: return self.consumer[key]
+        if key in self.setting: return self.setting[key].value
+        return None

trustgraph-base/trustgraph/base/flow_processor.py

@@ -11,124 +11,10 @@ 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, Subscriber
+from . async_processor import AsyncProcessor
-from . metrics import ConsumerMetrics, ProducerMetrics
+from . subscriber import Subscriber
+from . flow import Flow
 
-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.id = id
-        self.name = flow
-
-        self.producer = {}
-        self.consumer = {}
-        self.setting = {}
-
-        for spec in processor.specifications:
-            spec.add(self, processor, defn)
-
-    async def start(self):
-        for c in self.consumer.values():
-            await c.start()
-
-    async def stop(self):
-        for c in self.consumer.values():
-            await c.stop()
-                
 # Parent class for configurable processors, configured with flows by
 # the config service
 class FlowProcessor(AsyncProcessor):
@@ -151,18 +37,6 @@ class FlowProcessor(AsyncProcessor):
 
         print("Service initialised.")
 
-    # Register a new consumer name
-    def register_consumer(self, name, schema, handler):
-        self.specifications.append(ConsumerSpec(name, schema, handler))
-
-    # Register a producer name
-    def register_producer(self, name, schema):
-        self.specifications.append(ProducerSpec(name, schema))
-
-    # Register a configuration variable
-    def register_config(self, name):
-        self.specifications.append(SettingSpec(name))
-
     # Register a configuration variable
     def register_specification(self, spec):
         self.specifications.append(spec)

trustgraph-base/trustgraph/base/producer.py

@@ -1,11 +1,12 @@
 
+from pulsar.schema import JsonSchema
 import asyncio
 
 class Producer:
 
-    def __init__(self, client, queue, schema, metrics=None):
+    def __init__(self, client, topic, schema, metrics=None):
         self.client = client
-        self.queue = queue
+        self.topic = topic
         self.schema = schema
 
         self.metrics = metrics
@@ -34,9 +35,12 @@ class Producer:
         while self.running and self.producer is None:
 
             try:
-                print("Connect publisher to", self.queue, "...", flush=True)
+                print("Connect publisher to", self.topic, "...", flush=True)
-                self.producer = self.client.publish(self.queue, self.schema)
+                self.producer = self.client.create_producer(
-                print("Connected to", self.queue, flush=True)
+                    topic = self.topic,
+                    schema = JsonSchema(self.schema)
+                )
+                print("Connected to", self.topic, flush=True)
             except Exception as e:
                 print("Exception:", e, flush=True)
                 await asyncio.sleep(2)

trustgraph-base/trustgraph/base/producer_spec.py

@@ -0,0 +1,25 @@
+
+from . producer import Producer
+from . metrics import ProducerMetrics
+from . spec import Spec
+
+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 = Producer(
+            client = processor.client,
+            topic = definition[self.name],
+            schema = self.schema,
+            metrics = producer_metrics,
+        )
+
+        flow.producer[self.name] = producer
+

trustgraph-base/trustgraph/base/publisher.py

@@ -7,9 +7,9 @@ import pulsar
 
 class Publisher:
 
-    def __init__(self, pulsar_client, topic, schema=None, max_size=10,
+    def __init__(self, client, topic, schema=None, max_size=10,
                  chunking_enabled=True):
-        self.client = pulsar_client
+        self.client = client
         self.topic = topic
         self.schema = schema
         self.q = asyncio.Queue(maxsize=max_size)

trustgraph-base/trustgraph/base/pubsub.py

@@ -51,30 +51,6 @@ class PulsarClient:
             if self.client:
                 self.client.close()
 
-    def subscribe(self, queue, subscriber, schema,
-                  start_of_messages=False):
-
-        if start_of_messages:
-            pos = pulsar.InitialPosition.Earliest
-        else:
-            pos =  pulsar.InitialPosition.Latest
-
-        return self.client.subscribe(
-            topic = queue,
-            subscription_name = subscriber,
-            consumer_type = pulsar.ConsumerType.Shared,
-            schema = JsonSchema(schema),
-            initial_position = pos,
-        )
-
-    def publish(self, queue, schema):
-
-        return self.client.create_producer(
-            topic=queue,
-            schema=JsonSchema(schema),
-            chunking_enabled=True,
-        )
-
     @staticmethod
     def add_args(parser):
 

trustgraph-base/trustgraph/base/request_response.py

@@ -1,73 +0,0 @@
-
-import json
-from pulsar.schema import JsonSchema
-
-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 ProcessorMetrics, ConsumerMetrics, ProducerMetrics
-from . flow_processor import FlowProcessor
-
-class RequestResponseService(FlowProcessor):
-
-    def __init__(self, **params):
-
-        super(RequestResponseService, self).__init__(**params)
-
-        self.response_schema = params.get("responsedrequest_schema")
-
-        # These can be overriden by a derived class
-        self.consumer_spec = [
-            ("request", params.get("request_schema"), self.on_message)
-        ]
-        self.producer_spec = [
-            ("response", params.get("response_schema"))
-        ]
-
-        print("Service initialised.")
-
-    async def on_message(self, message, consumer, flow):
-        
-        v = message.value()
-
-        # Sender-produced ID
-        id = message.properties()["id"]
-
-        print(f"Handling input {id}...", flush=True)
-
-        try:
-            resp = await self.on_request(v, consumer, flow)
-
-            print("Send response...", flush=True)
-
-            await flow.producer["response"].send(resp, properties={"id": id})
-
-            return
-
-        except Exception as e:
-
-            print("Exception:", e, flush=True)
-            print("Send error response...", flush=True)
-            r = self.response_schema(
-                error=Error(
-                    type="internal-error",
-                    message = str(e)
-                )
-            )
-
-            await flow.producer["response"].send(r, properties={"id": id})
-
-        print("Done.", flush=True)
-
-    @staticmethod
-    def add_args(parser):
-
-        FlowProcessor.add_args(parser)
-
-def run():
-
-    Processor.launch(module, __doc__)
-

trustgraph-base/trustgraph/base/request_response_spec.py

@@ -0,0 +1,104 @@
+
+import uuid
+import asyncio
+
+from . subscriber import Subscriber
+from . producer import Producer
+from . spec import Spec
+from . metrics import ConsumerMetrics, ProducerMetrics
+
+class RequestResponse(Subscriber):
+
+    def __init__(
+            self, client, subscription, consumer_name,
+            request_topic, request_schema,
+            request_metrics,
+            response_topic, response_schema,
+            response_metrics,
+    ):
+
+        super(RequestResponse, self).__init__(
+            client = client,
+            subscription = subscription,
+            consumer_name = consumer_name,
+            topic = response_topic,
+            schema = response_schema,
+        )
+
+        self.producer = Producer(
+            client = client,
+            topic = request_topic,
+            schema = request_schema,
+            metrics = request_metrics,
+        )
+
+    async def start(self):
+        await self.producer.start()
+        await super(RequestResponse, self).start()
+
+    async def stop(self):
+        await self.producer.stop()
+        await super(RequestResponse, self).stop()
+
+    async def request(self, req, timeout=300):
+
+        id = str(uuid.uuid4())
+
+        q = await self.subscribe(id)
+
+        try:
+
+            await self.producer.send(
+                req,
+                properties={"id": id}
+            )
+
+            resp = await asyncio.wait_for(
+                q.get(),
+                timeout=timeout
+            )
+
+            return resp
+
+        except Exception as e:
+
+            print("Exception:", e)
+            raise e
+
+        finally:
+
+            await self.unsubscribe(id)
+
+# This deals with the request/response case.  The caller needs to
+# use another service in request/response mode.  Uses two topics:
+# - we send on the request topic as a producer
+# - we receive on the response topic as a subscriber
+class RequestResponseSpec(Spec):
+    def __init__(
+            self, request_name, request_schema, response_name, response_schema
+    ):
+        self.request_name = request_name
+        self.request_schema = request_schema
+        self.response_name = response_name
+        self.response_schema = response_schema
+
+    def add(self, flow, processor, definition):
+
+        producer_metrics = ProducerMetrics(
+            flow.id, f"{flow.name}-{self.response_name}"
+        )
+
+        rr = RequestResponse(
+            client = processor.client,
+            subscription = flow.id,
+            consumer_name = flow.id,
+            request_topic = definition[self.request_name],
+            request_schema = self.request_schema,
+            request_metrics = producer_metrics,
+            response_topic = definition[self.response_name],
+            response_schema = self.response_schema,
+            response_metrics = None,
+        )
+
+        flow.consumer[self.request_name] = rr
+

trustgraph-base/trustgraph/base/setting_spec.py

@@ -0,0 +1,19 @@
+
+from . spec import Spec
+
+class Setting:
+    def __init__(self, value):
+        self.value = value
+    async def start():
+        pass
+    async def stop():
+        pass
+        
+class SettingSpec(Spec):
+    def __init__(self, name):
+        self.name = name
+
+    def add(self, flow, processor, definition):
+
+        flow.config[self.name] = Setting(definition[self.name])
+

trustgraph-base/trustgraph/base/spec.py

@@ -0,0 +1,4 @@
+
+class Spec:
+    pass
+

trustgraph-base/trustgraph/base/subscriber.py

@@ -6,9 +6,9 @@ import time
 
 class Subscriber:
 
-    def __init__(self, pulsar_client, topic, subscription, consumer_name,
+    def __init__(self, client, topic, subscription, consumer_name,
                  schema=None, max_size=100):
-        self.client = pulsar_client
+        self.client = client
         self.topic = topic
         self.subscription = subscription
         self.consumer_name = consumer_name

trustgraph-base/trustgraph/base/subscriber_spec.py

@@ -0,0 +1,30 @@
+
+from . metrics import ConsumerMetrics
+from . subscriber import Subscriber
+from . spec import Spec
+
+class SubscriberSpec(Spec):
+
+    def __init__(self, name, schema):
+        self.name = name
+        self.schema = schema
+
+    def add(self, flow, processor, definition):
+
+        # FIXME: Metrics not used
+        subscriber_metrics = ConsumerMetrics(
+            flow.id, f"{flow.name}-{self.name}"
+        )
+
+        subscriber = Subscriber(
+            client = processor.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
+

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

@@ -10,7 +10,7 @@ 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 FlowProcessor
+from ... base import FlowProcessor, ConsumerSpec, ProducerSpec
 
 default_ident = "chunker"
 
@@ -41,15 +41,19 @@ class Processor(FlowProcessor):
             is_separator_regex=False,
         )
 
-        self.register_consumer(
+        self.register_specification(
-            name = "input",
+            ConsumerSpec(
-            schema = TextDocument,
+                name = "input",
-            handler = self.on_message,
+                schema = TextDocument,
+                handler = self.on_message,
+            )
         )
 
-        self.register_producer(
+        self.register_specification(
-            name = "output",
+            ProducerSpec(
-            schema = Chunk,
+                name = "output",
+                schema = Chunk,
+            )
         )
 
         print("Chunker initialised", flush=True)
@@ -76,7 +80,7 @@ class Processor(FlowProcessor):
                 id=consumer.id, flow=consumer.flow
             ).observe(len(chunk.page_content))
 
-            await flow.producer["output"].send(r)
+            await flow("output").send(r)
 
         print("Done.", flush=True)
 

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

@@ -40,15 +40,19 @@ class Processor(FlowProcessor):
             chunk_overlap=chunk_overlap,
         )
 
-        self.register_consumer(
+        self.register_specification(
-            name = "input",
+            ConsumerSpec(
-            schema = TextDocument,
+                name = "input",
-            handler = self.on_message,
+                schema = TextDocument,
+                handler = self.on_message,
+            )
         )
 
-        self.register_producer(
+        self.register_specification(
-            name = "output",
+            ProducerSpec(
-            schema = Chunk,
+                name = "output",
+                schema = Chunk,
+            )
         )
 
         print("Chunker initialised", flush=True)
@@ -75,7 +79,7 @@ class Processor(FlowProcessor):
                 id=consumer.id, flow=consumer.flow
             ).observe(len(chunk.page_content))
 
-            await flow.producer["output"].send(r)
+            await flow("output").send(r)
 
         print("Done.", flush=True)
 

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

@@ -14,6 +14,7 @@ from trustgraph.base import AsyncProcessor, Consumer, Producer
 
 from . config import Configuration
 from ... base import ProcessorMetrics, ConsumerMetrics, ProducerMetrics
+from ... base import Consumer, Producer
 
 default_ident = "config-svc"
 
@@ -42,29 +43,33 @@ class Processor(AsyncProcessor):
             }
         )
 
-        self.request_metrics = ConsumerMetrics(id + "-request")
+        request_metrics = ConsumerMetrics(id + "-request")
-        self.response_metrics = ProducerMetrics(id + "-response")
+        response_metrics = ProducerMetrics(id + "-response")
-        self.push_metrics = ProducerMetrics(id + "-push")
+        push_metrics = ProducerMetrics(id + "-push")
 
-        self.push_pub = self.publish(
+        self.push_pub = Producer(
-            queue = push_queue,
+            client = self.client,
+            topic = push_queue,
             schema = ConfigPush,
-            metrics = self.push_metrics,
+            metrics = push_metrics,
         )
 
-        self.response_pub = self.publish(
+        self.response_pub = Producer(
-            queue = response_queue,
+            client = self.client,
+            topic = response_queue,
             schema = ConfigResponse,
-            metrics = self.response_metrics,
+            metrics = response_metrics,
         )
 
-        self.subs = self.subscribe(
+        self.subs = Consumer(
+            taskgroup = self.taskgroup,
+            client = self.client,
             flow = None,
-            queue = request_queue,
+            topic = request_queue,
             subscriber = id,
             schema = request_schema,
             handler = self.on_message,
-            metrics = self.request_metrics,
+            metrics = request_metrics,
         )
 
         self.config = Configuration(self.push)

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

@@ -11,7 +11,7 @@ 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 FlowProcessor
+from ... base import FlowProcessor, ConsumerSpec, ProducerSpec
 
 default_ident = "pdf-decoder"
 
@@ -27,15 +27,19 @@ class Processor(FlowProcessor):
             }
         )
 
-        self.register_consumer(
+        self.register_specification(
-            name = "input",
+            ConsumerSpec(
-            schema = Document,
+                name = "input",
-            handler = self.on_message,
+                schema = Document,
+                handler = self.on_message,
+            )
         )
 
-        self.register_producer(
+        self.register_specification(
-            name = "output",
+            ProducerSpec(
-            schema = TextDocument,
+                name = "output",
+                schema = TextDocument,
+            )
         )
 
         print("PDF inited", flush=True)
@@ -67,7 +71,7 @@ class Processor(FlowProcessor):
                         text=page.page_content.encode("utf-8"),
                     )
 
-                    await flow.producer["output"].send(r)
+                    await flow("output").send(r)
 
         print("Done.", flush=True)
 

trustgraph-flow/trustgraph/extract/kg/definitions/extract.py

@@ -6,7 +6,6 @@ entity/context definitions for embedding.
 """
 
 import json
-import asyncio
 import urllib.parse
 from pulsar.schema import JsonSchema
 import uuid
@@ -18,7 +17,7 @@ from .... log_level import LogLevel
 from .... clients.prompt_client import PromptClient
 from .... rdf import TRUSTGRAPH_ENTITIES, DEFINITION, RDF_LABEL, SUBJECT_OF
 
-from .... base import FlowProcessor, SubscriberSpec, ConsumerSpec
+from .... base import FlowProcessor, RequestResponseSpec, ConsumerSpec
 from .... base import ProducerSpec
 
 DEFINITION_VALUE = Value(value=DEFINITION, is_uri=True)
@@ -48,16 +47,11 @@ class Processor(FlowProcessor):
         )
 
         self.register_specification(
-            ProducerSpec(
+            RequestResponseSpec(
-                name = "prompt-request",
+                request_name = "prompt-request",
-                schema = PromptRequest
+                request_schema = PromptRequest,
-            )
+                response_name = "prompt-response",
-        )
+                response_schema = PromptResponse,
-
-        self.register_specification(
-            SubscriberSpec(
-                name = "prompt-response",
-                schema = PromptResponse,
             )
         )
 
@@ -116,26 +110,16 @@ class Processor(FlowProcessor):
 
         try:
 
-            q = await flow.consumer["prompt-response"].subscribe(id)
-
             try:
 
-                await flow.producer["prompt-request"].send(
+                resp = await flow("prompt-request").request(
                     PromptRequest(
                         id="extract-definitions",
                         terms={
                             "text": json.dumps(chunk)
                         },
-                    ),
+                    )
-                    properties={"id": id}
                 )
-
-                # FIXME: hard-coded?
-                resp = await asyncio.wait_for(
-                    q.get(),
-                    timeout=600
-                )
-
                 print("Response", resp, flush=True)
 
                 if resp.error is not None:
@@ -150,8 +134,6 @@ class Processor(FlowProcessor):
             except Exception as e:
                 print("Prompt exception:", e, flush=True)
                 raise e
-            finally:
-                await flow.consumer["prompt-response"].unsubscribe(id)
 
             triples = []
             entities = []
@@ -201,7 +183,7 @@ class Processor(FlowProcessor):
                 entities.append(ec)
 
             await self.emit_triples(
-                flow.producer["triples"],
+                flow("triples"),
                 Metadata(
                     id=v.metadata.id,
                     metadata=[],
@@ -212,7 +194,7 @@ class Processor(FlowProcessor):
             )
 
             await self.emit_ecs(
-                flow.producer["entity-contexts"],
+                flow("entity-contexts"),
                 Metadata(
                     id=v.metadata.id,
                     metadata=[],

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

@@ -13,7 +13,7 @@ from .... schema import PromptRequest, PromptResponse, Error
 from .... schema import TextCompletionRequest, TextCompletionResponse
 
 from .... base import FlowProcessor
-from .... base import ProducerSpec, SubscriberSpec, ConsumerSpec
+from .... base import ProducerSpec, ConsumerSpec, RequestResponseSpec
 
 from . prompt_manager import PromptConfiguration, Prompt, PromptManager
 
@@ -43,16 +43,11 @@ class Processor(FlowProcessor):
         )
 
         self.register_specification(
-            ProducerSpec(
+            RequestResponseSpec(
-                name = "text-completion-request",
+                request_name = "text-completion-request",
-                schema = TextCompletionRequest
+                request_schema = TextCompletionRequest,
-            )
+                response_name = "text-completion-response",
-        )
+                response_schema = TextCompletionResponse,
-
-        self.register_specification(
-            SubscriberSpec(
-                name = "text-completion-response",
-                schema = TextCompletionResponse,
             )
         )
 
@@ -139,24 +134,15 @@ class Processor(FlowProcessor):
             
             print(f"Handling kind {kind}...", flush=True)
 
-            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(
+                resp = await flow("text-completion-request").request(
                     TextCompletionRequest(
                         system=system, prompt=prompt
                     ),
-                    properties={"id": id}
-                )
-
-                # FIXME: hard-coded?
-                resp = await asyncio.wait_for(
-                    q.get(),
-                    timeout=600
                 )
 
                 try:
@@ -170,8 +156,6 @@ class Processor(FlowProcessor):
             except Exception as e:
                 print("Invocation exception:", e, flush=True)
                 raise e
-            finally:
-                await flow.consumer["text-completion-response"].unsubscribe(id)
 
             print(resp, flush=True)
 
@@ -185,7 +169,7 @@ class Processor(FlowProcessor):
                     error=None,
                 )
 
-                await flow.response.send(r, properties={"id": id})
+                await flow("response").send(r, properties={"id": id})
 
                 return
 
@@ -200,7 +184,7 @@ class Processor(FlowProcessor):
                     error=None,
                 )
 
-                await flow.response.send(r, properties={"id": id})
+                await flow("response").send(r, properties={"id": id})
 
                 return
             
@@ -218,7 +202,7 @@ class Processor(FlowProcessor):
                 response=None,
             )
 
-            await flow.response.send(r, properties={"id": id})
+            await flow("response").send(r, properties={"id": id})
 
         except Exception as e: