From 44b1d439894a57cda7e674a0c5280b931cb59004 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Wed, 2 Oct 2024 14:29:11 +0100 Subject: [PATCH] Use user/collection in RAG --- .../clients/graph_embeddings_client.py | 6 +- .../clients/triples_query_client.py | 9 +- trustgraph-base/trustgraph/schema/graph.py | 4 + trustgraph-flow/trustgraph/graph_rag.py | 275 ++++++++++-------- .../query/doc_embeddings/qdrant/service.py | 5 +- .../query/graph_embeddings/qdrant/service.py | 5 +- .../trustgraph/retrieval/graph_rag/rag.py | 4 +- .../storage/doc_embeddings/qdrant/write.py | 5 +- .../storage/graph_embeddings/qdrant/write.py | 5 +- 9 files changed, 182 insertions(+), 136 deletions(-) diff --git a/trustgraph-base/trustgraph/clients/graph_embeddings_client.py b/trustgraph-base/trustgraph/clients/graph_embeddings_client.py index bb1358fc..401266bc 100644 --- a/trustgraph-base/trustgraph/clients/graph_embeddings_client.py +++ b/trustgraph-base/trustgraph/clients/graph_embeddings_client.py @@ -38,8 +38,12 @@ class GraphEmbeddingsClient(BaseClient): output_schema=GraphEmbeddingsResponse, ) - def request(self, vectors, limit=10, timeout=300): + def request( + self, vectors, user="trustgraph", collection="default", + limit=10, timeout=300 + ): return self.call( + user=user, collection=collection, vectors=vectors, limit=limit, timeout=timeout ).entities diff --git a/trustgraph-base/trustgraph/clients/triples_query_client.py b/trustgraph-base/trustgraph/clients/triples_query_client.py index 14b75151..6c18ac3f 100644 --- a/trustgraph-base/trustgraph/clients/triples_query_client.py +++ b/trustgraph-base/trustgraph/clients/triples_query_client.py @@ -48,11 +48,18 @@ class TriplesQueryClient(BaseClient): return Value(value=ent, is_uri=False) - def request(self, s, p, o, limit=10, timeout=60): + def request( + self, + s, p, o, + user="trustgraph", collection="default", + limit=10, timeout=60, + ): return self.call( s=self.create_value(s), p=self.create_value(p), o=self.create_value(o), + user=user, + collection=collection, limit=limit, timeout=timeout, ).triples diff --git a/trustgraph-base/trustgraph/schema/graph.py b/trustgraph-base/trustgraph/schema/graph.py index e4a7fa32..107478a4 100644 --- a/trustgraph-base/trustgraph/schema/graph.py +++ b/trustgraph-base/trustgraph/schema/graph.py @@ -23,6 +23,8 @@ graph_embeddings_store_queue = topic('graph-embeddings-store') class GraphEmbeddingsRequest(Record): vectors = Array(Array(Double())) limit = Integer() + user = String() + collection = String() class GraphEmbeddingsResponse(Record): error = Error() @@ -56,6 +58,8 @@ class TriplesQueryRequest(Record): p = Value() o = Value() limit = Integer() + user = String() + collection = String() class TriplesQueryResponse(Record): error = Error() diff --git a/trustgraph-flow/trustgraph/graph_rag.py b/trustgraph-flow/trustgraph/graph_rag.py index 15acb609..f69ebeb7 100644 --- a/trustgraph-flow/trustgraph/graph_rag.py +++ b/trustgraph-flow/trustgraph/graph_rag.py @@ -18,6 +18,144 @@ 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" +class Query: + + def __init__(self, rag, user, collection, verbose): + self.rag = rag + self.user = user + self.collection = collection + self.verbose = verbose + + def get_vector(self, query): + + if self.verbose: + print("Compute embeddings...", flush=True) + + qembeds = self.rag.embeddings.request(query) + + if self.verbose: + print("Done.", flush=True) + + return qembeds + + def get_entities(self, query): + + vectors = self.get_vector(query) + + if self.verbose: + print("Get entities...", flush=True) + + entities = self.rag.ge_client.request( + user=self.user, collection=self.collection, + vectors=vectors, limit=self.rag.entity_limit, + ) + + entities = [ + e.value + for e in entities + ] + + if self.verbose: + print("Entities:", flush=True) + for ent in entities: + print(" ", ent, flush=True) + + return entities + + def maybe_label(self, e): + + if e in self.rag.label_cache: + return self.rag.label_cache[e] + + res = self.rag.triples_client.request( + user=self.user, collection=self.collection, + s=e, p=LABEL, o=None, limit=1, + ) + + if len(res) == 0: + self.rag.label_cache[e] = e + return e + + self.rag.label_cache[e] = res[0].o.value + return self.rag.label_cache[e] + + def get_subgraph(self, query): + + entities = self.get_entities(query) + + subgraph = set() + + if self.verbose: + print("Get subgraph...", flush=True) + + for e in entities: + + res = self.rag.triples_client.request( + user=self.user, collection=self.collection, + s=e, p=None, o=None, + limit=self.rag.query_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.rag.query_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.rag.query_limit, + ) + + for triple in res: + subgraph.add( + (triple.s.value, triple.p.value, triple.o.value) + ) + + subgraph = list(subgraph) + + subgraph = subgraph[0:self.rag.max_subgraph_size] + + if self.verbose: + print("Subgraph:", flush=True) + for edge in subgraph: + print(" ", str(edge), flush=True) + + if self.verbose: + print("Done.", flush=True) + + return subgraph + + def get_labelgraph(self, query): + + subgraph = self.get_subgraph(query) + + sg2 = [] + + for edge in subgraph: + + if edge[1] == LABEL: + continue + + s = self.maybe_label(edge[0]) + p = self.maybe_label(edge[1]) + o = self.maybe_label(edge[2]) + + sg2.append((s, p, o)) + + return sg2 + class GraphRag: def __init__( @@ -94,7 +232,7 @@ class GraphRag: self.label_cache = {} - self.lang = PromptClient( + self.prompt = PromptClient( pulsar_host=pulsar_host, input_queue=pr_request_queue, output_queue=pr_response_queue, @@ -104,144 +242,23 @@ class GraphRag: if self.verbose: print("Initialised", flush=True) - def get_vector(self, query): - - if self.verbose: - print("Compute embeddings...", flush=True) - - qembeds = self.embeddings.request(query) - - if self.verbose: - print("Done.", flush=True) - - return qembeds - - def get_entities(self, query): - - vectors = self.get_vector(query) - - if self.verbose: - print("Get entities...", flush=True) - - entities = self.ge_client.request( - vectors, self.entity_limit - ) - - entities = [ - e.value - for e in entities - ] - - if self.verbose: - print("Entities:", flush=True) - for ent in entities: - print(" ", ent, flush=True) - - return entities - - def maybe_label(self, e): - - if e in self.label_cache: - return self.label_cache[e] - - 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].o.value - return self.label_cache[e] - - def get_subgraph(self, query): - - entities = self.get_entities(query) - - subgraph = set() - - if self.verbose: - print("Get subgraph...", flush=True) - - for e in entities: - - res = self.triples_client.request( - e, None, 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, 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) - - subgraph = subgraph[0:self.max_subgraph_size] - - if self.verbose: - print("Subgraph:", flush=True) - for edge in subgraph: - print(" ", str(edge), flush=True) - - if self.verbose: - print("Done.", flush=True) - - return subgraph - - def get_labelgraph(self, query): - - subgraph = self.get_subgraph(query) - - sg2 = [] - - for edge in subgraph: - - if edge[1] == LABEL: - continue - - s = self.maybe_label(edge[0]) - p = self.maybe_label(edge[1]) - o = self.maybe_label(edge[2]) - - sg2.append((s, p, o)) - - return sg2 - - def query(self, query): + def query(self, query, user="trustgraph", collection="default"): if self.verbose: print("Construct prompt...", flush=True) - kg = self.get_labelgraph(query) + q = Query( + rag=self, user=user, collection=collection, verbose=self.verbose + ) + + kg = q.get_labelgraph(query) if self.verbose: print("Invoke LLM...", flush=True) print(kg) print(query) - resp = self.lang.request_kg_prompt(query, kg) + resp = self.prompt.request_kg_prompt(query, kg) if self.verbose: print("Done", flush=True) diff --git a/trustgraph-flow/trustgraph/query/doc_embeddings/qdrant/service.py b/trustgraph-flow/trustgraph/query/doc_embeddings/qdrant/service.py index 5cc41437..cd350a32 100755 --- a/trustgraph-flow/trustgraph/query/doc_embeddings/qdrant/service.py +++ b/trustgraph-flow/trustgraph/query/doc_embeddings/qdrant/service.py @@ -60,7 +60,10 @@ class Processor(ConsumerProducer): for vec in v.vectors: dim = len(vec) - collection = "doc_" + str(dim) + collection = ( + "d_" + v.user + "_" + v.collection + + str(dim) + ) search_result = self.client.query_points( collection_name=collection, diff --git a/trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py b/trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py index e61a00a7..0f9a7b48 100755 --- a/trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py +++ b/trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py @@ -66,7 +66,10 @@ class Processor(ConsumerProducer): for vec in v.vectors: dim = len(vec) - collection = "triples_" + str(dim) + collection = ( + "t_" + v.user + "_" + v.collection + + str(dim) + ) search_result = self.client.query_points( collection_name=collection, diff --git a/trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py b/trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py index 87f8d24f..1219050e 100755 --- a/trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py +++ b/trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py @@ -108,7 +108,9 @@ class Processor(ConsumerProducer): print(f"Handling input {id}...", flush=True) - response = self.rag.query(v.query) + response = self.rag.query( + query=v.query, user=v.user, collection=v.collection + ) print("Send response...", flush=True) r = GraphRagResponse(response = response, error=None) diff --git a/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py b/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py index f22ae74a..8e2ee43f 100644 --- a/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py +++ b/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py @@ -52,7 +52,10 @@ class Processor(Consumer): for vec in v.vectors: dim = len(vec) - collection = "doc_" + str(dim) + collection = ( + "d_" + v.metadata.user + "_" + v.metadata.collection + + str(dim) + ) if dim != self.last_dim: diff --git a/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py b/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py index 95448750..8f9964f6 100755 --- a/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py +++ b/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py @@ -50,7 +50,10 @@ class Processor(Consumer): for vec in v.vectors: dim = len(vec) - collection = "triples_" + str(dim) + collection = ( + "t_" + v.metadata.user + "_" + v.metadata.collection + + str(dim) + ) if dim != self.last_dim: