Flow management CLI and API working

This commit is contained in:
Cyber MacGeddon 2025-04-24 18:54:38 +01:00
parent 3d284baff9
commit 58228c656f
14 changed files with 383 additions and 54 deletions

View file

@ -596,7 +596,7 @@ class Api:
# The input consists of system and prompt strings # The input consists of system and prompt strings
input = { input = {
"operation": "get-class", "operation": "get-class",
"flow-id": class_name, "class-name": class_name,
} }
url = f"{self.url}flow" url = f"{self.url}flow"
@ -617,8 +617,9 @@ class Api:
self.check_error(object) self.check_error(object)
try: try:
return json.load(object["class-definition"]) return json.loads(object["class-definition"])
except: except Exception as e:
print(e)
raise ProtocolException(f"Response not formatted correctly") raise ProtocolException(f"Response not formatted correctly")
def flow_put_class(self, class_name, definition): def flow_put_class(self, class_name, definition):
@ -626,8 +627,8 @@ class Api:
# The input consists of system and prompt strings # The input consists of system and prompt strings
input = { input = {
"operation": "put-class", "operation": "put-class",
"flow-id": class_name, "class-name": class_name,
"definition": json.dumps(definition), "class-definition": json.dumps(definition),
} }
url = f"{self.url}flow" url = f"{self.url}flow"
@ -654,7 +655,7 @@ class Api:
# The input consists of system and prompt strings # The input consists of system and prompt strings
input = { input = {
"operation": "delete-class", "operation": "delete-class",
"flow-id": class_name, "class-name": class_name,
} }
url = f"{self.url}flow" url = f"{self.url}flow"
@ -742,6 +743,7 @@ class Api:
"operation": "start-flow", "operation": "start-flow",
"flow-id": id, "flow-id": id,
"class-name": class_name, "class-name": class_name,
"description": description,
} }
url = f"{self.url}flow" url = f"{self.url}flow"

View file

