Output working but messy

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
+import pulsar
-from trustgraph.base.pubsub import get_pubsub
+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,
+def parse_raw_message(msg):
-                        shutdown_event):
+    """Parse a raw Pulsar message into (correlation_id, body_dict)."""
-    """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", "")
+        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,
+def receive_with_timeout(consumer, timeout_ms=500):
-                           response_queue_name, max_lines, max_width,
+    """Receive a message with timeout, returning None on timeout."""
-                           shutdown_event):
+    try:
-    """Process messages from central queue, correlating requests/responses."""
+        return consumer.receive(timeout_millis=timeout_ms)
-
+    except Exception:
-    # Track in-flight requests: id -> (timestamp, template_id)
+        return None
-    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,
+async def monitor(flow, queue_type, max_lines, max_width,
-                     pulsar_host, listener_name):
+                  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(
+    client = pulsar.Client(
-        pulsar_host=pulsar_host,
+        pulsar_host,
-        pulsar_listener=listener_name,
+        listener_name=listener_name,
-        pubsub_backend="pulsar",
     )
 
-    central_queue = asyncio.Queue()
+    req_consumer = client.subscribe(
-    shutdown_event = asyncio.Event()
+        request_queue,
-    subscribers = []
+        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]:
+        while True:
-            sub = Subscriber(
+            # Poll both queues
-                backend=backend,
+            msg = receive_with_timeout(req_consumer, 100)
-                topic=queue_name,
+            if msg:
-                subscription="prompt-monitor",
+                timestamp = datetime.now()
-                consumer_name=f"prompt-monitor-{queue_name}",
+                corr_id, body = parse_raw_message(msg)
-                schema=None,
+                time_str = timestamp.strftime("%H:%M:%S.%f")[:-3]
-            )
-            await sub.start()
-            subscribers.append((queue_name, sub))
 
-        tasks = []
+                template_id = body.get("id", "(unknown)")
-        for queue_name, sub in subscribers:
+                terms = body.get("terms", {})
-            task = asyncio.create_task(
+                streaming = body.get("streaming", False)
-                monitor_queue(sub, queue_name, central_queue,
-                              "prompt-monitor", shutdown_event)
-            )
-            tasks.append(task)
 
-        processor_task = asyncio.create_task(
+                in_flight[corr_id] = (timestamp, template_id)
-            process_messages(central_queue, request_queue, response_queue,
-                             max_lines, max_width, shutdown_event)
-        )
-        tasks.append(processor_task)
 
-        print("Listening...\n")
+                # Limit size
-        await asyncio.gather(*tasks)
+                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()
+        req_consumer.close()
-        try:
+        resp_consumer.close()
-            await asyncio.wait_for(
+        client.close()
-                asyncio.gather(*tasks, return_exceptions=True),
-                timeout=2.0,
-            )
-        except asyncio.TimeoutError:
-            pass
-
-        for _, sub in subscribers:
-            await sub.stop()
-        backend.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,