VertexAI LLM working

Makefile

@@ -51,14 +51,18 @@ container: update-package-versions
 	    -t ${CONTAINER_BASE}/trustgraph-base:${VERSION} .
 	${DOCKER} build -f containers/Containerfile.flow \
 	    -t ${CONTAINER_BASE}/trustgraph-flow:${VERSION} .
-#	${DOCKER} build -f containers/Containerfile.bedrock \
+	${DOCKER} build -f containers/Containerfile.bedrock \
-#	    -t ${CONTAINER_BASE}/trustgraph-bedrock:${VERSION} .
+	    -t ${CONTAINER_BASE}/trustgraph-bedrock:${VERSION} .
-#	${DOCKER} build -f containers/Containerfile.vertexai \
+	${DOCKER} build -f containers/Containerfile.vertexai \
-#	    -t ${CONTAINER_BASE}/trustgraph-vertexai:${VERSION} .
+	    -t ${CONTAINER_BASE}/trustgraph-vertexai:${VERSION} .
-#	${DOCKER} build -f containers/Containerfile.hf \
+	${DOCKER} build -f containers/Containerfile.hf \
-#	    -t ${CONTAINER_BASE}/trustgraph-hf:${VERSION} .
+	    -t ${CONTAINER_BASE}/trustgraph-hf:${VERSION} .
-#	${DOCKER} build -f containers/Containerfile.ocr \
+	${DOCKER} build -f containers/Containerfile.ocr \
-#	    -t ${CONTAINER_BASE}/trustgraph-ocr:${VERSION} .
+	    -t ${CONTAINER_BASE}/trustgraph-ocr:${VERSION} .
+
+some-containers:
+	${DOCKER} build -f containers/Containerfile.vertexai \
+	    -t ${CONTAINER_BASE}/trustgraph-vertexai:${VERSION} .
 
 basic-containers: update-package-versions
 	${DOCKER} build -f containers/Containerfile.base \

tests/test-embeddings

@@ -3,7 +3,12 @@
 import pulsar
 from trustgraph.clients.embeddings_client import EmbeddingsClient
 
-embed = EmbeddingsClient(pulsar_host="pulsar://localhost:6650")
+embed = EmbeddingsClient(
+    pulsar_host="pulsar://pulsar:6650",
+    input_queue="non-persistent://tg/request/embeddings:default",
+    output_queue="non-persistent://tg/response/embeddings:default",
+    subscriber="test1",
+)
 
 prompt="Write a funny limerick about a llama"
 
@@ -11,5 +16,3 @@ resp = embed.request(prompt)
 
 print(resp)
 
-
-

tests/test-llm

@@ -3,7 +3,12 @@
 import pulsar
 from trustgraph.clients.llm_client import LlmClient
 
-llm = LlmClient(pulsar_host="pulsar://localhost:6650")
+llm = LlmClient(
+    pulsar_host="pulsar://pulsar:6650",
+    input_queue="non-persistent://tg/request/text-completion:default",
+    output_queue="non-persistent://tg/response/text-completion:default",
+    subscriber="test1",
+)
 
 system = "You are a lovely assistant."
 prompt="Write a funny limerick about a llama"

trustgraph-base/trustgraph/base/__init__.py

@@ -7,4 +7,7 @@ from . publisher import Publisher
 from . subscriber import Subscriber
 from . metrics import ProcessorMetrics, ConsumerMetrics, ProducerMetrics
 from . flow_processor import FlowProcessor
+from . request_response import RequestResponseService
+
+
 

trustgraph-base/trustgraph/base/request_response.py

@@ -17,9 +17,11 @@ class RequestResponseService(FlowProcessor):
 
         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_request)
+            ("request", params.get("request_schema"), self.on_message)
         ]
         self.producer_spec = [
             ("response", params.get("response_schema"))
@@ -27,10 +29,43 @@ class RequestResponseService(FlowProcessor):
 
         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, default_subscriber):
 
-        FlowProcessor.add_args(parser)
+        FlowProcessor.add_args(parser, default_subscriber)
 
 def run():
 

trustgraph-flow/trustgraph/embeddings/fastembed/processor.py

@@ -5,53 +5,35 @@ Input is text, output is embeddings vector.
 """
 
 from ... schema import EmbeddingsRequest, EmbeddingsResponse
-from ... schema import embeddings_request_queue, embeddings_response_queue
+from ... base import RequestResponseService
-from ... log_level import LogLevel
+
-from ... base import ConsumerProducer
 from fastembed import TextEmbedding
 import os
 
 module = "embeddings"
 
-default_input_queue = embeddings_request_queue
-default_output_queue = embeddings_response_queue
 default_subscriber = module
 default_model="sentence-transformers/all-MiniLM-L6-v2"
 
-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)
-        subscriber = params.get("subscriber", default_subscriber)
-
         model = params.get("model", default_model)
 
         super(Processor, self).__init__(
             **params | {
-                "input_queue": input_queue,
-                "output_queue": output_queue,
-                "subscriber": subscriber,
-                "input_schema": EmbeddingsRequest,
-                "output_schema": EmbeddingsResponse,
                 "model": model,
+                "request_schema": EmbeddingsRequest,
+                "response_schema": EmbeddingsResponse,
             }
         )
 
         self.embeddings = TextEmbedding(model_name = model)
 
-    async def handle(self, msg):
+    async def on_request(self, request, consumer, flow):
 
-        v = msg.value()
+        text = request.text
-
-        # Sender-produced ID
-
-        id = msg.properties()["id"]
-
-        print(f"Handling input {id}...", flush=True)
-
-        text = v.text
         vecs = self.embeddings.embed([text])
 
         vecs = [
@@ -59,23 +41,15 @@ class Processor(ConsumerProducer):
             for v in vecs
         ]
 
-        print("Send response...", flush=True)
+        return EmbeddingsResponse(
-        r = EmbeddingsResponse(
             vectors=list(vecs),
             error=None,
         )
 
-        await self.send(r, properties={"id": id})
-
-        print("Done.", flush=True)
-
     @staticmethod
     def add_args(parser):
 
-        ConsumerProducer.add_args(
+        RequestResponseService.add_args(parser, default_subscriber)
-            parser, default_input_queue, default_subscriber,
-            default_output_queue,
-        )
 
         parser.add_argument(
             '-m', '--model',

trustgraph-flow/trustgraph/embeddings/graph_embeddings/embeddings.py

@@ -11,47 +11,52 @@ from ... schema import graph_embeddings_store_queue
 from ... schema import embeddings_request_queue, embeddings_response_queue
 from ... clients.embeddings_client import EmbeddingsClient
 from ... log_level import LogLevel
-from ... base import ConsumerProducer
+from ... base import FlowProcessor
 
 module = "graph-embeddings"
 
-default_input_queue = entity_contexts_ingest_queue
-default_output_queue = graph_embeddings_store_queue
 default_subscriber = module
 
-class Processor(ConsumerProducer):
+class Processor(FlowProcessor):
 
     def __init__(self, **params):
 
-        input_queue = params.get("input_queue", default_input_queue)
+
-        output_queue = params.get("output_queue", default_output_queue)
+
+                "input_schema": EntityContexts,
+                "output_schema": GraphEmbeddings,
+        
+
+
+
+        id = params.get("id")
         subscriber = params.get("subscriber", default_subscriber)
-        emb_request_queue = params.get(
-            "embeddings_request_queue", embeddings_request_queue
-        )
-        emb_response_queue = params.get(
-            "embeddings_response_queue", embeddings_response_queue
-        )
 
         super(Processor, self).__init__(
             **params | {
-                "input_queue": input_queue,
+                "id": id,
-                "output_queue": output_queue,
-                "embeddings_request_queue": emb_request_queue,
-                "embeddings_response_queue": emb_response_queue,
                 "subscriber": subscriber,
-                "input_schema": EntityContexts,
-                "output_schema": GraphEmbeddings,
             }
         )
 
-        self.embeddings = EmbeddingsClient(
+        self.register_consumer(
-            pulsar_host=self.pulsar_host,
+            name = "input",
-            input_queue=emb_request_queue,
+            schema = EntityContexts,
-            output_queue=emb_response_queue,
+            handler = self.on_message,
-            subscriber=module + "-emb",
         )
 
+        self.register_producer(
+            name = "output",
+            schema = GraphEmbeddings,
+        )
+
+        # self.embeddings = EmbeddingsClient(
+        #     pulsar_host=self.pulsar_host,
+        #     input_queue=emb_request_queue,
+        #     output_queue=emb_response_queue,
+        #     subscriber=module + "-emb",
+        # )
+
     async def handle(self, msg):
 
         v = msg.value()

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,
+    Content, FunctionDeclaration, GenerativeModel, GenerationConfig,
-    FunctionDeclaration,
+    HarmCategory, HarmBlockThreshold, Part, Tool,
-    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)
+        id = params.get("id")
-        output_queue = params.get("output_queue", default_output_queue)
         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,
+                "request_schema": TextCompletionRequest,
-                "output_queue": output_queue,
+                "response_schema": TextCompletionResponse,
-                "subscriber": subscriber,
-                "input_schema": TextCompletionRequest,
-                "output_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
+            with __class__.text_completion_metric.labels(
-
+                    id=self.id,
-            id = msg.properties()["id"]
+                    flow=f"{flow.name}-{consumer.name}",
-
+            ).time():
-            print(f"Handling prompt {id}...", flush=True)
-
-            prompt = v.system + "\n\n" + v.prompt
-
-            with __class__.text_completion_metric.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(
+        RequestResponseService.add_args(parser, default_subscriber)
-            parser, default_input_queue, default_subscriber,
-            default_output_queue,
-        )
 
         parser.add_argument(
             '-m', '--model',