diff --git a/trustgraph-base/trustgraph/base/base_processor.py b/trustgraph-base/trustgraph/base/base_processor.py index f258ff1a..8963e990 100644 --- a/trustgraph-base/trustgraph/base/base_processor.py +++ b/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 - - 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()) - ) + ) 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', diff --git a/trustgraph-base/trustgraph/clients/agent_client.py b/trustgraph-base/trustgraph/clients/agent_client.py index 2ef69274..b31b4e36 100644 --- a/trustgraph-base/trustgraph/clients/agent_client.py +++ b/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( diff --git a/trustgraph-base/trustgraph/clients/base.py b/trustgraph-base/trustgraph/clients/base.py index 78116f41..42de45a5 100644 --- a/trustgraph-base/trustgraph/clients/base.py +++ b/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, diff --git a/trustgraph-base/trustgraph/clients/document_embeddings_client.py b/trustgraph-base/trustgraph/clients/document_embeddings_client.py index d432991d..c27daab0 100644 --- a/trustgraph-base/trustgraph/clients/document_embeddings_client.py +++ b/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, ) diff --git a/trustgraph-base/trustgraph/clients/document_rag_client.py b/trustgraph-base/trustgraph/clients/document_rag_client.py index 103cbb69..6cbafa9b 100644 --- a/trustgraph-base/trustgraph/clients/document_rag_client.py +++ b/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, ) diff --git a/trustgraph-base/trustgraph/clients/embeddings_client.py b/trustgraph-base/trustgraph/clients/embeddings_client.py index 8d21bdec..811f6ed2 100644 --- a/trustgraph-base/trustgraph/clients/embeddings_client.py +++ b/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, ) diff --git a/trustgraph-base/trustgraph/clients/graph_embeddings_client.py b/trustgraph-base/trustgraph/clients/graph_embeddings_client.py index 401266bc..1a7a9512 100644 --- a/trustgraph-base/trustgraph/clients/graph_embeddings_client.py +++ b/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, ) diff --git a/trustgraph-base/trustgraph/clients/graph_rag_client.py b/trustgraph-base/trustgraph/clients/graph_rag_client.py index 9f8eff62..77102e36 100644 --- a/trustgraph-base/trustgraph/clients/graph_rag_client.py +++ b/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, ) diff --git a/trustgraph-base/trustgraph/clients/llm_client.py b/trustgraph-base/trustgraph/clients/llm_client.py index cfb0e606..a8894c8f 100644 --- a/trustgraph-base/trustgraph/clients/llm_client.py +++ b/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, ) diff --git a/trustgraph-base/trustgraph/clients/prompt_client.py b/trustgraph-base/trustgraph/clients/prompt_client.py index 4b026cf0..91707670 100644 --- a/trustgraph-base/trustgraph/clients/prompt_client.py +++ b/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, ) diff --git a/trustgraph-base/trustgraph/clients/triples_query_client.py b/trustgraph-base/trustgraph/clients/triples_query_client.py index fc1e4b26..8ed2ebb7 100644 --- a/trustgraph-base/trustgraph/clients/triples_query_client.py +++ b/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, diff --git a/trustgraph-cli/scripts/tg-graph-show b/trustgraph-cli/scripts/tg-graph-show index c09266fb..488b7dca 100755 --- a/trustgraph-cli/scripts/tg-graph-show +++ b/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: diff --git a/trustgraph-cli/scripts/tg-graph-to-turtle b/trustgraph-cli/scripts/tg-graph-to-turtle index 1d75478e..65c396f2 100755 --- a/trustgraph-cli/scripts/tg-graph-to-turtle +++ b/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: diff --git a/trustgraph-cli/scripts/tg-invoke-agent b/trustgraph-cli/scripts/tg-invoke-agent index 3f05071c..b23831cb 100755 --- a/trustgraph-cli/scripts/tg-invoke-agent +++ b/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: diff --git a/trustgraph-cli/scripts/tg-invoke-llm b/trustgraph-cli/scripts/tg-invoke-llm index d7289b5f..8e9a4bf2 100755 --- a/trustgraph-cli/scripts/tg-invoke-llm +++ b/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: diff --git a/trustgraph-cli/scripts/tg-invoke-prompt b/trustgraph-cli/scripts/tg-invoke-prompt index 19f30912..b165dde6 100755 --- a/trustgraph-cli/scripts/tg-invoke-prompt +++ b/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: diff --git a/trustgraph-cli/scripts/tg-load-pdf b/trustgraph-cli/scripts/tg-load-pdf index a0d2b3bc..0cdc909f 100755 --- a/trustgraph-cli/scripts/tg-load-pdf +++ b/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, diff --git a/trustgraph-cli/scripts/tg-load-text b/trustgraph-cli/scripts/tg-load-text index 51664a1b..c9dda8a8 100755 --- a/trustgraph-cli/scripts/tg-load-text +++ b/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, diff --git a/trustgraph-cli/scripts/tg-load-turtle b/trustgraph-cli/scripts/tg-load-turtle index 7c258fcc..3417a87d 100755 --- a/trustgraph-cli/scripts/tg-load-turtle +++ b/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, diff --git a/trustgraph-cli/scripts/tg-query-document-rag b/trustgraph-cli/scripts/tg-query-document-rag index 8d800629..ac4c98ca 100755 --- a/trustgraph-cli/scripts/tg-query-document-rag +++ b/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, diff --git a/trustgraph-cli/scripts/tg-query-graph-rag b/trustgraph-cli/scripts/tg-query-graph-rag index 8a865eea..88ab8142 100755 --- a/trustgraph-cli/scripts/tg-query-graph-rag +++ b/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, diff --git a/trustgraph-flow/trustgraph/agent/react/service.py b/trustgraph-flow/trustgraph/agent/react/service.py index 8799816b..a5c24563 100755 --- a/trustgraph-flow/trustgraph/agent/react/service.py +++ b/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 diff --git a/trustgraph-flow/trustgraph/document_rag.py b/trustgraph-flow/trustgraph/document_rag.py index f3c8b158..5c340150 100644 --- a/trustgraph-flow/trustgraph/document_rag.py +++ b/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: diff --git a/trustgraph-flow/trustgraph/embeddings/vectorize/vectorize.py b/trustgraph-flow/trustgraph/embeddings/vectorize/vectorize.py index 4cf2af05..c24b8251 100755 --- a/trustgraph-flow/trustgraph/embeddings/vectorize/vectorize.py +++ b/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", diff --git a/trustgraph-flow/trustgraph/extract/kg/definitions/extract.py b/trustgraph-flow/trustgraph/extract/kg/definitions/extract.py index eed34574..7b82e0d7 100755 --- a/trustgraph-flow/trustgraph/extract/kg/definitions/extract.py +++ b/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", diff --git a/trustgraph-flow/trustgraph/extract/kg/relationships/extract.py b/trustgraph-flow/trustgraph/extract/kg/relationships/extract.py index d2dea062..34f36ef6 100755 --- a/trustgraph-flow/trustgraph/extract/kg/relationships/extract.py +++ b/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", diff --git a/trustgraph-flow/trustgraph/extract/kg/topics/extract.py b/trustgraph-flow/trustgraph/extract/kg/topics/extract.py index 8dfc3e6e..d1c0c3bc 100755 --- a/trustgraph-flow/trustgraph/extract/kg/topics/extract.py +++ b/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", diff --git a/trustgraph-flow/trustgraph/extract/object/row/extract.py b/trustgraph-flow/trustgraph/extract/object/row/extract.py index 185a59c3..0e2a5701 100755 --- a/trustgraph-flow/trustgraph/extract/object/row/extract.py +++ b/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", diff --git a/trustgraph-flow/trustgraph/gateway/agent.py b/trustgraph-flow/trustgraph/gateway/agent.py index c7af947b..932613d4 100644 --- a/trustgraph-flow/trustgraph/gateway/agent.py +++ b/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, diff --git a/trustgraph-flow/trustgraph/gateway/dbpedia.py b/trustgraph-flow/trustgraph/gateway/dbpedia.py index 8ae4f695..54e4c97c 100644 --- a/trustgraph-flow/trustgraph/gateway/dbpedia.py +++ b/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, diff --git a/trustgraph-flow/trustgraph/gateway/document_load.py b/trustgraph-flow/trustgraph/gateway/document_load.py index 0fd9a0df..55ee2038 100644 --- a/trustgraph-flow/trustgraph/gateway/document_load.py +++ b/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, ) diff --git a/trustgraph-flow/trustgraph/gateway/embeddings.py b/trustgraph-flow/trustgraph/gateway/embeddings.py index 1efafa76..1d2815fc 100644 --- a/trustgraph-flow/trustgraph/gateway/embeddings.py +++ b/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, diff --git a/trustgraph-flow/trustgraph/gateway/encyclopedia.py b/trustgraph-flow/trustgraph/gateway/encyclopedia.py index 3f4dad79..c5bdebe1 100644 --- a/trustgraph-flow/trustgraph/gateway/encyclopedia.py +++ b/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, diff --git a/trustgraph-flow/trustgraph/gateway/graph_embeddings_load.py b/trustgraph-flow/trustgraph/gateway/graph_embeddings_load.py index 18a2e6fe..236a4ea6 100644 --- a/trustgraph-flow/trustgraph/gateway/graph_embeddings_load.py +++ b/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) ) diff --git a/trustgraph-flow/trustgraph/gateway/graph_embeddings_query.py b/trustgraph-flow/trustgraph/gateway/graph_embeddings_query.py index 5e3c0ce9..2e21e5c4 100644 --- a/trustgraph-flow/trustgraph/gateway/graph_embeddings_query.py +++ b/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, diff --git a/trustgraph-flow/trustgraph/gateway/graph_embeddings_stream.py b/trustgraph-flow/trustgraph/gateway/graph_embeddings_stream.py index f0b4dd86..ec7a3f9d 100644 --- a/trustgraph-flow/trustgraph/gateway/graph_embeddings_stream.py +++ b/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) ) diff --git a/trustgraph-flow/trustgraph/gateway/graph_rag.py b/trustgraph-flow/trustgraph/gateway/graph_rag.py index 55fd5d2f..8ae121af 100644 --- a/trustgraph-flow/trustgraph/gateway/graph_rag.py +++ b/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, diff --git a/trustgraph-flow/trustgraph/gateway/internet_search.py b/trustgraph-flow/trustgraph/gateway/internet_search.py index 127cd5d1..8f3a5581 100644 --- a/trustgraph-flow/trustgraph/gateway/internet_search.py +++ b/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, diff --git a/trustgraph-flow/trustgraph/gateway/mux.py b/trustgraph-flow/trustgraph/gateway/mux.py index ae699ae6..27138a39 100644 --- a/trustgraph-flow/trustgraph/gateway/mux.py +++ b/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__( diff --git a/trustgraph-flow/trustgraph/gateway/prompt.py b/trustgraph-flow/trustgraph/gateway/prompt.py index 080d5618..40f0bff3 100644 --- a/trustgraph-flow/trustgraph/gateway/prompt.py +++ b/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, diff --git a/trustgraph-flow/trustgraph/gateway/publisher.py b/trustgraph-flow/trustgraph/gateway/publisher.py index 89c612ce..ab752cba 100644 --- a/trustgraph-flow/trustgraph/gateway/publisher.py +++ b/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( - self.pulsar_host, - ) + + if self.pulsar_api_key: + 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, diff --git a/trustgraph-flow/trustgraph/gateway/requestor.py b/trustgraph-flow/trustgraph/gateway/requestor.py index 5f6e2692..8e2911d1 100644 --- a/trustgraph-flow/trustgraph/gateway/requestor.py +++ b/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) ) diff --git a/trustgraph-flow/trustgraph/gateway/sender.py b/trustgraph-flow/trustgraph/gateway/sender.py index 93f1164c..91e3469c 100644 --- a/trustgraph-flow/trustgraph/gateway/sender.py +++ b/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) ) diff --git a/trustgraph-flow/trustgraph/gateway/service.py b/trustgraph-flow/trustgraph/gateway/service.py index a260b631..9aea1acf 100755 --- a/trustgraph-flow/trustgraph/gateway/service.py +++ b/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', diff --git a/trustgraph-flow/trustgraph/gateway/subscriber.py b/trustgraph-flow/trustgraph/gateway/subscriber.py index cccfc5b4..09622b3e 100644 --- a/trustgraph-flow/trustgraph/gateway/subscriber.py +++ b/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, diff --git a/trustgraph-flow/trustgraph/gateway/text_completion.py b/trustgraph-flow/trustgraph/gateway/text_completion.py index 7291fc88..b9550ce6 100644 --- a/trustgraph-flow/trustgraph/gateway/text_completion.py +++ b/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, diff --git a/trustgraph-flow/trustgraph/gateway/text_load.py b/trustgraph-flow/trustgraph/gateway/text_load.py index ade6b1c1..562ca779 100644 --- a/trustgraph-flow/trustgraph/gateway/text_load.py +++ b/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, ) diff --git a/trustgraph-flow/trustgraph/gateway/triples_load.py b/trustgraph-flow/trustgraph/gateway/triples_load.py index 2689f3ad..05d329d0 100644 --- a/trustgraph-flow/trustgraph/gateway/triples_load.py +++ b/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): diff --git a/trustgraph-flow/trustgraph/gateway/triples_query.py b/trustgraph-flow/trustgraph/gateway/triples_query.py index 0ea7cd8d..173b75e6 100644 --- a/trustgraph-flow/trustgraph/gateway/triples_query.py +++ b/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, diff --git a/trustgraph-flow/trustgraph/gateway/triples_stream.py b/trustgraph-flow/trustgraph/gateway/triples_stream.py index 92ada132..2507a34d 100644 --- a/trustgraph-flow/trustgraph/gateway/triples_stream.py +++ b/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) ) diff --git a/trustgraph-flow/trustgraph/graph_rag.py b/trustgraph-flow/trustgraph/graph_rag.py index f69ebeb7..a40ba32c 100644 --- a/trustgraph-flow/trustgraph/graph_rag.py +++ b/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", diff --git a/trustgraph-flow/trustgraph/model/prompt/generic/service.py b/trustgraph-flow/trustgraph/model/prompt/generic/service.py index 96c9be57..89cc850c 100755 --- a/trustgraph-flow/trustgraph/model/prompt/generic/service.py +++ b/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): diff --git a/trustgraph-flow/trustgraph/model/prompt/template/service.py b/trustgraph-flow/trustgraph/model/prompt/template/service.py index 2e5416f4..b8aa489c 100755 --- a/trustgraph-flow/trustgraph/model/prompt/template/service.py +++ b/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 diff --git a/trustgraph-flow/trustgraph/processing/processing.py b/trustgraph-flow/trustgraph/processing/processing.py index 5e4c7c8a..e5e9593f 100644 --- a/trustgraph-flow/trustgraph/processing/processing.py +++ b/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', diff --git a/trustgraph-flow/trustgraph/retrieval/document_rag/rag.py b/trustgraph-flow/trustgraph/retrieval/document_rag/rag.py index 4310cdbd..0d2a3b86 100755 --- a/trustgraph-flow/trustgraph/retrieval/document_rag/rag.py +++ b/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, diff --git a/trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py b/trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py index 1219050e..5d5793d6 100755 --- a/trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py +++ b/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,