apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-18 01:31:02 +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

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