fix imports

Author: boczekbartek

examples/agents/state_based.py

@@ -12,8 +12,8 @@
 # See the License for the specific language goveself.rning permissions and
 # limitations under the License.
 
-from pydantic import config
-from rai.agents import ROS2StateBasedAgent, StateBasedConfig, wait_for_shutdown
+from rai.agents.ros2 import ROS2StateBasedAgent
+from rai.agents import StateBasedConfig, wait_for_shutdown
 from rai.aggregators.ros2 import (
     ROS2ImgVLMDiffAggregator,
     ROS2LogsAggregator,

src/rai_core/rai/agents/__init__.py

@@ -15,13 +15,13 @@
 from rai.agents.conversational_agent import create_conversational_agent
 from rai.agents.react_agent import ReActAgent
 from rai.agents.runner import AgentRunner, wait_for_shutdown
-from rai.agents.state_based_agent import ROS2StateBasedAgent, StateBasedConfig
+from rai.agents.base_state_based_agent import BaseStateBasedAgent, StateBasedConfig
 from rai.agents.tool_runner import ToolRunner
 
 __all__ = [
     "AgentRunner",
     "ReActAgent",
-    "ROS2StateBasedAgent",
+    "BaseStateBasedAgent",
     "StateBasedConfig",
     "ToolRunner",
     "create_conversational_agent",

src/rai_core/rai/agents/base_state_based_agent.py

@@ -44,7 +44,7 @@ class StateBasedConfig(BaseModel):
     )
 
 
-class StateBasedAgent(ReActAgent, ABC):
+class BaseStateBasedAgent(ReActAgent, ABC):
     """
     Agent that runs aggregators (config.aggregators) every config.time_interval seconds.
     Aggregators are registered to their sources using

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

@@ -0,0 +1,16 @@
+# Copyright (C) 2024 Robotec.AI
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from .state_based_agent import ROS2StateBasedAgent
+
+__all__ = ["ROS2StateBasedAgent"]

src/rai_core/rai/agents/ros2/state_based_agent.py

@@ -1,7 +1,20 @@
+# Copyright (C) 2025 Robotec.AI
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 
-from rai.agents.state_based_agent import StateBasedAgent
-from rai.communication.ros2.connectors.connector import ROS2Connector
+from rai.agents import BaseStateBasedAgent
+from rai.communication.ros2 import ROS2Connector
 
-class ROS2StateBasedAgent(StateBasedAgent):
+class ROS2StateBasedAgent(BaseStateBasedAgent):
     def setup_connector(self):
         return ROS2Connector()