trustgraph/trustgraph-cli/scripts/tg-invoke-llm
2025-05-24 12:27:56 +01:00

70 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, 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)
main()