From 8eaaab2d16cc93d6b008196abaa21cf9d223afa2 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Mon, 12 Aug 2024 21:56:34 +0100 Subject: [PATCH] Rework graph RAG implementation to work with new services --- scripts/graph-show | 2 +- scripts/query | 2 - trustgraph/graph_rag.py | 146 +++++++++++++++++++++++----------------- 3 files changed, 85 insertions(+), 65 deletions(-) diff --git a/scripts/graph-show b/scripts/graph-show index b0dcdb19..4340758f 100755 --- a/scripts/graph-show +++ b/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): diff --git a/scripts/query b/scripts/query index caffa602..c8d72aa2 100755 --- a/scripts/query +++ b/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 == "": diff --git a/trustgraph/graph_rag.py b/trustgraph/graph_rag.py index c4ebe158..833c080b 100644 --- a/trustgraph/graph_rag.py +++ b/trustgraph/graph_rag.py @@ -1,12 +1,19 @@ -from . direct.cassandra import TrustGraph -from . direct.milvus import TripleVectors +from . graph_embeddings_client import GraphEmbeddingsClient +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: - graph_hosts = ["cassandra"] + if ge_request_queue is None: + 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( - vector, - limit=self.entity_limit - ) - - print("Obtained", len(res), "entities") - - entities = set([ - item["entity"]["entity"] - for item in res - ]) - - everything.extend(entities) + entities = [ + e.value + for e in 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 = list(res) + res = self.triples_client.request( + 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) - for p, o in res: - subgraph.add((e, p, o)) + res = self.triples_client.request( + e, None, None, + limit=self.query_limit + ) - res = self.graph.get_p(e, limit=self.query_limit) - for s, o in res: - subgraph.add((s, e, o)) + for triple in res: + subgraph.add( + (triple.s.value, triple.p.value, triple.o.value) + ) - res = self.graph.get_o(e, limit=self.query_limit) - for s, p in res: - subgraph.add((s, p, e)) + res = self.triples_client.request( + None, e, None, + 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):