@ -17,7 +17,7 @@ from .. exceptions import TooManyRequests
from . pubsub import PulsarClient from . pubsub import PulsarClient
from . producer import Producer from . producer import Producer
from . consumer import Consumer from . consumer import Consumer
from . metrics import ProcessorMetrics from . metrics import ProcessorMetrics, ConsumerMetrics
default_config_queue = config_push_queue default_config_queue = config_push_queue
@ -57,6 +57,10 @@ class AsyncProcessor:
# service # service
config_subscriber_id = str(uuid.uuid4()) config_subscriber_id = str(uuid.uuid4())
config_consumer_metrics = ConsumerMetrics(
processor = self.id, flow = None, name = "config",
)
# Subscribe to config queue # Subscribe to config queue
self.config_sub_task = Consumer( self.config_sub_task = Consumer(
@ -70,6 +74,8 @@ class AsyncProcessor:
handler = self.on_config_change, handler = self.on_config_change,
metrics = config_consumer_metrics,
# This causes new subscriptions to view the entire history of # This causes new subscriptions to view the entire history of
# configuration # configuration
start_of_messages = True start_of_messages = True
@ -101,15 +107,12 @@ class AsyncProcessor:
self.config_handlers.append(handler) self.config_handlers.append(handler)
# Called when a new configuration message push occurs # Called when a new configuration message push occurs
async def on_config_change(self, message, consumer): async def on_config_change(self, message, consumer, flow):
# Get configuration data and version number # Get configuration data and version number
config = message.value().config config = message.value().config
version = message.value().version version = message.value().version
# Acknowledge the message
consumer.acknowledge(message)
# Invoke message handlers # Invoke message handlers
print("Config change event", config, version, flush=True) print("Config change event", config, version, flush=True)
for ch in self.config_handlers: for ch in self.config_handlers:

View file

@ -156,7 +156,7 @@ class Consumer:
await self.handler(msg, self, self.flow) await self.handler(msg, self, self.flow)
else: else:
await self.handler(msg, self.consumer) await self.handler(msg, self, self.flow)
print("Handled.", flush=True) print("Handled.", flush=True)

View file

@ -90,4 +90,47 @@ class ProcessorMetrics:
__class__.processor_metric.labels( __class__.processor_metric.labels(
processor = self.processor processor = self.processor
).info(info) ).info(info)
class SubscriberMetrics:
def __init__(self, processor, flow, name):
self.processor = processor
self.flow = flow
self.name = name
if not hasattr(__class__, "state_metric"):
__class__.state_metric = Enum(
'subscriber_state', 'Subscriber state',
["processor", "flow", "name"],
states=['stopped', 'running']
)
if not hasattr(__class__, "received_metric"):
__class__.received_metric = Counter(
'received_count', 'Received count',
["processor", "flow", "name"],
)
if not hasattr(__class__, "dropped_metric"):
__class__.dropped_metric = Counter(
'dropped_count', 'Dropped messages count',
["processor", "flow", "name"],
)
def received(self):
__class__.received_metric.labels(
processor = self.processor, flow = self.flow, name = self.name,
).inc()
def state(self, state):
__class__.state_metric.labels(
processor = self.processor, flow = self.flow, name = self.name,
).state(state)
def dropped(self, state):
__class__.dropped_metric.labels(
processor = self.processor, flow = self.flow, name = self.name,
).inc()

View file

@ -5,7 +5,7 @@ import asyncio
from . subscriber import Subscriber from . subscriber import Subscriber
from . producer import Producer from . producer import Producer
from . spec import Spec from . spec import Spec
from . metrics import ConsumerMetrics, ProducerMetrics from . metrics import ConsumerMetrics, ProducerMetrics, SubscriberMetrics
class RequestResponse(Subscriber): class RequestResponse(Subscriber):
@ -23,6 +23,7 @@ class RequestResponse(Subscriber):
consumer_name = consumer_name, consumer_name = consumer_name,
topic = response_topic, topic = response_topic,
schema = response_schema, schema = response_schema,
metrics = response_metrics,
) )
self.producer = Producer( self.producer = Producer(
@ -116,8 +117,12 @@ class RequestResponseSpec(Spec):
def add(self, flow, processor, definition): def add(self, flow, processor, definition):
producer_metrics = ProducerMetrics( request_metrics = ProducerMetrics(
processor = flow.id, flow = flow.name, name = self.response_name processor = flow.id, flow = flow.name, name = self.request_name
)
response_metrics = SubscriberMetrics(
processor = flow.id, flow = flow.name, name = self.request_name
) )
rr = self.impl( rr = self.impl(
@ -126,10 +131,10 @@ class RequestResponseSpec(Spec):
consumer_name = flow.id, consumer_name = flow.id,
request_topic = definition[self.request_name], request_topic = definition[self.request_name],
request_schema = self.request_schema, request_schema = self.request_schema,
request_metrics = producer_metrics, request_metrics = request_metrics,
response_topic = definition[self.response_name], response_topic = definition[self.response_name],
response_schema = self.response_schema, response_schema = self.response_schema,
response_metrics = None, response_metrics = response_metrics,
) )
flow.consumer[self.request_name] = rr flow.consumer[self.request_name] = rr

View file

@ -7,7 +7,7 @@ import time
class Subscriber: class Subscriber:
def __init__(self, client, topic, subscription, consumer_name, def __init__(self, client, topic, subscription, consumer_name,
schema=None, max_size=100): schema=None, max_size=100, metrics=None):
self.client = client self.client = client
self.topic = topic self.topic = topic
self.subscription = subscription self.subscription = subscription
@ -18,6 +18,7 @@ class Subscriber:
self.max_size = max_size self.max_size = max_size
self.lock = asyncio.Lock() self.lock = asyncio.Lock()
self.running = True self.running = True
self.metrics = metrics
async def __del__(self): async def __del__(self):
self.running = False self.running = False
@ -36,6 +37,9 @@ class Subscriber:
while self.running: while self.running:
if self.metrics:
self.metrics.state("stopped")
try: try:
consumer = self.client.subscribe( consumer = self.client.subscribe(
@ -45,6 +49,9 @@ class Subscriber:
schema = JsonSchema(self.schema), schema = JsonSchema(self.schema),
) )
if self.metrics:
self.metrics.state("running")
print("Subscriber running...", flush=True) print("Subscriber running...", flush=True)
while self.running: while self.running:
@ -61,6 +68,9 @@ class Subscriber:
print(type(e)) print(type(e))
raise e raise e
if self.metrics:
self.metrics.received()
# Acknowledge successful reception of the message # Acknowledge successful reception of the message
consumer.acknowledge(msg) consumer.acknowledge(msg)
@ -83,7 +93,9 @@ class Subscriber:
self.q[id].put(value), self.q[id].put(value),
timeout=2 timeout=2
) )
except Exception as e: except Exception as e:
self.metrics.dropped()
print("Q Put:", e, flush=True) print("Q Put:", e, flush=True)
for q in self.full.values(): for q in self.full.values():
@ -94,6 +106,7 @@ class Subscriber:
timeout=2 timeout=2
) )
except Exception as e: except Exception as e:
self.metrics.dropped()
print("Q Put:", e, flush=True) print("Q Put:", e, flush=True)
except Exception as e: except Exception as e:
@ -101,6 +114,9 @@ class Subscriber:
consumer.close() consumer.close()
if self.metrics:
self.metrics.state("stopped")
# If handler drops out, sleep a retry # If handler drops out, sleep a retry
time.sleep(2) time.sleep(2)

View file

