From 6c3cd607d12ea2c97c974b2798e761ca3837fc87 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Sat, 4 Jan 2025 21:35:22 +0000 Subject: [PATCH] Fix doc RAG processing, put graph-rag and doc-rag in appropriate component files. --- templates/components/document-rag.jsonnet | 30 ++++++ templates/components/graph-rag.jsonnet | 30 ++++++ templates/components/trustgraph.jsonnet | 60 ----------- trustgraph-base/trustgraph/api/api.py | 29 ++++++ .../clients/document_embeddings_client.py | 6 +- .../trustgraph/schema/documents.py | 2 + trustgraph-flow/trustgraph/document_rag.py | 78 ++++++++------- trustgraph-flow/trustgraph/gateway/service.py | 9 ++ .../storage/doc_embeddings/milvus/write.py | 13 ++- .../storage/doc_embeddings/pinecone/write.py | 99 ++++++++++--------- .../storage/doc_embeddings/qdrant/write.py | 69 ++++++------- 11 files changed, 245 insertions(+), 180 deletions(-) diff --git a/templates/components/document-rag.jsonnet b/templates/components/document-rag.jsonnet index 0a68dd52..11dc9c13 100644 --- a/templates/components/document-rag.jsonnet +++ b/templates/components/document-rag.jsonnet @@ -39,5 +39,35 @@ local prompts = import "prompts/mixtral.jsonnet"; }, + "document-embeddings" +: { + + create:: function(engine) + + local container = + engine.container("document-embeddings") + .with_image(images.trustgraph) + .with_command([ + "document-embeddings", + "-p", + url.pulsar, + ]) + .with_limits("1.0", "512M") + .with_reservations("0.5", "512M"); + + local containerSet = engine.containers( + "document-embeddings", [ container ] + ); + + local service = + engine.internalService(containerSet) + .with_port(8000, 8000, "metrics"); + + engine.resources([ + containerSet, + service, + ]) + + }, + } diff --git a/templates/components/graph-rag.jsonnet b/templates/components/graph-rag.jsonnet index 860152c9..eb72754e 100644 --- a/templates/components/graph-rag.jsonnet +++ b/templates/components/graph-rag.jsonnet @@ -138,5 +138,35 @@ local url = import "values/url.jsonnet"; }, + "graph-embeddings" +: { + + create:: function(engine) + + local container = + engine.container("graph-embeddings") + .with_image(images.trustgraph) + .with_command([ + "graph-embeddings", + "-p", + url.pulsar, + ]) + .with_limits("1.0", "512M") + .with_reservations("0.5", "512M"); + + local containerSet = engine.containers( + "graph-embeddings", [ container ] + ); + + local service = + engine.internalService(containerSet) + .with_port(8000, 8000, "metrics"); + + engine.resources([ + containerSet, + service, + ]) + + }, + } diff --git a/templates/components/trustgraph.jsonnet b/templates/components/trustgraph.jsonnet index 1de981a4..541beeed 100644 --- a/templates/components/trustgraph.jsonnet +++ b/templates/components/trustgraph.jsonnet @@ -119,66 +119,6 @@ local prompt = import "prompt-template.jsonnet"; }, - "graph-embeddings" +: { - - create:: function(engine) - - local container = - engine.container("graph-embeddings") - .with_image(images.trustgraph) - .with_command([ - "graph-embeddings", - "-p", - url.pulsar, - ]) - .with_limits("1.0", "512M") - .with_reservations("0.5", "512M"); - - local containerSet = engine.containers( - "graph-embeddings", [ container ] - ); - - local service = - engine.internalService(containerSet) - .with_port(8000, 8000, "metrics"); - - engine.resources([ - containerSet, - service, - ]) - - }, - - "document-embeddings" +: { - - create:: function(engine) - - local container = - engine.container("document-embeddings") - .with_image(images.trustgraph) - .with_command([ - "document-embeddings", - "-p", - url.pulsar, - ]) - .with_limits("1.0", "512M") - .with_reservations("0.5", "512M"); - - local containerSet = engine.containers( - "document-embeddings", [ container ] - ); - - local service = - engine.internalService(containerSet) - .with_port(8000, 8000, "metrics"); - - engine.resources([ - containerSet, - service, - ]) - - }, - "metering" +: { create:: function(engine) diff --git a/trustgraph-base/trustgraph/api/api.py b/trustgraph-base/trustgraph/api/api.py index de96499c..24207f32 100644 --- a/trustgraph-base/trustgraph/api/api.py +++ b/trustgraph-base/trustgraph/api/api.py @@ -131,6 +131,35 @@ class Api: except: raise ProtocolException(f"Response not formatted correctly") + def document_rag(self, question): + + # The input consists of a question + input = { + "query": question + } + + url = f"{self.url}document-rag" + + # Invoke the API, input is passed as JSON + resp = requests.post(url, json=input) + + # Should be a 200 status code + if resp.status_code != 200: + raise ProtocolException(f"Status code {resp.status_code}") + + try: + # Parse the response as JSON + object = resp.json() + except: + raise ProtocolException(f"Expected JSON response") + + self.check_error(resp) + + try: + return object["response"] + except: + raise ProtocolException(f"Response not formatted correctly") + def embeddings(self, text): # The input consists of a text block diff --git a/trustgraph-base/trustgraph/clients/document_embeddings_client.py b/trustgraph-base/trustgraph/clients/document_embeddings_client.py index d432991d..5b6d324e 100644 --- a/trustgraph-base/trustgraph/clients/document_embeddings_client.py +++ b/trustgraph-base/trustgraph/clients/document_embeddings_client.py @@ -38,8 +38,12 @@ class DocumentEmbeddingsClient(BaseClient): output_schema=DocumentEmbeddingsResponse, ) - 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 ).documents diff --git a/trustgraph-base/trustgraph/schema/documents.py b/trustgraph-base/trustgraph/schema/documents.py index 46eee250..fd0049ee 100644 --- a/trustgraph-base/trustgraph/schema/documents.py +++ b/trustgraph-base/trustgraph/schema/documents.py @@ -55,6 +55,8 @@ document_embeddings_store_queue = topic('document-embeddings-store') class DocumentEmbeddingsRequest(Record): vectors = Array(Array(Double())) limit = Integer() + user = String() + collection = String() class DocumentEmbeddingsResponse(Record): error = Error() diff --git a/trustgraph-flow/trustgraph/document_rag.py b/trustgraph-flow/trustgraph/document_rag.py index f3c8b158..86298783 100644 --- a/trustgraph-flow/trustgraph/document_rag.py +++ b/trustgraph-flow/trustgraph/document_rag.py @@ -16,6 +16,44 @@ from . schema import document_embeddings_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_docs(self, query): + + vectors = self.get_vector(query) + + if self.verbose: + print("Get entities...", flush=True) + + docs = self.rag.de_client.request( + vectors, limit=self.rag.doc_limit + ) + + if self.verbose: + print("Docs:", flush=True) + for doc in docs: + print(doc, flush=True) + + return docs + class DocumentRag: def __init__( @@ -55,7 +93,7 @@ class DocumentRag: print("Initialising...", flush=True) # FIXME: Configurable - self.entity_limit = 20 + self.doc_limit = 20 self.de_client = DocumentEmbeddingsClient( pulsar_host=pulsar_host, @@ -81,42 +119,16 @@ class DocumentRag: 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_docs(self, query): - - vectors = self.get_vector(query) - - if self.verbose: - print("Get entities...", flush=True) - - docs = self.de_client.request( - vectors, self.entity_limit - ) - - if self.verbose: - print("Docs:", flush=True) - for doc in docs: - print(doc, flush=True) - - return docs - - def query(self, query): + def query(self, query, user="trustgraph", collection="default"): if self.verbose: print("Construct prompt...", flush=True) - docs = self.get_docs(query) + q = Query( + rag=self, user=user, collection=collection, verbose=self.verbose + ) + + docs = q.get_docs(query) if self.verbose: print("Invoke LLM...", flush=True) diff --git a/trustgraph-flow/trustgraph/gateway/service.py b/trustgraph-flow/trustgraph/gateway/service.py index a260b631..d6306ac6 100755 --- a/trustgraph-flow/trustgraph/gateway/service.py +++ b/trustgraph-flow/trustgraph/gateway/service.py @@ -31,6 +31,7 @@ from . subscriber import Subscriber from . text_completion import TextCompletionRequestor from . prompt import PromptRequestor from . graph_rag import GraphRagRequestor +from . document_rag import DocumentRagRequestor from . triples_query import TriplesQueryRequestor from . graph_embeddings_query import GraphEmbeddingsQueryRequestor from . embeddings import EmbeddingsRequestor @@ -91,6 +92,10 @@ class Api: pulsar_host=self.pulsar_host, timeout=self.timeout, auth = self.auth, ), + "document-rag": DocumentRagRequestor( + pulsar_host=self.pulsar_host, timeout=self.timeout, + auth = self.auth, + ), "triples-query": TriplesQueryRequestor( pulsar_host=self.pulsar_host, timeout=self.timeout, auth = self.auth, @@ -140,6 +145,10 @@ class Api: endpoint_path = "/api/v1/graph-rag", auth=self.auth, requestor = self.services["graph-rag"], ), + ServiceEndpoint( + endpoint_path = "/api/v1/document-rag", auth=self.auth, + requestor = self.services["document-rag"], + ), ServiceEndpoint( endpoint_path = "/api/v1/triples-query", auth=self.auth, requestor = self.services["triples-query"], diff --git a/trustgraph-flow/trustgraph/storage/doc_embeddings/milvus/write.py b/trustgraph-flow/trustgraph/storage/doc_embeddings/milvus/write.py index 15836f56..bfa6c123 100755 --- a/trustgraph-flow/trustgraph/storage/doc_embeddings/milvus/write.py +++ b/trustgraph-flow/trustgraph/storage/doc_embeddings/milvus/write.py @@ -39,11 +39,16 @@ class Processor(Consumer): v = msg.value() - chunk = v.chunk.decode("utf-8") + for emb in v.chunks: - if v.chunk != "" and v.chunk is not None: - for vec in v.vectors: - self.vecstore.insert(vec, chunk) + chunk = emb.chunk.decode("utf-8") + if chunk == "" or chunk is None: continue + + for vec in emb.vectors: + + if chunk != "" and v.chunk is not None: + for vec in v.vectors: + self.vecstore.insert(vec, chunk) @staticmethod def add_args(parser): diff --git a/trustgraph-flow/trustgraph/storage/doc_embeddings/pinecone/write.py b/trustgraph-flow/trustgraph/storage/doc_embeddings/pinecone/write.py index fc87aa72..c59ecd7b 100644 --- a/trustgraph-flow/trustgraph/storage/doc_embeddings/pinecone/write.py +++ b/trustgraph-flow/trustgraph/storage/doc_embeddings/pinecone/write.py @@ -65,71 +65,74 @@ class Processor(Consumer): v = msg.value() - chunk = v.chunk.decode("utf-8") + for emb in v.chunks: - if chunk == "": return + chunk = emb.chunk.decode("utf-8") + if chunk == "" or chunk is None: continue - for vec in v.vectors: + for vec in emb.vectors: - dim = len(vec) - collection = ( - "d-" + v.metadata.user + "-" + str(dim) - ) + for vec in v.vectors: - if index_name != self.last_index_name: + dim = len(vec) + collection = ( + "d-" + v.metadata.user + "-" + str(dim) + ) - if not self.pinecone.has_index(index_name): + if index_name != self.last_index_name: - try: + if not self.pinecone.has_index(index_name): - self.pinecone.create_index( - name = index_name, - dimension = dim, - metric = "cosine", - spec = ServerlessSpec( - cloud = self.cloud, - region = self.region, - ) - ) + try: - for i in range(0, 1000): + self.pinecone.create_index( + name = index_name, + dimension = dim, + metric = "cosine", + spec = ServerlessSpec( + cloud = self.cloud, + region = self.region, + ) + ) - if self.pinecone.describe_index( - index_name - ).status["ready"]: - break + for i in range(0, 1000): - time.sleep(1) + if self.pinecone.describe_index( + index_name + ).status["ready"]: + break - if not self.pinecone.describe_index( - index_name - ).status["ready"]: - raise RuntimeError( - "Gave up waiting for index creation" - ) + time.sleep(1) - except Exception as e: - print("Pinecone index creation failed") - raise e + if not self.pinecone.describe_index( + index_name + ).status["ready"]: + raise RuntimeError( + "Gave up waiting for index creation" + ) - print(f"Index {index_name} created", flush=True) + except Exception as e: + print("Pinecone index creation failed") + raise e - self.last_index_name = index_name + print(f"Index {index_name} created", flush=True) - index = self.pinecone.Index(index_name) + self.last_index_name = index_name - records = [ - { - "id": id, - "values": vec, - "metadata": { "doc": chunk }, - } - ] + index = self.pinecone.Index(index_name) - index.upsert( - vectors = records, - namespace = v.metadata.collection, - ) + records = [ + { + "id": id, + "values": vec, + "metadata": { "doc": chunk }, + } + ] + + index.upsert( + vectors = records, + namespace = v.metadata.collection, + ) @staticmethod def add_args(parser): diff --git a/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py b/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py index 683d8534..f852e03b 100644 --- a/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py +++ b/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py @@ -44,47 +44,48 @@ class Processor(Consumer): v = msg.value() - chunk = v.chunk.decode("utf-8") + for emb in v.chunks: - if chunk == "": return + chunk = emb.chunk.decode("utf-8") + if chunk == "": return - for vec in v.vectors: + for vec in emb.vectors: - dim = len(vec) - collection = ( - "d_" + v.metadata.user + "_" + v.metadata.collection + "_" + - str(dim) - ) + dim = len(vec) + collection = ( + "d_" + v.metadata.user + "_" + v.metadata.collection + "_" + + str(dim) + ) - if collection != self.last_collection: + if collection != self.last_collection: - if not self.client.collection_exists(collection): + if not self.client.collection_exists(collection): - try: - self.client.create_collection( - collection_name=collection, - vectors_config=VectorParams( - size=dim, distance=Distance.COSINE - ), + try: + self.client.create_collection( + collection_name=collection, + vectors_config=VectorParams( + size=dim, distance=Distance.COSINE + ), + ) + except Exception as e: + print("Qdrant collection creation failed") + raise e + + self.last_collection = collection + + self.client.upsert( + collection_name=collection, + points=[ + PointStruct( + id=str(uuid.uuid4()), + vector=vec, + payload={ + "doc": chunk, + } ) - except Exception as e: - print("Qdrant collection creation failed") - raise e - - self.last_collection = collection - - self.client.upsert( - collection_name=collection, - points=[ - PointStruct( - id=str(uuid.uuid4()), - vector=vec, - payload={ - "doc": chunk, - } - ) - ] - ) + ] + ) @staticmethod def add_args(parser):