diff --git a/Makefile b/Makefile index 73edcf2f..1899e602 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/trustgraph-base/trustgraph/base/consumer.py b/trustgraph-base/trustgraph/base/consumer.py index 7a9f25ba..1ffa4ef1 100644 --- a/trustgraph-base/trustgraph/base/consumer.py +++ b/trustgraph-base/trustgraph/base/consumer.py @@ -9,9 +9,12 @@ class Consumer: def __init__( self, taskgroup, flow, client, queue, subscriber, schema, - handler, rate_limit_retry_time = 10, - rate_limit_timeout = 7200, metrics = None, + handler, + metrics = None, start_of_messages=False, + rate_limit_retry_time = 10, rate_limit_timeout = 7200, + reconnect_time = 5, + ): self.taskgroup = taskgroup @@ -21,8 +24,12 @@ class Consumer: self.subscriber = subscriber self.schema = schema self.handler = handler + self.rate_limit_retry_time = rate_limit_retry_time self.rate_limit_timeout = rate_limit_timeout + + self.reconnect_time = 5 + self.start_of_messages = start_of_messages self.running = True @@ -55,11 +62,11 @@ class Consumer: async def run(self): - if self.metrics: - self.metrics.state("stopped") - while self.running: + if self.metrics: + self.metrics.state("stopped") + try: print(self.queue, "subscribing...", flush=True) @@ -74,7 +81,7 @@ class Consumer: except Exception as e: print("Exception:", e, flush=True) - await asyncio.sleep(2) + await asyncio.sleep(self.reconnect_time) continue print(self.queue, "subscribed", flush=True) @@ -94,7 +101,7 @@ class Consumer: print("Exception:", e, flush=True) self.consumer.close() self.consumer = None - await asyncio.sleep(2) + await asyncio.sleep(self.reconnect_time) continue async def consume(self): @@ -161,7 +168,7 @@ class Consumer: self.metrics.rate_limit() # Sleep - await asyncio.sleep(self.rate_limit_retry) + await asyncio.sleep(self.rate_limit_retry_time) # Contine from retry loop, just causes a reprocessing continue diff --git a/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py b/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py index b99f772f..fadc42d9 100755 --- a/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py +++ b/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py @@ -18,26 +18,22 @@ from vertexai.preview.generative_models import ( ) from .... schema import TextCompletionRequest, TextCompletionResponse, Error -from .... schema import text_completion_request_queue -from .... schema import text_completion_response_queue -from .... log_level import LogLevel from .... exceptions import TooManyRequests -from .... base import RequestResponseService +from .... base import FlowProcessor, ConsumerSpec, ProducerSpec + +default_ident = "text-completion" -module = "text-completion" -default_subscriber = module default_model = 'gemini-1.0-pro-001' default_region = 'us-central1' default_temperature = 0.0 default_max_output = 8192 default_private_key = "private.json" -class Processor(RequestResponseService): +class Processor(FlowProcessor): def __init__(self, **params): id = params.get("id") - subscriber = params.get("subscriber", default_subscriber) region = params.get("region", default_region) model = params.get("model", default_model) private_key = params.get("private_key", default_private_key) @@ -54,6 +50,21 @@ class Processor(RequestResponseService): } ) + self.register_specification( + ConsumerSpec( + name = "request", + schema = TextCompletionRequest, + handler = self.on_request + ) + ) + + self.register_specification( + ProducerSpec( + name = "response", + schema = TextCompletionResponse + ) + ) + if not hasattr(__class__, "text_completion_metric"): __class__.text_completion_metric = Histogram( 'text_completion_duration', @@ -98,7 +109,11 @@ class Processor(RequestResponseService): print("Initialise VertexAI...", flush=True) if private_key: - credentials = service_account.Credentials.from_service_account_file(private_key) + credentials = ( + service_account.Credentials.from_service_account_file( + private_key + ) + ) else: credentials = None @@ -119,10 +134,16 @@ class Processor(RequestResponseService): print("Initialisation complete", flush=True) - async def on_request(self, request, consumer, flow): + async def on_request(self, msg, consumer, flow): try: + request = msg.value() + + # Sender-produced ID + + id = msg.properties()["id"] + prompt = request.system + "\n\n" + request.prompt with __class__.text_completion_metric.labels( @@ -144,12 +165,16 @@ class Processor(RequestResponseService): print("Send response...", flush=True) - return TextCompletionResponse( - error=None, - response=resp, - in_token=inputtokens, - out_token=outputtokens, - model=self.model + + await flow.producer["response"].send( + TextCompletionResponse( + error=None, + response=resp, + in_token=inputtokens, + out_token=outputtokens, + model=self.model + ), + properties={"id": id} ) except google.api_core.exceptions.ResourceExhausted as e: @@ -167,28 +192,30 @@ class Processor(RequestResponseService): print("Send error response...", flush=True) - return TextCompletionResponse( - error=Error( - type = "llm-error", - message = str(e), + await flow.producer["response"].send( + TextCompletionResponse( + error=Error( + type = "llm-error", + message = str(e), + ), + response=None, + in_token=None, + out_token=None, + model=None, ), - response=None, - in_token=None, - out_token=None, - model=None, + properties={"id": id} ) @staticmethod def add_args(parser): - RequestResponseService.add_args(parser) + FlowProcessor.add_args(parser) parser.add_argument( '-m', '--model', default=default_model, help=f'LLM model (default: {default_model})' ) - # Also: text-bison-32k parser.add_argument( '-k', '--private-key', @@ -217,5 +244,5 @@ class Processor(RequestResponseService): def run(): - Processor.launch(module, __doc__) + Processor.launch(default_ident, __doc__)