Rate limit working

Makefile

@@ -65,8 +65,8 @@ some-containers:
 	    -t ${CONTAINER_BASE}/trustgraph-base:${VERSION} .
 	${DOCKER} build -f containers/Containerfile.flow \
 	    -t ${CONTAINER_BASE}/trustgraph-flow:${VERSION} .
-#	${DOCKER} build -f containers/Containerfile.vertexai \
-#	    -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 \

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

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__)