build: merge

merge with riAgent

ref: N25B-208

src/control_backend/agents/bdi/bdi_core.py

@@ -13,7 +13,7 @@ from control_backend.core.config import settings
 
 class BDICoreAgent(BDIAgent):
     """
-    This is the Brain agent that does the belief inference with AgentSpeak. 
+    This is the Brain agent that does the belief inference with AgentSpeak.
     This is a continous process that happens automatically in the background.
     This class contains all the actions that can be called from AgentSpeak plans.
     It has the BeliefSetter behaviour and can aks and recieve requests from the LLM agent.

src/control_backend/agents/bdi/behaviours/belief_setter.py

@@ -8,15 +8,17 @@ from spade_bdi.bdi import BDIAgent
 
 from control_backend.core.config import settings
 
+
 class BeliefSetter(CyclicBehaviour):
     """
-    This is the behaviour that the BDI agent runs.
-    This behaviour waits for incoming message and processes it based on sender.
-    Currently, t only waits for messages containing beliefs from Belief Collector and adds these to its KB.
+    This is the behaviour that the BDI agent runs. This behaviour waits for incoming
+    message and processes it based on sender. Currently, it only waits for messages
+    containing beliefs from BeliefCollector and adds these to its KB.
     """
+
     agent: BDIAgent
     logger = logging.getLogger("BDI/Belief Setter")
-    
+
     async def run(self):
         msg = await self.receive(timeout=0.1)
         if msg:
@@ -37,7 +39,8 @@ class BeliefSetter(CyclicBehaviour):
                 pass
 
     def _process_belief_message(self, message: Message):
-        if not message.body: return
+        if not message.body:
+            return
 
         match message.thread:
             case "beliefs":
@@ -49,7 +52,6 @@ class BeliefSetter(CyclicBehaviour):
             case _:
                 pass
 
-
     def _set_beliefs(self, beliefs: dict[str, list[list[str]]]):
         if self.agent.bdi is None:
             self.logger.warning("Cannot set beliefs, since agent's BDI is not yet initialized.")

src/control_backend/agents/ri_command_agent.py

@@ -0,0 +1,74 @@
+import json
+import logging
+from spade.agent import Agent
+from spade.behaviour import CyclicBehaviour
+import zmq
+
+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__)
+
+
+class RICommandAgent(Agent):
+    subsocket: zmq.Socket
+    pubsocket: zmq.Socket
+    address = ""
+    bind = False
+
+    def __init__(
+        self,
+        jid: str,
+        password: str,
+        port: int = 5222,
+        verify_security: bool = False,
+        address="tcp://localhost:0000",
+        bind=False,
+    ):
+        super().__init__(jid, password, port, verify_security)
+        self.address = address
+        self.bind = bind
+
+    class SendCommandsBehaviour(CyclicBehaviour):
+        async def run(self):
+            """
+            Run the command publishing loop indefinetely.
+            """
+            assert self.agent is not None
+            # Get a message internally (with topic command)
+            topic, body = await self.agent.subsocket.recv_multipart()
+
+            # Try to get body
+            try:
+                body = json.loads(body)
+                message = SpeechCommand.model_validate(body)
+
+                # Send to the robot.
+                await self.agent.pubsocket.send_json(message.model_dump())
+            except Exception as e:
+                logger.error("Error processing message: %s", e)
+
+    async def setup(self):
+        """
+        Setup the command agent
+        """
+        logger.info("Setting up %s", self.jid)
+
+        # To the robot
+        self.pubsocket = context.socket(zmq.PUB)
+        if self.bind:
+            self.pubsocket.bind(self.address)
+        else:
+            self.pubsocket.connect(self.address)
+
+        # Receive internal topics regarding commands
+        self.subsocket = context.socket(zmq.SUB)
+        self.subsocket.connect(settings.zmq_settings.internal_comm_address)
+        self.subsocket.setsockopt(zmq.SUBSCRIBE, b"command")
+
+        # Add behaviour to our agent
+        commands_behaviour = self.SendCommandsBehaviour()
+        self.add_behaviour(commands_behaviour)
+
+        logger.info("Finished setting up %s", self.jid)

src/control_backend/agents/ri_communication_agent.py

@@ -0,0 +1,165 @@
+import asyncio
+import json
+import logging
+from spade.agent import Agent
+from spade.behaviour import CyclicBehaviour
+import zmq
+
+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
+
+logger = logging.getLogger(__name__)
+
+
+class RICommunicationAgent(Agent):
+    req_socket: zmq.Socket
+    _address = ""
+    _bind = True
+
+    def __init__(
+        self,
+        jid: str,
+        password: str,
+        port: int = 5222,
+        verify_security: bool = False,
+        address="tcp://localhost:0000",
+        bind=False,
+    ):
+        super().__init__(jid, password, port, verify_security)
+        self._address = address
+        self._bind = bind
+
+    class ListenBehaviour(CyclicBehaviour):
+        async def run(self):
+            """
+            Run the listening (ping) loop indefinetely.
+            """
+            assert self.agent is not None
+
+            # We need to listen and sent pings.
+            message = {"endpoint": "ping", "data": {"id": "e.g. some reference id"}}
+            await self.agent.req_socket.send_json(message)
+
+            # Wait up to three seconds for a reply:)
+            try:
+                message = await asyncio.wait_for(self.agent.req_socket.recv_json(), timeout=3.0)
+
+            # We didnt get a reply :(
+            except asyncio.TimeoutError as e:
+                logger.info("No ping retrieved in 3 seconds, killing myself.")
+                self.kill()
+
+            logger.debug('Received message "%s"', message)
+            if "endpoint" not in message:
+                logger.error("No received endpoint in message, excepted ping endpoint.")
+                return
+
+            # See what endpoint we received
+            match message["endpoint"]:
+                case "ping":
+                    await asyncio.sleep(1)
+                case _:
+                    logger.info(
+                        "Received message with topic different than ping, while ping expected."
+                    )
+
+    async def setup(self, max_retries: int = 5):
+        """
+        Try to setup the communication agent, we have 5 retries in case we dont have a response yet.
+        """
+        logger.info("Setting up %s", self.jid)
+        retries = 0
+
+        # 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)
+            if self._bind:
+                self.req_socket.bind(self._address)
+            else:
+                self.req_socket.connect(self._address)
+
+            # Send our message and receive one back:)
+            message = {"endpoint": "negotiate/ports", "data": None}
+            await self.req_socket.send_json(message)
+
+            try:
+                received_message = await asyncio.wait_for(self.req_socket.recv_json(), timeout=20.0)
+
+            except asyncio.TimeoutError:
+                logger.warning(
+                    "No connection established in 20 seconds (attempt %d/%d)",
+                    retries + 1,
+                    max_retries,
+                )
+                retries += 1
+                continue
+
+            except Exception as e:
+                logger.error("Unexpected error during negotiation: %s", e)
+                retries += 1
+                continue
+
+            # Validate endpoint
+            endpoint = received_message.get("endpoint")
+            if endpoint != "negotiate/ports":
+                # TODO: Should this send a message back?
+                logger.error(
+                    "Invalid endpoint '%s' received (attempt %d/%d)",
+                    endpoint,
+                    retries + 1,
+                    max_retries,
+                )
+                retries += 1
+                continue
+
+            # At this point, we have a valid response
+            try:
+                for port_data in received_message["data"]:
+                    id = port_data["id"]
+                    port = port_data["port"]
+                    bind = port_data["bind"]
+
+                    if not bind:
+                        addr = f"tcp://localhost:{port}"
+                    else:
+                        addr = f"tcp://*:{port}"
+
+                    match id:
+                        case "main":
+                            if addr != self._address:
+                                if not bind:
+                                    self.req_socket.connect(addr)
+                                else:
+                                    self.req_socket.bind(addr)
+                        case "actuation":
+                            ri_commands_agent = RICommandAgent(
+                                settings.agent_settings.ri_command_agent_name
+                                + "@"
+                                + settings.agent_settings.host,
+                                settings.agent_settings.ri_command_agent_name,
+                                address=addr,
+                                bind=bind,
+                            )
+                            await ri_commands_agent.start()
+                        case _:
+                            logger.warning("Unhandled negotiation id: %s", id)
+
+            except Exception as e:
+                logger.error("Error unpacking negotiation data: %s", e)
+                retries += 1
+                continue
+
+            # setup succeeded
+            break
+
+        else:
+            logger.error("Failed to set up RICommunicationAgent after %d retries", max_retries)
+            return
+
+        # Set up ping behaviour
+        listen_behaviour = self.ListenBehaviour()
+        self.add_behaviour(listen_behaviour)
+        logger.info("Finished setting up %s", self.jid)