Add pulsar API token check

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,
     ):
-
-        self.client = pulsar.Client(
-            pulsar_host,
-            logger=pulsar.ConsoleLogger(log_level.to_pulsar())
-        )
+        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,
@@ -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(
-            pulsar_host,
-            logger=pulsar.ConsoleLogger(log_level.to_pulsar())
-        )
+        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,
@@ -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,