From 7d7b2686f1ee8b1452a885605b83fbb318bc8a9c Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Thu, 13 Mar 2025 00:30:41 +0000 Subject: [PATCH] Multi-hop path traversal GraphRAG --- trustgraph-base/trustgraph/api/api.py | 4 +- .../trustgraph/schema/retrieval.py | 1 + trustgraph-cli/scripts/tg-invoke-graph-rag | 15 +++- .../trustgraph/gateway/graph_rag.py | 1 + trustgraph-flow/trustgraph/graph_rag.py | 85 +++++++++++-------- .../trustgraph/retrieval/graph_rag/rag.py | 23 ++++- 6 files changed, 86 insertions(+), 43 deletions(-) diff --git a/trustgraph-base/trustgraph/api/api.py b/trustgraph-base/trustgraph/api/api.py index 68c4b147..4c72d3ca 100644 --- a/trustgraph-base/trustgraph/api/api.py +++ b/trustgraph-base/trustgraph/api/api.py @@ -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" diff --git a/trustgraph-base/trustgraph/schema/retrieval.py b/trustgraph-base/trustgraph/schema/retrieval.py index a81d57cc..caeb8e67 100644 --- a/trustgraph-base/trustgraph/schema/retrieval.py +++ b/trustgraph-base/trustgraph/schema/retrieval.py @@ -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() diff --git a/trustgraph-cli/scripts/tg-invoke-graph-rag b/trustgraph-cli/scripts/tg-invoke-graph-rag index dedec5e6..5bbe5f59 100755 --- a/trustgraph-cli/scripts/tg-invoke-graph-rag +++ b/trustgraph-cli/scripts/tg-invoke-graph-rag @@ -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: diff --git a/trustgraph-flow/trustgraph/gateway/graph_rag.py b/trustgraph-flow/trustgraph/gateway/graph_rag.py index fc02a17f..b2b69758 100644 --- a/trustgraph-flow/trustgraph/gateway/graph_rag.py +++ b/trustgraph-flow/trustgraph/gateway/graph_rag.py @@ -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): diff --git a/trustgraph-flow/trustgraph/graph_rag.py b/trustgraph-flow/trustgraph/graph_rag.py index 3d52afe4..f4e4a69c 100644 --- a/trustgraph-flow/trustgraph/graph_rag.py +++ b/trustgraph-flow/trustgraph/graph_rag.py @@ -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) diff --git a/trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py b/trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py index d94d9a8f..2c45ecd4 100755 --- a/trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py +++ b/trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py @@ -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(