Add pulsar API token check

trustgraph-base/trustgraph/base/base_processor.py

@@ -11,6 +11,7 @@ from .. log_level import LogLevel
 class BaseProcessor:
 
     default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://pulsar:6650')
+    default_pulsar_api_key = os.getenv("PULSAR_API_KEY", None)
 
     def __init__(self, **params):
 
@@ -28,14 +29,22 @@ class BaseProcessor:
         })
 
         pulsar_host = params.get("pulsar_host", self.default_pulsar_host)
+        pulsar_api_key = params.get("pulsar_api_key", None)
         log_level = params.get("log_level", LogLevel.INFO)
 
         self.pulsar_host = pulsar_host
-
+        if pulsar_api_key:
-        self.client = pulsar.Client(
+            auth = pulsar.AuthenticationToken(pulsar_api_key)
+            self.client = pulsar.Client(
+            pulsar_host,
+            authentication=auth,
+            logger=pulsar.ConsoleLogger(log_level.to_pulsar())
+            )
+        else:
+            self.client = pulsar.Client(
             pulsar_host,
             logger=pulsar.ConsoleLogger(log_level.to_pulsar())
-        )
+            )
 
     def __del__(self):
 
@@ -51,6 +60,12 @@ class BaseProcessor:
             default=__class__.default_pulsar_host,
             help=f'Pulsar host (default: {__class__.default_pulsar_host})',
         )
+        
+        parser.add_argument(
+            '--pulsar-api-key',
+            default=__class__.default_pulsar_api_key,
+            help=f'Pulsar API key',
+        )
 
         parser.add_argument(
             '-l', '--log-level',

trustgraph-base/trustgraph/clients/agent_client.py

@@ -20,6 +20,7 @@ class AgentClient(BaseClient):
             input_queue=None,
             output_queue=None,
             pulsar_host="pulsar://pulsar:6650",
+            pulsar_api_key=None,
     ):
 
         if input_queue is None: input_queue = agent_request_queue
@@ -33,6 +34,7 @@ class AgentClient(BaseClient):
             pulsar_host=pulsar_host,
             input_schema=AgentRequest,
             output_schema=AgentResponse,
+            pulsar_api_key=pulsar_api_key
         )
 
     def request(

trustgraph-base/trustgraph/clients/base.py

@@ -27,6 +27,7 @@ class BaseClient:
             input_schema=None,
             output_schema=None,
             pulsar_host="pulsar://pulsar:6650",
+            pulsar_api_key=None,
     ):
 
         if input_queue == None: raise RuntimeError("Need input_queue")
@@ -37,10 +38,18 @@ class BaseClient:
         if subscriber == None:
             subscriber = str(uuid.uuid4())
 
-        self.client = pulsar.Client(
+        if pulsar_api_key:
+            auth = pulsar.AuthenticationToken(pulsar_api_key)
+            self.client = pulsar.Client(
             pulsar_host,
-            logger=pulsar.ConsoleLogger(log_level),
+            authentication=auth,
-        )
+            logger=pulsar.ConsoleLogger(log_level.to_pulsar())
+            )
+        else:
+            self.client = pulsar.Client(
+            pulsar_host,
+            logger=pulsar.ConsoleLogger(log_level.to_pulsar())
+            )
 
         self.producer = self.client.create_producer(
             topic=input_queue,

trustgraph-base/trustgraph/clients/document_embeddings_client.py

@@ -20,6 +20,7 @@ class DocumentEmbeddingsClient(BaseClient):
             input_queue=None,
             output_queue=None,
             pulsar_host="pulsar://pulsar:6650",
+            pulsar_api_key=None,
     ):
 
         if input_queue == None:
@@ -34,6 +35,7 @@ class DocumentEmbeddingsClient(BaseClient):
             input_queue=input_queue,
             output_queue=output_queue,
             pulsar_host=pulsar_host,
+            pulsar_api_key=pulsar_api_key,
             input_schema=DocumentEmbeddingsRequest,
             output_schema=DocumentEmbeddingsResponse,
         )

trustgraph-base/trustgraph/clients/document_rag_client.py

@@ -20,6 +20,7 @@ class DocumentRagClient(BaseClient):
             input_queue=None,
             output_queue=None,
             pulsar_host="pulsar://pulsar:6650",
+            pulsar_api_key=None,
     ):
 
         if input_queue == None:
@@ -34,6 +35,7 @@ class DocumentRagClient(BaseClient):
             input_queue=input_queue,
             output_queue=output_queue,
             pulsar_host=pulsar_host,
+            pulsar_api_key=pulsar_api_key,
             input_schema=DocumentRagQuery,
             output_schema=DocumentRagResponse,
         )

trustgraph-base/trustgraph/clients/embeddings_client.py

@@ -20,6 +20,7 @@ class EmbeddingsClient(BaseClient):
             output_queue=None,
             subscriber=None,
             pulsar_host="pulsar://pulsar:6650",
+            pulsar_api_key=None,
     ):
 
         if input_queue == None:
@@ -34,6 +35,7 @@ class EmbeddingsClient(BaseClient):
             input_queue=input_queue,
             output_queue=output_queue,
             pulsar_host=pulsar_host,
+            pulsar_api_key=pulsar_api_key,
             input_schema=EmbeddingsRequest,
             output_schema=EmbeddingsResponse,
         )

trustgraph-base/trustgraph/clients/graph_embeddings_client.py

@@ -20,6 +20,7 @@ class GraphEmbeddingsClient(BaseClient):
             input_queue=None,
             output_queue=None,
             pulsar_host="pulsar://pulsar:6650",
+            pulsar_api_key=None,
     ):
 
         if input_queue == None:
@@ -34,6 +35,7 @@ class GraphEmbeddingsClient(BaseClient):
             input_queue=input_queue,
             output_queue=output_queue,
             pulsar_host=pulsar_host,
+            pulsar_api_key=pulsar_api_key,
             input_schema=GraphEmbeddingsRequest,
             output_schema=GraphEmbeddingsResponse,
         )

trustgraph-base/trustgraph/clients/graph_rag_client.py

@@ -20,6 +20,7 @@ class GraphRagClient(BaseClient):
             input_queue=None,
             output_queue=None,
             pulsar_host="pulsar://pulsar:6650",
+            pulsar_api_key=None,
     ):
 
         if input_queue == None:
@@ -34,6 +35,7 @@ class GraphRagClient(BaseClient):
             input_queue=input_queue,
             output_queue=output_queue,
             pulsar_host=pulsar_host,
