diff --git a/trustgraph-cli/scripts/tg-invoke-prompt b/trustgraph-cli/scripts/tg-invoke-prompt new file mode 100755 index 00000000..8f62965c --- /dev/null +++ b/trustgraph-cli/scripts/tg-invoke-prompt @@ -0,0 +1,75 @@ +#!/usr/bin/env python3 + +""" +Uses the GraphRAG service to answer a query +""" + +import argparse +import os +import json +from trustgraph.clients.prompt_client import PromptClient + +default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650') + +def query(pulsar_host, id, terms): + + cli = PromptClient(pulsar_host=pulsar_host) + + resp = cli.request(id=id, terms=terms) + + if isinstance(resp, str): + print(resp) + else: + print(json.dumps(resp, indent=4)) + +def main(): + + parser = argparse.ArgumentParser( + prog='tg-graph-query-rag', + description=__doc__, + ) + + parser.add_argument( + '-p', '--pulsar-host', + default=default_pulsar_host, + help=f'Pulsar host (default: {default_pulsar_host})', + ) + + parser.add_argument( + 'id', + nargs=1, + help=f'Prompt identifier', + ) + + parser.add_argument( + 'term', + nargs='*', + help=f'Prompt terms', + ) + + args = parser.parse_args() + + terms = {} + + for term in args.term: + + toks = term.split("=", 1) + if len(toks) != 2: + raise RuntimeError(f"Malformed term: {term}") + + terms[toks[0]] = toks[1] + + try: + + query( + pulsar_host=args.pulsar_host, + id=args.id[0], + terms=terms, + ) + + except Exception as e: + + print("Exception:", e, flush=True) + +main() + diff --git a/trustgraph-cli/setup.py b/trustgraph-cli/setup.py index 13101855..392819e2 100644 --- a/trustgraph-cli/setup.py +++ b/trustgraph-cli/setup.py @@ -50,5 +50,6 @@ "scripts/tg-query-graph-rag", "scripts/tg-init-pulsar", "scripts/tg-processor-state", + "scripts/tg-invoke-prompt", ] )