mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-04-29 10:26:21 +02:00
Replace per-request websocket connections in SocketClient and AsyncSocketClient with a single persistent connection that multiplexes requests by ID via a background reader task. This eliminates repeated TCP+WS handshakes which caused significant latency over proxies. Convert show_flows, show_flow_blueprints, and show_parameter_types CLI tools from sequential HTTP requests to concurrent websocket requests using AsyncSocketClient, reducing round trips from O(N) sequential to a small number of parallel batches. Also fix describe_interfaces bug in show_flows where response queue was reading the request field instead of the response field.
173 lines
4.8 KiB
Python
173 lines
4.8 KiB
Python
"""
|
|
Shows all defined flow blueprints.
|
|
"""
|
|
|
|
import argparse
|
|
import asyncio
|
|
import os
|
|
import tabulate
|
|
from trustgraph.api import AsyncSocketClient
|
|
import json
|
|
|
|
default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
|
|
default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
|
|
|
|
def format_parameters(params_metadata, param_type_defs):
|
|
"""
|
|
Format parameter metadata for display.
|
|
|
|
param_type_defs is a dict of type_name -> parsed type definition,
|
|
pre-fetched concurrently.
|
|
"""
|
|
if not params_metadata:
|
|
return "None"
|
|
|
|
param_list = []
|
|
|
|
sorted_params = sorted(
|
|
params_metadata.items(),
|
|
key=lambda x: x[1].get("order", 999)
|
|
)
|
|
|
|
for param_name, param_meta in sorted_params:
|
|
description = param_meta.get("description", param_name)
|
|
param_type = param_meta.get("type", "unknown")
|
|
|
|
type_info = param_type
|
|
if param_type in param_type_defs:
|
|
param_type_def = param_type_defs[param_type]
|
|
default = param_type_def.get("default")
|
|
if default is not None:
|
|
type_info = f"{param_type} (default: {default})"
|
|
|
|
param_list.append(f" {param_name}: {description} [{type_info}]")
|
|
|
|
return "\n".join(param_list)
|
|
|
|
async def fetch_data(client):
|
|
"""Fetch all data needed for show_flow_blueprints concurrently."""
|
|
|
|
# Round 1: list blueprints
|
|
resp = await client._send_request("flow", None, {
|
|
"operation": "list-blueprints",
|
|
})
|
|
blueprint_names = resp.get("blueprint-names", [])
|
|
|
|
if not blueprint_names:
|
|
return [], {}, {}
|
|
|
|
# Round 2: get all blueprints in parallel
|
|
blueprint_tasks = [
|
|
client._send_request("flow", None, {
|
|
"operation": "get-blueprint",
|
|
"blueprint-name": name,
|
|
})
|
|
for name in blueprint_names
|
|
]
|
|
blueprint_results = await asyncio.gather(*blueprint_tasks)
|
|
|
|
blueprints = {}
|
|
for name, resp in zip(blueprint_names, blueprint_results):
|
|
bp_data = resp.get("blueprint-definition", "{}")
|
|
blueprints[name] = json.loads(bp_data) if isinstance(bp_data, str) else bp_data
|
|
|
|
# Round 3: get all parameter type definitions in parallel
|
|
param_types_needed = set()
|
|
for bp in blueprints.values():
|
|
for param_meta in bp.get("parameters", {}).values():
|
|
pt = param_meta.get("type", "")
|
|
if pt:
|
|
param_types_needed.add(pt)
|
|
|
|
param_type_defs = {}
|
|
if param_types_needed:
|
|
param_type_tasks = [
|
|
client._send_request("config", None, {
|
|
"operation": "get",
|
|
"keys": [{"type": "parameter-type", "key": pt}],
|
|
})
|
|
for pt in param_types_needed
|
|
]
|
|
param_type_results = await asyncio.gather(*param_type_tasks)
|
|
|
|
for pt, resp in zip(param_types_needed, param_type_results):
|
|
values = resp.get("values", [])
|
|
if values:
|
|
try:
|
|
param_type_defs[pt] = json.loads(values[0].get("value", "{}"))
|
|
except (json.JSONDecodeError, AttributeError):
|
|
pass
|
|
|
|
return blueprint_names, blueprints, param_type_defs
|
|
|
|
async def _show_flow_blueprints_async(url, token=None):
|
|
async with AsyncSocketClient(url, timeout=60, token=token) as client:
|
|
return await fetch_data(client)
|
|
|
|
def show_flow_blueprints(url, token=None):
|
|
|
|
blueprint_names, blueprints, param_type_defs = asyncio.run(
|
|
_show_flow_blueprints_async(url, token=token)
|
|
)
|
|
|
|
if not blueprint_names:
|
|
print("No flow blueprints.")
|
|
return
|
|
|
|
for blueprint_name in blueprint_names:
|
|
cls = blueprints[blueprint_name]
|
|
|
|
table = []
|
|
table.append(("name", blueprint_name))
|
|
table.append(("description", cls.get("description", "")))
|
|
|
|
tags = cls.get("tags", [])
|
|
if tags:
|
|
table.append(("tags", ", ".join(tags)))
|
|
|
|
parameters = cls.get("parameters", {})
|
|
if parameters:
|
|
param_str = format_parameters(parameters, param_type_defs)
|
|
table.append(("parameters", param_str))
|
|
|
|
print(tabulate.tabulate(
|
|
table,
|
|
tablefmt="pretty",
|
|
stralign="left",
|
|
))
|
|
print()
|
|
|
|
def main():
|
|
|
|
parser = argparse.ArgumentParser(
|
|
prog='tg-show-flow-blueprints',
|
|
description=__doc__,
|
|
)
|
|
|
|
parser.add_argument(
|
|
'-u', '--api-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)',
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
try:
|
|
|
|
show_flow_blueprints(
|
|
url=args.api_url,
|
|
token=args.token,
|
|
)
|
|
|
|
except Exception as e:
|
|
|
|
print("Exception:", e, flush=True)
|
|
|
|
if __name__ == "__main__":
|
|
main()
|