+            pulsar_api_key=pulsar_api_key,
             input_schema=GraphRagQuery,
             output_schema=GraphRagResponse,
         )

trustgraph-base/trustgraph/clients/llm_client.py

@@ -20,6 +20,7 @@ class LlmClient(BaseClient):
             input_queue=None,
             output_queue=None,
             pulsar_host="pulsar://pulsar:6650",
+            pulsar_api_key=None,
     ):
 
         if input_queue is None: input_queue = text_completion_request_queue
@@ -31,6 +32,7 @@ class LlmClient(BaseClient):
             input_queue=input_queue,
             output_queue=output_queue,
             pulsar_host=pulsar_host,
+            pulsar_api_key=pulsar_api_key,
             input_schema=TextCompletionRequest,
             output_schema=TextCompletionResponse,
         )

trustgraph-base/trustgraph/clients/prompt_client.py

@@ -39,6 +39,7 @@ class PromptClient(BaseClient):
             input_queue=None,
             output_queue=None,
             pulsar_host="pulsar://pulsar:6650",
+            pulsar_api_key=None,
     ):
 
         if input_queue == None:
@@ -53,6 +54,7 @@ class PromptClient(BaseClient):
             input_queue=input_queue,
             output_queue=output_queue,
             pulsar_host=pulsar_host,
+            pulsar_api_key=pulsar_api_key,
             input_schema=PromptRequest,
             output_schema=PromptResponse,
         )

trustgraph-base/trustgraph/clients/triples_query_client.py

@@ -21,6 +21,7 @@ class TriplesQueryClient(BaseClient):
             input_queue=None,
             output_queue=None,
             pulsar_host="pulsar://pulsar:6650",
+            pulsar_api_key=None,
     ):
 
         if input_queue == None:
@@ -34,6 +35,7 @@ class TriplesQueryClient(BaseClient):
             subscriber=subscriber,
             input_queue=input_queue,
             output_queue=output_queue,
+            pulsar_api_key=pulsar_api_key,
             pulsar_host=pulsar_host,
             input_schema=TriplesQueryRequest,
             output_schema=TriplesQueryResponse,

trustgraph-cli/scripts/tg-graph-show

@@ -9,12 +9,13 @@ import os
 from trustgraph.clients.triples_query_client import TriplesQueryClient
 
 default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650')
+default_pulsar_api_key = os.getenv("PULSAR_API_KEY", None)
 default_user = 'trustgraph'
 default_collection = 'default'
 
-def show_graph(pulsar, user, collection):
+def show_graph(pulsar, user, collection, pulsar_api_key=None):
 
-    tq = TriplesQueryClient(pulsar_host=pulsar)
+    tq = TriplesQueryClient(pulsar_host=pulsar, pulsar_api_key=pulsar_api_key)
 
     rows = tq.request(
         user=user, collection=collection,
@@ -48,7 +49,13 @@ def main():
         default=default_collection,
         help=f'Collection ID (default: {default_collection})'
     )
-
+        
+    parser.add_argument(
+        '--pulsar-api-key',
+        default=default_pulsar_api_key,
+        help=f'Pulsar API key',
+    )
+    
     args = parser.parse_args()
 
     try:
@@ -56,6 +63,7 @@ def main():
         show_graph(
             pulsar=args.pulsar_host, user=args.user,
             collection=args.collection,
+            pulsar_api_key=args.pulsar_api_key,
         )
 
     except Exception as e:

trustgraph-cli/scripts/tg-graph-to-turtle

@@ -13,10 +13,11 @@ import io
 import sys
 
 default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650')
+default_pulsar_api_key = os.getenv("PULSAR_API_KEY", None)
 
-def show_graph(pulsar):
+def show_graph(pulsar, pulsar_api_key=None):
 
-    tq = TriplesQueryClient(pulsar_host=pulsar)
+    tq = TriplesQueryClient(pulsar_host=pulsar, pulsar_api_key=pulsar_api_key)
 
     rows = tq.request(None, None, None, limit=10_000_000)
 
@@ -60,12 +61,18 @@ def main():
         default=default_pulsar_host,
         help=f'Pulsar host (default: {default_pulsar_host})',
     )
-
+    
+    parser.add_argument(
+        '--pulsar-api-key',
+        default=default_pulsar_api_key,
+        help=f'Pulsar API key',
+    )
+    
     args = parser.parse_args()
 
     try:
 
-        show_graph(args.pulsar_host)
+        show_graph(args.pulsar_host, pulsar_api_key=args.pulsar_api_key)
 
     except Exception as e:
 

trustgraph-cli/scripts/tg-invoke-agent

@@ -11,6 +11,7 @@ import textwrap
 from trustgraph.clients.agent_client import AgentClient
 
 default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650')
+default_pulsar_api_key = os.getenv("PULSAR_API_KEY", None)
 default_user = 'trustgraph'
 default_collection = 'default'
 
@@ -29,10 +30,10 @@ def output(text, prefix="> ", width=78):
 
 def query(
         pulsar_host, query, user, collection,
-        plan=None, state=None, verbose=False
+        plan=None, state=None, verbose=False, pulsar_api_key=None
 ):
 
-    am = AgentClient(pulsar_host=pulsar_host)
+    am = AgentClient(pulsar_host=pulsar_host, pulsar_api_key=pulsar_api_key)
 
     if verbose:
         output(wrap(query), "\U00002753 ")
@@ -100,6 +101,12 @@ def main():
         action="store_true",
         help=f'Output thinking/observations'
     )
+    
+    parser.add_argument(
+        '--pulsar-api-key',
+        default=default_pulsar_api_key,
+        help=f'Pulsar API key',
+    )
 
     args = parser.parse_args()
 
@@ -113,6 +120,7 @@ def main():
             plan=args.plan,
             state=args.state,
             verbose=args.verbose,
+            pulsar_api_key=args.pulsar_api_key,
         )
 
     except Exception as e:

trustgraph-cli/scripts/tg-invoke-llm

@@ -11,10 +11,12 @@ import json
 from trustgraph.clients.llm_client import LlmClient
 
 default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650')
+default_pulsar_api_key = os.getenv("PULSAR_API_KEY", None)
 
-def query(pulsar_host, system, prompt):
 
-    cli = LlmClient(pulsar_host=pulsar_host)
+def query(pulsar_host, system, prompt, pulsar_api_key=None):
+
+    cli = LlmClient(pulsar_host=pulsar_host, pulsar_api_key=pulsar_api_key)
 
     resp = cli.request(system=system, prompt=prompt)
 