@ -1,5 +1,5 @@
from . metrics import ConsumerMetrics from . metrics import SubscriberMetrics
from . subscriber import Subscriber from . subscriber import Subscriber
from . spec import Spec from . spec import Spec
@ -11,8 +11,7 @@ class SubscriberSpec(Spec):
def add(self, flow, processor, definition): def add(self, flow, processor, definition):
# FIXME: Metrics not used subscriber_metrics = SubscriberMetrics(
subscriber_metrics = ConsumerMetrics(
processor = flow.id, flow = flow.name, name = self.name processor = flow.id, flow = flow.name, name = self.name
) )
@ -22,6 +21,7 @@ class SubscriberSpec(Spec):
subscription = flow.id, subscription = flow.id,
consumer_name = flow.id, consumer_name = flow.id,
schema = self.schema, schema = self.schema,
metrics = subscriber_metrics,
) )
# Put it in the consumer map, does that work? # Put it in the consumer map, does that work?

View file

@ -0,0 +1,52 @@
#!/usr/bin/env python3
"""
"""
import argparse
import os
import tabulate
from trustgraph.api import Api
import json
default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
def delete_flow_class(url, class_name):
api = Api(url)
class_names = api.flow_delete_class(class_name)
def main():
parser = argparse.ArgumentParser(
prog='tg-delete-flow-class',
description=__doc__,
)
parser.add_argument(
'-u', '--api-url',
default=default_url,
help=f'API URL (default: {default_url})',
)
parser.add_argument(
'-n', '--class-name',
help=f'Flow class name',
)
args = parser.parse_args()
try:
delete_flow_class(
url=args.api_url,
class_name=args.class_name,
)
except Exception as e:
print("Exception:", e, flush=True)
main()

View file

@ -0,0 +1,56 @@
#!/usr/bin/env python3
"""
Dumps out the current configuration
"""
import argparse
import os
import tabulate
from trustgraph.api import Api
import json
default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
def get_flow_class(url, class_name):
api = Api(url)
cls = api.flow_get_class(class_name)
print(json.dumps(cls, indent=4))
def main():
parser = argparse.ArgumentParser(
prog='tg-get-flow-class',
description=__doc__,
)
parser.add_argument(
'-u', '--api-url',
default=default_url,
help=f'API URL (default: {default_url})',
)
parser.add_argument(
'-n', '--class-name',
required=True,
help=f'Flow class name',
)
args = parser.parse_args()
try:
get_flow_class(
url=args.api_url,
class_name=args.class_name,
)
except Exception as e:
print("Exception:", e, flush=True)
main()

View file

@ -0,0 +1,59 @@
#!/usr/bin/env python3
"""
Dumps out the current configuration
"""
import argparse
import os
import tabulate
from trustgraph.api import Api
import json
default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
def put_flow_class(url, class_name, config):
api = Api(url)
class_names = api.flow_put_class(class_name, config)
def main():
parser = argparse.ArgumentParser(
prog='tg-put-flow-class',
description=__doc__,
)
parser.add_argument(
'-u', '--api-url',
default=default_url,
help=f'API URL (default: {default_url})',
)
parser.add_argument(
'-n', '--class-name',
help=f'Flow class name',
)
parser.add_argument(
'-c', '--config',
help=f'Initial configuration to load',
)
args = parser.parse_args()
try:
put_flow_class(
url=args.api_url,
class_name=args.class_name,
config=json.loads(args.config),
)
except Exception as e:
print("Exception:", e, flush=True)
main()

View file

@ -0,0 +1,64 @@
#!/usr/bin/env python3
"""
"""
import argparse
import os
import tabulate
from trustgraph.api import Api
import json
default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
def show_flow_classes(url):
api = Api(url)
class_names = api.flow_list_classes()
classes = []
for class_name in class_names:
cls = api.flow_get_class(class_name)
classes.append((
class_name,
cls.get("description", ""),
", ".join(cls.get("tags", [])),
))
print(tabulate.tabulate(
classes,
tablefmt="pretty",
maxcolwidths=[None, 40, 20],
stralign="left",
headers = ["flow class", "description", "tags"],
))
def main():
parser = argparse.ArgumentParser(
prog='tg-show-flow-classes',
description=__doc__,
)
parser.add_argument(
'-u', '--api-url',
default=default_url,
help=f'API URL (default: {default_url})',
)
args = parser.parse_args()
try:
show_flow_classes(
url=args.api_url,
)
except Exception as e:
print("Exception:", e, flush=True)
main()

View file

@ -1,11 +1,11 @@
#!/usr/bin/env python3 #!/usr/bin/env python3
""" """
Dumps out the current configuration
""" """
import argparse import argparse
import os import os
import tabulate
from trustgraph.api import Api from trustgraph.api import Api
import json import json
@ -15,19 +15,31 @@ def show_flows(url):
api = Api(url) api = Api(url)
BROKEN flow_ids = api.flow_list()
print(flow_ids)
flows = []
config, version = api.config_all() for id in flow_ids:
flow = api.flow_get(id)
print("Version:", version) flows.append((
print(json.dumps(config, indent=4)) id,
flow.get("description", ""),
))
print(tabulate.tabulate(
flows,
tablefmt="pretty",
maxcolwidths=[None, 40],
stralign="left",
headers = ["id", "description"],
))
def main(): def main():
parser = argparse.ArgumentParser( parser = argparse.ArgumentParser(
prog='tg-show-config', prog='tg-show-flows',
description=__doc__, description=__doc__,
) )
@ -41,7 +53,7 @@ def main():
try: try:
show_config( show_flows(
url=args.api_url, url=args.api_url,
) )

View file

@ -63,6 +63,13 @@ setuptools.setup(
"scripts/tg-save-kg-core", "scripts/tg-save-kg-core",
"scripts/tg-save-doc-embeds", "scripts/tg-save-doc-embeds",
"scripts/tg-show-config", "scripts/tg-show-config",
"scripts/tg-show-flows",
"scripts/tg-show-flow-classes",
"scripts/tg-get-flow-class",
"scripts/tg-start-flow",
"scripts/tg-stop-flow",
"scripts/tg-delete-flow-class",
"scripts/tg-put-flow-class",
"scripts/tg-set-prompt", "scripts/tg-set-prompt",
"scripts/tg-show-tools", "scripts/tg-show-tools",
"scripts/tg-show-prompts", "scripts/tg-show-prompts",

View file

@ -35,6 +35,8 @@ class FlowConfig:
async def handle_delete_class(self, msg): async def handle_delete_class(self, msg):
print(msg)
del self.config["flow-classes"][msg.class_name] del self.config["flow-classes"][msg.class_name]
await self.config.push() await self.config.push()
@ -63,6 +65,21 @@ class FlowConfig:
async def handle_start_flow(self, msg): async def handle_start_flow(self, msg):
if msg.class_name is None:
raise RuntimeError("No class name")
if msg.flow_id is None:
raise RuntimeError("No flow ID")
if msg.flow_id in self.config["flows"]:
raise RuntimeError("Flow already exists")
if msg.description is None:
raise RuntimeError("No description")
if msg.class_name not in self.config["flow-classes"]:
raise RuntimeError("Class does not exist")
def repl_template(tmp): def repl_template(tmp):
return tmp.replace( return tmp.replace(
"{class}", msg.class_name "{class}", msg.class_name
@ -72,8 +89,6 @@ class FlowConfig:
cls = json.loads(self.config["flow-classes"][msg.class_name]) cls = json.loads(self.config["flow-classes"][msg.class_name])
plumb = {}
for kind in ("class", "flow"): for kind in ("class", "flow"):
for k, v in cls[kind].items(): for k, v in cls[kind].items():
@ -97,10 +112,10 @@ class FlowConfig:
self.config["flows-active"][processor] = json.dumps(target) self.config["flows-active"][processor] = json.dumps(target)
self.config["flows"][msg.flow_id] = { self.config["flows"][msg.flow_id] = json.dumps({
"description": msg.description, "description": msg.description,
"class-name": msg.class_name, "class-name": msg.class_name,
} })
await self.config.push() await self.config.push()
@ -110,7 +125,20 @@ class FlowConfig:
async def handle_stop_flow(self, msg): async def handle_stop_flow(self, msg):
class_name = self.config["flows"][msg.flow_id]["class-name"] if msg.flow_id is None:
raise RuntimeError("No flow ID")
if msg.flow_id not in self.config["flows"]:
raise RuntimeError("Flow ID invalid")
flow = json.loads(self.config["flows"][msg.flow_id])
if "class-name" not in flow:
raise RuntimeError("Internal error: flow has no flow class")
class_name = flow["class-name"]
cls = json.loads(self.config["flow-classes"][class_name])
def repl_template(tmp): def repl_template(tmp):
return tmp.replace( return tmp.replace(
@ -119,11 +147,7 @@ class FlowConfig:
"{id}", msg.flow_id "{id}", msg.flow_id
) )
cls = json.loads(self.config["flow-classes"][class_name]) for kind in ("flow",):
plumb = {}
for kind in ("flow"):
for k, v in cls[kind].items(): for k, v in cls[kind].items():
@ -150,19 +174,6 @@ class FlowConfig:
error = None, error = None,
) )
flow = self.config["flows"][msg.flow_id]
return FlowResponse(
error = None,
)
async def handle(self, msg): async def handle(self, msg):
print("Handle message ", msg.operation) print("Handle message ", msg.operation)
@ -197,4 +208,3 @@ class FlowConfig:
return resp return resp