diff --git a/tests/test-graph-embeddings b/tests/test-graph-embeddings new file mode 100755 index 00000000..742e7f81 --- /dev/null +++ b/tests/test-graph-embeddings @@ -0,0 +1,23 @@ +#!/usr/bin/env python3 + +import pulsar +from trustgraph.graph_embeddings_client import GraphEmbeddingsClient +from trustgraph.embeddings_client import EmbeddingsClient + +ec = EmbeddingsClient(pulsar_host="pulsar://localhost:6650") + +vectors = ec.request("What caused the space shuttle to explode?") + +print(vectors) + +llm = GraphEmbeddingsClient(pulsar_host="pulsar://localhost:6650") + +limit=10 + +resp = llm.request(vectors, limit) + +print("Response...") +for val in resp: + print(val.value) + + diff --git a/trustgraph/graph_embeddings_client.py b/trustgraph/graph_embeddings_client.py new file mode 100644 index 00000000..c3fbb58a --- /dev/null +++ b/trustgraph/graph_embeddings_client.py @@ -0,0 +1,89 @@ +#!/usr/bin/env python3 + +import pulsar +import _pulsar +from pulsar.schema import JsonSchema +import hashlib +import uuid + +from . schema import GraphEmbeddingsRequest, GraphEmbeddingsResponse +from . schema import graph_embeddings_request_queue +from . schema import graph_embeddings_response_queue + +# Ugly +ERROR=_pulsar.LoggerLevel.Error +WARN=_pulsar.LoggerLevel.Warn +INFO=_pulsar.LoggerLevel.Info +DEBUG=_pulsar.LoggerLevel.Debug + +class GraphEmbeddingsClient: + + def __init__( + self, log_level=ERROR, + subscriber=None, + input_queue=None, + output_queue=None, + pulsar_host="pulsar://pulsar:6650", + ): + + if input_queue == None: + input_queue = graph_embeddings_request_queue + + if output_queue == None: + output_queue = graph_embeddings_response_queue + + if subscriber == None: + subscriber = str(uuid.uuid4()) + + self.client = pulsar.Client( + pulsar_host, + logger=pulsar.ConsoleLogger(log_level), + ) + + self.producer = self.client.create_producer( + topic=input_queue, + schema=JsonSchema(GraphEmbeddingsRequest), + chunking_enabled=True, + ) + + self.consumer = self.client.subscribe( + output_queue, subscriber, + schema=JsonSchema(GraphEmbeddingsResponse), + ) + + def request(self, vectors, limit=10, timeout=500): + + id = str(uuid.uuid4()) + + r = GraphEmbeddingsRequest( + vectors=vectors, + limit=limit, + ) + + self.producer.send(r, properties={ "id": id }) + + while True: + + msg = self.consumer.receive(timeout_millis=timeout * 1000) + + mid = msg.properties()["id"] + + if mid == id: + resp = msg.value().entities + self.consumer.acknowledge(msg) + return resp + + # Ignore messages with wrong ID + self.consumer.acknowledge(msg) + + def __del__(self): + + if hasattr(self, "consumer"): + self.consumer.close() + + if hasattr(self, "producer"): + self.producer.flush() + self.producer.close() + + self.client.close() + diff --git a/trustgraph/query/graph_embeddings/milvus/service.py b/trustgraph/query/graph_embeddings/milvus/service.py index 456ae587..c95f2ef5 100755 --- a/trustgraph/query/graph_embeddings/milvus/service.py +++ b/trustgraph/query/graph_embeddings/milvus/service.py @@ -4,11 +4,11 @@ Graph embeddings query service. Input is vector, output is list of e """ -from ... direct.milvus import TripleVectors -from ... schema import GraphEmbeddingsRequest, GraphEmbeddingsResponse -from ... schema import graph_embeddings_request_queue -from ... schema import graph_embeddings_response_queue -from ... base import ConsumerProducer +from .... direct.milvus import TripleVectors +from .... schema import GraphEmbeddingsRequest, GraphEmbeddingsResponse, Value +from .... schema import graph_embeddings_request_queue +from .... schema import graph_embeddings_response_queue +from .... base import ConsumerProducer module = ".".join(__name__.split(".")[1:-1]) @@ -55,7 +55,21 @@ class Processor(ConsumerProducer): resp = self.vecstore.search(vec, limit=v.limit) for r in resp: - entities.add(r) + ent = r["entity"]["entity"] + entities.add(ent) + + # Convert set to list + entities = list(entities) + + ents2 = [] + + for ent in entities: + if ent.startswith("http://") or ent.startswith("https://"): + ents2.append(Value(value=ent, is_uri=True)) + else: + ents2.append(Value(value=ent, is_uri=False)) + + entities = ents2 print("Send response...", flush=True) r = GraphEmbeddingsResponse(entities=entities) diff --git a/trustgraph/schema.py b/trustgraph/schema.py index f0dbb8ff..537a5e6a 100644 --- a/trustgraph/schema.py +++ b/trustgraph/schema.py @@ -73,11 +73,11 @@ graph_embeddings_store_queue = topic('graph-embeddings-store') # Graph embeddings query -class GraphEmbeddingsQueryRequest(Record): +class GraphEmbeddingsRequest(Record): vectors = Array(Array(Double())) limit = Integer() -class GraphEmbeddingsQueryResponse(Record): +class GraphEmbeddingsResponse(Record): entities = Array(Value()) graph_embeddings_request_queue = topic(