@@ -32,7 +34,7 @@ def main():
         default=default_pulsar_host,
         help=f'Pulsar host (default: {default_pulsar_host})',
     )
-
+    
     parser.add_argument(
         'system',
         nargs=1,
@@ -44,6 +46,13 @@ def main():
         nargs=1,
         help='LLM prompt e.g. What is 2 + 2?',
     )
+    
+    parser.add_argument(
+        '--pulsar-api-key',
+        default=default_pulsar_api_key,
+        help=f'Pulsar API key',
+    )
+
 
     args = parser.parse_args()
 
@@ -53,6 +62,7 @@ def main():
             pulsar_host=args.pulsar_host,
             system=args.system[0],
             prompt=args.prompt[0],
+            pulsar_api_key=args.pulsar_api_key,
         )
 
     except Exception as e:

trustgraph-cli/scripts/tg-invoke-prompt

@@ -15,10 +15,12 @@ import json
 from trustgraph.clients.prompt_client import PromptClient
 
 default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650')
+default_pulsar_api_key = os.getenv("PULSAR_API_KEY", None)
 
-def query(pulsar_host, template_id, variables):
 
-    cli = PromptClient(pulsar_host=pulsar_host)
+def query(pulsar_host, template_id, variables, pulsar_api_key=None):
+
+    cli = PromptClient(pulsar_host=pulsar_host, pulsar_api_key=pulsar_api_key)
 
     resp = cli.request(id=template_id, variables=variables)
 
@@ -55,6 +57,13 @@ def main():
 specified multiple times''',
     )
 
+    
+    parser.add_argument(
+        '--pulsar-api-key',
+        default=default_pulsar_api_key,
+        help=f'Pulsar API key',
+    )
+
     args = parser.parse_args()
 
     variables = {}
@@ -73,6 +82,7 @@ specified multiple times''',
             pulsar_host=args.pulsar_host,
             template_id=args.id[0],
             variables=variables,
+            pulsar_api_key=args.pulsar_api_key,
         )
 
     except Exception as e:

trustgraph-cli/scripts/tg-load-pdf

@@ -34,13 +34,22 @@ class Loader:
             collection,
             log_level,
             metadata,
+            pulsar_api_key=None,
     ):
-
+        
-        self.client = pulsar.Client(
+        if pulsar_api_key:
+            auth = pulsar.AuthenticationToken(pulsar_api_key)
+            self.client = pulsar.Client(
+				pulsar_host,
+				authentication=auth,
+				logger=pulsar.ConsoleLogger(log_level.to_pulsar())
+            )
+        else:
+            self.client = pulsar.Client(
             pulsar_host,
             logger=pulsar.ConsoleLogger(log_level.to_pulsar())
-        )
+            )
-
+            
         self.producer = self.client.create_producer(
             topic=output_queue,
             schema=JsonSchema(Document),
@@ -120,6 +129,7 @@ def main():
     )
 
     default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650')
+    default_pulsar_api_key = os.getenv("PULSAR_API_KEY", None)
     default_output_queue = document_ingest_queue
 
     parser.add_argument(
@@ -127,6 +137,12 @@ def main():
         default=default_pulsar_host,
         help=f'Pulsar host (default: {default_pulsar_host})',
     )
