mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-06 03:42:11 +02:00
Rework graph RAG implementation to work with new services
This commit is contained in:
Cyber MacGeddon
parent
72f36b34b2
commit
8eaaab2d16
3 changed files with 85 additions and 65 deletions
|
|
@ -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):
|
||||
|
||||
|
|
|
|||
|
|
@ -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 == "":
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue