diff --git a/trustgraph-cli/pyproject.toml b/trustgraph-cli/pyproject.toml index 42b9fb9b..9fd6bed7 100644 --- a/trustgraph-cli/pyproject.toml +++ b/trustgraph-cli/pyproject.toml @@ -35,6 +35,7 @@ tg-delete-kg-core = "trustgraph.cli.delete_kg_core:main" tg-delete-tool = "trustgraph.cli.delete_tool:main" tg-dump-msgpack = "trustgraph.cli.dump_msgpack:main" tg-dump-queues = "trustgraph.cli.dump_queues:main" +tg-monitor-prompts = "trustgraph.cli.monitor_prompts:main" tg-get-flow-blueprint = "trustgraph.cli.get_flow_blueprint:main" tg-get-kg-core = "trustgraph.cli.get_kg_core:main" tg-get-document-content = "trustgraph.cli.get_document_content:main" diff --git a/trustgraph-cli/trustgraph/cli/monitor_prompts.py b/trustgraph-cli/trustgraph/cli/monitor_prompts.py new file mode 100644 index 00000000..9c48b5ac --- /dev/null +++ b/trustgraph-cli/trustgraph/cli/monitor_prompts.py @@ -0,0 +1,327 @@ +""" +Monitor prompt request/response queues and log activity with timing. + +Subscribes to prompt request and response Pulsar queues, correlates +them by message ID, and logs a summary of each request/response with +elapsed time. + +Examples: + tg-monitor-prompts + tg-monitor-prompts --flow default --max-lines 5 + tg-monitor-prompts --queue-type prompt-rag +""" + +import json +import asyncio +import sys +import argparse +from datetime import datetime +from collections import OrderedDict + +from trustgraph.base.subscriber import Subscriber +from trustgraph.base.pubsub import get_pubsub + +default_flow = "default" +default_queue_type = "prompt" +default_max_lines = 3 +default_max_width = 80 + + +def truncate_text(text, max_lines, max_width): + """Truncate text to max_lines lines, each at most max_width chars.""" + if not text: + return "(empty)" + + lines = text.splitlines() + result = [] + for line in lines[:max_lines]: + if len(line) > max_width: + result.append(line[:max_width - 3] + "...") + else: + result.append(line) + + if len(lines) > max_lines: + result.append(f" ... ({len(lines) - max_lines} more lines)") + + return "\n".join(result) + + +def format_terms(terms, max_lines, max_width): + """Format prompt terms for display.""" + if not terms: + return "(no terms)" + + parts = [] + for key, value in terms.items(): + # Try to parse JSON-encoded values + try: + parsed = json.loads(value) + if isinstance(parsed, str): + value = parsed + else: + value = json.dumps(parsed, indent=2) + except (json.JSONDecodeError, TypeError): + pass + + truncated = truncate_text(str(value), max_lines, max_width) + parts.append(f" {key}: {truncated}") + + 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).""" + # Get correlation ID from properties + try: + props = msg.properties() + msg_id = props.get("id", "") + except Exception: + msg_id = "" + + # Parse body + try: + value = msg.value() + if isinstance(value, bytes): + value = value.decode("utf-8") + body = json.loads(value) if isinstance(value, str) else {} + except Exception: + body = {} + + return msg_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}") + + +async def async_main(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}" + + print(f"Monitoring prompt queues:") + print(f" Request: {request_queue}") + 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", + ) + + central_queue = asyncio.Queue() + shutdown_event = asyncio.Event() + subscribers = [] + + 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)) + + 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) + + processor_task = asyncio.create_task( + process_messages(central_queue, request_queue, response_queue, + max_lines, max_width, shutdown_event) + ) + tasks.append(processor_task) + + print("Listening...\n") + await asyncio.gather(*tasks) + + 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() + + +def main(): + parser = argparse.ArgumentParser( + prog="tg-monitor-prompts", + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + + parser.add_argument( + "-f", "--flow", + default=default_flow, + help=f"Flow ID (default: {default_flow})", + ) + + parser.add_argument( + "-q", "--queue-type", + default=default_queue_type, + help=f"Queue type: prompt or prompt-rag (default: {default_queue_type})", + ) + + parser.add_argument( + "-l", "--max-lines", + type=int, + default=default_max_lines, + help=f"Max lines of text to show per term/response (default: {default_max_lines})", + ) + + parser.add_argument( + "-w", "--max-width", + type=int, + default=default_max_width, + help=f"Max width per line (default: {default_max_width})", + ) + + parser.add_argument( + "--pulsar-host", + default="pulsar://localhost:6650", + help="Pulsar host URL (default: pulsar://localhost:6650)", + ) + + parser.add_argument( + "--listener-name", + default="localhost", + help="Pulsar listener name (default: localhost)", + ) + + args = parser.parse_args() + + try: + asyncio.run(async_main( + flow=args.flow, + queue_type=args.queue_type, + max_lines=args.max_lines, + max_width=args.max_width, + pulsar_host=args.pulsar_host, + listener_name=args.listener_name, + )) + except KeyboardInterrupt: + pass + except Exception as e: + print(f"Fatal error: {e}", file=sys.stderr) + sys.exit(1) + + +if __name__ == "__main__": + main()