From a37f262da09f3cdbcc327ac6404d5ad31e615990 Mon Sep 17 00:00:00 2001 From: Paul Nechifor Date: Sat, 11 Oct 2025 23:32:07 +0300 Subject: [PATCH] forgotten context managers --- dimos/agents2/skills/osm.py | 11 +++++------ dimos/agents2/skills/ros_navigation.py | 11 +++++------ dimos/mapping/osm/demo_osm.py | 4 ++-- dimos/robot/unitree_webrtc/unitree_g1.py | 17 ++++++++++------- pyproject.toml | 1 + 5 files changed, 23 insertions(+), 21 deletions(-) diff --git a/dimos/agents2/skills/osm.py b/dimos/agents2/skills/osm.py index c76242fe87..6c609e87f4 100644 --- a/dimos/agents2/skills/osm.py +++ b/dimos/agents2/skills/osm.py @@ -22,13 +22,14 @@ from dimos.protocol.skill.skill import SkillContainer, skill from dimos.robot.robot import Robot from dimos.utils.logging_config import setup_logger +from dimos.core.resource import Resource from reactivex.disposable import CompositeDisposable logger = setup_logger(__file__) -class OsmSkillContainer(SkillContainer): +class OsmSkillContainer(SkillContainer, Resource): _robot: Robot _disposables: CompositeDisposable _latest_location: Optional[LatLon] @@ -45,15 +46,13 @@ def __init__(self, robot: Robot, position_stream: Observable[LatLon]): self._current_location_map = CurrentLocationMap(QwenVlModel()) self._started = False - def __enter__(self) -> "OsmSkillContainer": + def start(self): self._started = True self._disposables.add(self._position_stream.subscribe(self._on_gps_location)) - return self - def __exit__(self, exc_type, exc_val, exc_tb): + def stop(self): self._disposables.dispose() - self.stop() - return False + super().stop() def _on_gps_location(self, location: LatLon) -> None: self._latest_location = location diff --git a/dimos/agents2/skills/ros_navigation.py b/dimos/agents2/skills/ros_navigation.py index d989f95b5d..e751a5d7aa 100644 --- a/dimos/agents2/skills/ros_navigation.py +++ b/dimos/agents2/skills/ros_navigation.py @@ -19,6 +19,7 @@ from dimos.protocol.skill.skill import SkillContainer, skill from dimos.utils.logging_config import setup_logger from dimos.utils.transform_utils import euler_to_quaternion +from dimos.core.resource import Resource from typing import TYPE_CHECKING @@ -28,7 +29,7 @@ logger = setup_logger(__file__) -class RosNavigation(SkillContainer): +class RosNavigation(SkillContainer, Resource): _robot: "UnitreeG1" _started: bool @@ -37,13 +38,11 @@ def __init__(self, robot: "UnitreeG1"): self._similarity_threshold = 0.23 self._started = False - def __enter__(self) -> "RosNavigation": + def start(self) -> None: self._started = True - return self - def __exit__(self, exc_type, exc_val, exc_tb): - self._close_module() - return False + def stop(self) -> None: + super().stop() @skill() def navigate_with_text(self, query: str) -> str: diff --git a/dimos/mapping/osm/demo_osm.py b/dimos/mapping/osm/demo_osm.py index c791008510..7617a48b9f 100644 --- a/dimos/mapping/osm/demo_osm.py +++ b/dimos/mapping/osm/demo_osm.py @@ -51,7 +51,7 @@ def start(self) -> None: self._robot = FakeRobot() self._agent = Agent(system_prompt=SYSTEM_PROMPT) self._osm_skill_container = OsmSkillContainer(self._robot, _get_fake_location()) - self._osm_skill_container.__enter__() + self._osm_skill_container.start() self._agent.register_skills(self._osm_skill_container) self._agent.register_skills(HumanInput()) self._agent.run_implicit_skill("human") @@ -64,7 +64,7 @@ def stop(self) -> None: if self._robot_debugger: self._robot_debugger.stop() if self._osm_skill_container: - self._osm_skill_container.__exit__(None, None, None) + self._osm_skill_container.stop() if self._agent: self._agent.stop() diff --git a/dimos/robot/unitree_webrtc/unitree_g1.py b/dimos/robot/unitree_webrtc/unitree_g1.py index b39ddd6db6..f319e2c87c 100644 --- a/dimos/robot/unitree_webrtc/unitree_g1.py +++ b/dimos/robot/unitree_webrtc/unitree_g1.py @@ -201,6 +201,7 @@ def __init__( self.joystick = None self.ros_bridge = None self.camera = None + self._ros_nav = None self._setup_directories() def _setup_directories(self): @@ -223,7 +224,7 @@ def _setup_directories(self): os.makedirs(self.db_path, exist_ok=True) def _deploy_detection(self, goto): - detection = self.dimos.deploy( + detection = self._dimos.deploy( ObjectDBModule, goto=goto, camera_info=zed.CameraInfo.SingleWebcam ) @@ -269,7 +270,7 @@ def start(self): if self.enable_ros_bridge: self._deploy_ros_bridge() - self.nav = self.dimos.deploy(NavigationModule) + self.nav = self._dimos.deploy(NavigationModule) self.nav.goal_reached.transport = core.LCMTransport("/goal_reached", Bool) self.nav.goal_pose.transport = core.LCMTransport("/goal_pose", PoseStamped) self.nav.cancel_goal.transport = core.LCMTransport("/cancel_goal", Bool) @@ -294,16 +295,16 @@ def start(self): g1_skills = UnitreeG1SkillContainer(robot=self) agent.register_skills(g1_skills) - human_input = self.dimos.deploy(HumanInput) + human_input = self._dimos.deploy(HumanInput) agent.register_skills(human_input) if self.enable_perception: agent.register_skills(self.detection) # Register ROS navigation - ros_nav = RosNavigation(self) - ros_nav.__enter__() - agent.register_skills(ros_nav) + self._ros_nav = RosNavigation(self) + self._ros_nav.start() + agent.register_skills(self._ros_nav) agent.run_implicit_skill("human") agent.start() @@ -320,6 +321,8 @@ def start(self): def stop(self) -> None: self._dimos.stop() + if self._ros_nav: + self._ros_nav.stop() self.lcm.stop() def _deploy_connection(self): @@ -373,7 +376,7 @@ def _deploy_visualization(self): self.foxglove_bridge.start() def _deploy_perception(self): - self.spatial_memory_module = self.dimos.deploy( + self.spatial_memory_module = self._dimos.deploy( SpatialMemory, collection_name=self.spatial_memory_collection, db_path=self.db_path, diff --git a/pyproject.toml b/pyproject.toml index 0e4f8e0fe6..c6e442908a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -174,6 +174,7 @@ dev = [ "pytest-asyncio==0.26.0", "pytest-mock==3.15.0", "pytest-env==1.1.5", + "pytest-timeout==2.4.0", "textual==3.7.1", "requests-mock==1.12.1", ]