Hacking wth graph rag

tests/test-graph-rag

@@ -3,7 +3,12 @@
 import pulsar
 from trustgraph.clients.graph_rag_client import GraphRagClient
 
-rag = GraphRagClient(pulsar_host="pulsar://localhost:6650")
+rag = GraphRagClient(
+    pulsar_host="pulsar://localhost:6650",
+    subscriber="test1",
+    input_queue = "non-persistent://tg/request/graph-rag:default",
+    output_queue = "non-persistent://tg/response/graph-rag:default",
+)
 
 query="""
 This knowledge graph describes the Space Shuttle disaster.

trustgraph-base/trustgraph/base/prompt_client.py

@@ -38,6 +38,18 @@ class PromptClient(RequestResponse):
             variables = { "text": text }
         )
 
+    async def kg_prompt(self, query, kg, timeout=600):
+        return await self.prompt(
+            id = "kg-prompt",
+            variables = {
+                "query": query,
+                "knowledge": [
+                    { "s": v[0], "p": v[1], "o": v[2] }
+                    for v in kg
+                ]
+            }
+        )
+
 class PromptClientSpec(RequestResponseSpec):
     def __init__(
             self, request_name, response_name,

trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py

@@ -40,7 +40,7 @@ class Processor(GraphEmbeddingsQueryService):
         else:
             return Value(value=ent, is_uri=False)
         
-    async def handle(self, msg):
+    async def query_graph_embeddings(self, msg):
 
         try:
 

trustgraph-flow/trustgraph/retrieval/graph_rag/graph_rag.py

@@ -25,7 +25,7 @@ class Query:
         if self.verbose:
             print("Compute embeddings...", flush=True)
 
-        qembeds = await  self.rag.embeddings.request(query)
+        qembeds = await  self.rag.embeddings_client.embed(query)
 
         if self.verbose:
             print("Done.", flush=True)
@@ -39,11 +39,13 @@ class Query:
         if self.verbose:
             print("Get entities...", flush=True)
 
-        entities = await self.rag.ge_client.request(
-            user=self.user, collection=self.collection,
+        entities = await self.rag.graph_embeddings_client.query(
             vectors=vectors, limit=self.entity_limit,
+            user=self.user, collection=self.collection,
         )
 
+        print("ENT>", entities, flush=True)
+
         entities = [
             e.value
             for e in entities
@@ -83,10 +85,10 @@ class Query:
         if len(subgraph) >= self.max_subgraph_size:
             return
 
-        res = await self.rag.triples_client.request(
-            user=self.user, collection=self.collection,
+        res = await self.rag.triples_client.query(
             s=ent, p=None, o=None,
-            limit=self.triple_limit
+            limit=self.triple_limit,
+            user=self.user, collection=self.collection,
         )
 
         for triple in res:
@@ -208,7 +210,7 @@ class GraphRag:
             print(kg)
             print(query)
 
-        resp = await self.prompt.request_kg_prompt(query, kg)
+        resp = await self.prompt_client.request_kg_prompt(query, kg)
 
         if self.verbose:
             print("Done", flush=True)

trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py

@@ -154,10 +154,6 @@ class Processor(FlowProcessor):
                 properties = {"id": id}
             )
 
-            await self.send(r, properties={"id": id})
-
-            self.consumer.acknowledge(msg)
-
     @staticmethod
     def add_args(parser):