refactor: remove ROS 2 imports from non ROS 2 files in rai_core (#528)

Co-authored-by: Bartłomiej Boczek <[email protected]>

Author: maciejmajek

src/rai_core/rai/agents/conversational_agent.py

@@ -15,26 +15,23 @@
 
 import logging
 from functools import partial
-from typing import List, Optional, TypedDict, Union
+from typing import List, Optional, TypedDict
 
 from langchain.chat_models.base import BaseChatModel
 from langchain_core.messages import BaseMessage, SystemMessage
 from langchain_core.tools import BaseTool
 from langgraph.graph import START, StateGraph
 from langgraph.graph.state import CompiledStateGraph
 from langgraph.prebuilt.tool_node import tools_condition
-from rclpy.impl.rcutils_logger import RcutilsLogger
 
 from rai.agents.tool_runner import ToolRunner
 
-loggers_type = Union[RcutilsLogger, logging.Logger]
-
 
 class State(TypedDict):
     messages: List[BaseMessage]
 
 
-def agent(llm: BaseChatModel, logger: loggers_type, system_prompt: str, state: State):
+def agent(llm: BaseChatModel, logger: logging.Logger, system_prompt: str, state: State):
     logger.info("Running thinker")
 
     # If there are no messages, do nothing
@@ -53,7 +50,7 @@ def create_conversational_agent(
     llm: BaseChatModel,
     tools: List[BaseTool],
     system_prompt: str,
-    logger: Optional[RcutilsLogger | logging.Logger] = None,
+    logger: Optional[logging.Logger] = None,
     debug=False,
 ) -> CompiledStateGraph:
     _logger = None

src/rai_core/rai/agents/state_based.py

@@ -27,13 +27,10 @@
 from langgraph.graph.graph import CompiledGraph
 from langgraph.prebuilt.tool_node import msg_content_output
 from pydantic import BaseModel, Field, ValidationError
-from rclpy.impl.rcutils_logger import RcutilsLogger
 
 from rai.agents.tool_runner import ToolRunner
 from rai.messages import HumanMultimodalMessage
 
-loggers_type = Union[RcutilsLogger, logging.Logger]
-
 
 class State(TypedDict):
     messages: List[BaseMessage]
@@ -68,7 +65,7 @@ def tools_condition(
     return "reporter"
 
 
-def thinker(llm: BaseChatModel, logger: loggers_type, state: State):
+def thinker(llm: BaseChatModel, logger: logging.Logger, state: State):
     logger.info("Running thinker")
     prompt = (
         "Based on the data provided, reason about the situation. "
@@ -81,7 +78,7 @@ def thinker(llm: BaseChatModel, logger: loggers_type, state: State):
 
 
 def decider(
-    llm: Runnable[LanguageModelInput, BaseMessage], logger: loggers_type, state: State
+    llm: Runnable[LanguageModelInput, BaseMessage], logger: logging.Logger, state: State
 ):
     logger.info("Running decider")
     prompt = (
@@ -98,7 +95,7 @@ def decider(
     return state
 
 
-def reporter(llm: BaseChatModel, logger: loggers_type, state: State):
+def reporter(llm: BaseChatModel, logger: logging.Logger, state: State):
     logger.info("Summarizing the conversation")
     prompt = (
         "You are the reporter. Your task is to summarize what happened previously. "
@@ -126,7 +123,7 @@ def reporter(llm: BaseChatModel, logger: loggers_type, state: State):
 
 
 def retriever_wrapper(
-    state_retriever: Callable[[], Dict[str, Any]], logger: loggers_type, state: State
+    state_retriever: Callable[[], Dict[str, Any]], logger: logging.Logger, state: State
 ):
     """This wrapper is used to retrieve multimodal information from the output of state_retriever."""
     ts = time.perf_counter()
@@ -150,10 +147,10 @@ def create_state_based_agent(
     llm: BaseChatModel,
     tools: List[BaseTool],
     state_retriever: Callable[[], Dict[str, Any]],
-    logger: Optional[RcutilsLogger | logging.Logger] = None,
+    logger: Optional[logging.Logger] = None,
 ) -> CompiledGraph:
     _logger = None
-    if isinstance(logger, RcutilsLogger):
+    if isinstance(logger, logging.Logger):
         _logger = logger
     else:
         _logger = logging.getLogger(__name__)

src/rai_core/rai/agents/tool_runner.py

@@ -26,7 +26,6 @@
 from langgraph.prebuilt.tool_node import msg_content_output
 from langgraph.utils.runnable import RunnableCallable
 from pydantic import ValidationError
-from rclpy.impl.rcutils_logger import RcutilsLogger
 
 from rai.messages import MultimodalArtifact, ToolMultimodalMessage, store_artifacts
 
@@ -38,7 +37,7 @@ def __init__(
         *,
         name: str = "tools",
         tags: Optional[list[str]] = None,
-        logger: Optional[Union[RcutilsLogger, logging.Logger]] = None,
+        logger: Optional[logging.Logger] = None,
     ) -> None:
         super().__init__(self._func, name=name, tags=tags, trace=False)
         self.logger = logger or logging.getLogger(__name__)

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

@@ -12,6 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import importlib.util
+
+if importlib.util.find_spec("rclpy") is None:
+    raise ImportError(
+        "This is a ROS2 feature. Make sure ROS2 is installed and sourced."
+    )
+
 from .api import (
     IROS2Message,  # TODO: IROS2Message should not be a part of the public API
     TopicConfig,  # TODO: TopicConfig should not be a part of the public API

src/rai_core/rai/communication/sound_device/__init__.py

@@ -12,6 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import importlib.util
+
+if importlib.util.find_spec("sounddevice") is None:
+    raise ImportError(
+        "This feature is based on sounddevice. Make sure sounddevice is installed."
+    )
+
 from .api import SoundDeviceAPI, SoundDeviceConfig, SoundDeviceError
 from .connector import SoundDeviceConnector, SoundDeviceMessage
 

src/rai_core/rai/communication/sound_device/api.py

@@ -17,6 +17,7 @@
 from typing import Any, Callable, Optional
 
 import numpy as np
+import sounddevice as sd
 from numpy._typing import NDArray
 from pydub import AudioSegment
 from scipy.signal import resample
@@ -85,10 +86,6 @@ def __init__(self, config: SoundDeviceConfig):
         self.device_name = ""
 
         self.config = config
-        try:
-            import sounddevice as sd
-        except ImportError:
-            raise SoundDeviceError("SoundDeviceAPI requires sound_device module!")
         sd.default.latency = ("low", "low")  # type: ignore
         if config.device_name:
             self.device_name = config.device_name
@@ -143,10 +140,7 @@ def write(self, data: AudioSegment, blocking: bool = False, loop: bool = False):
         """
         if not self.write_flag:
             raise SoundDeviceError(f"{self.device_name} does not support writing!")
-        try:
-            import sounddevice as sd
-        except ImportError:
-            raise SoundDeviceError("SoundDeviceAPI requires sound_device module!")
+
         audio = np.array(data.get_array_of_samples())
         sd.play(
             audio,
@@ -187,10 +181,7 @@ def read(self, time: float, blocking: bool = False) -> AudioSegment:
 
         if not self.read_flag:
             raise SoundDeviceError(f"{self.device_name} does not support reading!")
-        try:
-            import sounddevice as sd
-        except ImportError:
-            raise SoundDeviceError("SoundDeviceAPI requires sound_device module!")
+
         frames = int(time * self.sample_rate)
         recording = sd.rec(
             frames=frames,
@@ -217,10 +208,6 @@ def stop(self):
         - This is a convenience function to stop the sound device from playing or recording.
         - It will stop any sound that is currently playing and any recording currently happening.
         """
-        try:
-            import sounddevice as sd
-        except ImportError:
-            raise SoundDeviceError("SoundDeviceAPI requires sound_device module!")
         sd.stop()
 
     def wait(self):
@@ -232,10 +219,6 @@ def wait(self):
         - This is a convenience function to wait for the sound device to finish playing or recording.
         - It will block until the sound is played or recorded.
         """
-        try:
-            import sounddevice as sd
-        except ImportError:
-            raise SoundDeviceError("SoundDeviceAPI requires sound_device module!")
         sd.wait()
 
     def open_write_stream(
@@ -250,10 +233,6 @@ def open_write_stream(
                 f"{self.device_name} does not support streaming writing!"
             )
 
-        try:
-            import sounddevice as sd
-        except ImportError:
-            raise SoundDeviceError("SoundDeviceAPI requires sound_device module!")
         from sounddevice import CallbackFlags
 
         def callback(indata: NDArray, frames: int, time: Any, status: CallbackFlags):
@@ -323,10 +302,6 @@ def callback(indata: NDArray, frames: int, _, status: CallbackFlags):
             }
             on_feedback(indata, flag_dict)
 
-        try:
-            import sounddevice as sd
-        except ImportError:
-            raise SoundDeviceError("SoundDeviceAPI requires sound_device module!")
         try:
             if self.config.consumer_sampling_rate is None:
                 window_size_samples = self.config.block_size * self.sample_rate

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

@@ -12,6 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import importlib.util
+
+if importlib.util.find_spec("rclpy") is None:
+    raise ImportError(
+        "This is a ROS2 feature. Make sure ROS2 is installed and sourced."
+    )
+
 from .cli import (
     ROS2CLIToolkit,
     ros2_action,

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

@@ -12,13 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import importlib.util
-
-if importlib.util.find_spec("rclpy") is None:
-    raise ImportError(
-        "This is a ROS2 feature. Make sure ROS2 is installed and sourced."
-    )
-
 import uuid
 from collections import defaultdict
 from functools import partial

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

@@ -12,13 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import importlib.util
-
-if importlib.util.find_spec("rclpy") is None:
-    raise ImportError(
-        "This is a ROS2 feature. Make sure ROS2 is installed and sourced."
-    )
-
 from typing import Any, Dict, List, Type
 
 from langchain_core.tools import BaseTool

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

@@ -12,13 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import importlib.util
-
-if importlib.util.find_spec("rclpy") is None:
-    raise ImportError(
-        "This is a ROS2 feature. Make sure ROS2 is installed and sourced."
-    )
-
 import json
 from typing import Any, Dict, List, Literal, Tuple, Type