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

117 lines
2.8 KiB
Python
Executable file

#!/usr/bin/env python3
"""
Uses the GraphRAG service to answer a question
"""
import argparse
import os
from trustgraph.api import Api
default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
default_user = 'trustgraph'
default_collection = 'default'
default_entity_limit = 50
default_triple_limit = 30
default_max_subgraph_size = 150
default_max_path_length = 2
def question(
url, flow_id, question, user, collection, entity_limit, triple_limit,
max_subgraph_size, max_path_length
):
api = Api(url).flow().id(flow_id)
resp = api.graph_rag(
question=question, user=user, collection=collection,
entity_limit=entity_limit, triple_limit=triple_limit,
max_subgraph_size=max_subgraph_size,
max_path_length=max_path_length
)
print(resp)
def main():
parser = argparse.ArgumentParser(
prog='tg-invoke-graph-rag',
description=__doc__,
)
parser.add_argument(
'-u', '--url',
default=default_url,
help=f'API URL (default: {default_url})',
)
parser.add_argument(
'-f', '--flow-id',
default="default",
help=f'Flow ID (default: default)'
)
parser.add_argument(
'-q', '--question',
required=True,
help=f'Question to answer',
)
parser.add_argument(
'-U', '--user',
default=default_user,
help=f'User ID (default: {default_user})'
)
parser.add_argument(
'-C', '--collection',
default=default_collection,
help=f'Collection ID (default: {default_collection})'
)
parser.add_argument(
'-e', '--entity-limit',
default=default_entity_limit,
help=f'Entity limit (default: {default_entity_limit})'
)
parser.add_argument(
'-t', '--triple-limit',
default=default_triple_limit,
help=f'Triple limit (default: {default_triple_limit})'
)
parser.add_argument(
'-s', '--max-subgraph-size',
default=default_max_subgraph_size,
help=f'Max subgraph size (default: {default_max_subgraph_size})'
)
parser.add_argument(
'-p', '--max-path-length',
default=default_max_path_length,
help=f'Max path length (default: {default_max_path_length})'
)
args = parser.parse_args()
try:
question(
url=args.url,
flow_id = args.flow_id,
question=args.question,
user=args.user,
collection=args.collection,
entity_limit=args.entity_limit,
triple_limit=args.triple_limit,
max_subgraph_size=args.max_subgraph_size,
max_path_length=args.max_path_length,
)
except Exception as e:
print("Exception:", e, flush=True)
main()