From 6be7b3063385c258dc40c679f051fd70896fef9c Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Tue, 22 Apr 2025 00:27:39 +0100 Subject: [PATCH] Graph RAG working I think --- tests/test-triples | 4 ++- .../trustgraph/base/prompt_client.py | 9 ++++--- .../trustgraph/base/triples_client.py | 12 +++++++-- trustgraph-base/trustgraph/knowledge/defs.py | 4 --- .../retrieval/graph_rag/graph_rag.py | 26 +++++++++---------- 5 files changed, 31 insertions(+), 24 deletions(-) diff --git a/tests/test-triples b/tests/test-triples index 05263d0d..e804d844 100755 --- a/tests/test-triples +++ b/tests/test-triples @@ -3,7 +3,9 @@ import pulsar from trustgraph.clients.triples_query_client import TriplesQueryClient -tq = TriplesQueryClient(pulsar_host="pulsar://localhost:6650") +tq = TriplesQueryClient( + pulsar_host="pulsar://localhost:6650", +) e = "http://trustgraph.ai/e/shuttle" diff --git a/trustgraph-base/trustgraph/base/prompt_client.py b/trustgraph-base/trustgraph/base/prompt_client.py index f1ffa25d..6758bff2 100644 --- a/trustgraph-base/trustgraph/base/prompt_client.py +++ b/trustgraph-base/trustgraph/base/prompt_client.py @@ -29,13 +29,15 @@ class PromptClient(RequestResponse): async def extract_definitions(self, text, timeout=600): return await self.prompt( id = "extract-definitions", - variables = { "text": text } + variables = { "text": text }, + timeout = timeout, ) async def extract_relationships(self, text, timeout=600): return await self.prompt( id = "extract-relationships", - variables = { "text": text } + variables = { "text": text }, + timeout = timeout, ) async def kg_prompt(self, query, kg, timeout=600): @@ -47,7 +49,8 @@ class PromptClient(RequestResponse): { "s": v[0], "p": v[1], "o": v[2] } for v in kg ] - } + }, + timeout = timeout, ) class PromptClientSpec(RequestResponseSpec): diff --git a/trustgraph-base/trustgraph/base/triples_client.py b/trustgraph-base/trustgraph/base/triples_client.py index 56996e91..c9f747b5 100644 --- a/trustgraph-base/trustgraph/base/triples_client.py +++ b/trustgraph-base/trustgraph/base/triples_client.py @@ -1,7 +1,13 @@ from . request_response_spec import RequestResponse, RequestResponseSpec from .. schema import TriplesQueryRequest, TriplesQueryResponse, Value -from .. knowledge import Uri, Literal, Triple +from .. knowledge import Uri, Literal + +class Triple: + def __init__(self, s, p, o): + self.s = s + self.p = p + self.o = o def to_value(x): if x.is_uri: return Uri(x.value) @@ -34,11 +40,13 @@ class TriplesClient(RequestResponse): if resp.error: raise RuntimeError(resp.error.message) - return [ + triples = [ Triple(to_value(v.s), to_value(v.p), to_value(v.o)) for v in resp.triples ] + return triples + class TriplesClientSpec(RequestResponseSpec): def __init__( self, request_name, response_name, diff --git a/trustgraph-base/trustgraph/knowledge/defs.py b/trustgraph-base/trustgraph/knowledge/defs.py index c55f5148..d6290930 100644 --- a/trustgraph-base/trustgraph/knowledge/defs.py +++ b/trustgraph-base/trustgraph/knowledge/defs.py @@ -23,10 +23,6 @@ URL = 'https://schema.org/url' IDENTIFIER = 'https://schema.org/identifier' KEYWORD = 'https://schema.org/keywords' -class Triple: - def __init__(self, s, p, o): - self.s, self.p, self.o = s, p, o - class Uri(str): def is_uri(self): return True def is_literal(self): return False diff --git a/trustgraph-flow/trustgraph/retrieval/graph_rag/graph_rag.py b/trustgraph-flow/trustgraph/retrieval/graph_rag/graph_rag.py index 28d32b6d..41518a61 100644 --- a/trustgraph-flow/trustgraph/retrieval/graph_rag/graph_rag.py +++ b/trustgraph-flow/trustgraph/retrieval/graph_rag/graph_rag.py @@ -39,15 +39,11 @@ class Query: if self.verbose: print("Get entities...", flush=True) - print("HERE>", query, flush=True) - entities = await self.rag.graph_embeddings_client.query( vectors=vectors, limit=self.entity_limit, user=self.user, collection=self.collection, ) - print("ENT>", entities, flush=True) - entities = [ str(e) for e in entities @@ -65,9 +61,9 @@ class Query: if e in self.rag.label_cache: return self.rag.label_cache[e] - res = await self.rag.triples_client.request( - user=self.user, collection=self.collection, + res = await self.rag.triples_client.query( s=e, p=LABEL, o=None, limit=1, + user=self.user, collection=self.collection, ) if len(res) == 0: @@ -100,10 +96,10 @@ class Query: if path_length > 1: await self.follow_edges(str(triple.o), subgraph, path_length-1) - res = await self.rag.triples_client.request( - user=self.user, collection=self.collection, + res = await self.rag.triples_client.query( s=None, p=ent, o=None, - limit=self.triple_limit + limit=self.triple_limit, + user=self.user, collection=self.collection, ) for triple in res: @@ -111,18 +107,20 @@ class Query: (str(triple.s), str(triple.p), str(triple.o)) ) - res = await self.rag.triples_client.request( - user=self.user, collection=self.collection, + res = await self.rag.triples_client.query( s=None, p=None, o=ent, limit=self.triple_limit, + user=self.user, collection=self.collection, ) for triple in res: subgraph.add( - (triple.s.value, triple.p.value, triple.o.value) + (str(triple.s), str(triple.p), str(triple.o)) ) if path_length > 1: - await self.follow_edges(triple.s.value, subgraph, path_length-1) + await self.follow_edges( + str(triple.s), subgraph, path_length-1 + ) async def get_subgraph(self, query): @@ -212,7 +210,7 @@ class GraphRag: print(kg) print(query) - resp = await self.prompt_client.request_kg_prompt(query, kg) + resp = await self.prompt_client.kg_prompt(query, kg) if self.verbose: print("Done", flush=True)