Add pulsar API token check

This commit is contained in:
Tyler O 2025-02-10 16:30:53 +00:00 committed by Cyber MacGeddon
parent f7df2df266
commit 0464423b04
42 changed files with 199 additions and 24 deletions

View file

@ -48,7 +48,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:
@ -57,6 +63,7 @@ def main():
url=args.api_url,
user=args.user,
collection=args.collection,
pulsar_api_key=args.pulsar_api_key,
)
except Exception as e:

View file

@ -80,7 +80,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:

View file

@ -137,6 +137,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()

View file

@ -34,8 +34,14 @@ def main():
help=f'API URL (default: {default_url})',
)
# parser.add_argument(
# '--pulsar-api-key',
# default=default_pulsar_api_key,
# help=f'Pulsar API key',
# )
parser.add_argument(
'-q', '--question',
'-q', '--query',
required=True,
help=f'Question to answer',
)

View file

@ -33,6 +33,12 @@ def main():
default=default_url,
help=f'API URL (default: {default_url})',
)
# parser.add_argument(
# '--pulsar-api-key',
# default=default_pulsar_api_key,
# help=f'Pulsar API key',
# )
parser.add_argument(
'-q', '--question',

View file

@ -32,7 +32,7 @@ def main():
default=default_url,
help=f'API URL (default: {default_url})',
)
parser.add_argument(
'system',
nargs=1,
@ -44,6 +44,12 @@ 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()

View file

@ -54,6 +54,12 @@ def main():
help='''Prompt template terms of the form variable=value, can be
specified multiple times''',
)
# parser.add_argument(
# '--pulsar-api-key',
# default=default_pulsar_api_key,
# help=f'Pulsar API key',
# )
args = parser.parse_args()

View file

@ -30,6 +30,7 @@ class Loader:
user,
collection,
metadata,
pulsar_api_key=None,
):
self.api = Api(url)
@ -80,6 +81,12 @@ def main():
default=default_url,
help=f'API URL (default: {default_url})',
)
# parser.add_argument(
# '--pulsar-api-key',
# default=default_pulsar_api_key,
# help=f'Pulsar API key',
# )
parser.add_argument(
'-U', '--user',

View file

@ -30,6 +30,7 @@ class Loader:
user,
collection,
metadata,
pulsar_api_key=None,
):
self.api = Api(url)
@ -80,6 +81,12 @@ def main():
default=default_url,
help=f'API URL (default: {default_url})',
)
# parser.add_argument(
# '--pulsar-api-key',
# default=default_pulsar_api_key,
# help=f'Pulsar API key',
# )
parser.add_argument(
'-U', '--user',

View file

@ -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,