feat: able to send to multiple receivers #42

Merged
k.marinus merged 4 commits from feat/multiple-receivers into dev 2026-01-15 09:26:14 +00:00
2 changed files with 69 additions and 3 deletions
Showing only changes of commit 70e05b6c92 - Show all commits

View File

@@ -60,6 +60,9 @@ class BaseAgent(ABC):
self._tasks: set[asyncio.Task] = set() self._tasks: set[asyncio.Task] = set()
self._running = False self._running = False
self._internal_pub_socket: None | azmq.Socket = None
self._internal_sub_socket: None | azmq.Socket = None
# Register immediately # Register immediately
AgentDirectory.register(name, self) AgentDirectory.register(name, self)
@@ -141,7 +144,7 @@ class BaseAgent(ABC):
self.logger.debug(f"Sent message {message.body} to {message.to} via regular inbox.") self.logger.debug(f"Sent message {message.body} to {message.to} via regular inbox.")
else: else:
# Apparently target agent is on a different process, send via ZMQ # Apparently target agent is on a different process, send via ZMQ
topic = f"internal/{message.to}".encode() topic = f"internal/{receiver}".encode()
body = message.model_dump_json().encode() body = message.model_dump_json().encode()
await self._internal_pub_socket.send_multipart([topic, body]) await self._internal_pub_socket.send_multipart([topic, body])
self.logger.debug(f"Sent message {message.body} to {message.to} via ZMQ.") self.logger.debug(f"Sent message {message.body} to {message.to} via ZMQ.")

View File

@@ -99,12 +99,75 @@ async def test_send_to_local_agent(monkeypatch):
# Patch inbox.put # Patch inbox.put
target.inbox.put = AsyncMock() target.inbox.put = AsyncMock()
message = InternalMessage(to="receiver", sender="sender", body="hello") message = InternalMessage(to=target.name, sender=sender.name, body="hello")
await sender.send(message) await sender.send(message)
target.inbox.put.assert_awaited_once_with(message) target.inbox.put.assert_awaited_once_with(message)
sender.logger.debug.assert_called_once()
@pytest.mark.asyncio
async def test_send_to_zmq_agent(monkeypatch):
sender = DummyAgent("sender")
target = "remote_receiver"
# Fake logger
sender.logger = MagicMock()
# Fake zmq
sender._internal_pub_socket = AsyncMock()
message = InternalMessage(to=target, sender=sender.name, body="hello")
await sender.send(message)
zmq_calls = sender._internal_pub_socket.send_multipart.call_args[0][0]
assert zmq_calls[0] == f"internal/{target}".encode()
@pytest.mark.asyncio
async def test_send_to_multiple_local_agents(monkeypatch):
sender = DummyAgent("sender")
target1 = DummyAgent("receiver1")
target2 = DummyAgent("receiver2")
# Fake logger
sender.logger = MagicMock()
# Patch inbox.put
target1.inbox.put = AsyncMock()
target2.inbox.put = AsyncMock()
message = InternalMessage(to=[target1.name, target2.name], sender=sender.name, body="hello")
await sender.send(message)
target1.inbox.put.assert_awaited_once_with(message)
target2.inbox.put.assert_awaited_once_with(message)
@pytest.mark.asyncio
async def test_send_to_multiple_agents(monkeypatch):
sender = DummyAgent("sender")
target1 = DummyAgent("receiver1")
target2 = "remote_receiver"
# Fake logger
sender.logger = MagicMock()
# Fake zmq
sender._internal_pub_socket = AsyncMock()
# Patch inbox.put
target1.inbox.put = AsyncMock()
message = InternalMessage(to=[target1.name, target2], sender=sender.name, body="hello")
await sender.send(message)
target1.inbox.put.assert_awaited_once_with(message)
zmq_calls = sender._internal_pub_socket.send_multipart.call_args[0][0]
assert zmq_calls[0] == f"internal/{target2}".encode()
@pytest.mark.asyncio @pytest.mark.asyncio