memory2 StreamModules, cleanup (#1682)

Author: leshy

data/.lfs/go2_bigoffice.db.tar.gz

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:2d48cb0b8250bb2878d1008093d45ea377406de00ad42f0f96d7b382e1a9617b
-size 191193336
+oid sha256:142f7a7d64d3b77c97acd0d15d53e9ea28c4f558776a6bb3919a4da32c2f4d37
+size 192241937

dimos/agents/agent_test_runner.py

@@ -48,8 +48,8 @@ def __init__(self, **kwargs: Any) -> None:
     @rpc
     def start(self) -> None:
         super().start()
-        self._disposables.add(Disposable(self.agent.subscribe(self._on_agent_message)))
-        self._disposables.add(Disposable(self.agent_idle.subscribe(self._on_agent_idle)))
+        self.register_disposable(Disposable(self.agent.subscribe(self._on_agent_message)))
+        self.register_disposable(Disposable(self.agent_idle.subscribe(self._on_agent_idle)))
         # Signal that subscription is ready
         self._subscription_ready.set()
 

dimos/agents/mcp/mcp_client.py

@@ -169,7 +169,7 @@ def start(self) -> None:
         def _on_human_input(string: str) -> None:
             self._message_queue.put(HumanMessage(content=string))
 
-        self._disposables.add(Disposable(self.human_input.subscribe(_on_human_input)))
+        self.register_disposable(Disposable(self.human_input.subscribe(_on_human_input)))
 
     @rpc
     def on_system_modules(self, _modules: list[RPCClient]) -> None:

dimos/agents/skills/demo_robot.py

@@ -25,7 +25,7 @@ class DemoRobot(Module):
 
     def start(self) -> None:
         super().start()
-        self._disposables.add(interval(1.0).subscribe(lambda _: self._publish_gps_location()))
+        self.register_disposable(interval(1.0).subscribe(lambda _: self._publish_gps_location()))
 
     def stop(self) -> None:
         super().stop()

dimos/agents/skills/google_maps_skill_container.py

@@ -15,6 +15,8 @@
 import json
 from typing import Any
 
+from reactivex.disposable import Disposable
+
 from dimos.agents.annotation import skill
 from dimos.core.core import rpc
 from dimos.core.module import Module
@@ -49,7 +51,7 @@ def __init__(self, **kwargs: Any) -> None:
     @rpc
     def start(self) -> None:
         super().start()
-        self._disposables.add(self.gps_location.subscribe(self._on_gps_location))  # type: ignore[arg-type]
+        self.register_disposable(Disposable(self.gps_location.subscribe(self._on_gps_location)))
 
     @rpc
     def stop(self) -> None:

dimos/agents/skills/gps_nav_skill.py

@@ -14,6 +14,8 @@
 
 import json
 
+from reactivex.disposable import Disposable
+
 from dimos.agents.annotation import skill
 from dimos.core.core import rpc
 from dimos.core.module import Module
@@ -38,7 +40,7 @@ class GpsNavSkillContainer(Module):
     @rpc
     def start(self) -> None:
         super().start()
-        self._disposables.add(self.gps_location.subscribe(self._on_gps_location))  # type: ignore[arg-type]
+        self.register_disposable(Disposable(self.gps_location.subscribe(self._on_gps_location)))
 
     @rpc
     def stop(self) -> None:

dimos/agents/skills/navigation.py

@@ -62,8 +62,8 @@ def __init__(self, **kwargs: Any) -> None:
     @rpc
     def start(self) -> None:
         super().start()
-        self._disposables.add(Disposable(self.color_image.subscribe(self._on_color_image)))
-        self._disposables.add(Disposable(self.odom.subscribe(self._on_odom)))
+        self.register_disposable(Disposable(self.color_image.subscribe(self._on_color_image)))
+        self.register_disposable(Disposable(self.odom.subscribe(self._on_odom)))
         self._skill_started = True
 
     @rpc

dimos/agents/skills/osm.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 
 
+from reactivex.disposable import Disposable
+
 from dimos.agents.annotation import skill
 from dimos.core.module import Module
 from dimos.core.stream import In
@@ -39,7 +41,7 @@ def __init__(self) -> None:
     def start(self) -> None:
         super().start()
         if hasattr(self.gps_location, "subscribe"):
-            self._disposables.add(self.gps_location.subscribe(self._on_gps_location))  # type: ignore[arg-type]
+            self.register_disposable(Disposable(self.gps_location.subscribe(self._on_gps_location)))
         else:
             logger.warning(
                 "OsmSkill: gps_location stream does not support direct subscribe (RemoteIn)"

dimos/agents/skills/person_follow.py

@@ -94,9 +94,9 @@ def __init__(self, **kwargs: Any) -> None:
     @rpc
     def start(self) -> None:
         super().start()
-        self._disposables.add(Disposable(self.color_image.subscribe(self._on_color_image)))
+        self.register_disposable(Disposable(self.color_image.subscribe(self._on_color_image)))
         if self.config.use_3d_navigation:
-            self._disposables.add(Disposable(self.global_map.subscribe(self._on_pointcloud)))
+            self.register_disposable(Disposable(self.global_map.subscribe(self._on_pointcloud)))
 
     @rpc
     def stop(self) -> None:

dimos/agents/vlm_agent.py

@@ -16,6 +16,7 @@
 
 from langchain.chat_models import init_chat_model
 from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
+from reactivex.disposable import Disposable
 
 from dimos.agents.system_prompt import SYSTEM_PROMPT
 from dimos.core.core import rpc
@@ -60,8 +61,8 @@ def __init__(self, **kwargs: Any) -> None:
     @rpc
     def start(self) -> None:
         super().start()
-        self._disposables.add(self.color_image.subscribe(self._on_image))  # type: ignore[arg-type]
-        self._disposables.add(self.query_stream.subscribe(self._on_query))  # type: ignore[arg-type]
+        self.register_disposable(Disposable(self.color_image.subscribe(self._on_image)))
+        self.register_disposable(Disposable(self.query_stream.subscribe(self._on_query)))
 
     @rpc
     def stop(self) -> None: