apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-20 18:51:03 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

Multi-hop path traversal GraphRAG

Browse source
This commit is contained in:
Cyber MacGeddon 2025-03-13 00:30:41 +00:00
parent c8df9a4fa2
commit 7d7b2686f1
6 changed files with 86 additions and 43 deletions
Download patch file Download diff file Expand all files Collapse all files

4
trustgraph-base/trustgraph/api/api.py
View file

@ -104,7 +104,8 @@ class Api:
def graph_rag(
self, question, user="trustgraph", collection="default",
entity_limit=50, triple_limit=30, max_subgraph_size=1000,
entity_limit=50, triple_limit=30, max_subgraph_size=150,
max_path_length=2,
):
# The input consists of a question
@ -115,6 +116,7 @@ class Api:
"entity-limit": entity_limit,
"triple-limit": triple_limit,
"max-subgraph-size": max_subgraph_size,
"max-path-length": max_path_length,
}
url = f"{self.url}graph-rag"

1
trustgraph-base/trustgraph/schema/retrieval.py
View file

@ -14,6 +14,7 @@ class GraphRagQuery(Record):
entity_limit = Integer()
triple_limit = Integer()
max_subgraph_size = Integer()
max_path_length = Integer()
class GraphRagResponse(Record):
error = Error()

15
trustgraph-cli/scripts/tg-invoke-graph-rag
View file

@ -13,11 +13,12 @@ default_user = 'trustgraph'
default_collection = 'default'
default_entity_limit = 50
default_triple_limit = 30
default_max_subgraph_size = 1000
default_max_subgraph_size = 150
default_max_path_length = 2
def question(
url, question, user, collection, entity_limit, triple_limit,
max_subgraph_size
max_subgraph_size, max_path_length
):
rag = Api(url)
@ -25,7 +26,8 @@ def question(
resp = rag.graph_rag(
question=question, user=user, collection=collection,
entity_limit=entity_limit, triple_limit=triple_limit,
max_subgraph_size=max_subgraph_size
max_subgraph_size=max_subgraph_size,
max_path_length=max_path_length
)
print(resp)
@ -79,6 +81,12 @@ def main():
help=f'Max subgraph size (default: {default_max_subgraph_size})'
)
parser.add_argument(
'-p', '--max-path-length',
default=default_max_path_length,
help=f'Max path length (default: {default_max_path_length})'
)
args = parser.parse_args()
try:
@ -91,6 +99,7 @@ def main():
entity_limit=args.entity_limit,
triple_limit=args.triple_limit,
max_subgraph_size=args.max_subgraph_size,
max_path_length=args.max_path_length,
)
except Exception as e:

1
trustgraph-flow/trustgraph/gateway/graph_rag.py
View file

@ -26,6 +26,7 @@ class GraphRagRequestor(ServiceRequestor):
entity_limit=int(body.get("entity-limit", 50)),
triple_limit=int(body.get("triple-limit", 30)),
max_subgraph_size=int(body.get("max-subgraph-size", 1000)),
max_path_length=int(body.get("max-path-length", 2)),
)
def from_response(self, message):

85
trustgraph-flow/trustgraph/graph_rag.py
View file

