refactor: ZMQ context and proxy

Use ZMQ's global context instance and setup an XPUB/XSUB proxy intermediary to allow for easier multi-pubs.

close: N25B-217

src/control_backend/agents/ri_command_agent.py

@@ -1,11 +1,12 @@
 import json
 import logging
+
+import zmq
 from spade.agent import Agent
 from spade.behaviour import CyclicBehaviour
-import zmq
+from zmq.asyncio import Context
 
 from control_backend.core.config import settings
-from control_backend.core.zmq_context import context
 from control_backend.schemas.ri_message import SpeechCommand
 
 logger = logging.getLogger(__name__)
@@ -55,6 +56,8 @@ class RICommandAgent(Agent):
         """
         logger.info("Setting up %s", self.jid)
 
+        context = Context.instance()
+
         # To the robot
         self.pubsocket = context.socket(zmq.PUB)
         if self.bind:
@@ -64,7 +67,7 @@ class RICommandAgent(Agent):
 
         # Receive internal topics regarding commands
         self.subsocket = context.socket(zmq.SUB)
-        self.subsocket.connect(settings.zmq_settings.internal_comm_address)
+        self.subsocket.connect(settings.zmq_settings.internal_sub_address)
         self.subsocket.setsockopt(zmq.SUBSCRIBE, b"command")
 
         # Add behaviour to our agent

src/control_backend/agents/ri_communication_agent.py

@@ -1,14 +1,13 @@
 import asyncio
-import json
 import logging
+
+import zmq
 from spade.agent import Agent
 from spade.behaviour import CyclicBehaviour
-import zmq
+from zmq.asyncio import Context
 
-from control_backend.core.config import settings
-from control_backend.core.zmq_context import context
-from control_backend.schemas.message import Message
 from control_backend.agents.ri_command_agent import RICommandAgent
+from control_backend.core.config import settings
 
 logger = logging.getLogger(__name__)
 
@@ -47,7 +46,7 @@ class RICommunicationAgent(Agent):
                 message = await asyncio.wait_for(self.agent.req_socket.recv_json(), timeout=3.0)
 
             # We didnt get a reply :(
-            except asyncio.TimeoutError as e:
+            except TimeoutError:
                 logger.info("No ping retrieved in 3 seconds, killing myself.")
                 self.kill()
 
@@ -75,7 +74,7 @@ class RICommunicationAgent(Agent):
         # Let's try a certain amount of times before failing connection
         while retries < max_retries:
             # Bind request socket
-            self.req_socket = context.socket(zmq.REQ)
+            self.req_socket = Context.instance().socket(zmq.REQ)
             if self._bind:
                 self.req_socket.bind(self._address)
             else:
@@ -88,7 +87,7 @@ class RICommunicationAgent(Agent):
             try:
                 received_message = await asyncio.wait_for(self.req_socket.recv_json(), timeout=20.0)
 
-            except asyncio.TimeoutError:
+            except TimeoutError:
                 logger.warning(
                     "No connection established in 20 seconds (attempt %d/%d)",
                     retries + 1,

src/control_backend/agents/transcription/transcription_agent.py

@@ -10,7 +10,6 @@ from spade.message import Message
 
 from control_backend.agents.transcription.speech_recognizer import SpeechRecognizer
 from control_backend.core.config import settings
-from control_backend.core.zmq_context import context as zmq_context
 
 logger = logging.getLogger(__name__)
 
@@ -47,7 +46,8 @@ class TranscriptionAgent(Agent):
             """Share a transcription to the other agents that depend on it."""
             receiver_jids = [
                 settings.agent_settings.text_belief_extractor_agent_name
-                + '@' + settings.agent_settings.host,
+                + "@"
+                + settings.agent_settings.host,
             ]  # Set message receivers here
 
             for receiver_jid in receiver_jids:
@@ -68,7 +68,7 @@ class TranscriptionAgent(Agent):
         return await super().stop()
 
     def _connect_audio_in_socket(self):
-        self.audio_in_socket = zmq_context.socket(zmq.SUB)
+        self.audio_in_socket = azmq.Context.instance().socket(zmq.SUB)
         self.audio_in_socket.setsockopt_string(zmq.SUBSCRIBE, "")
         self.audio_in_socket.connect(self.audio_in_address)
 

src/control_backend/agents/vad_agent.py

@@ -9,7 +9,6 @@ from spade.behaviour import CyclicBehaviour
 
 from control_backend.agents.transcription import TranscriptionAgent
 from control_backend.core.config import settings
-from control_backend.core.zmq_context import context as zmq_context
 
 logger = logging.getLogger(__name__)
 
@@ -121,7 +120,7 @@ class VADAgent(Agent):
         return await super().stop()
 
     def _connect_audio_in_socket(self):
-        self.audio_in_socket = zmq_context.socket(zmq.SUB)
+        self.audio_in_socket = azmq.Context.instance().socket(zmq.SUB)
         self.audio_in_socket.setsockopt_string(zmq.SUBSCRIBE, "")
         if self.audio_in_bind:
             self.audio_in_socket.bind(self.audio_in_address)
@@ -132,7 +131,7 @@ class VADAgent(Agent):
     def _connect_audio_out_socket(self) -> int | None:
         """Returns the port bound, or None if binding failed."""
         try:
-            self.audio_out_socket = zmq_context.socket(zmq.PUB)
+            self.audio_out_socket = azmq.Context.instance().socket(zmq.PUB)
             return self.audio_out_socket.bind_to_random_port("tcp://*", max_tries=100)
         except zmq.ZMQBindError:
             logger.error("Failed to bind an audio output socket after 100 tries.")