From b110af41fb4c64a205b8ea829342951a6b34b5da Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Wed, 16 Apr 2025 23:48:11 +0100 Subject: [PATCH] Prompt template working --- Makefile | 4 +- tests/test-llm | 4 +- tests/test-prompt-question | 9 +- trustgraph-base/trustgraph/base/__init__.py | 3 +- trustgraph-base/trustgraph/base/consumer.py | 62 ++++-- .../trustgraph/base/consumer.py.OLD | 187 ------------------ .../trustgraph/base/consumer_producer.py | 62 ------ .../trustgraph/base/flow_processor.py | 178 ++++++++++------- trustgraph-base/trustgraph/base/producer.py | 53 ++++- .../trustgraph/base/producer.py.OLD | 57 ------ trustgraph-base/trustgraph/base/publisher.py | 35 ++-- trustgraph-base/trustgraph/base/subscriber.py | 96 +++++---- .../model/prompt/template/prompt_manager.py | 14 +- .../model/prompt/template/service.py | 140 +++++++------ .../model/text_completion/vertexai/llm.py | 2 +- 15 files changed, 380 insertions(+), 526 deletions(-) delete mode 100644 trustgraph-base/trustgraph/base/consumer.py.OLD delete mode 100644 trustgraph-base/trustgraph/base/consumer_producer.py delete mode 100644 trustgraph-base/trustgraph/base/producer.py.OLD diff --git a/Makefile b/Makefile index 1899e602..73edcf2f 100644 --- a/Makefile +++ b/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 \ - -t ${CONTAINER_BASE}/trustgraph-vertexai:${VERSION} . +# ${DOCKER} build -f containers/Containerfile.vertexai \ +# -t ${CONTAINER_BASE}/trustgraph-vertexai:${VERSION} . basic-containers: update-package-versions ${DOCKER} build -f containers/Containerfile.base \ diff --git a/tests/test-llm b/tests/test-llm index 61f0da6b..aaae30a6 100755 --- a/tests/test-llm +++ b/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) - - diff --git a/tests/test-prompt-question b/tests/test-prompt-question index 50660965..78ba72aa 100755 --- a/tests/test-prompt-question +++ b/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 } ) diff --git a/trustgraph-base/trustgraph/base/__init__.py b/trustgraph-base/trustgraph/base/__init__.py index f6dd670c..1006cd25 100644 --- a/trustgraph-base/trustgraph/base/__init__.py +++ b/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 diff --git a/trustgraph-base/trustgraph/base/consumer.py b/trustgraph-base/trustgraph/base/consumer.py index adc33de6..7a9f25ba 100644 --- a/trustgraph-base/trustgraph/base/consumer.py +++ b/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") - diff --git a/trustgraph-base/trustgraph/base/consumer.py.OLD b/trustgraph-base/trustgraph/base/consumer.py.OLD deleted file mode 100644 index 6ba2650e..00000000 --- a/trustgraph-base/trustgraph/base/consumer.py.OLD +++ /dev/null @@ -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})' - ) - diff --git a/trustgraph-base/trustgraph/base/consumer_producer.py b/trustgraph-base/trustgraph/base/consumer_producer.py deleted file mode 100644 index 1006f9b5..00000000 --- a/trustgraph-base/trustgraph/base/consumer_producer.py +++ /dev/null @@ -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})' - ) - diff --git a/trustgraph-base/trustgraph/base/flow_processor.py b/trustgraph-base/trustgraph/base/flow_processor.py index 25d9985f..ea0935ea 100644 --- a/trustgraph-base/trustgraph/base/flow_processor.py +++ b/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.consumer = {} - self.config = {} + self.id = id self.name = flow - for spec in processor.config_spec: - name = spec - self.config[name] = defn[name] + self.producer = {} + self.consumer = {} + self.setting = {} - if not hasattr(self, name): - setattr(self, name, defn[name]) - - 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) + for spec in processor.specifications: + spec.add(self, processor, defn) 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) - self.consumer_spec = [] - - # Producer specification, array of ("name", SchemaType) - self.producer_spec = [] - - # Configuration specification, collects some flow variables from - # config, array of "name" - self.config_spec = [] + # Array of specifications: ConsumerSpec, ProducerSpec, SettingSpec + self.specifications = [] 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): diff --git a/trustgraph-base/trustgraph/base/producer.py b/trustgraph-base/trustgraph/base/producer.py index 30bfabfc..5abc6796 100644 --- a/trustgraph-base/trustgraph/base/producer.py +++ b/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: - self.metrics.inc() + if not self.running: return - 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 diff --git a/trustgraph-base/trustgraph/base/producer.py.OLD b/trustgraph-base/trustgraph/base/producer.py.OLD deleted file mode 100644 index 619670db..00000000 --- a/trustgraph-base/trustgraph/base/producer.py.OLD +++ /dev/null @@ -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})' - ) - diff --git a/trustgraph-base/trustgraph/base/publisher.py b/trustgraph-base/trustgraph/base/publisher.py index 2da63331..99a27c19 100644 --- a/trustgraph-base/trustgraph/base/publisher.py +++ b/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): - self.task = threading.Thread(target=self.run) - self.task.start() + async def start(self): + self.task = asyncio.create_task(self.run()) + await self.task.start() - def stop(self): + async def stop(self): self.running = False - def join(self): - self.stop() - self.task.join() + async def join(self): + await self.stop() + 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) - except queue.Empty: + id, item = asyncio.wait_for( + 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): - self.q.put((id, msg)) + async def send(self, id, item): + await self.q.put((id, item)) - diff --git a/trustgraph-base/trustgraph/base/subscriber.py b/trustgraph-base/trustgraph/base/subscriber.py index 3b1f158d..79c1b918 100644 --- a/trustgraph-base/trustgraph/base/subscriber.py +++ b/trustgraph-base/trustgraph/base/subscriber.py @@ -1,7 +1,7 @@ -import queue -import pulsar -import threading +from pulsar.schema import JsonSchema +import asyncio +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): - self.task = threading.Thread(target=self.run) - self.task.start() - - def stop(self): + async def __del__(self): self.running = False - def join(self): - self.task.join() + async def start(self): + 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, - subscription_name=self.subscription, - consumer_name=self.consumer_name, - schema=self.schema, + topic = self.topic, + subscription_name = self.subscription, + consumer_name = self.consumer_name, + schema = JsonSchema(self.schema), ) - print("SUBSCRIBER RUNNING...", flush=True) + print("Subscriber running...", flush=True) while self.running: - msg = consumer.receive() - - print("GOT MESSAGE...", flush=True) + try: + msg = await asyncio.to_thread( + 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) - except: - pass + await asyncio.wait_for( + self.q[id].put(value), + 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) - except: - pass + await asyncio.wait_for( + q.put(value), + 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) diff --git a/trustgraph-flow/trustgraph/model/prompt/template/prompt_manager.py b/trustgraph-flow/trustgraph/model/prompt/template/prompt_manager.py index d8a032ca..7c765acc 100644 --- a/trustgraph-flow/trustgraph/model/prompt/template/prompt_manager.py +++ b/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): - self.llm = llm + def __init__(self, config): 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) - - print(resp, flush=True) + resp = await llm(**prompt) 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}") diff --git a/trustgraph-flow/trustgraph/model/prompt/template/service.py b/trustgraph-flow/trustgraph/model/prompt/template/service.py index 87b10a7f..79fc5f3b 100755 --- a/trustgraph-flow/trustgraph/model/prompt/template/service.py +++ b/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 .... schema import text_completion_response_queue -from .... schema import prompt_request_queue, prompt_response_queue -from .... clients.llm_client import LlmClient -from .... base import RequestResponseService +from .... base import FlowProcessor, SubscriberSpec, ConsumerSpec +from .... base import ProducerSpec from . prompt_manager import PromptConfiguration, Prompt, PromptManager -module = "prompt" -default_subscriber = module +default_ident = "prompt" -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( - subscriber=subscriber, - input_queue=tc_request_queue, - output_queue=tc_response_queue, - pulsar_host = self.pulsar_host, - pulsar_api_key=self.pulsar_api_key, + # self.llm = LlmClient( + # subscriber=subscriber, + # input_queue=tc_request_queue, + # output_queue=tc_response_queue, + # pulsar_host = self.pulsar_host, + # pulsar_api_key=self.pulsar_api_key, + # ) + + self.register_specification( + ConsumerSpec( + name = "request", + schema = PromptRequest, + handler = self.on_request + ) ) - # System prompt hack - class Llm: - def __init__(self, llm): - self.llm = llm - def request(self, system, prompt): - print(system) - print(prompt, flush=True) - return self.llm.request(system, prompt) + self.register_specification( + ProducerSpec( + name = "text-completion-request", + schema = TextCompletionRequest + ) + ) - 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( - 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})', - ) + FlowProcessor.add_args(parser) parser.add_argument( '--config-type', @@ -229,5 +255,5 @@ class Processor(RequestResponseService): def run(): - Processor.launch(module, __doc__) + Processor.launch(default_ident, __doc__) diff --git a/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py b/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py index a18c4f82..b99f772f 100755 --- a/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py +++ b/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',