trustgraph/trustgraph-cli/scripts/tg-invoke-llm
cybermaggedon 617eb7efd5
Feature/pulsar api key support (#308)
* Add pulsar API token check

* Added missing api_key references

---------

Co-authored-by: Tyler O <4535788+toliver38@users.noreply.github.com>
2025-02-15 11:22:48 +00:00

69 lines
1.3 KiB
Python
Executable file

#!/usr/bin/env python3
"""
Invokes the text completion service by specifying an LLM system prompt
and user prompt. Both arguments are required.
"""
import argparse
import os
import json
from trustgraph.api import Api
default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
def query(url, system, prompt):
api = Api(url)
resp = api.text_completion(system=system, prompt=prompt)
print(resp)
def main():
parser = argparse.ArgumentParser(
prog='tg-invoke-llm',
description=__doc__,
)
parser.add_argument(
'-u', '--url',
default=default_url,
help=f'API URL (default: {default_url})',
)
parser.add_argument(
'system',
nargs=1,
help='LLM system prompt e.g. You are a helpful assistant',
)
parser.add_argument(
'prompt',
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()
try:
query(
url=args.url,
system=args.system[0],
prompt=args.prompt[0],
)
except Exception as e:
print("Exception:", e, flush=True)
main()