+    
+    parser.add_argument(
+        '--pulsar-api-key',
+        default=default_pulsar_api_key,
+        help=f'Pulsar API key',
+    )
 
     parser.add_argument(
         '-o', '--output-queue',
@@ -240,6 +256,7 @@ def main():
 
             p = Loader(
                 pulsar_host=args.pulsar_host,
+                pulsar_api_key=args.pulsar_api_key,
                 output_queue=args.output_queue,
                 user=args.user,
                 collection=args.collection,

trustgraph-cli/scripts/tg-load-text

@@ -33,12 +33,20 @@ class Loader:
             collection,
             log_level,
             metadata,
+            pulsar_api_key=None,
     ):
-
+        if pulsar_api_key:
-        self.client = pulsar.Client(
+            auth = pulsar.AuthenticationToken(pulsar_api_key)
-            pulsar_host,
+            self.client = pulsar.Client(
-            logger=pulsar.ConsoleLogger(log_level.to_pulsar())
+                pulsar_host,
-        )
+                authentication=auth,
+                logger=pulsar.ConsoleLogger(log_level.to_pulsar())
+            )
+        else:
+            self.client = pulsar.Client(
+                pulsar_host,
+                logger=pulsar.ConsoleLogger(log_level.to_pulsar())
+            )
 
         self.producer = self.client.create_producer(
             topic=output_queue,
@@ -119,6 +127,8 @@ def main():
     )
 
     default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650')
+    default_pulsar_api_key = os.getenv("PULSAR_API_KEY", None)
+
     default_output_queue = text_ingest_queue
 
     parser.add_argument(
@@ -126,6 +136,12 @@ def main():
         default=default_pulsar_host,
         help=f'Pulsar host (default: {default_pulsar_host})',
     )
+    
+    parser.add_argument(
+        '--pulsar-api-key',
+        default=default_pulsar_api_key,
+        help=f'Pulsar API key',
+    )
 
     parser.add_argument(
         '-o', '--output-queue',
@@ -239,6 +255,7 @@ def main():
 
             p = Loader(
                 pulsar_host=args.pulsar_host,
+                pulsar_api_key=args.pulsar_api_key,
                 output_queue=args.output_queue,
                 user=args.user,
                 collection=args.collection,

trustgraph-cli/scripts/tg-load-turtle

@@ -19,6 +19,8 @@ from trustgraph.log_level import LogLevel
 default_user = 'trustgraph'
 default_collection = 'default'
 default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650')
+default_pulsar_api_key = os.getenv("PULSAR_API_KEY", None)
+
 default_output_queue = triples_store_queue
 
 class Loader:
@@ -31,12 +33,21 @@ class Loader:
             files,
             user,
             collection,
+            pulsar_api_key=None,
     ):
 
-        self.client = pulsar.Client(
+        if pulsar_api_key:
-            pulsar_host,
+            auth = pulsar.AuthenticationToken(pulsar_api_key)
-            logger=pulsar.ConsoleLogger(log_level.to_pulsar())
+            self.client = pulsar.Client(
-        )
+				pulsar_host,
+				authentication=auth,
+				logger=pulsar.ConsoleLogger(log_level.to_pulsar())
+            )
+        else:
+            self.client = pulsar.Client(
+				pulsar_host,
+				logger=pulsar.ConsoleLogger(log_level.to_pulsar())
+            )
 
         self.producer = self.client.create_producer(
             topic=output_queue,
@@ -98,6 +109,12 @@ def main():
         default=default_pulsar_host,
         help=f'Pulsar host (default: {default_pulsar_host})',
     )
+    
+    parser.add_argument(
+        '--pulsar-api-key',
+        default=default_pulsar_api_key,
+        help=f'Pulsar API key',
+    )
 
     parser.add_argument(
         '-o', '--output-queue',
@@ -137,6 +154,7 @@ def main():
         try:
             p = Loader(
                 pulsar_host=args.pulsar_host,
+                pulsar_api_key=args.pulsar_api_key,
                 output_queue=args.output_queue,
                 log_level=args.log_level,
                 files=args.files,

trustgraph-cli/scripts/tg-query-document-rag

@@ -9,12 +9,14 @@ import os
 from trustgraph.clients.document_rag_client import DocumentRagClient
 
 default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650')
+default_pulsar_api_key = os.getenv("PULSAR_API_KEY", None)
+
 default_user = 'trustgraph'
 default_collection = 'default'
 
-def query(pulsar_host, query, user, collection):
+def query(pulsar_host, query, user, collection, pulsar_api_key=None):
 
-    rag = DocumentRagClient(pulsar_host=pulsar)
+    rag = DocumentRagClient(pulsar_host=pulsar_host, pulsar_api_key=pulsar_api_key)
     resp = rag.request(user=user, collection=collection, query=query)
     print(resp)
 
@@ -30,7 +32,12 @@ def main():
         default=default_pulsar_host,
         help=f'Pulsar host (default: {default_pulsar_host})',
     )
-
+    
+    parser.add_argument(
+        '--pulsar-api-key',
+        default=default_pulsar_api_key,
+        help=f'Pulsar API key',
+    )
     parser.add_argument(
         '-q', '--query',
         required=True,
@@ -55,6 +62,7 @@ def main():
 
         query(
             pulsar_host=args.pulsar_host,
+            pulsar_api_key=args.pulsar_api_key,
             query=args.query,
             user=args.user,
             collection=args.collection,

trustgraph-cli/scripts/tg-query-graph-rag

@@ -9,12 +9,14 @@ import os
 from trustgraph.clients.graph_rag_client import GraphRagClient
 
 default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650')
+default_pulsar_api_key = os.getenv("PULSAR_API_KEY", None)
+
 default_user = 'trustgraph'
 default_collection = 'default'
 
-def query(pulsar_host, query, user, collection):
+def query(pulsar_host, query, user, collection, pulsar_api_key=None):
 
-    rag = GraphRagClient(pulsar_host=pulsar_host)
+    rag = GraphRagClient(pulsar_host=pulsar_host, pulsar_api_key=pulsar_api_key)
     resp = rag.request(user=user, collection=collection, query=query)
     print(resp)
 
@@ -31,6 +33,12 @@ def main():
         help=f'Pulsar host (default: {default_pulsar_host})',
     )
 
+    parser.add_argument(
+        '--pulsar-api-key',
+        default=default_pulsar_api_key,
+        help=f'Pulsar API key',
+    )
+
     parser.add_argument(
         '-q', '--query',
         required=True,
@@ -55,6 +63,7 @@ def main():
 
         query(
             pulsar_host=args.pulsar_host,
+            pulsar_api_key=args.pulsar_api_key,
             query=args.query,
             user=args.user,
             collection=args.collection,

trustgraph-flow/trustgraph/agent/react/service.py

@@ -166,21 +166,24 @@ class Processor(ConsumerProducer):
             subscriber=subscriber,
             input_queue=prompt_request_queue,
             output_queue=prompt_response_queue,
-            pulsar_host = self.pulsar_host
+            pulsar_host = self.pulsar_host,
+            pulsar_api_key=self.pulsar_api_key,
         )
 
         self.llm = LlmClient(
             subscriber=subscriber,
             input_queue=text_completion_request_queue,
             output_queue=text_completion_response_queue,
-            pulsar_host = self.pulsar_host
+            pulsar_host = self.pulsar_host,
+            pulsar_api_key=self.pulsar_api_key,
         )
 
         self.graph_rag = GraphRagClient(
             subscriber=subscriber,
             input_queue=graph_rag_request_queue,
             output_queue=graph_rag_response_queue,
-            pulsar_host = self.pulsar_host
+            pulsar_host = self.pulsar_host,
+            pulsar_api_key=self.pulsar_api_key,
         )
 
         # Need to be able to feed requests to myself

trustgraph-flow/trustgraph/document_rag.py

@@ -21,6 +21,7 @@ class DocumentRag:
     def __init__(
             self,
             pulsar_host="pulsar://pulsar:6650",
+            pulsar_api_key=None,
             pr_request_queue=None,
             pr_response_queue=None,
             emb_request_queue=None,
@@ -62,6 +63,7 @@ class DocumentRag:
             subscriber=module + "-de",
             input_queue=de_request_queue,
             output_queue=de_response_queue,
+            pulsar_api_key=pulsar_api_key,
         )            
 
         self.embeddings = EmbeddingsClient(
@@ -69,6 +71,7 @@ class DocumentRag:
             input_queue=emb_request_queue,
             output_queue=emb_response_queue,
             subscriber=module + "-emb",
+            pulsar_api_key=pulsar_api_key,
         )
 
         self.lang = PromptClient(
@@ -76,6 +79,7 @@ class DocumentRag:
             input_queue=pr_request_queue,
             output_queue=pr_response_queue,
             subscriber=module + "-de-prompt",
+            pulsar_api_key=pulsar_api_key,
         )
 
         if self.verbose:

trustgraph-flow/trustgraph/embeddings/vectorize/vectorize.py

@@ -45,6 +45,7 @@ class Processor(ConsumerProducer):
 
         self.embeddings = EmbeddingsClient(
             pulsar_host=self.pulsar_host,
+			pulsar_api_key=self.pulsar_api_key,
             input_queue=emb_request_queue,
             output_queue=emb_response_queue,
             subscriber=module + "-emb",

trustgraph-flow/trustgraph/extract/kg/definitions/extract.py

@@ -54,6 +54,7 @@ class Processor(ConsumerProducer):
 
         self.prompt = PromptClient(
             pulsar_host=self.pulsar_host,
+            pulsar_api_key=self.pulsar_api_key,
             input_queue=pr_request_queue,
             output_queue=pr_response_queue,
             subscriber = module + "-prompt",

trustgraph-flow/trustgraph/extract/kg/relationships/extract.py

@@ -76,6 +76,7 @@ class Processor(ConsumerProducer):
 
         self.prompt = PromptClient(
             pulsar_host=self.pulsar_host,
+            pulsar_api_key=self.pulsar_api_key,
             input_queue=pr_request_queue,
             output_queue=pr_response_queue,
             subscriber = module + "-prompt",

trustgraph-flow/trustgraph/extract/kg/topics/extract.py

@@ -52,6 +52,7 @@ class Processor(ConsumerProducer):
 
         self.prompt = PromptClient(
             pulsar_host=self.pulsar_host,
+            pulsar_api_key=self.pulsar_api_key,
             input_queue=pr_request_queue,
             output_queue=pr_response_queue,
             subscriber = module + "-prompt",

trustgraph-flow/trustgraph/extract/object/row/extract.py

@@ -112,6 +112,7 @@ class Processor(ConsumerProducer):
 
         self.prompt = PromptClient(
             pulsar_host=self.pulsar_host,
+            pulsar_api_key=self.pulsar_api_key,
             input_queue=pr_request_queue,
             output_queue=pr_response_queue,
             subscriber = module + "-prompt",

trustgraph-flow/trustgraph/gateway/agent.py

@@ -7,10 +7,11 @@ from . endpoint import ServiceEndpoint
 from . requestor import ServiceRequestor
 
 class AgentRequestor(ServiceRequestor):
-    def __init__(self, pulsar_host, timeout, auth):
+    def __init__(self, pulsar_host, timeout, auth, pulsar_api_key=None):
 
         super(AgentRequestor, self).__init__(
             pulsar_host=pulsar_host,
+            pulsar_api_key=pulsar_api_key,
             request_queue=agent_request_queue,
             response_queue=agent_response_queue,
             request_schema=AgentRequest,

trustgraph-flow/trustgraph/gateway/dbpedia.py

@@ -7,10 +7,11 @@ from . endpoint import ServiceEndpoint
 from . requestor import ServiceRequestor
 
 class DbpediaRequestor(ServiceRequestor):
-    def __init__(self, pulsar_host, timeout, auth):
+    def __init__(self, pulsar_host, timeout, auth, pulsar_api_key=None):
 
         super(DbpediaRequestor, self).__init__(
             pulsar_host=pulsar_host,
+            pulsar_api_key=pulsar_api_key,
             request_queue=dbpedia_lookup_request_queue,
             response_queue=dbpedia_lookup_response_queue,
             request_schema=LookupRequest,

trustgraph-flow/trustgraph/gateway/document_load.py

@@ -8,10 +8,11 @@ from . sender import ServiceSender
 from . serialize import to_subgraph
 
 class DocumentLoadSender(ServiceSender):
-    def __init__(self, pulsar_host):
+    def __init__(self, pulsar_host, pulsar_api_key=None):
 
         super(DocumentLoadSender, self).__init__(
             pulsar_host=pulsar_host,
+            pulsar_api_key=pulsar_api_key,
             request_queue=document_ingest_queue,
             request_schema=Document,
         )

trustgraph-flow/trustgraph/gateway/embeddings.py

@@ -7,10 +7,11 @@ from . endpoint import ServiceEndpoint
 from . requestor import ServiceRequestor
 
 class EmbeddingsRequestor(ServiceRequestor):
-    def __init__(self, pulsar_host, timeout, auth):
+    def __init__(self, pulsar_host, timeout, auth, pulsar_api_key=None):
 
         super(EmbeddingsRequestor, self).__init__(
             pulsar_host=pulsar_host,
+            pulsar_api_key=pulsar_api_key,
             request_queue=embeddings_request_queue,
             response_queue=embeddings_response_queue,
             request_schema=EmbeddingsRequest,

trustgraph-flow/trustgraph/gateway/encyclopedia.py

@@ -7,10 +7,11 @@ from . endpoint import ServiceEndpoint
 from . requestor import ServiceRequestor
 
 class EncyclopediaRequestor(ServiceRequestor):
-    def __init__(self, pulsar_host, timeout, auth):
+    def __init__(self, pulsar_host, timeout, auth, pulsar_api_key=None):
 
         super(EncyclopediaRequestor, self).__init__(
             pulsar_host=pulsar_host,
+            pulsar_api_key=pulsar_api_key,
             request_queue=encyclopedia_lookup_request_queue,
             response_queue=encyclopedia_lookup_response_queue,
             request_schema=LookupRequest,

trustgraph-flow/trustgraph/gateway/graph_embeddings_load.py

@@ -15,7 +15,7 @@ from . serialize import to_subgraph, to_value
 class GraphEmbeddingsLoadEndpoint(SocketEndpoint):
 
     def __init__(
-            self, pulsar_host, auth, path="/api/v1/load/graph-embeddings",
+            self, pulsar_host, auth, pulsar_api_key=None, path="/api/v1/load/graph-embeddings",
     ):
 
         super(GraphEmbeddingsLoadEndpoint, self).__init__(
@@ -23,9 +23,11 @@ class GraphEmbeddingsLoadEndpoint(SocketEndpoint):
         )
 
         self.pulsar_host=pulsar_host
+        self.pulsar_api_key=pulsar_api_key
 
         self.publisher = Publisher(
             self.pulsar_host, graph_embeddings_store_queue,
+            self.pulsar_api_key,
             schema=JsonSchema(GraphEmbeddings)
         )
 

trustgraph-flow/trustgraph/gateway/graph_embeddings_query.py

@@ -8,10 +8,11 @@ from . requestor import ServiceRequestor
 from . serialize import serialize_value
 
 class GraphEmbeddingsQueryRequestor(ServiceRequestor):
-    def __init__(self, pulsar_host, timeout, auth):
+    def __init__(self, pulsar_host, timeout, auth, pulsar_api_key=None):
 
         super(GraphEmbeddingsQueryRequestor, self).__init__(
             pulsar_host=pulsar_host,
+            pulsar_api_key=pulsar_api_key,
             request_queue=graph_embeddings_request_queue,
             response_queue=graph_embeddings_response_queue,
             request_schema=GraphEmbeddingsRequest,

trustgraph-flow/trustgraph/gateway/graph_embeddings_stream.py

@@ -14,7 +14,7 @@ from . serialize import serialize_graph_embeddings
 class GraphEmbeddingsStreamEndpoint(SocketEndpoint):
 
     def __init__(
-            self, pulsar_host, auth, path="/api/v1/stream/graph-embeddings"
+            self, pulsar_host, auth, path="/api/v1/stream/graph-embeddings", pulsar_api_key=None
     ):
 
         super(GraphEmbeddingsStreamEndpoint, self).__init__(
@@ -22,10 +22,12 @@ class GraphEmbeddingsStreamEndpoint(SocketEndpoint):
         )
 
         self.pulsar_host=pulsar_host
+        self.pulsar_api_key=pulsar_api_key
 
         self.subscriber = Subscriber(
             self.pulsar_host, graph_embeddings_store_queue,
             "api-gateway", "api-gateway",
+            pulsar_api_key=self.pulsar_api_key,
             schema=JsonSchema(GraphEmbeddings)
         )
 

trustgraph-flow/trustgraph/gateway/graph_rag.py

@@ -7,10 +7,11 @@ from . endpoint import ServiceEndpoint
 from . requestor import ServiceRequestor
 
 class GraphRagRequestor(ServiceRequestor):
-    def __init__(self, pulsar_host, timeout, auth):
+    def __init__(self, pulsar_host, timeout, auth, pulsar_api_key=None):
 
         super(GraphRagRequestor, self).__init__(
             pulsar_host=pulsar_host,
+            pulsar_api_key=pulsar_api_key,
             request_queue=graph_rag_request_queue,
             response_queue=graph_rag_response_queue,
             request_schema=GraphRagQuery,

trustgraph-flow/trustgraph/gateway/internet_search.py

@@ -7,10 +7,11 @@ from . endpoint import ServiceEndpoint
 from . requestor import ServiceRequestor
 
 class InternetSearchRequestor(ServiceRequestor):
-    def __init__(self, pulsar_host, timeout, auth):
+    def __init__(self, pulsar_host, timeout, auth, pulsar_api_key=None):
 
         super(InternetSearchRequestor, self).__init__(
             pulsar_host=pulsar_host,
+            pulsar_api_key=pulsar_api_key,
             request_queue=internet_search_request_queue,
             response_queue=internet_search_response_queue,
             request_schema=LookupRequest,

trustgraph-flow/trustgraph/gateway/mux.py

@@ -21,6 +21,7 @@ class MuxEndpoint(SocketEndpoint):
             self, pulsar_host, auth,
             services,
             path="/api/v1/socket",
+            pulsar_api_key=None
     ):
 
         super(MuxEndpoint, self).__init__(

trustgraph-flow/trustgraph/gateway/prompt.py

@@ -9,10 +9,11 @@ from . endpoint import ServiceEndpoint
 from . requestor import ServiceRequestor
 
 class PromptRequestor(ServiceRequestor):
-    def __init__(self, pulsar_host, timeout, auth):
+    def __init__(self, pulsar_host, timeout, auth, pulsar_api_key=None):
 
         super(PromptRequestor, self).__init__(
             pulsar_host=pulsar_host,
+            pulsar_api_key=pulsar_api_key,
             request_queue=prompt_request_queue,
             response_queue=prompt_response_queue,
             request_schema=PromptRequest,

trustgraph-flow/trustgraph/gateway/publisher.py

@@ -7,8 +7,9 @@ import threading
 class Publisher:
 
     def __init__(self, pulsar_host, topic, schema=None, max_size=10,
-                 chunking_enabled=False):
+                 chunking_enabled=False, pulsar_api_key=None):
         self.pulsar_host = pulsar_host
+        self.pulsar_api_key = pulsar_api_key,
         self.topic = topic
         self.schema = schema
         self.q = queue.Queue(maxsize=max_size)
@@ -23,10 +24,16 @@ class Publisher:
         while True:
 
             try:
-
+                
-                client = pulsar.Client(
+                if self.pulsar_api_key:
-                    self.pulsar_host,
+                    client = pulsar.Client(
-                )
+                        self.pulsar_host,
+                        authentication=pulsar.AuthenticationToken(self.pulsar_api_key)
+                    )
+                else:
+                    client = pulsar.Client(
+                        self.pulsar_host,
+                    )
 
                 producer = client.create_producer(
                     topic=self.topic,

trustgraph-flow/trustgraph/gateway/requestor.py

@@ -19,16 +19,19 @@ class ServiceRequestor:
             response_queue, response_schema,
             subscription="api-gateway", consumer_name="api-gateway",
             timeout=600,
+            pulsar_api_key=None,
     ):
 
         self.pub = Publisher(
             pulsar_host, request_queue,
-            schema=JsonSchema(request_schema)
+            pulsar_api_key,
+            schema=JsonSchema(request_schema),
         )
 
         self.sub = Subscriber(
             pulsar_host, response_queue,
             subscription, consumer_name,
+            pulsar_api_key,
             JsonSchema(response_schema)
         )
 

trustgraph-flow/trustgraph/gateway/sender.py

@@ -17,10 +17,12 @@ class ServiceSender:
             self,
             pulsar_host,
             request_queue, request_schema,
+            pulsar_api_key=None,
     ):
 
         self.pub = Publisher(
             pulsar_host, request_queue,
+            pulsar_api_key,
             schema=JsonSchema(request_schema)
         )
 

trustgraph-flow/trustgraph/gateway/service.py

@@ -53,6 +53,7 @@ logger = logging.getLogger("api")
 logger.setLevel(logging.INFO)
 
 default_pulsar_host = os.getenv("PULSAR_HOST", "pulsar://pulsar:6650")
+default_pulsar_api_key = os.getenv("PULSAR_API_KEY", None)
 default_timeout = 600
 default_port = 8088
 default_api_token = os.getenv("GATEWAY_SECRET", "")
@@ -69,6 +70,7 @@ class Api:
         self.port = int(config.get("port", default_port))
         self.timeout = int(config.get("timeout", default_timeout))
         self.pulsar_host = config.get("pulsar_host", default_pulsar_host)
+        self.pulsar_api_key = config.get("pulsar_api_key", default_pulsar_api_key)
 
         api_token = config.get("api_token", default_api_token)
 
@@ -81,49 +83,49 @@ class Api:
         self.services = {
             "text-completion": TextCompletionRequestor(
                 pulsar_host=self.pulsar_host, timeout=self.timeout,
-                auth = self.auth,
+                auth = self.auth, pulsar_api_key=self.pulsar_api_key,
             ),
             "prompt": PromptRequestor(
                 pulsar_host=self.pulsar_host, timeout=self.timeout,
-                auth = self.auth,
+                auth = self.auth, pulsar_api_key=self.pulsar_api_key,
             ),
             "graph-rag": GraphRagRequestor(
                 pulsar_host=self.pulsar_host, timeout=self.timeout,
-                auth = self.auth,
+                auth = self.auth, pulsar_api_key=self.pulsar_api_key,
             ),
             "triples-query": TriplesQueryRequestor(
                 pulsar_host=self.pulsar_host, timeout=self.timeout,
-                auth = self.auth,
+                auth = self.auth, pulsar_api_key=self.pulsar_api_key,
             ),
             "graph-embeddings-query": GraphEmbeddingsQueryRequestor(
                 pulsar_host=self.pulsar_host, timeout=self.timeout,
-                auth = self.auth,
+                auth = self.auth, pulsar_api_key=self.pulsar_api_key,
             ),
             "embeddings": EmbeddingsRequestor(
                 pulsar_host=self.pulsar_host, timeout=self.timeout,
-                auth = self.auth,
+                auth = self.auth, pulsar_api_key=self.pulsar_api_key,
             ),
             "agent": AgentRequestor(
                 pulsar_host=self.pulsar_host, timeout=self.timeout,
-                auth = self.auth,
+                auth = self.auth, pulsar_api_key=self.pulsar_api_key,
             ),
             "encyclopedia": EncyclopediaRequestor(
                 pulsar_host=self.pulsar_host, timeout=self.timeout,
-                auth = self.auth,
+                auth = self.auth, pulsar_api_key=self.pulsar_api_key,
             ),
             "dbpedia": DbpediaRequestor(
                 pulsar_host=self.pulsar_host, timeout=self.timeout,
-                auth = self.auth,
+                auth = self.auth, pulsar_api_key=self.pulsar_api_key,
             ),
             "internet-search": InternetSearchRequestor(
                 pulsar_host=self.pulsar_host, timeout=self.timeout,
-                auth = self.auth,
+                auth = self.auth, pulsar_api_key=self.pulsar_api_key,
             ),
             "document-load": DocumentLoadSender(
-                pulsar_host=self.pulsar_host,
+                pulsar_host=self.pulsar_host, pulsar_api_key=self.pulsar_api_key,
             ),
             "text-load": TextLoadSender(
-                pulsar_host=self.pulsar_host,
+                pulsar_host=self.pulsar_host, pulsar_api_key=self.pulsar_api_key,
             ),
         }
 
@@ -179,24 +181,29 @@ class Api:
             ),
             TriplesStreamEndpoint(
                 pulsar_host=self.pulsar_host,
+                pulsar_api_key=self.pulsar_api_key,
                 auth = self.auth,
             ),
             GraphEmbeddingsStreamEndpoint(
                 pulsar_host=self.pulsar_host,
+                pulsar_api_key=self.pulsar_api_key,
                 auth = self.auth,
             ),
             TriplesLoadEndpoint(
                 pulsar_host=self.pulsar_host,
                 auth = self.auth,
+                pulsar_api_key=self.pulsar_api_key,
             ),
             GraphEmbeddingsLoadEndpoint(
                 pulsar_host=self.pulsar_host,
+                pulsar_api_key=self.pulsar_api_key,
                 auth = self.auth,
             ),
             MuxEndpoint(
                 pulsar_host=self.pulsar_host,
                 auth = self.auth,
                 services = self.services,
+                pulsar_api_key=self.pulsar_api_key,
             ),
         ]
 
@@ -225,6 +232,12 @@ def run():
         default=default_pulsar_host,
         help=f'Pulsar host (default: {default_pulsar_host})',
     )
+    
+    parser.add_argument(
+        '--pulsar-api-key',
+        default=default_pulsar_api_key,
+        help=f'Pulsar API key',
+    )
 
     parser.add_argument(
         '--port',

trustgraph-flow/trustgraph/gateway/subscriber.py

@@ -6,9 +6,10 @@ import time
 
 class Subscriber:
 
-    def __init__(self, pulsar_host, topic, subscription, consumer_name,
+    def __init__(self, pulsar_host, topic, subscription, consumer_name, pulsar_api_key=None,
                  schema=None, max_size=100):
         self.pulsar_host = pulsar_host
+        self.pulsar_api_key = pulsar_api_key
         self.topic = topic
         self.subscription = subscription
         self.consumer_name = consumer_name
@@ -28,9 +29,16 @@ class Subscriber:
 
             try:
 
-                client = pulsar.Client(
+                if self.pulsar_api_key:
+                    auth = pulsar.AuthenticationToken(self.pulsar_api_key)
+                    client = pulsar.Client(
                     self.pulsar_host,
-                )
+                    authentication=auth,
+                    )
+                else:
+                    client = pulsar.Client(
+                    self.pulsar_host,
+                    )
 
                 consumer = client.subscribe(
                     topic=self.topic,

trustgraph-flow/trustgraph/gateway/text_completion.py

@@ -7,10 +7,11 @@ from . endpoint import ServiceEndpoint
 from . requestor import ServiceRequestor
 
 class TextCompletionRequestor(ServiceRequestor):
-    def __init__(self, pulsar_host, timeout, auth):
+    def __init__(self, pulsar_host, timeout, auth, pulsar_api_key=None):
 
         super(TextCompletionRequestor, self).__init__(
             pulsar_host=pulsar_host,
+            pulsar_api_key=pulsar_api_key,
             request_queue=text_completion_request_queue,
             response_queue=text_completion_response_queue,
             request_schema=TextCompletionRequest,

trustgraph-flow/trustgraph/gateway/text_load.py

@@ -8,10 +8,11 @@ from . sender import ServiceSender
 from . serialize import to_subgraph
 
 class TextLoadSender(ServiceSender):
-    def __init__(self, pulsar_host):
+    def __init__(self, pulsar_host, pulsar_api_key=None):
 
         super(TextLoadSender, self).__init__(
             pulsar_host=pulsar_host,
+            pulsar_api_key=pulsar_api_key,
             request_queue=text_ingest_queue,
             request_schema=TextDocument,
         )

trustgraph-flow/trustgraph/gateway/triples_load.py

@@ -14,17 +14,19 @@ from . serialize import to_subgraph
 
 class TriplesLoadEndpoint(SocketEndpoint):
 
-    def __init__(self, pulsar_host, auth, path="/api/v1/load/triples"):
+    def __init__(self, pulsar_host, auth, path="/api/v1/load/triples", pulsar_api_key=None):
 
         super(TriplesLoadEndpoint, self).__init__(
             endpoint_path=path, auth=auth,
         )
 
         self.pulsar_host=pulsar_host
+        self.pulsar_api_key=pulsar_api_key
 
         self.publisher = Publisher(
             self.pulsar_host, triples_store_queue,
-            schema=JsonSchema(Triples)
+            schema=JsonSchema(Triples),
+            pulsar_api_key=self.pulsar_api_key
         )
 
     async def start(self):

trustgraph-flow/trustgraph/gateway/triples_query.py

@@ -8,10 +8,11 @@ from . requestor import ServiceRequestor
 from . serialize import to_value, serialize_subgraph
 
 class TriplesQueryRequestor(ServiceRequestor):
-    def __init__(self, pulsar_host, timeout, auth):
+    def __init__(self, pulsar_host, timeout, auth, pulsar_api_key=None):
 
         super(TriplesQueryRequestor, self).__init__(
             pulsar_host=pulsar_host,
+            pulsar_api_key=pulsar_api_key,
             request_queue=triples_request_queue,
             response_queue=triples_response_queue,
             request_schema=TriplesQueryRequest,

trustgraph-flow/trustgraph/gateway/triples_stream.py

@@ -13,17 +13,19 @@ from . serialize import serialize_triples
 
 class TriplesStreamEndpoint(SocketEndpoint):
 
-    def __init__(self, pulsar_host, auth, path="/api/v1/stream/triples"):
+    def __init__(self, pulsar_host, auth, path="/api/v1/stream/triples", pulsar_api_key=None):
 
         super(TriplesStreamEndpoint, self).__init__(
             endpoint_path=path, auth=auth,
         )
 
         self.pulsar_host=pulsar_host
+        self.pulsar_api_key=pulsar_api_key
 
         self.subscriber = Subscriber(
             self.pulsar_host, triples_store_queue,
             "api-gateway", "api-gateway",
+            pulsar_api_key=self.pulsar_api_key,
             schema=JsonSchema(Triples)
         )
 

trustgraph-flow/trustgraph/graph_rag.py

@@ -161,6 +161,7 @@ class GraphRag:
     def __init__(
             self,
             pulsar_host="pulsar://pulsar:6650",
+            pulsar_api_key=None,
             pr_request_queue=None,
             pr_response_queue=None,
             emb_request_queue=None,
@@ -207,6 +208,7 @@ class GraphRag:
 
         self.ge_client = GraphEmbeddingsClient(
             pulsar_host=pulsar_host,
+            pulsar_api_key=-pulsar_api_key,
             subscriber=module + "-ge",
             input_queue=ge_request_queue,
             output_queue=ge_response_queue,
@@ -214,6 +216,7 @@ class GraphRag:
 
         self.triples_client = TriplesQueryClient(
             pulsar_host=pulsar_host,
+            pulsar_api_key=-pulsar_api_key,
             subscriber=module + "-tpl",
             input_queue=tpl_request_queue,
             output_queue=tpl_response_queue
@@ -221,6 +224,7 @@ class GraphRag:
 
         self.embeddings = EmbeddingsClient(
             pulsar_host=pulsar_host,
+            pulsar_api_key=-pulsar_api_key,
             input_queue=emb_request_queue,
             output_queue=emb_response_queue,
             subscriber=module + "-emb",
@@ -234,6 +238,7 @@ class GraphRag:
 
         self.prompt = PromptClient(
             pulsar_host=pulsar_host,
+            pulsar_api_key=-pulsar_api_key,
             input_queue=pr_request_queue,
             output_queue=pr_response_queue,
             subscriber=module + "-prompt",

trustgraph-flow/trustgraph/model/prompt/generic/service.py

@@ -63,7 +63,8 @@ class Processor(ConsumerProducer):
             subscriber=subscriber,
             input_queue=tc_request_queue,
             output_queue=tc_response_queue,
-            pulsar_host = self.pulsar_host
+            pulsar_host = self.pulsar_host,
+            pulsar_api_key=self.pulsar_api_key,
         )
 
     def parse_json(self, text):

trustgraph-flow/trustgraph/model/prompt/template/service.py

@@ -136,7 +136,8 @@ class Processor(ConsumerProducer):
             subscriber=subscriber,
             input_queue=tc_request_queue,
             output_queue=tc_response_queue,
-            pulsar_host = self.pulsar_host
+            pulsar_host = self.pulsar_host,
+            pulsar_api_key=self.pulsar_api_key,
         )
 
         # System prompt hack

trustgraph-flow/trustgraph/processing/processing.py

@@ -49,11 +49,12 @@ class Processing:
             pulsar_host,
             log_level,
             file,
+            pulsar_api_key=None,
     ):
         self.pulsar_host = pulsar_host
         self.log_level = log_level
         self.file = file
-
+        self.pulsar_api_key = pulsar_api_key
         self.defs = load(open(file, "r"), Loader=Loader)
 
     def run(self):
@@ -125,12 +126,19 @@ def run():
     )
 
     default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://pulsar:6650')
+    default_pulsar_api_key = os.getenv("PULSAR_API_KEY", None)
 
     parser.add_argument(
         '-p', '--pulsar-host',
         default=default_pulsar_host,
         help=f'Pulsar host (default: {default_pulsar_host})',
     )
+    
+    parser.add_argument(
+        '--pulsar-api-key',
+        default=default_pulsar_api_key,
+        help=f'Pulsar API key',
+    )
 
     parser.add_argument(
         '-l', '--log-level',

trustgraph-flow/trustgraph/retrieval/document_rag/rag.py

@@ -68,6 +68,7 @@ class Processor(ConsumerProducer):
 
         self.rag = DocumentRag(
             pulsar_host=self.pulsar_host,
+            pulsar_api_key=self.pulsar_api_key,
             pr_request_queue=pr_request_queue,
             pr_response_queue=pr_response_queue,
             emb_request_queue=emb_request_queue,

trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py

@@ -82,6 +82,7 @@ class Processor(ConsumerProducer):
 
         self.rag = GraphRag(
             pulsar_host=self.pulsar_host,
+            pulsar_api_key=self.pulsar_api_key,
             pr_request_queue=pr_request_queue,
             pr_response_queue=pr_response_queue,
             emb_request_queue=emb_request_queue,