trustgraph/trustgraph-cli/trustgraph/cli/invoke_nlp_query.py
Cyber MacGeddon db05427d0e IAM tech spec: Auth and access management current state and proposed
changes.

Support for separate workspaces

Addition of workspace CLI support for test purposes
2026-04-21 10:50:52 +01:00

129 lines
No EOL
3.5 KiB
Python

"""
Uses the NLP Query service to convert natural language questions to GraphQL queries
"""
import argparse
import os
import json
import sys
from trustgraph.api import Api
default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
default_workspace = os.getenv("TRUSTGRAPH_WORKSPACE", "default")
def nlp_query(url, flow_id, question, max_results, output_format='json', token=None, workspace="default"):
api = Api(url, token=token, workspace=workspace).flow().id(flow_id)
resp = api.nlp_query(
question=question,
max_results=max_results
)
# Check for errors
if "error" in resp and resp["error"]:
print("Error:", resp["error"].get("message", "Unknown error"), file=sys.stderr)
sys.exit(1)
# Format output based on requested format
if output_format == 'json':
print(json.dumps(resp, indent=2))
elif output_format == 'graphql':
# Just print the GraphQL query
if "graphql_query" in resp:
print(resp["graphql_query"])
else:
print("No GraphQL query generated", file=sys.stderr)
sys.exit(1)
elif output_format == 'summary':
# Print a human-readable summary
if "graphql_query" in resp:
print(f"Generated GraphQL Query:")
print("-" * 40)
print(resp["graphql_query"])
print("-" * 40)
if "detected_schemas" in resp and resp["detected_schemas"]:
print(f"Detected Schemas: {', '.join(resp['detected_schemas'])}")
if "confidence" in resp:
print(f"Confidence: {resp['confidence']:.2%}")
if "variables" in resp and resp["variables"]:
print(f"Variables: {json.dumps(resp['variables'], indent=2)}")
else:
print("No GraphQL query generated", file=sys.stderr)
sys.exit(1)
def main():
parser = argparse.ArgumentParser(
prog='tg-invoke-nlp-query',
description=__doc__,
)
parser.add_argument(
'-u', '--url',
default=default_url,
help=f'API URL (default: {default_url})',
)
parser.add_argument(
'-t', '--token',
default=default_token,
help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
)
parser.add_argument(
'-w', '--workspace',
default=default_workspace,
help=f'Workspace (default: {default_workspace})',
)
parser.add_argument(
'-f', '--flow-id',
default="default",
help=f'Flow ID (default: default)'
)
parser.add_argument(
'-q', '--question',
required=True,
help='Natural language question to convert to GraphQL',
)
parser.add_argument(
'-m', '--max-results',
type=int,
default=100,
help='Maximum number of results (default: 100)'
)
parser.add_argument(
'--format',
choices=['json', 'graphql', 'summary'],
default='summary',
help='Output format (default: summary)'
)
args = parser.parse_args()
try:
nlp_query(
url=args.url,
flow_id=args.flow_id,
question=args.question,
max_results=args.max_results,
output_format=args.format,
token = args.token,
workspace = args.workspace,
)
except Exception as e:
print("Exception:", e, flush=True, file=sys.stderr)
sys.exit(1)
if __name__ == "__main__":
main()