liqiang
b119135836
Build latest book artifacts / build (push) Canceled after 0s
dependency resolution / resolve (3.11) (push) Canceled after 0s
dependency resolution / resolve (3.13) (push) Canceled after 0s
deploy-pages / build (push) Canceled after 0s
deploy-pages / deploy (push) Canceled after 0s
i18n consistency check / check (push) Canceled after 0s
provider adoption tests / test (chapter2/context-compression) (push) Canceled after 0s
provider adoption tests / test (chapter2/prompt-injection) (push) Canceled after 0s
provider adoption tests / test (chapter2/system-hint) (push) Canceled after 0s
provider adoption tests / test (chapter3/log-sanitization) (push) Canceled after 0s
web-search-agent tests / test (push) Canceled after 0s
web-search-agent tests / agentbook (push) Canceled after 0s
76 lines
2.6 KiB
Python
76 lines
2.6 KiB
Python
"""Asynchronous, timestamped point-to-point bus for the two live Agents."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import time
|
|
from collections import defaultdict
|
|
from pathlib import Path
|
|
from typing import DefaultDict, List, Optional
|
|
|
|
from models import AgentMessage
|
|
|
|
|
|
class MessageBus:
|
|
def __init__(self, trace_path: Optional[str] = None):
|
|
self.started = time.monotonic()
|
|
self._sequence = 0
|
|
self._queues: DefaultDict[str, asyncio.Queue[AgentMessage]] = defaultdict(asyncio.Queue)
|
|
self.history: List[AgentMessage] = []
|
|
self.trace_path = Path(trace_path) if trace_path else None
|
|
|
|
async def send(
|
|
self,
|
|
sender: str,
|
|
recipient: str,
|
|
type: str,
|
|
*,
|
|
sensitive_keys: tuple[str, ...] = (),
|
|
**payload,
|
|
) -> AgentMessage:
|
|
self._sequence += 1
|
|
message = AgentMessage(
|
|
sender=sender,
|
|
recipient=recipient,
|
|
type=type,
|
|
payload=payload,
|
|
sequence=self._sequence,
|
|
monotonic_seconds=round(time.monotonic() - self.started, 6),
|
|
wall_time=time.strftime("%Y-%m-%dT%H:%M:%S%z"),
|
|
)
|
|
self.history.append(message)
|
|
await self._queues[recipient].put(message)
|
|
printable = {k: ("<redacted>" if k in sensitive_keys else v) for k, v in payload.items()}
|
|
# Keep the redaction policy beside the in-memory envelope. The receiver gets
|
|
# the value, while console/disk traces never retain spoken personal data.
|
|
setattr(message, "_sensitive_keys", sensitive_keys)
|
|
print(
|
|
f"[t={message.monotonic_seconds:8.3f}s #{message.sequence:03d}] "
|
|
f"{sender} -> {recipient} | {type} | "
|
|
f"{json.dumps(printable, ensure_ascii=False)}"
|
|
)
|
|
self.flush()
|
|
return message
|
|
|
|
async def receive(self, recipient: str, timeout: Optional[float] = None) -> AgentMessage:
|
|
get = self._queues[recipient].get()
|
|
return await asyncio.wait_for(get, timeout) if timeout else await get
|
|
|
|
def flush(self) -> None:
|
|
if not self.trace_path:
|
|
return
|
|
self.trace_path.parent.mkdir(parents=True, exist_ok=True)
|
|
rows = []
|
|
for message in self.history:
|
|
row = message.to_dict()
|
|
keys = getattr(message, "_sensitive_keys", ())
|
|
row["payload"] = {
|
|
k: ("<redacted>" if k in keys else v) for k, v in row["payload"].items()
|
|
}
|
|
rows.append(row)
|
|
self.trace_path.write_text(
|
|
json.dumps(rows, ensure_ascii=False, indent=2),
|
|
encoding="utf-8",
|
|
)
|