refactor: connector api (#532)

Author: maciejmajek

docs/developer_guide/tools.md

@@ -152,14 +152,14 @@ Refer to the [BaseROS2Tool source code](../../src/rai_core/rai/tools/ros2/base.p
 Tools can be initialized with parameters such as a connector, enabling custom configurations for ROS 2 environments.
 
 ```python
-from rai.communication.ros2 import ROS2ARIConnector
+from rai.communication.ros2 import ROS2Connector
 from rai.tools.ros2 import (
     GetROS2ImageTool,
     GetROS2TopicsNamesAndTypesTool,
     PublishROS2MessageTool,
 )
 
-def initialize_tools(connector: ROS2ARIConnector):
+def initialize_tools(connector: ROS2Connector):
     """Initialize and configure ROS 2 tools.
 
     Returns:
@@ -193,7 +193,7 @@ TODO(docs): add link to the BaseAgent docs (regarding distributed setup)
 
 ```python
 from rai.agents import ReActAgent
-from rai.communication import ROS2ARIConnector, ROS2HRIConnector
+from rai.communication import ROS2Connector, ROS2HRIConnector
 from rai.tools.ros2 import ROS2Toolkit
 from rai.communication.ros2 import ROS2Context
 from rai import AgentRunner
@@ -202,10 +202,10 @@ from rai import AgentRunner
 def main() -> None:
     """Initialize and run the RAI agent with configured tools."""
     connector = ROS2HRIConnector(sources=["/from_human"], targets=["/to_human"])
-    ari_connector = ROS2ARIConnector()
+    ros2_connector = ROS2Connector()
     agent = ReActAgent(
         connectors={"hri": connector},
-        tools=initialize_tools(connector=ari_connector),
+        tools=initialize_tools(connector=ros2_connector),
     )
     runner = AgentRunner([agent])
     runner.run_and_wait_for_shutdown()
@@ -226,9 +226,9 @@ from rai.communication.ros2 import ROS2Context
 
 @ROS2Context()
 def main():
-    ari_connector = ROS2ARIConnector()
+    ros2_connector = ROS2Connector()
     agent = create_react_runnable(
-        tools=initialize_tools(connector=ari_connector),
+        tools=initialize_tools(connector=ros2_connector),
         system_prompt="You are a helpful assistant that can answer questions and help with tasks.",
     )
     state = {'messages': []}
@@ -244,5 +244,5 @@ def main():
 ## Related Topics
 
 - [Connectors](../communication/connectors.md)
-- [ROS2ARIConnector](../communication/ros2.md)
+- [ROS2Connector](../communication/ros2.md)
 - [ROS2HRIConnector](../communication/ros2.md)

examples/agents/react.py

@@ -13,17 +13,17 @@
 # limitations under the License.
 
 from rai.agents import AgentRunner, ReActAgent
-from rai.communication.ros2 import ROS2ARIConnector, ROS2Context, ROS2HRIConnector
+from rai.communication.ros2 import ROS2Connector, ROS2Context, ROS2HRIConnector
 from rai.tools.ros2 import ROS2Toolkit
 
 
 @ROS2Context()
 def main():
     connector = ROS2HRIConnector(sources=["/from_human"], targets=["/to_human"])
-    ari_connector = ROS2ARIConnector()
+    ros2_connector = ROS2Connector()
     agent = ReActAgent(
         connectors={"hri": connector},
-        tools=ROS2Toolkit(connector=ari_connector).get_tools(),
+        tools=ROS2Toolkit(connector=ros2_connector).get_tools(),
     )  # type: ignore
     runner = AgentRunner([agent])
     runner.run_and_wait_for_shutdown()

examples/agents/streamlit_react.py

@@ -18,16 +18,16 @@
 from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
 from rai.agents.integrations.streamlit import get_streamlit_cb, streamlit_invoke
 from rai.agents.react_agent import ReActAgent
-from rai.communication.ros2 import ROS2ARIConnector
+from rai.communication.ros2 import ROS2Connector
 from rai.messages import HumanMultimodalMessage
 from rai.tools.ros2 import ROS2Toolkit
 
 
 @st.cache_resource
 def initialize_graph():
     rclpy.init()
-    ari_connector = ROS2ARIConnector()
-    tools = ROS2Toolkit(connector=ari_connector).get_tools()
+    ros2_connector = ROS2Connector()
+    tools = ROS2Toolkit(connector=ros2_connector).get_tools()
     agent = ReActAgent(connectors={}, tools=tools).agent
     return agent
 

examples/agriculture-demo.py

@@ -19,7 +19,7 @@
 from langchain_core.runnables import Runnable
 from rai import get_llm_model
 from rai.agents.conversational_agent import State, create_conversational_agent
-from rai.communication.ros2.connectors import ROS2ARIConnector
+from rai.communication.ros2.connectors import ROS2Connector
 from rai.tools.ros2 import ROS2ServicesToolkit, ROS2TopicsToolkit
 from rai.tools.time import WaitForSecondsTool
 from rclpy.callback_groups import ReentrantCallbackGroup
@@ -103,7 +103,7 @@ def main():
 
     Important: You must call only one service. The tractor can only handle one service call.
     """
-    connector = ROS2ARIConnector()
+    connector = ROS2Connector()
     agent = create_conversational_agent(
         llm=get_llm_model("complex_model"),
         system_prompt=SYSTEM_PROMPT,

examples/manipulation-demo.py

@@ -18,15 +18,15 @@
 from langchain_core.messages import HumanMessage
 from rai import get_llm_model
 from rai.agents.conversational_agent import create_conversational_agent
-from rai.communication.ros2.connectors import ROS2ARIConnector
+from rai.communication.ros2.connectors import ROS2Connector
 from rai.tools.ros2 import GetROS2ImageTool, GetROS2TopicsNamesAndTypesTool
 from rai.tools.ros2.manipulation import GetObjectPositionsTool, MoveToPointTool
 from rai_open_set_vision.tools import GetGrabbingPointTool
 
 
 def create_agent():
     rclpy.init()
-    connector = ROS2ARIConnector()
+    connector = ROS2Connector()
     node = connector.node
     node.declare_parameter("conversion_ratio", 1.0)
 

examples/rosbot-xl-demo.py

@@ -19,7 +19,7 @@
 from langchain_core.tools import BaseTool
 from rai import get_llm_model
 from rai.agents import ReActAgent
-from rai.communication.ros2 import ROS2ARIConnector
+from rai.communication.ros2 import ROS2Connector
 from rai.frontend.streamlit import run_streamlit_app
 from rai.tools.ros2 import (
     GetObjectPositionsTool,
@@ -56,7 +56,7 @@ def initialize_agent():
     Your job is to transform user intent into meaningful, goal-driven behavior within the physical world.
     """
 
-    connector = ROS2ARIConnector()
+    connector = ROS2Connector()
     tools: List[BaseTool] = [
         GetROS2TransformConfiguredTool(
             connector=connector,

src/rai_asr/rai_asr/agents/asr_agent.py

@@ -23,10 +23,10 @@
 from numpy.typing import NDArray
 from rai.agents.base import BaseAgent
 from rai.communication.ros2 import (
-    ROS2ARIConnector,
-    ROS2ARIMessage,
+    ROS2Connector,
     ROS2HRIConnector,
     ROS2HRIMessage,
+    ROS2Message,
 )
 from rai.communication.sound_device import (
     SoundDeviceConfig,
@@ -81,11 +81,11 @@ def __init__(
             targets=[], sources=[("microphone", microphone_config)]
         )
         ros2_hri_connector = ROS2HRIConnector(ros2_name, targets=["/from_human"])
-        ros2_ari_connector = ROS2ARIConnector(ros2_name + "ari")
+        ros2_connector = ROS2Connector(ros2_name + "ari")
         self.connectors = {
             "microphone": microphone,
             "ros2_hri": ros2_hri_connector,
-            "ros2_ari": ros2_ari_connector,
+            "ros2": ros2_connector,
         }
         super().__init__()
         self.should_record_pipeline: List[BaseVoiceDetectionModel] = []
@@ -266,9 +266,9 @@ def _transcription_thread(self, identifier: str):
     def _send_ros2_message(self, data: str, topic: str):
         self.logger.debug(f"Sending message to {topic}: {data}")
         if topic == "/voice_commands":
-            msg = ROS2ARIMessage({"data": data})
+            msg = ROS2Message({"data": data})
             try:
-                self.connectors["ros2_ari"].send_message(
+                self.connectors["ros2"].send_message(
                     msg, topic, msg_type="std_msgs/msg/String"
                 )
             except Exception as e:

src/rai_bench/rai_bench/examples/o3de_test_benchmark.py

@@ -22,7 +22,7 @@
 import rclpy
 from langchain.tools import BaseTool
 from rai.agents.conversational_agent import create_conversational_agent
-from rai.communication.ros2.connectors import ROS2ARIConnector
+from rai.communication.ros2.connectors import ROS2Connector
 from rai.initialization import get_llm_model
 from rai.tools.ros2 import (
     GetObjectPositionsTool,
@@ -46,7 +46,7 @@
 
 if __name__ == "__main__":
     rclpy.init()
-    connector = ROS2ARIConnector()
+    connector = ROS2Connector()
     node = connector.node
     node.declare_parameter("conversion_ratio", 1.0)
 

src/rai_bench/rai_bench/tool_calling_agent_bench/mocked_tools.py

@@ -17,8 +17,8 @@
 
 import numpy as np
 import numpy.typing as npt
-from rai.communication.ros2.connectors import ROS2ARIConnector
-from rai.communication.ros2.messages import ROS2ARIMessage
+from rai.communication.ros2.connectors import ROS2Connector
+from rai.communication.ros2.messages import ROS2Message
 from rai.messages import MultimodalArtifact, preprocess_image
 from rai.tools.ros2 import (
     GetObjectPositionsTool,
@@ -31,7 +31,7 @@
 
 
 class MockGetROS2TopicsNamesAndTypesTool(GetROS2TopicsNamesAndTypesTool):
-    connector: ROS2ARIConnector = MagicMock(spec=ROS2ARIConnector)
+    connector: ROS2Connector = MagicMock(spec=ROS2Connector)
     mock_topics_names_and_types: list[str]
 
     def _run(self) -> str:
@@ -46,7 +46,7 @@ def _run(self) -> str:
 
 
 class MockGetROS2ImageTool(GetROS2ImageTool):
-    connector: ROS2ARIConnector = MagicMock(spec=ROS2ARIConnector)
+    connector: ROS2Connector = MagicMock(spec=ROS2Connector)
     expected_topics: List[str]
 
     def _run(
@@ -95,7 +95,7 @@ def generate_mock_image() -> npt.NDArray[np.uint8]:
 
 
 class MockReceiveROS2MessageTool(ReceiveROS2MessageTool):
-    connector: ROS2ARIConnector = MagicMock(spec=ROS2ARIConnector)
+    connector: ROS2Connector = MagicMock(spec=ROS2Connector)
     expected_topics: List[str]
 
     def _run(self, topic: str) -> str:
@@ -120,14 +120,14 @@ def _run(self, topic: str) -> str:
             raise ValueError(
                 f"Topic {topic} is not available within 1.0 seconds. Check if the topic exists."
             )
-        message: ROS2ARIMessage = MagicMock(spec=ROS2ARIMessage)
+        message: ROS2Message = MagicMock(spec=ROS2Message)
         message.payload = {"mock": "payload"}
         message.metadata = {"mock": "metadata"}
         return str({"payload": message.payload, "metadata": message.metadata})
 
 
 class MockMoveToPointTool(MoveToPointTool):
-    connector: ROS2ARIConnector = MagicMock(spec=ROS2ARIConnector)
+    connector: ROS2Connector = MagicMock(spec=ROS2Connector)
 
     def _run(self, x: float, y: float, z: float, task: str) -> str:
         """Method that return a mock message with the end effector position.
@@ -149,7 +149,7 @@ def _run(self, x: float, y: float, z: float, task: str) -> str:
 
 
 class MockGetObjectPositionsTool(GetObjectPositionsTool):
-    connector: ROS2ARIConnector = MagicMock(spec=ROS2ARIConnector)
+    connector: ROS2Connector = MagicMock(spec=ROS2Connector)
 
     # Create mock instances for the arguments
     target_frame: str = MagicMock(spec=str)

src/rai_core/rai/communication/__init__.py

@@ -12,13 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .ari_connector import ARIConnector, ARIMessage
 from .base_connector import BaseConnector, BaseMessage
 from .hri_connector import HRIConnector, HRIMessage
 
 __all__ = [
-    "ARIConnector",
-    "ARIMessage",
     "BaseConnector",
     "BaseMessage",
     "HRIConnector",