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

Use user/collection in RAG

Browse source
This commit is contained in:
Cyber MacGeddon 2024-10-02 14:29:11 +01:00
parent ac91df91d0
commit 44b1d43989
9 changed files with 182 additions and 136 deletions
Download patch file Download diff file Expand all files Collapse all files

6
trustgraph-base/trustgraph/clients/graph_embeddings_client.py
View file

@ -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

9
trustgraph-base/trustgraph/clients/triples_query_client.py
View file

@ -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

4
trustgraph-base/trustgraph/schema/graph.py
View file

@ -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()

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

@ -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)

5
trustgraph-flow/trustgraph/query/doc_embeddings/qdrant/service.py
View file

@ -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,

5
trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py
View file

@ -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,

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

@ -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)

5
trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py
View file

@ -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:

5
trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py
View file

@ -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: