Unified TimeSeriesStore with pluggable backends, global rewrite of timed event storage (#1080)

* embeddings not generic, lsp added to dev deps

* Fix return type annotations for reactive pipe operators

- quality_barrier: Callable[[Observable[T]], Observable[T]]
- sharpness_barrier: Callable[[Observable[Image]], Observable[Image]]

* sketching a new embedding db

* benchmark prints loss heatmap

* typing

* get_data supports nested paths

* sensor storage sketch and pickle implementation

* Implement SensorStore convenience methods with drift-free streaming

- Implement find_closest(), first_timestamp(), iterate(), iterate_ts(),
  iterate_realtime() methods using abstract _iter_items/_find_closest_timestamp
- Add scheduler-based stream() with absolute time reference to prevent timing
  drift during long playback (ported from replay.py)
- Move imports to top of file, add proper typing throughout
- Fix pickledir.py mypy error (pickle.load returns Any)

* Add SqliteStore backend for sensor data

- Single-file SQLite storage with indexed timestamp queries
- BLOB storage for pickled sensor data
- INSERT OR REPLACE for duplicate timestamp handling
- Supports multiple tables per database (different sensors)
- Added to parametrized tests (15 tests across 3 backends)

* Add PostgresStore backend for sensor data

- PostgresStore implements SensorStore[T] + Resource for lifecycle management
- Multiple stores can share same database with different tables
- Tables created automatically on first save
- Tests are optional - skip gracefully if PostgreSQL not available
- Added psycopg2-binary and types-psycopg2 dependencies
- Includes reset_db() helper for simple migrations (drop/recreate)

* Add SQL identifier validation and fix mypy issues

- Validate table/database names in SqliteStore and PostgresStore
  using regex (alphanumeric/underscore, not starting with digit)
- Fix Transform.to_pose() return type using TYPE_CHECKING import
- Add return type annotation to TF.get_pose()
- Fix ambiguous doclink in transports.md

* Make SqliteStore use get_data/get_data_dir like PickleDirStore

SqliteStore now accepts a name (e.g. "recordings/lidar") that gets
resolved via get_data_dir to data/recordings/lidar.db. Still supports
absolute paths and :memory: for backward compatibility.

* Require T to be Timestamped subclass in SensorStore

- Add TypeVar bound: T = TypeVar("T", bound=Timestamped)
- Simplify save() to always use data.ts (no more optional timestamp)
- Update tests to use SampleData(Timestamped) instead of strings
- SqliteStore accepts str | Path for backward compatibility

* consolidating old replay and new sensor store

* Add LegacyPickleStore backend for TimedSensorReplay compatibility

* legacy sensor store implemented, replay shim implemented

* replaced legacy replay with new system

* Remove import-untyped type ignores for consistency with dev

* Restore import-untyped type ignores for local mypy compatibility

Required when cupy/contact_graspnet are installed locally without type stubs.

* record/replay on modules

* Fix mypy errors: add psycopg2 stubs and xacro type ignore

- Add import-untyped to xacro type: ignore comment in mesh_utils.py
- Remove unused record/replay RPC methods from ModuleBase

* Decouple SensorStore from Timestamped type constraint

Remove the Timestamped bound from SensorStore's TypeVar, enabling storage
of arbitrary data types. Timestamps are now provided explicitly via
save(ts, data), with Timestamped convenience methods (save_ts, pipe_save_ts,
consume_stream_ts) as opt-in helpers. iterate_realtime() and stream() now
use stored timestamps instead of data.ts.

* CI code cleanup

* Fix ruff errors in embedding models

Fix import sorting in treid.py and suppress B027 for optional warmup()
in base.py.

* Remove unused EmbeddingModel.warmup method

* Remove unused warmup(), fix consume_stream callers

Remove dead EmbeddingModel.warmup() method. Update go2.py to use
consume_stream_ts() for the new SensorStore API.

* Rename SensorStore methods: Timestamped as default, raw for explicit ts

save/pipe_save/consume_stream now work with Timestamped data by default.
save_raw/pipe_save_raw/consume_stream_raw take explicit timestamps for
non-Timestamped data.

* Rename SensorStore to TimeSeriesStore

* Add _delete to all backends, fix find_closest and add/get/prune_old

Implement _delete for InMemoryStore, SqliteStore, PickleDirStore,
PostgresStore (LegacyPickleStore raises NotImplementedError). Fix
find_closest docstring placement and add get/add/prune_old convenience
methods.

* Add collection API to TimeSeriesStore, rewrite InMemoryStore with SortedKeyList

Replace InMemoryStore's dict + sorted-cache (O(n log n) rebuild on every write)
with SortedKeyList for O(log n) insert, delete, and range queries. Add collection
methods to TimeSeriesStore base: __len__, __iter__, last/last_timestamp, start_ts/
end_ts, time_range, duration, find_before/find_after, slice_by_time. Implement
backing abstract methods (_count, _last_timestamp, _find_before, _find_after) in
all five backends. Performance benchmarks confirm InMemoryStore matches
TimestampedCollection on 100k items.

* Rename memory/sensor to memory/timeseries, extract InMemoryStore to inmemory.py

Rename the module to better reflect its purpose. Extract InMemoryStore from
base.py into its own file (inmemory.py) to keep base.py focused on the abstract
TimeSeriesStore class. Update all internal and external imports.

* Simplify TimeSeriesStore: bound T to Timestamped, remove _raw methods, store T directly

- Bound T to Timestamped — no more raw/non-Timestamped data paths
- Removed save_raw, pipe_save_raw, consume_stream_raw
- InMemoryStore stores T directly in SortedKeyList (no _Entry wrapper)
- Removed duplicate-check on insert (same semantics as TimestampedCollection)
- Performance now at parity with TimestampedCollection

* tests reorganization

* replacing memory store in tf.py

* Fix mypy types: remove base get/add (TBuffer overrides), add None guards in tests

* Replace TimestampedCollection with InMemoryStore, remove dead code

- Delete TimestampedCollection class (replaced by InMemoryStore)
- Rewrite TimestampedBufferCollection to inherit InMemoryStore
- Remove TBuffer_old dead code from tf.py
- Fix prune_old mutation-during-iteration bug in base.py
- Break circular import with TYPE_CHECKING guard in base.py
- Update Image.py to use public API instead of _items access
- Update tests to use InMemoryStore directly

* delete legacy code

* removed small comment

* test fix, psql mypy

* fix(memory): import ModuleConfig from concrete module to fix mypy error

Import Module and ModuleConfig from dimos.core.module instead of the
lazy-loaded dimos.core namespace, which mypy sees as type Any.

Author: leshy

.gitignore

@@ -66,4 +66,5 @@ yolo11n.pt
 *mobileclip*
 /results
 
+CLAUDE.MD
 /assets/teleop_certs/

dimos/core/stream.py

@@ -69,7 +69,7 @@ def _subscribe(observer, scheduler=None):  # type: ignore[no-untyped-def]
 
     # default return is backpressured because most
     # use cases will want this by default
-    def observable(self):  # type: ignore[no-untyped-def]
+    def observable(self) -> Observable[T]:
         return backpressure(self.pure_observable())
 
 

dimos/core/testing.py

@@ -79,6 +79,5 @@ def odomloop(self) -> None:
                 self.odometry.publish(odom)
 
                 lidarmsg = next(lidariter)
-                lidarmsg.pubtime = time.perf_counter()  # type: ignore[union-attr]
                 self.lidar.publish(lidarmsg)
                 time.sleep(0.1)

dimos/mapping/pointclouds/test_occupancy_speed.py

@@ -20,19 +20,18 @@
 from dimos.mapping.pointclouds.occupancy import OCCUPANCY_ALGOS
 from dimos.mapping.voxels import VoxelGridMapper
 from dimos.utils.cli.plot import bar
-from dimos.utils.data import _get_data_dir, get_data
+from dimos.utils.data import get_data, get_data_dir
 from dimos.utils.testing import TimedSensorReplay
 
 
 @pytest.mark.tool
 def test_build_map():
     mapper = VoxelGridMapper(publish_interval=-1)
 
-    for ts, frame in TimedSensorReplay("unitree_go2_bigoffice/lidar").iterate_duration():
-        print(ts, frame)
+    for _ts, frame in TimedSensorReplay("unitree_go2_bigoffice/lidar").iterate():
         mapper.add_frame(frame)
 
-    pickle_file = _get_data_dir() / "unitree_go2_bigoffice_map.pickle"
+    pickle_file = get_data_dir() / "unitree_go2_bigoffice_map.pickle"
     global_pcd = mapper.get_global_pointcloud2()
 
     with open(pickle_file, "wb") as f:

dimos/memory/embedding.py

@@ -0,0 +1,105 @@
+# Copyright 2026 Dimensional Inc.
+#
+# 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 collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import cast
+
+import reactivex as rx
+from reactivex import operators as ops
+from reactivex.observable import Observable
+
+from dimos.core import In, rpc
+from dimos.core.module import Module, ModuleConfig
+from dimos.models.embedding.base import Embedding, EmbeddingModel
+from dimos.models.embedding.clip import CLIPModel
+from dimos.msgs.geometry_msgs import PoseStamped
+from dimos.msgs.nav_msgs import OccupancyGrid
+from dimos.msgs.sensor_msgs import Image
+from dimos.msgs.sensor_msgs.Image import Image, sharpness_barrier
+from dimos.utils.reactive import getter_hot
+
+
+@dataclass
+class Config(ModuleConfig):
+    embedding_model: EmbeddingModel = field(default_factory=CLIPModel)
+
+
+@dataclass
+class SpatialEntry:
+    image: Image
+    pose: PoseStamped
+
+
+@dataclass
+class SpatialEmbedding(SpatialEntry):
+    embedding: Embedding
+
+
+class EmbeddingMemory(Module[Config]):
+    default_config = Config
+    config: Config
+    color_image: In[Image]
+    global_costmap: In[OccupancyGrid]
+
+    _costmap_getter: Callable[[], OccupancyGrid] | None = None
+
+    def get_costmap(self) -> OccupancyGrid:
+        if self._costmap_getter is None:
+            self._costmap_getter = getter_hot(self.global_costmap.pure_observable())
+            self._disposables.add(self._costmap_getter)
+        return self._costmap_getter()
+
+    @rpc
+    def query_costmap(self, text: str) -> OccupancyGrid:
+        costmap = self.get_costmap()
+        # overlay costmap with embedding heat
+        return costmap
+
+    @rpc
+    def start(self) -> None:
+        # would be cool if this sharpness_barrier was somehow self-calibrating
+        #
+        # we need a Governor system, sharpness_barrier frequency shouldn't
+        # be a fixed float but an observable that adjusts based on downstream load
+        #
+        # (also voxel size for mapper for example would benefit from this)
+        self.color_image.pure_observable().pipe(
+            sharpness_barrier(0.5),
+            ops.flat_map(self._try_create_spatial_entry),
+            ops.map(self._embed_spatial_entry),
+            ops.map(self._store_spatial_entry),
+        ).subscribe(print)
+
+    def _try_create_spatial_entry(self, img: Image) -> Observable[SpatialEntry]:
+        pose = self.tf.get_pose("world", "base_link")
+        if not pose:
+            return rx.empty()
+        return rx.of(SpatialEntry(image=img, pose=pose))
+
+    def _embed_spatial_entry(self, spatial_entry: SpatialEntry) -> SpatialEmbedding:
+        embedding = cast("Embedding", self.config.embedding_model.embed(spatial_entry.image))
+        return SpatialEmbedding(
+            image=spatial_entry.image,
+            pose=spatial_entry.pose,
+            embedding=embedding,
+        )
+
+    def _store_spatial_entry(self, spatial_embedding: SpatialEmbedding) -> SpatialEmbedding:
+        return spatial_embedding
+
+    def query_text(self, query: str) -> list[SpatialEmbedding]:
+        self.config.embedding_model.embed_text(query)
+        results: list[SpatialEmbedding] = []
+        return results

dimos/memory/test_embedding.py

@@ -0,0 +1,53 @@
+# Copyright 2026 Dimensional Inc.
+#
+# 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.
+
+import pytest
+
+from dimos.memory.embedding import EmbeddingMemory, SpatialEntry
+from dimos.msgs.geometry_msgs import PoseStamped
+from dimos.utils.data import get_data
+from dimos.utils.testing import TimedSensorReplay
+
+dir_name = "unitree_go2_bigoffice"
+
+
+@pytest.mark.skip
+def test_embed_frame() -> None:
+    """Test embedding a single frame."""
+    # Load a frame from recorded data
+    video = TimedSensorReplay(get_data(dir_name) / "video")
+    frame = video.find_closest_seek(10)
+
+    # Create memory and embed
+    memory = EmbeddingMemory()
+
+    try:
+        # Create a spatial entry with dummy pose (no TF needed for this test)
+        dummy_pose = PoseStamped(
+            position=[0, 0, 0],
+            orientation=[0, 0, 0, 1],  # identity quaternion
+        )
+        spatial_entry = SpatialEntry(image=frame, pose=dummy_pose)
+
+        # Embed the frame
+        result = memory._embed_spatial_entry(spatial_entry)
+
+        # Verify
+        assert result is not None
+        assert result.embedding is not None
+        assert result.embedding.vector is not None
+        print(f"Embedding shape: {result.embedding.vector.shape}")
+        print(f"Embedding vector (first 5): {result.embedding.vector[:5]}")
+    finally:
+        memory.stop()

dimos/memory/timeseries/__init__.py

@@ -0,0 +1,41 @@
+# Copyright 2025-2026 Dimensional Inc.
+#
+# 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.
+"""Time series storage and replay."""
+
+from dimos.memory.timeseries.base import TimeSeriesStore
+from dimos.memory.timeseries.inmemory import InMemoryStore
+from dimos.memory.timeseries.pickledir import PickleDirStore
+from dimos.memory.timeseries.sqlite import SqliteStore
+
+
+def __getattr__(name: str):  # type: ignore[no-untyped-def]
+    if name == "PostgresStore":
+        from dimos.memory.timeseries.postgres import PostgresStore
+
+        return PostgresStore
+    if name == "reset_db":
+        from dimos.memory.timeseries.postgres import reset_db
+
+        return reset_db
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = [
+    "InMemoryStore",
+    "PickleDirStore",
+    "PostgresStore",
+    "SqliteStore",
+    "TimeSeriesStore",
+    "reset_db",
+]