@ -23,6 +23,7 @@ class Query:
def __init__(
self, rag, user, collection, verbose,
entity_limit=50, triple_limit=30, max_subgraph_size=1000,
max_path_length=2,
):
self.rag = rag
self.user = user
@ -31,6 +32,7 @@ class Query:
self.entity_limit = entity_limit
self.triple_limit = triple_limit
self.max_subgraph_size = max_subgraph_size
self.max_path_length = max_path_length
def get_vector(self, query):
@ -85,49 +87,60 @@ class Query:
self.rag.label_cache[e] = res[0].o.value
return self.rag.label_cache[e]
def follow_edges(self, ent, subgraph, path_length):
# Not needed?
if path_length <= 0:
return
res = self.rag.triples_client.request(
user=self.user, collection=self.collection,
s=ent, p=None, o=None,
limit=self.triple_limit
)
for triple in res:
subgraph.add(
(triple.s.value, triple.p.value, triple.o.value)
)
if path_length > 1:
self.follow_edges(triple.o.value, subgraph, path_length-1)
res = self.rag.triples_client.request(
user=self.user, collection=self.collection,
s=None, p=ent, o=None,
limit=self.triple_limit
)
for triple in res:
subgraph.add(
(triple.s.value, triple.p.value, triple.o.value)
)
res = self.rag.triples_client.request(
user=self.user, collection=self.collection,
s=None, p=None, o=ent,
limit=self.triple_limit,
)
for triple in res:
subgraph.add(
(triple.s.value, triple.p.value, triple.o.value)
)
if path_length > 1:
self.follow_edges(triple.s.value, subgraph, path_length-1)
def get_subgraph(self, query):
entities = self.get_entities(query)
subgraph = set()
if self.verbose:
print("Get subgraph...", flush=True)
for e in entities:
subgraph = set()
res = self.rag.triples_client.request(
user=self.user, collection=self.collection,
s=e, p=None, o=None,
limit=self.triple_limit
)
for triple in res:
subgraph.add(
(triple.s.value, triple.p.value, triple.o.value)
)
res = self.rag.triples_client.request(
user=self.user, collection=self.collection,
s=None, p=e, o=None,
limit=self.triple_limit
)
for triple in res:
subgraph.add(
(triple.s.value, triple.p.value, triple.o.value)
)
res = self.rag.triples_client.request(
user=self.user, collection=self.collection,
s=None, p=None, o=e,
limit=self.triple_limit,
)
for triple in res:
subgraph.add(
(triple.s.value, triple.p.value, triple.o.value)
)
for ent in entities:
self.follow_edges(ent, subgraph, self.max_path_length)
subgraph = list(subgraph)
@ -249,6 +262,7 @@ class GraphRag:
def query(
self, query, user="trustgraph", collection="default",
entity_limit=50, triple_limit=30, max_subgraph_size=1000,
max_path_length=2,
):
if self.verbose:
@ -258,6 +272,7 @@ class GraphRag:
rag=self, user=user, collection=collection, verbose=self.verbose,
entity_limit=entity_limit, triple_limit=triple_limit,
max_subgraph_size=max_subgraph_size,
max_path_length=max_path_length,
)
kg = q.get_labelgraph(query)

23
trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py
View file

@ -59,7 +59,8 @@ class Processor(ConsumerProducer):
entity_limit = params.get("entity_limit", 50)
triple_limit = params.get("triple_limit", 30)
max_subgraph_size = params.get("max_subgraph_size", 1000)
max_subgraph_size = params.get("max_subgraph_size", 150)
max_path_length = params.get("max_path_length", 2)
super(Processor, self).__init__(
**params | {
@ -100,6 +101,7 @@ class Processor(ConsumerProducer):
self.default_entity_limit = entity_limit
self.default_triple_limit = triple_limit
self.default_max_subgraph_size = max_subgraph_size
self.default_max_path_length = max_path_length
async def handle(self, msg):
@ -127,10 +129,16 @@ class Processor(ConsumerProducer):
else:
max_subgraph_size = self.default_max_subgraph_size
if v.max_path_length:
max_path_length = v.max_path_length
else:
max_path_length = self.default_max_path_length
response = self.rag.query(
query=v.query, user=v.user, collection=v.collection,
entity_limit=entity_limit, triple_limit=triple_limit,
max_subgraph_size=max_subgraph_size
max_subgraph_size=max_subgraph_size,
max_path_length=max_path_length,
)
print("Send response...", flush=True)
@ -182,8 +190,15 @@ class Processor(ConsumerProducer):
parser.add_argument(
'-u', '--max-subgraph-size',
type=int,
default=1000,
help=f'Default max subgraph size (default: 1000)'
default=150,
help=f'Default max subgraph size (default: 150)'
)
parser.add_argument(
'-a', '--max-path-length',
type=int,
default=2,
help=f'Default max path length (default: 2)'
)
parser.add_argument(