From 38cea4c26d63fa0e3c651d6b3b6e13fa53b1b09d Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Wed, 16 Apr 2025 17:43:49 +0100 Subject: [PATCH] VertexAI LLM working --- Makefile | 20 +++--- tests/test-embeddings | 9 ++- tests/test-llm | 7 +- trustgraph-base/trustgraph/base/__init__.py | 3 + .../trustgraph/base/request_response.py | 39 ++++++++++- .../embeddings/fastembed/processor.py | 44 +++---------- .../embeddings/graph_embeddings/embeddings.py | 51 ++++++++------- .../model/text_completion/vertexai/llm.py | 64 +++++-------------- 8 files changed, 118 insertions(+), 119 deletions(-) diff --git a/Makefile b/Makefile index 78e0ff4d..a449d36b 100644 --- a/Makefile +++ b/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 \ -# -t ${CONTAINER_BASE}/trustgraph-bedrock:${VERSION} . -# ${DOCKER} build -f containers/Containerfile.vertexai \ -# -t ${CONTAINER_BASE}/trustgraph-vertexai:${VERSION} . -# ${DOCKER} build -f containers/Containerfile.hf \ -# -t ${CONTAINER_BASE}/trustgraph-hf:${VERSION} . -# ${DOCKER} build -f containers/Containerfile.ocr \ -# -t ${CONTAINER_BASE}/trustgraph-ocr:${VERSION} . + ${DOCKER} build -f containers/Containerfile.bedrock \ + -t ${CONTAINER_BASE}/trustgraph-bedrock:${VERSION} . + ${DOCKER} build -f containers/Containerfile.vertexai \ + -t ${CONTAINER_BASE}/trustgraph-vertexai:${VERSION} . + ${DOCKER} build -f containers/Containerfile.hf \ + -t ${CONTAINER_BASE}/trustgraph-hf:${VERSION} . + ${DOCKER} build -f containers/Containerfile.ocr \ + -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 \ diff --git a/tests/test-embeddings b/tests/test-embeddings index 3855fcf0..5fcd31e6 100755 --- a/tests/test-embeddings +++ b/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) - - diff --git a/tests/test-llm b/tests/test-llm index 4e86387a..61f0da6b 100755 --- a/tests/test-llm +++ b/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" diff --git a/trustgraph-base/trustgraph/base/__init__.py b/trustgraph-base/trustgraph/base/__init__.py index 1b04d576..f6dd670c 100644 --- a/trustgraph-base/trustgraph/base/__init__.py +++ b/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 + + diff --git a/trustgraph-base/trustgraph/base/request_response.py b/trustgraph-base/trustgraph/base/request_response.py index 0f652603..a97f0cd0 100644 --- a/trustgraph-base/trustgraph/base/request_response.py +++ b/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(): diff --git a/trustgraph-flow/trustgraph/embeddings/fastembed/processor.py b/trustgraph-flow/trustgraph/embeddings/fastembed/processor.py index f90771d4..46716364 100755 --- a/trustgraph-flow/trustgraph/embeddings/fastembed/processor.py +++ b/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 ... log_level import LogLevel -from ... base import ConsumerProducer +from ... base import RequestResponseService + 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() - - # Sender-produced ID - - id = msg.properties()["id"] - - print(f"Handling input {id}...", flush=True) - - text = v.text + text = request.text vecs = self.embeddings.embed([text]) vecs = [ @@ -59,23 +41,15 @@ class Processor(ConsumerProducer): for v in vecs ] - print("Send response...", flush=True) - r = EmbeddingsResponse( + return 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( - parser, default_input_queue, default_subscriber, - default_output_queue, - ) + RequestResponseService.add_args(parser, default_subscriber) parser.add_argument( '-m', '--model', diff --git a/trustgraph-flow/trustgraph/embeddings/graph_embeddings/embeddings.py b/trustgraph-flow/trustgraph/embeddings/graph_embeddings/embeddings.py index a61bc9c1..61606405 100755 --- a/trustgraph-flow/trustgraph/embeddings/graph_embeddings/embeddings.py +++ b/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, - "output_queue": output_queue, - "embeddings_request_queue": emb_request_queue, - "embeddings_response_queue": emb_response_queue, + "id": id, "subscriber": subscriber, - "input_schema": EntityContexts, - "output_schema": GraphEmbeddings, } ) - self.embeddings = EmbeddingsClient( - pulsar_host=self.pulsar_host, - input_queue=emb_request_queue, - output_queue=emb_response_queue, - subscriber=module + "-emb", + self.register_consumer( + name = "input", + schema = EntityContexts, + handler = self.on_message, ) + 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() diff --git a/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py b/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py index f8351c7c..bc774997 100755 --- a/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py +++ b/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',