Rework graph RAG implementation to work with new services

scripts/graph-show

@@ -7,7 +7,7 @@ Connects to the trustgraph graph hosts and dumps all graph edges.
 import argparse
 import time
 
-from trustgraph.trustgraph import TrustGraph
+from trustgraph.direct.cassandra import TrustGraph
 
 def show_graph(graph_hosts):
 

scripts/query

@@ -7,9 +7,7 @@ query = " ".join(sys.argv[1:])
 
 gr = GraphRag(
     verbose=True,
-    vector_store="http://localhost:19530",
     pulsar_host="pulsar://localhost:6650",
-    graph_hosts=["localhost"],
 )
 
 if query == "":

trustgraph/graph_rag.py

@@ -1,12 +1,19 @@
 
-from . direct.cassandra import TrustGraph
+from . graph_embeddings_client import GraphEmbeddingsClient
-from . direct.milvus import TripleVectors
+from . triples_query_client import TriplesQueryClient
 from . llm_client import LlmClient
 from . embeddings_client import EmbeddingsClient
+
+from . schema import GraphEmbeddingsRequest, GraphEmbeddingsResponse
+from . schema import TriplesQueryRequest, TriplesQueryResponse
 from . schema import text_completion_request_queue
 from . schema import text_completion_response_queue
 from . schema import embeddings_request_queue
 from . schema import embeddings_response_queue
+from . schema import graph_embeddings_request_queue
+from . schema import graph_embeddings_response_queue
+from . schema import triples_request_queue
+from . schema import triples_response_queue
 
 LABEL="http://www.w3.org/2000/01/rdf-schema#label"
 DEFINITION="http://www.w3.org/2004/02/skos/core#definition"
@@ -15,13 +22,15 @@ class GraphRag:
 
     def __init__(
             self,
-            graph_hosts=None,
             pulsar_host="pulsar://pulsar:6650",
-            vector_store="http://milvus:19530",
             completion_request_queue=None,
             completion_response_queue=None,
             emb_request_queue=None,
             emb_response_queue=None,
+            ge_request_queue=None,
+            ge_response_queue=None,
+            tpl_request_queue=None,
+            tpl_response_queue=None,
             verbose=False,
             entity_limit=50,
             triple_limit=30,
@@ -31,25 +40,46 @@ class GraphRag:
 
         self.verbose=verbose
 
-        if completion_request_queue == None:
+        if completion_request_queue is None:
             completion_request_queue = text_completion_request_queue
 
-        if completion_response_queue == None:
+        if completion_response_queue is None:
             completion_response_queue = text_completion_response_queue
 
-        if emb_request_queue == None:
+        if emb_request_queue is None:
             emb_request_queue = embeddings_request_queue
 
-        if emb_response_queue == None:
+        if emb_response_queue is None:
             emb_response_queue = embeddings_response_queue
 
-        if graph_hosts == None:
+        if ge_request_queue is None:
-            graph_hosts = ["cassandra"]
+            ge_request_queue = graph_embeddings_request_queue
+
+        if ge_response_queue is None:
+            ge_response_queue = graph_embeddings_response_queue
+
+        if tpl_request_queue is None:
+            tpl_request_queue = triples_request_queue
+
+        if tpl_response_queue is None:
+            tpl_response_queue = triples_response_queue
 
         if self.verbose:
             print("Initialising...", flush=True)
 
-        self.graph = TrustGraph(graph_hosts)
+        self.ge_client = GraphEmbeddingsClient(
+            pulsar_host=pulsar_host,
+            subscriber=module + "-ge",
+            input_queue=ge_request_queue,
+            output_queue=ge_response_queue,
+        )            
+
+        self.triples_client = TriplesQueryClient(
+            pulsar_host=pulsar_host,
+            subscriber=module + "-tpl",
+            input_queue=tpl_request_queue,
+            output_queue=tpl_response_queue
+        )
 
         self.embeddings = EmbeddingsClient(
             pulsar_host=pulsar_host,
@@ -58,8 +88,6 @@ class GraphRag:
             subscriber=module + "-emb",
         )
 
-        self.vecstore = TripleVectors(vector_store)
-
         self.entity_limit=entity_limit
         self.query_limit=triple_limit
         self.max_subgraph_size=max_subgraph_size
@@ -90,70 +118,43 @@ class GraphRag:
 
     def get_entities(self, query):
 
-        everything = []
-
         vectors = self.get_vector(query)
 
         if self.verbose:
             print("Get entities...", flush=True)
 
-        for vector in vectors:
+        entities = self.ge_client.request(
+            vectors, self.entity_limit
+        )
 
-            res = self.vecstore.search(
+        entities = [
-                vector,
+            e.value
-                limit=self.entity_limit
+            for e in entities
-            )
+        ]
-
-            print("Obtained", len(res), "entities")
-
-            entities = set([
-                item["entity"]["entity"]
-                for item in res
-            ])
-
-            everything.extend(entities)
 
         if self.verbose:
             print("Entities:", flush=True)
-            for ent in everything:
+            for ent in entities:
                 print(" ", ent, flush=True)
 
-        return everything
+        return entities
         
     def maybe_label(self, e):
 
         if e in self.label_cache:
             return self.label_cache[e]
 
-        res = self.graph.get_sp(e, LABEL)
+        res = self.triples_client.request(
-        res = list(res)
+            e, LABEL, None, limit=1
+        )
 
         if len(res) == 0:
             self.label_cache[e] = e
             return e
 
-        self.label_cache[e] = res[0][0]
+        self.label_cache[e] = res[0].o.value
         return self.label_cache[e]
 
-    def get_nodes(self, query):
-
-        ents = self.get_entities(query)
-
-        if self.verbose:
-            print("Get labels...", flush=True)
-
-        nodes = [
-            self.maybe_label(e)
-            for e in ents
-        ]
-
-        if self.verbose:
-            print("Nodes:", flush=True)
-            for node in nodes:
-                print(" ", node, flush=True)
-
-        return nodes
-
     def get_subgraph(self, query):
 
         entities = self.get_entities(query)
@@ -165,17 +166,35 @@ class GraphRag:
 
         for e in entities:
 
-            res = self.graph.get_s(e, limit=self.query_limit)
+            res = self.triples_client.request(
-            for p, o in res:
+                e, None, None,
-                subgraph.add((e, p, o))
+                limit=self.query_limit
+            )
 
-            res = self.graph.get_p(e, limit=self.query_limit)
+            for triple in res:
-            for s, o in res:
+                subgraph.add(
-                subgraph.add((s, e, o))
+                    (triple.s.value, triple.p.value, triple.o.value)
+                )
 
-            res = self.graph.get_o(e, limit=self.query_limit)
+            res = self.triples_client.request(
-            for s, p in res:
+                None, e, None,
-                subgraph.add((s, p, e))
+                limit=self.query_limit
+            )
+
+            for triple in res:
+                subgraph.add(
+                    (triple.s.value, triple.p.value, triple.o.value)
+                )
+
+            res = self.triples_client.request(
+                None, None, e,
+                limit=self.query_limit
+            )
+
+            for triple in res:
+                subgraph.add(
+                    (triple.s.value, triple.p.value, triple.o.value)
+                )
 
         subgraph = list(subgraph)
 
@@ -223,6 +242,9 @@ class GraphRag:
         kg = "\n".join(sg2)
         kg = kg.replace("\\", "-")
 
+        if self.verbose:
+            print(kg)
+
         return kg
 
     def get_graph_prompt(self, query):