VertexAI LLM working

trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py

@@ -13,27 +13,19 @@ from google.oauth2 import service_account
 import google
 
 from vertexai.preview.generative_models import (
-    Content,
-    FunctionDeclaration,
-    GenerativeModel,
-    GenerationConfig,
-    HarmCategory,
-    HarmBlockThreshold,
-    Part,
-    Tool,
+    Content, FunctionDeclaration, GenerativeModel, GenerationConfig,
+    HarmCategory, HarmBlockThreshold, Part, Tool,
 )
 
 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 .... base import ConsumerProducer
 from .... exceptions import TooManyRequests
+from .... base import RequestResponseService
 
 module = "text-completion"
 
-default_input_queue = text_completion_request_queue
-default_output_queue = text_completion_response_queue
 default_subscriber = module
 default_model = 'gemini-1.0-pro-001'
 default_region = 'us-central1'
@@ -41,12 +33,11 @@ default_temperature = 0.0
 default_max_output = 8192
 default_private_key = "private.json"
 
-class Processor(ConsumerProducer):
+class Processor(RequestResponseService):
 
     def __init__(self, **params):
 
-        input_queue = params.get("input_queue", default_input_queue)
-        output_queue = params.get("output_queue", default_output_queue)
+        id = params.get("id")
         subscriber = params.get("subscriber", default_subscriber)
         region = params.get("region", default_region)
         model = params.get("model", default_model)
@@ -59,11 +50,8 @@ class Processor(ConsumerProducer):
 
         super(Processor, self).__init__(
             **params | {
-                "input_queue": input_queue,
-                "output_queue": output_queue,
-                "subscriber": subscriber,
-                "input_schema": TextCompletionRequest,
-                "output_schema": TextCompletionResponse,
+                "request_schema": TextCompletionRequest,
+                "response_schema": TextCompletionResponse,
             }
         )
 
@@ -71,6 +59,7 @@ class Processor(ConsumerProducer):
             __class__.text_completion_metric = Histogram(
                 'text_completion_duration',
                 'Text completion duration (seconds)',
+                ["id", "flow"],
                 buckets=[
                     0.25, 0.5, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0,
                     8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0,
@@ -131,21 +120,16 @@ class Processor(ConsumerProducer):
 
         print("Initialisation complete", flush=True)
 
-    async def handle(self, msg):
+    async def on_request(self, request, consumer, flow):
 
         try:
 
-            v = msg.value()
+            prompt = request.system + "\n\n" + request.prompt
 
-            # Sender-produced ID
-
-            id = msg.properties()["id"]
-
-            print(f"Handling prompt {id}...", flush=True)
-
-            prompt = v.system + "\n\n" + v.prompt
-
-            with __class__.text_completion_metric.time():
+            with __class__.text_completion_metric.labels(
+                    id=self.id,
+                    flow=f"{flow.name}-{consumer.name}",
+            ).time():
 
                 response = self.llm.generate_content(
                     prompt, generation_config=self.generation_config,
@@ -161,7 +145,7 @@ class Processor(ConsumerProducer):
 
             print("Send response...", flush=True)
 
-            r = TextCompletionResponse(
+            return TextCompletionResponse(
                 error=None,
                 response=resp,
                 in_token=inputtokens,
@@ -169,13 +153,6 @@ class Processor(ConsumerProducer):
                 model=self.model
             )
 
-            await self.send(r, properties={"id": id})
-
-            print("Done.", flush=True)
-
-            # Acknowledge successful processing of the message
-            self.consumer.acknowledge(msg)
-
         except google.api_core.exceptions.ResourceExhausted as e:
 
             print("Hit rate limit:", e, flush=True)
@@ -191,7 +168,7 @@ class Processor(ConsumerProducer):
 
             print("Send error response...", flush=True)
 
-            r = TextCompletionResponse(
+            return TextCompletionResponse(
                 error=Error(
                     type = "llm-error",
                     message = str(e),
@@ -202,17 +179,10 @@ class Processor(ConsumerProducer):
                 model=None,
             )
 
-            await self.send(r, properties={"id": id})
-
-            self.consumer.acknowledge(msg)
-
     @staticmethod
     def add_args(parser):
 
-        ConsumerProducer.add_args(
-            parser, default_input_queue, default_subscriber,
-            default_output_queue,
-        )
+        RequestResponseService.add_args(parser, default_subscriber)
 
         parser.add_argument(
             '-m', '--model',