mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-12 14:52:11 +02:00
Output working but messy
This commit is contained in:
Cyber MacGeddon
parent
596c7a8514
commit
734382c44c
1 changed files with 117 additions and 156 deletions
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue