mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-06-25 06:38:06 +02:00
Persistent websocket connections for socket clients and CLI tools (#723)
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.
This commit is contained in:
parent
1ec081f42f
commit
9c55a0a0ff
6 changed files with 654 additions and 1067 deletions
|
|
@ -3,31 +3,27 @@ Shows all defined flow blueprints.
|
|||
"""
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import os
|
||||
import tabulate
|
||||
from trustgraph.api import Api, ConfigKey
|
||||
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, config_api):
|
||||
def format_parameters(params_metadata, param_type_defs):
|
||||
"""
|
||||
Format parameter metadata for display
|
||||
Format parameter metadata for display.
|
||||
|
||||
Args:
|
||||
params_metadata: Parameter definitions from flow blueprint
|
||||
config_api: API client to get parameter type information
|
||||
|
||||
Returns:
|
||||
Formatted string describing parameters
|
||||
param_type_defs is a dict of type_name -> parsed type definition,
|
||||
pre-fetched concurrently.
|
||||
"""
|
||||
if not params_metadata:
|
||||
return "None"
|
||||
|
||||
param_list = []
|
||||
|
||||
# Sort parameters by order if available
|
||||
sorted_params = sorted(
|
||||
params_metadata.items(),
|
||||
key=lambda x: x[1].get("order", 999)
|
||||
|
|
@ -37,41 +33,89 @@ def format_parameters(params_metadata, config_api):
|
|||
description = param_meta.get("description", param_name)
|
||||
param_type = param_meta.get("type", "unknown")
|
||||
|
||||
# Get type information if available
|
||||
type_info = param_type
|
||||
if config_api:
|
||||
try:
|
||||
key = ConfigKey("parameter-type", param_type)
|
||||
type_def_value = config_api.get([key])[0].value
|
||||
param_type_def = json.loads(type_def_value)
|
||||
|
||||
# Add default value if available
|
||||
default = param_type_def.get("default")
|
||||
if default is not None:
|
||||
type_info = f"{param_type} (default: {default})"
|
||||
|
||||
except:
|
||||
# If we can't get type definition, just show the type name
|
||||
pass
|
||||
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):
|
||||
|
||||
api = Api(url, token=token)
|
||||
flow_api = api.flow()
|
||||
config_api = api.config()
|
||||
blueprint_names, blueprints, param_type_defs = asyncio.run(
|
||||
_show_flow_blueprints_async(url, token=token)
|
||||
)
|
||||
|
||||
blueprint_names = flow_api.list_blueprints()
|
||||
|
||||
if len(blueprint_names) == 0:
|
||||
if not blueprint_names:
|
||||
print("No flow blueprints.")
|
||||
return
|
||||
|
||||
for blueprint_name in blueprint_names:
|
||||
cls = flow_api.get_blueprint(blueprint_name)
|
||||
cls = blueprints[blueprint_name]
|
||||
|
||||
table = []
|
||||
table.append(("name", blueprint_name))
|
||||
|
|
@ -81,10 +125,9 @@ def show_flow_blueprints(url, token=None):
|
|||
if tags:
|
||||
table.append(("tags", ", ".join(tags)))
|
||||
|
||||
# Show parameters if they exist
|
||||
parameters = cls.get("parameters", {})
|
||||
if parameters:
|
||||
param_str = format_parameters(parameters, config_api)
|
||||
param_str = format_parameters(parameters, param_type_defs)
|
||||
table.append(("parameters", param_str))
|
||||
|
||||
print(tabulate.tabulate(
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue