From 1047cd2fa1eaa320f01c6d449ba4f170352f53d0 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Fri, 18 Apr 2025 18:35:17 +0100 Subject: [PATCH] Request/response spec working, implemented in extract-defs and vertexai --- trustgraph-base/trustgraph/base/__init__.py | 8 +- .../trustgraph/base/async_processor.py | 46 ++---- trustgraph-base/trustgraph/base/consumer.py | 26 ++-- .../trustgraph/base/consumer_spec.py | 36 +++++ trustgraph-base/trustgraph/base/flow.py | 32 +++++ .../trustgraph/base/flow_processor.py | 132 +----------------- trustgraph-base/trustgraph/base/producer.py | 14 +- .../trustgraph/base/producer_spec.py | 25 ++++ trustgraph-base/trustgraph/base/publisher.py | 4 +- trustgraph-base/trustgraph/base/pubsub.py | 24 ---- .../trustgraph/base/request_response.py | 73 ---------- .../trustgraph/base/request_response_spec.py | 104 ++++++++++++++ .../trustgraph/base/setting_spec.py | 19 +++ trustgraph-base/trustgraph/base/spec.py | 4 + trustgraph-base/trustgraph/base/subscriber.py | 4 +- .../trustgraph/base/subscriber_spec.py | 30 ++++ .../trustgraph/chunking/recursive/chunker.py | 22 +-- .../trustgraph/chunking/token/chunker.py | 20 +-- .../trustgraph/config/service/service.py | 29 ++-- .../trustgraph/decoding/pdf/pdf_decoder.py | 22 +-- .../extract/kg/definitions/extract.py | 38 ++--- .../model/prompt/template/service.py | 36 ++--- 22 files changed, 378 insertions(+), 370 deletions(-) create mode 100644 trustgraph-base/trustgraph/base/consumer_spec.py create mode 100644 trustgraph-base/trustgraph/base/flow.py create mode 100644 trustgraph-base/trustgraph/base/producer_spec.py delete mode 100644 trustgraph-base/trustgraph/base/request_response.py create mode 100644 trustgraph-base/trustgraph/base/request_response_spec.py create mode 100644 trustgraph-base/trustgraph/base/setting_spec.py create mode 100644 trustgraph-base/trustgraph/base/spec.py create mode 100644 trustgraph-base/trustgraph/base/subscriber_spec.py diff --git a/trustgraph-base/trustgraph/base/__init__.py b/trustgraph-base/trustgraph/base/__init__.py index d8ae20df..9e706bf5 100644 --- a/trustgraph-base/trustgraph/base/__init__.py +++ b/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 . flow_processor import SubscriberSpec, ProducerSpec -from . request_response import RequestResponseService +from . consumer_spec import ConsumerSpec +from . setting_spec import SettingSpec +from . producer_spec import ProducerSpec +from . subscriber_spec import SubscriberSpec +from . request_response_spec import RequestResponseSpec diff --git a/trustgraph-base/trustgraph/base/async_processor.py b/trustgraph-base/trustgraph/base/async_processor.py index a4f4312a..0f21b491 100644 --- a/trustgraph-base/trustgraph/base/async_processor.py +++ b/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 diff --git a/trustgraph-base/trustgraph/base/consumer.py b/trustgraph-base/trustgraph/base/consumer.py index 1ffa4ef1..d24e8f6f 100644 --- a/trustgraph-base/trustgraph/base/consumer.py +++ b/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, - schema = self.schema, - start_of_messages = self.start_of_messages + topic = self.topic, + subscription_name = self.subscriber, + 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") diff --git a/trustgraph-base/trustgraph/base/consumer_spec.py b/trustgraph-base/trustgraph/base/consumer_spec.py new file mode 100644 index 00000000..3d01be8f --- /dev/null +++ b/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 + diff --git a/trustgraph-base/trustgraph/base/flow.py b/trustgraph-base/trustgraph/base/flow.py new file mode 100644 index 00000000..9cda34a0 --- /dev/null +++ b/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 diff --git a/trustgraph-base/trustgraph/base/flow_processor.py b/trustgraph-base/trustgraph/base/flow_processor.py index ea0935ea..e6757e97 100644 --- a/trustgraph-base/trustgraph/base/flow_processor.py +++ b/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 . metrics import ConsumerMetrics, ProducerMetrics +from . async_processor import AsyncProcessor +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) diff --git a/trustgraph-base/trustgraph/base/producer.py b/trustgraph-base/trustgraph/base/producer.py index 5abc6796..bb665924 100644 --- a/trustgraph-base/trustgraph/base/producer.py +++ b/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) - self.producer = self.client.publish(self.queue, self.schema) - print("Connected to", self.queue, flush=True) + print("Connect publisher to", self.topic, "...", flush=True) + self.producer = self.client.create_producer( + 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) diff --git a/trustgraph-base/trustgraph/base/producer_spec.py b/trustgraph-base/trustgraph/base/producer_spec.py new file mode 100644 index 00000000..9007f48b --- /dev/null +++ b/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 + diff --git a/trustgraph-base/trustgraph/base/publisher.py b/trustgraph-base/trustgraph/base/publisher.py index f40fbb30..ce9e364e 100644 --- a/trustgraph-base/trustgraph/base/publisher.py +++ b/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) diff --git a/trustgraph-base/trustgraph/base/pubsub.py b/trustgraph-base/trustgraph/base/pubsub.py index 1ad08287..b9f233d4 100644 --- a/trustgraph-base/trustgraph/base/pubsub.py +++ b/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): diff --git a/trustgraph-base/trustgraph/base/request_response.py b/trustgraph-base/trustgraph/base/request_response.py deleted file mode 100644 index 44259c2c..00000000 --- a/trustgraph-base/trustgraph/base/request_response.py +++ /dev/null @@ -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__) - diff --git a/trustgraph-base/trustgraph/base/request_response_spec.py b/trustgraph-base/trustgraph/base/request_response_spec.py new file mode 100644 index 00000000..876a8349 --- /dev/null +++ b/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 + diff --git a/trustgraph-base/trustgraph/base/setting_spec.py b/trustgraph-base/trustgraph/base/setting_spec.py new file mode 100644 index 00000000..5c5152b2 --- /dev/null +++ b/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]) + diff --git a/trustgraph-base/trustgraph/base/spec.py b/trustgraph-base/trustgraph/base/spec.py new file mode 100644 index 00000000..4d0d937b --- /dev/null +++ b/trustgraph-base/trustgraph/base/spec.py @@ -0,0 +1,4 @@ + +class Spec: + pass + diff --git a/trustgraph-base/trustgraph/base/subscriber.py b/trustgraph-base/trustgraph/base/subscriber.py index 79c1b918..a8ff58f7 100644 --- a/trustgraph-base/trustgraph/base/subscriber.py +++ b/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 diff --git a/trustgraph-base/trustgraph/base/subscriber_spec.py b/trustgraph-base/trustgraph/base/subscriber_spec.py new file mode 100644 index 00000000..2f89290b --- /dev/null +++ b/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 + diff --git a/trustgraph-flow/trustgraph/chunking/recursive/chunker.py b/trustgraph-flow/trustgraph/chunking/recursive/chunker.py index d03af3b1..6540d0b0 100755 --- a/trustgraph-flow/trustgraph/chunking/recursive/chunker.py +++ b/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( - name = "input", - schema = TextDocument, - handler = self.on_message, + self.register_specification( + ConsumerSpec( + name = "input", + schema = TextDocument, + handler = self.on_message, + ) ) - self.register_producer( - name = "output", - schema = Chunk, + self.register_specification( + ProducerSpec( + 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) diff --git a/trustgraph-flow/trustgraph/chunking/token/chunker.py b/trustgraph-flow/trustgraph/chunking/token/chunker.py index 76e49e87..5ae0f4f4 100755 --- a/trustgraph-flow/trustgraph/chunking/token/chunker.py +++ b/trustgraph-flow/trustgraph/chunking/token/chunker.py @@ -40,15 +40,19 @@ class Processor(FlowProcessor): chunk_overlap=chunk_overlap, ) - self.register_consumer( - name = "input", - schema = TextDocument, - handler = self.on_message, + self.register_specification( + ConsumerSpec( + name = "input", + schema = TextDocument, + handler = self.on_message, + ) ) - self.register_producer( - name = "output", - schema = Chunk, + self.register_specification( + ProducerSpec( + 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) diff --git a/trustgraph-flow/trustgraph/config/service/service.py b/trustgraph-flow/trustgraph/config/service/service.py index bd1b1cb4..47bac828 100644 --- a/trustgraph-flow/trustgraph/config/service/service.py +++ b/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") - self.response_metrics = ProducerMetrics(id + "-response") - self.push_metrics = ProducerMetrics(id + "-push") + request_metrics = ConsumerMetrics(id + "-request") + response_metrics = ProducerMetrics(id + "-response") + push_metrics = ProducerMetrics(id + "-push") - self.push_pub = self.publish( - queue = push_queue, + self.push_pub = Producer( + client = self.client, + topic = push_queue, schema = ConfigPush, - metrics = self.push_metrics, + metrics = push_metrics, ) - self.response_pub = self.publish( - queue = response_queue, + self.response_pub = Producer( + 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) diff --git a/trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py b/trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py index 72ebe3dc..c9b2a4d6 100755 --- a/trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py +++ b/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( - name = "input", - schema = Document, - handler = self.on_message, + self.register_specification( + ConsumerSpec( + name = "input", + schema = Document, + handler = self.on_message, + ) ) - self.register_producer( - name = "output", - schema = TextDocument, + self.register_specification( + ProducerSpec( + 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) diff --git a/trustgraph-flow/trustgraph/extract/kg/definitions/extract.py b/trustgraph-flow/trustgraph/extract/kg/definitions/extract.py index cbf6a55b..a0f981d8 100755 --- a/trustgraph-flow/trustgraph/extract/kg/definitions/extract.py +++ b/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( - name = "prompt-request", - schema = PromptRequest - ) - ) - - self.register_specification( - SubscriberSpec( - name = "prompt-response", - schema = PromptResponse, + RequestResponseSpec( + request_name = "prompt-request", + request_schema = PromptRequest, + response_name = "prompt-response", + 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=[], diff --git a/trustgraph-flow/trustgraph/model/prompt/template/service.py b/trustgraph-flow/trustgraph/model/prompt/template/service.py index ba72fdd3..70731245 100755 --- a/trustgraph-flow/trustgraph/model/prompt/template/service.py +++ b/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( - name = "text-completion-request", - schema = TextCompletionRequest - ) - ) - - self.register_specification( - SubscriberSpec( - name = "text-completion-response", - schema = TextCompletionResponse, + RequestResponseSpec( + request_name = "text-completion-request", + request_schema = TextCompletionRequest, + response_name = "text-completion-response", + 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: