diff --git a/trustgraph-flow/trustgraph/gateway/document_embeddings_load.py b/trustgraph-flow/trustgraph/gateway/document_embeddings_load.py index 1a7f635d..6de41a35 100644 --- a/trustgraph-flow/trustgraph/gateway/document_embeddings_load.py +++ b/trustgraph-flow/trustgraph/gateway/document_embeddings_load.py @@ -15,7 +15,7 @@ from . serialize import to_subgraph class DocumentEmbeddingsLoadEndpoint(SocketEndpoint): def __init__( - self, pulsar_host, auth, path="/api/v1/load/document-embeddings", + self, pulsar_host, auth, path="/api/v1/load/document-embeddings", pulsar_api_key=None ): super(DocumentEmbeddingsLoadEndpoint, self).__init__( @@ -23,10 +23,12 @@ class DocumentEmbeddingsLoadEndpoint(SocketEndpoint): ) self.pulsar_host=pulsar_host + self.pulsar_api_key=pulsar_api_key self.publisher = Publisher( self.pulsar_host, document_embeddings_store_queue, - schema=JsonSchema(DocumentEmbeddings) + schema=JsonSchema(DocumentEmbeddings), + pulsar_api_key=self.pulsar_api_key ) async def start(self): diff --git a/trustgraph-flow/trustgraph/gateway/document_embeddings_stream.py b/trustgraph-flow/trustgraph/gateway/document_embeddings_stream.py index 99cfb0a9..252a4926 100644 --- a/trustgraph-flow/trustgraph/gateway/document_embeddings_stream.py +++ b/trustgraph-flow/trustgraph/gateway/document_embeddings_stream.py @@ -14,7 +14,7 @@ from . serialize import serialize_document_embeddings class DocumentEmbeddingsStreamEndpoint(SocketEndpoint): def __init__( - self, pulsar_host, auth, path="/api/v1/stream/document-embeddings" + self, pulsar_host, auth, path="/api/v1/stream/document-embeddings", pulsar_api_key=None ): super(DocumentEmbeddingsStreamEndpoint, self).__init__( @@ -22,11 +22,13 @@ class DocumentEmbeddingsStreamEndpoint(SocketEndpoint): ) self.pulsar_host=pulsar_host + self.pulsar_api_key=pulsar_api_key self.subscriber = Subscriber( self.pulsar_host, document_embeddings_store_queue, "api-gateway", "api-gateway", - schema=JsonSchema(DocumentEmbeddings) + schema=JsonSchema(DocumentEmbeddings), + pulsar_api_key=self.pulsar_api_key ) async def listener(self, ws, running): diff --git a/trustgraph-flow/trustgraph/gateway/document_rag.py b/trustgraph-flow/trustgraph/gateway/document_rag.py index 6665a7dc..d3231a5c 100644 --- a/trustgraph-flow/trustgraph/gateway/document_rag.py +++ b/trustgraph-flow/trustgraph/gateway/document_rag.py @@ -7,10 +7,11 @@ from . endpoint import ServiceEndpoint from . requestor import ServiceRequestor class DocumentRagRequestor(ServiceRequestor): - def __init__(self, pulsar_host, timeout, auth): + def __init__(self, pulsar_host, timeout, auth, pulsar_api_key=None): super(DocumentRagRequestor, self).__init__( pulsar_host=pulsar_host, + pulsar_api_key=pulsar_api_key, request_queue=document_rag_request_queue, response_queue=document_rag_response_queue, request_schema=DocumentRagQuery, diff --git a/trustgraph-flow/trustgraph/gateway/graph_embeddings_load.py b/trustgraph-flow/trustgraph/gateway/graph_embeddings_load.py index 0d4b12e9..b2d5b80c 100644 --- a/trustgraph-flow/trustgraph/gateway/graph_embeddings_load.py +++ b/trustgraph-flow/trustgraph/gateway/graph_embeddings_load.py @@ -27,7 +27,7 @@ class GraphEmbeddingsLoadEndpoint(SocketEndpoint): self.publisher = Publisher( self.pulsar_host, graph_embeddings_store_queue, - self.pulsar_api_key, + pulsar_api_key=self.pulsar_api_key, schema=JsonSchema(GraphEmbeddings) ) diff --git a/trustgraph-flow/trustgraph/processing/processing.py b/trustgraph-flow/trustgraph/processing/processing.py index e5e9593f..5352776a 100644 --- a/trustgraph-flow/trustgraph/processing/processing.py +++ b/trustgraph-flow/trustgraph/processing/processing.py @@ -69,6 +69,7 @@ class Processing: params = { "pulsar_host": self.pulsar_host, + "pulsar_api_key": self.pulsar_api_key, "log_level": str(self.log_level), } @@ -161,6 +162,7 @@ def run(): try: p = Processing( pulsar_host=args.pulsar_host, + pulsar_api_key=args.pulsar_api_key, file=args.file, log_level=args.log_level, )