multi_agent_lab/bus.py

from __future__ import annotations

import json
import time
from dataclasses import dataclass
from typing import Any, Dict, Iterable, List, Optional


@dataclass(frozen=True)
class Message:
    ts: float
    sender: str
    receiver: str
    kind: str
    payload: Dict[str, Any]

    def to_dict(self) -> Dict[str, Any]:
        return {
            "ts": self.ts,
            "sender": self.sender,
            "receiver": self.receiver,
            "kind": self.kind,
            "payload": self.payload,
        }


class MessageBus:
    def __init__(self) -> None:
        self._messages: List[Message] = []

    def send(
        self,
        *,
        sender: str,
        receiver: str,
        kind: str,
        payload: Optional[Dict[str, Any]] = None,
    ) -> None:
        self._messages.append(
            Message(
                ts=time.time(),
                sender=sender,
                receiver=receiver,
                kind=kind,
                payload=payload or {},
            )
        )

    def iter(self) -> Iterable[Message]:
        return iter(self._messages)

    def to_jsonl(self) -> str:
        return "\n".join(json.dumps(m.to_dict(), ensure_ascii=False) for m in self._messages) + "\n"