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

View file

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

View file

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

View file

@ -90,4 +90,47 @@ class ProcessorMetrics:
__class__.processor_metric.labels(
processor = self.processor
).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 . producer import Producer
from . spec import Spec
from . metrics import ConsumerMetrics, ProducerMetrics
from . metrics import ConsumerMetrics, ProducerMetrics, SubscriberMetrics
class RequestResponse(Subscriber):
@ -23,6 +23,7 @@ class RequestResponse(Subscriber):
consumer_name = consumer_name,
topic = response_topic,
schema = response_schema,
metrics = response_metrics,
)
self.producer = Producer(
@ -116,8 +117,12 @@ class RequestResponseSpec(Spec):
def add(self, flow, processor, definition):
producer_metrics = ProducerMetrics(
processor = flow.id, flow = flow.name, name = self.response_name
request_metrics = ProducerMetrics(
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(
@ -126,10 +131,10 @@ class RequestResponseSpec(Spec):
consumer_name = flow.id,
request_topic = definition[self.request_name],
request_schema = self.request_schema,
request_metrics = producer_metrics,
request_metrics = request_metrics,
response_topic = definition[self.response_name],
response_schema = self.response_schema,
response_metrics = None,
response_metrics = response_metrics,
)
flow.consumer[self.request_name] = rr

View file

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

View file

@ -1,5 +1,5 @@
from . metrics import ConsumerMetrics
from . metrics import SubscriberMetrics
from . subscriber import Subscriber
from . spec import Spec
@ -11,8 +11,7 @@ class SubscriberSpec(Spec):
def add(self, flow, processor, definition):
# FIXME: Metrics not used
subscriber_metrics = ConsumerMetrics(
subscriber_metrics = SubscriberMetrics(
processor = flow.id, flow = flow.name, name = self.name
)
@ -22,6 +21,7 @@ class SubscriberSpec(Spec):
subscription = flow.id,
consumer_name = flow.id,
schema = self.schema,
metrics = subscriber_metrics,
)
# 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
"""
Dumps out the current configuration
"""
import argparse
import os
import tabulate
from trustgraph.api import Api
import json
@ -15,19 +15,31 @@ def show_flows(url):
api = Api(url)
BROKEN
flow_ids = api.flow_list()
print(flow_ids)
flows = []
config, version = api.config_all()
print("Version:", version)
print(json.dumps(config, indent=4))
for id in flow_ids:
flow = api.flow_get(id)
flows.append((
id,
flow.get("description", ""),
))
print(tabulate.tabulate(
flows,
tablefmt="pretty",
maxcolwidths=[None, 40],
stralign="left",
headers = ["id", "description"],
))
def main():
parser = argparse.ArgumentParser(
prog='tg-show-config',
prog='tg-show-flows',
description=__doc__,
)
@ -41,7 +53,7 @@ def main():
try:
show_config(
show_flows(
url=args.api_url,
)

View file

@ -63,6 +63,13 @@ setuptools.setup(
"scripts/tg-save-kg-core",
"scripts/tg-save-doc-embeds",
"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-show-tools",
"scripts/tg-show-prompts",

View file

@ -35,6 +35,8 @@ class FlowConfig:
async def handle_delete_class(self, msg):
print(msg)
del self.config["flow-classes"][msg.class_name]
await self.config.push()
@ -63,6 +65,21 @@ class FlowConfig:
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):
return tmp.replace(
"{class}", msg.class_name
@ -72,8 +89,6 @@ class FlowConfig:
cls = json.loads(self.config["flow-classes"][msg.class_name])
plumb = {}
for kind in ("class", "flow"):
for k, v in cls[kind].items():
@ -97,10 +112,10 @@ class FlowConfig:
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,
"class-name": msg.class_name,
}
})
await self.config.push()
@ -110,7 +125,20 @@ class FlowConfig:
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):
return tmp.replace(
@ -119,11 +147,7 @@ class FlowConfig:
"{id}", msg.flow_id
)
cls = json.loads(self.config["flow-classes"][class_name])
plumb = {}
for kind in ("flow"):
for kind in ("flow",):
for k, v in cls[kind].items():
@ -150,19 +174,6 @@ class FlowConfig:
error = None,
)
flow = self.config["flows"][msg.flow_id]
return FlowResponse(
error = None,
)
async def handle(self, msg):
print("Handle message ", msg.operation)
@ -197,4 +208,3 @@ class FlowConfig:
return resp