refactor(state_based): adapt to register_callback api

Author: boczekbartek

examples/agents/state_based.py

@@ -13,12 +13,12 @@
 # limitations under the License.
 
 from rai.agents import StateBasedAgent, StateBasedConfig, wait_for_shutdown
-from rai.agents.postprocessors.ros2 import (
-    ROS2ImgVLMDiffPostprocessor,
-    ROS2LogsPostprocessor,
+from rai.aggregators.ros2 import (
+    ROS2ImgVLMDiffAggregator,
+    ROS2LogsAggregator,
 )
 from rai.communication.ros2 import (
-    ROS2ARIConnector,
+    ROS2Connector,
     ROS2Context,
     ROS2HRIConnector,
 )
@@ -29,12 +29,12 @@
 @ROS2Context()
 def main():
     connector = ROS2HRIConnector(sources=["/from_human"], targets=["/to_human"])
-    ari_connector = ROS2ARIConnector()
+    ros2_connector = ROS2Connector()
 
     config = StateBasedConfig(
-        postprocessors={
-            "/camera/camera/color/image_raw": [ROS2ImgVLMDiffPostprocessor()],
-            "/rosout": [ROS2LogsPostprocessor()],
+        aggregators={
+            "/camera/camera/color/image_raw": [ROS2ImgVLMDiffAggregator()],
+            "/rosout": [ROS2LogsAggregator()],
         },
         sources={
             "/camera/camera/color/image_raw": TopicConfig(
@@ -47,7 +47,7 @@ def main():
     agent = StateBasedAgent(
         config=config,
         connectors={"hri": connector},
-        tools=ROS2Toolkit(connector=ari_connector).get_tools(),
+        tools=ROS2Toolkit(connector=ros2_connector).get_tools(),
     )  # type: ignore
     agent.run()
     wait_for_shutdown([agent])

src/rai_core/rai/agents/state_based_agent.py

@@ -23,11 +23,11 @@
 from langchain_core.messages import HumanMessage
 from langchain_core.tools import BaseTool
 from pydantic import BaseModel, ConfigDict, Field
+from rai.aggregators import BaseAggregator
 from rclpy.callback_groups import ReentrantCallbackGroup
 from rclpy.subscription import Subscription
 
 from rai.agents.langchain.runnables import create_state_based_runnable
-from rai.agents.postprocessors import BaseStatePostprocessor
 from rai.communication.hri_connector import HRIConnector, HRIMessage
 from rai.communication.ros2.api.conversion import import_message_from_str
 from rai.communication.ros2.api.topic import TopicConfig
@@ -38,7 +38,7 @@
 
 
 class StateBasedConfig(BaseModel):
-    postprocessors: Dict[str, List[BaseStatePostprocessor]]
+    aggregators: Dict[str, List[BaseAggregator]]
     sources: Dict[str, TopicConfig]
     time_interval: float = Field(default=5.0)
     max_workers: int = 8
@@ -96,26 +96,13 @@ def __init__(
         self._aggregation_thread: threading.Thread | None = None
 
     def _configure_state_sources(self):
-        for topic, config in self.config.sources.items():
-            if topic in self._subscriptions:
-                continue
-            # NOTE(boczekbartek): refactor to use confugired_callbacks once implemented
-            # in the connector
-            qos_profile = self._ros2_connector._topic_api._resolve_qos_profile(
-                topic, config.auto_qos_matching, config.qos_profile, for_publisher=False
-            )
-            msg_type = import_message_from_str(config.msg_type)
-            self._subscriptions[topic] = self._ros2_connector.node.create_subscription(
-                msg_type=msg_type,
-                topic=topic,
-                callback=partial(self._state_topic_callback, topic),
-                qos_profile=qos_profile,
-                callback_group=self._callback_group,
-            )
-
-    def _state_topic_callback(self, topic_name: str, msg: Any):
-        with self._db_lock:
-            self._db[topic_name].append(msg)
+        for source, aggregators in self.config.aggregators.items():
+            for aggregator in aggregators:
+                self._ros2_connector.register_callback(
+                    source,
+                    aggregator,
+                    raw=True
+                )
 
     def run(self):
         super().run()
@@ -146,30 +133,25 @@ def _on_aggregation_interval(self):
         """Runs aggregation on collected data"""
 
         def process_aggregator(
-            source: str, postprocessor: BaseStatePostprocessor
+            source: str, aggregator: BaseAggregator
         ) -> Tuple[str, HumanMessage | HumanMultimodalMessage | None]:
-            with self._db_lock:
-                msgs = tuple(self._db[source])  # shallow, immutable copy
             self.logger.info(
-                f"Running postprocessor: {postprocessor}(source={source}) on {len(msgs)} messages"
+                f"Running postprocessor: {aggregator}(source={source}) on {len(aggregator.get_buffer())} messages"
             )
-            return source, postprocessor(msgs)
+            return source, aggregator.get()
 
         with ThreadPoolExecutor(max_workers=self.config.max_workers) as executor:
             futures = list()
-            for source, postprocessors in self.config.postprocessors.items():
-                for postprocessor in postprocessors:
-                    future = executor.submit(process_aggregator, source, postprocessor)
+            for source, aggregators in self.config.aggregators.items():
+                for aggregator in aggregators:
+                    future = executor.submit(process_aggregator, source, aggregator)
                     futures.append(future)
 
             for future in as_completed(futures):
                 source, output = future.result()
                 if output is not None:
                     self._aggregation_results[source] = output
 
-        with self._db_lock:
-            self._db.clear()
-
     def stop(self):
         """Stop the agent's execution loop."""
         self.logger.info("Stopping the agent. Please wait...")

src/rai_core/rai/aggregators/base.py

@@ -16,7 +16,8 @@
 from collections import deque
 from typing import Deque, Generic, TypeVar, List
 
-from langchain_core.messages import BaseMessage
+from langchain_core.messages import HumanMessage
+from rai.messages.multimodal import HumanMultimodalMessage
 
 
 T = TypeVar("T")
@@ -42,7 +43,7 @@ def __call__(
         self._buffer.append(msg)
 
     @abstractmethod
-    def get(self) -> BaseMessage | None:
+    def get(self) -> HumanMessage | HumanMultimodalMessage | None:
         """ Returns the aggregated message """
         pass
 

src/rai_core/rai/aggregators/ros2/__init__.py

@@ -12,16 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .postprocessors import (
-    ROS2GetLastImagePostprocessor,
-    ROS2ImgVLMDescriptionPostprocessor,
-    ROS2ImgVLMDiffPostprocessor,
-    ROS2LogsPostprocessor,
+from .aggregators import (
+    ROS2GetLastImageAggregator,
+    ROS2ImgVLMDescriptionAggregator,
+    ROS2ImgVLMDiffAggregator,
+    ROS2LogsAggregator,
 )
 
 __all__ = [
-    "ROS2GetLastImagePostprocessor",
-    "ROS2ImgVLMDescriptionPostprocessor",
-    "ROS2ImgVLMDiffPostprocessor",
-    "ROS2LogsPostprocessor",
+    "ROS2GetLastImageAggregator",
+    "ROS2ImgVLMDescriptionAggregator",
+    "ROS2ImgVLMDiffAggregator",
+    "ROS2LogsAggregator",
 ]

src/rai_core/rai/aggregators/ros2/aggregators.py

@@ -25,7 +25,7 @@
 from rai.messages import HumanMultimodalMessage
 
 
-class ROS2LogsPostprocessor(BaseAggregator[Log]):
+class ROS2LogsAggregator(BaseAggregator[Log]):
     """Returns only unique messages while keeping their order"""
 
     levels = {10: "DEBUG", 20: "INFO", 30: "WARNING", 40: "ERROR", 50: "FATAL"}
@@ -51,7 +51,7 @@ def get(self) -> HumanMessage:
         return HumanMessage(content=result)
 
 
-class ROS2GetLastImagePostprocessor(BaseAggregator[Image | CompressedImage]):
+class ROS2GetLastImageAggregator(BaseAggregator[Image | CompressedImage]):
     """Returns the last image from the buffer as base64 encoded string"""
 
     def get(self) -> HumanMultimodalMessage | None:
@@ -63,7 +63,7 @@ def get(self) -> HumanMultimodalMessage | None:
         return HumanMultimodalMessage(content="", images=[b64_image])
 
 
-class ROS2ImgVLMDescriptionPostprocessor(
+class ROS2ImgVLMDescriptionAggregator(
     BaseAggregator[Image | CompressedImage]
 ):
     """
@@ -104,7 +104,7 @@ class ROS2ImgDescription(BaseModel):
         )
 
 
-class ROS2ImgVLMDiffPostprocessor(BaseAggregator[Image | CompressedImage]):
+class ROS2ImgVLMDiffAggregator(BaseAggregator[Image | CompressedImage]):
     """
     Returns the LLM analysis of the differences between 3 images in the
     aggregation buffer: 1st, midden, last
@@ -125,10 +125,11 @@ def get_key_elements(elements: List[Any]) -> List[Any]:
         return [elements[0], elements[middle_index], elements[-1]]
 
     def get(self) -> HumanMessage | None:
-        if len(self.get_buffer()) == 0:
+        msgs = self.get_buffer()
+        if len(msgs) == 0:
             return None
 
-        b64_images = [encode_ros2_img_to_base64(msg) for msg in self._buffer]
+        b64_images = [encode_ros2_img_to_base64(msg) for msg in msgs]
 
         self.clear()
 
@@ -154,5 +155,5 @@ class ROS2ImgDiffOutput(BaseModel):
         llm = self.llm.with_structured_output(ROS2ImgDiffOutput)
         response = cast(ROS2ImgDiffOutput, llm.invoke(task))
         return HumanMessage(
-            content=f"Result of the analysis of the {len(b64_images)} keyframes selected from {len(msgs)} last images:\n{response}"
+            content=f"Result of the analysis of the {len(b64_images)} keyframes selected from {len(b64_images)} last images:\n{response}"
         )