From 179285e1865345dbf980d56974b472c7663870d1 Mon Sep 17 00:00:00 2001
From: Cyber MacGeddon <cybermaggedon@gmail.com>
Date: Mon, 21 Apr 2025 18:52:52 +0100
Subject: [PATCH] Hacking wth graph rag

---
 tests/test-graph-rag                             |  7 ++++++-
 trustgraph-base/trustgraph/base/prompt_client.py | 12 ++++++++++++
 .../query/graph_embeddings/qdrant/service.py     |  2 +-
 .../trustgraph/retrieval/graph_rag/graph_rag.py  | 16 +++++++++-------
 .../trustgraph/retrieval/graph_rag/rag.py        |  4 ----
 5 files changed, 28 insertions(+), 13 deletions(-)

diff --git a/tests/test-graph-rag b/tests/test-graph-rag
index 036f73f4..ac407f42 100755
--- a/tests/test-graph-rag
+++ b/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.
diff --git a/trustgraph-base/trustgraph/base/prompt_client.py b/trustgraph-base/trustgraph/base/prompt_client.py
index 74cec56b..f1ffa25d 100644
--- a/trustgraph-base/trustgraph/base/prompt_client.py
+++ b/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,
diff --git a/trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py b/trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py
index d818f672..d6ab67d2 100755
--- a/trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py
+++ b/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:
 
diff --git a/trustgraph-flow/trustgraph/retrieval/graph_rag/graph_rag.py b/trustgraph-flow/trustgraph/retrieval/graph_rag/graph_rag.py
index d115a51e..c5e60760 100644
--- a/trustgraph-flow/trustgraph/retrieval/graph_rag/graph_rag.py
+++ b/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)
diff --git a/trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py b/trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py
index 99143118..5d3cc2f4 100755
--- a/trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py
+++ b/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):