refactor: registering on source in agent

- rollback registering using __call__ method
- add subscribe_source method

Author: boczekbartek

examples/agents/react.py

@@ -12,6 +12,7 @@
 # See the License for the specific language goveself.rning permissions and
 # limitations under the License.
 
+
 from rai.agents import AgentRunner
 from rai.agents.langchain.react_agent import ReActAgent
 from rai.communication.ros2 import ROS2Connector, ROS2Context
@@ -23,14 +24,15 @@
 def main():
     ros2_connector = ROS2Connector()
     hri_connector = ROS2HRIConnector()
-    target_connectors = {"/to_human": hri_connector}
+
     agent = ReActAgent(
-        target_connectors=target_connectors,
+        target_connectors={
+            "/to_human": hri_connector,
+        },
         tools=ROS2Toolkit(connector=ros2_connector).get_tools(),
     )
-    hri_connector.register_callback(
-        "/from_human", agent, msg_type="rai_interfaces/msg/HRIMessage"
-    )
+    agent.subscribe_source("/from_human", hri_connector)
+
     runner = AgentRunner([agent])
     runner.run_and_wait_for_shutdown()
 

src/rai_core/rai/agents/langchain/agent.py

@@ -21,29 +21,22 @@
 
 from langchain_core.messages import BaseMessage
 from langchain_core.runnables import Runnable
-from pydantic import BaseModel
 
 from rai.agents.base import BaseAgent
 from rai.agents.langchain import HRICallbackHandler
 from rai.agents.langchain.runnables import ReActAgentState
-from rai.communication.base_connector import BaseConnector
-from rai.communication.hri_connector import HRIMessage
+from rai.communication.hri_connector import HRIConnector, HRIMessage
 from rai.initialization import get_tracing_callbacks
 
 
 class BaseState(TypedDict):
     messages: List[BaseMessage]
 
 
-class HRIConfig(BaseModel):
-    source: str
-    targets: List[str]
-
-
 class LangChainAgent(BaseAgent):
     def __init__(
         self,
-        target_connectors: Dict[str, BaseConnector],
+        target_connectors: Dict[str, HRIConnector[HRIMessage]],
         runnable: Runnable,
         state: BaseState | None = None,
         new_message_behavior: Literal[
@@ -61,7 +54,7 @@ def __init__(
         self.new_message_behavior = new_message_behavior
         self.tracing_callbacks = get_tracing_callbacks()
         self.state = state or ReActAgentState(messages=[])
-        self.callback = HRICallbackHandler(
+        self._langchain_callback = HRICallbackHandler(
             connectors=target_connectors,
             aggregate_chunks=True,
             logger=self.logger,
@@ -76,7 +69,13 @@ def __init__(
         self._interupt_event = threading.Event()
         self._agent_ready_event = threading.Event()
 
-    def __call__(self, msg: HRIMessage):
+    def subscribe_source(self, source: str, connector: HRIConnector[HRIMessage]):
+        connector.register_callback(
+            source,
+            self.source_callback,
+        )
+
+    def source_callback(self, msg: HRIMessage):
         if self.max_size is not None and len(self._received_messages) >= self.max_size:
             self.logger.warning("Buffer overflow. Dropping olders message")
             self._received_messages.popleft()
@@ -117,7 +116,9 @@ def run_agent(self):
             self.state["messages"].append(langchain_message)
             for _ in self.agent.stream(
                 self.state,
-                config={"callbacks": [self.callback, *self.tracing_callbacks]},
+                config={
+                    "callbacks": [self._langchain_callback, *self.tracing_callbacks]
+                },
             ):
                 if self._interupt_event.is_set():
                     break

src/rai_core/rai/agents/langchain/callback.py

@@ -50,7 +50,7 @@ def _send_all_targets(self, tokens: str, done: bool = False):
         for target, connector in self.connectors.items():
             self.logger.info(f"Sending {len(tokens)} tokens to targer: {target}")
             try:
-                to_send = connector.T_class.from_langchain(
+                to_send: HRIMessage = connector.build_message(
                     AIMessage(content=tokens),
                     self.current_conversation_id,
                     self.current_chunk_id,

src/rai_core/rai/agents/langchain/react_agent.py

@@ -20,13 +20,13 @@
 from rai.agents.langchain import create_react_runnable
 from rai.agents.langchain.agent import LangChainAgent
 from rai.agents.langchain.runnables import ReActAgentState
-from rai.communication.hri_connector import HRIConnector
+from rai.communication.hri_connector import HRIConnector, HRIMessage
 
 
 class ReActAgent(LangChainAgent):
     def __init__(
         self,
-        target_connectors: Dict[str, HRIConnector],
+        target_connectors: Dict[str, HRIConnector[HRIMessage]],
         llm: Optional[BaseChatModel] = None,
         tools: Optional[List[BaseTool]] = None,
         state: Optional[ReActAgentState] = None,

src/rai_core/rai/communication/hri_connector.py

@@ -143,7 +143,7 @@ class HRIConnector(Generic[T], BaseConnector[T]):
     Used for sending and receiving messages between human and robot from various sources.
     """
 
-    def _build_message(
+    def build_message(
         self,
         message: LangchainBaseMessage | RAIMultimodalMessage,
         communication_id: Optional[str] = None,

src/rai_core/rai/communication/ros2/connectors/hri_connector.py

@@ -19,7 +19,7 @@
 
 from rclpy.qos import QoSProfile
 
-from rai.communication.hri_connector import HRIConnector
+from rai.communication import HRIConnector
 from rai.communication.ros2.connectors.base import ROS2BaseConnector
 from rai.communication.ros2.messages import ROS2HRIMessage
 

src/rai_core/rai/communication/ros2/messages.py

@@ -65,7 +65,7 @@ def from_ros2(
             for audio_msg in cast(List[ROS2HRIMessage__Audio], msg.audios)
         ]
         communication_id = msg.communication_id if msg.communication_id != "" else None
-        return ROS2HRIMessage(
+        return cls(
             text=msg.text,
             images=pil_images,
             audios=audio_segments,