- Added graph_embeddings_client.py

- Added tests/test-graph-embeddings to demonstrate client works
- Fixed schema

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

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

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 .... direct.milvus import TripleVectors
-from ... schema import GraphEmbeddingsRequest, GraphEmbeddingsResponse
+from .... schema import GraphEmbeddingsRequest, GraphEmbeddingsResponse, Value
-from ... schema import graph_embeddings_request_queue
+from .... schema import graph_embeddings_request_queue
-from ... schema import graph_embeddings_response_queue
+from .... schema import graph_embeddings_response_queue
-from ... base import ConsumerProducer
+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)

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(