trustgraph/trustgraph-cli/trustgraph/cli/invoke_llm.py
2025-07-22 22:00:06 +01:00

68 lines
No EOL
1.3 KiB
Python

"""
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, flow_id, system, prompt):
api = Api(url).flow().id(flow_id)
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(
'-f', '--flow-id',
default="default",
help=f'Flow ID (default: default)'
)
args = parser.parse_args()
try:
query(
url=args.url,
flow_id = args.flow_id,
system=args.system[0],
prompt=args.prompt[0],
)
except Exception as e:
print("Exception:", e, flush=True)
if __name__ == "__main__":
main()