diff --git a/.gitignore b/.gitignore index 686327e990..44966af366 100644 --- a/.gitignore +++ b/.gitignore @@ -66,4 +66,5 @@ yolo11n.pt *mobileclip* /results +CLAUDE.MD /assets/teleop_certs/ diff --git a/dimos/core/stream.py b/dimos/core/stream.py index 89871dabc1..77edf45417 100644 --- a/dimos/core/stream.py +++ b/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()) diff --git a/dimos/core/testing.py b/dimos/core/testing.py index 38774ef327..c08f2477c2 100644 --- a/dimos/core/testing.py +++ b/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) diff --git a/dimos/manipulation/planning/utils/mesh_utils.py b/dimos/manipulation/planning/utils/mesh_utils.py index 46d14eb1b4..7525cc29a3 100644 --- a/dimos/manipulation/planning/utils/mesh_utils.py +++ b/dimos/manipulation/planning/utils/mesh_utils.py @@ -133,7 +133,7 @@ def _process_xacro( ) -> str: """Process xacro file to URDF.""" try: - import xacro # type: ignore[import-not-found] + import xacro # type: ignore[import-not-found,import-untyped] except ImportError: raise ImportError( "xacro is required for processing .xacro files. Install with: pip install xacro" diff --git a/dimos/mapping/pointclouds/test_occupancy_speed.py b/dimos/mapping/pointclouds/test_occupancy_speed.py index c34c2865f2..2def839dd5 100644 --- a/dimos/mapping/pointclouds/test_occupancy_speed.py +++ b/dimos/mapping/pointclouds/test_occupancy_speed.py @@ -20,7 +20,7 @@ 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 @@ -28,11 +28,10 @@ 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: diff --git a/dimos/memory/embedding.py b/dimos/memory/embedding.py new file mode 100644 index 0000000000..e92a0edb7e --- /dev/null +++ b/dimos/memory/embedding.py @@ -0,0 +1,104 @@ +# 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, Module, ModuleConfig, rpc +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 diff --git a/dimos/memory/test_embedding.py b/dimos/memory/test_embedding.py new file mode 100644 index 0000000000..b7e7fbb294 --- /dev/null +++ b/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() diff --git a/dimos/memory/timeseries/__init__.py b/dimos/memory/timeseries/__init__.py new file mode 100644 index 0000000000..debc14ab3a --- /dev/null +++ b/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", +] diff --git a/dimos/memory/timeseries/base.py b/dimos/memory/timeseries/base.py new file mode 100644 index 0000000000..0d88355b5b --- /dev/null +++ b/dimos/memory/timeseries/base.py @@ -0,0 +1,367 @@ +# 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. +"""Unified time series storage and replay.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +import time +from typing import TYPE_CHECKING, Generic, TypeVar + +import reactivex as rx +from reactivex import operators as ops +from reactivex.disposable import CompositeDisposable, Disposable +from reactivex.scheduler import TimeoutScheduler + +if TYPE_CHECKING: + from collections.abc import Iterator + + from reactivex.observable import Observable + + from dimos.types.timestamped import Timestamped + +T = TypeVar("T", bound="Timestamped") + + +class TimeSeriesStore(Generic[T], ABC): + """Unified storage + replay for sensor data. + + Implement abstract methods for your backend (in-memory, pickle, sqlite, etc.). + All iteration, streaming, and seek logic comes free from the base class. + + T must be a Timestamped subclass — timestamps are taken from .ts attribute. + """ + + @abstractmethod + def _save(self, timestamp: float, data: T) -> None: + """Save data at timestamp.""" + ... + + @abstractmethod + def _load(self, timestamp: float) -> T | None: + """Load data at exact timestamp. Returns None if not found.""" + ... + + @abstractmethod + def _delete(self, timestamp: float) -> T | None: + """Delete data at exact timestamp. Returns the deleted item or None.""" + ... + + @abstractmethod + def _iter_items( + self, start: float | None = None, end: float | None = None + ) -> Iterator[tuple[float, T]]: + """Lazy iteration of (timestamp, data) in range.""" + ... + + @abstractmethod + def _find_closest_timestamp( + self, timestamp: float, tolerance: float | None = None + ) -> float | None: + """Find closest timestamp. Backend can optimize (binary search, db index, etc.).""" + ... + + @abstractmethod + def _count(self) -> int: + """Return number of stored items.""" + ... + + @abstractmethod + def _last_timestamp(self) -> float | None: + """Return the last (largest) timestamp, or None if empty.""" + ... + + @abstractmethod + def _find_before(self, timestamp: float) -> tuple[float, T] | None: + """Find the last (ts, data) strictly before the given timestamp.""" + ... + + @abstractmethod + def _find_after(self, timestamp: float) -> tuple[float, T] | None: + """Find the first (ts, data) strictly after the given timestamp.""" + ... + + # --- Collection API (built on abstract methods) --- + + def __len__(self) -> int: + return self._count() + + def __iter__(self) -> Iterator[T]: + """Iterate over data items in timestamp order.""" + for _, data in self._iter_items(): + yield data + + def last_timestamp(self) -> float | None: + """Get the last timestamp in the store.""" + return self._last_timestamp() + + def last(self) -> T | None: + """Get the last data item in the store.""" + ts = self._last_timestamp() + if ts is None: + return None + return self._load(ts) + + @property + def start_ts(self) -> float | None: + """Get the start timestamp of the store.""" + return self.first_timestamp() + + @property + def end_ts(self) -> float | None: + """Get the end timestamp of the store.""" + return self._last_timestamp() + + def time_range(self) -> tuple[float, float] | None: + """Get the time range (start, end) of the store.""" + s = self.first_timestamp() + e = self._last_timestamp() + if s is None or e is None: + return None + return (s, e) + + def duration(self) -> float: + """Get the duration of the store in seconds.""" + r = self.time_range() + return (r[1] - r[0]) if r else 0.0 + + def find_before(self, timestamp: float) -> T | None: + """Find the last item strictly before the given timestamp.""" + result = self._find_before(timestamp) + return result[1] if result else None + + def find_after(self, timestamp: float) -> T | None: + """Find the first item strictly after the given timestamp.""" + result = self._find_after(timestamp) + return result[1] if result else None + + def slice_by_time(self, start: float, end: float) -> list[T]: + """Return items in [start, end) range.""" + return [data for _, data in self._iter_items(start=start, end=end)] + + def save(self, *data: T) -> None: + """Save one or more Timestamped items.""" + for item in data: + self._save(item.ts, item) + + def pipe_save(self, source: Observable[T]) -> Observable[T]: + """Operator for Observable.pipe() — saves items using .ts. + + Usage: + observable.pipe(store.pipe_save).subscribe(...) + """ + + def _save_and_return(data: T) -> T: + self._save(data.ts, data) + return data + + return source.pipe(ops.map(_save_and_return)) + + def consume_stream(self, observable: Observable[T]) -> rx.abc.DisposableBase: + """Subscribe to an observable and save items using .ts. + + Usage: + disposable = store.consume_stream(observable) + """ + return observable.subscribe(on_next=lambda data: self._save(data.ts, data)) + + def load(self, timestamp: float) -> T | None: + """Load data at exact timestamp.""" + return self._load(timestamp) + + def prune_old(self, cutoff: float) -> None: + """Prune items older than cutoff timestamp.""" + to_delete = [ts for ts, _ in self._iter_items(end=cutoff)] + for ts in to_delete: + self._delete(ts) + + def find_closest( + self, + timestamp: float, + tolerance: float | None = None, + ) -> T | None: + """Find data closest to the given absolute timestamp.""" + closest_ts = self._find_closest_timestamp(timestamp, tolerance) + if closest_ts is None: + return None + return self._load(closest_ts) + + def find_closest_seek( + self, + relative_seconds: float, + tolerance: float | None = None, + ) -> T | None: + """Find data closest to a time relative to the start.""" + first = self.first_timestamp() + if first is None: + return None + return self.find_closest(first + relative_seconds, tolerance) + + def first_timestamp(self) -> float | None: + """Get the first timestamp in the store.""" + for ts, _ in self._iter_items(): + return ts + return None + + def first(self) -> T | None: + """Get the first data item in the store.""" + for _, data in self._iter_items(): + return data + return None + + def iterate_items( + self, + seek: float | None = None, + duration: float | None = None, + from_timestamp: float | None = None, + loop: bool = False, + ) -> Iterator[tuple[float, T]]: + """Iterate over (timestamp, data) tuples with optional seek/duration.""" + first = self.first_timestamp() + if first is None: + return + + if from_timestamp is not None: + start = from_timestamp + elif seek is not None: + start = first + seek + else: + start = None + + end = None + if duration is not None: + start_ts = start if start is not None else first + end = start_ts + duration + + while True: + yield from self._iter_items(start=start, end=end) + if not loop: + break + + def iterate( + self, + seek: float | None = None, + duration: float | None = None, + from_timestamp: float | None = None, + loop: bool = False, + ) -> Iterator[T]: + """Iterate over data items with optional seek/duration.""" + for _, data in self.iterate_items( + seek=seek, duration=duration, from_timestamp=from_timestamp, loop=loop + ): + yield data + + def iterate_realtime( + self, + speed: float = 1.0, + seek: float | None = None, + duration: float | None = None, + from_timestamp: float | None = None, + loop: bool = False, + ) -> Iterator[T]: + """Iterate data, sleeping to match original timing.""" + prev_ts: float | None = None + for ts, data in self.iterate_items( + seek=seek, duration=duration, from_timestamp=from_timestamp, loop=loop + ): + if prev_ts is not None: + delay = (ts - prev_ts) / speed + if delay > 0: + time.sleep(delay) + prev_ts = ts + yield data + + def stream( + self, + speed: float = 1.0, + seek: float | None = None, + duration: float | None = None, + from_timestamp: float | None = None, + loop: bool = False, + ) -> Observable[T]: + """Stream data as Observable with timing control. + + Uses scheduler-based timing with absolute time reference to prevent drift. + """ + + def subscribe( + observer: rx.abc.ObserverBase[T], + scheduler: rx.abc.SchedulerBase | None = None, + ) -> rx.abc.DisposableBase: + sched = scheduler or TimeoutScheduler() + disp = CompositeDisposable() + is_disposed = False + + iterator = self.iterate_items( + seek=seek, duration=duration, from_timestamp=from_timestamp, loop=loop + ) + + try: + first_ts, first_data = next(iterator) + except StopIteration: + observer.on_completed() + return Disposable() + + start_local_time = time.time() + start_replay_time = first_ts + + observer.on_next(first_data) + + try: + next_message: tuple[float, T] | None = next(iterator) + except StopIteration: + observer.on_completed() + return disp + + def schedule_emission(message: tuple[float, T]) -> None: + nonlocal next_message, is_disposed + + if is_disposed: + return + + msg_ts, msg_data = message + + try: + next_message = next(iterator) + except StopIteration: + next_message = None + + target_time = start_local_time + (msg_ts - start_replay_time) / speed + delay = max(0.0, target_time - time.time()) + + def emit( + _scheduler: rx.abc.SchedulerBase, _state: object + ) -> rx.abc.DisposableBase | None: + if is_disposed: + return None + observer.on_next(msg_data) + if next_message is not None: + schedule_emission(next_message) + else: + observer.on_completed() + return None + + sched.schedule_relative(delay, emit) + + if next_message is not None: + schedule_emission(next_message) + + def dispose() -> None: + nonlocal is_disposed + is_disposed = True + disp.dispose() + + return Disposable(dispose) + + return rx.create(subscribe) diff --git a/dimos/memory/timeseries/inmemory.py b/dimos/memory/timeseries/inmemory.py new file mode 100644 index 0000000000..b67faca644 --- /dev/null +++ b/dimos/memory/timeseries/inmemory.py @@ -0,0 +1,119 @@ +# 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. +"""In-memory backend for TimeSeriesStore.""" + +from collections.abc import Iterator + +from sortedcontainers import SortedKeyList # type: ignore[import-untyped] + +from dimos.memory.timeseries.base import T, TimeSeriesStore + + +class InMemoryStore(TimeSeriesStore[T]): + """In-memory storage using SortedKeyList. O(log n) insert, lookup, and range queries.""" + + def __init__(self) -> None: + self._entries: SortedKeyList = SortedKeyList(key=lambda e: e.ts) + + def _bisect_exact(self, timestamp: float) -> int | None: + """Return index of entry with exact timestamp, or None.""" + pos = self._entries.bisect_key_left(timestamp) + if pos < len(self._entries) and self._entries[pos].ts == timestamp: + return pos # type: ignore[no-any-return] + return None + + def _save(self, timestamp: float, data: T) -> None: + self._entries.add(data) + + def _load(self, timestamp: float) -> T | None: + idx = self._bisect_exact(timestamp) + if idx is not None: + return self._entries[idx] # type: ignore[no-any-return] + return None + + def _delete(self, timestamp: float) -> T | None: + idx = self._bisect_exact(timestamp) + if idx is not None: + data = self._entries[idx] + del self._entries[idx] + return data # type: ignore[no-any-return] + return None + + def __iter__(self) -> Iterator[T]: + yield from self._entries + + def _iter_items( + self, start: float | None = None, end: float | None = None + ) -> Iterator[tuple[float, T]]: + if start is not None and end is not None: + it = self._entries.irange_key(start, end, (True, False)) + elif start is not None: + it = self._entries.irange_key(min_key=start) + elif end is not None: + it = self._entries.irange_key(max_key=end, inclusive=(True, False)) + else: + it = iter(self._entries) + for e in it: + yield (e.ts, e) + + def _find_closest_timestamp( + self, timestamp: float, tolerance: float | None = None + ) -> float | None: + if not self._entries: + return None + + pos = self._entries.bisect_key_left(timestamp) + + candidates: list[float] = [] + if pos > 0: + candidates.append(self._entries[pos - 1].ts) + if pos < len(self._entries): + candidates.append(self._entries[pos].ts) + + if not candidates: + return None + + # On ties, prefer the later timestamp (more recent data) + closest = max(candidates, key=lambda ts: (-abs(ts - timestamp), ts)) + + if tolerance is not None and abs(closest - timestamp) > tolerance: + return None + + return closest + + def _count(self) -> int: + return len(self._entries) + + def _last_timestamp(self) -> float | None: + if not self._entries: + return None + return self._entries[-1].ts # type: ignore[no-any-return] + + def _find_before(self, timestamp: float) -> tuple[float, T] | None: + if not self._entries: + return None + pos = self._entries.bisect_key_left(timestamp) + if pos > 0: + e = self._entries[pos - 1] + return (e.ts, e) + return None + + def _find_after(self, timestamp: float) -> tuple[float, T] | None: + if not self._entries: + return None + pos = self._entries.bisect_key_right(timestamp) + if pos < len(self._entries): + e = self._entries[pos] + return (e.ts, e) + return None diff --git a/dimos/memory/timeseries/legacy.py b/dimos/memory/timeseries/legacy.py new file mode 100644 index 0000000000..821d306d2d --- /dev/null +++ b/dimos/memory/timeseries/legacy.py @@ -0,0 +1,398 @@ +# 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. +"""Legacy pickle directory backend for TimeSeriesStore. + +Compatible with TimedSensorReplay/TimedSensorStorage file format. +""" + +from collections.abc import Callable, Iterator +import glob +import os +from pathlib import Path +import pickle +import re +import time +from typing import Any, cast + +import reactivex as rx +from reactivex.disposable import CompositeDisposable, Disposable +from reactivex.observable import Observable +from reactivex.scheduler import TimeoutScheduler + +from dimos.memory.timeseries.base import T, TimeSeriesStore +from dimos.utils.data import get_data, get_data_dir + + +class LegacyPickleStore(TimeSeriesStore[T]): + """Legacy pickle backend compatible with TimedSensorReplay/TimedSensorStorage. + + File format: + {name}/ + 000.pickle # contains (timestamp, data) tuple + 001.pickle + ... + + Files are assumed to be in chronological order (timestamps increase with file number). + No index is built - iteration is lazy and memory-efficient for large datasets. + + Usage: + # Load existing recording (auto-downloads from LFS if needed) + store = LegacyPickleStore("unitree_go2_bigoffice/lidar") + data = store.find_closest_seek(10.0) + + # Create new recording (directory created on first save) + store = LegacyPickleStore("my_recording/images") + store.save_ts(image) # uses image.ts for timestamp + + Backward compatibility: + This class also supports the old TimedSensorReplay/SensorReplay API: + - iterate_ts() - iterate returning (timestamp, data) tuples + - files - property returning list of file paths + - load_one() - load a single pickle file + """ + + def __init__(self, name: str | Path, autocast: Callable[[Any], T] | None = None) -> None: + """ + Args: + name: Data directory name (e.g. "unitree_go2_bigoffice/lidar") or absolute path. + autocast: Optional function to transform data after loading (for replay) or + before saving (for storage). E.g., `Odometry.from_msg`. + """ + self._name = str(name) + self._root_dir: Path | None = None + self._counter: int = 0 + self._autocast = autocast + + def _get_root_dir(self, for_write: bool = False) -> Path: + """Get root directory, creating on first write if needed.""" + if self._root_dir is not None: + # Ensure directory exists if writing + if for_write: + self._root_dir.mkdir(parents=True, exist_ok=True) + return self._root_dir + + # If absolute path, use directly + if Path(self._name).is_absolute(): + self._root_dir = Path(self._name) + if for_write: + self._root_dir.mkdir(parents=True, exist_ok=True) + elif for_write: + # For writing: use get_data_dir and create if needed + self._root_dir = get_data_dir(self._name) + self._root_dir.mkdir(parents=True, exist_ok=True) + else: + # For reading: use get_data (handles LFS download) + self._root_dir = get_data(self._name) + + return self._root_dir + + def _iter_files(self) -> Iterator[Path]: + """Iterate pickle files in sorted order (by number in filename).""" + + def extract_number(filepath: str) -> int: + basename = os.path.basename(filepath) + match = re.search(r"(\d+)\.pickle$", basename) + return int(match.group(1)) if match else 0 + + root_dir = self._get_root_dir() + files = sorted( + glob.glob(os.path.join(root_dir, "*.pickle")), + key=extract_number, + ) + for f in files: + yield Path(f) + + def _save(self, timestamp: float, data: T) -> None: + root_dir = self._get_root_dir(for_write=True) + + # Initialize counter from existing files if needed + if self._counter == 0: + existing = list(root_dir.glob("*.pickle")) + if existing: + # Find highest existing counter + max_num = 0 + for filepath in existing: + match = re.search(r"(\d+)\.pickle$", filepath.name) + if match: + max_num = max(max_num, int(match.group(1))) + self._counter = max_num + 1 + + full_path = root_dir / f"{self._counter:03d}.pickle" + + if full_path.exists(): + raise RuntimeError(f"File {full_path} already exists") + + # Save as (timestamp, data) tuple for legacy compatibility + with open(full_path, "wb") as f: + pickle.dump((timestamp, data), f) + + self._counter += 1 + + def _load(self, timestamp: float) -> T | None: + """Load data at exact timestamp (linear scan).""" + for ts, data in self._iter_items(): + if ts == timestamp: + return data + return None + + def _delete(self, timestamp: float) -> T | None: + """Delete not supported for legacy pickle format.""" + raise NotImplementedError("LegacyPickleStore does not support deletion") + + def _iter_items( + self, start: float | None = None, end: float | None = None + ) -> Iterator[tuple[float, T]]: + """Lazy iteration - loads one file at a time. + + Handles both timed format (timestamp, data) and non-timed format (just data). + For non-timed data, uses file index as synthetic timestamp. + """ + for idx, filepath in enumerate(self._iter_files()): + try: + with open(filepath, "rb") as f: + raw = pickle.load(f) + + # Handle both timed (timestamp, data) and non-timed (just data) formats + if isinstance(raw, tuple) and len(raw) == 2: + ts, data = raw + ts = float(ts) + else: + # Non-timed format: use index as synthetic timestamp + ts = float(idx) + data = raw + except Exception: + continue + + if start is not None and ts < start: + continue + if end is not None and ts >= end: + break + + if self._autocast is not None: + data = self._autocast(data) + yield (ts, cast("T", data)) + + def _find_closest_timestamp( + self, timestamp: float, tolerance: float | None = None + ) -> float | None: + """Linear scan with early exit (assumes timestamps are monotonically increasing).""" + closest_ts: float | None = None + closest_diff = float("inf") + + for ts, _ in self._iter_items(): + diff = abs(ts - timestamp) + + if diff < closest_diff: + closest_diff = diff + closest_ts = ts + elif diff > closest_diff: + # Moving away from target, can stop + break + + if closest_ts is None: + return None + + if tolerance is not None and closest_diff > tolerance: + return None + + return closest_ts + + def _count(self) -> int: + return sum(1 for _ in self._iter_files()) + + def _last_timestamp(self) -> float | None: + last_ts: float | None = None + for ts, _ in self._iter_items(): + last_ts = ts + return last_ts + + def _find_before(self, timestamp: float) -> tuple[float, T] | None: + result: tuple[float, T] | None = None + for ts, data in self._iter_items(): + if ts < timestamp: + result = (ts, data) + else: + break + return result + + def _find_after(self, timestamp: float) -> tuple[float, T] | None: + for ts, data in self._iter_items(): + if ts > timestamp: + return (ts, data) + return None + + # === Backward-compatible API (TimedSensorReplay/SensorReplay) === + + @property + def files(self) -> list[Path]: + """Return list of pickle files (backward compatibility with SensorReplay).""" + return list(self._iter_files()) + + def load_one(self, name: int | str | Path) -> T | Any: + """Load a single pickle file (backward compatibility with SensorReplay). + + Args: + name: File index (int), filename without extension (str), or full path (Path) + + Returns: + For TimedSensorReplay: (timestamp, data) tuple + For SensorReplay: just the data + """ + root_dir = self._get_root_dir() + + if isinstance(name, int): + full_path = root_dir / f"{name:03d}.pickle" + elif isinstance(name, Path): + full_path = name + else: + full_path = root_dir / Path(f"{name}.pickle") + + with open(full_path, "rb") as f: + data = pickle.load(f) + + # Legacy format: (timestamp, data) tuple + if isinstance(data, tuple) and len(data) == 2: + ts, payload = data + if self._autocast is not None: + payload = self._autocast(payload) + return (ts, payload) + + # Non-timed format: just data + if self._autocast is not None: + data = self._autocast(data) + return data + + def iterate_ts( + self, + seek: float | None = None, + duration: float | None = None, + from_timestamp: float | None = None, + loop: bool = False, + ) -> Iterator[tuple[float, T]]: + """Iterate with timestamps (backward compatibility with TimedSensorReplay). + + Args: + seek: Relative seconds from start + duration: Duration window in seconds + from_timestamp: Absolute timestamp to start from + loop: Whether to loop the data + + Yields: + (timestamp, data) tuples + """ + first = self.first_timestamp() + if first is None: + return + + # Calculate start timestamp + start: float | None = None + if from_timestamp is not None: + start = from_timestamp + elif seek is not None: + start = first + seek + + # Calculate end timestamp + end: float | None = None + if duration is not None: + start_ts = start if start is not None else first + end = start_ts + duration + + while True: + yield from self._iter_items(start=start, end=end) + if not loop: + break + + def stream( + self, + speed: float = 1.0, + seek: float | None = None, + duration: float | None = None, + from_timestamp: float | None = None, + loop: bool = False, + ) -> Observable[T]: + """Stream data as Observable with timing control. + + Uses stored timestamps from pickle files for timing (not data.ts). + """ + + def subscribe( + observer: rx.abc.ObserverBase[T], + scheduler: rx.abc.SchedulerBase | None = None, + ) -> rx.abc.DisposableBase: + sched = scheduler or TimeoutScheduler() + disp = CompositeDisposable() + is_disposed = False + + iterator = self.iterate_ts( + seek=seek, duration=duration, from_timestamp=from_timestamp, loop=loop + ) + + try: + first_ts, first_data = next(iterator) + except StopIteration: + observer.on_completed() + return Disposable() + + start_local_time = time.time() + start_replay_time = first_ts + + observer.on_next(first_data) + + try: + next_message: tuple[float, T] | None = next(iterator) + except StopIteration: + observer.on_completed() + return disp + + def schedule_emission(message: tuple[float, T]) -> None: + nonlocal next_message, is_disposed + + if is_disposed: + return + + ts, data = message + + try: + next_message = next(iterator) + except StopIteration: + next_message = None + + target_time = start_local_time + (ts - start_replay_time) / speed + delay = max(0.0, target_time - time.time()) + + def emit( + _scheduler: rx.abc.SchedulerBase, _state: object + ) -> rx.abc.DisposableBase | None: + if is_disposed: + return None + observer.on_next(data) + if next_message is not None: + schedule_emission(next_message) + else: + observer.on_completed() + return None + + sched.schedule_relative(delay, emit) + + if next_message is not None: + schedule_emission(next_message) + + def dispose() -> None: + nonlocal is_disposed + is_disposed = True + disp.dispose() + + return Disposable(dispose) + + return rx.create(subscribe) diff --git a/dimos/memory/timeseries/pickledir.py b/dimos/memory/timeseries/pickledir.py new file mode 100644 index 0000000000..9e8cd5a249 --- /dev/null +++ b/dimos/memory/timeseries/pickledir.py @@ -0,0 +1,198 @@ +# 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. +"""Pickle directory backend for TimeSeriesStore.""" + +import bisect +from collections.abc import Iterator +import glob +import os +from pathlib import Path +import pickle + +from dimos.memory.timeseries.base import T, TimeSeriesStore +from dimos.utils.data import get_data, get_data_dir + + +class PickleDirStore(TimeSeriesStore[T]): + """Pickle directory backend. Files named by timestamp. + + Directory structure: + {name}/ + 1704067200.123.pickle + 1704067200.456.pickle + ... + + Usage: + # Load existing recording (auto-downloads from LFS if needed) + store = PickleDirStore("unitree_go2_bigoffice/lidar") + data = store.find_closest_seek(10.0) + + # Create new recording (directory created on first save) + store = PickleDirStore("my_recording/images") + store.save(image) # saves using image.ts + """ + + def __init__(self, name: str) -> None: + """ + Args: + name: Data directory name (e.g. "unitree_go2_bigoffice/lidar") + """ + self._name = name + self._root_dir: Path | None = None + + # Cached sorted timestamps for find_closest + self._timestamps: list[float] | None = None + + def _get_root_dir(self, for_write: bool = False) -> Path: + """Get root directory, creating on first write if needed.""" + if self._root_dir is not None: + return self._root_dir + + # If absolute path, use directly + if Path(self._name).is_absolute(): + self._root_dir = Path(self._name) + if for_write: + self._root_dir.mkdir(parents=True, exist_ok=True) + elif for_write: + # For writing: use get_data_dir and create if needed + self._root_dir = get_data_dir(self._name) + self._root_dir.mkdir(parents=True, exist_ok=True) + else: + # For reading: use get_data (handles LFS download) + self._root_dir = get_data(self._name) + + return self._root_dir + + def _save(self, timestamp: float, data: T) -> None: + root_dir = self._get_root_dir(for_write=True) + full_path = root_dir / f"{timestamp}.pickle" + + if full_path.exists(): + raise RuntimeError(f"File {full_path} already exists") + + with open(full_path, "wb") as f: + pickle.dump(data, f) + + self._timestamps = None # Invalidate cache + + def _load(self, timestamp: float) -> T | None: + filepath = self._get_root_dir() / f"{timestamp}.pickle" + if filepath.exists(): + return self._load_file(filepath) + return None + + def _delete(self, timestamp: float) -> T | None: + filepath = self._get_root_dir() / f"{timestamp}.pickle" + if filepath.exists(): + data = self._load_file(filepath) + filepath.unlink() + self._timestamps = None # Invalidate cache + return data + return None + + def _iter_items( + self, start: float | None = None, end: float | None = None + ) -> Iterator[tuple[float, T]]: + for ts in self._get_timestamps(): + if start is not None and ts < start: + continue + if end is not None and ts >= end: + break + data = self._load(ts) + if data is not None: + yield (ts, data) + + def _find_closest_timestamp( + self, timestamp: float, tolerance: float | None = None + ) -> float | None: + timestamps = self._get_timestamps() + if not timestamps: + return None + + pos = bisect.bisect_left(timestamps, timestamp) + + # Check neighbors + candidates = [] + if pos > 0: + candidates.append(timestamps[pos - 1]) + if pos < len(timestamps): + candidates.append(timestamps[pos]) + + if not candidates: + return None + + closest = min(candidates, key=lambda ts: abs(ts - timestamp)) + + if tolerance is not None and abs(closest - timestamp) > tolerance: + return None + + return closest + + def _get_timestamps(self) -> list[float]: + """Get sorted list of all timestamps.""" + if self._timestamps is not None: + return self._timestamps + + timestamps: list[float] = [] + root_dir = self._get_root_dir() + for filepath in glob.glob(os.path.join(root_dir, "*.pickle")): + try: + ts = float(Path(filepath).stem) + timestamps.append(ts) + except ValueError: + continue + + timestamps.sort() + self._timestamps = timestamps + return timestamps + + def _count(self) -> int: + return len(self._get_timestamps()) + + def _last_timestamp(self) -> float | None: + timestamps = self._get_timestamps() + return timestamps[-1] if timestamps else None + + def _find_before(self, timestamp: float) -> tuple[float, T] | None: + timestamps = self._get_timestamps() + if not timestamps: + return None + pos = bisect.bisect_left(timestamps, timestamp) + if pos > 0: + ts = timestamps[pos - 1] + data = self._load(ts) + if data is not None: + return (ts, data) + return None + + def _find_after(self, timestamp: float) -> tuple[float, T] | None: + timestamps = self._get_timestamps() + if not timestamps: + return None + pos = bisect.bisect_right(timestamps, timestamp) + if pos < len(timestamps): + ts = timestamps[pos] + data = self._load(ts) + if data is not None: + return (ts, data) + return None + + def _load_file(self, filepath: Path) -> T | None: + """Load data from a pickle file (LRU cached).""" + try: + with open(filepath, "rb") as f: + data: T = pickle.load(f) + return data + except Exception: + return None diff --git a/dimos/memory/timeseries/postgres.py b/dimos/memory/timeseries/postgres.py new file mode 100644 index 0000000000..a6a853bbd3 --- /dev/null +++ b/dimos/memory/timeseries/postgres.py @@ -0,0 +1,312 @@ +# 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. +"""PostgreSQL backend for TimeSeriesStore.""" + +from collections.abc import Iterator +import pickle +import re + +import psycopg2 # type: ignore[import-not-found] +import psycopg2.extensions # type: ignore[import-not-found] + +from dimos.core.resource import Resource +from dimos.memory.timeseries.base import T, TimeSeriesStore + +# Valid SQL identifier: alphanumeric and underscores, not starting with digit +_VALID_IDENTIFIER = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_]*$") + + +def _validate_identifier(name: str) -> str: + """Validate SQL identifier to prevent injection.""" + if not _VALID_IDENTIFIER.match(name): + raise ValueError( + f"Invalid identifier '{name}': must be alphanumeric/underscore, not start with digit" + ) + if len(name) > 128: + raise ValueError(f"Identifier too long: {len(name)} > 128") + return name + + +class PostgresStore(TimeSeriesStore[T], Resource): + """PostgreSQL backend for sensor data. + + Multiple stores can share the same database with different tables. + Implements Resource for lifecycle management (start/stop/dispose). + + Usage: + # Create store + store = PostgresStore("lidar") + store.start() # open connection + + # Use store + store.save(data) # saves using data.ts + data = store.find_closest_seek(10.0) + + # Cleanup + store.stop() # close connection + + # Multiple sensors in same db + lidar = PostgresStore("lidar") + images = PostgresStore("images") + + # Manual run management via table naming + run1_lidar = PostgresStore("run1_lidar") + """ + + def __init__( + self, + table: str, + db: str = "dimensional", + host: str = "localhost", + port: int = 5432, + user: str | None = None, + ) -> None: + """ + Args: + table: Table name for this sensor's data (alphanumeric/underscore only). + db: Database name (alphanumeric/underscore only). + host: PostgreSQL host. + port: PostgreSQL port. + user: PostgreSQL user. Defaults to current system user. + """ + self._table = _validate_identifier(table) + self._db = _validate_identifier(db) + self._host = host + self._port = port + self._user = user + self._conn: psycopg2.extensions.connection | None = None + self._table_created = False + + def start(self) -> None: + """Open database connection.""" + if self._conn is not None: + return + self._conn = psycopg2.connect( + dbname=self._db, + host=self._host, + port=self._port, + user=self._user, + ) + + def stop(self) -> None: + """Close database connection.""" + if self._conn is not None: + self._conn.close() + self._conn = None + + def _get_conn(self) -> psycopg2.extensions.connection: + """Get connection, starting if needed.""" + if self._conn is None: + self.start() + assert self._conn is not None + return self._conn + + def _ensure_table(self) -> None: + """Create table if it doesn't exist.""" + if self._table_created: + return + conn = self._get_conn() + with conn.cursor() as cur: + cur.execute(f""" + CREATE TABLE IF NOT EXISTS {self._table} ( + timestamp DOUBLE PRECISION PRIMARY KEY, + data BYTEA NOT NULL + ) + """) + cur.execute(f""" + CREATE INDEX IF NOT EXISTS idx_{self._table}_ts + ON {self._table}(timestamp) + """) + conn.commit() + self._table_created = True + + def _save(self, timestamp: float, data: T) -> None: + self._ensure_table() + conn = self._get_conn() + blob = pickle.dumps(data) + with conn.cursor() as cur: + cur.execute( + f""" + INSERT INTO {self._table} (timestamp, data) VALUES (%s, %s) + ON CONFLICT (timestamp) DO UPDATE SET data = EXCLUDED.data + """, + (timestamp, psycopg2.Binary(blob)), + ) + conn.commit() + + def _load(self, timestamp: float) -> T | None: + self._ensure_table() + conn = self._get_conn() + with conn.cursor() as cur: + cur.execute(f"SELECT data FROM {self._table} WHERE timestamp = %s", (timestamp,)) + row = cur.fetchone() + if row is None: + return None + data: T = pickle.loads(row[0]) + return data + + def _delete(self, timestamp: float) -> T | None: + data = self._load(timestamp) + if data is not None: + conn = self._get_conn() + with conn.cursor() as cur: + cur.execute(f"DELETE FROM {self._table} WHERE timestamp = %s", (timestamp,)) + conn.commit() + return data + + def _iter_items( + self, start: float | None = None, end: float | None = None + ) -> Iterator[tuple[float, T]]: + self._ensure_table() + conn = self._get_conn() + + query = f"SELECT timestamp, data FROM {self._table}" + params: list[float] = [] + conditions = [] + + if start is not None: + conditions.append("timestamp >= %s") + params.append(start) + if end is not None: + conditions.append("timestamp < %s") + params.append(end) + + if conditions: + query += " WHERE " + " AND ".join(conditions) + query += " ORDER BY timestamp" + + with conn.cursor() as cur: + cur.execute(query, params) + for row in cur: + ts: float = row[0] + data: T = pickle.loads(row[1]) + yield (ts, data) + + def _find_closest_timestamp( + self, timestamp: float, tolerance: float | None = None + ) -> float | None: + self._ensure_table() + conn = self._get_conn() + + with conn.cursor() as cur: + # Get closest timestamp <= target + cur.execute( + f""" + SELECT timestamp FROM {self._table} + WHERE timestamp <= %s + ORDER BY timestamp DESC LIMIT 1 + """, + (timestamp,), + ) + before = cur.fetchone() + + # Get closest timestamp >= target + cur.execute( + f""" + SELECT timestamp FROM {self._table} + WHERE timestamp >= %s + ORDER BY timestamp ASC LIMIT 1 + """, + (timestamp,), + ) + after = cur.fetchone() + + candidates: list[float] = [] + if before: + candidates.append(before[0]) + if after: + candidates.append(after[0]) + + if not candidates: + return None + + closest = min(candidates, key=lambda ts: abs(ts - timestamp)) + + if tolerance is not None and abs(closest - timestamp) > tolerance: + return None + + return closest + + def _count(self) -> int: + self._ensure_table() + conn = self._get_conn() + with conn.cursor() as cur: + cur.execute(f"SELECT COUNT(*) FROM {self._table}") + row = cur.fetchone() + return row[0] if row else 0 # type: ignore[no-any-return] + + def _last_timestamp(self) -> float | None: + self._ensure_table() + conn = self._get_conn() + with conn.cursor() as cur: + cur.execute(f"SELECT MAX(timestamp) FROM {self._table}") + row = cur.fetchone() + if row is None or row[0] is None: + return None + return row[0] # type: ignore[no-any-return] + + def _find_before(self, timestamp: float) -> tuple[float, T] | None: + self._ensure_table() + conn = self._get_conn() + with conn.cursor() as cur: + cur.execute( + f"SELECT timestamp, data FROM {self._table} WHERE timestamp < %s ORDER BY timestamp DESC LIMIT 1", + (timestamp,), + ) + row = cur.fetchone() + if row is None: + return None + return (row[0], pickle.loads(row[1])) + + def _find_after(self, timestamp: float) -> tuple[float, T] | None: + self._ensure_table() + conn = self._get_conn() + with conn.cursor() as cur: + cur.execute( + f"SELECT timestamp, data FROM {self._table} WHERE timestamp > %s ORDER BY timestamp ASC LIMIT 1", + (timestamp,), + ) + row = cur.fetchone() + if row is None: + return None + return (row[0], pickle.loads(row[1])) + + +def reset_db(db: str = "dimensional", host: str = "localhost", port: int = 5432) -> None: + """Drop and recreate database. Simple migration strategy. + + WARNING: This deletes all data in the database! + + Args: + db: Database name to reset (alphanumeric/underscore only). + host: PostgreSQL host. + port: PostgreSQL port. + """ + db = _validate_identifier(db) + # Connect to 'postgres' database to drop/create + conn = psycopg2.connect(dbname="postgres", host=host, port=port) + conn.autocommit = True + with conn.cursor() as cur: + # Terminate existing connections + cur.execute( + """ + SELECT pg_terminate_backend(pid) + FROM pg_stat_activity + WHERE datname = %s AND pid <> pg_backend_pid() + """, + (db,), + ) + cur.execute(f"DROP DATABASE IF EXISTS {db}") + cur.execute(f"CREATE DATABASE {db}") + conn.close() diff --git a/dimos/memory/timeseries/sqlite.py b/dimos/memory/timeseries/sqlite.py new file mode 100644 index 0000000000..6e2ac7a7f5 --- /dev/null +++ b/dimos/memory/timeseries/sqlite.py @@ -0,0 +1,268 @@ +# 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. +"""SQLite backend for TimeSeriesStore.""" + +from collections.abc import Iterator +from pathlib import Path +import pickle +import re +import sqlite3 + +from dimos.memory.timeseries.base import T, TimeSeriesStore +from dimos.utils.data import get_data, get_data_dir + +# Valid SQL identifier: alphanumeric and underscores, not starting with digit +_VALID_IDENTIFIER = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_]*$") + + +def _validate_identifier(name: str) -> str: + """Validate SQL identifier to prevent injection.""" + if not _VALID_IDENTIFIER.match(name): + raise ValueError( + f"Invalid identifier '{name}': must be alphanumeric/underscore, not start with digit" + ) + if len(name) > 128: + raise ValueError(f"Identifier too long: {len(name)} > 128") + return name + + +class SqliteStore(TimeSeriesStore[T]): + """SQLite backend for sensor data. Good for indexed queries and single-file storage. + + Data is stored as pickled BLOBs with timestamp as indexed column. + + Usage: + # Named store (uses data/ directory, auto-downloads from LFS if needed) + store = SqliteStore("recordings/lidar") # -> data/recordings/lidar.db + store.save(data) # saves using data.ts + + # Absolute path + store = SqliteStore("/path/to/sensors.db") + + # In-memory (for testing) + store = SqliteStore(":memory:") + + # Multiple tables in one DB + store = SqliteStore("recordings/sensors", table="lidar") + """ + + def __init__(self, name: str | Path, table: str = "sensor_data") -> None: + """ + Args: + name: Data name (e.g. "recordings/lidar") resolved via get_data, + absolute path, or ":memory:" for in-memory. + table: Table name for this sensor's data (alphanumeric/underscore only). + """ + self._name = str(name) + self._table = _validate_identifier(table) + self._db_path: str | None = None + self._conn: sqlite3.Connection | None = None + + def _get_db_path(self, for_write: bool = False) -> str: + """Get database path, resolving via get_data if needed.""" + if self._db_path is not None: + return self._db_path + + # Special case for in-memory + if self._name == ":memory:": + self._db_path = ":memory:" + return self._db_path + + # If absolute path, use directly + if Path(self._name).is_absolute(): + self._db_path = self._name + elif for_write: + # For writing: use get_data_dir + db_file = get_data_dir(self._name + ".db") + db_file.parent.mkdir(parents=True, exist_ok=True) + self._db_path = str(db_file) + else: + # For reading: use get_data (handles LFS download) + # Try with .db extension first + try: + db_file = get_data(self._name + ".db") + self._db_path = str(db_file) + except FileNotFoundError: + # Fall back to get_data_dir for new databases + db_file = get_data_dir(self._name + ".db") + db_file.parent.mkdir(parents=True, exist_ok=True) + self._db_path = str(db_file) + + return self._db_path + + def _get_conn(self) -> sqlite3.Connection: + """Get or create database connection.""" + if self._conn is None: + db_path = self._get_db_path(for_write=True) + self._conn = sqlite3.connect(db_path, check_same_thread=False) + self._create_table() + return self._conn + + def _create_table(self) -> None: + """Create table if it doesn't exist.""" + conn = self._conn + assert conn is not None + conn.execute(f""" + CREATE TABLE IF NOT EXISTS {self._table} ( + timestamp REAL PRIMARY KEY, + data BLOB NOT NULL + ) + """) + conn.execute(f""" + CREATE INDEX IF NOT EXISTS idx_{self._table}_timestamp + ON {self._table}(timestamp) + """) + conn.commit() + + def _save(self, timestamp: float, data: T) -> None: + conn = self._get_conn() + blob = pickle.dumps(data) + conn.execute( + f"INSERT OR REPLACE INTO {self._table} (timestamp, data) VALUES (?, ?)", + (timestamp, blob), + ) + conn.commit() + + def _load(self, timestamp: float) -> T | None: + conn = self._get_conn() + cursor = conn.execute(f"SELECT data FROM {self._table} WHERE timestamp = ?", (timestamp,)) + row = cursor.fetchone() + if row is None: + return None + data: T = pickle.loads(row[0]) + return data + + def _delete(self, timestamp: float) -> T | None: + data = self._load(timestamp) + if data is not None: + conn = self._get_conn() + conn.execute(f"DELETE FROM {self._table} WHERE timestamp = ?", (timestamp,)) + conn.commit() + return data + + def _iter_items( + self, start: float | None = None, end: float | None = None + ) -> Iterator[tuple[float, T]]: + conn = self._get_conn() + + # Build query with optional range filters + query = f"SELECT timestamp, data FROM {self._table}" + params: list[float] = [] + conditions = [] + + if start is not None: + conditions.append("timestamp >= ?") + params.append(start) + if end is not None: + conditions.append("timestamp < ?") + params.append(end) + + if conditions: + query += " WHERE " + " AND ".join(conditions) + query += " ORDER BY timestamp" + + cursor = conn.execute(query, params) + for row in cursor: + ts: float = row[0] + data: T = pickle.loads(row[1]) + yield (ts, data) + + def _find_closest_timestamp( + self, timestamp: float, tolerance: float | None = None + ) -> float | None: + conn = self._get_conn() + + # Find closest timestamp using SQL + # Get the closest timestamp <= target + cursor = conn.execute( + f""" + SELECT timestamp FROM {self._table} + WHERE timestamp <= ? + ORDER BY timestamp DESC LIMIT 1 + """, + (timestamp,), + ) + before = cursor.fetchone() + + # Get the closest timestamp >= target + cursor = conn.execute( + f""" + SELECT timestamp FROM {self._table} + WHERE timestamp >= ? + ORDER BY timestamp ASC LIMIT 1 + """, + (timestamp,), + ) + after = cursor.fetchone() + + # Find the closest of the two + candidates: list[float] = [] + if before: + candidates.append(before[0]) + if after: + candidates.append(after[0]) + + if not candidates: + return None + + closest = min(candidates, key=lambda ts: abs(ts - timestamp)) + + if tolerance is not None and abs(closest - timestamp) > tolerance: + return None + + return closest + + def _count(self) -> int: + conn = self._get_conn() + cursor = conn.execute(f"SELECT COUNT(*) FROM {self._table}") + return cursor.fetchone()[0] # type: ignore[no-any-return] + + def _last_timestamp(self) -> float | None: + conn = self._get_conn() + cursor = conn.execute(f"SELECT MAX(timestamp) FROM {self._table}") + row = cursor.fetchone() + if row is None or row[0] is None: + return None + return row[0] # type: ignore[no-any-return] + + def _find_before(self, timestamp: float) -> tuple[float, T] | None: + conn = self._get_conn() + cursor = conn.execute( + f"SELECT timestamp, data FROM {self._table} WHERE timestamp < ? ORDER BY timestamp DESC LIMIT 1", + (timestamp,), + ) + row = cursor.fetchone() + if row is None: + return None + return (row[0], pickle.loads(row[1])) + + def _find_after(self, timestamp: float) -> tuple[float, T] | None: + conn = self._get_conn() + cursor = conn.execute( + f"SELECT timestamp, data FROM {self._table} WHERE timestamp > ? ORDER BY timestamp ASC LIMIT 1", + (timestamp,), + ) + row = cursor.fetchone() + if row is None: + return None + return (row[0], pickle.loads(row[1])) + + def close(self) -> None: + """Close the database connection.""" + if self._conn is not None: + self._conn.close() + self._conn = None + + def __del__(self) -> None: + self.close() diff --git a/dimos/memory/timeseries/test_base.py b/dimos/memory/timeseries/test_base.py new file mode 100644 index 0000000000..9491d2c93c --- /dev/null +++ b/dimos/memory/timeseries/test_base.py @@ -0,0 +1,468 @@ +# 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. +"""Tests for TimeSeriesStore implementations.""" + +from dataclasses import dataclass +from pathlib import Path +import tempfile +import uuid + +import pytest + +from dimos.memory.timeseries.base import TimeSeriesStore +from dimos.memory.timeseries.inmemory import InMemoryStore +from dimos.memory.timeseries.legacy import LegacyPickleStore +from dimos.memory.timeseries.pickledir import PickleDirStore +from dimos.memory.timeseries.sqlite import SqliteStore +from dimos.types.timestamped import Timestamped + + +@dataclass +class SampleData(Timestamped): + """Simple timestamped data for testing.""" + + value: str + + def __init__(self, value: str, ts: float) -> None: + super().__init__(ts) + self.value = value + + def __eq__(self, other: object) -> bool: + if isinstance(other, SampleData): + return self.value == other.value and self.ts == other.ts + return False + + +@pytest.fixture +def temp_dir(): + """Create a temporary directory for file-based store tests.""" + with tempfile.TemporaryDirectory() as tmpdir: + yield tmpdir + + +def make_in_memory_store() -> TimeSeriesStore[SampleData]: + return InMemoryStore[SampleData]() + + +def make_pickle_dir_store(tmpdir: str) -> TimeSeriesStore[SampleData]: + return PickleDirStore[SampleData](tmpdir) + + +def make_sqlite_store(tmpdir: str) -> TimeSeriesStore[SampleData]: + return SqliteStore[SampleData](Path(tmpdir) / "test.db") + + +def make_legacy_pickle_store(tmpdir: str) -> TimeSeriesStore[SampleData]: + return LegacyPickleStore[SampleData](Path(tmpdir) / "legacy") + + +# Base test data (always available) +testdata: list[tuple[object, str]] = [ + (lambda _: make_in_memory_store(), "InMemoryStore"), + (lambda tmpdir: make_pickle_dir_store(tmpdir), "PickleDirStore"), + (lambda tmpdir: make_sqlite_store(tmpdir), "SqliteStore"), + (lambda tmpdir: make_legacy_pickle_store(tmpdir), "LegacyPickleStore"), +] + +# Track postgres tables to clean up +_postgres_tables: list[str] = [] + +try: + import psycopg2 + + from dimos.memory.timeseries.postgres import PostgresStore + + # Test connection + _test_conn = psycopg2.connect(dbname="dimensional") + _test_conn.close() + + def make_postgres_store(_tmpdir: str) -> TimeSeriesStore[SampleData]: + """Create PostgresStore with unique table name.""" + table = f"test_{uuid.uuid4().hex[:8]}" + _postgres_tables.append(table) + store = PostgresStore[SampleData](table) + store.start() + return store + + testdata.append((lambda tmpdir: make_postgres_store(tmpdir), "PostgresStore")) + + @pytest.fixture(autouse=True) + def cleanup_postgres_tables(): + """Clean up postgres test tables after each test.""" + yield + if _postgres_tables: + try: + conn = psycopg2.connect(dbname="dimensional") + conn.autocommit = True + with conn.cursor() as cur: + for table in _postgres_tables: + cur.execute(f"DROP TABLE IF EXISTS {table}") + conn.close() + except Exception: + pass # Ignore cleanup errors + _postgres_tables.clear() + +except Exception: + print("PostgreSQL not available") + + +@pytest.mark.parametrize("store_factory,store_name", testdata) +class TestTimeSeriesStore: + """Parametrized tests for all TimeSeriesStore implementations.""" + + def test_save_and_load(self, store_factory, store_name, temp_dir): + store = store_factory(temp_dir) + store.save(SampleData("data_at_1", 1.0)) + store.save(SampleData("data_at_2", 2.0)) + + assert store.load(1.0) == SampleData("data_at_1", 1.0) + assert store.load(2.0) == SampleData("data_at_2", 2.0) + assert store.load(3.0) is None + + def test_find_closest_timestamp(self, store_factory, store_name, temp_dir): + store = store_factory(temp_dir) + store.save(SampleData("a", 1.0), SampleData("b", 2.0), SampleData("c", 3.0)) + + # Exact match + assert store._find_closest_timestamp(2.0) == 2.0 + + # Closest to 1.4 is 1.0 + assert store._find_closest_timestamp(1.4) == 1.0 + + # Closest to 1.6 is 2.0 + assert store._find_closest_timestamp(1.6) == 2.0 + + # With tolerance + assert store._find_closest_timestamp(1.4, tolerance=0.5) == 1.0 + assert store._find_closest_timestamp(1.4, tolerance=0.3) is None + + def test_iter_items(self, store_factory, store_name, temp_dir): + store = store_factory(temp_dir) + store.save(SampleData("a", 1.0), SampleData("b", 2.0), SampleData("c", 3.0)) + + # Should iterate in timestamp order + items = list(store._iter_items()) + assert items == [ + (1.0, SampleData("a", 1.0)), + (2.0, SampleData("b", 2.0)), + (3.0, SampleData("c", 3.0)), + ] + + def test_iter_items_with_range(self, store_factory, store_name, temp_dir): + store = store_factory(temp_dir) + store.save( + SampleData("a", 1.0), + SampleData("b", 2.0), + SampleData("c", 3.0), + SampleData("d", 4.0), + ) + + # Start only + items = list(store._iter_items(start=2.0)) + assert items == [ + (2.0, SampleData("b", 2.0)), + (3.0, SampleData("c", 3.0)), + (4.0, SampleData("d", 4.0)), + ] + + # End only + items = list(store._iter_items(end=3.0)) + assert items == [(1.0, SampleData("a", 1.0)), (2.0, SampleData("b", 2.0))] + + # Both + items = list(store._iter_items(start=2.0, end=4.0)) + assert items == [(2.0, SampleData("b", 2.0)), (3.0, SampleData("c", 3.0))] + + def test_empty_store(self, store_factory, store_name, temp_dir): + store = store_factory(temp_dir) + + assert store.load(1.0) is None + assert store._find_closest_timestamp(1.0) is None + assert list(store._iter_items()) == [] + + def test_first_and_first_timestamp(self, store_factory, store_name, temp_dir): + store = store_factory(temp_dir) + + # Empty store + assert store.first() is None + assert store.first_timestamp() is None + + # Add data (in chronological order) + store.save(SampleData("a", 1.0), SampleData("b", 2.0), SampleData("c", 3.0)) + + # Should return first by timestamp + assert store.first_timestamp() == 1.0 + assert store.first() == SampleData("a", 1.0) + + def test_find_closest(self, store_factory, store_name, temp_dir): + store = store_factory(temp_dir) + store.save(SampleData("a", 1.0), SampleData("b", 2.0), SampleData("c", 3.0)) + + # Exact match + assert store.find_closest(2.0) == SampleData("b", 2.0) + + # Closest to 1.4 is 1.0 + assert store.find_closest(1.4) == SampleData("a", 1.0) + + # Closest to 1.6 is 2.0 + assert store.find_closest(1.6) == SampleData("b", 2.0) + + # With tolerance + assert store.find_closest(1.4, tolerance=0.5) == SampleData("a", 1.0) + assert store.find_closest(1.4, tolerance=0.3) is None + + def test_find_closest_seek(self, store_factory, store_name, temp_dir): + store = store_factory(temp_dir) + store.save(SampleData("a", 10.0), SampleData("b", 11.0), SampleData("c", 12.0)) + + # Seek 0 = first item (10.0) + assert store.find_closest_seek(0.0) == SampleData("a", 10.0) + + # Seek 1.0 = 11.0 + assert store.find_closest_seek(1.0) == SampleData("b", 11.0) + + # Seek 1.4 -> closest to 11.4 is 11.0 + assert store.find_closest_seek(1.4) == SampleData("b", 11.0) + + # Seek 1.6 -> closest to 11.6 is 12.0 + assert store.find_closest_seek(1.6) == SampleData("c", 12.0) + + # With tolerance + assert store.find_closest_seek(1.4, tolerance=0.5) == SampleData("b", 11.0) + assert store.find_closest_seek(1.4, tolerance=0.3) is None + + def test_iterate(self, store_factory, store_name, temp_dir): + store = store_factory(temp_dir) + store.save(SampleData("a", 1.0), SampleData("b", 2.0), SampleData("c", 3.0)) + + # Should iterate in timestamp order, returning data only (not tuples) + items = list(store.iterate()) + assert items == [ + SampleData("a", 1.0), + SampleData("b", 2.0), + SampleData("c", 3.0), + ] + + def test_iterate_with_seek_and_duration(self, store_factory, store_name, temp_dir): + store = store_factory(temp_dir) + store.save( + SampleData("a", 10.0), + SampleData("b", 11.0), + SampleData("c", 12.0), + SampleData("d", 13.0), + ) + + # Seek from start + items = list(store.iterate(seek=1.0)) + assert items == [ + SampleData("b", 11.0), + SampleData("c", 12.0), + SampleData("d", 13.0), + ] + + # Duration + items = list(store.iterate(duration=2.0)) + assert items == [SampleData("a", 10.0), SampleData("b", 11.0)] + + # Seek + duration + items = list(store.iterate(seek=1.0, duration=2.0)) + assert items == [SampleData("b", 11.0), SampleData("c", 12.0)] + + # from_timestamp + items = list(store.iterate(from_timestamp=12.0)) + assert items == [SampleData("c", 12.0), SampleData("d", 13.0)] + + def test_variadic_save(self, store_factory, store_name, temp_dir): + store = store_factory(temp_dir) + + # Save multiple items at once + store.save( + SampleData("a", 1.0), + SampleData("b", 2.0), + SampleData("c", 3.0), + ) + + assert store.load(1.0) == SampleData("a", 1.0) + assert store.load(2.0) == SampleData("b", 2.0) + assert store.load(3.0) == SampleData("c", 3.0) + + def test_pipe_save(self, store_factory, store_name, temp_dir): + import reactivex as rx + + store = store_factory(temp_dir) + + # Create observable with test data + source = rx.of( + SampleData("a", 1.0), + SampleData("b", 2.0), + SampleData("c", 3.0), + ) + + # Pipe through store.pipe_save — should save and pass through + results: list[SampleData] = [] + source.pipe(store.pipe_save).subscribe(results.append) + + # Data should be saved + assert store.load(1.0) == SampleData("a", 1.0) + assert store.load(2.0) == SampleData("b", 2.0) + assert store.load(3.0) == SampleData("c", 3.0) + + # Data should also pass through + assert results == [ + SampleData("a", 1.0), + SampleData("b", 2.0), + SampleData("c", 3.0), + ] + + def test_consume_stream(self, store_factory, store_name, temp_dir): + import reactivex as rx + + store = store_factory(temp_dir) + + # Create observable with test data + source = rx.of( + SampleData("a", 1.0), + SampleData("b", 2.0), + SampleData("c", 3.0), + ) + + # Consume stream — should save all items + disposable = store.consume_stream(source) + + # Data should be saved + assert store.load(1.0) == SampleData("a", 1.0) + assert store.load(2.0) == SampleData("b", 2.0) + assert store.load(3.0) == SampleData("c", 3.0) + + disposable.dispose() + + def test_iterate_items(self, store_factory, store_name, temp_dir): + store = store_factory(temp_dir) + store.save(SampleData("a", 1.0), SampleData("b", 2.0), SampleData("c", 3.0)) + + items = list(store.iterate_items()) + assert items == [ + (1.0, SampleData("a", 1.0)), + (2.0, SampleData("b", 2.0)), + (3.0, SampleData("c", 3.0)), + ] + + # With seek + items = list(store.iterate_items(seek=1.0)) + assert len(items) == 2 + assert items[0] == (2.0, SampleData("b", 2.0)) + + def test_stream_basic(self, store_factory, store_name, temp_dir): + store = store_factory(temp_dir) + store.save(SampleData("a", 1.0), SampleData("b", 2.0), SampleData("c", 3.0)) + + # Stream at high speed (essentially instant) + results: list[SampleData] = [] + store.stream(speed=1000.0).subscribe( + on_next=results.append, + on_completed=lambda: None, + ) + + # Give it a moment to complete + import time + + time.sleep(0.1) + + assert results == [ + SampleData("a", 1.0), + SampleData("b", 2.0), + SampleData("c", 3.0), + ] + + +@pytest.mark.parametrize("store_factory,store_name", testdata) +class TestCollectionAPI: + """Test new collection API methods on all backends.""" + + def test_len(self, store_factory, store_name, temp_dir): + store = store_factory(temp_dir) + assert len(store) == 0 + store.save(SampleData("a", 1.0), SampleData("b", 2.0), SampleData("c", 3.0)) + assert len(store) == 3 + + def test_iter(self, store_factory, store_name, temp_dir): + store = store_factory(temp_dir) + store.save(SampleData("a", 1.0), SampleData("b", 2.0)) + items = list(store) + assert items == [SampleData("a", 1.0), SampleData("b", 2.0)] + + def test_last_timestamp(self, store_factory, store_name, temp_dir): + store = store_factory(temp_dir) + assert store.last_timestamp() is None + store.save(SampleData("a", 1.0), SampleData("b", 2.0), SampleData("c", 3.0)) + assert store.last_timestamp() == 3.0 + + def test_last(self, store_factory, store_name, temp_dir): + store = store_factory(temp_dir) + assert store.last() is None + store.save(SampleData("a", 1.0), SampleData("b", 2.0), SampleData("c", 3.0)) + assert store.last() == SampleData("c", 3.0) + + def test_start_end_ts(self, store_factory, store_name, temp_dir): + store = store_factory(temp_dir) + assert store.start_ts is None + assert store.end_ts is None + store.save(SampleData("a", 1.0), SampleData("b", 2.0), SampleData("c", 3.0)) + assert store.start_ts == 1.0 + assert store.end_ts == 3.0 + + def test_time_range(self, store_factory, store_name, temp_dir): + store = store_factory(temp_dir) + assert store.time_range() is None + store.save(SampleData("a", 1.0), SampleData("b", 5.0)) + assert store.time_range() == (1.0, 5.0) + + def test_duration(self, store_factory, store_name, temp_dir): + store = store_factory(temp_dir) + assert store.duration() == 0.0 + store.save(SampleData("a", 1.0), SampleData("b", 5.0)) + assert store.duration() == 4.0 + + def test_find_before(self, store_factory, store_name, temp_dir): + store = store_factory(temp_dir) + store.save(SampleData("a", 1.0), SampleData("b", 2.0), SampleData("c", 3.0)) + + assert store.find_before(0.5) is None + assert store.find_before(1.0) is None # strictly before + assert store.find_before(1.5) == SampleData("a", 1.0) + assert store.find_before(2.5) == SampleData("b", 2.0) + assert store.find_before(10.0) == SampleData("c", 3.0) + + def test_find_after(self, store_factory, store_name, temp_dir): + store = store_factory(temp_dir) + store.save(SampleData("a", 1.0), SampleData("b", 2.0), SampleData("c", 3.0)) + + assert store.find_after(0.5) == SampleData("a", 1.0) + assert store.find_after(1.0) == SampleData("b", 2.0) # strictly after + assert store.find_after(2.5) == SampleData("c", 3.0) + assert store.find_after(3.0) is None # strictly after + assert store.find_after(10.0) is None + + def test_slice_by_time(self, store_factory, store_name, temp_dir): + store = store_factory(temp_dir) + store.save( + SampleData("a", 1.0), + SampleData("b", 2.0), + SampleData("c", 3.0), + SampleData("d", 4.0), + ) + + # [2.0, 4.0) should include b, c + result = store.slice_by_time(2.0, 4.0) + assert result == [SampleData("b", 2.0), SampleData("c", 3.0)] diff --git a/dimos/memory/timeseries/test_legacy.py b/dimos/memory/timeseries/test_legacy.py new file mode 100644 index 0000000000..aaad962a95 --- /dev/null +++ b/dimos/memory/timeseries/test_legacy.py @@ -0,0 +1,48 @@ +# 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. +"""Tests specific to LegacyPickleStore.""" + +from dimos.memory.timeseries.legacy import LegacyPickleStore + + +class TestLegacyPickleStoreRealData: + """Test LegacyPickleStore with real recorded data.""" + + def test_read_lidar_recording(self) -> None: + """Test reading from unitree_go2_bigoffice/lidar recording.""" + store = LegacyPickleStore("unitree_go2_bigoffice/lidar") + + # Check first timestamp exists + first_ts = store.first_timestamp() + assert first_ts is not None + assert first_ts > 0 + + # Check first data + first = store.first() + assert first is not None + assert hasattr(first, "ts") + + # Check find_closest_seek works + data_at_10s = store.find_closest_seek(10.0) + assert data_at_10s is not None + + # Check iteration returns monotonically increasing timestamps + prev_ts = None + for i, item in enumerate(store.iterate()): + assert item.ts is not None + if prev_ts is not None: + assert item.ts >= prev_ts, "Timestamps should be monotonically increasing" + prev_ts = item.ts + if i >= 10: # Only check first 10 items + break diff --git a/dimos/models/embedding/__init__.py b/dimos/models/embedding/__init__.py index a780f92246..050d35467e 100644 --- a/dimos/models/embedding/__init__.py +++ b/dimos/models/embedding/__init__.py @@ -7,24 +7,24 @@ # Optional: CLIP support try: - from dimos.models.embedding.clip import CLIPEmbedding, CLIPModel + from dimos.models.embedding.clip import CLIPModel + + __all__.append("CLIPModel") except ImportError: pass -else: - __all__ += ["CLIPEmbedding", "CLIPModel"] # Optional: MobileCLIP support try: - from dimos.models.embedding.mobileclip import MobileCLIPEmbedding, MobileCLIPModel + from dimos.models.embedding.mobileclip import MobileCLIPModel + + __all__.append("MobileCLIPModel") except ImportError: pass -else: - __all__ += ["MobileCLIPEmbedding", "MobileCLIPModel"] # Optional: TorchReID support try: - from dimos.models.embedding.treid import TorchReIDEmbedding, TorchReIDModel + from dimos.models.embedding.treid import TorchReIDModel + + __all__.append("TorchReIDModel") except ImportError: pass -else: - __all__ += ["TorchReIDEmbedding", "TorchReIDModel"] diff --git a/dimos/models/embedding/base.py b/dimos/models/embedding/base.py index eba5e45894..c6b78fcf2c 100644 --- a/dimos/models/embedding/base.py +++ b/dimos/models/embedding/base.py @@ -17,7 +17,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass import time -from typing import TYPE_CHECKING, Generic, TypeVar +from typing import TYPE_CHECKING import numpy as np import torch @@ -90,16 +90,13 @@ def to_cpu(self) -> Embedding: return self -E = TypeVar("E", bound="Embedding") - - -class EmbeddingModel(ABC, Generic[E]): +class EmbeddingModel(ABC): """Abstract base class for embedding models supporting vision and language.""" device: str @abstractmethod - def embed(self, *images: Image) -> E | list[E]: + def embed(self, *images: Image) -> Embedding | list[Embedding]: """ Embed one or more images. Returns single Embedding if one image, list if multiple. @@ -107,14 +104,14 @@ def embed(self, *images: Image) -> E | list[E]: pass @abstractmethod - def embed_text(self, *texts: str) -> E | list[E]: + def embed_text(self, *texts: str) -> Embedding | list[Embedding]: """ Embed one or more text strings. Returns single Embedding if one text, list if multiple. """ pass - def compare_one_to_many(self, query: E, candidates: list[E]) -> torch.Tensor: + def compare_one_to_many(self, query: Embedding, candidates: list[Embedding]) -> torch.Tensor: """ Efficiently compare one query against many candidates on GPU. @@ -129,7 +126,9 @@ def compare_one_to_many(self, query: E, candidates: list[E]) -> torch.Tensor: candidate_tensors = torch.stack([c.to_torch(self.device) for c in candidates]) return query_tensor @ candidate_tensors.T - def compare_many_to_many(self, queries: list[E], candidates: list[E]) -> torch.Tensor: + def compare_many_to_many( + self, queries: list[Embedding], candidates: list[Embedding] + ) -> torch.Tensor: """ Efficiently compare all queries against all candidates on GPU. @@ -144,7 +143,9 @@ def compare_many_to_many(self, queries: list[E], candidates: list[E]) -> torch.T candidate_tensors = torch.stack([c.to_torch(self.device) for c in candidates]) return query_tensors @ candidate_tensors.T - def query(self, query_emb: E, candidates: list[E], top_k: int = 5) -> list[tuple[int, float]]: + def query( + self, query_emb: Embedding, candidates: list[Embedding], top_k: int = 5 + ) -> list[tuple[int, float]]: """ Find top-k most similar candidates to query (GPU accelerated). @@ -160,6 +161,5 @@ def query(self, query_emb: E, candidates: list[E], top_k: int = 5) -> list[tuple top_values, top_indices = similarities.topk(k=min(top_k, len(candidates))) return [(idx.item(), val.item()) for idx, val in zip(top_indices, top_values, strict=False)] - def warmup(self) -> None: - """Optional warmup method to pre-load model.""" - pass + + ... diff --git a/dimos/models/embedding/clip.py b/dimos/models/embedding/clip.py index a6512568d2..1b8d3e68bb 100644 --- a/dimos/models/embedding/clip.py +++ b/dimos/models/embedding/clip.py @@ -25,16 +25,13 @@ from dimos.msgs.sensor_msgs import Image -class CLIPEmbedding(Embedding): ... - - @dataclass class CLIPModelConfig(HuggingFaceEmbeddingModelConfig): model_name: str = "openai/clip-vit-base-patch32" dtype: torch.dtype = torch.float32 -class CLIPModel(EmbeddingModel[CLIPEmbedding], HuggingFaceModel): +class CLIPModel(EmbeddingModel, HuggingFaceModel): """CLIP embedding model for vision-language re-identification.""" default_config = CLIPModelConfig @@ -50,7 +47,7 @@ def _model(self) -> HFCLIPModel: def _processor(self) -> CLIPProcessor: return CLIPProcessor.from_pretrained(self.config.model_name) - def embed(self, *images: Image) -> CLIPEmbedding | list[CLIPEmbedding]: + def embed(self, *images: Image) -> Embedding | list[Embedding]: """Embed one or more images. Returns embeddings as torch.Tensor on device for efficient GPU comparisons. @@ -67,14 +64,14 @@ def embed(self, *images: Image) -> CLIPEmbedding | list[CLIPEmbedding]: image_features = functional.normalize(image_features, dim=-1) # Create embeddings (keep as torch.Tensor on device) - embeddings = [] + embeddings: list[Embedding] = [] for i, feat in enumerate(image_features): timestamp = images[i].ts - embeddings.append(CLIPEmbedding(vector=feat, timestamp=timestamp)) + embeddings.append(Embedding(vector=feat, timestamp=timestamp)) return embeddings[0] if len(images) == 1 else embeddings - def embed_text(self, *texts: str) -> CLIPEmbedding | list[CLIPEmbedding]: + def embed_text(self, *texts: str) -> Embedding | list[Embedding]: """Embed one or more text strings. Returns embeddings as torch.Tensor on device for efficient GPU comparisons. @@ -89,9 +86,9 @@ def embed_text(self, *texts: str) -> CLIPEmbedding | list[CLIPEmbedding]: text_features = functional.normalize(text_features, dim=-1) # Create embeddings (keep as torch.Tensor on device) - embeddings = [] + embeddings: list[Embedding] = [] for feat in text_features: - embeddings.append(CLIPEmbedding(vector=feat)) + embeddings.append(Embedding(vector=feat)) return embeddings[0] if len(texts) == 1 else embeddings diff --git a/dimos/models/embedding/mobileclip.py b/dimos/models/embedding/mobileclip.py index 7c3d7adc69..c02361b367 100644 --- a/dimos/models/embedding/mobileclip.py +++ b/dimos/models/embedding/mobileclip.py @@ -27,15 +27,12 @@ from dimos.utils.data import get_data -class MobileCLIPEmbedding(Embedding): ... - - @dataclass class MobileCLIPModelConfig(EmbeddingModelConfig): model_name: str = "MobileCLIP2-S4" -class MobileCLIPModel(EmbeddingModel[MobileCLIPEmbedding], LocalModel): +class MobileCLIPModel(EmbeddingModel, LocalModel): """MobileCLIP embedding model for vision-language re-identification.""" default_config = MobileCLIPModelConfig @@ -62,7 +59,7 @@ def _preprocess(self) -> Any: def _tokenizer(self) -> Any: return open_clip.get_tokenizer(self.config.model_name) - def embed(self, *images: Image) -> MobileCLIPEmbedding | list[MobileCLIPEmbedding]: + def embed(self, *images: Image) -> Embedding | list[Embedding]: """Embed one or more images. Returns embeddings as torch.Tensor on device for efficient GPU comparisons. @@ -83,11 +80,11 @@ def embed(self, *images: Image) -> MobileCLIPEmbedding | list[MobileCLIPEmbeddin embeddings = [] for i, feat in enumerate(feats): timestamp = images[i].ts - embeddings.append(MobileCLIPEmbedding(vector=feat, timestamp=timestamp)) + embeddings.append(Embedding(vector=feat, timestamp=timestamp)) return embeddings[0] if len(images) == 1 else embeddings - def embed_text(self, *texts: str) -> MobileCLIPEmbedding | list[MobileCLIPEmbedding]: + def embed_text(self, *texts: str) -> Embedding | list[Embedding]: """Embed one or more text strings. Returns embeddings as torch.Tensor on device for efficient GPU comparisons. @@ -101,7 +98,7 @@ def embed_text(self, *texts: str) -> MobileCLIPEmbedding | list[MobileCLIPEmbedd # Create embeddings (keep as torch.Tensor on device) embeddings = [] for feat in feats: - embeddings.append(MobileCLIPEmbedding(vector=feat)) + embeddings.append(Embedding(vector=feat)) return embeddings[0] if len(texts) == 1 else embeddings diff --git a/dimos/models/embedding/treid.py b/dimos/models/embedding/treid.py index 23da1effe2..85e32cd39b 100644 --- a/dimos/models/embedding/treid.py +++ b/dimos/models/embedding/treid.py @@ -13,13 +13,14 @@ # limitations under the License. import warnings + +warnings.filterwarnings("ignore", message="Cython evaluation.*unavailable", category=UserWarning) + from dataclasses import dataclass from functools import cached_property import torch import torch.nn.functional as functional - -warnings.filterwarnings("ignore", message="Cython evaluation.*unavailable", category=UserWarning) from torchreid import utils as torchreid_utils from dimos.models.base import LocalModel @@ -28,9 +29,6 @@ from dimos.utils.data import get_data -class TorchReIDEmbedding(Embedding): ... - - # osnet models downloaded from https://kaiyangzhou.github.io/deep-person-reid/MODEL_ZOO.html # into dimos/data/models_torchreid/ # feel free to add more @@ -39,7 +37,7 @@ class TorchReIDModelConfig(EmbeddingModelConfig): model_name: str = "osnet_x1_0" -class TorchReIDModel(EmbeddingModel[TorchReIDEmbedding], LocalModel): +class TorchReIDModel(EmbeddingModel, LocalModel): """TorchReID embedding model for person re-identification.""" default_config = TorchReIDModelConfig @@ -54,7 +52,7 @@ def _model(self) -> torchreid_utils.FeatureExtractor: device=self.config.device, ) - def embed(self, *images: Image) -> TorchReIDEmbedding | list[TorchReIDEmbedding]: + def embed(self, *images: Image) -> Embedding | list[Embedding]: """Embed one or more images. Returns embeddings as torch.Tensor on device for efficient GPU comparisons. @@ -79,11 +77,11 @@ def embed(self, *images: Image) -> TorchReIDEmbedding | list[TorchReIDEmbedding] embeddings = [] for i, feat in enumerate(features_tensor): timestamp = images[i].ts - embeddings.append(TorchReIDEmbedding(vector=feat, timestamp=timestamp)) + embeddings.append(Embedding(vector=feat, timestamp=timestamp)) return embeddings[0] if len(images) == 1 else embeddings - def embed_text(self, *texts: str) -> TorchReIDEmbedding | list[TorchReIDEmbedding]: + def embed_text(self, *texts: str) -> Embedding | list[Embedding]: """Text embedding not supported for ReID models. TorchReID models are vision-only person re-identification models diff --git a/dimos/msgs/geometry_msgs/Transform.py b/dimos/msgs/geometry_msgs/Transform.py index 58bab1ec4a..5f50f9b9d1 100644 --- a/dimos/msgs/geometry_msgs/Transform.py +++ b/dimos/msgs/geometry_msgs/Transform.py @@ -20,6 +20,8 @@ if TYPE_CHECKING: import rerun as rr + from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped + from dimos_lcm.geometry_msgs import ( Transform as LCMTransform, TransformStamped as LCMTransformStamped, @@ -191,7 +193,7 @@ def from_pose(cls, frame_id: str, pose: Pose | PoseStamped) -> Transform: # typ else: raise TypeError(f"Expected Pose or PoseStamped, got {type(pose).__name__}") - def to_pose(self, **kwargs) -> PoseStamped: # type: ignore[name-defined, no-untyped-def] + def to_pose(self, **kwargs: object) -> PoseStamped: """Create a Transform from a Pose or PoseStamped. Args: @@ -201,10 +203,10 @@ def to_pose(self, **kwargs) -> PoseStamped: # type: ignore[name-defined, no-unt A Transform with the same translation and rotation as the pose """ # Import locally to avoid circular imports - from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped + from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped as _PoseStamped # Handle both Pose and PoseStamped - return PoseStamped( + result: PoseStamped = _PoseStamped( **{ "position": self.translation, "orientation": self.rotation, @@ -212,6 +214,7 @@ def to_pose(self, **kwargs) -> PoseStamped: # type: ignore[name-defined, no-unt }, **kwargs, ) + return result def to_matrix(self) -> np.ndarray: # type: ignore[name-defined] """Convert Transform to a 4x4 transformation matrix. diff --git a/dimos/msgs/sensor_msgs/Image.py b/dimos/msgs/sensor_msgs/Image.py index 2a8aa2c017..c74001a369 100644 --- a/dimos/msgs/sensor_msgs/Image.py +++ b/dimos/msgs/sensor_msgs/Image.py @@ -33,6 +33,7 @@ from dimos.utils.reactive import quality_barrier if TYPE_CHECKING: + from collections.abc import Callable import os from reactivex.observable import Observable @@ -592,9 +593,9 @@ def sharpness_window(target_frequency: float, source: Observable[Image]) -> Obse thread_scheduler = ThreadPoolScheduler(max_workers=1) def find_best(*_args: Any) -> Image | None: - if not window._items: + if len(window) == 0: return None - return max(window._items, key=lambda img: img.sharpness) # type: ignore[no-any-return] + return max(window, key=lambda img: img.sharpness) # type: ignore[no-any-return] return rx.interval(1.0 / target_frequency).pipe( # type: ignore[misc] ops.observe_on(thread_scheduler), @@ -603,7 +604,7 @@ def find_best(*_args: Any) -> Image | None: ) -def sharpness_barrier(target_frequency: float) -> Any: +def sharpness_barrier(target_frequency: float) -> Callable[[Observable[Image]], Observable[Image]]: """Select the sharpest Image within each time window.""" if target_frequency <= 0: raise ValueError("target_frequency must be positive") diff --git a/dimos/perception/detection/reid/embedding_id_system.py b/dimos/perception/detection/reid/embedding_id_system.py index 9b57e1eb6c..15bb491f5c 100644 --- a/dimos/perception/detection/reid/embedding_id_system.py +++ b/dimos/perception/detection/reid/embedding_id_system.py @@ -33,7 +33,7 @@ class EmbeddingIDSystem(IDSystem): def __init__( self, - model: Callable[[], EmbeddingModel[Embedding]], + model: Callable[[], EmbeddingModel], padding: int = 0, similarity_threshold: float = 0.63, comparison_mode: Literal["max", "mean", "top_k_mean"] = "top_k_mean", diff --git a/dimos/protocol/pubsub/benchmark/test_benchmark.py b/dimos/protocol/pubsub/benchmark/test_benchmark.py index 865c4ee324..39a4421c35 100644 --- a/dimos/protocol/pubsub/benchmark/test_benchmark.py +++ b/dimos/protocol/pubsub/benchmark/test_benchmark.py @@ -82,6 +82,7 @@ def benchmark_results() -> Generator[BenchmarkResults, None, None]: results.print_heatmap() results.print_bandwidth_heatmap() results.print_latency_heatmap() + results.print_loss_heatmap() @pytest.mark.tool diff --git a/dimos/protocol/pubsub/benchmark/type.py b/dimos/protocol/pubsub/benchmark/type.py index 425f2a480a..a9ef80fe7a 100644 --- a/dimos/protocol/pubsub/benchmark/type.py +++ b/dimos/protocol/pubsub/benchmark/type.py @@ -263,3 +263,11 @@ def fmt(v: float) -> str: fmt, high_is_good=False, ) + + def print_loss_heatmap(self) -> None: + """Print message loss percentage heatmap.""" + + def fmt(v: float) -> str: + return f"{v:.1f}%" + + self._print_heatmap("Loss %", lambda r: r.loss_pct, fmt, high_is_good=False) diff --git a/dimos/protocol/tf/test_tf.py b/dimos/protocol/tf/test_tf.py index 0b5b332c3d..bdbd808cbb 100644 --- a/dimos/protocol/tf/test_tf.py +++ b/dimos/protocol/tf/test_tf.py @@ -48,6 +48,7 @@ def test_tf_ros_example() -> None: time.sleep(0.2) end_effector_global_pose = tf.get("base_link", "end_effector") + assert end_effector_global_pose is not None assert end_effector_global_pose.translation.x == pytest.approx(1.366, abs=1e-3) assert end_effector_global_pose.translation.y == pytest.approx(0.366, abs=1e-3) @@ -116,6 +117,7 @@ def test_tf_main() -> None: # The chain should compose: world->robot (1,2,3) + robot->sensor (0.5,0,0.2) # Expected translation: (1.5, 2.0, 3.2) + assert chain_transform is not None assert abs(chain_transform.translation.x - 1.5) < 0.001 assert abs(chain_transform.translation.y - 2.0) < 0.001 assert abs(chain_transform.translation.z - 3.2) < 0.001 @@ -163,12 +165,14 @@ def test_tf_main() -> None: # if you have "diagon" https://diagon.arthursonzogni.com/ installed you can draw a graph print(broadcaster.graph()) + assert world_object is not None assert abs(world_object.translation.x - 1.5) < 0.001 assert abs(world_object.translation.y - 3.0) < 0.001 assert abs(world_object.translation.z - 3.2) < 0.001 # this doesn't work atm robot_to_charger = broadcaster.get("robot", "charger") + assert robot_to_charger is not None # Expected: robot->world->charger print(f"robot_to_charger translation: {robot_to_charger.translation}") @@ -196,7 +200,7 @@ def test_add_transform(self) -> None: buffer.add(transform) assert len(buffer) == 1 - assert buffer[0] == transform + assert buffer.first() == transform def test_get(self) -> None: buffer = TBuffer() @@ -250,7 +254,9 @@ def test_buffer_pruning(self) -> None: # Old transform should be pruned assert len(buffer) == 1 - assert buffer[0].translation.x == 2.0 + first = buffer.first() + assert first is not None + assert first.translation.x == 2.0 class TestMultiTBuffer: diff --git a/dimos/protocol/tf/tf.py b/dimos/protocol/tf/tf.py index 44d5303a05..825e89fc8c 100644 --- a/dimos/protocol/tf/tf.py +++ b/dimos/protocol/tf/tf.py @@ -20,12 +20,12 @@ from functools import reduce from typing import TypeVar -from dimos.msgs.geometry_msgs import Transform +from dimos.memory.timeseries.inmemory import InMemoryStore +from dimos.msgs.geometry_msgs import PoseStamped, Transform from dimos.msgs.tf2_msgs import TFMessage from dimos.protocol.pubsub.impl.lcmpubsub import LCM, Topic from dimos.protocol.pubsub.spec import PubSub from dimos.protocol.service.lcmservice import Service # type: ignore[attr-defined] -from dimos.types.timestamped import TimestampedCollection CONFIG = TypeVar("CONFIG") @@ -52,13 +52,13 @@ def get_frames(self) -> set[str]: return set() @abstractmethod - def get( # type: ignore[no-untyped-def] + def get( self, parent_frame: str, child_frame: str, time_point: float | None = None, time_tolerance: float | None = None, - ): ... + ) -> Transform | None: ... def receive_transform(self, *args: Transform) -> None: ... @@ -71,64 +71,44 @@ def receive_tfmessage(self, msg: TFMessage) -> None: TopicT = TypeVar("TopicT") -# stores a single transform -class TBuffer(TimestampedCollection[Transform]): +class TBuffer(InMemoryStore[Transform]): def __init__(self, buffer_size: float = 10.0) -> None: super().__init__() self.buffer_size = buffer_size def add(self, transform: Transform) -> None: - super().add(transform) - self._prune_old_transforms(transform.ts) - - def _prune_old_transforms(self, current_time) -> None: # type: ignore[no-untyped-def] - if not self._items: - return - - cutoff_time = current_time - self.buffer_size - - while self._items and self._items[0].ts < cutoff_time: - self._items.pop(0) + self.save(transform) + self.prune_old(transform.ts - self.buffer_size) def get(self, time_point: float | None = None, time_tolerance: float = 1.0) -> Transform | None: """Get transform at specified time or latest if no time given.""" if time_point is None: - # Return the latest transform - return self[-1] if len(self) > 0 else None - + return self.last() return self.find_closest(time_point, time_tolerance) def __str__(self) -> str: - if not self._items: + if len(self) == 0: return "TBuffer(empty)" - # Get unique frame info from the transforms - frame_pairs = set() - if self._items: - frame_pairs.add((self._items[0].frame_id, self._items[0].child_frame_id)) - + first_item = self.first() time_range = self.time_range() - if time_range: + if time_range and first_item: from dimos.types.timestamped import to_human_readable start_time = to_human_readable(time_range[0]) end_time = to_human_readable(time_range[1]) duration = time_range[1] - time_range[0] - frame_str = ( - f"{self._items[0].frame_id} -> {self._items[0].child_frame_id}" - if self._items - else "unknown" - ) + frame_str = f"{first_item.frame_id} -> {first_item.child_frame_id}" return ( f"TBuffer(" f"{frame_str}, " - f"{len(self._items)} msgs, " + f"{len(self)} msgs, " f"{duration:.2f}s [{start_time} - {end_time}])" ) - return f"TBuffer({len(self._items)} msgs)" + return f"TBuffer({len(self)} msgs)" # stores multiple transform buffers @@ -334,6 +314,18 @@ def get( ) -> Transform | None: return super().get(parent_frame, child_frame, time_point, time_tolerance) + def get_pose( + self, + parent_frame: str, + child_frame: str, + time_point: float | None = None, + time_tolerance: float | None = None, + ) -> PoseStamped | None: + tf = self.get(parent_frame, child_frame, time_point, time_tolerance) + if not tf: + return None + return tf.to_pose() + def receive_msg(self, msg: TFMessage, topic: Topic) -> None: self.receive_tfmessage(msg) diff --git a/dimos/robot/unitree/connection/go2.py b/dimos/robot/unitree/connection/go2.py index 58145c9371..8e021b3510 100644 --- a/dimos/robot/unitree/connection/go2.py +++ b/dimos/robot/unitree/connection/go2.py @@ -194,13 +194,13 @@ def __init__( # type: ignore[no-untyped-def] @rpc def record(self, recording_name: str) -> None: lidar_store: TimedSensorStorage = TimedSensorStorage(f"{recording_name}/lidar") # type: ignore[type-arg] - lidar_store.save_stream(self.connection.lidar_stream()).subscribe(lambda x: x) # type: ignore[arg-type] + lidar_store.consume_stream(self.connection.lidar_stream()) odom_store: TimedSensorStorage = TimedSensorStorage(f"{recording_name}/odom") # type: ignore[type-arg] - odom_store.save_stream(self.connection.odom_stream()).subscribe(lambda x: x) # type: ignore[arg-type] + odom_store.consume_stream(self.connection.odom_stream()) video_store: TimedSensorStorage = TimedSensorStorage(f"{recording_name}/video") # type: ignore[type-arg] - video_store.save_stream(self.connection.video_stream()).subscribe(lambda x: x) # type: ignore[arg-type] + video_store.consume_stream(self.connection.video_stream()) @rpc def start(self) -> None: diff --git a/dimos/robot/unitree_webrtc/type/test_odometry.py b/dimos/robot/unitree_webrtc/type/test_odometry.py index e277455cdd..489cd523b5 100644 --- a/dimos/robot/unitree_webrtc/type/test_odometry.py +++ b/dimos/robot/unitree_webrtc/type/test_odometry.py @@ -14,10 +14,7 @@ from __future__ import annotations -from operator import add, sub - import pytest -import reactivex.operators as ops from dimos.robot.unitree_webrtc.type.odometry import Odometry from dimos.utils.testing import SensorReplay @@ -38,19 +35,6 @@ def test_odometry_conversion_and_count() -> None: assert isinstance(odom, Odometry) -def test_last_yaw_value() -> None: - """Verify yaw of the final message (regression guard).""" - last_msg = SensorReplay(name="raw_odometry_rotate_walk").stream().pipe(ops.last()).run() - - assert last_msg is not None, "Replay is empty" - assert last_msg["data"]["pose"]["orientation"] == { - "x": 0.01077, - "y": 0.008505, - "z": 0.499171, - "w": -0.866395, - } - - def test_total_rotation_travel_iterate() -> None: total_rad = 0.0 prev_yaw: float | None = None @@ -63,19 +47,3 @@ def test_total_rotation_travel_iterate() -> None: prev_yaw = yaw assert total_rad == pytest.approx(_EXPECTED_TOTAL_RAD, abs=0.001) - - -def test_total_rotation_travel_rxpy() -> None: - total_rad = ( - SensorReplay(name="raw_odometry_rotate_walk", autocast=Odometry.from_msg) - .stream() - .pipe( - ops.map(lambda odom: odom.orientation.radians.z), - ops.pairwise(), # [1,2,3,4] -> [[1,2], [2,3], [3,4]] - ops.starmap(sub), # [sub(1,2), sub(2,3), sub(3,4)] - ops.reduce(add), - ) - .run() - ) - - assert total_rad == pytest.approx(4.05, abs=0.01) diff --git a/dimos/types/test_timestamped.py b/dimos/types/test_timestamped.py index 88a8d65102..7de82e8f9a 100644 --- a/dimos/types/test_timestamped.py +++ b/dimos/types/test_timestamped.py @@ -19,11 +19,11 @@ from reactivex import operators as ops from reactivex.scheduler import ThreadPoolScheduler +from dimos.memory.timeseries.inmemory import InMemoryStore from dimos.msgs.sensor_msgs import Image from dimos.types.timestamped import ( Timestamped, TimestampedBufferCollection, - TimestampedCollection, align_timestamped, to_datetime, to_ros_stamp, @@ -133,13 +133,20 @@ def sample_items(): ] +def make_store(items: list[SimpleTimestamped] | None = None) -> InMemoryStore[SimpleTimestamped]: + store: InMemoryStore[SimpleTimestamped] = InMemoryStore() + if items: + store.save(*items) + return store + + @pytest.fixture def collection(sample_items): - return TimestampedCollection(sample_items) + return make_store(sample_items) def test_empty_collection() -> None: - collection = TimestampedCollection() + collection = make_store() assert len(collection) == 0 assert collection.duration() == 0.0 assert collection.time_range() is None @@ -147,16 +154,17 @@ def test_empty_collection() -> None: def test_add_items() -> None: - collection = TimestampedCollection() + collection = make_store() item1 = SimpleTimestamped(2.0, "two") item2 = SimpleTimestamped(1.0, "one") - collection.add(item1) - collection.add(item2) + collection.save(item1) + collection.save(item2) assert len(collection) == 2 - assert collection[0].data == "one" # Should be sorted by timestamp - assert collection[1].data == "two" + items = list(collection) + assert items[0].data == "one" # Should be sorted by timestamp + assert items[1].data == "two" def test_find_closest(collection) -> None: @@ -196,21 +204,13 @@ def test_find_before_after(collection) -> None: assert collection.find_after(7.0) is None # Nothing after last item -def test_merge_collections() -> None: - collection1 = TimestampedCollection( - [ - SimpleTimestamped(1.0, "a"), - SimpleTimestamped(3.0, "c"), - ] - ) - collection2 = TimestampedCollection( - [ - SimpleTimestamped(2.0, "b"), - SimpleTimestamped(4.0, "d"), - ] - ) +def test_save_from_multiple_stores() -> None: + store1 = make_store([SimpleTimestamped(1.0, "a"), SimpleTimestamped(3.0, "c")]) + store2 = make_store([SimpleTimestamped(2.0, "b"), SimpleTimestamped(4.0, "d")]) - merged = collection1.merge(collection2) + merged = make_store() + merged.save(*store1) + merged.save(*store2) assert len(merged) == 4 assert [item.data for item in merged] == ["a", "b", "c", "d"] @@ -244,7 +244,7 @@ def test_iteration(collection) -> None: def test_single_item_collection() -> None: - single = TimestampedCollection([SimpleTimestamped(5.0, "only")]) + single = make_store([SimpleTimestamped(5.0, "only")]) assert single.duration() == 0.0 assert single.time_range() == (5.0, 5.0) @@ -264,14 +264,17 @@ def test_time_window_collection() -> None: # Add a message at t=4.0, should keep messages from t=2.0 onwards window.add(SimpleTimestamped(4.0, "msg4")) assert len(window) == 3 # msg1 should be dropped - assert window[0].data == "msg2" # oldest is now msg2 - assert window[-1].data == "msg4" # newest is msg4 + first = window.first() + last = window.last() + assert first is not None and first.data == "msg2" # oldest is now msg2 + assert last is not None and last.data == "msg4" # newest is msg4 # Add a message at t=5.5, should drop msg2 and msg3 window.add(SimpleTimestamped(5.5, "msg5")) assert len(window) == 2 # only msg4 and msg5 remain - assert window[0].data == "msg4" - assert window[1].data == "msg5" + items = list(window) + assert items[0].data == "msg4" + assert items[1].data == "msg5" # Verify time range assert window.start_ts == 4.0 diff --git a/dimos/types/timestamped.py b/dimos/types/timestamped.py index 765b1adbcb..b229a2478e 100644 --- a/dimos/types/timestamped.py +++ b/dimos/types/timestamped.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. from collections import defaultdict -from collections.abc import Iterable, Iterator from datetime import datetime, timezone from typing import Generic, TypeVar, Union @@ -22,8 +21,8 @@ # from dimos_lcm.std_msgs import Time as ROSTime from reactivex.observable import Observable -from sortedcontainers import SortedKeyList # type: ignore[import-untyped] +from dimos.memory.timeseries.inmemory import InMemoryStore from dimos.types.weaklist import WeakList from dimos.utils.logging_config import setup_logger @@ -117,152 +116,29 @@ def ros_timestamp(self) -> list[int]: T = TypeVar("T", bound=Timestamped) -class TimestampedCollection(Generic[T]): - """A collection of timestamped objects with efficient time-based operations.""" - - def __init__(self, items: Iterable[T] | None = None) -> None: - self._items = SortedKeyList(items or [], key=lambda x: x.ts) - - def add(self, item: T) -> None: - """Add a timestamped item to the collection.""" - self._items.add(item) - - def find_closest(self, timestamp: float, tolerance: float | None = None) -> T | None: - """Find the timestamped object closest to the given timestamp.""" - if not self._items: - return None - - # Use binary search to find insertion point - idx = self._items.bisect_key_left(timestamp) - - # Check exact match - if idx < len(self._items) and self._items[idx].ts == timestamp: - return self._items[idx] # type: ignore[no-any-return] - - # Find candidates: item before and after - candidates = [] - - # Item before - if idx > 0: - candidates.append((idx - 1, abs(self._items[idx - 1].ts - timestamp))) - - # Item after - if idx < len(self._items): - candidates.append((idx, abs(self._items[idx].ts - timestamp))) - - if not candidates: - return None - - # Find closest - # When distances are equal, prefer the later item (higher index) - closest_idx, closest_distance = min(candidates, key=lambda x: (x[1], -x[0])) - - # Check tolerance if provided - if tolerance is not None and closest_distance > tolerance: - return None - - return self._items[closest_idx] # type: ignore[no-any-return] - - def find_before(self, timestamp: float) -> T | None: - """Find the last item before the given timestamp.""" - idx = self._items.bisect_key_left(timestamp) - return self._items[idx - 1] if idx > 0 else None - - def find_after(self, timestamp: float) -> T | None: - """Find the first item after the given timestamp.""" - idx = self._items.bisect_key_right(timestamp) - return self._items[idx] if idx < len(self._items) else None - - def merge(self, other: "TimestampedCollection[T]") -> "TimestampedCollection[T]": - """Merge two timestamped collections into a new one.""" - result = TimestampedCollection[T]() - result._items = SortedKeyList(self._items + other._items, key=lambda x: x.ts) - return result - - def duration(self) -> float: - """Get the duration of the collection in seconds.""" - if len(self._items) < 2: - return 0.0 - return self._items[-1].ts - self._items[0].ts # type: ignore[no-any-return] - - def time_range(self) -> tuple[float, float] | None: - """Get the time range (start, end) of the collection.""" - if not self._items: - return None - return (self._items[0].ts, self._items[-1].ts) - - def slice_by_time(self, start: float, end: float) -> "TimestampedCollection[T]": - """Get a subset of items within the given time range.""" - start_idx = self._items.bisect_key_left(start) - end_idx = self._items.bisect_key_right(end) - return TimestampedCollection(self._items[start_idx:end_idx]) - - @property - def start_ts(self) -> float | None: - """Get the start timestamp of the collection.""" - return self._items[0].ts if self._items else None - - @property - def end_ts(self) -> float | None: - """Get the end timestamp of the collection.""" - return self._items[-1].ts if self._items else None - - def __len__(self) -> int: - return len(self._items) - - def __iter__(self) -> Iterator: # type: ignore[type-arg] - return iter(self._items) - - def __getitem__(self, idx: int) -> T: - return self._items[idx] # type: ignore[no-any-return] - - PRIMARY = TypeVar("PRIMARY", bound=Timestamped) SECONDARY = TypeVar("SECONDARY", bound=Timestamped) -class TimestampedBufferCollection(TimestampedCollection[T]): - """A timestamped collection that maintains a sliding time window, dropping old messages.""" - - def __init__(self, window_duration: float, items: Iterable[T] | None = None) -> None: - """ - Initialize with a time window duration in seconds. +class TimestampedBufferCollection(InMemoryStore[T]): + """A sliding time window buffer backed by InMemoryStore.""" - Args: - window_duration: Maximum age of messages to keep in seconds - items: Optional initial items - """ - super().__init__(items) + def __init__(self, window_duration: float) -> None: + super().__init__() self.window_duration = window_duration def add(self, item: T) -> None: - """Add a timestamped item and remove any items outside the time window.""" - super().add(item) - self._prune_old_messages(item.ts) - - def _prune_old_messages(self, current_ts: float) -> None: - """Remove messages older than window_duration from the given timestamp.""" - cutoff_ts = current_ts - self.window_duration - - # Find the index of the first item that should be kept - keep_idx = self._items.bisect_key_left(cutoff_ts) + """Add a timestamped item and prune items outside the time window.""" + self.save(item) + self.prune_old(item.ts - self.window_duration) - # Remove old items - if keep_idx > 0: - del self._items[:keep_idx] + def remove(self, item: T) -> bool: + """Remove a timestamped item. Returns True if found and removed.""" + return self._delete(item.ts) is not None def remove_by_timestamp(self, timestamp: float) -> bool: - """Remove an item with the given timestamp. Returns True if item was found and removed.""" - idx = self._items.bisect_key_left(timestamp) - - if idx < len(self._items) and self._items[idx].ts == timestamp: - del self._items[idx] - return True - return False - - def remove(self, item: T) -> bool: - """Remove a timestamped item from the collection. Returns True if item was found and removed.""" - return self.remove_by_timestamp(item.ts) + """Remove an item by timestamp. Returns True if found and removed.""" + return self._delete(timestamp) is not None class MatchContainer(Timestamped, Generic[PRIMARY, SECONDARY]): diff --git a/dimos/utils/data.py b/dimos/utils/data.py index 7666343dbb..094d4a6a7c 100644 --- a/dimos/utils/data.py +++ b/dimos/utils/data.py @@ -105,7 +105,7 @@ def _get_repo_root() -> Path: @cache -def _get_data_dir(extra_path: str | None = None) -> Path: +def get_data_dir(extra_path: str | None = None) -> Path: if extra_path: return _get_repo_root() / "data" / extra_path return _get_repo_root() / "data" @@ -113,7 +113,7 @@ def _get_data_dir(extra_path: str | None = None) -> Path: @cache def _get_lfs_dir() -> Path: - return _get_data_dir() / ".lfs" + return get_data_dir() / ".lfs" def _check_git_lfs_available() -> bool: @@ -174,7 +174,7 @@ def _lfs_pull(file_path: Path, repo_root: Path) -> None: def _decompress_archive(filename: str | Path) -> Path: - target_dir = _get_data_dir() + target_dir = get_data_dir() filename_path = Path(filename) with tarfile.open(filename_path, "r:gz") as tar: tar.extractall(target_dir) @@ -211,7 +211,7 @@ def _pull_lfs_archive(filename: str | Path) -> Path: return file_path -def get_data(filename: str | Path) -> Path: +def get_data(name: str | Path) -> Path: """ Get the path to a test data, downloading from LFS if needed. @@ -222,29 +222,43 @@ def get_data(filename: str | Path) -> Path: 4. Download the file from LFS if it's a pointer file 5. Return the Path object to the actual file or dir + Supports nested paths like "dataset/subdir/file.jpg" - will download and + decompress "dataset" archive but return the full nested path. + Args: - filename: Name of the test file (e.g., "lidar_sample.bin") + name: Name of the test file or dir, optionally with nested path + (e.g., "lidar_sample.bin" or "dataset/frames/001.png") Returns: - Path: Path object to the test file + Path: Path object to the test file or dir Raises: RuntimeError: If Git LFS is not available or LFS operations fail FileNotFoundError: If the test file doesn't exist Usage: - # As string path - file_path = str(testFile("sample.bin")) + # Simple file/dir + file_path = get_data("sample.bin") - # As context manager for file operations - with testFile("sample.bin").open('rb') as f: - data = f.read() + # Nested path - downloads "dataset" archive, returns path to nested file + frame = get_data("dataset/frames/001.png") """ - data_dir = _get_data_dir() - file_path = data_dir / filename + data_dir = get_data_dir() + file_path = data_dir / name # already pulled and decompressed, return it directly if file_path.exists(): return file_path - return _decompress_archive(_pull_lfs_archive(filename)) + # extract archive root (first path component) and nested path + path_parts = Path(name).parts + archive_name = path_parts[0] + nested_path = Path(*path_parts[1:]) if len(path_parts) > 1 else None + + # download and decompress the archive root + archive_path = _decompress_archive(_pull_lfs_archive(archive_name)) + + # return full path including nested components + if nested_path: + return archive_path / nested_path + return archive_path diff --git a/dimos/utils/reactive.py b/dimos/utils/reactive.py index d76bd518cc..4397e0171e 100644 --- a/dimos/utils/reactive.py +++ b/dimos/utils/reactive.py @@ -19,6 +19,7 @@ import reactivex as rx from reactivex import operators as ops +from reactivex.abc import DisposableBase from reactivex.disposable import Disposable from reactivex.observable import Observable from reactivex.scheduler import ThreadPoolScheduler @@ -64,7 +65,7 @@ def _subscribe(observer, sch=None): # type: ignore[no-untyped-def] return rx.defer(lambda *_: per_sub()) # type: ignore[no-untyped-call] -class LatestReader(Generic[T]): +class LatestReader(DisposableBase, Generic[T]): """A callable object that returns the latest value from an observable.""" def __init__(self, initial_value: T, subscription, connection=None) -> None: # type: ignore[no-untyped-def] @@ -259,7 +260,9 @@ def spyfun(x): # type: ignore[no-untyped-def] return ops.map(spyfun) -def quality_barrier(quality_func: Callable[[T], float], target_frequency: float): # type: ignore[no-untyped-def] +def quality_barrier( + quality_func: Callable[[T], float], target_frequency: float +) -> Callable[[Observable[T]], Observable[T]]: """ RxPY pipe operator that selects the highest quality item within each time window. diff --git a/dimos/utils/test_data.py b/dimos/utils/test_data.py index 01f145f60c..ba135fe255 100644 --- a/dimos/utils/test_data.py +++ b/dimos/utils/test_data.py @@ -26,7 +26,7 @@ def test_pull_file() -> None: repo_root = data._get_repo_root() test_file_name = "cafe.jpg" test_file_compressed = data._get_lfs_dir() / (test_file_name + ".tar.gz") - test_file_decompressed = data._get_data_dir() / test_file_name + test_file_decompressed = data.get_data_dir() / test_file_name # delete decompressed test file if it exists if test_file_decompressed.exists(): @@ -82,7 +82,7 @@ def test_pull_dir() -> None: repo_root = data._get_repo_root() test_dir_name = "ab_lidar_frames" test_dir_compressed = data._get_lfs_dir() / (test_dir_name + ".tar.gz") - test_dir_decompressed = data._get_data_dir() / test_dir_name + test_dir_decompressed = data.get_data_dir() / test_dir_name # delete decompressed test directory if it exists if test_dir_decompressed.exists(): diff --git a/dimos/utils/testing/moment.py b/dimos/utils/testing/moment.py index 436240a48b..e92d771687 100644 --- a/dimos/utils/testing/moment.py +++ b/dimos/utils/testing/moment.py @@ -17,12 +17,13 @@ from typing import TYPE_CHECKING, Any, Generic, TypeVar from dimos.core.resource import Resource +from dimos.types.timestamped import Timestamped from dimos.utils.testing.replay import TimedSensorReplay if TYPE_CHECKING: from dimos.core import Transport -T = TypeVar("T") +T = TypeVar("T", bound=Timestamped) class SensorMoment(Generic[T], Resource): diff --git a/dimos/utils/testing/replay.py b/dimos/utils/testing/replay.py index 89225c322e..588b63e099 100644 --- a/dimos/utils/testing/replay.py +++ b/dimos/utils/testing/replay.py @@ -11,399 +11,12 @@ # 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, Iterator -import functools -import glob -import os -from pathlib import Path -import pickle -import re -import time -from typing import Any, Generic, TypeVar -from reactivex import ( - from_iterable, - interval, - operators as ops, -) -from reactivex.observable import Observable -from reactivex.scheduler import TimeoutScheduler +"""Shim for TimedSensorReplay/TimedSensorStorage.""" -from dimos.utils.data import _get_data_dir, get_data +from dimos.memory.timeseries.legacy import LegacyPickleStore -T = TypeVar("T") - - -class SensorReplay(Generic[T]): - """Generic sensor data replay utility. - - Args: - name: The name of the test dataset - autocast: Optional function that takes unpickled data and returns a processed result. - For example: pointcloud2_from_webrtc_lidar - """ - - def __init__(self, name: str, autocast: Callable[[Any], T] | None = None) -> None: - self.root_dir = get_data(name) - self.autocast = autocast - - def load(self, *names: int | str) -> T | Any | list[T] | list[Any]: - if len(names) == 1: - return self.load_one(names[0]) - return list(map(lambda name: self.load_one(name), names)) - - def load_one(self, name: int | str | Path) -> T | Any: - if isinstance(name, int): - full_path = self.root_dir / f"/{name:03d}.pickle" - elif isinstance(name, Path): - full_path = name - else: - full_path = self.root_dir / Path(f"{name}.pickle") - - with open(full_path, "rb") as f: - data = pickle.load(f) - if self.autocast: - return self.autocast(data) - return data - - def first(self) -> T | Any | None: - try: - return next(self.iterate()) - except StopIteration: - return None - - @functools.cached_property - def files(self) -> list[Path]: - def extract_number(filepath): # type: ignore[no-untyped-def] - """Extract last digits before .pickle extension""" - basename = os.path.basename(filepath) - match = re.search(r"(\d+)\.pickle$", basename) - return int(match.group(1)) if match else 0 - - return sorted( - glob.glob(os.path.join(self.root_dir, "*")), # type: ignore[arg-type] - key=extract_number, - ) - - def iterate(self, loop: bool = False) -> Iterator[T | Any]: - while True: - for file_path in self.files: - yield self.load_one(Path(file_path)) - if not loop: - break - - def stream(self, rate_hz: float | None = None, loop: bool = False) -> Observable[T | Any]: - if rate_hz is None: - return from_iterable(self.iterate(loop=loop)) - - sleep_time = 1.0 / rate_hz - - return from_iterable(self.iterate(loop=loop)).pipe( - ops.zip(interval(sleep_time)), - ops.map(lambda x: x[0] if isinstance(x, tuple) else x), - ) - - -class SensorStorage(Generic[T]): - """Generic sensor data storage utility - . - Creates a directory in the test data directory and stores pickled sensor data. - - Args: - name: The name of the storage directory - autocast: Optional function that takes data and returns a processed result before storage. - """ - - def __init__(self, name: str, autocast: Callable[[T], Any] | None = None) -> None: - self.name = name - self.autocast = autocast - self.cnt = 0 - - # Create storage directory in the data dir - self.root_dir = _get_data_dir() / name - - # Check if directory exists and is not empty - if self.root_dir.exists(): - existing_files = list(self.root_dir.glob("*.pickle")) - if existing_files: - raise RuntimeError( - f"Storage directory '{name}' already exists and contains {len(existing_files)} files. " - f"Please use a different name or clean the directory first." - ) - else: - # Create the directory - self.root_dir.mkdir(parents=True, exist_ok=True) - - def consume_stream(self, observable: Observable[T | Any]) -> None: - """Consume an observable stream of sensor data without saving.""" - return observable.subscribe(self.save_one) # type: ignore[arg-type, return-value] - - def save_stream(self, observable: Observable[T | Any]) -> Observable[int]: - """Save an observable stream of sensor data to pickle files.""" - return observable.pipe(ops.map(lambda frame: self.save_one(frame))) - - def save(self, *frames) -> int: # type: ignore[no-untyped-def] - """Save one or more frames to pickle files.""" - for frame in frames: - self.save_one(frame) - return self.cnt - - def save_one(self, frame) -> int: # type: ignore[no-untyped-def] - """Save a single frame to a pickle file.""" - file_name = f"{self.cnt:03d}.pickle" - full_path = self.root_dir / file_name - - if full_path.exists(): - raise RuntimeError(f"File {full_path} already exists") - - # Apply autocast if provided - data_to_save = frame - if self.autocast: - data_to_save = self.autocast(frame) - # Convert to raw message if frame has a raw_msg attribute - elif hasattr(frame, "raw_msg"): - data_to_save = frame.raw_msg - - with open(full_path, "wb") as f: - pickle.dump(data_to_save, f) - - self.cnt += 1 - return self.cnt - - -class TimedSensorStorage(SensorStorage[T]): - def save_one(self, frame: T) -> int: - return super().save_one((time.time(), frame)) - - -class TimedSensorReplay(SensorReplay[T]): - def load_one(self, name: int | str | Path) -> T | Any: - if isinstance(name, int): - full_path = self.root_dir / f"/{name:03d}.pickle" - elif isinstance(name, Path): - full_path = name - else: - full_path = self.root_dir / Path(f"{name}.pickle") - - with open(full_path, "rb") as f: - data = pickle.load(f) - if self.autocast: - return (data[0], self.autocast(data[1])) - return data - - def find_closest(self, timestamp: float, tolerance: float | None = None) -> T | Any | None: - """Find the frame closest to the given timestamp. - - Args: - timestamp: The target timestamp to search for - tolerance: Optional maximum time difference allowed - - Returns: - The data frame closest to the timestamp, or None if no match within tolerance - """ - closest_data = None - closest_diff = float("inf") - - # Check frames before and after the timestamp - for ts, data in self.iterate_ts(): - diff = abs(ts - timestamp) - - if diff < closest_diff: - closest_diff = diff - closest_data = data - elif diff > closest_diff: - # We're moving away from the target, can stop - break - - if tolerance is not None and closest_diff > tolerance: - return None - - return closest_data - - def find_closest_seek( - self, relative_seconds: float, tolerance: float | None = None - ) -> T | Any | None: - """Find the frame closest to a time relative to the start. - - Args: - relative_seconds: Seconds from the start of the dataset - tolerance: Optional maximum time difference allowed - - Returns: - The data frame closest to the relative timestamp, or None if no match within tolerance - """ - # Get the first timestamp - first_ts = self.first_timestamp() - if first_ts is None: - return None - - # Calculate absolute timestamp and use find_closest - target_timestamp = first_ts + relative_seconds - return self.find_closest(target_timestamp, tolerance) - - def first_timestamp(self) -> float | None: - """Get the timestamp of the first item in the dataset. - - Returns: - The first timestamp, or None if dataset is empty - """ - try: - ts, _ = next(self.iterate_ts()) - return ts - except StopIteration: - return None - - def iterate(self, loop: bool = False) -> Iterator[T | Any]: - return (x[1] for x in super().iterate(loop=loop)) # type: ignore[index] - - def iterate_duration(self, **kwargs: Any) -> Iterator[tuple[float, T] | Any]: - """Iterate with timestamps relative to the start of the dataset.""" - first_ts = self.first_timestamp() - if first_ts is None: - return - for ts, data in self.iterate_ts(**kwargs): - yield (ts - first_ts, data) - - def iterate_realtime(self, speed: float = 1.0, **kwargs: Any) -> Iterator[T | Any]: - """Iterate data, sleeping to match original timing. - - Args: - speed: Playback speed multiplier (1.0 = realtime, 2.0 = 2x speed) - **kwargs: Passed to iterate_ts (seek, duration, from_timestamp, loop) - """ - iterator = self.iterate_ts(**kwargs) - - try: - first_ts, first_data = next(iterator) - except StopIteration: - return - - start_time = time.time() - start_ts = first_ts - yield first_data - - for ts, data in iterator: - target_time = start_time + (ts - start_ts) / speed - sleep_duration = target_time - time.time() - if sleep_duration > 0: - time.sleep(sleep_duration) - yield data - - def iterate_ts( - self, - seek: float | None = None, - duration: float | None = None, - from_timestamp: float | None = None, - loop: bool = False, - ) -> Iterator[tuple[float, T] | Any]: - """Iterate with absolute timestamps, with optional seek and duration.""" - first_ts = None - if (seek is not None) or (duration is not None): - first_ts = self.first_timestamp() - if first_ts is None: - return - - if seek is not None: - from_timestamp = first_ts + seek # type: ignore[operator] - - end_timestamp = None - if duration is not None: - end_timestamp = (from_timestamp if from_timestamp else first_ts) + duration # type: ignore[operator] - - while True: - for ts, data in super().iterate(): # type: ignore[misc] - if from_timestamp is None or ts >= from_timestamp: - if end_timestamp is not None and ts >= end_timestamp: - break - yield (ts, data) - if not loop: - break - - def stream( # type: ignore[override] - self, - speed: float = 1.0, - seek: float | None = None, - duration: float | None = None, - from_timestamp: float | None = None, - loop: bool = False, - ) -> Observable[T | Any]: - def _subscribe(observer, scheduler=None): # type: ignore[no-untyped-def] - from reactivex.disposable import CompositeDisposable, Disposable - - scheduler = scheduler or TimeoutScheduler() - disp = CompositeDisposable() - is_disposed = False - - iterator = self.iterate_ts( - seek=seek, duration=duration, from_timestamp=from_timestamp, loop=loop - ) - - # Get first message - try: - first_ts, first_data = next(iterator) - except StopIteration: - observer.on_completed() - return Disposable() - - # Establish timing reference - start_local_time = time.time() - start_replay_time = first_ts - - # Emit first sample immediately - observer.on_next(first_data) - - # Pre-load next message - try: - next_message = next(iterator) - except StopIteration: - observer.on_completed() - return disp - - def schedule_emission(message) -> None: # type: ignore[no-untyped-def] - nonlocal next_message, is_disposed - - if is_disposed: - return - - ts, data = message - - # Pre-load the following message while we have time - try: - next_message = next(iterator) - except StopIteration: - next_message = None - - # Calculate absolute emission time - target_time = start_local_time + (ts - start_replay_time) / speed - delay = max(0.0, target_time - time.time()) - - def emit() -> None: - if is_disposed: - return - observer.on_next(data) - if next_message is not None: - schedule_emission(next_message) - else: - observer.on_completed() - # Dispose of the scheduler to clean up threads - if hasattr(scheduler, "dispose"): - scheduler.dispose() - - scheduler.schedule_relative(delay, lambda sc, _: emit()) - - schedule_emission(next_message) - - # Create a custom disposable that properly cleans up - def dispose() -> None: - nonlocal is_disposed - is_disposed = True - disp.dispose() - # Ensure scheduler is disposed to clean up any threads - if hasattr(scheduler, "dispose"): - scheduler.dispose() - - return Disposable(dispose) - - from reactivex import create - - return create(_subscribe) +SensorReplay = LegacyPickleStore +SensorStorage = LegacyPickleStore +TimedSensorReplay = LegacyPickleStore +TimedSensorStorage = LegacyPickleStore diff --git a/dimos/utils/testing/test_replay.py b/dimos/utils/testing/test_replay.py index 640fe92979..8336c6a2b9 100644 --- a/dimos/utils/testing/test_replay.py +++ b/dimos/utils/testing/test_replay.py @@ -23,27 +23,9 @@ from dimos.utils.testing import replay -def test_sensor_replay() -> None: - counter = 0 - for message in replay.SensorReplay(name="office_lidar").iterate(): - counter += 1 - assert isinstance(message, dict) - assert counter == 500 - - -def test_sensor_replay_cast() -> None: - counter = 0 - for message in replay.SensorReplay( - name="office_lidar", autocast=pointcloud2_from_webrtc_lidar - ).iterate(): - counter += 1 - assert isinstance(message, PointCloud2) - assert counter == 500 - - def test_timed_sensor_replay() -> None: get_data("unitree_office_walk") - odom_store = replay.TimedSensorReplay("unitree_office_walk/odom", autocast=Odometry.from_msg) + odom_store = replay.TimedSensorReplay("unitree_office_walk/odom") itermsgs = [] for msg in odom_store.iterate(): @@ -87,7 +69,7 @@ def test_iterate_ts_no_seek() -> None: def test_iterate_ts_with_from_timestamp() -> None: """Test iterate_ts with from_timestamp (absolute timestamp)""" - odom_store = replay.TimedSensorReplay("unitree_office_walk/odom", autocast=Odometry.from_msg) + odom_store = replay.TimedSensorReplay("unitree_office_walk/odom") # First get all messages to find a good seek point all_msgs = [] @@ -115,7 +97,7 @@ def test_iterate_ts_with_from_timestamp() -> None: def test_iterate_ts_with_relative_seek() -> None: """Test iterate_ts with seek (relative seconds after first timestamp)""" - odom_store = replay.TimedSensorReplay("unitree_office_walk/odom", autocast=Odometry.from_msg) + odom_store = replay.TimedSensorReplay("unitree_office_walk/odom") # Get first few messages to understand timing all_msgs = [] @@ -144,7 +126,7 @@ def test_iterate_ts_with_relative_seek() -> None: def test_stream_with_seek() -> None: """Test stream method with seek parameters""" - odom_store = replay.TimedSensorReplay("unitree_office_walk/odom", autocast=Odometry.from_msg) + odom_store = replay.TimedSensorReplay("unitree_office_walk/odom") # Test stream with relative seek msgs_with_seek = [] @@ -170,7 +152,7 @@ def test_stream_with_seek() -> None: def test_duration_with_loop() -> None: """Test duration parameter with looping in TimedSensorReplay""" - odom_store = replay.TimedSensorReplay("unitree_office_walk/odom", autocast=Odometry.from_msg) + odom_store = replay.TimedSensorReplay("unitree_office_walk/odom") # Collect timestamps from a small duration window collected_ts = [] diff --git a/pyproject.toml b/pyproject.toml index e6b542a65e..aa8bb492be 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -253,11 +253,16 @@ dev = [ "types-tabulate>=0.9.0.20241207,<1", "types-tensorflow>=2.18.0.20251008,<3", "types-tqdm>=4.67.0.20250809,<5", + "types-psycopg2>=2.9.21.20251012", # Tools "py-spy", ] +psql = [ + "psycopg2-binary>=2.9.11" +] + sim = [ # Simulation "mujoco>=3.3.4", @@ -303,6 +308,7 @@ base = [ "dimos[agents,web,perception,visualization,sim]", ] + [tool.ruff] line-length = 100 exclude = [ diff --git a/uv.lock b/uv.lock index 75f839303f..c6d378b9f7 100644 --- a/uv.lock +++ b/uv.lock @@ -1850,6 +1850,7 @@ dev = [ { name = "types-networkx" }, { name = "types-protobuf" }, { name = "types-psutil" }, + { name = "types-psycopg2" }, { name = "types-pysocks" }, { name = "types-pytz" }, { name = "types-pyyaml" }, @@ -1937,6 +1938,9 @@ perception = [ { name = "transformers", extra = ["torch"] }, { name = "ultralytics" }, ] +psql = [ + { name = "psycopg2-binary" }, +] sim = [ { name = "mujoco" }, { name = "playground" }, @@ -2068,6 +2072,7 @@ requires-dist = [ { name = "plum-dispatch", specifier = "==2.5.7" }, { name = "plum-dispatch", marker = "extra == 'docker'", specifier = "==2.5.7" }, { name = "pre-commit", marker = "extra == 'dev'", specifier = "==4.2.0" }, + { name = "psycopg2-binary", marker = "extra == 'psql'", specifier = ">=2.9.11" }, { name = "py-spy", marker = "extra == 'dev'" }, { name = "pydantic" }, { name = "pydantic", marker = "extra == 'docker'" }, @@ -2133,6 +2138,7 @@ requires-dist = [ { name = "types-networkx", marker = "extra == 'dev'", specifier = ">=3.5.0.20251001,<4" }, { name = "types-protobuf", marker = "extra == 'dev'", specifier = ">=6.32.1.20250918,<7" }, { name = "types-psutil", marker = "extra == 'dev'", specifier = ">=7.0.0.20251001,<8" }, + { name = "types-psycopg2", marker = "extra == 'dev'", specifier = ">=2.9.21.20251012" }, { name = "types-pysocks", marker = "extra == 'dev'", specifier = ">=1.7.1.20251001,<2" }, { name = "types-pytz", marker = "extra == 'dev'", specifier = ">=2025.2.0.20250809,<2026" }, { name = "types-pyyaml", marker = "extra == 'dev'", specifier = ">=6.0.12.20250915,<7" }, @@ -2150,7 +2156,7 @@ requires-dist = [ { name = "xformers", marker = "extra == 'cuda'", specifier = ">=0.0.20" }, { name = "yapf", marker = "extra == 'misc'", specifier = "==0.40.2" }, ] -provides-extras = ["misc", "visualization", "agents", "web", "perception", "unitree", "manipulation", "cpu", "cuda", "dev", "sim", "drone", "docker", "base"] +provides-extras = ["misc", "visualization", "agents", "web", "perception", "unitree", "manipulation", "cpu", "cuda", "dev", "psql", "sim", "drone", "docker", "base"] [[package]] name = "dimos-lcm" @@ -6845,6 +6851,69 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/3e/73/2ce007f4198c80fcf2cb24c169884f833fe93fbc03d55d302627b094ee91/psutil-7.2.1-cp37-abi3-win_arm64.whl", hash = "sha256:0d67c1822c355aa6f7314d92018fb4268a76668a536f133599b91edd48759442", size = 133836, upload-time = "2025-12-29T08:26:43.086Z" }, ] +[[package]] +name = "psycopg2-binary" +version = "2.9.11" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/ac/6c/8767aaa597ba424643dc87348c6f1754dd9f48e80fdc1b9f7ca5c3a7c213/psycopg2-binary-2.9.11.tar.gz", hash = "sha256:b6aed9e096bf63f9e75edf2581aa9a7e7186d97ab5c177aa6c87797cd591236c", size = 379620, upload-time = "2025-10-10T11:14:48.041Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6a/f2/8e377d29c2ecf99f6062d35ea606b036e8800720eccfec5fe3dd672c2b24/psycopg2_binary-2.9.11-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:d6fe6b47d0b42ce1c9f1fa3e35bb365011ca22e39db37074458f27921dca40f2", size = 3756506, upload-time = "2025-10-10T11:10:30.144Z" }, + { url = "https://files.pythonhosted.org/packages/24/cc/dc143ea88e4ec9d386106cac05023b69668bd0be20794c613446eaefafe5/psycopg2_binary-2.9.11-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:a6c0e4262e089516603a09474ee13eabf09cb65c332277e39af68f6233911087", size = 3863943, upload-time = "2025-10-10T11:10:34.586Z" }, + { url = "https://files.pythonhosted.org/packages/8c/df/16848771155e7c419c60afeb24950b8aaa3ab09c0a091ec3ccca26a574d0/psycopg2_binary-2.9.11-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:c47676e5b485393f069b4d7a811267d3168ce46f988fa602658b8bb901e9e64d", size = 4410873, upload-time = "2025-10-10T11:10:38.951Z" }, + { url = "https://files.pythonhosted.org/packages/43/79/5ef5f32621abd5a541b89b04231fe959a9b327c874a1d41156041c75494b/psycopg2_binary-2.9.11-cp310-cp310-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:a28d8c01a7b27a1e3265b11250ba7557e5f72b5ee9e5f3a2fa8d2949c29bf5d2", size = 4468016, upload-time = "2025-10-10T11:10:43.319Z" }, + { url = "https://files.pythonhosted.org/packages/f0/9b/d7542d0f7ad78f57385971f426704776d7b310f5219ed58da5d605b1892e/psycopg2_binary-2.9.11-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5f3f2732cf504a1aa9e9609d02f79bea1067d99edf844ab92c247bbca143303b", size = 4164996, upload-time = "2025-10-10T11:10:46.705Z" }, + { url = "https://files.pythonhosted.org/packages/14/ed/e409388b537fa7414330687936917c522f6a77a13474e4238219fcfd9a84/psycopg2_binary-2.9.11-cp310-cp310-manylinux_2_38_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:865f9945ed1b3950d968ec4690ce68c55019d79e4497366d36e090327ce7db14", size = 3981881, upload-time = "2025-10-30T02:54:57.182Z" }, + { url = "https://files.pythonhosted.org/packages/bf/30/50e330e63bb05efc6fa7c1447df3e08954894025ca3dcb396ecc6739bc26/psycopg2_binary-2.9.11-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:91537a8df2bde69b1c1db01d6d944c831ca793952e4f57892600e96cee95f2cd", size = 3650857, upload-time = "2025-10-10T11:10:50.112Z" }, + { url = "https://files.pythonhosted.org/packages/f0/e0/4026e4c12bb49dd028756c5b0bc4c572319f2d8f1c9008e0dad8cc9addd7/psycopg2_binary-2.9.11-cp310-cp310-musllinux_1_2_ppc64le.whl", hash = "sha256:4dca1f356a67ecb68c81a7bc7809f1569ad9e152ce7fd02c2f2036862ca9f66b", size = 3296063, upload-time = "2025-10-10T11:10:54.089Z" }, + { url = "https://files.pythonhosted.org/packages/2c/34/eb172be293c886fef5299fe5c3fcf180a05478be89856067881007934a7c/psycopg2_binary-2.9.11-cp310-cp310-musllinux_1_2_riscv64.whl", hash = "sha256:0da4de5c1ac69d94ed4364b6cbe7190c1a70d325f112ba783d83f8440285f152", size = 3043464, upload-time = "2025-10-30T02:55:02.483Z" }, + { url = "https://files.pythonhosted.org/packages/18/1c/532c5d2cb11986372f14b798a95f2eaafe5779334f6a80589a68b5fcf769/psycopg2_binary-2.9.11-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:37d8412565a7267f7d79e29ab66876e55cb5e8e7b3bbf94f8206f6795f8f7e7e", size = 3345378, upload-time = "2025-10-10T11:11:01.039Z" }, + { url = "https://files.pythonhosted.org/packages/70/e7/de420e1cf16f838e1fa17b1120e83afff374c7c0130d088dba6286fcf8ea/psycopg2_binary-2.9.11-cp310-cp310-win_amd64.whl", hash = "sha256:c665f01ec8ab273a61c62beeb8cce3014c214429ced8a308ca1fc410ecac3a39", size = 2713904, upload-time = "2025-10-10T11:11:04.81Z" }, + { url = "https://files.pythonhosted.org/packages/c7/ae/8d8266f6dd183ab4d48b95b9674034e1b482a3f8619b33a0d86438694577/psycopg2_binary-2.9.11-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:0e8480afd62362d0a6a27dd09e4ca2def6fa50ed3a4e7c09165266106b2ffa10", size = 3756452, upload-time = "2025-10-10T11:11:11.583Z" }, + { url = "https://files.pythonhosted.org/packages/4b/34/aa03d327739c1be70e09d01182619aca8ebab5970cd0cfa50dd8b9cec2ac/psycopg2_binary-2.9.11-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:763c93ef1df3da6d1a90f86ea7f3f806dc06b21c198fa87c3c25504abec9404a", size = 3863957, upload-time = "2025-10-10T11:11:16.932Z" }, + { url = "https://files.pythonhosted.org/packages/48/89/3fdb5902bdab8868bbedc1c6e6023a4e08112ceac5db97fc2012060e0c9a/psycopg2_binary-2.9.11-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2e164359396576a3cc701ba8af4751ae68a07235d7a380c631184a611220d9a4", size = 4410955, upload-time = "2025-10-10T11:11:21.21Z" }, + { url = "https://files.pythonhosted.org/packages/ce/24/e18339c407a13c72b336e0d9013fbbbde77b6fd13e853979019a1269519c/psycopg2_binary-2.9.11-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:d57c9c387660b8893093459738b6abddbb30a7eab058b77b0d0d1c7d521ddfd7", size = 4468007, upload-time = "2025-10-10T11:11:24.831Z" }, + { url = "https://files.pythonhosted.org/packages/91/7e/b8441e831a0f16c159b5381698f9f7f7ed54b77d57bc9c5f99144cc78232/psycopg2_binary-2.9.11-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:2c226ef95eb2250974bf6fa7a842082b31f68385c4f3268370e3f3870e7859ee", size = 4165012, upload-time = "2025-10-10T11:11:29.51Z" }, + { url = "https://files.pythonhosted.org/packages/0d/61/4aa89eeb6d751f05178a13da95516c036e27468c5d4d2509bb1e15341c81/psycopg2_binary-2.9.11-cp311-cp311-manylinux_2_38_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:a311f1edc9967723d3511ea7d2708e2c3592e3405677bf53d5c7246753591fbb", size = 3981881, upload-time = "2025-10-30T02:55:07.332Z" }, + { url = "https://files.pythonhosted.org/packages/76/a1/2f5841cae4c635a9459fe7aca8ed771336e9383b6429e05c01267b0774cf/psycopg2_binary-2.9.11-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:ebb415404821b6d1c47353ebe9c8645967a5235e6d88f914147e7fd411419e6f", size = 3650985, upload-time = "2025-10-10T11:11:34.975Z" }, + { url = "https://files.pythonhosted.org/packages/84/74/4defcac9d002bca5709951b975173c8c2fa968e1a95dc713f61b3a8d3b6a/psycopg2_binary-2.9.11-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:f07c9c4a5093258a03b28fab9b4f151aa376989e7f35f855088234e656ee6a94", size = 3296039, upload-time = "2025-10-10T11:11:40.432Z" }, + { url = "https://files.pythonhosted.org/packages/6d/c2/782a3c64403d8ce35b5c50e1b684412cf94f171dc18111be8c976abd2de1/psycopg2_binary-2.9.11-cp311-cp311-musllinux_1_2_riscv64.whl", hash = "sha256:00ce1830d971f43b667abe4a56e42c1e2d594b32da4802e44a73bacacb25535f", size = 3043477, upload-time = "2025-10-30T02:55:11.182Z" }, + { url = "https://files.pythonhosted.org/packages/c8/31/36a1d8e702aa35c38fc117c2b8be3f182613faa25d794b8aeaab948d4c03/psycopg2_binary-2.9.11-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:cffe9d7697ae7456649617e8bb8d7a45afb71cd13f7ab22af3e5c61f04840908", size = 3345842, upload-time = "2025-10-10T11:11:45.366Z" }, + { url = "https://files.pythonhosted.org/packages/6e/b4/a5375cda5b54cb95ee9b836930fea30ae5a8f14aa97da7821722323d979b/psycopg2_binary-2.9.11-cp311-cp311-win_amd64.whl", hash = "sha256:304fd7b7f97eef30e91b8f7e720b3db75fee010b520e434ea35ed1ff22501d03", size = 2713894, upload-time = "2025-10-10T11:11:48.775Z" }, + { url = "https://files.pythonhosted.org/packages/d8/91/f870a02f51be4a65987b45a7de4c2e1897dd0d01051e2b559a38fa634e3e/psycopg2_binary-2.9.11-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:be9b840ac0525a283a96b556616f5b4820e0526addb8dcf6525a0fa162730be4", size = 3756603, upload-time = "2025-10-10T11:11:52.213Z" }, + { url = "https://files.pythonhosted.org/packages/27/fa/cae40e06849b6c9a95eb5c04d419942f00d9eaac8d81626107461e268821/psycopg2_binary-2.9.11-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:f090b7ddd13ca842ebfe301cd587a76a4cf0913b1e429eb92c1be5dbeb1a19bc", size = 3864509, upload-time = "2025-10-10T11:11:56.452Z" }, + { url = "https://files.pythonhosted.org/packages/2d/75/364847b879eb630b3ac8293798e380e441a957c53657995053c5ec39a316/psycopg2_binary-2.9.11-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:ab8905b5dcb05bf3fb22e0cf90e10f469563486ffb6a96569e51f897c750a76a", size = 4411159, upload-time = "2025-10-10T11:12:00.49Z" }, + { url = "https://files.pythonhosted.org/packages/6f/a0/567f7ea38b6e1c62aafd58375665a547c00c608a471620c0edc364733e13/psycopg2_binary-2.9.11-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:bf940cd7e7fec19181fdbc29d76911741153d51cab52e5c21165f3262125685e", size = 4468234, upload-time = "2025-10-10T11:12:04.892Z" }, + { url = "https://files.pythonhosted.org/packages/30/da/4e42788fb811bbbfd7b7f045570c062f49e350e1d1f3df056c3fb5763353/psycopg2_binary-2.9.11-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:fa0f693d3c68ae925966f0b14b8edda71696608039f4ed61b1fe9ffa468d16db", size = 4166236, upload-time = "2025-10-10T11:12:11.674Z" }, + { url = "https://files.pythonhosted.org/packages/3c/94/c1777c355bc560992af848d98216148be5f1be001af06e06fc49cbded578/psycopg2_binary-2.9.11-cp312-cp312-manylinux_2_38_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:a1cf393f1cdaf6a9b57c0a719a1068ba1069f022a59b8b1fe44b006745b59757", size = 3983083, upload-time = "2025-10-30T02:55:15.73Z" }, + { url = "https://files.pythonhosted.org/packages/bd/42/c9a21edf0e3daa7825ed04a4a8588686c6c14904344344a039556d78aa58/psycopg2_binary-2.9.11-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:ef7a6beb4beaa62f88592ccc65df20328029d721db309cb3250b0aae0fa146c3", size = 3652281, upload-time = "2025-10-10T11:12:17.713Z" }, + { url = "https://files.pythonhosted.org/packages/12/22/dedfbcfa97917982301496b6b5e5e6c5531d1f35dd2b488b08d1ebc52482/psycopg2_binary-2.9.11-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:31b32c457a6025e74d233957cc9736742ac5a6cb196c6b68499f6bb51390bd6a", size = 3298010, upload-time = "2025-10-10T11:12:22.671Z" }, + { url = "https://files.pythonhosted.org/packages/66/ea/d3390e6696276078bd01b2ece417deac954dfdd552d2edc3d03204416c0c/psycopg2_binary-2.9.11-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:edcb3aeb11cb4bf13a2af3c53a15b3d612edeb6409047ea0b5d6a21a9d744b34", size = 3044641, upload-time = "2025-10-30T02:55:19.929Z" }, + { url = "https://files.pythonhosted.org/packages/12/9a/0402ded6cbd321da0c0ba7d34dc12b29b14f5764c2fc10750daa38e825fc/psycopg2_binary-2.9.11-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:62b6d93d7c0b61a1dd6197d208ab613eb7dcfdcca0a49c42ceb082257991de9d", size = 3347940, upload-time = "2025-10-10T11:12:26.529Z" }, + { url = "https://files.pythonhosted.org/packages/b1/d2/99b55e85832ccde77b211738ff3925a5d73ad183c0b37bcbbe5a8ff04978/psycopg2_binary-2.9.11-cp312-cp312-win_amd64.whl", hash = "sha256:b33fabeb1fde21180479b2d4667e994de7bbf0eec22832ba5d9b5e4cf65b6c6d", size = 2714147, upload-time = "2025-10-10T11:12:29.535Z" }, + { url = "https://files.pythonhosted.org/packages/ff/a8/a2709681b3ac11b0b1786def10006b8995125ba268c9a54bea6f5ae8bd3e/psycopg2_binary-2.9.11-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:b8fb3db325435d34235b044b199e56cdf9ff41223a4b9752e8576465170bb38c", size = 3756572, upload-time = "2025-10-10T11:12:32.873Z" }, + { url = "https://files.pythonhosted.org/packages/62/e1/c2b38d256d0dafd32713e9f31982a5b028f4a3651f446be70785f484f472/psycopg2_binary-2.9.11-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:366df99e710a2acd90efed3764bb1e28df6c675d33a7fb40df9b7281694432ee", size = 3864529, upload-time = "2025-10-10T11:12:36.791Z" }, + { url = "https://files.pythonhosted.org/packages/11/32/b2ffe8f3853c181e88f0a157c5fb4e383102238d73c52ac6d93a5c8bffe6/psycopg2_binary-2.9.11-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:8c55b385daa2f92cb64b12ec4536c66954ac53654c7f15a203578da4e78105c0", size = 4411242, upload-time = "2025-10-10T11:12:42.388Z" }, + { url = "https://files.pythonhosted.org/packages/10/04/6ca7477e6160ae258dc96f67c371157776564679aefd247b66f4661501a2/psycopg2_binary-2.9.11-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:c0377174bf1dd416993d16edc15357f6eb17ac998244cca19bc67cdc0e2e5766", size = 4468258, upload-time = "2025-10-10T11:12:48.654Z" }, + { url = "https://files.pythonhosted.org/packages/3c/7e/6a1a38f86412df101435809f225d57c1a021307dd0689f7a5e7fe83588b1/psycopg2_binary-2.9.11-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5c6ff3335ce08c75afaed19e08699e8aacf95d4a260b495a4a8545244fe2ceb3", size = 4166295, upload-time = "2025-10-10T11:12:52.525Z" }, + { url = "https://files.pythonhosted.org/packages/f2/7d/c07374c501b45f3579a9eb761cbf2604ddef3d96ad48679112c2c5aa9c25/psycopg2_binary-2.9.11-cp313-cp313-manylinux_2_38_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:84011ba3109e06ac412f95399b704d3d6950e386b7994475b231cf61eec2fc1f", size = 3983133, upload-time = "2025-10-30T02:55:24.329Z" }, + { url = "https://files.pythonhosted.org/packages/82/56/993b7104cb8345ad7d4516538ccf8f0d0ac640b1ebd8c754a7b024e76878/psycopg2_binary-2.9.11-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:ba34475ceb08cccbdd98f6b46916917ae6eeb92b5ae111df10b544c3a4621dc4", size = 3652383, upload-time = "2025-10-10T11:12:56.387Z" }, + { url = "https://files.pythonhosted.org/packages/2d/ac/eaeb6029362fd8d454a27374d84c6866c82c33bfc24587b4face5a8e43ef/psycopg2_binary-2.9.11-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:b31e90fdd0f968c2de3b26ab014314fe814225b6c324f770952f7d38abf17e3c", size = 3298168, upload-time = "2025-10-10T11:13:00.403Z" }, + { url = "https://files.pythonhosted.org/packages/2b/39/50c3facc66bded9ada5cbc0de867499a703dc6bca6be03070b4e3b65da6c/psycopg2_binary-2.9.11-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:d526864e0f67f74937a8fce859bd56c979f5e2ec57ca7c627f5f1071ef7fee60", size = 3044712, upload-time = "2025-10-30T02:55:27.975Z" }, + { url = "https://files.pythonhosted.org/packages/9c/8e/b7de019a1f562f72ada81081a12823d3c1590bedc48d7d2559410a2763fe/psycopg2_binary-2.9.11-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:04195548662fa544626c8ea0f06561eb6203f1984ba5b4562764fbeb4c3d14b1", size = 3347549, upload-time = "2025-10-10T11:13:03.971Z" }, + { url = "https://files.pythonhosted.org/packages/80/2d/1bb683f64737bbb1f86c82b7359db1eb2be4e2c0c13b947f80efefa7d3e5/psycopg2_binary-2.9.11-cp313-cp313-win_amd64.whl", hash = "sha256:efff12b432179443f54e230fdf60de1f6cc726b6c832db8701227d089310e8aa", size = 2714215, upload-time = "2025-10-10T11:13:07.14Z" }, + { url = "https://files.pythonhosted.org/packages/64/12/93ef0098590cf51d9732b4f139533732565704f45bdc1ffa741b7c95fb54/psycopg2_binary-2.9.11-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:92e3b669236327083a2e33ccfa0d320dd01b9803b3e14dd986a4fc54aa00f4e1", size = 3756567, upload-time = "2025-10-10T11:13:11.885Z" }, + { url = "https://files.pythonhosted.org/packages/7c/a9/9d55c614a891288f15ca4b5209b09f0f01e3124056924e17b81b9fa054cc/psycopg2_binary-2.9.11-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:e0deeb03da539fa3577fcb0b3f2554a97f7e5477c246098dbb18091a4a01c16f", size = 3864755, upload-time = "2025-10-10T11:13:17.727Z" }, + { url = "https://files.pythonhosted.org/packages/13/1e/98874ce72fd29cbde93209977b196a2edae03f8490d1bd8158e7f1daf3a0/psycopg2_binary-2.9.11-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:9b52a3f9bb540a3e4ec0f6ba6d31339727b2950c9772850d6545b7eae0b9d7c5", size = 4411646, upload-time = "2025-10-10T11:13:24.432Z" }, + { url = "https://files.pythonhosted.org/packages/5a/bd/a335ce6645334fb8d758cc358810defca14a1d19ffbc8a10bd38a2328565/psycopg2_binary-2.9.11-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:db4fd476874ccfdbb630a54426964959e58da4c61c9feba73e6094d51303d7d8", size = 4468701, upload-time = "2025-10-10T11:13:29.266Z" }, + { url = "https://files.pythonhosted.org/packages/44/d6/c8b4f53f34e295e45709b7568bf9b9407a612ea30387d35eb9fa84f269b4/psycopg2_binary-2.9.11-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:47f212c1d3be608a12937cc131bd85502954398aaa1320cb4c14421a0ffccf4c", size = 4166293, upload-time = "2025-10-10T11:13:33.336Z" }, + { url = "https://files.pythonhosted.org/packages/4b/e0/f8cc36eadd1b716ab36bb290618a3292e009867e5c97ce4aba908cb99644/psycopg2_binary-2.9.11-cp314-cp314-manylinux_2_38_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:e35b7abae2b0adab776add56111df1735ccc71406e56203515e228a8dc07089f", size = 3983184, upload-time = "2025-10-30T02:55:32.483Z" }, + { url = "https://files.pythonhosted.org/packages/53/3e/2a8fe18a4e61cfb3417da67b6318e12691772c0696d79434184a511906dc/psycopg2_binary-2.9.11-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:fcf21be3ce5f5659daefd2b3b3b6e4727b028221ddc94e6c1523425579664747", size = 3652650, upload-time = "2025-10-10T11:13:38.181Z" }, + { url = "https://files.pythonhosted.org/packages/76/36/03801461b31b29fe58d228c24388f999fe814dfc302856e0d17f97d7c54d/psycopg2_binary-2.9.11-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:9bd81e64e8de111237737b29d68039b9c813bdf520156af36d26819c9a979e5f", size = 3298663, upload-time = "2025-10-10T11:13:44.878Z" }, + { url = "https://files.pythonhosted.org/packages/97/77/21b0ea2e1a73aa5fa9222b2a6b8ba325c43c3a8d54272839c991f2345656/psycopg2_binary-2.9.11-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:32770a4d666fbdafab017086655bcddab791d7cb260a16679cc5a7338b64343b", size = 3044737, upload-time = "2025-10-30T02:55:35.69Z" }, + { url = "https://files.pythonhosted.org/packages/67/69/f36abe5f118c1dca6d3726ceae164b9356985805480731ac6712a63f24f0/psycopg2_binary-2.9.11-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:c3cb3a676873d7506825221045bd70e0427c905b9c8ee8d6acd70cfcbd6e576d", size = 3347643, upload-time = "2025-10-10T11:13:53.499Z" }, + { url = "https://files.pythonhosted.org/packages/e1/36/9c0c326fe3a4227953dfb29f5d0c8ae3b8eb8c1cd2967aa569f50cb3c61f/psycopg2_binary-2.9.11-cp314-cp314-win_amd64.whl", hash = "sha256:4012c9c954dfaccd28f94e84ab9f94e12df76b4afb22331b1f0d3154893a6316", size = 2803913, upload-time = "2025-10-10T11:13:57.058Z" }, +] + [[package]] name = "ptyprocess" version = "0.7.0" @@ -9819,6 +9888,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/12/61/81f180ffbcd0b3516fa3e0e95588dcd48200b6a08e3df53c6c0941a688fe/types_psutil-7.2.1.20251231-py3-none-any.whl", hash = "sha256:40735ca2fc818aed9dcbff7acb3317a774896615e3f4a7bd356afa224b9178e3", size = 32426, upload-time = "2025-12-31T03:18:28.14Z" }, ] +[[package]] +name = "types-psycopg2" +version = "2.9.21.20251012" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/9b/b3/2d09eaf35a084cffd329c584970a3fa07101ca465c13cad1576d7c392587/types_psycopg2-2.9.21.20251012.tar.gz", hash = "sha256:4cdafd38927da0cfde49804f39ab85afd9c6e9c492800e42f1f0c1a1b0312935", size = 26710, upload-time = "2025-10-12T02:55:39.5Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ec/0c/05feaf8cb51159f2c0af04b871dab7e98a2f83a3622f5f216331d2dd924c/types_psycopg2-2.9.21.20251012-py3-none-any.whl", hash = "sha256:712bad5c423fe979e357edbf40a07ca40ef775d74043de72bd4544ca328cc57e", size = 24883, upload-time = "2025-10-12T02:55:38.439Z" }, +] + [[package]] name = "types-pysocks" version = "1.7.1.20251001"