From 734382c44c45db0332a391d07fde5647ec09f05f Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Mon, 30 Mar 2026 10:37:07 +0100 Subject: [PATCH] Output working but messy --- .../trustgraph/cli/monitor_prompts.py | 273 ++++++++---------- 1 file changed, 117 insertions(+), 156 deletions(-) diff --git a/trustgraph-cli/trustgraph/cli/monitor_prompts.py b/trustgraph-cli/trustgraph/cli/monitor_prompts.py index 9c48b5ac..f7922cce 100644 --- a/trustgraph-cli/trustgraph/cli/monitor_prompts.py +++ b/trustgraph-cli/trustgraph/cli/monitor_prompts.py @@ -18,8 +18,9 @@ import argparse from datetime import datetime from collections import OrderedDict -from trustgraph.base.subscriber import Subscriber -from trustgraph.base.pubsub import get_pubsub +import pulsar +from pulsar.schema import BytesSchema + default_flow = "default" default_queue_type = "prompt" @@ -49,7 +50,7 @@ def truncate_text(text, max_lines, max_width): def format_terms(terms, max_lines, max_width): """Format prompt terms for display.""" if not terms: - return "(no terms)" + return "" parts = [] for key, value in terms.items(): @@ -69,40 +70,14 @@ def format_terms(terms, max_lines, max_width): return "\n".join(parts) -async def monitor_queue(subscriber, queue_name, central_queue, monitor_id, - shutdown_event): - """Monitor a single queue and forward messages to central queue.""" - msg_queue = None - try: - msg_queue = await subscriber.subscribe_all(monitor_id) - - while not shutdown_event.is_set(): - try: - msg = await asyncio.wait_for(msg_queue.get(), timeout=0.5) - timestamp = datetime.now() - await central_queue.put((timestamp, queue_name, msg)) - except asyncio.TimeoutError: - continue - - except Exception as e: - if not shutdown_event.is_set(): - print(f"Error monitoring {queue_name}: {e}", file=sys.stderr) - finally: - if msg_queue is not None: - try: - await subscriber.unsubscribe_all(monitor_id) - except Exception: - pass - - -def parse_message(msg): - """Parse a Pulsar message into (id, body_dict).""" +def parse_raw_message(msg): + """Parse a raw Pulsar message into (correlation_id, body_dict).""" # Get correlation ID from properties try: props = msg.properties() - msg_id = props.get("id", "") + corr_id = props.get("id", "") except Exception: - msg_id = "" + corr_id = "" # Parse body try: @@ -113,88 +88,19 @@ def parse_message(msg): except Exception: body = {} - return msg_id, body + return corr_id, body -async def process_messages(central_queue, request_queue_name, - response_queue_name, max_lines, max_width, - shutdown_event): - """Process messages from central queue, correlating requests/responses.""" - - # Track in-flight requests: id -> (timestamp, template_id) - in_flight = OrderedDict() - - while not shutdown_event.is_set(): - try: - timestamp, queue_name, msg = await asyncio.wait_for( - central_queue.get(), timeout=0.5 - ) - except asyncio.TimeoutError: - continue - - msg_id, body = parse_message(msg) - time_str = timestamp.strftime("%H:%M:%S.%f")[:-3] - - if queue_name == request_queue_name: - # Request - template_id = body.get("id", "(unknown)") - terms = body.get("terms", {}) - streaming = body.get("streaming", False) - - in_flight[msg_id] = (timestamp, template_id) - - # Limit in_flight size to avoid unbounded growth - while len(in_flight) > 1000: - in_flight.popitem(last=False) - - stream_flag = " [streaming]" if streaming else "" - print(f"[{time_str}] REQ {msg_id[:8]} " - f"template={template_id}{stream_flag}") - - if terms: - print(format_terms(terms, max_lines, max_width)) - - elif queue_name == response_queue_name: - # Response - elapsed_str = "" - template_id = "" - - if msg_id in in_flight: - req_timestamp, template_id = in_flight.pop(msg_id) - elapsed = (timestamp - req_timestamp).total_seconds() - elapsed_str = f" ({elapsed:.3f}s)" - - error = body.get("error") - text = body.get("text", "") - obj = body.get("object", "") - eos = body.get("end_of_stream", False) - - if error: - err_msg = error - if isinstance(error, dict): - err_msg = error.get("message", str(error)) - print(f"[{time_str}] ERR {msg_id[:8]} " - f"{err_msg}{elapsed_str}") - elif eos: - print(f"[{time_str}] END {msg_id[:8]}" - f"{elapsed_str}") - elif text: - truncated = truncate_text(text, max_lines, max_width) - print(f"[{time_str}] RESP {msg_id[:8]}" - f"{elapsed_str}") - print(f" {truncated}") - elif obj: - truncated = truncate_text(obj, max_lines, max_width) - print(f"[{time_str}] RESP {msg_id[:8]} " - f"(object){elapsed_str}") - print(f" {truncated}") - else: - print(f"[{time_str}] RESP {msg_id[:8]} " - f"(empty){elapsed_str}") +def receive_with_timeout(consumer, timeout_ms=500): + """Receive a message with timeout, returning None on timeout.""" + try: + return consumer.receive(timeout_millis=timeout_ms) + except Exception: + return None -async def async_main(flow, queue_type, max_lines, max_width, - pulsar_host, listener_name): +async def monitor(flow, queue_type, max_lines, max_width, + pulsar_host, listener_name): request_queue = f"non-persistent://tg/request/{queue_type}:{flow}" response_queue = f"non-persistent://tg/response/{queue_type}:{flow}" @@ -204,60 +110,115 @@ async def async_main(flow, queue_type, max_lines, max_width, print(f" Response: {response_queue}") print(f"Press Ctrl+C to stop\n") - backend = get_pubsub( - pulsar_host=pulsar_host, - pulsar_listener=listener_name, - pubsub_backend="pulsar", + client = pulsar.Client( + pulsar_host, + listener_name=listener_name, ) - central_queue = asyncio.Queue() - shutdown_event = asyncio.Event() - subscribers = [] + req_consumer = client.subscribe( + request_queue, + subscription_name="prompt-monitor-req", + consumer_type=pulsar.ConsumerType.Shared, + schema=BytesSchema(), + initial_position=pulsar.InitialPosition.Latest, + ) + + resp_consumer = client.subscribe( + response_queue, + subscription_name="prompt-monitor-resp", + consumer_type=pulsar.ConsumerType.Shared, + schema=BytesSchema(), + initial_position=pulsar.InitialPosition.Latest, + ) + + # Track in-flight requests: corr_id -> (timestamp, template_id) + in_flight = OrderedDict() + + print("Listening...\n") try: - for queue_name in [request_queue, response_queue]: - sub = Subscriber( - backend=backend, - topic=queue_name, - subscription="prompt-monitor", - consumer_name=f"prompt-monitor-{queue_name}", - schema=None, - ) - await sub.start() - subscribers.append((queue_name, sub)) + while True: + # Poll both queues + msg = receive_with_timeout(req_consumer, 100) + if msg: + timestamp = datetime.now() + corr_id, body = parse_raw_message(msg) + time_str = timestamp.strftime("%H:%M:%S.%f")[:-3] - tasks = [] - for queue_name, sub in subscribers: - task = asyncio.create_task( - monitor_queue(sub, queue_name, central_queue, - "prompt-monitor", shutdown_event) - ) - tasks.append(task) + template_id = body.get("id", "(unknown)") + terms = body.get("terms", {}) + streaming = body.get("streaming", False) - processor_task = asyncio.create_task( - process_messages(central_queue, request_queue, response_queue, - max_lines, max_width, shutdown_event) - ) - tasks.append(processor_task) + in_flight[corr_id] = (timestamp, template_id) - print("Listening...\n") - await asyncio.gather(*tasks) + # Limit size + while len(in_flight) > 1000: + in_flight.popitem(last=False) + + stream_flag = " [streaming]" if streaming else "" + id_display = corr_id[:8] if corr_id else "--------" + print(f"[{time_str}] REQ {id_display} " + f"template={template_id}{stream_flag}") + + if terms: + print(format_terms(terms, max_lines, max_width)) + + req_consumer.acknowledge(msg) + + msg = receive_with_timeout(resp_consumer, 100) + if msg: + timestamp = datetime.now() + corr_id, body = parse_raw_message(msg) + time_str = timestamp.strftime("%H:%M:%S.%f")[:-3] + + elapsed_str = "" + if corr_id in in_flight: + req_timestamp, template_id = in_flight.pop(corr_id) + elapsed = (timestamp - req_timestamp).total_seconds() + elapsed_str = f" ({elapsed:.3f}s)" + + id_display = corr_id[:8] if corr_id else "--------" + + error = body.get("error") + text = body.get("text", "") + obj = body.get("object", "") + eos = body.get("end_of_stream", False) + + if error: + err_msg = error + if isinstance(error, dict): + err_msg = error.get("message", str(error)) + print(f"[{time_str}] ERR {id_display} " + f"{err_msg}{elapsed_str}") + elif eos: + print(f"[{time_str}] END {id_display}" + f"{elapsed_str}") + elif text: + truncated = truncate_text(text, max_lines, max_width) + print(f"[{time_str}] RESP {id_display}" + f"{elapsed_str}") + print(f" {truncated}") + elif obj: + truncated = truncate_text(obj, max_lines, max_width) + print(f"[{time_str}] RESP {id_display} " + f"(object){elapsed_str}") + print(f" {truncated}") + else: + print(f"[{time_str}] RESP {id_display} " + f"(empty){elapsed_str}") + + resp_consumer.acknowledge(msg) + + # Small sleep if no messages from either queue + if not msg: + await asyncio.sleep(0.05) except KeyboardInterrupt: print("\nStopping...") finally: - shutdown_event.set() - try: - await asyncio.wait_for( - asyncio.gather(*tasks, return_exceptions=True), - timeout=2.0, - ) - except asyncio.TimeoutError: - pass - - for _, sub in subscribers: - await sub.stop() - backend.close() + req_consumer.close() + resp_consumer.close() + client.close() def main(): @@ -283,7 +244,7 @@ def main(): "-l", "--max-lines", type=int, default=default_max_lines, - help=f"Max lines of text to show per term/response (default: {default_max_lines})", + help=f"Max lines of text per term/response (default: {default_max_lines})", ) parser.add_argument( @@ -308,7 +269,7 @@ def main(): args = parser.parse_args() try: - asyncio.run(async_main( + asyncio.run(monitor( flow=args.flow, queue_type=args.queue_type, max_lines=args.max_lines,