cleanup

Author: boczekbartek

examples/agents/state_based.py

@@ -30,18 +30,14 @@ def main():
     connector = ROS2HRIConnector(sources=["/from_human"], targets=["/to_human"])
     ros2_connector = ROS2Connector()
 
-    aggregators={
+    aggregators = {
         "/camera/camera/color/image_raw": [ROS2ImgVLMDiffAggregator()],
         "/rosout": [ROS2LogsAggregator()],
     }
 
-    for source, aggregators in self.config.aggregators.items():
+    for source, aggregators in aggregators.items():
         for aggregator in aggregators:
-            ros2_connector.register_callback(
-                source,
-                aggregator,
-                raw=True
-            )
+            ros2_connector.register_callback(source, aggregator, raw=True)
 
     agent = StateBasedAgent(
         config=StateBasedConfig(aggregators=aggregators),

src/rai_core/rai/agents/state_based_agent.py

@@ -14,22 +14,19 @@
 
 import threading
 import time
-from collections import defaultdict
 from concurrent.futures import ThreadPoolExecutor, as_completed
-from functools import partial
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Dict, List, Optional, Tuple
 
 from langchain_core.language_models import BaseChatModel
 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.aggregators import BaseAggregator
 from rai.communication.hri_connector import HRIConnector, HRIMessage
-from rai.communication.ros2.connectors import ROS2Connector
 from rai.messages.multimodal import HumanMultimodalMessage
 
 from .react_agent import ReActAgent, ReActAgentState
@@ -48,7 +45,7 @@ class StateBasedConfig(BaseModel):
 class StateBasedAgent(ReActAgent):
     """
     Agent that runs aggregators (config.aggregators) every config.time_interval seconds.
-    Aggregators are registered to their sources using 
+    Aggregators are registered to their sources using
     :py:class:`~rai.communication.ros2.connectors.ROS2Connector`
 
     Output from aggragators is called `state`. Such state is saved and can be
@@ -77,27 +74,19 @@ def __init__(
             connectors, llm, tools, state, system_prompt, runnable=runnable
         )
         self.config = config
-        self.ros2_connector = ros2_connector
 
         self._callback_group = ReentrantCallbackGroup()
         self._subscriptions: Dict[str, Subscription] = dict()
 
         self._aggregation_results: Dict[str, HumanMessage | HumanMultimodalMessage] = (
             dict()
         )
-        self._db: Dict[str, List[Any]] = defaultdict(list)
-        self._db_lock = threading.Lock()
-
         self._aggregation_thread: threading.Thread | None = None
 
     def _configure_state_sources(self):
         for source, aggregators in self.config.aggregators.items():
             for aggregator in aggregators:
-                self.ros2_connector.register_callback(
-                    source,
-                    aggregator,
-                    raw=True
-                )
+                self.ros2_connector.register_callback(source, aggregator, raw=True)
 
     def run(self):
         super().run()
@@ -157,6 +146,6 @@ def stop(self):
         if self._aggregation_thread is not None:
             self._aggregation_thread.join()
             self._aggregation_thread = None
-        #TODO(boczekbartek): deregister aggregators
+        # TODO(boczekbartek): deregister aggregators
         self._stop_event.clear()
         self.logger.info("Agent stopped")

src/rai_core/rai/tools/ros2/generic/topics.py

@@ -22,8 +22,8 @@
 from pydantic import BaseModel, Field
 
 from rai.communication.ros2 import ROS2Connector, ROS2Message
-from rai.communication.ros2.api.conversion import ros2_message_to_dict
-from rai.messages import MultimodalArtifact, preprocess_image
+from rai.communication.ros2.api.conversion import encode_ros2_img_to_base64, ros2_message_to_dict
+from rai.messages import MultimodalArtifact
 from rai.tools.ros2.base import BaseROS2Tool, BaseROS2Toolkit