apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-17 09:11:03 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

VertexAI LLM working

Browse source
This commit is contained in:
Cyber MacGeddon 2025-04-16 17:43:49 +01:00
parent b94b4b7389
commit 38cea4c26d
8 changed files with 118 additions and 119 deletions
Download patch file Download diff file Expand all files Collapse all files

20
Makefile
View file

@ -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 \

9
tests/test-embeddings
View file

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

7
tests/test-llm
View file

@ -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"

3
trustgraph-base/trustgraph/base/__init__.py
View file

@ -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

39
trustgraph-base/trustgraph/base/request_response.py
View file

@ -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():

44
trustgraph-flow/trustgraph/embeddings/fastembed/processor.py
View file

@ -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',

51
trustgraph-flow/trustgraph/embeddings/graph_embeddings/embeddings.py
View file

@ -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()

64
trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py
View file

@ -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',