mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-24 12:41:02 +02:00
317 lines
10 KiB
Python
317 lines
10 KiB
Python
"""
|
|
Multi-queue Pulsar message dumper for debugging TrustGraph message flows.
|
|
|
|
This utility monitors multiple Pulsar queues simultaneously and logs all messages
|
|
to a file with timestamps and pretty-printed formatting. Useful for debugging
|
|
message flows, diagnosing stuck services, and understanding system behavior.
|
|
|
|
Uses TrustGraph's Subscriber abstraction for future-proof pub/sub compatibility.
|
|
"""
|
|
|
|
import pulsar
|
|
from pulsar.schema import BytesSchema
|
|
import sys
|
|
import json
|
|
import asyncio
|
|
from datetime import datetime
|
|
import argparse
|
|
|
|
from trustgraph.base.subscriber import Subscriber
|
|
|
|
def format_message(queue_name, msg):
|
|
"""Format a message with timestamp and queue name."""
|
|
timestamp = datetime.now().isoformat()
|
|
|
|
# Try to parse as JSON and pretty-print
|
|
try:
|
|
value = msg.value()
|
|
|
|
# If it's bytes, decode it
|
|
if isinstance(value, bytes):
|
|
value = value.decode('utf-8')
|
|
|
|
# If it's a string, try to parse as JSON
|
|
if isinstance(value, str):
|
|
try:
|
|
parsed = json.loads(value)
|
|
body = json.dumps(parsed, indent=2)
|
|
except (json.JSONDecodeError, TypeError):
|
|
body = value
|
|
else:
|
|
# Try to convert to dict for pretty printing
|
|
try:
|
|
# Pulsar schema objects have __dict__ or similar
|
|
if hasattr(value, '__dict__'):
|
|
parsed = {k: v for k, v in value.__dict__.items()
|
|
if not k.startswith('_')}
|
|
else:
|
|
parsed = str(value)
|
|
body = json.dumps(parsed, indent=2, default=str)
|
|
except (TypeError, AttributeError):
|
|
body = str(value)
|
|
|
|
except Exception as e:
|
|
body = f"<Error formatting message: {e}>\n{msg.value()}"
|
|
|
|
# Format the output
|
|
header = f"\n{'='*80}\n[{timestamp}] Queue: {queue_name}\n{'='*80}\n"
|
|
return header + body + "\n"
|
|
|
|
|
|
async def monitor_queue(subscriber, queue_name, central_queue, monitor_id):
|
|
"""
|
|
Monitor a single queue via Subscriber and forward messages to central queue.
|
|
|
|
Args:
|
|
subscriber: Subscriber instance for this queue
|
|
queue_name: Name of the queue (for logging)
|
|
central_queue: asyncio.Queue to forward messages to
|
|
monitor_id: Unique ID for this monitor's subscription
|
|
"""
|
|
try:
|
|
# Subscribe to all messages from this Subscriber
|
|
msg_queue = await subscriber.subscribe_all(monitor_id)
|
|
|
|
while True:
|
|
# Read from Subscriber's internal queue
|
|
msg = await msg_queue.get()
|
|
timestamp = datetime.now()
|
|
formatted = format_message(queue_name, msg)
|
|
|
|
# Forward to central queue for writing
|
|
await central_queue.put((timestamp, queue_name, formatted))
|
|
|
|
except asyncio.CancelledError:
|
|
# Task cancelled during shutdown
|
|
await subscriber.unsubscribe_all(monitor_id)
|
|
raise
|
|
except Exception as e:
|
|
error_msg = f"\n{'='*80}\n[{datetime.now().isoformat()}] ERROR in monitor for {queue_name}\n{'='*80}\n{e}\n"
|
|
await central_queue.put((datetime.now(), queue_name, error_msg))
|
|
|
|
|
|
async def log_writer(central_queue, file_handle, console_output=True):
|
|
"""
|
|
Write messages from central queue to file.
|
|
|
|
Args:
|
|
central_queue: asyncio.Queue containing (timestamp, queue_name, formatted_msg) tuples
|
|
file_handle: Open file handle to write to
|
|
console_output: Whether to print abbreviated messages to console
|
|
"""
|
|
try:
|
|
while True:
|
|
timestamp, queue_name, formatted_msg = await central_queue.get()
|
|
|
|
# Write to file
|
|
file_handle.write(formatted_msg)
|
|
file_handle.flush()
|
|
|
|
# Print abbreviated message to console
|
|
if console_output:
|
|
time_str = timestamp.strftime('%H:%M:%S')
|
|
print(f"[{time_str}] {queue_name}: Message received")
|
|
|
|
except asyncio.CancelledError:
|
|
# Flush remaining messages before shutdown
|
|
while not central_queue.empty():
|
|
try:
|
|
timestamp, queue_name, formatted_msg = central_queue.get_nowait()
|
|
file_handle.write(formatted_msg)
|
|
file_handle.flush()
|
|
except asyncio.QueueEmpty:
|
|
break
|
|
raise
|
|
|
|
|
|
async def async_main(queues, output_file, pulsar_host, listener_name, subscriber_name, append_mode):
|
|
"""
|
|
Main async function to monitor multiple queues concurrently.
|
|
|
|
Args:
|
|
queues: List of queue names to monitor
|
|
output_file: Path to output file
|
|
pulsar_host: Pulsar connection URL
|
|
listener_name: Pulsar listener name
|
|
subscriber_name: Base name for subscribers
|
|
append_mode: Whether to append to existing file
|
|
"""
|
|
print(f"TrustGraph Queue Dumper")
|
|
print(f"Monitoring {len(queues)} queue(s):")
|
|
for q in queues:
|
|
print(f" - {q}")
|
|
print(f"Output file: {output_file}")
|
|
print(f"Mode: {'append' if append_mode else 'overwrite'}")
|
|
print(f"Press Ctrl+C to stop\n")
|
|
|
|
# Connect to Pulsar
|
|
try:
|
|
client = pulsar.Client(pulsar_host, listener_name=listener_name)
|
|
except Exception as e:
|
|
print(f"Error connecting to Pulsar at {pulsar_host}: {e}", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
# Create Subscribers and central queue
|
|
central_queue = asyncio.Queue()
|
|
subscribers = []
|
|
|
|
for queue_name in queues:
|
|
try:
|
|
sub = Subscriber(
|
|
client=client,
|
|
topic=queue_name,
|
|
subscription=subscriber_name,
|
|
consumer_name=f"{subscriber_name}-{queue_name}",
|
|
schema=None, # No schema - accept any message type
|
|
)
|
|
await sub.start()
|
|
subscribers.append((queue_name, sub))
|
|
print(f"✓ Subscribed to: {queue_name}")
|
|
except Exception as e:
|
|
print(f"✗ Error subscribing to {queue_name}: {e}", file=sys.stderr)
|
|
|
|
if not subscribers:
|
|
print("\nNo subscribers created. Exiting.", file=sys.stderr)
|
|
client.close()
|
|
sys.exit(1)
|
|
|
|
print(f"\nListening for messages...\n")
|
|
|
|
# Open output file
|
|
mode = 'a' if append_mode else 'w'
|
|
try:
|
|
with open(output_file, mode) as f:
|
|
f.write(f"\n{'#'*80}\n")
|
|
f.write(f"# Session started: {datetime.now().isoformat()}\n")
|
|
f.write(f"# Monitoring queues: {', '.join(queues)}\n")
|
|
f.write(f"{'#'*80}\n")
|
|
f.flush()
|
|
|
|
# Start monitoring tasks
|
|
try:
|
|
try:
|
|
async with asyncio.TaskGroup() as tg:
|
|
# Create one monitor task per subscriber
|
|
for queue_name, sub in subscribers:
|
|
tg.create_task(monitor_queue(sub, queue_name, central_queue, "logger"))
|
|
|
|
# Create single writer task
|
|
tg.create_task(log_writer(central_queue, f))
|
|
|
|
except* Exception as eg:
|
|
# TaskGroup exception group
|
|
for exc in eg.exceptions:
|
|
if not isinstance(exc, asyncio.CancelledError):
|
|
print(f"Task error: {exc}", file=sys.stderr)
|
|
|
|
except KeyboardInterrupt:
|
|
print("\n\nStopping...")
|
|
finally:
|
|
# Write session end marker
|
|
f.write(f"\n{'#'*80}\n")
|
|
f.write(f"# Session ended: {datetime.now().isoformat()}\n")
|
|
f.write(f"{'#'*80}\n")
|
|
|
|
except IOError as e:
|
|
print(f"Error writing to {output_file}: {e}", file=sys.stderr)
|
|
sys.exit(1)
|
|
finally:
|
|
# Cleanup subscribers
|
|
for _, sub in subscribers:
|
|
await sub.stop()
|
|
client.close()
|
|
|
|
print(f"\nMessages logged to: {output_file}")
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
prog='tg-dump-queues',
|
|
description='Monitor and dump messages from multiple Pulsar queues',
|
|
epilog="""
|
|
Examples:
|
|
# Monitor agent and prompt queues
|
|
tg-dump-queues non-persistent://tg/request/agent:default \\
|
|
non-persistent://tg/request/prompt:default
|
|
|
|
# Monitor with custom output file
|
|
tg-dump-queues non-persistent://tg/request/agent:default \\
|
|
--output debug.log
|
|
|
|
# Append to existing log file
|
|
tg-dump-queues non-persistent://tg/request/agent:default \\
|
|
--output queue.log --append
|
|
|
|
Common queue patterns:
|
|
- Agent requests: non-persistent://tg/request/agent:default
|
|
- Agent responses: non-persistent://tg/response/agent:default
|
|
- Prompt requests: non-persistent://tg/request/prompt:default
|
|
- Prompt responses: non-persistent://tg/response/prompt:default
|
|
- LLM requests: non-persistent://tg/request/text-completion:default
|
|
- LLM responses: non-persistent://tg/response/text-completion:default
|
|
""",
|
|
formatter_class=argparse.RawDescriptionHelpFormatter
|
|
)
|
|
|
|
parser.add_argument(
|
|
'queues',
|
|
nargs='+',
|
|
help='Pulsar queue names to monitor'
|
|
)
|
|
|
|
parser.add_argument(
|
|
'--output', '-o',
|
|
default='queue.log',
|
|
help='Output file (default: queue.log)'
|
|
)
|
|
|
|
parser.add_argument(
|
|
'--append', '-a',
|
|
action='store_true',
|
|
help='Append to output file instead of overwriting'
|
|
)
|
|
|
|
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)'
|
|
)
|
|
|
|
parser.add_argument(
|
|
'--subscriber',
|
|
default='debug',
|
|
help='Subscriber name for queue subscription (default: debug)'
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
# Filter out any accidentally included flags
|
|
queues = [q for q in args.queues if not q.startswith('--')]
|
|
|
|
if not queues:
|
|
parser.error("No queues specified")
|
|
|
|
# Run async main
|
|
try:
|
|
asyncio.run(async_main(
|
|
queues=queues,
|
|
output_file=args.output,
|
|
pulsar_host=args.pulsar_host,
|
|
listener_name=args.listener_name,
|
|
subscriber_name=args.subscriber,
|
|
append_mode=args.append
|
|
))
|
|
except KeyboardInterrupt:
|
|
# Already handled in async_main
|
|
pass
|
|
except Exception as e:
|
|
print(f"Fatal error: {e}", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
if __name__ == '__main__':
|
|
main()
|