From 5f0f19426633bb2045e2921078c850be9303778b Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Wed, 21 Jan 2026 17:17:29 +0800 Subject: [PATCH 01/39] embeddings not generic, lsp added to dev deps --- dimos/models/embedding/__init__.py | 12 +- dimos/models/embedding/base.py | 17 +- dimos/models/embedding/clip.py | 17 +- dimos/models/embedding/mobileclip.py | 13 +- dimos/models/embedding/treid.py | 11 +- .../detection/reid/embedding_id_system.py | 2 +- pyproject.toml | 4 + uv.lock | 438 ++++++++++++++++++ 8 files changed, 472 insertions(+), 42 deletions(-) diff --git a/dimos/models/embedding/__init__.py b/dimos/models/embedding/__init__.py index 981e25e5c2..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__.extend(["CLIPEmbedding", "CLIPModel"]) + __all__.append("CLIPModel") except ImportError: pass # Optional: MobileCLIP support try: - from dimos.models.embedding.mobileclip import MobileCLIPEmbedding, MobileCLIPModel + from dimos.models.embedding.mobileclip import MobileCLIPModel - __all__.extend(["MobileCLIPEmbedding", "MobileCLIPModel"]) + __all__.append("MobileCLIPModel") except ImportError: pass # Optional: TorchReID support try: - from dimos.models.embedding.treid import TorchReIDEmbedding, TorchReIDModel + from dimos.models.embedding.treid import TorchReIDModel - __all__.extend(["TorchReIDEmbedding", "TorchReIDModel"]) + __all__.append("TorchReIDModel") except ImportError: pass diff --git a/dimos/models/embedding/base.py b/dimos/models/embedding/base.py index eba5e45894..7242a5efc6 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,7 @@ 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 +141,7 @@ 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). diff --git a/dimos/models/embedding/clip.py b/dimos/models/embedding/clip.py index d8a62efcb2..282ca7fa33 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 = F.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 = F.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 a8893d38e4..07f4f272c1 100644 --- a/dimos/models/embedding/treid.py +++ b/dimos/models/embedding/treid.py @@ -25,9 +25,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 @@ -36,7 +33,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 @@ -51,7 +48,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. @@ -76,11 +73,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/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/pyproject.toml b/pyproject.toml index 91b1a1288d..e04dd4bbfe 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -226,6 +226,10 @@ dev = [ # docs "md-babel-py==1.1.1", + # lsp + "python-lsp-server[all]==1.14.0", + "python-lsp-ruff==2.3.0", + # Types "lxml-stubs>=0.5.1,<1", "pandas-stubs>=2.3.2.250926,<3", diff --git a/uv.lock b/uv.lock index b1c2bece00..486e216e20 100644 --- a/uv.lock +++ b/uv.lock @@ -163,6 +163,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/81/29/5ecc3a15d5a33e31b26c11426c45c501e439cb865d0bff96315d86443b78/appnope-0.1.4-py2.py3-none-any.whl", hash = "sha256:502575ee11cd7a28c0205f379b525beefebab9d161b7c964670864014ed7213c", size = 4321, upload-time = "2024-02-06T09:43:09.663Z" }, ] +[[package]] +name = "astroid" +version = "4.0.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions", marker = "python_full_version < '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a1/ca/c17d0f83016532a1ad87d1de96837164c99d47a3b6bbba28bd597c25b37a/astroid-4.0.3.tar.gz", hash = "sha256:08d1de40d251cc3dc4a7a12726721d475ac189e4e583d596ece7422bc176bda3", size = 406224, upload-time = "2026-01-03T22:14:26.096Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ce/66/686ac4fc6ef48f5bacde625adac698f41d5316a9753c2b20bb0931c9d4e2/astroid-4.0.3-py3-none-any.whl", hash = "sha256:864a0a34af1bd70e1049ba1e61cee843a7252c826d97825fcee9b2fcbd9e1b14", size = 276443, upload-time = "2026-01-03T22:14:24.412Z" }, +] + [[package]] name = "asttokens" version = "3.0.1" @@ -190,6 +202,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/3a/2a/7cc015f5b9f5db42b7d48157e23356022889fc354a2813c15934b7cb5c0e/attrs-25.4.0-py3-none-any.whl", hash = "sha256:adcf7e2a1fb3b36ac48d97835bb6d8ade15b8dcce26aba8bf1d14847b57a3373", size = 67615, upload-time = "2025-10-06T13:54:43.17Z" }, ] +[[package]] +name = "autopep8" +version = "2.0.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pycodestyle" }, + { name = "tomli", marker = "python_full_version < '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e0/8a/9be661f5400867a09706e29f5ab99a59987fd3a4c337757365e7491fa90b/autopep8-2.0.4.tar.gz", hash = "sha256:2913064abd97b3419d1cc83ea71f042cb821f87e45b9c88cad5ad3c4ea87fe0c", size = 116472, upload-time = "2023-08-26T13:49:59.375Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d8/f2/e63c9f9c485cd90df8e4e7ae90fa3be2469c9641888558c7b45fa98a76f8/autopep8-2.0.4-py2.py3-none-any.whl", hash = "sha256:067959ca4a07b24dbd5345efa8325f5f58da4298dab0dde0443d5ed765de80cb", size = 45340, upload-time = "2023-08-26T13:49:56.111Z" }, +] + [[package]] name = "av" version = "16.0.1" @@ -372,6 +397,50 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/1d/4f/02d3cb62a1b0b5a1ca7ff03dce3606be1bf3ead4744f47eb762dbf471069/bitsandbytes-0.49.1-py3-none-manylinux_2_24_x86_64.whl", hash = "sha256:e7940bf32457dc2e553685285b2a86e82f5ec10b2ae39776c408714f9ae6983c", size = 59054193, upload-time = "2026-01-08T14:31:31.743Z" }, ] +[[package]] +name = "black" +version = "25.12.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "click" }, + { name = "mypy-extensions" }, + { name = "packaging" }, + { name = "pathspec" }, + { name = "platformdirs" }, + { name = "pytokens" }, + { name = "tomli", marker = "python_full_version < '3.11'" }, + { name = "typing-extensions", marker = "python_full_version < '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c4/d9/07b458a3f1c525ac392b5edc6b191ff140b596f9d77092429417a54e249d/black-25.12.0.tar.gz", hash = "sha256:8d3dd9cea14bff7ddc0eb243c811cdb1a011ebb4800a5f0335a01a68654796a7", size = 659264, upload-time = "2025-12-08T01:40:52.501Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/37/d5/8d3145999d380e5d09bb00b0f7024bf0a8ccb5c07b5648e9295f02ec1d98/black-25.12.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:f85ba1ad15d446756b4ab5f3044731bf68b777f8f9ac9cdabd2425b97cd9c4e8", size = 1895720, upload-time = "2025-12-08T01:46:58.197Z" }, + { url = "https://files.pythonhosted.org/packages/06/97/7acc85c4add41098f4f076b21e3e4e383ad6ed0a3da26b2c89627241fc11/black-25.12.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:546eecfe9a3a6b46f9d69d8a642585a6eaf348bcbbc4d87a19635570e02d9f4a", size = 1727193, upload-time = "2025-12-08T01:52:26.674Z" }, + { url = "https://files.pythonhosted.org/packages/24/f0/fdf0eb8ba907ddeb62255227d29d349e8256ef03558fbcadfbc26ecfe3b2/black-25.12.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:17dcc893da8d73d8f74a596f64b7c98ef5239c2cd2b053c0f25912c4494bf9ea", size = 1774506, upload-time = "2025-12-08T01:46:25.721Z" }, + { url = "https://files.pythonhosted.org/packages/e4/f5/9203a78efe00d13336786b133c6180a9303d46908a9aa72d1104ca214222/black-25.12.0-cp310-cp310-win_amd64.whl", hash = "sha256:09524b0e6af8ba7a3ffabdfc7a9922fb9adef60fed008c7cd2fc01f3048e6e6f", size = 1416085, upload-time = "2025-12-08T01:46:06.073Z" }, + { url = "https://files.pythonhosted.org/packages/ba/cc/7a6090e6b081c3316282c05c546e76affdce7bf7a3b7d2c3a2a69438bd01/black-25.12.0-cp310-cp310-win_arm64.whl", hash = "sha256:b162653ed89eb942758efeb29d5e333ca5bb90e5130216f8369857db5955a7da", size = 1226038, upload-time = "2025-12-08T01:45:29.388Z" }, + { url = "https://files.pythonhosted.org/packages/60/ad/7ac0d0e1e0612788dbc48e62aef8a8e8feffac7eb3d787db4e43b8462fa8/black-25.12.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:d0cfa263e85caea2cff57d8f917f9f51adae8e20b610e2b23de35b5b11ce691a", size = 1877003, upload-time = "2025-12-08T01:43:29.967Z" }, + { url = "https://files.pythonhosted.org/packages/e8/dd/a237e9f565f3617a88b49284b59cbca2a4f56ebe68676c1aad0ce36a54a7/black-25.12.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:1a2f578ae20c19c50a382286ba78bfbeafdf788579b053d8e4980afb079ab9be", size = 1712639, upload-time = "2025-12-08T01:52:46.756Z" }, + { url = "https://files.pythonhosted.org/packages/12/80/e187079df1ea4c12a0c63282ddd8b81d5107db6d642f7d7b75a6bcd6fc21/black-25.12.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d3e1b65634b0e471d07ff86ec338819e2ef860689859ef4501ab7ac290431f9b", size = 1758143, upload-time = "2025-12-08T01:45:29.137Z" }, + { url = "https://files.pythonhosted.org/packages/93/b5/3096ccee4f29dc2c3aac57274326c4d2d929a77e629f695f544e159bfae4/black-25.12.0-cp311-cp311-win_amd64.whl", hash = "sha256:a3fa71e3b8dd9f7c6ac4d818345237dfb4175ed3bf37cd5a581dbc4c034f1ec5", size = 1420698, upload-time = "2025-12-08T01:45:53.379Z" }, + { url = "https://files.pythonhosted.org/packages/7e/39/f81c0ffbc25ffbe61c7d0385bf277e62ffc3e52f5ee668d7369d9854fadf/black-25.12.0-cp311-cp311-win_arm64.whl", hash = "sha256:51e267458f7e650afed8445dc7edb3187143003d52a1b710c7321aef22aa9655", size = 1229317, upload-time = "2025-12-08T01:46:35.606Z" }, + { url = "https://files.pythonhosted.org/packages/d1/bd/26083f805115db17fda9877b3c7321d08c647df39d0df4c4ca8f8450593e/black-25.12.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:31f96b7c98c1ddaeb07dc0f56c652e25bdedaac76d5b68a059d998b57c55594a", size = 1924178, upload-time = "2025-12-08T01:49:51.048Z" }, + { url = "https://files.pythonhosted.org/packages/89/6b/ea00d6651561e2bdd9231c4177f4f2ae19cc13a0b0574f47602a7519b6ca/black-25.12.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:05dd459a19e218078a1f98178c13f861fe6a9a5f88fc969ca4d9b49eb1809783", size = 1742643, upload-time = "2025-12-08T01:49:59.09Z" }, + { url = "https://files.pythonhosted.org/packages/6d/f3/360fa4182e36e9875fabcf3a9717db9d27a8d11870f21cff97725c54f35b/black-25.12.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c1f68c5eff61f226934be6b5b80296cf6939e5d2f0c2f7d543ea08b204bfaf59", size = 1800158, upload-time = "2025-12-08T01:44:27.301Z" }, + { url = "https://files.pythonhosted.org/packages/f8/08/2c64830cb6616278067e040acca21d4f79727b23077633953081c9445d61/black-25.12.0-cp312-cp312-win_amd64.whl", hash = "sha256:274f940c147ddab4442d316b27f9e332ca586d39c85ecf59ebdea82cc9ee8892", size = 1426197, upload-time = "2025-12-08T01:45:51.198Z" }, + { url = "https://files.pythonhosted.org/packages/d4/60/a93f55fd9b9816b7432cf6842f0e3000fdd5b7869492a04b9011a133ee37/black-25.12.0-cp312-cp312-win_arm64.whl", hash = "sha256:169506ba91ef21e2e0591563deda7f00030cb466e747c4b09cb0a9dae5db2f43", size = 1237266, upload-time = "2025-12-08T01:45:10.556Z" }, + { url = "https://files.pythonhosted.org/packages/c8/52/c551e36bc95495d2aa1a37d50566267aa47608c81a53f91daa809e03293f/black-25.12.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:a05ddeb656534c3e27a05a29196c962877c83fa5503db89e68857d1161ad08a5", size = 1923809, upload-time = "2025-12-08T01:46:55.126Z" }, + { url = "https://files.pythonhosted.org/packages/a0/f7/aac9b014140ee56d247e707af8db0aae2e9efc28d4a8aba92d0abd7ae9d1/black-25.12.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:9ec77439ef3e34896995503865a85732c94396edcc739f302c5673a2315e1e7f", size = 1742384, upload-time = "2025-12-08T01:49:37.022Z" }, + { url = "https://files.pythonhosted.org/packages/74/98/38aaa018b2ab06a863974c12b14a6266badc192b20603a81b738c47e902e/black-25.12.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0e509c858adf63aa61d908061b52e580c40eae0dfa72415fa47ac01b12e29baf", size = 1798761, upload-time = "2025-12-08T01:46:05.386Z" }, + { url = "https://files.pythonhosted.org/packages/16/3a/a8ac542125f61574a3f015b521ca83b47321ed19bb63fe6d7560f348bfe1/black-25.12.0-cp313-cp313-win_amd64.whl", hash = "sha256:252678f07f5bac4ff0d0e9b261fbb029fa530cfa206d0a636a34ab445ef8ca9d", size = 1429180, upload-time = "2025-12-08T01:45:34.903Z" }, + { url = "https://files.pythonhosted.org/packages/e6/2d/bdc466a3db9145e946762d52cd55b1385509d9f9004fec1c97bdc8debbfb/black-25.12.0-cp313-cp313-win_arm64.whl", hash = "sha256:bc5b1c09fe3c931ddd20ee548511c64ebf964ada7e6f0763d443947fd1c603ce", size = 1239350, upload-time = "2025-12-08T01:46:09.458Z" }, + { url = "https://files.pythonhosted.org/packages/35/46/1d8f2542210c502e2ae1060b2e09e47af6a5e5963cb78e22ec1a11170b28/black-25.12.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:0a0953b134f9335c2434864a643c842c44fba562155c738a2a37a4d61f00cad5", size = 1917015, upload-time = "2025-12-08T01:53:27.987Z" }, + { url = "https://files.pythonhosted.org/packages/41/37/68accadf977672beb8e2c64e080f568c74159c1aaa6414b4cd2aef2d7906/black-25.12.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:2355bbb6c3b76062870942d8cc450d4f8ac71f9c93c40122762c8784df49543f", size = 1741830, upload-time = "2025-12-08T01:54:36.861Z" }, + { url = "https://files.pythonhosted.org/packages/ac/76/03608a9d8f0faad47a3af3a3c8c53af3367f6c0dd2d23a84710456c7ac56/black-25.12.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9678bd991cc793e81d19aeeae57966ee02909877cb65838ccffef24c3ebac08f", size = 1791450, upload-time = "2025-12-08T01:44:52.581Z" }, + { url = "https://files.pythonhosted.org/packages/06/99/b2a4bd7dfaea7964974f947e1c76d6886d65fe5d24f687df2d85406b2609/black-25.12.0-cp314-cp314-win_amd64.whl", hash = "sha256:97596189949a8aad13ad12fcbb4ae89330039b96ad6742e6f6b45e75ad5cfd83", size = 1452042, upload-time = "2025-12-08T01:46:13.188Z" }, + { url = "https://files.pythonhosted.org/packages/b2/7c/d9825de75ae5dd7795d007681b752275ea85a1c5d83269b4b9c754c2aaab/black-25.12.0-cp314-cp314-win_arm64.whl", hash = "sha256:778285d9ea197f34704e3791ea9404cd6d07595745907dd2ce3da7a13627b29b", size = 1267446, upload-time = "2025-12-08T01:46:14.497Z" }, + { url = "https://files.pythonhosted.org/packages/68/11/21331aed19145a952ad28fca2756a1433ee9308079bd03bd898e903a2e53/black-25.12.0-py3-none-any.whl", hash = "sha256:48ceb36c16dbc84062740049eef990bb2ce07598272e673c17d1a7720c71c828", size = 206191, upload-time = "2025-12-08T01:40:50.963Z" }, +] + [[package]] name = "blinker" version = "1.9.0" @@ -472,6 +541,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/99/1b/50316bd6f95c50686b35799abebb6168d90ee18b7c03e3065f587f010f7c/catkin_pkg-1.1.0-py3-none-any.whl", hash = "sha256:7f5486b4f5681b5f043316ce10fc638c8d0ba8127146e797c85f4024e4356027", size = 76369, upload-time = "2025-09-10T17:34:35.639Z" }, ] +[[package]] +name = "cattrs" +version = "25.3.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "attrs" }, + { name = "exceptiongroup", marker = "python_full_version < '3.11'" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/6e/00/2432bb2d445b39b5407f0a90e01b9a271475eea7caf913d7a86bcb956385/cattrs-25.3.0.tar.gz", hash = "sha256:1ac88d9e5eda10436c4517e390a4142d88638fe682c436c93db7ce4a277b884a", size = 509321, upload-time = "2025-10-07T12:26:08.737Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d8/2b/a40e1488fdfa02d3f9a653a61a5935ea08b3c2225ee818db6a76c7ba9695/cattrs-25.3.0-py3-none-any.whl", hash = "sha256:9896e84e0a5bf723bc7b4b68f4481785367ce07a8a02e7e9ee6eb2819bc306ff", size = 70738, upload-time = "2025-10-07T12:26:06.603Z" }, +] + [[package]] name = "cerebras-cloud-sdk" version = "1.64.1" @@ -1379,6 +1462,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/4e/8c/f3147f5c4b73e7550fe5f9352eaa956ae838d5c51eb58e7a25b9f3e2643b/decorator-5.2.1-py3-none-any.whl", hash = "sha256:d316bb415a2d9e2d2b3abcc4084c6502fc09240e292cd76a76afc106a1c8e04a", size = 9190, upload-time = "2025-02-24T04:41:32.565Z" }, ] +[[package]] +name = "dill" +version = "0.4.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/81/e1/56027a71e31b02ddc53c7d65b01e68edf64dea2932122fe7746a516f75d5/dill-0.4.1.tar.gz", hash = "sha256:423092df4182177d4d8ba8290c8a5b640c66ab35ec7da59ccfa00f6fa3eea5fa", size = 187315, upload-time = "2026-01-19T02:36:56.85Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1e/77/dc8c558f7593132cf8fefec57c4f60c83b16941c574ac5f619abb3ae7933/dill-0.4.1-py3-none-any.whl", hash = "sha256:1e1ce33e978ae97fcfcff5638477032b801c46c7c65cf717f95fbc2248f79a9d", size = 120019, upload-time = "2026-01-19T02:36:55.663Z" }, +] + [[package]] name = "dimos" version = "0.0.8" @@ -1484,6 +1576,8 @@ dev = [ { name = "pytest-env" }, { name = "pytest-mock" }, { name = "pytest-timeout" }, + { name = "python-lsp-ruff" }, + { name = "python-lsp-server", extra = ["all"] }, { name = "requests-mock" }, { name = "ruff" }, { name = "terminaltexteffects" }, @@ -1692,6 +1786,8 @@ requires-dist = [ { name = "pytest-timeout", marker = "extra == 'dev'", specifier = "==2.4.0" }, { name = "python-dotenv" }, { name = "python-fcl", marker = "extra == 'manipulation'", specifier = ">=0.7.0.4" }, + { name = "python-lsp-ruff", marker = "extra == 'dev'", specifier = "==2.3.0" }, + { name = "python-lsp-server", extras = ["all"], marker = "extra == 'dev'", specifier = "==1.14.0" }, { name = "python-multipart", marker = "extra == 'misc'", specifier = "==0.0.20" }, { name = "pyturbojpeg", specifier = "==1.8.2" }, { name = "pyyaml", marker = "extra == 'manipulation'", specifier = ">=6.0" }, @@ -1825,6 +1921,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/55/e2/2537ebcff11c1ee1ff17d8d0b6f4db75873e3b0fb32c2d4a2ee31ecb310a/docstring_parser-0.17.0-py3-none-any.whl", hash = "sha256:cf2569abd23dce8099b300f9b4fa8191e9582dda731fd533daf54c4551658708", size = 36896, upload-time = "2025-07-21T07:35:00.684Z" }, ] +[[package]] +name = "docstring-to-markdown" +version = "0.17" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "importlib-metadata" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/52/d8/8abe80d62c5dce1075578031bcfde07e735bcf0afe2886dd48b470162ab4/docstring_to_markdown-0.17.tar.gz", hash = "sha256:df72a112294c7492487c9da2451cae0faeee06e86008245c188c5761c9590ca3", size = 32260, upload-time = "2025-05-02T15:09:07.932Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/56/7b/af3d0da15bed3a8665419bb3a630585756920f4ad67abfdfef26240ebcc0/docstring_to_markdown-0.17-py3-none-any.whl", hash = "sha256:fd7d5094aa83943bf5f9e1a13701866b7c452eac19765380dead666e36d3711c", size = 23479, upload-time = "2025-05-02T15:09:06.676Z" }, +] + [[package]] name = "docutils" version = "0.22.4" @@ -2124,6 +2233,20 @@ dependencies = [ ] sdist = { url = "https://files.pythonhosted.org/packages/f6/1d/ac8914360460fafa1990890259b7fa5ef7ba4cd59014e782e4ab3ab144d8/filterpy-1.4.5.zip", hash = "sha256:4f2a4d39e4ea601b9ab42b2db08b5918a9538c168cff1c6895ae26646f3d73b1", size = 177985, upload-time = "2018-10-10T22:38:24.63Z" } +[[package]] +name = "flake8" +version = "7.1.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mccabe" }, + { name = "pycodestyle" }, + { name = "pyflakes" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/58/16/3f2a0bb700ad65ac9663262905a025917c020a3f92f014d2ba8964b4602c/flake8-7.1.2.tar.gz", hash = "sha256:c586ffd0b41540951ae41af572e6790dbd49fc12b3aa2541685d253d9bd504bd", size = 48119, upload-time = "2025-02-16T18:45:44.296Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/35/f8/08d37b2cd89da306e3520bd27f8a85692122b42b56c0c2c3784ff09c022f/flake8-7.1.2-py2.py3-none-any.whl", hash = "sha256:1cbc62e65536f65e6d754dfe6f1bada7f5cf392d6f5db3c2b85892466c3e7c1a", size = 57745, upload-time = "2025-02-16T18:45:42.351Z" }, +] + [[package]] name = "flask" version = "3.1.2" @@ -2916,6 +3039,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d9/33/1f075bf72b0b747cb3288d011319aaf64083cf2efef8354174e3ed4540e2/ipython_pygments_lexers-1.1.1-py3-none-any.whl", hash = "sha256:a9462224a505ade19a605f71f8fa63c2048833ce50abc86768a0d81d876dc81c", size = 8074, upload-time = "2025-01-17T11:24:33.271Z" }, ] +[[package]] +name = "isort" +version = "7.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/63/53/4f3c058e3bace40282876f9b553343376ee687f3c35a525dc79dbd450f88/isort-7.0.0.tar.gz", hash = "sha256:5513527951aadb3ac4292a41a16cbc50dd1642432f5e8c20057d414bdafb4187", size = 805049, upload-time = "2025-10-11T13:30:59.107Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7f/ed/e3705d6d02b4f7aea715a353c8ce193efd0b5db13e204df895d38734c244/isort-7.0.0-py3-none-any.whl", hash = "sha256:1bcabac8bc3c36c7fb7b98a76c8abb18e0f841a3ba81decac7691008592499c1", size = 94672, upload-time = "2025-10-11T13:30:57.665Z" }, +] + [[package]] name = "itsdangerous" version = "2.2.0" @@ -3849,6 +3981,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/54/20/6aa79ba3570bddd1bf7e951c6123f806751e58e8cce736bad77b2cf348d7/logistro-2.0.1-py3-none-any.whl", hash = "sha256:06ffa127b9fb4ac8b1972ae6b2a9d7fde57598bf5939cd708f43ec5bba2d31eb", size = 8555, upload-time = "2025-11-01T02:41:17.587Z" }, ] +[[package]] +name = "lsprotocol" +version = "2025.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "attrs" }, + { name = "cattrs" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e9/26/67b84e6ec1402f0e6764ef3d2a0aaf9a79522cc1d37738f4e5bb0b21521a/lsprotocol-2025.0.0.tar.gz", hash = "sha256:e879da2b9301e82cfc3e60d805630487ac2f7ab17492f4f5ba5aaba94fe56c29", size = 74896, upload-time = "2025-06-17T21:30:18.156Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7b/f0/92f2d609d6642b5f30cb50a885d2bf1483301c69d5786286500d15651ef2/lsprotocol-2025.0.0-py3-none-any.whl", hash = "sha256:f9d78f25221f2a60eaa4a96d3b4ffae011b107537facee61d3da3313880995c7", size = 76250, upload-time = "2025-06-17T21:30:19.455Z" }, +] + [[package]] name = "lxml" version = "6.0.2" @@ -4239,6 +4384,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/af/33/ee4519fa02ed11a94aef9559552f3b17bb863f2ecfe1a35dc7f548cde231/matplotlib_inline-0.2.1-py3-none-any.whl", hash = "sha256:d56ce5156ba6085e00a9d54fead6ed29a9c47e215cd1bba2e976ef39f5710a76", size = 9516, upload-time = "2025-10-23T09:00:20.675Z" }, ] +[[package]] +name = "mccabe" +version = "0.7.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e7/ff/0ffefdcac38932a54d2b5eed4e0ba8a408f215002cd178ad1df0f2806ff8/mccabe-0.7.0.tar.gz", hash = "sha256:348e0240c33b60bbdf4e523192ef919f28cb2c3d7d5c7794f74009290f236325", size = 9658, upload-time = "2022-01-24T01:14:51.113Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/27/1a/1f68f9ba0c207934b35b86a8ca3aad8395a3d6dd7921c0686e23853ff5a9/mccabe-0.7.0-py2.py3-none-any.whl", hash = "sha256:6c2d30ab6be0e4a46919781807b4f0d834ebdd6c6e3dca0bda5a15f863427b6e", size = 7350, upload-time = "2022-01-24T01:14:49.62Z" }, +] + [[package]] name = "mcp" version = "1.25.0" @@ -6493,6 +6647,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ee/2b/e18ee7c5ee508a82897f021c1981533eca2940b5f072fc6ed0906c03a7a7/pybase64-1.4.3-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:debf737e09b8bf832ba86f5ecc3d3dbd0e3021d6cd86ba4abe962d6a5a77adb3", size = 36134, upload-time = "2025-12-06T13:26:47.35Z" }, ] +[[package]] +name = "pycodestyle" +version = "2.12.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/43/aa/210b2c9aedd8c1cbeea31a50e42050ad56187754b34eb214c46709445801/pycodestyle-2.12.1.tar.gz", hash = "sha256:6838eae08bbce4f6accd5d5572075c63626a15ee3e6f842df996bf62f6d73521", size = 39232, upload-time = "2024-08-04T20:26:54.576Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3a/d8/a211b3f85e99a0daa2ddec96c949cac6824bd305b040571b82a03dd62636/pycodestyle-2.12.1-py2.py3-none-any.whl", hash = "sha256:46f0fb92069a7c28ab7bb558f05bfc0110dac69a0cd23c61ea0040283a9d78b3", size = 31284, upload-time = "2024-08-04T20:26:53.173Z" }, +] + [[package]] name = "pycparser" version = "2.23" @@ -6684,6 +6847,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c1/60/5d4751ba3f4a40a6891f24eec885f51afd78d208498268c734e256fb13c4/pydantic_settings-2.12.0-py3-none-any.whl", hash = "sha256:fddb9fd99a5b18da837b29710391e945b1e30c135477f484084ee513adb93809", size = 51880, upload-time = "2025-11-10T14:25:45.546Z" }, ] +[[package]] +name = "pydocstyle" +version = "6.3.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "snowballstemmer" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e9/5c/d5385ca59fd065e3c6a5fe19f9bc9d5ea7f2509fa8c9c22fb6b2031dd953/pydocstyle-6.3.0.tar.gz", hash = "sha256:7ce43f0c0ac87b07494eb9c0b462c0b73e6ff276807f204d6b53edc72b7e44e1", size = 36796, upload-time = "2023-01-17T20:29:19.838Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/36/ea/99ddefac41971acad68f14114f38261c1f27dac0b3ec529824ebc739bdaa/pydocstyle-6.3.0-py3-none-any.whl", hash = "sha256:118762d452a49d6b05e194ef344a55822987a462831ade91ec5c06fd2169d019", size = 38038, upload-time = "2023-01-17T20:29:18.094Z" }, +] + [[package]] name = "pydub" version = "0.25.1" @@ -6705,6 +6880,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/9b/4d/b9add7c84060d4c1906abe9a7e5359f2a60f7a9a4f67268b2766673427d8/pyee-13.0.0-py3-none-any.whl", hash = "sha256:48195a3cddb3b1515ce0695ed76036b5ccc2ef3a9f963ff9f77aec0139845498", size = 15730, upload-time = "2025-03-17T18:53:14.532Z" }, ] +[[package]] +name = "pyflakes" +version = "3.2.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/57/f9/669d8c9c86613c9d568757c7f5824bd3197d7b1c6c27553bc5618a27cce2/pyflakes-3.2.0.tar.gz", hash = "sha256:1c61603ff154621fb2a9172037d84dca3500def8c8b630657d1701f026f8af3f", size = 63788, upload-time = "2024-01-05T00:28:47.703Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d4/d7/f1b7db88d8e4417c5d47adad627a93547f44bdc9028372dbd2313f34a855/pyflakes-3.2.0-py2.py3-none-any.whl", hash = "sha256:84b5be138a2dfbb40689ca07e2152deb896a65c3a3e24c251c5c62489568074a", size = 62725, upload-time = "2024-01-05T00:28:45.903Z" }, +] + [[package]] name = "pygame" version = "2.6.1" @@ -6795,6 +6979,25 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/8e/ec/6e02b2561d056ea5b33046e3cad21238e6a9097b97d6ccc0fbe52b50c858/pylibsrtp-1.0.0-cp310-abi3-win_arm64.whl", hash = "sha256:2696bdb2180d53ac55d0eb7b58048a2aa30cd4836dd2ca683669889137a94d2a", size = 1159246, upload-time = "2025-10-13T16:12:30.285Z" }, ] +[[package]] +name = "pylint" +version = "4.0.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "astroid" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "dill" }, + { name = "isort" }, + { name = "mccabe" }, + { name = "platformdirs" }, + { name = "tomli", marker = "python_full_version < '3.11'" }, + { name = "tomlkit" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5a/d2/b081da1a8930d00e3fc06352a1d449aaf815d4982319fab5d8cdb2e9ab35/pylint-4.0.4.tar.gz", hash = "sha256:d9b71674e19b1c36d79265b5887bf8e55278cbe236c9e95d22dc82cf044fdbd2", size = 1571735, upload-time = "2025-11-30T13:29:04.315Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a6/92/d40f5d937517cc489ad848fc4414ecccc7592e4686b9071e09e64f5e378e/pylint-4.0.4-py3-none-any.whl", hash = "sha256:63e06a37d5922555ee2c20963eb42559918c20bd2b21244e4ef426e7c43b92e0", size = 536425, upload-time = "2025-11-30T13:29:02.53Z" }, +] + [[package]] name = "pymavlink" version = "2.4.49" @@ -7102,6 +7305,65 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b9/06/a4ddfd46794c7d6e175c34e8c10554949d1c17aeb78c188050b4746d4b48/python_fcl-0.7.0.10-cp314-cp314t-win_amd64.whl", hash = "sha256:6ab961f459c294695385d518f7a6eb3a2577029ca008698045dac2b7253fa3f7", size = 1140958, upload-time = "2025-10-22T06:28:44.586Z" }, ] +[[package]] +name = "python-lsp-jsonrpc" +version = "1.1.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "ujson" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/48/b6/fd92e2ea4635d88966bb42c20198df1a981340f07843b5e3c6694ba3557b/python-lsp-jsonrpc-1.1.2.tar.gz", hash = "sha256:4688e453eef55cd952bff762c705cedefa12055c0aec17a06f595bcc002cc912", size = 15298, upload-time = "2023-09-23T17:48:30.451Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cb/d9/656659d5b5d5f402b2b174cd0ba9bc827e07ce3c0bf88da65424baf64af8/python_lsp_jsonrpc-1.1.2-py3-none-any.whl", hash = "sha256:7339c2e9630ae98903fdaea1ace8c47fba0484983794d6aafd0bd8989be2b03c", size = 8805, upload-time = "2023-09-23T17:48:28.804Z" }, +] + +[[package]] +name = "python-lsp-ruff" +version = "2.3.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cattrs" }, + { name = "lsprotocol" }, + { name = "python-lsp-server" }, + { name = "ruff" }, + { name = "tomli", marker = "python_full_version < '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/af/79/2f6322c47bd2956447e0a6787084b4110b4473e3d2501b86aa47c802e6a0/python_lsp_ruff-2.3.0.tar.gz", hash = "sha256:647745b7f3010ac101e3c53a797b8f9deb1f52228b608d70ad0e8e056978c3b7", size = 17268, upload-time = "2025-09-29T20:14:02.994Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/85/c0/761e359e255fce641c263a3c3e43f7685d1667139e9d35a376c1cc9f6f70/python_lsp_ruff-2.3.0-py3-none-any.whl", hash = "sha256:b858b698fbaff5670f6d5e6c66afc632908f78639d73dc85dedd33ae5fdd204f", size = 12039, upload-time = "2025-09-29T20:14:01.56Z" }, +] + +[[package]] +name = "python-lsp-server" +version = "1.14.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "black" }, + { name = "docstring-to-markdown" }, + { name = "jedi" }, + { name = "pluggy" }, + { name = "python-lsp-jsonrpc" }, + { name = "ujson" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b4/b5/b989d41c63390dfc2bf63275ab543b82fed076723d912055e77ccbae1422/python_lsp_server-1.14.0.tar.gz", hash = "sha256:509c445fc667f41ffd3191cb7512a497bf7dd76c14ceb1ee2f6c13ebe71f9a6b", size = 121536, upload-time = "2025-12-06T16:12:20.86Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/08/cf/587f913335e3855e0ddca2aee7c3f9d5de2d75a1e23434891e9f74783bcd/python_lsp_server-1.14.0-py3-none-any.whl", hash = "sha256:a71a917464effc48f4c70363f90b8520e5e3ba8201428da80b97a7ceb259e32a", size = 77060, upload-time = "2025-12-06T16:12:19.46Z" }, +] + +[package.optional-dependencies] +all = [ + { name = "autopep8" }, + { name = "flake8" }, + { name = "mccabe" }, + { name = "pycodestyle" }, + { name = "pydocstyle" }, + { name = "pyflakes" }, + { name = "pylint" }, + { name = "rope" }, + { name = "whatthepatch" }, + { name = "yapf" }, +] + [[package]] name = "python-multipart" version = "0.0.20" @@ -7124,6 +7386,63 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/28/d2/2ccc2b69a187b80fda3152745670cfba936704f296a9fa54c6c8ac694d12/python_socketio-5.16.0-py3-none-any.whl", hash = "sha256:d95802961e15c7bd54ecf884c6e7644f81be8460f0a02ee66b473df58088ee8a", size = 79607, upload-time = "2025-12-24T23:51:47.2Z" }, ] +[[package]] +name = "pytokens" +version = "0.4.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e5/16/4b9cfd90d55e66ffdb277d7ebe3bc25250c2311336ec3fc73b2673c794d5/pytokens-0.4.0.tar.gz", hash = "sha256:6b0b03e6ea7c9f9d47c5c61164b69ad30f4f0d70a5d9fe7eac4d19f24f77af2d", size = 15039, upload-time = "2026-01-19T07:59:50.623Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3f/c5/c20818fef16c4ab5f9fd7bad699268ba21bf24f655711df4e33bb7a9ab47/pytokens-0.4.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:af0c3166aea367a9e755a283171befb92dd3043858b94ae9b3b7efbe9def26a3", size = 160682, upload-time = "2026-01-19T07:58:51.583Z" }, + { url = "https://files.pythonhosted.org/packages/46/c4/ad03e4abe05c6af57c4d7f8f031fafe80f0074796d09ab5a73bf2fac895f/pytokens-0.4.0-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:daae524ed14ca459932cbf51d74325bea643701ba8a8b0cc2d10f7cd4b3e2b63", size = 245748, upload-time = "2026-01-19T07:58:53.944Z" }, + { url = "https://files.pythonhosted.org/packages/6b/b9/4a7ee0a692603b16d8fdfbc5c44e0f6910d45eec6b2c2188daa4670f179d/pytokens-0.4.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e95cb158c44d642ed62f555bf8136bbe780dbd64d2fb0b9169e11ffb944664c3", size = 258671, upload-time = "2026-01-19T07:58:55.667Z" }, + { url = "https://files.pythonhosted.org/packages/fa/a3/02bb29dc4985fb8d759d9c96f189c3a828e74f0879fdb843e9fb7a1db637/pytokens-0.4.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:df58d44630eaf25f587540e94bdf1fc50b4e6d5f212c786de0fb024bfcb8753a", size = 261749, upload-time = "2026-01-19T07:58:57.442Z" }, + { url = "https://files.pythonhosted.org/packages/10/f2/9a8bdcc5444d85d4dba4aa1b530d81af3edc4a9ab76bf1d53ea8bfe8479d/pytokens-0.4.0-cp310-cp310-win_amd64.whl", hash = "sha256:55efcc36f9a2e0e930cfba0ce7f83445306b02f8326745585ed5551864eba73a", size = 102805, upload-time = "2026-01-19T07:58:59.068Z" }, + { url = "https://files.pythonhosted.org/packages/b4/05/3196399a353dd4cd99138a88f662810979ee2f1a1cdb0b417cb2f4507836/pytokens-0.4.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:92eb3ef88f27c22dc9dbab966ace4d61f6826e02ba04dac8e2d65ea31df56c8e", size = 160075, upload-time = "2026-01-19T07:59:00.316Z" }, + { url = "https://files.pythonhosted.org/packages/28/1d/c8fc4ed0a1c4f660391b201cda00b1d5bbcc00e2998e8bcd48b15eefd708/pytokens-0.4.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f4b77858a680635ee9904306f54b0ee4781effb89e211ba0a773d76539537165", size = 247318, upload-time = "2026-01-19T07:59:01.636Z" }, + { url = "https://files.pythonhosted.org/packages/8e/0e/53e55ba01f3e858d229cd84b02481542f42ba59050483a78bf2447ee1af7/pytokens-0.4.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:25cacc20c2ad90acb56f3739d87905473c54ca1fa5967ffcd675463fe965865e", size = 259752, upload-time = "2026-01-19T07:59:04.229Z" }, + { url = "https://files.pythonhosted.org/packages/dc/56/2d930d7f899e3f21868ca6e8ec739ac31e8fc532f66e09cbe45d3df0a84f/pytokens-0.4.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:628fab535ebc9079e4db35cd63cb401901c7ce8720a9834f9ad44b9eb4e0f1d4", size = 262842, upload-time = "2026-01-19T07:59:06.14Z" }, + { url = "https://files.pythonhosted.org/packages/42/dd/4e7e6920d23deffaf66e6f40d45f7610dcbc132ca5d90ab4faccef22f624/pytokens-0.4.0-cp311-cp311-win_amd64.whl", hash = "sha256:4d0f568d7e82b7e96be56d03b5081de40e43c904eb6492bf09aaca47cd55f35b", size = 102620, upload-time = "2026-01-19T07:59:07.839Z" }, + { url = "https://files.pythonhosted.org/packages/3d/65/65460ebbfefd0bc1b160457904370d44f269e6e4582e0a9b6cba7c267b04/pytokens-0.4.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:cd8da894e5a29ba6b6da8be06a4f7589d7220c099b5e363cb0643234b9b38c2a", size = 159864, upload-time = "2026-01-19T07:59:08.908Z" }, + { url = "https://files.pythonhosted.org/packages/25/70/a46669ec55876c392036b4da9808b5c3b1c5870bbca3d4cc923bf68bdbc1/pytokens-0.4.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:237ba7cfb677dbd3b01b09860810aceb448871150566b93cd24501d5734a04b1", size = 254448, upload-time = "2026-01-19T07:59:10.594Z" }, + { url = "https://files.pythonhosted.org/packages/62/0b/c486fc61299c2fc3b7f88ee4e115d4c8b6ffd1a7f88dc94b398b5b1bc4b8/pytokens-0.4.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:01d1a61e36812e4e971cfe2c0e4c1f2d66d8311031dac8bf168af8a249fa04dd", size = 268863, upload-time = "2026-01-19T07:59:12.31Z" }, + { url = "https://files.pythonhosted.org/packages/79/92/b036af846707d25feaff7cafbd5280f1bd6a1034c16bb06a7c910209c1ab/pytokens-0.4.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:e47e2ef3ec6ee86909e520d79f965f9b23389fda47460303cf715d510a6fe544", size = 267181, upload-time = "2026-01-19T07:59:13.856Z" }, + { url = "https://files.pythonhosted.org/packages/0d/c0/6d011fc00fefa74ce34816c84a923d2dd7c46b8dbc6ee52d13419786834c/pytokens-0.4.0-cp312-cp312-win_amd64.whl", hash = "sha256:3d36954aba4557fd5a418a03cf595ecbb1cdcce119f91a49b19ef09d691a22ae", size = 102814, upload-time = "2026-01-19T07:59:15.288Z" }, + { url = "https://files.pythonhosted.org/packages/98/63/627b7e71d557383da5a97f473ad50f8d9c2c1f55c7d3c2531a120c796f6e/pytokens-0.4.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:73eff3bdd8ad08da679867992782568db0529b887bed4c85694f84cdf35eafc6", size = 159744, upload-time = "2026-01-19T07:59:16.88Z" }, + { url = "https://files.pythonhosted.org/packages/28/d7/16f434c37ec3824eba6bcb6e798e5381a8dc83af7a1eda0f95c16fe3ade5/pytokens-0.4.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d97cc1f91b1a8e8ebccf31c367f28225699bea26592df27141deade771ed0afb", size = 253207, upload-time = "2026-01-19T07:59:18.069Z" }, + { url = "https://files.pythonhosted.org/packages/ab/96/04102856b9527701ae57d74a6393d1aca5bad18a1b1ca48ccffb3c93b392/pytokens-0.4.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a2c8952c537cb73a1a74369501a83b7f9d208c3cf92c41dd88a17814e68d48ce", size = 267452, upload-time = "2026-01-19T07:59:19.328Z" }, + { url = "https://files.pythonhosted.org/packages/0e/ef/0936eb472b89ab2d2c2c24bb81c50417e803fa89c731930d9fb01176fe9f/pytokens-0.4.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:5dbf56f3c748aed9310b310d5b8b14e2c96d3ad682ad5a943f381bdbbdddf753", size = 265965, upload-time = "2026-01-19T07:59:20.613Z" }, + { url = "https://files.pythonhosted.org/packages/ae/f5/64f3d6f7df4a9e92ebda35ee85061f6260e16eac82df9396020eebbca775/pytokens-0.4.0-cp313-cp313-win_amd64.whl", hash = "sha256:e131804513597f2dff2b18f9911d9b6276e21ef3699abeffc1c087c65a3d975e", size = 102813, upload-time = "2026-01-19T07:59:22.012Z" }, + { url = "https://files.pythonhosted.org/packages/5f/f1/d07e6209f18ef378fc2ae9dee8d1dfe91fd2447c2e2dbfa32867b6dd30cf/pytokens-0.4.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:0d7374c917197106d3c4761374718bc55ea2e9ac0fb94171588ef5840ee1f016", size = 159968, upload-time = "2026-01-19T07:59:23.07Z" }, + { url = "https://files.pythonhosted.org/packages/0a/73/0eb111400abd382a04f253b269819db9fcc748aa40748441cebdcb6d068f/pytokens-0.4.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0cd3fa1caf9e47a72ee134a29ca6b5bea84712724bba165d6628baa190c6ea5b", size = 253373, upload-time = "2026-01-19T07:59:24.381Z" }, + { url = "https://files.pythonhosted.org/packages/bd/8d/9e4e2fdb5bcaba679e54afcc304e9f13f488eb4d626e6b613f9553e03dbd/pytokens-0.4.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9c6986576b7b07fe9791854caa5347923005a80b079d45b63b0be70d50cce5f1", size = 267024, upload-time = "2026-01-19T07:59:25.74Z" }, + { url = "https://files.pythonhosted.org/packages/cb/b7/e0a370321af2deb772cff14ff337e1140d1eac2c29a8876bfee995f486f0/pytokens-0.4.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:9940f7c2e2f54fb1cb5fe17d0803c54da7a2bf62222704eb4217433664a186a7", size = 270912, upload-time = "2026-01-19T07:59:27.072Z" }, + { url = "https://files.pythonhosted.org/packages/7c/54/4348f916c440d4c3e68b53b4ed0e66b292d119e799fa07afa159566dcc86/pytokens-0.4.0-cp314-cp314-win_amd64.whl", hash = "sha256:54691cf8f299e7efabcc25adb4ce715d3cef1491e1c930eaf555182f898ef66a", size = 103836, upload-time = "2026-01-19T07:59:28.112Z" }, + { url = "https://files.pythonhosted.org/packages/e8/f8/a693c0cfa9c783a2a8c4500b7b2a8bab420f8ca4f2d496153226bf1c12e3/pytokens-0.4.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:94ff5db97a0d3cd7248a5b07ba2167bd3edc1db92f76c6db00137bbaf068ddf8", size = 167643, upload-time = "2026-01-19T07:59:29.292Z" }, + { url = "https://files.pythonhosted.org/packages/c0/dd/a64eb1e9f3ec277b69b33ef1b40ffbcc8f0a3bafcde120997efc7bdefebf/pytokens-0.4.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d0dd6261cd9cc95fae1227b1b6ebee023a5fd4a4b6330b071c73a516f5f59b63", size = 289553, upload-time = "2026-01-19T07:59:30.537Z" }, + { url = "https://files.pythonhosted.org/packages/df/22/06c1079d93dbc3bca5d013e1795f3d8b9ed6c87290acd6913c1c526a6bb2/pytokens-0.4.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0cdca8159df407dbd669145af4171a0d967006e0be25f3b520896bc7068f02c4", size = 302490, upload-time = "2026-01-19T07:59:32.352Z" }, + { url = "https://files.pythonhosted.org/packages/8d/de/a6f5e43115b4fbf4b93aa87d6c83c79932cdb084f9711daae04549e1e4ad/pytokens-0.4.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:4b5770abeb2a24347380a1164a558f0ebe06e98aedbd54c45f7929527a5fb26e", size = 305652, upload-time = "2026-01-19T07:59:33.685Z" }, + { url = "https://files.pythonhosted.org/packages/ab/3d/c136e057cb622e36e0c3ff7a8aaa19ff9720050c4078235691da885fe6ee/pytokens-0.4.0-cp314-cp314t-win_amd64.whl", hash = "sha256:74500d72c561dad14c037a9e86a657afd63e277dd5a3bb7570932ab7a3b12551", size = 115472, upload-time = "2026-01-19T07:59:34.734Z" }, + { url = "https://files.pythonhosted.org/packages/7c/3c/6941a82f4f130af6e1c68c076b6789069ef10c04559bd4733650f902fd3b/pytokens-0.4.0-py3-none-any.whl", hash = "sha256:0508d11b4de157ee12063901603be87fb0253e8f4cb9305eb168b1202ab92068", size = 13224, upload-time = "2026-01-19T07:59:49.822Z" }, +] + +[[package]] +name = "pytoolconfig" +version = "1.3.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "packaging" }, + { name = "tomli", marker = "python_full_version < '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/18/dc/abf70d2c2bcac20e8c71a7cdf6d44e4ddba4edf65acb179248d554d743db/pytoolconfig-1.3.1.tar.gz", hash = "sha256:51e6bd1a6f108238ae6aab6a65e5eed5e75d456be1c2bf29b04e5c1e7d7adbae", size = 16655, upload-time = "2024-01-11T16:25:11.914Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/92/44/da239917f5711ca7105f7d7f9e2765716dd883b241529beafc0f28504725/pytoolconfig-1.3.1-py3-none-any.whl", hash = "sha256:5d8cea8ae1996938ec3eaf44567bbc5ef1bc900742190c439a44a704d6e1b62b", size = 17022, upload-time = "2024-01-11T16:25:10.589Z" }, +] + +[package.optional-dependencies] +global = [ + { name = "platformdirs" }, +] + [[package]] name = "pyturbojpeg" version = "1.8.2" @@ -7533,6 +7852,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/25/7a/b0178788f8dc6cafce37a212c99565fa1fe7872c70c6c9c1e1a372d9d88f/rich-14.2.0-py3-none-any.whl", hash = "sha256:76bc51fe2e57d2b1be1f96c524b890b816e334ab4c1e45888799bfaab0021edd", size = 243393, upload-time = "2025-10-09T14:16:51.245Z" }, ] +[[package]] +name = "rope" +version = "1.14.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pytoolconfig", extra = ["global"] }, +] +sdist = { url = "https://files.pythonhosted.org/packages/74/3a/85e60d154f26ecdc1d47a63ac58bd9f32a5a9f3f771f6672197f02a00ade/rope-1.14.0.tar.gz", hash = "sha256:8803e3b667315044f6270b0c69a10c0679f9f322ed8efe6245a93ceb7658da69", size = 296801, upload-time = "2025-07-12T17:46:07.786Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/75/35/130469d1901da2b3a5a377539b4ffcd8a5c983f1c9e3ba5ffdd8d71ae314/rope-1.14.0-py3-none-any.whl", hash = "sha256:00a7ea8c0c376fc0b053b2f2f8ef3bfb8b50fecf1ebf3eb80e4f8bd7f1941918", size = 207143, upload-time = "2025-07-12T17:46:05.928Z" }, +] + [[package]] name = "rpds-py" version = "0.30.0" @@ -8125,6 +8456,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e9/44/75a9c9421471a6c4805dbf2356f7c181a29c1879239abab1ea2cc8f38b40/sniffio-1.3.1-py3-none-any.whl", hash = "sha256:2f6da418d1f1e0fddd844478f41680e794e6051915791a034ff65e5f100525a2", size = 10235, upload-time = "2024-02-25T23:20:01.196Z" }, ] +[[package]] +name = "snowballstemmer" +version = "3.0.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/75/a7/9810d872919697c9d01295633f5d574fb416d47e535f258272ca1f01f447/snowballstemmer-3.0.1.tar.gz", hash = "sha256:6d5eeeec8e9f84d4d56b847692bacf79bc2c8e90c7f80ca4444ff8b6f2e52895", size = 105575, upload-time = "2025-05-09T16:34:51.843Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c8/78/3565d011c61f5a43488987ee32b6f3f656e7f107ac2782dd57bdd7d91d9a/snowballstemmer-3.0.1-py3-none-any.whl", hash = "sha256:6cd7b3897da8d6c9ffb968a6781fa6532dce9c3618a4b127d920dab764a19064", size = 103274, upload-time = "2025-05-09T16:34:50.371Z" }, +] + [[package]] name = "sortedcontainers" version = "2.4.0" @@ -8607,6 +8947,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/77/b8/0135fadc89e73be292b473cb820b4f5a08197779206b33191e801feeae40/tomli-2.3.0-py3-none-any.whl", hash = "sha256:e95b1af3c5b07d9e643909b5abbec77cd9f1217e6d0bca72b0234736b9fb1f1b", size = 14408, upload-time = "2025-10-08T22:01:46.04Z" }, ] +[[package]] +name = "tomlkit" +version = "0.14.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/c3/af/14b24e41977adb296d6bd1fb59402cf7d60ce364f90c890bd2ec65c43b5a/tomlkit-0.14.0.tar.gz", hash = "sha256:cf00efca415dbd57575befb1f6634c4f42d2d87dbba376128adb42c121b87064", size = 187167, upload-time = "2026-01-13T01:14:53.304Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b5/11/87d6d29fb5d237229d67973a6c9e06e048f01cf4994dee194ab0ea841814/tomlkit-0.14.0-py3-none-any.whl", hash = "sha256:592064ed85b40fa213469f81ac584f67a4f2992509a7c3ea2d632208623a3680", size = 39310, upload-time = "2026-01-13T01:14:51.965Z" }, +] + [[package]] name = "toolz" version = "1.1.0" @@ -9074,6 +9423,86 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/37/87/1f677586e8ac487e29672e4b17455758fce261de06a0d086167bb760361a/uc_micro_py-1.0.3-py3-none-any.whl", hash = "sha256:db1dffff340817673d7b466ec86114a9dc0e9d4d9b5ba229d9d60e5c12600cd5", size = 6229, upload-time = "2024-02-09T16:52:00.371Z" }, ] +[[package]] +name = "ujson" +version = "5.11.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/43/d9/3f17e3c5773fb4941c68d9a37a47b1a79c9649d6c56aefbed87cc409d18a/ujson-5.11.0.tar.gz", hash = "sha256:e204ae6f909f099ba6b6b942131cee359ddda2b6e4ea39c12eb8b991fe2010e0", size = 7156583, upload-time = "2025-08-20T11:57:02.452Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/86/0c/8bf7a4fabfd01c7eed92d9b290930ce6d14910dec708e73538baa38885d1/ujson-5.11.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:446e8c11c06048611c9d29ef1237065de0af07cabdd97e6b5b527b957692ec25", size = 55248, upload-time = "2025-08-20T11:55:02.368Z" }, + { url = "https://files.pythonhosted.org/packages/7b/2e/eeab0b8b641817031ede4f790db4c4942df44a12f44d72b3954f39c6a115/ujson-5.11.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:16ccb973b7ada0455201808ff11d48fe9c3f034a6ab5bd93b944443c88299f89", size = 53157, upload-time = "2025-08-20T11:55:04.012Z" }, + { url = "https://files.pythonhosted.org/packages/21/1b/a4e7a41870797633423ea79618526747353fd7be9191f3acfbdee0bf264b/ujson-5.11.0-cp310-cp310-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3134b783ab314d2298d58cda7e47e7a0f7f71fc6ade6ac86d5dbeaf4b9770fa6", size = 57657, upload-time = "2025-08-20T11:55:05.169Z" }, + { url = "https://files.pythonhosted.org/packages/94/ae/4e0d91b8f6db7c9b76423b3649612189506d5a06ddd3b6334b6d37f77a01/ujson-5.11.0-cp310-cp310-manylinux_2_24_i686.manylinux_2_28_i686.whl", hash = "sha256:185f93ebccffebc8baf8302c869fac70dd5dd78694f3b875d03a31b03b062cdb", size = 59780, upload-time = "2025-08-20T11:55:06.325Z" }, + { url = "https://files.pythonhosted.org/packages/b3/cc/46b124c2697ca2da7c65c4931ed3cb670646978157aa57a7a60f741c530f/ujson-5.11.0-cp310-cp310-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d06e87eded62ff0e5f5178c916337d2262fdbc03b31688142a3433eabb6511db", size = 57307, upload-time = "2025-08-20T11:55:07.493Z" }, + { url = "https://files.pythonhosted.org/packages/39/eb/20dd1282bc85dede2f1c62c45b4040bc4c389c80a05983515ab99771bca7/ujson-5.11.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:181fb5b15703a8b9370b25345d2a1fd1359f0f18776b3643d24e13ed9c036d4c", size = 1036369, upload-time = "2025-08-20T11:55:09.192Z" }, + { url = "https://files.pythonhosted.org/packages/64/a2/80072439065d493e3a4b1fbeec991724419a1b4c232e2d1147d257cac193/ujson-5.11.0-cp310-cp310-musllinux_1_2_i686.whl", hash = "sha256:a4df61a6df0a4a8eb5b9b1ffd673429811f50b235539dac586bb7e9e91994138", size = 1195738, upload-time = "2025-08-20T11:55:11.402Z" }, + { url = "https://files.pythonhosted.org/packages/5d/7e/d77f9e9c039d58299c350c978e086a804d1fceae4fd4a1cc6e8d0133f838/ujson-5.11.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:6eff24e1abd79e0ec6d7eae651dd675ddbc41f9e43e29ef81e16b421da896915", size = 1088718, upload-time = "2025-08-20T11:55:13.297Z" }, + { url = "https://files.pythonhosted.org/packages/ab/f1/697559d45acc849cada6b3571d53522951b1a64027400507aabc6a710178/ujson-5.11.0-cp310-cp310-win32.whl", hash = "sha256:30f607c70091483550fbd669a0b37471e5165b317d6c16e75dba2aa967608723", size = 39653, upload-time = "2025-08-20T11:55:14.869Z" }, + { url = "https://files.pythonhosted.org/packages/86/a2/70b73a0f55abe0e6b8046d365d74230c20c5691373e6902a599b2dc79ba1/ujson-5.11.0-cp310-cp310-win_amd64.whl", hash = "sha256:3d2720e9785f84312b8e2cb0c2b87f1a0b1c53aaab3b2af3ab817d54409012e0", size = 43720, upload-time = "2025-08-20T11:55:15.897Z" }, + { url = "https://files.pythonhosted.org/packages/1c/5f/b19104afa455630b43efcad3a24495b9c635d92aa8f2da4f30e375deb1a2/ujson-5.11.0-cp310-cp310-win_arm64.whl", hash = "sha256:85e6796631165f719084a9af00c79195d3ebf108151452fefdcb1c8bb50f0105", size = 38410, upload-time = "2025-08-20T11:55:17.556Z" }, + { url = "https://files.pythonhosted.org/packages/da/ea/80346b826349d60ca4d612a47cdf3533694e49b45e9d1c07071bb867a184/ujson-5.11.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:d7c46cb0fe5e7056b9acb748a4c35aa1b428025853032540bb7e41f46767321f", size = 55248, upload-time = "2025-08-20T11:55:19.033Z" }, + { url = "https://files.pythonhosted.org/packages/57/df/b53e747562c89515e18156513cc7c8ced2e5e3fd6c654acaa8752ffd7cd9/ujson-5.11.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:d8951bb7a505ab2a700e26f691bdfacf395bc7e3111e3416d325b513eea03a58", size = 53156, upload-time = "2025-08-20T11:55:20.174Z" }, + { url = "https://files.pythonhosted.org/packages/41/b8/ab67ec8c01b8a3721fd13e5cb9d85ab2a6066a3a5e9148d661a6870d6293/ujson-5.11.0-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:952c0be400229940248c0f5356514123d428cba1946af6fa2bbd7503395fef26", size = 57657, upload-time = "2025-08-20T11:55:21.296Z" }, + { url = "https://files.pythonhosted.org/packages/7b/c7/fb84f27cd80a2c7e2d3c6012367aecade0da936790429801803fa8d4bffc/ujson-5.11.0-cp311-cp311-manylinux_2_24_i686.manylinux_2_28_i686.whl", hash = "sha256:94fcae844f1e302f6f8095c5d1c45a2f0bfb928cccf9f1b99e3ace634b980a2a", size = 59779, upload-time = "2025-08-20T11:55:22.772Z" }, + { url = "https://files.pythonhosted.org/packages/5d/7c/48706f7c1e917ecb97ddcfb7b1d756040b86ed38290e28579d63bd3fcc48/ujson-5.11.0-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7e0ec1646db172beb8d3df4c32a9d78015e671d2000af548252769e33079d9a6", size = 57284, upload-time = "2025-08-20T11:55:24.01Z" }, + { url = "https://files.pythonhosted.org/packages/ec/ce/48877c6eb4afddfd6bd1db6be34456538c07ca2d6ed233d3f6c6efc2efe8/ujson-5.11.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:da473b23e3a54448b008d33f742bcd6d5fb2a897e42d1fc6e7bf306ea5d18b1b", size = 1036395, upload-time = "2025-08-20T11:55:25.725Z" }, + { url = "https://files.pythonhosted.org/packages/8b/7a/2c20dc97ad70cd7c31ad0596ba8e2cf8794d77191ba4d1e0bded69865477/ujson-5.11.0-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:aa6b3d4f1c0d3f82930f4cbd7fe46d905a4a9205a7c13279789c1263faf06dba", size = 1195731, upload-time = "2025-08-20T11:55:27.915Z" }, + { url = "https://files.pythonhosted.org/packages/15/f5/ca454f2f6a2c840394b6f162fff2801450803f4ff56c7af8ce37640b8a2a/ujson-5.11.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:4843f3ab4fe1cc596bb7e02228ef4c25d35b4bb0809d6a260852a4bfcab37ba3", size = 1088710, upload-time = "2025-08-20T11:55:29.426Z" }, + { url = "https://files.pythonhosted.org/packages/fe/d3/9ba310e07969bc9906eb7548731e33a0f448b122ad9705fed699c9b29345/ujson-5.11.0-cp311-cp311-win32.whl", hash = "sha256:e979fbc469a7f77f04ec2f4e853ba00c441bf2b06720aa259f0f720561335e34", size = 39648, upload-time = "2025-08-20T11:55:31.194Z" }, + { url = "https://files.pythonhosted.org/packages/57/f7/da05b4a8819f1360be9e71fb20182f0bb3ec611a36c3f213f4d20709e099/ujson-5.11.0-cp311-cp311-win_amd64.whl", hash = "sha256:683f57f0dd3acdd7d9aff1de0528d603aafcb0e6d126e3dc7ce8b020a28f5d01", size = 43717, upload-time = "2025-08-20T11:55:32.241Z" }, + { url = "https://files.pythonhosted.org/packages/9a/cc/f3f9ac0f24f00a623a48d97dc3814df5c2dc368cfb00031aa4141527a24b/ujson-5.11.0-cp311-cp311-win_arm64.whl", hash = "sha256:7855ccea3f8dad5e66d8445d754fc1cf80265a4272b5f8059ebc7ec29b8d0835", size = 38402, upload-time = "2025-08-20T11:55:33.641Z" }, + { url = "https://files.pythonhosted.org/packages/b9/ef/a9cb1fce38f699123ff012161599fb9f2ff3f8d482b4b18c43a2dc35073f/ujson-5.11.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:7895f0d2d53bd6aea11743bd56e3cb82d729980636cd0ed9b89418bf66591702", size = 55434, upload-time = "2025-08-20T11:55:34.987Z" }, + { url = "https://files.pythonhosted.org/packages/b1/05/dba51a00eb30bd947791b173766cbed3492269c150a7771d2750000c965f/ujson-5.11.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:12b5e7e22a1fe01058000d1b317d3b65cc3daf61bd2ea7a2b76721fe160fa74d", size = 53190, upload-time = "2025-08-20T11:55:36.384Z" }, + { url = "https://files.pythonhosted.org/packages/03/3c/fd11a224f73fbffa299fb9644e425f38b38b30231f7923a088dd513aabb4/ujson-5.11.0-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0180a480a7d099082501cad1fe85252e4d4bf926b40960fb3d9e87a3a6fbbc80", size = 57600, upload-time = "2025-08-20T11:55:37.692Z" }, + { url = "https://files.pythonhosted.org/packages/55/b9/405103cae24899df688a3431c776e00528bd4799e7d68820e7ebcf824f92/ujson-5.11.0-cp312-cp312-manylinux_2_24_i686.manylinux_2_28_i686.whl", hash = "sha256:fa79fdb47701942c2132a9dd2297a1a85941d966d8c87bfd9e29b0cf423f26cc", size = 59791, upload-time = "2025-08-20T11:55:38.877Z" }, + { url = "https://files.pythonhosted.org/packages/17/7b/2dcbc2bbfdbf68f2368fb21ab0f6735e872290bb604c75f6e06b81edcb3f/ujson-5.11.0-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8254e858437c00f17cb72e7a644fc42dad0ebb21ea981b71df6e84b1072aaa7c", size = 57356, upload-time = "2025-08-20T11:55:40.036Z" }, + { url = "https://files.pythonhosted.org/packages/d1/71/fea2ca18986a366c750767b694430d5ded6b20b6985fddca72f74af38a4c/ujson-5.11.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:1aa8a2ab482f09f6c10fba37112af5f957689a79ea598399c85009f2f29898b5", size = 1036313, upload-time = "2025-08-20T11:55:41.408Z" }, + { url = "https://files.pythonhosted.org/packages/a3/bb/d4220bd7532eac6288d8115db51710fa2d7d271250797b0bfba9f1e755af/ujson-5.11.0-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:a638425d3c6eed0318df663df44480f4a40dc87cc7c6da44d221418312f6413b", size = 1195782, upload-time = "2025-08-20T11:55:43.357Z" }, + { url = "https://files.pythonhosted.org/packages/80/47/226e540aa38878ce1194454385701d82df538ccb5ff8db2cf1641dde849a/ujson-5.11.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:7e3cff632c1d78023b15f7e3a81c3745cd3f94c044d1e8fa8efbd6b161997bbc", size = 1088817, upload-time = "2025-08-20T11:55:45.262Z" }, + { url = "https://files.pythonhosted.org/packages/7e/81/546042f0b23c9040d61d46ea5ca76f0cc5e0d399180ddfb2ae976ebff5b5/ujson-5.11.0-cp312-cp312-win32.whl", hash = "sha256:be6b0eaf92cae8cdee4d4c9e074bde43ef1c590ed5ba037ea26c9632fb479c88", size = 39757, upload-time = "2025-08-20T11:55:46.522Z" }, + { url = "https://files.pythonhosted.org/packages/44/1b/27c05dc8c9728f44875d74b5bfa948ce91f6c33349232619279f35c6e817/ujson-5.11.0-cp312-cp312-win_amd64.whl", hash = "sha256:b7b136cc6abc7619124fd897ef75f8e63105298b5ca9bdf43ebd0e1fa0ee105f", size = 43859, upload-time = "2025-08-20T11:55:47.987Z" }, + { url = "https://files.pythonhosted.org/packages/22/2d/37b6557c97c3409c202c838aa9c960ca3896843b4295c4b7bb2bbd260664/ujson-5.11.0-cp312-cp312-win_arm64.whl", hash = "sha256:6cd2df62f24c506a0ba322d5e4fe4466d47a9467b57e881ee15a31f7ecf68ff6", size = 38361, upload-time = "2025-08-20T11:55:49.122Z" }, + { url = "https://files.pythonhosted.org/packages/1c/ec/2de9dd371d52c377abc05d2b725645326c4562fc87296a8907c7bcdf2db7/ujson-5.11.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:109f59885041b14ee9569bf0bb3f98579c3fa0652317b355669939e5fc5ede53", size = 55435, upload-time = "2025-08-20T11:55:50.243Z" }, + { url = "https://files.pythonhosted.org/packages/5b/a4/f611f816eac3a581d8a4372f6967c3ed41eddbae4008d1d77f223f1a4e0a/ujson-5.11.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:a31c6b8004438e8c20fc55ac1c0e07dad42941db24176fe9acf2815971f8e752", size = 53193, upload-time = "2025-08-20T11:55:51.373Z" }, + { url = "https://files.pythonhosted.org/packages/e9/c5/c161940967184de96f5cbbbcce45b562a4bf851d60f4c677704b1770136d/ujson-5.11.0-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:78c684fb21255b9b90320ba7e199780f653e03f6c2528663768965f4126a5b50", size = 57603, upload-time = "2025-08-20T11:55:52.583Z" }, + { url = "https://files.pythonhosted.org/packages/2b/d6/c7b2444238f5b2e2d0e3dab300b9ddc3606e4b1f0e4bed5a48157cebc792/ujson-5.11.0-cp313-cp313-manylinux_2_24_i686.manylinux_2_28_i686.whl", hash = "sha256:4c9f5d6a27d035dd90a146f7761c2272cf7103de5127c9ab9c4cd39ea61e878a", size = 59794, upload-time = "2025-08-20T11:55:53.69Z" }, + { url = "https://files.pythonhosted.org/packages/fe/a3/292551f936d3d02d9af148f53e1bc04306b00a7cf1fcbb86fa0d1c887242/ujson-5.11.0-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:837da4d27fed5fdc1b630bd18f519744b23a0b5ada1bbde1a36ba463f2900c03", size = 57363, upload-time = "2025-08-20T11:55:54.843Z" }, + { url = "https://files.pythonhosted.org/packages/90/a6/82cfa70448831b1a9e73f882225980b5c689bf539ec6400b31656a60ea46/ujson-5.11.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:787aff4a84da301b7f3bac09bc696e2e5670df829c6f8ecf39916b4e7e24e701", size = 1036311, upload-time = "2025-08-20T11:55:56.197Z" }, + { url = "https://files.pythonhosted.org/packages/84/5c/96e2266be50f21e9b27acaee8ca8f23ea0b85cb998c33d4f53147687839b/ujson-5.11.0-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:6dd703c3e86dc6f7044c5ac0b3ae079ed96bf297974598116aa5fb7f655c3a60", size = 1195783, upload-time = "2025-08-20T11:55:58.081Z" }, + { url = "https://files.pythonhosted.org/packages/8d/20/78abe3d808cf3bb3e76f71fca46cd208317bf461c905d79f0d26b9df20f1/ujson-5.11.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:3772e4fe6b0c1e025ba3c50841a0ca4786825a4894c8411bf8d3afe3a8061328", size = 1088822, upload-time = "2025-08-20T11:55:59.469Z" }, + { url = "https://files.pythonhosted.org/packages/d8/50/8856e24bec5e2fc7f775d867aeb7a3f137359356200ac44658f1f2c834b2/ujson-5.11.0-cp313-cp313-win32.whl", hash = "sha256:8fa2af7c1459204b7a42e98263b069bd535ea0cd978b4d6982f35af5a04a4241", size = 39753, upload-time = "2025-08-20T11:56:01.345Z" }, + { url = "https://files.pythonhosted.org/packages/5b/d8/1baee0f4179a4d0f5ce086832147b6cc9b7731c24ca08e14a3fdb8d39c32/ujson-5.11.0-cp313-cp313-win_amd64.whl", hash = "sha256:34032aeca4510a7c7102bd5933f59a37f63891f30a0706fb46487ab6f0edf8f0", size = 43866, upload-time = "2025-08-20T11:56:02.552Z" }, + { url = "https://files.pythonhosted.org/packages/a9/8c/6d85ef5be82c6d66adced3ec5ef23353ed710a11f70b0b6a836878396334/ujson-5.11.0-cp313-cp313-win_arm64.whl", hash = "sha256:ce076f2df2e1aa62b685086fbad67f2b1d3048369664b4cdccc50707325401f9", size = 38363, upload-time = "2025-08-20T11:56:03.688Z" }, + { url = "https://files.pythonhosted.org/packages/28/08/4518146f4984d112764b1dfa6fb7bad691c44a401adadaa5e23ccd930053/ujson-5.11.0-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:65724738c73645db88f70ba1f2e6fb678f913281804d5da2fd02c8c5839af302", size = 55462, upload-time = "2025-08-20T11:56:04.873Z" }, + { url = "https://files.pythonhosted.org/packages/29/37/2107b9a62168867a692654d8766b81bd2fd1e1ba13e2ec90555861e02b0c/ujson-5.11.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:29113c003ca33ab71b1b480bde952fbab2a0b6b03a4ee4c3d71687cdcbd1a29d", size = 53246, upload-time = "2025-08-20T11:56:06.054Z" }, + { url = "https://files.pythonhosted.org/packages/9b/f8/25583c70f83788edbe3ca62ce6c1b79eff465d78dec5eb2b2b56b3e98b33/ujson-5.11.0-cp314-cp314-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c44c703842024d796b4c78542a6fcd5c3cb948b9fc2a73ee65b9c86a22ee3638", size = 57631, upload-time = "2025-08-20T11:56:07.374Z" }, + { url = "https://files.pythonhosted.org/packages/ed/ca/19b3a632933a09d696f10dc1b0dfa1d692e65ad507d12340116ce4f67967/ujson-5.11.0-cp314-cp314-manylinux_2_24_i686.manylinux_2_28_i686.whl", hash = "sha256:e750c436fb90edf85585f5c62a35b35082502383840962c6983403d1bd96a02c", size = 59877, upload-time = "2025-08-20T11:56:08.534Z" }, + { url = "https://files.pythonhosted.org/packages/55/7a/4572af5324ad4b2bfdd2321e898a527050290147b4ea337a79a0e4e87ec7/ujson-5.11.0-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f278b31a7c52eb0947b2db55a5133fbc46b6f0ef49972cd1a80843b72e135aba", size = 57363, upload-time = "2025-08-20T11:56:09.758Z" }, + { url = "https://files.pythonhosted.org/packages/7b/71/a2b8c19cf4e1efe53cf439cdf7198ac60ae15471d2f1040b490c1f0f831f/ujson-5.11.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:ab2cb8351d976e788669c8281465d44d4e94413718af497b4e7342d7b2f78018", size = 1036394, upload-time = "2025-08-20T11:56:11.168Z" }, + { url = "https://files.pythonhosted.org/packages/7a/3e/7b98668cba3bb3735929c31b999b374ebc02c19dfa98dfebaeeb5c8597ca/ujson-5.11.0-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:090b4d11b380ae25453100b722d0609d5051ffe98f80ec52853ccf8249dfd840", size = 1195837, upload-time = "2025-08-20T11:56:12.6Z" }, + { url = "https://files.pythonhosted.org/packages/a1/ea/8870f208c20b43571a5c409ebb2fe9b9dba5f494e9e60f9314ac01ea8f78/ujson-5.11.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:80017e870d882d5517d28995b62e4e518a894f932f1e242cbc802a2fd64d365c", size = 1088837, upload-time = "2025-08-20T11:56:14.15Z" }, + { url = "https://files.pythonhosted.org/packages/63/b6/c0e6607e37fa47929920a685a968c6b990a802dec65e9c5181e97845985d/ujson-5.11.0-cp314-cp314-win32.whl", hash = "sha256:1d663b96eb34c93392e9caae19c099ec4133ba21654b081956613327f0e973ac", size = 41022, upload-time = "2025-08-20T11:56:15.509Z" }, + { url = "https://files.pythonhosted.org/packages/4e/56/f4fe86b4c9000affd63e9219e59b222dc48b01c534533093e798bf617a7e/ujson-5.11.0-cp314-cp314-win_amd64.whl", hash = "sha256:849e65b696f0d242833f1df4182096cedc50d414215d1371fca85c541fbff629", size = 45111, upload-time = "2025-08-20T11:56:16.597Z" }, + { url = "https://files.pythonhosted.org/packages/0a/f3/669437f0280308db4783b12a6d88c00730b394327d8334cc7a32ef218e64/ujson-5.11.0-cp314-cp314-win_arm64.whl", hash = "sha256:e73df8648c9470af2b6a6bf5250d4744ad2cf3d774dcf8c6e31f018bdd04d764", size = 39682, upload-time = "2025-08-20T11:56:17.763Z" }, + { url = "https://files.pythonhosted.org/packages/6e/cd/e9809b064a89fe5c4184649adeb13c1b98652db3f8518980b04227358574/ujson-5.11.0-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:de6e88f62796372fba1de973c11138f197d3e0e1d80bcb2b8aae1e826096d433", size = 55759, upload-time = "2025-08-20T11:56:18.882Z" }, + { url = "https://files.pythonhosted.org/packages/1b/be/ae26a6321179ebbb3a2e2685b9007c71bcda41ad7a77bbbe164005e956fc/ujson-5.11.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:49e56ef8066f11b80d620985ae36869a3ff7e4b74c3b6129182ec5d1df0255f3", size = 53634, upload-time = "2025-08-20T11:56:20.012Z" }, + { url = "https://files.pythonhosted.org/packages/ae/e9/fb4a220ee6939db099f4cfeeae796ecb91e7584ad4d445d4ca7f994a9135/ujson-5.11.0-cp314-cp314t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1a325fd2c3a056cf6c8e023f74a0c478dd282a93141356ae7f16d5309f5ff823", size = 58547, upload-time = "2025-08-20T11:56:21.175Z" }, + { url = "https://files.pythonhosted.org/packages/bd/f8/fc4b952b8f5fea09ea3397a0bd0ad019e474b204cabcb947cead5d4d1ffc/ujson-5.11.0-cp314-cp314t-manylinux_2_24_i686.manylinux_2_28_i686.whl", hash = "sha256:a0af6574fc1d9d53f4ff371f58c96673e6d988ed2b5bf666a6143c782fa007e9", size = 60489, upload-time = "2025-08-20T11:56:22.342Z" }, + { url = "https://files.pythonhosted.org/packages/2e/e5/af5491dfda4f8b77e24cf3da68ee0d1552f99a13e5c622f4cef1380925c3/ujson-5.11.0-cp314-cp314t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:10f29e71ecf4ecd93a6610bd8efa8e7b6467454a363c3d6416db65de883eb076", size = 58035, upload-time = "2025-08-20T11:56:23.92Z" }, + { url = "https://files.pythonhosted.org/packages/c4/09/0945349dd41f25cc8c38d78ace49f14c5052c5bbb7257d2f466fa7bdb533/ujson-5.11.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:1a0a9b76a89827a592656fe12e000cf4f12da9692f51a841a4a07aa4c7ecc41c", size = 1037212, upload-time = "2025-08-20T11:56:25.274Z" }, + { url = "https://files.pythonhosted.org/packages/49/44/8e04496acb3d5a1cbee3a54828d9652f67a37523efa3d3b18a347339680a/ujson-5.11.0-cp314-cp314t-musllinux_1_2_i686.whl", hash = "sha256:b16930f6a0753cdc7d637b33b4e8f10d5e351e1fb83872ba6375f1e87be39746", size = 1196500, upload-time = "2025-08-20T11:56:27.517Z" }, + { url = "https://files.pythonhosted.org/packages/64/ae/4bc825860d679a0f208a19af2f39206dfd804ace2403330fdc3170334a2f/ujson-5.11.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:04c41afc195fd477a59db3a84d5b83a871bd648ef371cf8c6f43072d89144eef", size = 1089487, upload-time = "2025-08-20T11:56:29.07Z" }, + { url = "https://files.pythonhosted.org/packages/30/ed/5a057199fb0a5deabe0957073a1c1c1c02a3e99476cd03daee98ea21fa57/ujson-5.11.0-cp314-cp314t-win32.whl", hash = "sha256:aa6d7a5e09217ff93234e050e3e380da62b084e26b9f2e277d2606406a2fc2e5", size = 41859, upload-time = "2025-08-20T11:56:30.495Z" }, + { url = "https://files.pythonhosted.org/packages/aa/03/b19c6176bdf1dc13ed84b886e99677a52764861b6cc023d5e7b6ebda249d/ujson-5.11.0-cp314-cp314t-win_amd64.whl", hash = "sha256:48055e1061c1bb1f79e75b4ac39e821f3f35a9b82de17fce92c3140149009bec", size = 46183, upload-time = "2025-08-20T11:56:31.574Z" }, + { url = "https://files.pythonhosted.org/packages/5d/ca/a0413a3874b2dc1708b8796ca895bf363292f9c70b2e8ca482b7dbc0259d/ujson-5.11.0-cp314-cp314t-win_arm64.whl", hash = "sha256:1194b943e951092db611011cb8dbdb6cf94a3b816ed07906e14d3bc6ce0e90ab", size = 40264, upload-time = "2025-08-20T11:56:32.773Z" }, + { url = "https://files.pythonhosted.org/packages/50/17/30275aa2933430d8c0c4ead951cc4fdb922f575a349aa0b48a6f35449e97/ujson-5.11.0-pp311-pypy311_pp73-macosx_10_15_x86_64.whl", hash = "sha256:abae0fb58cc820092a0e9e8ba0051ac4583958495bfa5262a12f628249e3b362", size = 51206, upload-time = "2025-08-20T11:56:48.797Z" }, + { url = "https://files.pythonhosted.org/packages/c3/15/42b3924258eac2551f8f33fa4e35da20a06a53857ccf3d4deb5e5d7c0b6c/ujson-5.11.0-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:fac6c0649d6b7c3682a0a6e18d3de6857977378dce8d419f57a0b20e3d775b39", size = 48907, upload-time = "2025-08-20T11:56:50.136Z" }, + { url = "https://files.pythonhosted.org/packages/94/7e/0519ff7955aba581d1fe1fb1ca0e452471250455d182f686db5ac9e46119/ujson-5.11.0-pp311-pypy311_pp73-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4b42c115c7c6012506e8168315150d1e3f76e7ba0f4f95616f4ee599a1372bbc", size = 50319, upload-time = "2025-08-20T11:56:51.63Z" }, + { url = "https://files.pythonhosted.org/packages/74/cf/209d90506b7d6c5873f82c5a226d7aad1a1da153364e9ebf61eff0740c33/ujson-5.11.0-pp311-pypy311_pp73-manylinux_2_24_i686.manylinux_2_28_i686.whl", hash = "sha256:86baf341d90b566d61a394869ce77188cc8668f76d7bb2c311d77a00f4bdf844", size = 56584, upload-time = "2025-08-20T11:56:52.89Z" }, + { url = "https://files.pythonhosted.org/packages/e9/97/bd939bb76943cb0e1d2b692d7e68629f51c711ef60425fa5bb6968037ecd/ujson-5.11.0-pp311-pypy311_pp73-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4598bf3965fc1a936bd84034312bcbe00ba87880ef1ee33e33c1e88f2c398b49", size = 51588, upload-time = "2025-08-20T11:56:54.054Z" }, + { url = "https://files.pythonhosted.org/packages/52/5b/8c5e33228f7f83f05719964db59f3f9f276d272dc43752fa3bbf0df53e7b/ujson-5.11.0-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:416389ec19ef5f2013592f791486bef712ebce0cd59299bf9df1ba40bb2f6e04", size = 43835, upload-time = "2025-08-20T11:56:55.237Z" }, +] + [[package]] name = "ultralytics" version = "8.3.250" @@ -9517,6 +9946,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ad/e4/8d97cca767bcc1be76d16fb76951608305561c6e056811587f36cb1316a8/werkzeug-3.1.5-py3-none-any.whl", hash = "sha256:5111e36e91086ece91f93268bb39b4a35c1e6f1feac762c9c822ded0a4e322dc", size = 225025, upload-time = "2026-01-08T17:49:21.859Z" }, ] +[[package]] +name = "whatthepatch" +version = "1.0.7" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/06/28/55bc3e107a56fdcf7d5022cb32b8c21d98a9cc2df5cd9f3b93e10419099e/whatthepatch-1.0.7.tar.gz", hash = "sha256:9eefb4ebea5200408e02d413d2b4bc28daea6b78bb4b4d53431af7245f7d7edf", size = 34612, upload-time = "2024-11-16T17:21:22.153Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8e/93/af1d6ccb69ab6b5a00e03fa0cefa563f9862412667776ea15dd4eece3a90/whatthepatch-1.0.7-py3-none-any.whl", hash = "sha256:1b6f655fd31091c001c209529dfaabbabdbad438f5de14e3951266ea0fc6e7ed", size = 11964, upload-time = "2024-11-16T17:21:20.761Z" }, +] + [[package]] name = "wrapt" version = "1.17.3" From b63fba35ea0f7a2ee5aff5390533f49ec2ceabb1 Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Wed, 21 Jan 2026 17:48:29 +0800 Subject: [PATCH 02/39] Fix return type annotations for reactive pipe operators - quality_barrier: Callable[[Observable[T]], Observable[T]] - sharpness_barrier: Callable[[Observable[Image]], Observable[Image]] --- dimos/msgs/sensor_msgs/Image.py | 3 ++- dimos/utils/reactive.py | 4 +++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/dimos/msgs/sensor_msgs/Image.py b/dimos/msgs/sensor_msgs/Image.py index de3e7abeca..2bb99427bf 100644 --- a/dimos/msgs/sensor_msgs/Image.py +++ b/dimos/msgs/sensor_msgs/Image.py @@ -38,6 +38,7 @@ from dimos.utils.reactive import quality_barrier if TYPE_CHECKING: + from collections.abc import Callable import os from reactivex.observable import Observable @@ -724,7 +725,7 @@ def find_best(*_args): # type: ignore[no-untyped-def] ) -def sharpness_barrier(target_frequency: float): # type: ignore[no-untyped-def] +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/utils/reactive.py b/dimos/utils/reactive.py index bfc9cd0465..5cdd09ee03 100644 --- a/dimos/utils/reactive.py +++ b/dimos/utils/reactive.py @@ -211,7 +211,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. From d1bca7ce8ac698a8fe124b1fa2324a04bdd29602 Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Wed, 21 Jan 2026 19:14:11 +0800 Subject: [PATCH 03/39] sketching a new embedding db --- dimos/core/stream.py | 2 +- dimos/memory/embedding.py | 101 +++++++++++++++++++++++++++++++++ dimos/memory/test_embedding.py | 50 ++++++++++++++++ dimos/protocol/tf/tf.py | 12 ++++ 4 files changed, 164 insertions(+), 1 deletion(-) create mode 100644 dimos/memory/embedding.py create mode 100644 dimos/memory/test_embedding.py diff --git a/dimos/core/stream.py b/dimos/core/stream.py index 64a1e0edce..64d4870721 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/memory/embedding.py b/dimos/memory/embedding.py new file mode 100644 index 0000000000..d4fe4a1534 --- /dev/null +++ b/dimos/memory/embedding.py @@ -0,0 +1,101 @@ +# 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 dataclasses import dataclass, field +from typing import cast + +from reactivex import operators as ops + +from dimos.core import In, Module, ModuleConfig, Out, 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.sensor_msgs import Image +from dimos.msgs.sensor_msgs.Image import Image, sharpness_barrier +from dimos.utils.reactive import backpressure, 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 + # taking into account message processing frequency downstream + # + # 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.map(self._to_spatial_entry), + ops.filter(self._has_pose), + ops.map(self._embed_spatial_entry), + ops.map(self._store_spatial_entry), + ).subscribe(print) + + def _has_pose(self, entry: SpatialEntry) -> bool: + return entry.pose is not None + + def _to_spatial_entry(self, img: Image) -> SpatialEntry: + return SpatialEntry( + image=img, + pose=self.tf.get_pose("world", "base_link"), + ) + + 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 | None: + return spatial_embedding diff --git a/dimos/memory/test_embedding.py b/dimos/memory/test_embedding.py new file mode 100644 index 0000000000..779d1af314 --- /dev/null +++ b/dimos/memory/test_embedding.py @@ -0,0 +1,50 @@ +# 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 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" + + +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.process_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/protocol/tf/tf.py b/dimos/protocol/tf/tf.py index 3688b013cf..a52b4c5c3d 100644 --- a/dimos/protocol/tf/tf.py +++ b/dimos/protocol/tf/tf.py @@ -334,6 +334,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, + ): + 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) From 7098c773928fc52f51fd113517314c25fa2d8c5a Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Wed, 21 Jan 2026 19:37:30 +0800 Subject: [PATCH 04/39] benchmark prints loss heatmap --- dimos/memory/embedding.py | 2 -- dimos/protocol/pubsub/benchmark/test_benchmark.py | 1 + dimos/protocol/pubsub/benchmark/type.py | 8 ++++++++ 3 files changed, 9 insertions(+), 2 deletions(-) diff --git a/dimos/memory/embedding.py b/dimos/memory/embedding.py index d4fe4a1534..9a3925e8ab 100644 --- a/dimos/memory/embedding.py +++ b/dimos/memory/embedding.py @@ -60,13 +60,11 @@ def get_costmap(self) -> OccupancyGrid: 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 - # taking into account message processing frequency downstream # # we need a Governor system, sharpness_barrier frequency shouldn't # be a fixed float but an observable that adjusts based on downstream load 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 d27da2fde6..d7b5ff3c8b 100644 --- a/dimos/protocol/pubsub/benchmark/type.py +++ b/dimos/protocol/pubsub/benchmark/type.py @@ -275,3 +275,11 @@ def fmt(v: float) -> str: return f"{v * 1000:.0f}ms" self._print_heatmap("Latency", lambda r: r.receive_time, 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) From 472e2793188e34e7012e8b93f846c02041b46e54 Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Wed, 21 Jan 2026 20:41:48 +0800 Subject: [PATCH 05/39] typing --- dimos/memory/embedding.py | 27 ++++++++++++++++----------- dimos/memory/test_embedding.py | 2 +- dimos/utils/reactive.py | 3 ++- 3 files changed, 19 insertions(+), 13 deletions(-) diff --git a/dimos/memory/embedding.py b/dimos/memory/embedding.py index 9a3925e8ab..c9dc6f60d2 100644 --- a/dimos/memory/embedding.py +++ b/dimos/memory/embedding.py @@ -12,15 +12,19 @@ # 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, Out, 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 backpressure, getter_hot @@ -72,20 +76,16 @@ def start(self) -> None: # (also voxel size for mapper for example would benefit from this) self.color_image.pure_observable().pipe( sharpness_barrier(0.5), - ops.map(self._to_spatial_entry), - ops.filter(self._has_pose), + ops.flat_map(self._try_create_spatial_entry), ops.map(self._embed_spatial_entry), ops.map(self._store_spatial_entry), ).subscribe(print) - def _has_pose(self, entry: SpatialEntry) -> bool: - return entry.pose is not None - - def _to_spatial_entry(self, img: Image) -> SpatialEntry: - return SpatialEntry( - image=img, - pose=self.tf.get_pose("world", "base_link"), - ) + 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)) @@ -95,5 +95,10 @@ def _embed_spatial_entry(self, spatial_entry: SpatialEntry) -> SpatialEmbedding: embedding=embedding, ) - def _store_spatial_entry(self, spatial_embedding: SpatialEmbedding) -> SpatialEmbedding | None: + 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 index 779d1af314..ffd9e21a4f 100644 --- a/dimos/memory/test_embedding.py +++ b/dimos/memory/test_embedding.py @@ -38,7 +38,7 @@ def test_embed_frame() -> None: spatial_entry = SpatialEntry(image=frame, pose=dummy_pose) # Embed the frame - result = memory.process_spatial_entry(spatial_entry) + result = memory._embed_spatial_entry(spatial_entry) # Verify assert result is not None diff --git a/dimos/utils/reactive.py b/dimos/utils/reactive.py index 5cdd09ee03..c6decb92b4 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] From faabbf0d15f7b7630d4e7939b42eda051d02bf1e Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Wed, 21 Jan 2026 21:44:28 +0800 Subject: [PATCH 06/39] get_data supports nested paths --- .../pointclouds/test_occupancy_speed.py | 4 +- dimos/utils/data.py | 42 ++++++++++++------- dimos/utils/test_data.py | 4 +- dimos/utils/testing/replay.py | 4 +- 4 files changed, 34 insertions(+), 20 deletions(-) diff --git a/dimos/mapping/pointclouds/test_occupancy_speed.py b/dimos/mapping/pointclouds/test_occupancy_speed.py index c34c2865f2..e296494a45 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 @@ -32,7 +32,7 @@ def test_build_map(): print(ts, frame) 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/utils/data.py b/dimos/utils/data.py index 4ba9c73b0c..70f41e6edb 100644 --- a/dimos/utils/data.py +++ b/dimos/utils/data.py @@ -98,7 +98,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" @@ -106,7 +106,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: @@ -167,7 +167,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) @@ -204,7 +204,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. @@ -215,29 +215,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/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/replay.py b/dimos/utils/testing/replay.py index 89225c322e..be8d394f76 100644 --- a/dimos/utils/testing/replay.py +++ b/dimos/utils/testing/replay.py @@ -29,7 +29,7 @@ from reactivex.observable import Observable from reactivex.scheduler import TimeoutScheduler -from dimos.utils.data import _get_data_dir, get_data +from dimos.utils.data import get_data, get_data_dir T = TypeVar("T") @@ -120,7 +120,7 @@ def __init__(self, name: str, autocast: Callable[[T], Any] | None = None) -> Non self.cnt = 0 # Create storage directory in the data dir - self.root_dir = _get_data_dir() / name + self.root_dir = get_data_dir() / name # Check if directory exists and is not empty if self.root_dir.exists(): From eef2b3112278736f09ad450f7d94b78e7307bc8f Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Wed, 21 Jan 2026 21:44:42 +0800 Subject: [PATCH 07/39] sensor storage sketch and pickle implementation --- dimos/memory/sensor/__init__.py | 19 ++++ dimos/memory/sensor/base.py | 170 +++++++++++++++++++++++++++++++ dimos/memory/sensor/pickledir.py | 159 +++++++++++++++++++++++++++++ dimos/memory/sensor/test_base.py | 112 ++++++++++++++++++++ 4 files changed, 460 insertions(+) create mode 100644 dimos/memory/sensor/__init__.py create mode 100644 dimos/memory/sensor/base.py create mode 100644 dimos/memory/sensor/pickledir.py create mode 100644 dimos/memory/sensor/test_base.py diff --git a/dimos/memory/sensor/__init__.py b/dimos/memory/sensor/__init__.py new file mode 100644 index 0000000000..67746f6b13 --- /dev/null +++ b/dimos/memory/sensor/__init__.py @@ -0,0 +1,19 @@ +# 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. +"""Sensor storage and replay.""" + +from dimos.memory.sensor.base import InMemoryStore, SensorStore +from dimos.memory.sensor.pickledir import PickleDirStore + +__all__ = ["InMemoryStore", "PickleDirStore", "SensorStore"] diff --git a/dimos/memory/sensor/base.py b/dimos/memory/sensor/base.py new file mode 100644 index 0000000000..9767312845 --- /dev/null +++ b/dimos/memory/sensor/base.py @@ -0,0 +1,170 @@ +# 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 timestamped sensor storage and replay.""" + +from abc import ABC, abstractmethod +from collections.abc import Iterator +from typing import Generic, TypeVar + +from reactivex.observable import Observable + +T = TypeVar("T") + + +class SensorStore(Generic[T], ABC): + """Unified storage + replay for timestamped sensor data. + + Implement 4 abstract methods for your backend (in-memory, pickle, sqlite, etc.). + All iteration, streaming, and seek logic comes free from the base class. + """ + + # === Abstract - implement for your backend === + + @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 _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.).""" + ... + + def save(self, data: T, timestamp: float | None = None) -> None: + """Save data. Uses data.ts if available, otherwise timestamp arg, otherwise now.""" + import time + + if timestamp is None: + if hasattr(data, "ts"): + timestamp = data.ts # type: ignore[union-attr] + else: + timestamp = time.time() + self._save(timestamp, data) + + def load(self, timestamp: float) -> T | None: + """Load data at exact timestamp.""" + return self._load(timestamp) + + def find_closest( + self, + timestamp: float | None = None, + seek: float | None = None, + tolerance: float | None = None, + ) -> T | None: + """Find data closest to timestamp (absolute) or seek (relative to start).""" + ... + + def first_timestamp(self) -> float | None: + """Get the first timestamp in the store.""" + ... + + def iterate(self, loop: bool = False) -> Iterator[tuple[float, T]]: + """Iterate over (timestamp, data) pairs.""" + ... + + 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 seek/duration options.""" + ... + + def iterate_realtime(self, speed: float = 1.0, **kwargs) -> Iterator[T]: + """Iterate data, sleeping to match original timing.""" + ... + + 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.""" + ... + + +class InMemoryStore(SensorStore[T]): + """In-memory storage using dict. Good for live use.""" + + def __init__(self) -> None: + self._data: dict[float, T] = {} + self._sorted_timestamps: list[float] | None = None + + def _save(self, timestamp: float, data: T) -> None: + self._data[timestamp] = data + self._sorted_timestamps = None # Invalidate cache + + def _load(self, timestamp: float) -> T | None: + return self._data.get(timestamp) + + def _iter_items( + self, start: float | None = None, end: float | None = None + ) -> Iterator[tuple[float, T]]: + for ts in self._get_sorted_timestamps(): + if start is not None and ts < start: + continue + if end is not None and ts >= end: + break + yield (ts, self._data[ts]) + + def _find_closest_timestamp( + self, timestamp: float, tolerance: float | None = None + ) -> float | None: + import bisect + + timestamps = self._get_sorted_timestamps() + if not timestamps: + return None + + pos = bisect.bisect_left(timestamps, timestamp) + + 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_sorted_timestamps(self) -> list[float]: + if self._sorted_timestamps is None: + self._sorted_timestamps = sorted(self._data.keys()) + return self._sorted_timestamps diff --git a/dimos/memory/sensor/pickledir.py b/dimos/memory/sensor/pickledir.py new file mode 100644 index 0000000000..cce8d6389a --- /dev/null +++ b/dimos/memory/sensor/pickledir.py @@ -0,0 +1,159 @@ +# 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 SensorStore.""" + +import bisect +from collections.abc import Iterator +from functools import lru_cache +import glob +import os +from pathlib import Path +import pickle + +from dimos.memory.sensor.base import SensorStore, T +from dimos.utils.data import get_data, get_data_dir + + +class PickleDirStore(SensorStore[T]): + """Pickle directory backend. Compatible with TimedSensorStorage/TimedSensorReplay format. + + Directory structure: + {name}/ + 000.pickle # (timestamp, data) + 001.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.store(timestamp, image) + """ + + 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 _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 + + @lru_cache(maxsize=128) + def _load_file(self, filepath: Path) -> T | None: + """Load data from a pickle file (LRU cached).""" + try: + with open(filepath, "rb") as f: + return pickle.load(f) + except Exception: + return None diff --git a/dimos/memory/sensor/test_base.py b/dimos/memory/sensor/test_base.py new file mode 100644 index 0000000000..7acc4d2d53 --- /dev/null +++ b/dimos/memory/sensor/test_base.py @@ -0,0 +1,112 @@ +# 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 SensorStore implementations.""" + +from pathlib import Path +import tempfile + +import pytest + +from dimos.memory.sensor.base import InMemoryStore, SensorStore +from dimos.memory.sensor.pickledir import PickleDirStore + + +@pytest.fixture +def temp_pickle_dir(): + """Create a temporary directory for PickleDirStore tests.""" + with tempfile.TemporaryDirectory() as tmpdir: + yield tmpdir + + +def make_in_memory_store() -> SensorStore[str]: + return InMemoryStore[str]() + + +def make_pickle_dir_store(tmpdir: str) -> SensorStore[str]: + return PickleDirStore[str](tmpdir) + + +@pytest.mark.parametrize( + "store_factory,store_name", + [ + (lambda _: make_in_memory_store(), "InMemoryStore"), + (lambda tmpdir: make_pickle_dir_store(tmpdir), "PickleDirStore"), + ], +) +class TestSensorStore: + """Parametrized tests for all SensorStore implementations.""" + + def test_save_and_load(self, store_factory, store_name, temp_pickle_dir): + store = store_factory(temp_pickle_dir) + store.save("data_at_1", 1.0) + store.save("data_at_2", 2.0) + + assert store.load(1.0) == "data_at_1" + assert store.load(2.0) == "data_at_2" + assert store.load(3.0) is None + + def test_find_closest_timestamp(self, store_factory, store_name, temp_pickle_dir): + store = store_factory(temp_pickle_dir) + store.save("a", 1.0) + store.save("b", 2.0) + store.save("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_pickle_dir): + store = store_factory(temp_pickle_dir) + store.save("a", 1.0) + store.save("c", 3.0) + store.save("b", 2.0) + + # Should iterate in timestamp order + items = list(store._iter_items()) + assert items == [(1.0, "a"), (2.0, "b"), (3.0, "c")] + + def test_iter_items_with_range(self, store_factory, store_name, temp_pickle_dir): + store = store_factory(temp_pickle_dir) + store.save("a", 1.0) + store.save("b", 2.0) + store.save("c", 3.0) + store.save("d", 4.0) + + # Start only + items = list(store._iter_items(start=2.0)) + assert items == [(2.0, "b"), (3.0, "c"), (4.0, "d")] + + # End only + items = list(store._iter_items(end=3.0)) + assert items == [(1.0, "a"), (2.0, "b")] + + # Both + items = list(store._iter_items(start=2.0, end=4.0)) + assert items == [(2.0, "b"), (3.0, "c")] + + def test_empty_store(self, store_factory, store_name, temp_pickle_dir): + store = store_factory(temp_pickle_dir) + + assert store.load(1.0) is None + assert store._find_closest_timestamp(1.0) is None + assert list(store._iter_items()) == [] From f3c672afb197f3e9a57c7d7a38dfc4fb5024c16b Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Wed, 21 Jan 2026 21:51:18 +0800 Subject: [PATCH 08/39] Implement SensorStore convenience methods with drift-free streaming - Implement find_closest(), first_timestamp(), iterate(), iterate_ts(), iterate_realtime() methods using abstract _iter_items/_find_closest_timestamp - Add scheduler-based stream() with absolute time reference to prevent timing drift during long playback (ported from replay.py) - Move imports to top of file, add proper typing throughout - Fix pickledir.py mypy error (pickle.load returns Any) --- dimos/memory/sensor/base.py | 151 ++++++++++++++++++++++++++++--- dimos/memory/sensor/pickledir.py | 3 +- 2 files changed, 141 insertions(+), 13 deletions(-) diff --git a/dimos/memory/sensor/base.py b/dimos/memory/sensor/base.py index 9767312845..8a2948d966 100644 --- a/dimos/memory/sensor/base.py +++ b/dimos/memory/sensor/base.py @@ -14,10 +14,16 @@ """Unified timestamped sensor storage and replay.""" from abc import ABC, abstractmethod +import bisect from collections.abc import Iterator +import threading +import time from typing import Generic, TypeVar +import reactivex as rx +from reactivex.disposable import CompositeDisposable, Disposable from reactivex.observable import Observable +from reactivex.scheduler import TimeoutScheduler T = TypeVar("T") @@ -57,8 +63,6 @@ def _find_closest_timestamp( def save(self, data: T, timestamp: float | None = None) -> None: """Save data. Uses data.ts if available, otherwise timestamp arg, otherwise now.""" - import time - if timestamp is None: if hasattr(data, "ts"): timestamp = data.ts # type: ignore[union-attr] @@ -77,15 +81,33 @@ def find_closest( tolerance: float | None = None, ) -> T | None: """Find data closest to timestamp (absolute) or seek (relative to start).""" - ... + if timestamp is None and seek is None: + raise ValueError("Must provide either timestamp or seek") + + if seek is not None: + first = self.first_timestamp() + if first is None: + return None + timestamp = first + seek + + assert timestamp is not None + closest_ts = self._find_closest_timestamp(timestamp, tolerance) + if closest_ts is None: + return None + return self._load(closest_ts) def first_timestamp(self) -> float | None: """Get the first timestamp in the store.""" - ... + for ts, _ in self._iter_items(): + return ts + return None def iterate(self, loop: bool = False) -> Iterator[tuple[float, T]]: """Iterate over (timestamp, data) pairs.""" - ... + while True: + yield from self._iter_items() + if not loop: + break def iterate_ts( self, @@ -95,11 +117,39 @@ def iterate_ts( loop: bool = False, ) -> Iterator[tuple[float, T]]: """Iterate with seek/duration options.""" - ... + first = self.first_timestamp() + if first is None: + return + + # Calculate start timestamp + if from_timestamp is not None: + start = from_timestamp + elif seek is not None: + start = first + seek + else: + start = None + + # Calculate end timestamp + 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_realtime(self, speed: float = 1.0, **kwargs) -> Iterator[T]: + def iterate_realtime(self, speed: float = 1.0, **kwargs: float | bool | None) -> Iterator[T]: """Iterate data, sleeping to match original timing.""" - ... + prev_ts: float | None = None + for ts, data in self.iterate_ts(**kwargs): # type: ignore[arg-type] + 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, @@ -109,8 +159,87 @@ def stream( from_timestamp: float | None = None, loop: bool = False, ) -> Observable[T]: - """Stream data as Observable with timing control.""" - ... + """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_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 (absolute time prevents drift) + 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: 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 + + # 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( + _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) class InMemoryStore(SensorStore[T]): @@ -140,8 +269,6 @@ def _iter_items( def _find_closest_timestamp( self, timestamp: float, tolerance: float | None = None ) -> float | None: - import bisect - timestamps = self._get_sorted_timestamps() if not timestamps: return None diff --git a/dimos/memory/sensor/pickledir.py b/dimos/memory/sensor/pickledir.py index cce8d6389a..0f3e8388d4 100644 --- a/dimos/memory/sensor/pickledir.py +++ b/dimos/memory/sensor/pickledir.py @@ -154,6 +154,7 @@ def _load_file(self, filepath: Path) -> T | None: """Load data from a pickle file (LRU cached).""" try: with open(filepath, "rb") as f: - return pickle.load(f) + data: T = pickle.load(f) + return data except Exception: return None From b55f058e58ff8b970425492ffe4881a90c7b6699 Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Wed, 21 Jan 2026 21:59:06 +0800 Subject: [PATCH 09/39] Add SqliteStore backend for sensor data - Single-file SQLite storage with indexed timestamp queries - BLOB storage for pickled sensor data - INSERT OR REPLACE for duplicate timestamp handling - Supports multiple tables per database (different sensors) - Added to parametrized tests (15 tests across 3 backends) --- dimos/memory/sensor/__init__.py | 3 +- dimos/memory/sensor/sqlite.py | 175 +++++++++++++++++++++++++++++++ dimos/memory/sensor/test_base.py | 30 +++--- 3 files changed, 195 insertions(+), 13 deletions(-) create mode 100644 dimos/memory/sensor/sqlite.py diff --git a/dimos/memory/sensor/__init__.py b/dimos/memory/sensor/__init__.py index 67746f6b13..2dba973255 100644 --- a/dimos/memory/sensor/__init__.py +++ b/dimos/memory/sensor/__init__.py @@ -15,5 +15,6 @@ from dimos.memory.sensor.base import InMemoryStore, SensorStore from dimos.memory.sensor.pickledir import PickleDirStore +from dimos.memory.sensor.sqlite import SqliteStore -__all__ = ["InMemoryStore", "PickleDirStore", "SensorStore"] +__all__ = ["InMemoryStore", "PickleDirStore", "SensorStore", "SqliteStore"] diff --git a/dimos/memory/sensor/sqlite.py b/dimos/memory/sensor/sqlite.py new file mode 100644 index 0000000000..55a650e5a7 --- /dev/null +++ b/dimos/memory/sensor/sqlite.py @@ -0,0 +1,175 @@ +# 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 SensorStore.""" + +from collections.abc import Iterator +from pathlib import Path +import pickle +import sqlite3 + +from dimos.memory.sensor.base import SensorStore, T + + +class SqliteStore(SensorStore[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: + # File-based (persistent) + store = SqliteStore("sensors.db", table="lidar") + store.save(data, timestamp) + + # In-memory (for testing) + store = SqliteStore(":memory:") + + # Query by timestamp + data = store.find_closest(seek=10.0) + """ + + def __init__(self, db_path: str | Path, table: str = "sensor_data") -> None: + """ + Args: + db_path: Path to SQLite database file, or ":memory:" for in-memory. + table: Table name for this sensor's data. + """ + self._db_path = str(db_path) + self._table = table + self._conn: sqlite3.Connection | None = None + + def _get_conn(self) -> sqlite3.Connection: + """Get or create database connection.""" + if self._conn is None: + # Create parent directory if needed (for file-based DBs) + if self._db_path != ":memory:": + Path(self._db_path).parent.mkdir(parents=True, exist_ok=True) + + self._conn = sqlite3.connect(self._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 _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 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/sensor/test_base.py b/dimos/memory/sensor/test_base.py index 7acc4d2d53..19c1557a8a 100644 --- a/dimos/memory/sensor/test_base.py +++ b/dimos/memory/sensor/test_base.py @@ -20,11 +20,12 @@ from dimos.memory.sensor.base import InMemoryStore, SensorStore from dimos.memory.sensor.pickledir import PickleDirStore +from dimos.memory.sensor.sqlite import SqliteStore @pytest.fixture -def temp_pickle_dir(): - """Create a temporary directory for PickleDirStore tests.""" +def temp_dir(): + """Create a temporary directory for file-based store tests.""" with tempfile.TemporaryDirectory() as tmpdir: yield tmpdir @@ -37,18 +38,23 @@ def make_pickle_dir_store(tmpdir: str) -> SensorStore[str]: return PickleDirStore[str](tmpdir) +def make_sqlite_store(tmpdir: str) -> SensorStore[str]: + return SqliteStore[str](Path(tmpdir) / "test.db") + + @pytest.mark.parametrize( "store_factory,store_name", [ (lambda _: make_in_memory_store(), "InMemoryStore"), (lambda tmpdir: make_pickle_dir_store(tmpdir), "PickleDirStore"), + (lambda tmpdir: make_sqlite_store(tmpdir), "SqliteStore"), ], ) class TestSensorStore: """Parametrized tests for all SensorStore implementations.""" - def test_save_and_load(self, store_factory, store_name, temp_pickle_dir): - store = store_factory(temp_pickle_dir) + def test_save_and_load(self, store_factory, store_name, temp_dir): + store = store_factory(temp_dir) store.save("data_at_1", 1.0) store.save("data_at_2", 2.0) @@ -56,8 +62,8 @@ def test_save_and_load(self, store_factory, store_name, temp_pickle_dir): assert store.load(2.0) == "data_at_2" assert store.load(3.0) is None - def test_find_closest_timestamp(self, store_factory, store_name, temp_pickle_dir): - store = store_factory(temp_pickle_dir) + def test_find_closest_timestamp(self, store_factory, store_name, temp_dir): + store = store_factory(temp_dir) store.save("a", 1.0) store.save("b", 2.0) store.save("c", 3.0) @@ -75,8 +81,8 @@ def test_find_closest_timestamp(self, store_factory, store_name, temp_pickle_dir 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_pickle_dir): - store = store_factory(temp_pickle_dir) + def test_iter_items(self, store_factory, store_name, temp_dir): + store = store_factory(temp_dir) store.save("a", 1.0) store.save("c", 3.0) store.save("b", 2.0) @@ -85,8 +91,8 @@ def test_iter_items(self, store_factory, store_name, temp_pickle_dir): items = list(store._iter_items()) assert items == [(1.0, "a"), (2.0, "b"), (3.0, "c")] - def test_iter_items_with_range(self, store_factory, store_name, temp_pickle_dir): - store = store_factory(temp_pickle_dir) + def test_iter_items_with_range(self, store_factory, store_name, temp_dir): + store = store_factory(temp_dir) store.save("a", 1.0) store.save("b", 2.0) store.save("c", 3.0) @@ -104,8 +110,8 @@ def test_iter_items_with_range(self, store_factory, store_name, temp_pickle_dir) items = list(store._iter_items(start=2.0, end=4.0)) assert items == [(2.0, "b"), (3.0, "c")] - def test_empty_store(self, store_factory, store_name, temp_pickle_dir): - store = store_factory(temp_pickle_dir) + 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 From 493d70f4cb93f3db6ae18e6590c7053f1b3a4e17 Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Wed, 21 Jan 2026 22:16:02 +0800 Subject: [PATCH 10/39] Add PostgresStore backend for sensor data - PostgresStore implements SensorStore[T] + Resource for lifecycle management - Multiple stores can share same database with different tables - Tables created automatically on first save - Tests are optional - skip gracefully if PostgreSQL not available - Added psycopg2-binary and types-psycopg2 dependencies - Includes reset_db() helper for simple migrations (drop/recreate) --- dimos/memory/sensor/__init__.py | 10 +- dimos/memory/sensor/postgres.py | 243 +++++++++++++++++++++++++++++++ dimos/memory/sensor/test_base.py | 59 +++++++- pyproject.toml | 13 +- uv.lock | 82 +++++++++++ 5 files changed, 391 insertions(+), 16 deletions(-) create mode 100644 dimos/memory/sensor/postgres.py diff --git a/dimos/memory/sensor/__init__.py b/dimos/memory/sensor/__init__.py index 2dba973255..7be151f2c4 100644 --- a/dimos/memory/sensor/__init__.py +++ b/dimos/memory/sensor/__init__.py @@ -15,6 +15,14 @@ from dimos.memory.sensor.base import InMemoryStore, SensorStore from dimos.memory.sensor.pickledir import PickleDirStore +from dimos.memory.sensor.postgres import PostgresStore, reset_db from dimos.memory.sensor.sqlite import SqliteStore -__all__ = ["InMemoryStore", "PickleDirStore", "SensorStore", "SqliteStore"] +__all__ = [ + "InMemoryStore", + "PickleDirStore", + "PostgresStore", + "SensorStore", + "SqliteStore", + "reset_db", +] diff --git a/dimos/memory/sensor/postgres.py b/dimos/memory/sensor/postgres.py new file mode 100644 index 0000000000..511f4b151a --- /dev/null +++ b/dimos/memory/sensor/postgres.py @@ -0,0 +1,243 @@ +# 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 SensorStore.""" + +from collections.abc import Iterator +import pickle + +import psycopg2 +from psycopg2.extensions import connection as PgConnection + +from dimos.core.resource import Resource +from dimos.memory.sensor.base import SensorStore, T + + +class PostgresStore(SensorStore[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, timestamp) + 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. + db: Database name. + host: PostgreSQL host. + port: PostgreSQL port. + user: PostgreSQL user. Defaults to current system user. + """ + self._table = table + self._db = db + self._host = host + self._port = port + self._user = user + self._conn: PgConnection | 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) -> PgConnection: + """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 _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 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. + host: PostgreSQL host. + port: PostgreSQL port. + """ + # 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/sensor/test_base.py b/dimos/memory/sensor/test_base.py index 19c1557a8a..675701fbbd 100644 --- a/dimos/memory/sensor/test_base.py +++ b/dimos/memory/sensor/test_base.py @@ -15,6 +15,7 @@ from pathlib import Path import tempfile +import uuid import pytest @@ -42,14 +43,56 @@ def make_sqlite_store(tmpdir: str) -> SensorStore[str]: return SqliteStore[str](Path(tmpdir) / "test.db") -@pytest.mark.parametrize( - "store_factory,store_name", - [ - (lambda _: make_in_memory_store(), "InMemoryStore"), - (lambda tmpdir: make_pickle_dir_store(tmpdir), "PickleDirStore"), - (lambda tmpdir: make_sqlite_store(tmpdir), "SqliteStore"), - ], -) +# 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"), +] + +# Track postgres tables to clean up +_postgres_tables: list[str] = [] + +try: + import psycopg2 + + from dimos.memory.sensor.postgres import PostgresStore + + # Test connection + _test_conn = psycopg2.connect(dbname="dimensional") + _test_conn.close() + + def make_postgres_store(_tmpdir: str) -> SensorStore[str]: + """Create PostgresStore with unique table name.""" + table = f"test_{uuid.uuid4().hex[:8]}" + _postgres_tables.append(table) + store = PostgresStore[str](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 TestSensorStore: """Parametrized tests for all SensorStore implementations.""" diff --git a/pyproject.toml b/pyproject.toml index e04dd4bbfe..23f783399e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,7 +32,6 @@ dependencies = [ # Transport Protocols "dimos-lcm", "PyTurboJPEG==1.8.2", - # Core "numpy>=1.26.4", "scipy>=1.15.1", @@ -41,33 +40,28 @@ dependencies = [ "sortedcontainers==2.4.0", "pydantic", "python-dotenv", - # Multiprocess "dask[complete]==2025.5.1", "plum-dispatch==2.5.7", - # Logging "structlog>=25.5.0,<26", "colorlog==6.9.0", - # Core Msgs "opencv-python", "open3d", - # CLI "pydantic-settings>=2.11.0,<3", "textual==3.7.1", "terminaltexteffects==0.12.2", "typer>=0.19.2,<1", "plotext==5.3.2", - # Used for calculating the occupancy map. "numba>=0.60.0", # First version supporting Python 3.12 "llvmlite>=0.42.0", # Required by numba 0.60+ - # TODO: rerun shouldn't be required but rn its in core (there is NO WAY to use dimos without rerun rn) # remove this once rerun is optional in core "rerun-sdk>=0.20.0", + "psycopg2-binary>=2.9.11", ] @@ -276,6 +270,11 @@ base = [ "dimos[agents,web,perception,visualization,sim]", ] +[dependency-groups] +dev = [ + "types-psycopg2>=2.9.21.20251012", +] + [tool.ruff] line-length = 100 exclude = [ diff --git a/uv.lock b/uv.lock index 486e216e20..31c94e992f 100644 --- a/uv.lock +++ b/uv.lock @@ -1488,6 +1488,7 @@ dependencies = [ { name = "opencv-python" }, { name = "plotext" }, { name = "plum-dispatch" }, + { name = "psycopg2-binary" }, { name = "pydantic" }, { name = "pydantic-settings" }, { name = "python-dotenv" }, @@ -1702,6 +1703,11 @@ web = [ { name = "uvicorn" }, ] +[package.dev-dependencies] +dev = [ + { name = "types-psycopg2" }, +] + [package.metadata] requires-dist = [ { name = "anthropic", marker = "extra == 'agents'", specifier = ">=0.19.0" }, @@ -1773,6 +1779,7 @@ requires-dist = [ { name = "plotly", marker = "extra == 'manipulation'", specifier = ">=5.9.0" }, { name = "plum-dispatch", specifier = "==2.5.7" }, { name = "pre-commit", marker = "extra == 'dev'", specifier = "==4.2.0" }, + { name = "psycopg2-binary", specifier = ">=2.9.11" }, { name = "pydantic" }, { name = "pydantic-settings", specifier = ">=2.11.0,<3" }, { name = "pygame", marker = "extra == 'sim'", specifier = ">=2.6.1" }, @@ -1844,6 +1851,9 @@ requires-dist = [ ] provides-extras = ["misc", "visualization", "agents", "web", "perception", "unitree", "manipulation", "cpu", "cuda", "dev", "sim", "drone", "base"] +[package.metadata.requires-dev] +dev = [{ name = "types-psycopg2", specifier = ">=2.9.21.20251012" }] + [[package]] name = "dimos-lcm" version = "0.1.1" @@ -6353,6 +6363,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" @@ -9300,6 +9373,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" From ce1eab7c87a8921b8e0909b5c0fbc4e8b1877365 Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Thu, 22 Jan 2026 11:33:35 +0800 Subject: [PATCH 11/39] Add SQL identifier validation and fix mypy issues - Validate table/database names in SqliteStore and PostgresStore using regex (alphanumeric/underscore, not starting with digit) - Fix Transform.to_pose() return type using TYPE_CHECKING import - Add return type annotation to TF.get_pose() - Fix ambiguous doclink in transports.md --- dimos/memory/sensor/postgres.py | 26 +++++++++++++++++++++----- dimos/memory/sensor/sqlite.py | 19 +++++++++++++++++-- dimos/msgs/geometry_msgs/Transform.py | 12 ++++++++---- dimos/protocol/tf/tf.py | 4 ++-- docs/concepts/transports.md | 2 +- 5 files changed, 49 insertions(+), 14 deletions(-) diff --git a/dimos/memory/sensor/postgres.py b/dimos/memory/sensor/postgres.py index 511f4b151a..4b33c25158 100644 --- a/dimos/memory/sensor/postgres.py +++ b/dimos/memory/sensor/postgres.py @@ -15,6 +15,7 @@ from collections.abc import Iterator import pickle +import re import psycopg2 from psycopg2.extensions import connection as PgConnection @@ -22,6 +23,20 @@ from dimos.core.resource import Resource from dimos.memory.sensor.base import SensorStore, T +# 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(SensorStore[T], Resource): """PostgreSQL backend for sensor data. @@ -59,14 +74,14 @@ def __init__( ) -> None: """ Args: - table: Table name for this sensor's data. - db: Database name. + 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 = table - self._db = db + self._table = _validate_identifier(table) + self._db = _validate_identifier(db) self._host = host self._port = port self._user = user @@ -221,10 +236,11 @@ def reset_db(db: str = "dimensional", host: str = "localhost", port: int = 5432) WARNING: This deletes all data in the database! Args: - db: Database name to reset. + 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 diff --git a/dimos/memory/sensor/sqlite.py b/dimos/memory/sensor/sqlite.py index 55a650e5a7..a2376b89af 100644 --- a/dimos/memory/sensor/sqlite.py +++ b/dimos/memory/sensor/sqlite.py @@ -16,10 +16,25 @@ from collections.abc import Iterator from pathlib import Path import pickle +import re import sqlite3 from dimos.memory.sensor.base import SensorStore, T +# 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(SensorStore[T]): """SQLite backend for sensor data. Good for indexed queries and single-file storage. @@ -42,10 +57,10 @@ def __init__(self, db_path: str | Path, table: str = "sensor_data") -> None: """ Args: db_path: Path to SQLite database file, or ":memory:" for in-memory. - table: Table name for this sensor's data. + table: Table name for this sensor's data (alphanumeric/underscore only). """ self._db_path = str(db_path) - self._table = table + self._table = _validate_identifier(table) self._conn: sqlite3.Connection | None = None def _get_conn(self) -> sqlite3.Connection: diff --git a/dimos/msgs/geometry_msgs/Transform.py b/dimos/msgs/geometry_msgs/Transform.py index 3a52f5a8c0..8038184f84 100644 --- a/dimos/msgs/geometry_msgs/Transform.py +++ b/dimos/msgs/geometry_msgs/Transform.py @@ -15,7 +15,10 @@ from __future__ import annotations import time -from typing import BinaryIO +from typing import TYPE_CHECKING, BinaryIO + +if TYPE_CHECKING: + from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos_lcm.geometry_msgs import ( Transform as LCMTransform, @@ -266,7 +269,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: @@ -276,10 +279,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, @@ -287,6 +290,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/protocol/tf/tf.py b/dimos/protocol/tf/tf.py index a52b4c5c3d..3dbe28694d 100644 --- a/dimos/protocol/tf/tf.py +++ b/dimos/protocol/tf/tf.py @@ -20,7 +20,7 @@ from functools import reduce from typing import TypeVar -from dimos.msgs.geometry_msgs import Transform +from dimos.msgs.geometry_msgs import PoseStamped, Transform from dimos.msgs.tf2_msgs import TFMessage from dimos.protocol.pubsub.lcmpubsub import LCM, Topic from dimos.protocol.pubsub.spec import PubSub @@ -340,7 +340,7 @@ def get_pose( 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 diff --git a/docs/concepts/transports.md b/docs/concepts/transports.md index 62279b6baf..46a1ab6d6b 100644 --- a/docs/concepts/transports.md +++ b/docs/concepts/transports.md @@ -64,7 +64,7 @@ Received 2 messages: {'temperature': 23.0} ``` -The full implementation is minimal. See [`memory.py`](/dimos/protocol/pubsub/memory.py) for the complete source. +The full implementation is minimal. See [`pubsub/memory.py`](/dimos/protocol/pubsub/memory.py) for the complete source. ## Available Transports From 239fccc0685e3d4f9b3017e17ceb4028c2be1bd7 Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Thu, 22 Jan 2026 16:26:37 +0800 Subject: [PATCH 12/39] Make SqliteStore use get_data/get_data_dir like PickleDirStore SqliteStore now accepts a name (e.g. "recordings/lidar") that gets resolved via get_data_dir to data/recordings/lidar.db. Still supports absolute paths and :memory: for backward compatibility. --- dimos/memory/sensor/sqlite.py | 59 ++++++++++++++++++++++++++++------- 1 file changed, 47 insertions(+), 12 deletions(-) diff --git a/dimos/memory/sensor/sqlite.py b/dimos/memory/sensor/sqlite.py index a2376b89af..8bc70da5e8 100644 --- a/dimos/memory/sensor/sqlite.py +++ b/dimos/memory/sensor/sqlite.py @@ -20,6 +20,7 @@ import sqlite3 from dimos.memory.sensor.base import SensorStore, T +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_]*$") @@ -42,35 +43,69 @@ class SqliteStore(SensorStore[T]): Data is stored as pickled BLOBs with timestamp as indexed column. Usage: - # File-based (persistent) - store = SqliteStore("sensors.db", table="lidar") + # Named store (uses data/ directory, auto-downloads from LFS if needed) + store = SqliteStore("recordings/lidar") # -> data/recordings/lidar.db store.save(data, timestamp) + # Absolute path + store = SqliteStore("/path/to/sensors.db") + # In-memory (for testing) store = SqliteStore(":memory:") - # Query by timestamp - data = store.find_closest(seek=10.0) + # Multiple tables in one DB + store = SqliteStore("recordings/sensors", table="lidar") """ - def __init__(self, db_path: str | Path, table: str = "sensor_data") -> None: + def __init__(self, name: str, table: str = "sensor_data") -> None: """ Args: - db_path: Path to SQLite database file, or ":memory:" for in-memory. + 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._db_path = str(db_path) + self._name = 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: - # Create parent directory if needed (for file-based DBs) - if self._db_path != ":memory:": - Path(self._db_path).parent.mkdir(parents=True, exist_ok=True) - - self._conn = sqlite3.connect(self._db_path, check_same_thread=False) + 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 From 368ed9571aad74bc97fba15425447e7195bef75f Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Thu, 22 Jan 2026 17:01:31 +0800 Subject: [PATCH 13/39] Require T to be Timestamped subclass in SensorStore - Add TypeVar bound: T = TypeVar("T", bound=Timestamped) - Simplify save() to always use data.ts (no more optional timestamp) - Update tests to use SampleData(Timestamped) instead of strings - SqliteStore accepts str | Path for backward compatibility --- dimos/memory/sensor/base.py | 15 +++--- dimos/memory/sensor/sqlite.py | 4 +- dimos/memory/sensor/test_base.py | 78 +++++++++++++++++++++----------- 3 files changed, 60 insertions(+), 37 deletions(-) diff --git a/dimos/memory/sensor/base.py b/dimos/memory/sensor/base.py index 8a2948d966..9b411fc165 100644 --- a/dimos/memory/sensor/base.py +++ b/dimos/memory/sensor/base.py @@ -25,7 +25,9 @@ from reactivex.observable import Observable from reactivex.scheduler import TimeoutScheduler -T = TypeVar("T") +from dimos.types.timestamped import Timestamped + +T = TypeVar("T", bound=Timestamped) class SensorStore(Generic[T], ABC): @@ -61,14 +63,9 @@ def _find_closest_timestamp( """Find closest timestamp. Backend can optimize (binary search, db index, etc.).""" ... - def save(self, data: T, timestamp: float | None = None) -> None: - """Save data. Uses data.ts if available, otherwise timestamp arg, otherwise now.""" - if timestamp is None: - if hasattr(data, "ts"): - timestamp = data.ts # type: ignore[union-attr] - else: - timestamp = time.time() - self._save(timestamp, data) + def save(self, data: T) -> None: + """Save timestamped data using its .ts attribute.""" + self._save(data.ts, data) def load(self, timestamp: float) -> T | None: """Load data at exact timestamp.""" diff --git a/dimos/memory/sensor/sqlite.py b/dimos/memory/sensor/sqlite.py index 8bc70da5e8..971964b860 100644 --- a/dimos/memory/sensor/sqlite.py +++ b/dimos/memory/sensor/sqlite.py @@ -57,14 +57,14 @@ class SqliteStore(SensorStore[T]): store = SqliteStore("recordings/sensors", table="lidar") """ - def __init__(self, name: str, table: str = "sensor_data") -> None: + 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 = name + self._name = str(name) self._table = _validate_identifier(table) self._db_path: str | None = None self._conn: sqlite3.Connection | None = None diff --git a/dimos/memory/sensor/test_base.py b/dimos/memory/sensor/test_base.py index 675701fbbd..c4a0b7a91e 100644 --- a/dimos/memory/sensor/test_base.py +++ b/dimos/memory/sensor/test_base.py @@ -13,6 +13,7 @@ # limitations under the License. """Tests for SensorStore implementations.""" +from dataclasses import dataclass from pathlib import Path import tempfile import uuid @@ -22,6 +23,23 @@ from dimos.memory.sensor.base import InMemoryStore, SensorStore from dimos.memory.sensor.pickledir import PickleDirStore from dimos.memory.sensor.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 @@ -31,16 +49,16 @@ def temp_dir(): yield tmpdir -def make_in_memory_store() -> SensorStore[str]: - return InMemoryStore[str]() +def make_in_memory_store() -> SensorStore[SampleData]: + return InMemoryStore[SampleData]() -def make_pickle_dir_store(tmpdir: str) -> SensorStore[str]: - return PickleDirStore[str](tmpdir) +def make_pickle_dir_store(tmpdir: str) -> SensorStore[SampleData]: + return PickleDirStore[SampleData](tmpdir) -def make_sqlite_store(tmpdir: str) -> SensorStore[str]: - return SqliteStore[str](Path(tmpdir) / "test.db") +def make_sqlite_store(tmpdir: str) -> SensorStore[SampleData]: + return SqliteStore[SampleData](Path(tmpdir) / "test.db") # Base test data (always available) @@ -62,11 +80,11 @@ def make_sqlite_store(tmpdir: str) -> SensorStore[str]: _test_conn = psycopg2.connect(dbname="dimensional") _test_conn.close() - def make_postgres_store(_tmpdir: str) -> SensorStore[str]: + def make_postgres_store(_tmpdir: str) -> SensorStore[SampleData]: """Create PostgresStore with unique table name.""" table = f"test_{uuid.uuid4().hex[:8]}" _postgres_tables.append(table) - store = PostgresStore[str](table) + store = PostgresStore[SampleData](table) store.start() return store @@ -98,18 +116,18 @@ class TestSensorStore: def test_save_and_load(self, store_factory, store_name, temp_dir): store = store_factory(temp_dir) - store.save("data_at_1", 1.0) - store.save("data_at_2", 2.0) + store.save(SampleData("data_at_1", 1.0)) + store.save(SampleData("data_at_2", 2.0)) - assert store.load(1.0) == "data_at_1" - assert store.load(2.0) == "data_at_2" + 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("a", 1.0) - store.save("b", 2.0) - store.save("c", 3.0) + store.save(SampleData("a", 1.0)) + store.save(SampleData("b", 2.0)) + store.save(SampleData("c", 3.0)) # Exact match assert store._find_closest_timestamp(2.0) == 2.0 @@ -126,32 +144,40 @@ def test_find_closest_timestamp(self, store_factory, store_name, temp_dir): def test_iter_items(self, store_factory, store_name, temp_dir): store = store_factory(temp_dir) - store.save("a", 1.0) - store.save("c", 3.0) - store.save("b", 2.0) + store.save(SampleData("a", 1.0)) + store.save(SampleData("c", 3.0)) + store.save(SampleData("b", 2.0)) # Should iterate in timestamp order items = list(store._iter_items()) - assert items == [(1.0, "a"), (2.0, "b"), (3.0, "c")] + 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("a", 1.0) - store.save("b", 2.0) - store.save("c", 3.0) - store.save("d", 4.0) + store.save(SampleData("a", 1.0)) + store.save(SampleData("b", 2.0)) + store.save(SampleData("c", 3.0)) + store.save(SampleData("d", 4.0)) # Start only items = list(store._iter_items(start=2.0)) - assert items == [(2.0, "b"), (3.0, "c"), (4.0, "d")] + 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, "a"), (2.0, "b")] + 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, "b"), (3.0, "c")] + 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) From dbf4752e23425588f5aa65342c2d29db4bcfeca5 Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Thu, 22 Jan 2026 17:21:40 +0800 Subject: [PATCH 14/39] consolidating old replay and new sensor store --- dimos/memory/sensor/base.py | 107 +++++++++++------- dimos/memory/sensor/pickledir.py | 12 +- dimos/memory/sensor/postgres.py | 4 +- dimos/memory/sensor/sqlite.py | 2 +- dimos/memory/sensor/test_base.py | 186 +++++++++++++++++++++++++++++++ dimos/utils/testing/replay.py | 8 -- 6 files changed, 262 insertions(+), 57 deletions(-) diff --git a/dimos/memory/sensor/base.py b/dimos/memory/sensor/base.py index 9b411fc165..ca541d65cb 100644 --- a/dimos/memory/sensor/base.py +++ b/dimos/memory/sensor/base.py @@ -16,11 +16,11 @@ from abc import ABC, abstractmethod import bisect from collections.abc import Iterator -import threading import time from typing import Generic, TypeVar import reactivex as rx +from reactivex import operators as ops from reactivex.disposable import CompositeDisposable, Disposable from reactivex.observable import Observable from reactivex.scheduler import TimeoutScheduler @@ -63,9 +63,31 @@ def _find_closest_timestamp( """Find closest timestamp. Backend can optimize (binary search, db index, etc.).""" ... - def save(self, data: T) -> None: - """Save timestamped data using its .ts attribute.""" - self._save(data.ts, data) + def save(self, *data: T) -> None: + """Save one or more timestamped data items using their .ts attribute.""" + for item in data: + self._save(item.ts, item) + + def pipe_save(self, source: Observable[T]) -> Observable[T]: + """Operator for use with .pipe() - saves each item and passes through. + + Usage: + observable.pipe(store.pipe_save).subscribe(...) + """ + + def _save_and_return(data: T) -> T: + self.save(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 each item. + + Usage: + disposable = store.consume_stream(observable) + """ + return observable.subscribe(on_next=self.save) def load(self, timestamp: float) -> T | None: """Load data at exact timestamp.""" @@ -73,47 +95,46 @@ def load(self, timestamp: float) -> T | None: def find_closest( self, - timestamp: float | None = None, - seek: float | None = None, + timestamp: float, tolerance: float | None = None, ) -> T | None: - """Find data closest to timestamp (absolute) or seek (relative to start).""" - if timestamp is None and seek is None: - raise ValueError("Must provide either timestamp or seek") - - if seek is not None: - first = self.first_timestamp() - if first is None: - return None - timestamp = first + seek - - assert timestamp is not 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 iterate(self, loop: bool = False) -> Iterator[tuple[float, T]]: - """Iterate over (timestamp, data) pairs.""" - while True: - yield from self._iter_items() - if not loop: - break + 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_ts( + def iterate( self, seek: float | None = None, duration: float | None = None, from_timestamp: float | None = None, loop: bool = False, - ) -> Iterator[tuple[float, T]]: - """Iterate with seek/duration options.""" + ) -> Iterator[T]: + """Iterate over data items with optional seek/duration.""" first = self.first_timestamp() if first is None: return @@ -133,19 +154,29 @@ def iterate_ts( end = start_ts + duration while True: - yield from self._iter_items(start=start, end=end) + for _, data in self._iter_items(start=start, end=end): + yield data if not loop: break - def iterate_realtime(self, speed: float = 1.0, **kwargs: float | bool | None) -> Iterator[T]: + 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_ts(**kwargs): # type: ignore[arg-type] + for data in self.iterate( + seek=seek, duration=duration, from_timestamp=from_timestamp, loop=loop + ): if prev_ts is not None: - delay = (ts - prev_ts) / speed + delay = (data.ts - prev_ts) / speed if delay > 0: time.sleep(delay) - prev_ts = ts + prev_ts = data.ts yield data def stream( @@ -169,39 +200,37 @@ def subscribe( disp = CompositeDisposable() is_disposed = False - iterator = self.iterate_ts( + iterator = self.iterate( seek=seek, duration=duration, from_timestamp=from_timestamp, loop=loop ) # Get first message try: - first_ts, first_data = next(iterator) + first_data = next(iterator) except StopIteration: observer.on_completed() return Disposable() # Establish timing reference (absolute time prevents drift) start_local_time = time.time() - start_replay_time = first_ts + start_replay_time = first_data.ts # Emit first sample immediately observer.on_next(first_data) # Pre-load next message try: - next_message: tuple[float, T] | None = next(iterator) + next_message: T | None = next(iterator) except StopIteration: observer.on_completed() return disp - def schedule_emission(message: tuple[float, T]) -> None: + def schedule_emission(data: T) -> None: 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) @@ -209,7 +238,7 @@ def schedule_emission(message: tuple[float, T]) -> None: next_message = None # Calculate absolute emission time - target_time = start_local_time + (ts - start_replay_time) / speed + target_time = start_local_time + (data.ts - start_replay_time) / speed delay = max(0.0, target_time - time.time()) def emit( diff --git a/dimos/memory/sensor/pickledir.py b/dimos/memory/sensor/pickledir.py index 0f3e8388d4..d6b1ca282f 100644 --- a/dimos/memory/sensor/pickledir.py +++ b/dimos/memory/sensor/pickledir.py @@ -15,7 +15,6 @@ import bisect from collections.abc import Iterator -from functools import lru_cache import glob import os from pathlib import Path @@ -26,22 +25,22 @@ class PickleDirStore(SensorStore[T]): - """Pickle directory backend. Compatible with TimedSensorStorage/TimedSensorReplay format. + """Pickle directory backend. Files named by timestamp. Directory structure: {name}/ - 000.pickle # (timestamp, data) - 001.pickle + 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) + data = store.find_closest_seek(10.0) # Create new recording (directory created on first save) store = PickleDirStore("my_recording/images") - store.store(timestamp, image) + store.save(image) # uses image.ts for timestamp """ def __init__(self, name: str) -> None: @@ -149,7 +148,6 @@ def _get_timestamps(self) -> list[float]: self._timestamps = timestamps return timestamps - @lru_cache(maxsize=128) def _load_file(self, filepath: Path) -> T | None: """Load data from a pickle file (LRU cached).""" try: diff --git a/dimos/memory/sensor/postgres.py b/dimos/memory/sensor/postgres.py index 4b33c25158..6ee710d38f 100644 --- a/dimos/memory/sensor/postgres.py +++ b/dimos/memory/sensor/postgres.py @@ -50,8 +50,8 @@ class PostgresStore(SensorStore[T], Resource): store.start() # open connection # Use store - store.save(data, timestamp) - data = store.find_closest(seek=10.0) + store.save(data) # uses data.ts for timestamp + data = store.find_closest_seek(10.0) # Cleanup store.stop() # close connection diff --git a/dimos/memory/sensor/sqlite.py b/dimos/memory/sensor/sqlite.py index 971964b860..874641c4f5 100644 --- a/dimos/memory/sensor/sqlite.py +++ b/dimos/memory/sensor/sqlite.py @@ -45,7 +45,7 @@ class SqliteStore(SensorStore[T]): Usage: # Named store (uses data/ directory, auto-downloads from LFS if needed) store = SqliteStore("recordings/lidar") # -> data/recordings/lidar.db - store.save(data, timestamp) + store.save(data) # uses data.ts for timestamp # Absolute path store = SqliteStore("/path/to/sensors.db") diff --git a/dimos/memory/sensor/test_base.py b/dimos/memory/sensor/test_base.py index c4a0b7a91e..a21800b303 100644 --- a/dimos/memory/sensor/test_base.py +++ b/dimos/memory/sensor/test_base.py @@ -185,3 +185,189 @@ def test_empty_store(self, store_factory, store_name, 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 + store.save(SampleData("b", 2.0)) + store.save(SampleData("a", 1.0)) + store.save(SampleData("c", 3.0)) + + # Should return first by timestamp, not insertion order + 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)) + store.save(SampleData("b", 2.0)) + store.save(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)) + store.save(SampleData("b", 11.0)) + store.save(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)) + store.save(SampleData("c", 3.0)) + store.save(SampleData("b", 2.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)) + store.save(SampleData("b", 11.0)) + store.save(SampleData("c", 12.0)) + store.save(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_stream_basic(self, store_factory, store_name, temp_dir): + store = store_factory(temp_dir) + store.save(SampleData("a", 1.0)) + store.save(SampleData("b", 2.0)) + store.save(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), + ] diff --git a/dimos/utils/testing/replay.py b/dimos/utils/testing/replay.py index be8d394f76..0b28b3ec08 100644 --- a/dimos/utils/testing/replay.py +++ b/dimos/utils/testing/replay.py @@ -256,14 +256,6 @@ def first_timestamp(self) -> float | 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. From 2c7f5713a5dfa2d960a1cbcd82906e25ad31015f Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Thu, 22 Jan 2026 17:34:48 +0800 Subject: [PATCH 15/39] Add LegacyPickleStore backend for TimedSensorReplay compatibility --- dimos/memory/sensor/legacy.py | 175 +++++++++++++++++++++++++++++++ dimos/memory/sensor/test_base.py | 48 ++++++++- 2 files changed, 218 insertions(+), 5 deletions(-) create mode 100644 dimos/memory/sensor/legacy.py diff --git a/dimos/memory/sensor/legacy.py b/dimos/memory/sensor/legacy.py new file mode 100644 index 0000000000..40c7d13d40 --- /dev/null +++ b/dimos/memory/sensor/legacy.py @@ -0,0 +1,175 @@ +# 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 SensorStore. + +Compatible with TimedSensorReplay/TimedSensorStorage file format. +""" + +from collections.abc import Iterator +import glob +import os +from pathlib import Path +import pickle +import re +from typing import cast + +from dimos.memory.sensor.base import SensorStore, T +from dimos.utils.data import get_data, get_data_dir + + +class LegacyPickleStore(SensorStore[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(image) # uses image.ts for timestamp + """ + + def __init__(self, name: str | Path) -> None: + """ + Args: + name: Data directory name (e.g. "unitree_go2_bigoffice/lidar") or absolute path. + """ + self._name = str(name) + self._root_dir: Path | None = None + self._counter: int = 0 + + 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 _iter_items( + self, start: float | None = None, end: float | None = None + ) -> Iterator[tuple[float, T]]: + """Lazy iteration - loads one file at a time.""" + for filepath in self._iter_files(): + try: + with open(filepath, "rb") as f: + ts, data = pickle.load(f) + ts = float(ts) + except Exception: + continue + + if start is not None and ts < start: + continue + if end is not None and ts >= end: + break + 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 diff --git a/dimos/memory/sensor/test_base.py b/dimos/memory/sensor/test_base.py index a21800b303..2506f4edcf 100644 --- a/dimos/memory/sensor/test_base.py +++ b/dimos/memory/sensor/test_base.py @@ -21,11 +21,44 @@ import pytest from dimos.memory.sensor.base import InMemoryStore, SensorStore +from dimos.memory.sensor.legacy import LegacyPickleStore from dimos.memory.sensor.pickledir import PickleDirStore from dimos.memory.sensor.sqlite import SqliteStore from dimos.types.timestamped import Timestamped +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 + + @dataclass class SampleData(Timestamped): """Simple timestamped data for testing.""" @@ -61,11 +94,16 @@ def make_sqlite_store(tmpdir: str) -> SensorStore[SampleData]: return SqliteStore[SampleData](Path(tmpdir) / "test.db") +def make_legacy_pickle_store(tmpdir: str) -> SensorStore[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 @@ -145,8 +183,8 @@ def test_find_closest_timestamp(self, store_factory, store_name, temp_dir): def test_iter_items(self, store_factory, store_name, temp_dir): store = store_factory(temp_dir) store.save(SampleData("a", 1.0)) - store.save(SampleData("c", 3.0)) store.save(SampleData("b", 2.0)) + store.save(SampleData("c", 3.0)) # Should iterate in timestamp order items = list(store._iter_items()) @@ -193,12 +231,12 @@ def test_first_and_first_timestamp(self, store_factory, store_name, temp_dir): assert store.first() is None assert store.first_timestamp() is None - # Add data - store.save(SampleData("b", 2.0)) + # Add data (in chronological order) store.save(SampleData("a", 1.0)) + store.save(SampleData("b", 2.0)) store.save(SampleData("c", 3.0)) - # Should return first by timestamp, not insertion order + # Should return first by timestamp assert store.first_timestamp() == 1.0 assert store.first() == SampleData("a", 1.0) @@ -246,8 +284,8 @@ def test_find_closest_seek(self, store_factory, store_name, temp_dir): def test_iterate(self, store_factory, store_name, temp_dir): store = store_factory(temp_dir) store.save(SampleData("a", 1.0)) - store.save(SampleData("c", 3.0)) store.save(SampleData("b", 2.0)) + store.save(SampleData("c", 3.0)) # Should iterate in timestamp order, returning data only (not tuples) items = list(store.iterate()) From 3b64dc6bf8aaeaf5be05bec7f7fe88df0ac89819 Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Thu, 22 Jan 2026 19:26:43 +0800 Subject: [PATCH 16/39] legacy sensor store implemented, replay shim implemented --- dimos/memory/sensor/legacy.py | 12 +- dimos/utils/testing/replay.py | 393 +------------------------- dimos/utils/testing/replay_legacy.py | 401 +++++++++++++++++++++++++++ dimos/utils/testing/test_replay.py | 34 +-- 4 files changed, 426 insertions(+), 414 deletions(-) create mode 100644 dimos/utils/testing/replay_legacy.py diff --git a/dimos/memory/sensor/legacy.py b/dimos/memory/sensor/legacy.py index 40c7d13d40..6b01c5bc7b 100644 --- a/dimos/memory/sensor/legacy.py +++ b/dimos/memory/sensor/legacy.py @@ -16,13 +16,13 @@ Compatible with TimedSensorReplay/TimedSensorStorage file format. """ -from collections.abc import Iterator +from collections.abc import Callable, Iterator import glob import os from pathlib import Path import pickle import re -from typing import cast +from typing import Any, cast from dimos.memory.sensor.base import SensorStore, T from dimos.utils.data import get_data, get_data_dir @@ -50,14 +50,17 @@ class LegacyPickleStore(SensorStore[T]): store.save(image) # uses image.ts for timestamp """ - def __init__(self, name: str | Path) -> None: + 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.""" @@ -147,6 +150,9 @@ def _iter_items( 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( diff --git a/dimos/utils/testing/replay.py b/dimos/utils/testing/replay.py index 0b28b3ec08..47d9fcee57 100644 --- a/dimos/utils/testing/replay.py +++ b/dimos/utils/testing/replay.py @@ -11,391 +11,14 @@ # 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 +"""Shim for TimedSensorReplay/TimedSensorStorage. -from reactivex import ( - from_iterable, - interval, - operators as ops, -) -from reactivex.observable import Observable -from reactivex.scheduler import TimeoutScheduler +For the original implementation, see replay_legacy.py. +""" -from dimos.utils.data import get_data, get_data_dir +from dimos.memory.sensor.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_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/replay_legacy.py b/dimos/utils/testing/replay_legacy.py new file mode 100644 index 0000000000..0b28b3ec08 --- /dev/null +++ b/dimos/utils/testing/replay_legacy.py @@ -0,0 +1,401 @@ +# 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. +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 + +from dimos.utils.data import get_data, get_data_dir + +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_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) diff --git a/dimos/utils/testing/test_replay.py b/dimos/utils/testing/test_replay.py index 640fe92979..c6a319513a 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(): @@ -69,7 +51,7 @@ def test_timed_sensor_replay() -> None: def test_iterate_ts_no_seek() -> None: """Test iterate_ts without seek (start_timestamp=None)""" - odom_store = replay.TimedSensorReplay("unitree_office_walk/odom", autocast=Odometry.from_msg) + odom_store = replay.TimedSensorReplay("unitree_office_walk/odom") # Test without seek ts_msgs = [] @@ -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 = [] @@ -225,7 +207,7 @@ def test_first_methods() -> None: assert abs(first_msg.ts - first_from_iterate.ts) < 1.0 # Within 1 second tolerance # Test TimedSensorReplay.first_timestamp() - odom_store = replay.TimedSensorReplay("unitree_office_walk/odom", autocast=Odometry.from_msg) + odom_store = replay.TimedSensorReplay("unitree_office_walk/odom") first_ts = odom_store.first_timestamp() assert first_ts is not None assert isinstance(first_ts, float) @@ -242,7 +224,7 @@ def test_first_methods() -> None: def test_find_closest() -> None: """Test find_closest method in TimedSensorReplay""" - odom_store = replay.TimedSensorReplay("unitree_office_walk/odom", autocast=Odometry.from_msg) + odom_store = replay.TimedSensorReplay("unitree_office_walk/odom") # Get some reference timestamps timestamps = [] From 6430d70f513fb5288a908924af04ccb96cbb9025 Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Thu, 22 Jan 2026 20:08:59 +0800 Subject: [PATCH 17/39] replaced legacy replay with new system --- dimos/core/__init__.py | 2 +- dimos/core/testing.py | 1 - .../pointclouds/test_occupancy_speed.py | 3 +- dimos/memory/sensor/legacy.py | 195 +++++++++++++++++- dimos/memory/test_embedding.py | 3 + .../contact_graspnet_pytorch/inference.py | 8 +- dimos/msgs/sensor_msgs/Image.py | 2 +- .../sensor_msgs/image_impls/AbstractImage.py | 2 +- .../msgs/sensor_msgs/image_impls/CudaImage.py | 4 +- dimos/perception/common/utils.py | 2 +- dimos/robot/unitree/connection/go2.py | 6 +- .../unitree_webrtc/type/test_odometry.py | 29 --- dimos/utils/testing/moment.py | 3 +- dimos/utils/testing/test_replay.py | 6 +- 14 files changed, 214 insertions(+), 52 deletions(-) diff --git a/dimos/core/__init__.py b/dimos/core/__init__.py index b56fe74f4f..a6dde6c0e1 100644 --- a/dimos/core/__init__.py +++ b/dimos/core/__init__.py @@ -68,7 +68,7 @@ def teardown(self, worker) -> None: # type: ignore[no-untyped-def] import sys if "cupy" in sys.modules: - import cupy as cp # type: ignore[import-not-found] + import cupy as cp # type: ignore[import-not-found,import-untyped] # Clear memory pools mempool = cp.get_default_memory_pool() 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/mapping/pointclouds/test_occupancy_speed.py b/dimos/mapping/pointclouds/test_occupancy_speed.py index e296494a45..2def839dd5 100644 --- a/dimos/mapping/pointclouds/test_occupancy_speed.py +++ b/dimos/mapping/pointclouds/test_occupancy_speed.py @@ -28,8 +28,7 @@ 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" diff --git a/dimos/memory/sensor/legacy.py b/dimos/memory/sensor/legacy.py index 6b01c5bc7b..8ba4540f0a 100644 --- a/dimos/memory/sensor/legacy.py +++ b/dimos/memory/sensor/legacy.py @@ -22,8 +22,14 @@ 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.sensor.base import SensorStore, T from dimos.utils.data import get_data, get_data_dir @@ -48,6 +54,12 @@ class LegacyPickleStore(SensorStore[T]): # Create new recording (directory created on first save) store = LegacyPickleStore("my_recording/images") store.save(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: @@ -137,12 +149,24 @@ def _load(self, timestamp: float) -> T | None: def _iter_items( self, start: float | None = None, end: float | None = None ) -> Iterator[tuple[float, T]]: - """Lazy iteration - loads one file at a time.""" - for filepath in self._iter_files(): + """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: - ts, data = pickle.load(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 @@ -179,3 +203,168 @@ def _find_closest_timestamp( return None return closest_ts + + # === 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/test_embedding.py b/dimos/memory/test_embedding.py index ffd9e21a4f..b7e7fbb294 100644 --- a/dimos/memory/test_embedding.py +++ b/dimos/memory/test_embedding.py @@ -12,6 +12,8 @@ # 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 @@ -20,6 +22,7 @@ dir_name = "unitree_go2_bigoffice" +@pytest.mark.skip def test_embed_frame() -> None: """Test embedding a single frame.""" # Load a frame from recorded data diff --git a/dimos/models/manipulation/contact_graspnet_pytorch/inference.py b/dimos/models/manipulation/contact_graspnet_pytorch/inference.py index 0769fc150d..2586b15e66 100644 --- a/dimos/models/manipulation/contact_graspnet_pytorch/inference.py +++ b/dimos/models/manipulation/contact_graspnet_pytorch/inference.py @@ -2,12 +2,12 @@ import glob import os -from contact_graspnet_pytorch import config_utils # type: ignore[import-not-found] -from contact_graspnet_pytorch.checkpoints import CheckpointIO # type: ignore[import-not-found] -from contact_graspnet_pytorch.contact_grasp_estimator import ( # type: ignore[import-not-found] +from contact_graspnet_pytorch import config_utils # type: ignore[import-not-found,import-untyped] +from contact_graspnet_pytorch.checkpoints import CheckpointIO # type: ignore[import-not-found,import-untyped] +from contact_graspnet_pytorch.contact_grasp_estimator import ( # type: ignore[import-not-found,import-untyped] GraspEstimator, ) -from contact_graspnet_pytorch.data import ( # type: ignore[import-not-found] +from contact_graspnet_pytorch.data import ( # type: ignore[import-not-found,import-untyped] load_available_input_data, ) import numpy as np diff --git a/dimos/msgs/sensor_msgs/Image.py b/dimos/msgs/sensor_msgs/Image.py index 2bb99427bf..7fa38e0a0b 100644 --- a/dimos/msgs/sensor_msgs/Image.py +++ b/dimos/msgs/sensor_msgs/Image.py @@ -48,7 +48,7 @@ ) try: - import cupy as cp # type: ignore[import-not-found] + import cupy as cp # type: ignore[import-not-found,import-untyped] except Exception: cp = None diff --git a/dimos/msgs/sensor_msgs/image_impls/AbstractImage.py b/dimos/msgs/sensor_msgs/image_impls/AbstractImage.py index f5d92a3bc6..b20381bfef 100644 --- a/dimos/msgs/sensor_msgs/image_impls/AbstractImage.py +++ b/dimos/msgs/sensor_msgs/image_impls/AbstractImage.py @@ -25,7 +25,7 @@ import rerun as rr try: - import cupy as cp # type: ignore[import-not-found] + import cupy as cp # type: ignore[import-not-found,import-untyped] HAS_CUDA = True except Exception: # pragma: no cover - optional dependency diff --git a/dimos/msgs/sensor_msgs/image_impls/CudaImage.py b/dimos/msgs/sensor_msgs/image_impls/CudaImage.py index 8230daae29..a1d0dffa70 100644 --- a/dimos/msgs/sensor_msgs/image_impls/CudaImage.py +++ b/dimos/msgs/sensor_msgs/image_impls/CudaImage.py @@ -30,8 +30,8 @@ ) try: - import cupy as cp # type: ignore[import-not-found] - from cupyx.scipy import ( # type: ignore[import-not-found] + import cupy as cp # type: ignore[import-not-found,import-untyped] + from cupyx.scipy import ( # type: ignore[import-not-found,import-untyped] ndimage as cndimage, signal as csignal, ) diff --git a/dimos/perception/common/utils.py b/dimos/perception/common/utils.py index 1144234d71..3ca6a14416 100644 --- a/dimos/perception/common/utils.py +++ b/dimos/perception/common/utils.py @@ -36,7 +36,7 @@ # Optional CuPy support try: # pragma: no cover - optional dependency - import cupy as cp # type: ignore[import-not-found] + import cupy as cp # type: ignore[import-not-found,import-untyped] _HAS_CUDA = True except Exception: # pragma: no cover - optional dependency diff --git a/dimos/robot/unitree/connection/go2.py b/dimos/robot/unitree/connection/go2.py index 96a54117c8..187cfe638b 100644 --- a/dimos/robot/unitree/connection/go2.py +++ b/dimos/robot/unitree/connection/go2.py @@ -192,13 +192,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..13c10fc820 100644 --- a/dimos/robot/unitree_webrtc/type/test_odometry.py +++ b/dimos/robot/unitree_webrtc/type/test_odometry.py @@ -38,19 +38,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 +50,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/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/test_replay.py b/dimos/utils/testing/test_replay.py index c6a319513a..8336c6a2b9 100644 --- a/dimos/utils/testing/test_replay.py +++ b/dimos/utils/testing/test_replay.py @@ -51,7 +51,7 @@ def test_timed_sensor_replay() -> None: def test_iterate_ts_no_seek() -> None: """Test iterate_ts without seek (start_timestamp=None)""" - odom_store = replay.TimedSensorReplay("unitree_office_walk/odom") + odom_store = replay.TimedSensorReplay("unitree_office_walk/odom", autocast=Odometry.from_msg) # Test without seek ts_msgs = [] @@ -207,7 +207,7 @@ def test_first_methods() -> None: assert abs(first_msg.ts - first_from_iterate.ts) < 1.0 # Within 1 second tolerance # Test TimedSensorReplay.first_timestamp() - odom_store = replay.TimedSensorReplay("unitree_office_walk/odom") + odom_store = replay.TimedSensorReplay("unitree_office_walk/odom", autocast=Odometry.from_msg) first_ts = odom_store.first_timestamp() assert first_ts is not None assert isinstance(first_ts, float) @@ -224,7 +224,7 @@ def test_first_methods() -> None: def test_find_closest() -> None: """Test find_closest method in TimedSensorReplay""" - odom_store = replay.TimedSensorReplay("unitree_office_walk/odom") + odom_store = replay.TimedSensorReplay("unitree_office_walk/odom", autocast=Odometry.from_msg) # Get some reference timestamps timestamps = [] From bb7be311a46d431ae3d3eacc49acfb2161e5bff4 Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Sun, 25 Jan 2026 16:47:43 +0800 Subject: [PATCH 18/39] Remove import-untyped type ignores for consistency with dev --- dimos/core/__init__.py | 2 +- .../manipulation/contact_graspnet_pytorch/inference.py | 6 +++--- dimos/msgs/sensor_msgs/Image.py | 2 +- dimos/msgs/sensor_msgs/image_impls/AbstractImage.py | 2 +- dimos/msgs/sensor_msgs/image_impls/CudaImage.py | 4 ++-- dimos/perception/common/utils.py | 2 +- 6 files changed, 9 insertions(+), 9 deletions(-) diff --git a/dimos/core/__init__.py b/dimos/core/__init__.py index a6dde6c0e1..b56fe74f4f 100644 --- a/dimos/core/__init__.py +++ b/dimos/core/__init__.py @@ -68,7 +68,7 @@ def teardown(self, worker) -> None: # type: ignore[no-untyped-def] import sys if "cupy" in sys.modules: - import cupy as cp # type: ignore[import-not-found,import-untyped] + import cupy as cp # type: ignore[import-not-found] # Clear memory pools mempool = cp.get_default_memory_pool() diff --git a/dimos/models/manipulation/contact_graspnet_pytorch/inference.py b/dimos/models/manipulation/contact_graspnet_pytorch/inference.py index 07dc48f258..76bb377869 100644 --- a/dimos/models/manipulation/contact_graspnet_pytorch/inference.py +++ b/dimos/models/manipulation/contact_graspnet_pytorch/inference.py @@ -2,11 +2,11 @@ import glob import os -from contact_graspnet_pytorch import config_utils # type: ignore[import-not-found,import-untyped] -from contact_graspnet_pytorch.contact_grasp_estimator import ( # type: ignore[import-not-found,import-untyped] +from contact_graspnet_pytorch import config_utils # type: ignore[import-not-found] +from contact_graspnet_pytorch.contact_grasp_estimator import ( # type: ignore[import-not-found] GraspEstimator, ) -from contact_graspnet_pytorch.data import ( # type: ignore[import-not-found,import-untyped] +from contact_graspnet_pytorch.data import ( # type: ignore[import-not-found] load_available_input_data, ) import numpy as np diff --git a/dimos/msgs/sensor_msgs/Image.py b/dimos/msgs/sensor_msgs/Image.py index 7fa38e0a0b..2bb99427bf 100644 --- a/dimos/msgs/sensor_msgs/Image.py +++ b/dimos/msgs/sensor_msgs/Image.py @@ -48,7 +48,7 @@ ) try: - import cupy as cp # type: ignore[import-not-found,import-untyped] + import cupy as cp # type: ignore[import-not-found] except Exception: cp = None diff --git a/dimos/msgs/sensor_msgs/image_impls/AbstractImage.py b/dimos/msgs/sensor_msgs/image_impls/AbstractImage.py index c5ff75ba5b..b71c5476fc 100644 --- a/dimos/msgs/sensor_msgs/image_impls/AbstractImage.py +++ b/dimos/msgs/sensor_msgs/image_impls/AbstractImage.py @@ -25,7 +25,7 @@ import rerun as rr try: - import cupy as cp # type: ignore[import-not-found,import-untyped] + import cupy as cp # type: ignore[import-not-found] HAS_CUDA = True except Exception: # pragma: no cover - optional dependency diff --git a/dimos/msgs/sensor_msgs/image_impls/CudaImage.py b/dimos/msgs/sensor_msgs/image_impls/CudaImage.py index 05049f1e45..cdfa1bf088 100644 --- a/dimos/msgs/sensor_msgs/image_impls/CudaImage.py +++ b/dimos/msgs/sensor_msgs/image_impls/CudaImage.py @@ -30,8 +30,8 @@ ) try: - import cupy as cp # type: ignore[import-not-found,import-untyped] - from cupyx.scipy import ( # type: ignore[import-not-found,import-untyped] + import cupy as cp # type: ignore[import-not-found] + from cupyx.scipy import ( # type: ignore[import-not-found] ndimage as cndimage, signal as csignal, ) diff --git a/dimos/perception/common/utils.py b/dimos/perception/common/utils.py index 3ca6a14416..1144234d71 100644 --- a/dimos/perception/common/utils.py +++ b/dimos/perception/common/utils.py @@ -36,7 +36,7 @@ # Optional CuPy support try: # pragma: no cover - optional dependency - import cupy as cp # type: ignore[import-not-found,import-untyped] + import cupy as cp # type: ignore[import-not-found] _HAS_CUDA = True except Exception: # pragma: no cover - optional dependency From 697db1595a06ebb55813bd3ff6c8561fc4d283c0 Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Sun, 25 Jan 2026 16:57:05 +0800 Subject: [PATCH 19/39] Restore import-untyped type ignores for local mypy compatibility Required when cupy/contact_graspnet are installed locally without type stubs. --- dimos/core/__init__.py | 2 +- .../manipulation/contact_graspnet_pytorch/inference.py | 6 +++--- dimos/msgs/sensor_msgs/Image.py | 2 +- dimos/msgs/sensor_msgs/image_impls/AbstractImage.py | 4 ++-- dimos/msgs/sensor_msgs/image_impls/CudaImage.py | 4 ++-- dimos/perception/common/utils.py | 2 +- 6 files changed, 10 insertions(+), 10 deletions(-) diff --git a/dimos/core/__init__.py b/dimos/core/__init__.py index b56fe74f4f..a6dde6c0e1 100644 --- a/dimos/core/__init__.py +++ b/dimos/core/__init__.py @@ -68,7 +68,7 @@ def teardown(self, worker) -> None: # type: ignore[no-untyped-def] import sys if "cupy" in sys.modules: - import cupy as cp # type: ignore[import-not-found] + import cupy as cp # type: ignore[import-not-found,import-untyped] # Clear memory pools mempool = cp.get_default_memory_pool() diff --git a/dimos/models/manipulation/contact_graspnet_pytorch/inference.py b/dimos/models/manipulation/contact_graspnet_pytorch/inference.py index 76bb377869..07dc48f258 100644 --- a/dimos/models/manipulation/contact_graspnet_pytorch/inference.py +++ b/dimos/models/manipulation/contact_graspnet_pytorch/inference.py @@ -2,11 +2,11 @@ import glob import os -from contact_graspnet_pytorch import config_utils # type: ignore[import-not-found] -from contact_graspnet_pytorch.contact_grasp_estimator import ( # type: ignore[import-not-found] +from contact_graspnet_pytorch import config_utils # type: ignore[import-not-found,import-untyped] +from contact_graspnet_pytorch.contact_grasp_estimator import ( # type: ignore[import-not-found,import-untyped] GraspEstimator, ) -from contact_graspnet_pytorch.data import ( # type: ignore[import-not-found] +from contact_graspnet_pytorch.data import ( # type: ignore[import-not-found,import-untyped] load_available_input_data, ) import numpy as np diff --git a/dimos/msgs/sensor_msgs/Image.py b/dimos/msgs/sensor_msgs/Image.py index 2bb99427bf..7fa38e0a0b 100644 --- a/dimos/msgs/sensor_msgs/Image.py +++ b/dimos/msgs/sensor_msgs/Image.py @@ -48,7 +48,7 @@ ) try: - import cupy as cp # type: ignore[import-not-found] + import cupy as cp # type: ignore[import-not-found,import-untyped] except Exception: cp = None diff --git a/dimos/msgs/sensor_msgs/image_impls/AbstractImage.py b/dimos/msgs/sensor_msgs/image_impls/AbstractImage.py index b71c5476fc..9d40dc41ee 100644 --- a/dimos/msgs/sensor_msgs/image_impls/AbstractImage.py +++ b/dimos/msgs/sensor_msgs/image_impls/AbstractImage.py @@ -25,7 +25,7 @@ import rerun as rr try: - import cupy as cp # type: ignore[import-not-found] + import cupy as cp # type: ignore[import-not-found,import-untyped] HAS_CUDA = True except Exception: # pragma: no cover - optional dependency @@ -35,7 +35,7 @@ # NVRTC defaults to C++11; libcu++ in recent CUDA requires at least C++17. if HAS_CUDA: try: - import cupy.cuda.compiler as _cupy_compiler # type: ignore[import-not-found] + import cupy.cuda.compiler as _cupy_compiler # type: ignore[import-not-found,import-untyped] if not getattr(_cupy_compiler, "_dimos_force_cxx17", False): _orig_compile_using_nvrtc = _cupy_compiler.compile_using_nvrtc diff --git a/dimos/msgs/sensor_msgs/image_impls/CudaImage.py b/dimos/msgs/sensor_msgs/image_impls/CudaImage.py index cdfa1bf088..05049f1e45 100644 --- a/dimos/msgs/sensor_msgs/image_impls/CudaImage.py +++ b/dimos/msgs/sensor_msgs/image_impls/CudaImage.py @@ -30,8 +30,8 @@ ) try: - import cupy as cp # type: ignore[import-not-found] - from cupyx.scipy import ( # type: ignore[import-not-found] + import cupy as cp # type: ignore[import-not-found,import-untyped] + from cupyx.scipy import ( # type: ignore[import-not-found,import-untyped] ndimage as cndimage, signal as csignal, ) diff --git a/dimos/perception/common/utils.py b/dimos/perception/common/utils.py index 1144234d71..3ca6a14416 100644 --- a/dimos/perception/common/utils.py +++ b/dimos/perception/common/utils.py @@ -36,7 +36,7 @@ # Optional CuPy support try: # pragma: no cover - optional dependency - import cupy as cp # type: ignore[import-not-found] + import cupy as cp # type: ignore[import-not-found,import-untyped] _HAS_CUDA = True except Exception: # pragma: no cover - optional dependency From 0009e3028f1d425dc9515199892c92fccc00e719 Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Mon, 26 Jan 2026 00:46:45 +0800 Subject: [PATCH 20/39] record/replay on modules --- .gitignore | 2 ++ dimos/core/module.py | 12 ++++++++++++ dimos/memory/sensor/base.py | 2 -- 3 files changed, 14 insertions(+), 2 deletions(-) diff --git a/.gitignore b/.gitignore index e52d08ba32..877eef1460 100644 --- a/.gitignore +++ b/.gitignore @@ -64,3 +64,5 @@ yolo11n.pt *mobileclip* /results + +CLAUDE.MD diff --git a/dimos/core/module.py b/dimos/core/module.py index 08e428d3c7..b308df706b 100644 --- a/dimos/core/module.py +++ b/dimos/core/module.py @@ -119,6 +119,18 @@ def frame_id(self) -> str: return f"{self.config.frame_id_prefix}/{base}" return base + # called during record sessions + # sensor ingress modules are recording by default? + @rpc + def record(self) -> None: + for name, output in self.outputs.items(): + output.observable().subscribe(lambda v, n=name: print(f"RECORD {n}: {v}")) + + # called instead of start during replay sessions + @rpc + def replay(self) -> None: + raise NotImplementedError("replay() not implemented for this module") + @rpc def start(self) -> None: pass diff --git a/dimos/memory/sensor/base.py b/dimos/memory/sensor/base.py index ca541d65cb..21784a7c7f 100644 --- a/dimos/memory/sensor/base.py +++ b/dimos/memory/sensor/base.py @@ -37,8 +37,6 @@ class SensorStore(Generic[T], ABC): All iteration, streaming, and seek logic comes free from the base class. """ - # === Abstract - implement for your backend === - @abstractmethod def _save(self, timestamp: float, data: T) -> None: """Save data at timestamp.""" From cc633173f919d3110b7dbd72150663065a86c3e6 Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Wed, 11 Feb 2026 17:27:35 +0800 Subject: [PATCH 21/39] Fix mypy errors: add psycopg2 stubs and xacro type ignore - Add import-untyped to xacro type: ignore comment in mesh_utils.py - Remove unused record/replay RPC methods from ModuleBase --- dimos/core/module.py | 12 ------------ dimos/manipulation/planning/utils/mesh_utils.py | 2 +- 2 files changed, 1 insertion(+), 13 deletions(-) diff --git a/dimos/core/module.py b/dimos/core/module.py index f9db1b04ed..a1f60fed8f 100644 --- a/dimos/core/module.py +++ b/dimos/core/module.py @@ -127,18 +127,6 @@ def frame_id(self) -> str: return f"{self.config.frame_id_prefix}/{base}" return base - # called during record sessions - # sensor ingress modules are recording by default? - @rpc - def record(self) -> None: - for name, output in self.outputs.items(): - output.observable().subscribe(lambda v, n=name: print(f"RECORD {n}: {v}")) - - # called instead of start during replay sessions - @rpc - def replay(self) -> None: - raise NotImplementedError("replay() not implemented for this module") - @rpc def start(self) -> None: pass 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" From 42b9e099546f8aaebc0adeb86f649ee7fd240b37 Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Wed, 11 Feb 2026 18:39:08 +0800 Subject: [PATCH 22/39] Decouple SensorStore from Timestamped type constraint Remove the Timestamped bound from SensorStore's TypeVar, enabling storage of arbitrary data types. Timestamps are now provided explicitly via save(ts, data), with Timestamped convenience methods (save_ts, pipe_save_ts, consume_stream_ts) as opt-in helpers. iterate_realtime() and stream() now use stored timestamps instead of data.ts. --- dimos/memory/sensor/base.py | 122 +++++++++++++------ dimos/memory/sensor/legacy.py | 2 +- dimos/memory/sensor/pickledir.py | 2 +- dimos/memory/sensor/postgres.py | 8 +- dimos/memory/sensor/sqlite.py | 2 +- dimos/memory/sensor/test_base.py | 197 ++++++++++++++++++++++++------- 6 files changed, 248 insertions(+), 85 deletions(-) diff --git a/dimos/memory/sensor/base.py b/dimos/memory/sensor/base.py index 21784a7c7f..edd3852564 100644 --- a/dimos/memory/sensor/base.py +++ b/dimos/memory/sensor/base.py @@ -11,11 +11,11 @@ # 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 timestamped sensor storage and replay.""" +"""Unified sensor storage and replay.""" from abc import ABC, abstractmethod import bisect -from collections.abc import Iterator +from collections.abc import Callable, Iterator import time from typing import Generic, TypeVar @@ -27,14 +27,18 @@ from dimos.types.timestamped import Timestamped -T = TypeVar("T", bound=Timestamped) +T = TypeVar("T") class SensorStore(Generic[T], ABC): - """Unified storage + replay for timestamped sensor data. + """Unified storage + replay for sensor data. Implement 4 abstract methods for your backend (in-memory, pickle, sqlite, etc.). All iteration, streaming, and seek logic comes free from the base class. + + T can be any type — timestamps are provided explicitly. For Timestamped + subclasses, convenience methods (save_ts, pipe_save_ts, consume_stream_ts) + automatically extract the .ts attribute. """ @abstractmethod @@ -61,31 +65,67 @@ def _find_closest_timestamp( """Find closest timestamp. Backend can optimize (binary search, db index, etc.).""" ... - def save(self, *data: T) -> None: - """Save one or more timestamped data items using their .ts attribute.""" + def save(self, timestamp: float, data: T) -> None: + """Save a single data item at the given timestamp.""" + self._save(timestamp, data) + + def save_ts(self, *data: Timestamped) -> None: + """Save one or more Timestamped items using their .ts attribute.""" for item in data: - self._save(item.ts, item) + self._save(item.ts, item) # type: ignore[arg-type] + + def pipe_save(self, key: Callable[[T], float]) -> Callable[[Observable[T]], Observable[T]]: + """Operator for Observable.pipe() — saves each item using key(item) as timestamp. + + Usage: + observable.pipe(store.pipe_save(lambda x: x.ts)).subscribe(...) + """ + + def _operator(source: Observable[T]) -> Observable[T]: + def _save_and_return(data: T) -> T: + self._save(key(data), data) + return data + + return source.pipe(ops.map(_save_and_return)) - def pipe_save(self, source: Observable[T]) -> Observable[T]: - """Operator for use with .pipe() - saves each item and passes through. + return _operator + + def pipe_save_ts(self, source: Observable[T]) -> Observable[T]: + """Operator for Observable.pipe() — saves Timestamped items using .ts. Usage: - observable.pipe(store.pipe_save).subscribe(...) + observable.pipe(store.pipe_save_ts).subscribe(...) """ def _save_and_return(data: T) -> T: - self.save(data) + ts_data: Timestamped = data # type: ignore[assignment] + self._save(ts_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 each item. + def consume_stream( + self, observable: Observable[T], key: Callable[[T], float] + ) -> rx.abc.DisposableBase: + """Subscribe to an observable and save each item using key(item) as timestamp. + + Usage: + disposable = store.consume_stream(observable, key=lambda x: x.ts) + """ + return observable.subscribe(on_next=lambda data: self._save(key(data), data)) + + def consume_stream_ts(self, observable: Observable[T]) -> rx.abc.DisposableBase: + """Subscribe to an observable and save Timestamped items using .ts. Usage: - disposable = store.consume_stream(observable) + disposable = store.consume_stream_ts(observable) """ - return observable.subscribe(on_next=self.save) + + def _save_item(data: T) -> None: + ts_data: Timestamped = data # type: ignore[assignment] + self._save(ts_data.ts, data) + + return observable.subscribe(on_next=_save_item) def load(self, timestamp: float) -> T | None: """Load data at exact timestamp.""" @@ -125,19 +165,18 @@ def first(self) -> T | None: return data return None - def iterate( + def iterate_items( 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.""" + ) -> Iterator[tuple[float, T]]: + """Iterate over (timestamp, data) tuples with optional seek/duration.""" first = self.first_timestamp() if first is None: return - # Calculate start timestamp if from_timestamp is not None: start = from_timestamp elif seek is not None: @@ -145,18 +184,29 @@ def iterate( else: start = None - # Calculate end timestamp end = None if duration is not None: start_ts = start if start is not None else first end = start_ts + duration while True: - for _, data in self._iter_items(start=start, end=end): - yield data + 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, @@ -167,14 +217,14 @@ def iterate_realtime( ) -> Iterator[T]: """Iterate data, sleeping to match original timing.""" prev_ts: float | None = None - for data in self.iterate( + for ts, data in self.iterate_items( seek=seek, duration=duration, from_timestamp=from_timestamp, loop=loop ): if prev_ts is not None: - delay = (data.ts - prev_ts) / speed + delay = (ts - prev_ts) / speed if delay > 0: time.sleep(delay) - prev_ts = data.ts + prev_ts = ts yield data def stream( @@ -198,45 +248,41 @@ def subscribe( disp = CompositeDisposable() is_disposed = False - iterator = self.iterate( + iterator = self.iterate_items( seek=seek, duration=duration, from_timestamp=from_timestamp, loop=loop ) - # Get first message try: - first_data = next(iterator) + first_ts, first_data = next(iterator) except StopIteration: observer.on_completed() return Disposable() - # Establish timing reference (absolute time prevents drift) start_local_time = time.time() - start_replay_time = first_data.ts + start_replay_time = first_ts - # Emit first sample immediately observer.on_next(first_data) - # Pre-load next message try: - next_message: T | None = next(iterator) + next_message: tuple[float, T] | None = next(iterator) except StopIteration: observer.on_completed() return disp - def schedule_emission(data: T) -> None: + def schedule_emission(message: tuple[float, T]) -> None: nonlocal next_message, is_disposed if is_disposed: return - # Pre-load the following message while we have time + msg_ts, msg_data = message + try: next_message = next(iterator) except StopIteration: next_message = None - # Calculate absolute emission time - target_time = start_local_time + (data.ts - start_replay_time) / speed + target_time = start_local_time + (msg_ts - start_replay_time) / speed delay = max(0.0, target_time - time.time()) def emit( @@ -244,7 +290,7 @@ def emit( ) -> rx.abc.DisposableBase | None: if is_disposed: return None - observer.on_next(data) + observer.on_next(msg_data) if next_message is not None: schedule_emission(next_message) else: diff --git a/dimos/memory/sensor/legacy.py b/dimos/memory/sensor/legacy.py index 8ba4540f0a..c3b8d7c7f1 100644 --- a/dimos/memory/sensor/legacy.py +++ b/dimos/memory/sensor/legacy.py @@ -53,7 +53,7 @@ class LegacyPickleStore(SensorStore[T]): # Create new recording (directory created on first save) store = LegacyPickleStore("my_recording/images") - store.save(image) # uses image.ts for timestamp + store.save_ts(image) # uses image.ts for timestamp Backward compatibility: This class also supports the old TimedSensorReplay/SensorReplay API: diff --git a/dimos/memory/sensor/pickledir.py b/dimos/memory/sensor/pickledir.py index d6b1ca282f..fa0fa9446b 100644 --- a/dimos/memory/sensor/pickledir.py +++ b/dimos/memory/sensor/pickledir.py @@ -40,7 +40,7 @@ class PickleDirStore(SensorStore[T]): # Create new recording (directory created on first save) store = PickleDirStore("my_recording/images") - store.save(image) # uses image.ts for timestamp + store.save(image.ts, image) # explicit timestamp """ def __init__(self, name: str) -> None: diff --git a/dimos/memory/sensor/postgres.py b/dimos/memory/sensor/postgres.py index 6ee710d38f..8e01d04c84 100644 --- a/dimos/memory/sensor/postgres.py +++ b/dimos/memory/sensor/postgres.py @@ -18,7 +18,7 @@ import re import psycopg2 -from psycopg2.extensions import connection as PgConnection +import psycopg2.extensions from dimos.core.resource import Resource from dimos.memory.sensor.base import SensorStore, T @@ -50,7 +50,7 @@ class PostgresStore(SensorStore[T], Resource): store.start() # open connection # Use store - store.save(data) # uses data.ts for timestamp + store.save(data.ts, data) # explicit timestamp data = store.find_closest_seek(10.0) # Cleanup @@ -85,7 +85,7 @@ def __init__( self._host = host self._port = port self._user = user - self._conn: PgConnection | None = None + self._conn: psycopg2.extensions.connection | None = None self._table_created = False def start(self) -> None: @@ -105,7 +105,7 @@ def stop(self) -> None: self._conn.close() self._conn = None - def _get_conn(self) -> PgConnection: + def _get_conn(self) -> psycopg2.extensions.connection: """Get connection, starting if needed.""" if self._conn is None: self.start() diff --git a/dimos/memory/sensor/sqlite.py b/dimos/memory/sensor/sqlite.py index 874641c4f5..f406c05862 100644 --- a/dimos/memory/sensor/sqlite.py +++ b/dimos/memory/sensor/sqlite.py @@ -45,7 +45,7 @@ class SqliteStore(SensorStore[T]): Usage: # Named store (uses data/ directory, auto-downloads from LFS if needed) store = SqliteStore("recordings/lidar") # -> data/recordings/lidar.db - store.save(data) # uses data.ts for timestamp + store.save(data.ts, data) # explicit timestamp # Absolute path store = SqliteStore("/path/to/sensors.db") diff --git a/dimos/memory/sensor/test_base.py b/dimos/memory/sensor/test_base.py index 2506f4edcf..6f1eb35bcb 100644 --- a/dimos/memory/sensor/test_base.py +++ b/dimos/memory/sensor/test_base.py @@ -154,8 +154,8 @@ class TestSensorStore: 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)) + store.save(1.0, SampleData("data_at_1", 1.0)) + store.save(2.0, 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) @@ -163,9 +163,9 @@ def test_save_and_load(self, store_factory, store_name, temp_dir): def test_find_closest_timestamp(self, store_factory, store_name, temp_dir): store = store_factory(temp_dir) - store.save(SampleData("a", 1.0)) - store.save(SampleData("b", 2.0)) - store.save(SampleData("c", 3.0)) + store.save(1.0, SampleData("a", 1.0)) + store.save(2.0, SampleData("b", 2.0)) + store.save(3.0, SampleData("c", 3.0)) # Exact match assert store._find_closest_timestamp(2.0) == 2.0 @@ -182,9 +182,9 @@ def test_find_closest_timestamp(self, store_factory, store_name, temp_dir): def test_iter_items(self, store_factory, store_name, temp_dir): store = store_factory(temp_dir) - store.save(SampleData("a", 1.0)) - store.save(SampleData("b", 2.0)) - store.save(SampleData("c", 3.0)) + store.save(1.0, SampleData("a", 1.0)) + store.save(2.0, SampleData("b", 2.0)) + store.save(3.0, SampleData("c", 3.0)) # Should iterate in timestamp order items = list(store._iter_items()) @@ -196,10 +196,10 @@ def test_iter_items(self, store_factory, store_name, temp_dir): def test_iter_items_with_range(self, store_factory, store_name, temp_dir): store = store_factory(temp_dir) - store.save(SampleData("a", 1.0)) - store.save(SampleData("b", 2.0)) - store.save(SampleData("c", 3.0)) - store.save(SampleData("d", 4.0)) + store.save(1.0, SampleData("a", 1.0)) + store.save(2.0, SampleData("b", 2.0)) + store.save(3.0, SampleData("c", 3.0)) + store.save(4.0, SampleData("d", 4.0)) # Start only items = list(store._iter_items(start=2.0)) @@ -232,9 +232,9 @@ def test_first_and_first_timestamp(self, store_factory, store_name, temp_dir): assert store.first_timestamp() is None # Add data (in chronological order) - store.save(SampleData("a", 1.0)) - store.save(SampleData("b", 2.0)) - store.save(SampleData("c", 3.0)) + store.save(1.0, SampleData("a", 1.0)) + store.save(2.0, SampleData("b", 2.0)) + store.save(3.0, SampleData("c", 3.0)) # Should return first by timestamp assert store.first_timestamp() == 1.0 @@ -242,9 +242,9 @@ def test_first_and_first_timestamp(self, store_factory, store_name, temp_dir): def test_find_closest(self, store_factory, store_name, temp_dir): store = store_factory(temp_dir) - store.save(SampleData("a", 1.0)) - store.save(SampleData("b", 2.0)) - store.save(SampleData("c", 3.0)) + store.save(1.0, SampleData("a", 1.0)) + store.save(2.0, SampleData("b", 2.0)) + store.save(3.0, SampleData("c", 3.0)) # Exact match assert store.find_closest(2.0) == SampleData("b", 2.0) @@ -261,9 +261,9 @@ def test_find_closest(self, store_factory, store_name, temp_dir): def test_find_closest_seek(self, store_factory, store_name, temp_dir): store = store_factory(temp_dir) - store.save(SampleData("a", 10.0)) - store.save(SampleData("b", 11.0)) - store.save(SampleData("c", 12.0)) + store.save(10.0, SampleData("a", 10.0)) + store.save(11.0, SampleData("b", 11.0)) + store.save(12.0, SampleData("c", 12.0)) # Seek 0 = first item (10.0) assert store.find_closest_seek(0.0) == SampleData("a", 10.0) @@ -283,9 +283,9 @@ def test_find_closest_seek(self, store_factory, store_name, temp_dir): def test_iterate(self, store_factory, store_name, temp_dir): store = store_factory(temp_dir) - store.save(SampleData("a", 1.0)) - store.save(SampleData("b", 2.0)) - store.save(SampleData("c", 3.0)) + store.save(1.0, SampleData("a", 1.0)) + store.save(2.0, SampleData("b", 2.0)) + store.save(3.0, SampleData("c", 3.0)) # Should iterate in timestamp order, returning data only (not tuples) items = list(store.iterate()) @@ -297,10 +297,10 @@ def test_iterate(self, store_factory, store_name, temp_dir): 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)) - store.save(SampleData("b", 11.0)) - store.save(SampleData("c", 12.0)) - store.save(SampleData("d", 13.0)) + store.save(10.0, SampleData("a", 10.0)) + store.save(11.0, SampleData("b", 11.0)) + store.save(12.0, SampleData("c", 12.0)) + store.save(13.0, SampleData("d", 13.0)) # Seek from start items = list(store.iterate(seek=1.0)) @@ -322,11 +322,11 @@ def test_iterate_with_seek_and_duration(self, store_factory, store_name, temp_di 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): + def test_save_ts(self, store_factory, store_name, temp_dir): store = store_factory(temp_dir) - # Save multiple items at once - store.save( + # Save multiple Timestamped items using save_ts + store.save_ts( SampleData("a", 1.0), SampleData("b", 2.0), SampleData("c", 3.0), @@ -336,7 +336,7 @@ def test_variadic_save(self, store_factory, store_name, temp_dir): 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): + def test_pipe_save_ts(self, store_factory, store_name, temp_dir): import reactivex as rx store = store_factory(temp_dir) @@ -348,9 +348,9 @@ def test_pipe_save(self, store_factory, store_name, temp_dir): SampleData("c", 3.0), ) - # Pipe through store.pipe_save - should save and pass through + # Pipe through store.pipe_save_ts — should save and pass through results: list[SampleData] = [] - source.pipe(store.pipe_save).subscribe(results.append) + source.pipe(store.pipe_save_ts).subscribe(results.append) # Data should be saved assert store.load(1.0) == SampleData("a", 1.0) @@ -364,7 +364,24 @@ def test_pipe_save(self, store_factory, store_name, temp_dir): SampleData("c", 3.0), ] - def test_consume_stream(self, store_factory, store_name, temp_dir): + def test_pipe_save_with_key(self, store_factory, store_name, temp_dir): + import reactivex as rx + + store = store_factory(temp_dir) + + source = rx.of( + SampleData("a", 1.0), + SampleData("b", 2.0), + ) + + results: list[SampleData] = [] + source.pipe(store.pipe_save(lambda d: d.ts)).subscribe(results.append) + + assert store.load(1.0) == SampleData("a", 1.0) + assert store.load(2.0) == SampleData("b", 2.0) + assert len(results) == 2 + + def test_consume_stream_ts(self, store_factory, store_name, temp_dir): import reactivex as rx store = store_factory(temp_dir) @@ -376,8 +393,8 @@ def test_consume_stream(self, store_factory, store_name, temp_dir): SampleData("c", 3.0), ) - # Consume stream - should save all items - disposable = store.consume_stream(source) + # Consume stream — should save all items + disposable = store.consume_stream_ts(source) # Data should be saved assert store.load(1.0) == SampleData("a", 1.0) @@ -386,11 +403,46 @@ def test_consume_stream(self, store_factory, store_name, temp_dir): disposable.dispose() + def test_consume_stream_with_key(self, store_factory, store_name, temp_dir): + import reactivex as rx + + store = store_factory(temp_dir) + + source = rx.of( + SampleData("a", 1.0), + SampleData("b", 2.0), + ) + + disposable = store.consume_stream(source, key=lambda d: d.ts) + + assert store.load(1.0) == SampleData("a", 1.0) + assert store.load(2.0) == SampleData("b", 2.0) + + disposable.dispose() + + def test_iterate_items(self, store_factory, store_name, temp_dir): + store = store_factory(temp_dir) + store.save(1.0, SampleData("a", 1.0)) + store.save(2.0, SampleData("b", 2.0)) + store.save(3.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)) - store.save(SampleData("b", 2.0)) - store.save(SampleData("c", 3.0)) + store.save(1.0, SampleData("a", 1.0)) + store.save(2.0, SampleData("b", 2.0)) + store.save(3.0, SampleData("c", 3.0)) # Stream at high speed (essentially instant) results: list[SampleData] = [] @@ -409,3 +461,68 @@ def test_stream_basic(self, store_factory, store_name, temp_dir): SampleData("b", 2.0), SampleData("c", 3.0), ] + + +class TestNonTimestampedData: + """Test SensorStore with plain (non-Timestamped) data types.""" + + def test_store_plain_strings(self): + store: InMemoryStore[str] = InMemoryStore() + store.save(1.0, "hello") + store.save(2.0, "world") + + assert store.load(1.0) == "hello" + assert store.load(2.0) == "world" + assert store.first() == "hello" + assert store.first_timestamp() == 1.0 + + def test_iterate_plain_data(self): + store: InMemoryStore[int] = InMemoryStore() + store.save(10.0, 100) + store.save(20.0, 200) + store.save(30.0, 300) + + assert list(store.iterate()) == [100, 200, 300] + assert list(store.iterate_items()) == [(10.0, 100), (20.0, 200), (30.0, 300)] + + def test_find_closest_plain_data(self): + store: InMemoryStore[str] = InMemoryStore() + store.save(1.0, "a") + store.save(3.0, "b") + store.save(5.0, "c") + + assert store.find_closest(2.0) == "a" + assert store.find_closest(4.0) == "b" + assert store.find_closest_seek(2.0) == "b" + + def test_pipe_save_with_key_plain_data(self): + import reactivex as rx + + store: InMemoryStore[dict] = InMemoryStore() + source = rx.of( + {"ts": 1.0, "val": "x"}, + {"ts": 2.0, "val": "y"}, + ) + + results: list[dict] = [] + source.pipe(store.pipe_save(lambda d: d["ts"])).subscribe(results.append) + + assert store.load(1.0) == {"ts": 1.0, "val": "x"} + assert store.load(2.0) == {"ts": 2.0, "val": "y"} + assert len(results) == 2 + + def test_consume_stream_with_key_plain_data(self): + import reactivex as rx + + store: InMemoryStore[list] = InMemoryStore() + source = rx.of( + [1.0, "data_a"], + [2.0, "data_b"], + ) + + disposable = store.consume_stream(source, key=lambda d: d[0]) + + assert store.load(1.0) == [1.0, "data_a"] + assert store.load(2.0) == [2.0, "data_b"] + + disposable.dispose() From fa7752f4c7d0d678bbf7f636ada760192e5b3973 Mon Sep 17 00:00:00 2001 From: leshy <681516+leshy@users.noreply.github.com> Date: Wed, 11 Feb 2026 10:43:21 +0000 Subject: [PATCH 23/39] CI code cleanup --- dimos/memory/embedding.py | 4 ++-- dimos/robot/unitree_webrtc/type/test_odometry.py | 3 --- 2 files changed, 2 insertions(+), 5 deletions(-) diff --git a/dimos/memory/embedding.py b/dimos/memory/embedding.py index c9dc6f60d2..e92a0edb7e 100644 --- a/dimos/memory/embedding.py +++ b/dimos/memory/embedding.py @@ -20,14 +20,14 @@ from reactivex import operators as ops from reactivex.observable import Observable -from dimos.core import In, Module, ModuleConfig, Out, rpc +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 backpressure, getter_hot +from dimos.utils.reactive import getter_hot @dataclass diff --git a/dimos/robot/unitree_webrtc/type/test_odometry.py b/dimos/robot/unitree_webrtc/type/test_odometry.py index 13c10fc820..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 From 1cf5b0263d45fc5b2515044a66fffd74f27bba7f Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Wed, 11 Feb 2026 18:53:23 +0800 Subject: [PATCH 24/39] Fix ruff errors in embedding models Fix import sorting in treid.py and suppress B027 for optional warmup() in base.py. --- dimos/models/embedding/base.py | 12 ++++++++---- dimos/models/embedding/treid.py | 5 +++-- 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/dimos/models/embedding/base.py b/dimos/models/embedding/base.py index 7242a5efc6..a330600300 100644 --- a/dimos/models/embedding/base.py +++ b/dimos/models/embedding/base.py @@ -126,7 +126,9 @@ def compare_one_to_many(self, query: Embedding, candidates: list[Embedding]) -> 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[Embedding], candidates: list[Embedding]) -> 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. @@ -141,7 +143,9 @@ def compare_many_to_many(self, queries: list[Embedding], candidates: list[Embedd candidate_tensors = torch.stack([c.to_torch(self.device) for c in candidates]) return query_tensors @ candidate_tensors.T - def query(self, query_emb: Embedding, candidates: list[Embedding], 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). @@ -157,6 +161,6 @@ def query(self, query_emb: Embedding, candidates: list[Embedding], top_k: int = 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: + def warmup(self) -> None: # noqa: B027 """Optional warmup method to pre-load model.""" - pass + ... diff --git a/dimos/models/embedding/treid.py b/dimos/models/embedding/treid.py index ef444b0729..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 From 8d61809fdb360d54a5cc132a7a9c4b5524f8a12c Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Wed, 11 Feb 2026 18:55:33 +0800 Subject: [PATCH 25/39] Remove unused EmbeddingModel.warmup method --- dimos/models/embedding/base.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/dimos/models/embedding/base.py b/dimos/models/embedding/base.py index a330600300..c6b78fcf2c 100644 --- a/dimos/models/embedding/base.py +++ b/dimos/models/embedding/base.py @@ -161,6 +161,5 @@ def query( 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: # noqa: B027 - """Optional warmup method to pre-load model.""" + ... From 2608c9e796b00248fecf13a73fbbcfbd9b517300 Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Wed, 11 Feb 2026 18:57:47 +0800 Subject: [PATCH 26/39] Remove unused warmup(), fix consume_stream callers Remove dead EmbeddingModel.warmup() method. Update go2.py to use consume_stream_ts() for the new SensorStore API. --- dimos/robot/unitree/connection/go2.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/dimos/robot/unitree/connection/go2.py b/dimos/robot/unitree/connection/go2.py index 8e021b3510..57351f7ca0 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.consume_stream(self.connection.lidar_stream()) + lidar_store.consume_stream_ts(self.connection.lidar_stream()) odom_store: TimedSensorStorage = TimedSensorStorage(f"{recording_name}/odom") # type: ignore[type-arg] - odom_store.consume_stream(self.connection.odom_stream()) + odom_store.consume_stream_ts(self.connection.odom_stream()) video_store: TimedSensorStorage = TimedSensorStorage(f"{recording_name}/video") # type: ignore[type-arg] - video_store.consume_stream(self.connection.video_stream()) + video_store.consume_stream_ts(self.connection.video_stream()) @rpc def start(self) -> None: From 480dd9b619f712b830ad25f3d4fe315032ad1939 Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Wed, 11 Feb 2026 19:05:27 +0800 Subject: [PATCH 27/39] Rename SensorStore methods: Timestamped as default, raw for explicit ts save/pipe_save/consume_stream now work with Timestamped data by default. save_raw/pipe_save_raw/consume_stream_raw take explicit timestamps for non-Timestamped data. --- dimos/memory/sensor/base.py | 62 +++++++-------- dimos/memory/sensor/test_base.py | 104 ++++++++++++-------------- dimos/robot/unitree/connection/go2.py | 6 +- 3 files changed, 80 insertions(+), 92 deletions(-) diff --git a/dimos/memory/sensor/base.py b/dimos/memory/sensor/base.py index edd3852564..8a6576588d 100644 --- a/dimos/memory/sensor/base.py +++ b/dimos/memory/sensor/base.py @@ -36,9 +36,9 @@ class SensorStore(Generic[T], ABC): Implement 4 abstract methods for your backend (in-memory, pickle, sqlite, etc.). All iteration, streaming, and seek logic comes free from the base class. - T can be any type — timestamps are provided explicitly. For Timestamped - subclasses, convenience methods (save_ts, pipe_save_ts, consume_stream_ts) - automatically extract the .ts attribute. + T can be any type — timestamps are provided explicitly via the _raw methods. + The default save/pipe_save/consume_stream methods work with Timestamped data. + Use save_raw/pipe_save_raw/consume_stream_raw for non-Timestamped data. """ @abstractmethod @@ -65,20 +65,34 @@ def _find_closest_timestamp( """Find closest timestamp. Backend can optimize (binary search, db index, etc.).""" ... - def save(self, timestamp: float, data: T) -> None: - """Save a single data item at the given timestamp.""" - self._save(timestamp, data) - - def save_ts(self, *data: Timestamped) -> None: + def save(self, *data: Timestamped) -> None: """Save one or more Timestamped items using their .ts attribute.""" for item in data: self._save(item.ts, item) # type: ignore[arg-type] - def pipe_save(self, key: Callable[[T], float]) -> Callable[[Observable[T]], Observable[T]]: + def save_raw(self, timestamp: float, data: T) -> None: + """Save a single data item at the given timestamp.""" + self._save(timestamp, data) + + def pipe_save(self, source: Observable[T]) -> Observable[T]: + """Operator for Observable.pipe() — saves Timestamped items using .ts. + + Usage: + observable.pipe(store.pipe_save).subscribe(...) + """ + + def _save_and_return(data: T) -> T: + ts_data: Timestamped = data # type: ignore[assignment] + self._save(ts_data.ts, data) + return data + + return source.pipe(ops.map(_save_and_return)) + + def pipe_save_raw(self, key: Callable[[T], float]) -> Callable[[Observable[T]], Observable[T]]: """Operator for Observable.pipe() — saves each item using key(item) as timestamp. Usage: - observable.pipe(store.pipe_save(lambda x: x.ts)).subscribe(...) + observable.pipe(store.pipe_save_raw(lambda x: x.ts)).subscribe(...) """ def _operator(source: Observable[T]) -> Observable[T]: @@ -90,43 +104,29 @@ def _save_and_return(data: T) -> T: return _operator - def pipe_save_ts(self, source: Observable[T]) -> Observable[T]: - """Operator for Observable.pipe() — saves Timestamped items using .ts. + def consume_stream(self, observable: Observable[T]) -> rx.abc.DisposableBase: + """Subscribe to an observable and save Timestamped items using .ts. Usage: - observable.pipe(store.pipe_save_ts).subscribe(...) + disposable = store.consume_stream(observable) """ - def _save_and_return(data: T) -> T: + def _save_item(data: T) -> None: ts_data: Timestamped = data # type: ignore[assignment] self._save(ts_data.ts, data) - return data - return source.pipe(ops.map(_save_and_return)) + return observable.subscribe(on_next=_save_item) - def consume_stream( + def consume_stream_raw( self, observable: Observable[T], key: Callable[[T], float] ) -> rx.abc.DisposableBase: """Subscribe to an observable and save each item using key(item) as timestamp. Usage: - disposable = store.consume_stream(observable, key=lambda x: x.ts) + disposable = store.consume_stream_raw(observable, key=lambda x: x.ts) """ return observable.subscribe(on_next=lambda data: self._save(key(data), data)) - def consume_stream_ts(self, observable: Observable[T]) -> rx.abc.DisposableBase: - """Subscribe to an observable and save Timestamped items using .ts. - - Usage: - disposable = store.consume_stream_ts(observable) - """ - - def _save_item(data: T) -> None: - ts_data: Timestamped = data # type: ignore[assignment] - self._save(ts_data.ts, data) - - return observable.subscribe(on_next=_save_item) - def load(self, timestamp: float) -> T | None: """Load data at exact timestamp.""" return self._load(timestamp) diff --git a/dimos/memory/sensor/test_base.py b/dimos/memory/sensor/test_base.py index 6f1eb35bcb..7192060250 100644 --- a/dimos/memory/sensor/test_base.py +++ b/dimos/memory/sensor/test_base.py @@ -154,8 +154,8 @@ class TestSensorStore: def test_save_and_load(self, store_factory, store_name, temp_dir): store = store_factory(temp_dir) - store.save(1.0, SampleData("data_at_1", 1.0)) - store.save(2.0, SampleData("data_at_2", 2.0)) + 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) @@ -163,9 +163,7 @@ def test_save_and_load(self, store_factory, store_name, temp_dir): def test_find_closest_timestamp(self, store_factory, store_name, temp_dir): store = store_factory(temp_dir) - store.save(1.0, SampleData("a", 1.0)) - store.save(2.0, SampleData("b", 2.0)) - store.save(3.0, SampleData("c", 3.0)) + 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 @@ -182,9 +180,7 @@ def test_find_closest_timestamp(self, store_factory, store_name, temp_dir): def test_iter_items(self, store_factory, store_name, temp_dir): store = store_factory(temp_dir) - store.save(1.0, SampleData("a", 1.0)) - store.save(2.0, SampleData("b", 2.0)) - store.save(3.0, SampleData("c", 3.0)) + store.save(SampleData("a", 1.0), SampleData("b", 2.0), SampleData("c", 3.0)) # Should iterate in timestamp order items = list(store._iter_items()) @@ -196,10 +192,12 @@ def test_iter_items(self, store_factory, store_name, temp_dir): def test_iter_items_with_range(self, store_factory, store_name, temp_dir): store = store_factory(temp_dir) - store.save(1.0, SampleData("a", 1.0)) - store.save(2.0, SampleData("b", 2.0)) - store.save(3.0, SampleData("c", 3.0)) - store.save(4.0, SampleData("d", 4.0)) + 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)) @@ -232,9 +230,7 @@ def test_first_and_first_timestamp(self, store_factory, store_name, temp_dir): assert store.first_timestamp() is None # Add data (in chronological order) - store.save(1.0, SampleData("a", 1.0)) - store.save(2.0, SampleData("b", 2.0)) - store.save(3.0, SampleData("c", 3.0)) + 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 @@ -242,9 +238,7 @@ def test_first_and_first_timestamp(self, store_factory, store_name, temp_dir): def test_find_closest(self, store_factory, store_name, temp_dir): store = store_factory(temp_dir) - store.save(1.0, SampleData("a", 1.0)) - store.save(2.0, SampleData("b", 2.0)) - store.save(3.0, SampleData("c", 3.0)) + 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) @@ -261,9 +255,7 @@ def test_find_closest(self, store_factory, store_name, temp_dir): def test_find_closest_seek(self, store_factory, store_name, temp_dir): store = store_factory(temp_dir) - store.save(10.0, SampleData("a", 10.0)) - store.save(11.0, SampleData("b", 11.0)) - store.save(12.0, SampleData("c", 12.0)) + 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) @@ -283,9 +275,7 @@ def test_find_closest_seek(self, store_factory, store_name, temp_dir): def test_iterate(self, store_factory, store_name, temp_dir): store = store_factory(temp_dir) - store.save(1.0, SampleData("a", 1.0)) - store.save(2.0, SampleData("b", 2.0)) - store.save(3.0, SampleData("c", 3.0)) + 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()) @@ -297,10 +287,12 @@ def test_iterate(self, store_factory, store_name, temp_dir): def test_iterate_with_seek_and_duration(self, store_factory, store_name, temp_dir): store = store_factory(temp_dir) - store.save(10.0, SampleData("a", 10.0)) - store.save(11.0, SampleData("b", 11.0)) - store.save(12.0, SampleData("c", 12.0)) - store.save(13.0, SampleData("d", 13.0)) + 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)) @@ -322,11 +314,11 @@ def test_iterate_with_seek_and_duration(self, store_factory, store_name, temp_di items = list(store.iterate(from_timestamp=12.0)) assert items == [SampleData("c", 12.0), SampleData("d", 13.0)] - def test_save_ts(self, store_factory, store_name, temp_dir): + def test_variadic_save(self, store_factory, store_name, temp_dir): store = store_factory(temp_dir) - # Save multiple Timestamped items using save_ts - store.save_ts( + # Save multiple items at once + store.save( SampleData("a", 1.0), SampleData("b", 2.0), SampleData("c", 3.0), @@ -336,7 +328,7 @@ def test_save_ts(self, store_factory, store_name, temp_dir): assert store.load(2.0) == SampleData("b", 2.0) assert store.load(3.0) == SampleData("c", 3.0) - def test_pipe_save_ts(self, store_factory, store_name, temp_dir): + def test_pipe_save(self, store_factory, store_name, temp_dir): import reactivex as rx store = store_factory(temp_dir) @@ -348,9 +340,9 @@ def test_pipe_save_ts(self, store_factory, store_name, temp_dir): SampleData("c", 3.0), ) - # Pipe through store.pipe_save_ts — should save and pass through + # Pipe through store.pipe_save — should save and pass through results: list[SampleData] = [] - source.pipe(store.pipe_save_ts).subscribe(results.append) + source.pipe(store.pipe_save).subscribe(results.append) # Data should be saved assert store.load(1.0) == SampleData("a", 1.0) @@ -364,7 +356,7 @@ def test_pipe_save_ts(self, store_factory, store_name, temp_dir): SampleData("c", 3.0), ] - def test_pipe_save_with_key(self, store_factory, store_name, temp_dir): + def test_pipe_save_raw(self, store_factory, store_name, temp_dir): import reactivex as rx store = store_factory(temp_dir) @@ -375,13 +367,13 @@ def test_pipe_save_with_key(self, store_factory, store_name, temp_dir): ) results: list[SampleData] = [] - source.pipe(store.pipe_save(lambda d: d.ts)).subscribe(results.append) + source.pipe(store.pipe_save_raw(lambda d: d.ts)).subscribe(results.append) assert store.load(1.0) == SampleData("a", 1.0) assert store.load(2.0) == SampleData("b", 2.0) assert len(results) == 2 - def test_consume_stream_ts(self, store_factory, store_name, temp_dir): + def test_consume_stream(self, store_factory, store_name, temp_dir): import reactivex as rx store = store_factory(temp_dir) @@ -394,7 +386,7 @@ def test_consume_stream_ts(self, store_factory, store_name, temp_dir): ) # Consume stream — should save all items - disposable = store.consume_stream_ts(source) + disposable = store.consume_stream(source) # Data should be saved assert store.load(1.0) == SampleData("a", 1.0) @@ -403,7 +395,7 @@ def test_consume_stream_ts(self, store_factory, store_name, temp_dir): disposable.dispose() - def test_consume_stream_with_key(self, store_factory, store_name, temp_dir): + def test_consume_stream_raw(self, store_factory, store_name, temp_dir): import reactivex as rx store = store_factory(temp_dir) @@ -413,7 +405,7 @@ def test_consume_stream_with_key(self, store_factory, store_name, temp_dir): SampleData("b", 2.0), ) - disposable = store.consume_stream(source, key=lambda d: d.ts) + disposable = store.consume_stream_raw(source, key=lambda d: d.ts) assert store.load(1.0) == SampleData("a", 1.0) assert store.load(2.0) == SampleData("b", 2.0) @@ -422,9 +414,7 @@ def test_consume_stream_with_key(self, store_factory, store_name, temp_dir): def test_iterate_items(self, store_factory, store_name, temp_dir): store = store_factory(temp_dir) - store.save(1.0, SampleData("a", 1.0)) - store.save(2.0, SampleData("b", 2.0)) - store.save(3.0, SampleData("c", 3.0)) + store.save(SampleData("a", 1.0), SampleData("b", 2.0), SampleData("c", 3.0)) items = list(store.iterate_items()) assert items == [ @@ -440,9 +430,7 @@ def test_iterate_items(self, store_factory, store_name, temp_dir): def test_stream_basic(self, store_factory, store_name, temp_dir): store = store_factory(temp_dir) - store.save(1.0, SampleData("a", 1.0)) - store.save(2.0, SampleData("b", 2.0)) - store.save(3.0, SampleData("c", 3.0)) + store.save(SampleData("a", 1.0), SampleData("b", 2.0), SampleData("c", 3.0)) # Stream at high speed (essentially instant) results: list[SampleData] = [] @@ -468,8 +456,8 @@ class TestNonTimestampedData: def test_store_plain_strings(self): store: InMemoryStore[str] = InMemoryStore() - store.save(1.0, "hello") - store.save(2.0, "world") + store.save_raw(1.0, "hello") + store.save_raw(2.0, "world") assert store.load(1.0) == "hello" assert store.load(2.0) == "world" @@ -478,24 +466,24 @@ def test_store_plain_strings(self): def test_iterate_plain_data(self): store: InMemoryStore[int] = InMemoryStore() - store.save(10.0, 100) - store.save(20.0, 200) - store.save(30.0, 300) + store.save_raw(10.0, 100) + store.save_raw(20.0, 200) + store.save_raw(30.0, 300) assert list(store.iterate()) == [100, 200, 300] assert list(store.iterate_items()) == [(10.0, 100), (20.0, 200), (30.0, 300)] def test_find_closest_plain_data(self): store: InMemoryStore[str] = InMemoryStore() - store.save(1.0, "a") - store.save(3.0, "b") - store.save(5.0, "c") + store.save_raw(1.0, "a") + store.save_raw(3.0, "b") + store.save_raw(5.0, "c") assert store.find_closest(2.0) == "a" assert store.find_closest(4.0) == "b" assert store.find_closest_seek(2.0) == "b" - def test_pipe_save_with_key_plain_data(self): + def test_pipe_save_raw_plain_data(self): import reactivex as rx store: InMemoryStore[dict] = InMemoryStore() @@ -505,13 +493,13 @@ def test_pipe_save_with_key_plain_data(self): ) results: list[dict] = [] - source.pipe(store.pipe_save(lambda d: d["ts"])).subscribe(results.append) + source.pipe(store.pipe_save_raw(lambda d: d["ts"])).subscribe(results.append) assert store.load(1.0) == {"ts": 1.0, "val": "x"} assert store.load(2.0) == {"ts": 2.0, "val": "y"} assert len(results) == 2 - def test_consume_stream_with_key_plain_data(self): + def test_consume_stream_raw_plain_data(self): import reactivex as rx store: InMemoryStore[list] = InMemoryStore() @@ -520,7 +508,7 @@ def test_consume_stream_with_key_plain_data(self): [2.0, "data_b"], ) - disposable = store.consume_stream(source, key=lambda d: d[0]) + disposable = store.consume_stream_raw(source, key=lambda d: d[0]) assert store.load(1.0) == [1.0, "data_a"] assert store.load(2.0) == [2.0, "data_b"] diff --git a/dimos/robot/unitree/connection/go2.py b/dimos/robot/unitree/connection/go2.py index 57351f7ca0..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.consume_stream_ts(self.connection.lidar_stream()) + lidar_store.consume_stream(self.connection.lidar_stream()) odom_store: TimedSensorStorage = TimedSensorStorage(f"{recording_name}/odom") # type: ignore[type-arg] - odom_store.consume_stream_ts(self.connection.odom_stream()) + odom_store.consume_stream(self.connection.odom_stream()) video_store: TimedSensorStorage = TimedSensorStorage(f"{recording_name}/video") # type: ignore[type-arg] - video_store.consume_stream_ts(self.connection.video_stream()) + video_store.consume_stream(self.connection.video_stream()) @rpc def start(self) -> None: From e8fc574925f27ff3f7d9eb7356fecf5fc63945ec Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Wed, 11 Feb 2026 19:15:13 +0800 Subject: [PATCH 28/39] Rename SensorStore to TimeSeriesStore --- dimos/memory/sensor/__init__.py | 4 ++-- dimos/memory/sensor/base.py | 8 ++++++-- dimos/memory/sensor/legacy.py | 6 +++--- dimos/memory/sensor/pickledir.py | 6 +++--- dimos/memory/sensor/postgres.py | 6 +++--- dimos/memory/sensor/sqlite.py | 6 +++--- dimos/memory/sensor/test_base.py | 20 ++++++++++---------- 7 files changed, 30 insertions(+), 26 deletions(-) diff --git a/dimos/memory/sensor/__init__.py b/dimos/memory/sensor/__init__.py index 7be151f2c4..f7756fd4fc 100644 --- a/dimos/memory/sensor/__init__.py +++ b/dimos/memory/sensor/__init__.py @@ -13,7 +13,7 @@ # limitations under the License. """Sensor storage and replay.""" -from dimos.memory.sensor.base import InMemoryStore, SensorStore +from dimos.memory.sensor.base import InMemoryStore, TimeSeriesStore from dimos.memory.sensor.pickledir import PickleDirStore from dimos.memory.sensor.postgres import PostgresStore, reset_db from dimos.memory.sensor.sqlite import SqliteStore @@ -22,7 +22,7 @@ "InMemoryStore", "PickleDirStore", "PostgresStore", - "SensorStore", "SqliteStore", + "TimeSeriesStore", "reset_db", ] diff --git a/dimos/memory/sensor/base.py b/dimos/memory/sensor/base.py index 8a6576588d..03b3569070 100644 --- a/dimos/memory/sensor/base.py +++ b/dimos/memory/sensor/base.py @@ -30,7 +30,7 @@ T = TypeVar("T") -class SensorStore(Generic[T], ABC): +class TimeSeriesStore(Generic[T], ABC): """Unified storage + replay for sensor data. Implement 4 abstract methods for your backend (in-memory, pickle, sqlite, etc.). @@ -131,6 +131,10 @@ def load(self, timestamp: float) -> T | None: """Load data at exact timestamp.""" return self._load(timestamp) + def get(self, timestamp: float, tolerance: float | None = None) -> T | None: + """Get data at exact timestamp or closest within tolerance.""" + return self.find_closest(timestamp, tolerance) + def find_closest( self, timestamp: float, @@ -312,7 +316,7 @@ def dispose() -> None: return rx.create(subscribe) -class InMemoryStore(SensorStore[T]): +class InMemoryStore(TimeSeriesStore[T]): """In-memory storage using dict. Good for live use.""" def __init__(self) -> None: diff --git a/dimos/memory/sensor/legacy.py b/dimos/memory/sensor/legacy.py index c3b8d7c7f1..1bdef9ce1f 100644 --- a/dimos/memory/sensor/legacy.py +++ b/dimos/memory/sensor/legacy.py @@ -11,7 +11,7 @@ # 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 SensorStore. +"""Legacy pickle directory backend for TimeSeriesStore. Compatible with TimedSensorReplay/TimedSensorStorage file format. """ @@ -30,11 +30,11 @@ from reactivex.observable import Observable from reactivex.scheduler import TimeoutScheduler -from dimos.memory.sensor.base import SensorStore, T +from dimos.memory.sensor.base import T, TimeSeriesStore from dimos.utils.data import get_data, get_data_dir -class LegacyPickleStore(SensorStore[T]): +class LegacyPickleStore(TimeSeriesStore[T]): """Legacy pickle backend compatible with TimedSensorReplay/TimedSensorStorage. File format: diff --git a/dimos/memory/sensor/pickledir.py b/dimos/memory/sensor/pickledir.py index fa0fa9446b..c92258f157 100644 --- a/dimos/memory/sensor/pickledir.py +++ b/dimos/memory/sensor/pickledir.py @@ -11,7 +11,7 @@ # 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 SensorStore.""" +"""Pickle directory backend for TimeSeriesStore.""" import bisect from collections.abc import Iterator @@ -20,11 +20,11 @@ from pathlib import Path import pickle -from dimos.memory.sensor.base import SensorStore, T +from dimos.memory.sensor.base import T, TimeSeriesStore from dimos.utils.data import get_data, get_data_dir -class PickleDirStore(SensorStore[T]): +class PickleDirStore(TimeSeriesStore[T]): """Pickle directory backend. Files named by timestamp. Directory structure: diff --git a/dimos/memory/sensor/postgres.py b/dimos/memory/sensor/postgres.py index 8e01d04c84..31b559f761 100644 --- a/dimos/memory/sensor/postgres.py +++ b/dimos/memory/sensor/postgres.py @@ -11,7 +11,7 @@ # 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 SensorStore.""" +"""PostgreSQL backend for TimeSeriesStore.""" from collections.abc import Iterator import pickle @@ -21,7 +21,7 @@ import psycopg2.extensions from dimos.core.resource import Resource -from dimos.memory.sensor.base import SensorStore, T +from dimos.memory.sensor.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_]*$") @@ -38,7 +38,7 @@ def _validate_identifier(name: str) -> str: return name -class PostgresStore(SensorStore[T], Resource): +class PostgresStore(TimeSeriesStore[T], Resource): """PostgreSQL backend for sensor data. Multiple stores can share the same database with different tables. diff --git a/dimos/memory/sensor/sqlite.py b/dimos/memory/sensor/sqlite.py index f406c05862..08ba53d4b1 100644 --- a/dimos/memory/sensor/sqlite.py +++ b/dimos/memory/sensor/sqlite.py @@ -11,7 +11,7 @@ # 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 SensorStore.""" +"""SQLite backend for TimeSeriesStore.""" from collections.abc import Iterator from pathlib import Path @@ -19,7 +19,7 @@ import re import sqlite3 -from dimos.memory.sensor.base import SensorStore, T +from dimos.memory.sensor.base import T, TimeSeriesStore from dimos.utils.data import get_data, get_data_dir # Valid SQL identifier: alphanumeric and underscores, not starting with digit @@ -37,7 +37,7 @@ def _validate_identifier(name: str) -> str: return name -class SqliteStore(SensorStore[T]): +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. diff --git a/dimos/memory/sensor/test_base.py b/dimos/memory/sensor/test_base.py index 7192060250..e6c6e638dc 100644 --- a/dimos/memory/sensor/test_base.py +++ b/dimos/memory/sensor/test_base.py @@ -11,7 +11,7 @@ # 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 SensorStore implementations.""" +"""Tests for TimeSeriesStore implementations.""" from dataclasses import dataclass from pathlib import Path @@ -20,7 +20,7 @@ import pytest -from dimos.memory.sensor.base import InMemoryStore, SensorStore +from dimos.memory.sensor.base import InMemoryStore, TimeSeriesStore from dimos.memory.sensor.legacy import LegacyPickleStore from dimos.memory.sensor.pickledir import PickleDirStore from dimos.memory.sensor.sqlite import SqliteStore @@ -82,19 +82,19 @@ def temp_dir(): yield tmpdir -def make_in_memory_store() -> SensorStore[SampleData]: +def make_in_memory_store() -> TimeSeriesStore[SampleData]: return InMemoryStore[SampleData]() -def make_pickle_dir_store(tmpdir: str) -> SensorStore[SampleData]: +def make_pickle_dir_store(tmpdir: str) -> TimeSeriesStore[SampleData]: return PickleDirStore[SampleData](tmpdir) -def make_sqlite_store(tmpdir: str) -> SensorStore[SampleData]: +def make_sqlite_store(tmpdir: str) -> TimeSeriesStore[SampleData]: return SqliteStore[SampleData](Path(tmpdir) / "test.db") -def make_legacy_pickle_store(tmpdir: str) -> SensorStore[SampleData]: +def make_legacy_pickle_store(tmpdir: str) -> TimeSeriesStore[SampleData]: return LegacyPickleStore[SampleData](Path(tmpdir) / "legacy") @@ -118,7 +118,7 @@ def make_legacy_pickle_store(tmpdir: str) -> SensorStore[SampleData]: _test_conn = psycopg2.connect(dbname="dimensional") _test_conn.close() - def make_postgres_store(_tmpdir: str) -> SensorStore[SampleData]: + 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) @@ -149,8 +149,8 @@ def cleanup_postgres_tables(): @pytest.mark.parametrize("store_factory,store_name", testdata) -class TestSensorStore: - """Parametrized tests for all SensorStore implementations.""" +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) @@ -452,7 +452,7 @@ def test_stream_basic(self, store_factory, store_name, temp_dir): class TestNonTimestampedData: - """Test SensorStore with plain (non-Timestamped) data types.""" + """Test TimeSeriesStore with plain (non-Timestamped) data types.""" def test_store_plain_strings(self): store: InMemoryStore[str] = InMemoryStore() From a8b472b9d7b8d9ac3d221663583fb67f16411dbb Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Wed, 11 Feb 2026 19:20:08 +0800 Subject: [PATCH 29/39] Add _delete to all backends, fix find_closest and add/get/prune_old Implement _delete for InMemoryStore, SqliteStore, PickleDirStore, PostgresStore (LegacyPickleStore raises NotImplementedError). Fix find_closest docstring placement and add get/add/prune_old convenience methods. --- dimos/memory/sensor/base.py | 20 +++++++++++++++++++- dimos/memory/sensor/legacy.py | 4 ++++ dimos/memory/sensor/pickledir.py | 9 +++++++++ dimos/memory/sensor/postgres.py | 9 +++++++++ dimos/memory/sensor/sqlite.py | 8 ++++++++ 5 files changed, 49 insertions(+), 1 deletion(-) diff --git a/dimos/memory/sensor/base.py b/dimos/memory/sensor/base.py index 03b3569070..c76d764f27 100644 --- a/dimos/memory/sensor/base.py +++ b/dimos/memory/sensor/base.py @@ -51,6 +51,11 @@ 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 @@ -131,10 +136,19 @@ def load(self, timestamp: float) -> T | None: """Load data at exact timestamp.""" return self._load(timestamp) - def get(self, timestamp: float, tolerance: float | None = None) -> T | None: + def get(self, timestamp: float, tolerance: float | None = 1.0) -> T | None: """Get data at exact timestamp or closest within tolerance.""" return self.find_closest(timestamp, tolerance) + def add(self, data: Timestamped) -> None: + """Add a single Timestamped item using its .ts attribute.""" + self._save(data.ts, data) # type: ignore[arg-type] + + def prune_old(self, cutoff: float) -> None: + """Prune items older than cutoff timestamp.""" + for ts, _ in self._iter_items(end=cutoff): + self._delete(ts) + def find_closest( self, timestamp: float, @@ -330,6 +344,10 @@ def _save(self, timestamp: float, data: T) -> None: def _load(self, timestamp: float) -> T | None: return self._data.get(timestamp) + def _delete(self, timestamp: float) -> T | None: + self._sorted_timestamps = None # Invalidate cache + return self._data.pop(timestamp, None) + def _iter_items( self, start: float | None = None, end: float | None = None ) -> Iterator[tuple[float, T]]: diff --git a/dimos/memory/sensor/legacy.py b/dimos/memory/sensor/legacy.py index 1bdef9ce1f..fdb57dbdd1 100644 --- a/dimos/memory/sensor/legacy.py +++ b/dimos/memory/sensor/legacy.py @@ -146,6 +146,10 @@ def _load(self, timestamp: float) -> T | None: 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]]: diff --git a/dimos/memory/sensor/pickledir.py b/dimos/memory/sensor/pickledir.py index c92258f157..5a3831e547 100644 --- a/dimos/memory/sensor/pickledir.py +++ b/dimos/memory/sensor/pickledir.py @@ -92,6 +92,15 @@ def _load(self, timestamp: float) -> T | None: 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]]: diff --git a/dimos/memory/sensor/postgres.py b/dimos/memory/sensor/postgres.py index 31b559f761..d6183872c5 100644 --- a/dimos/memory/sensor/postgres.py +++ b/dimos/memory/sensor/postgres.py @@ -156,6 +156,15 @@ def _load(self, timestamp: float) -> T | 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]]: diff --git a/dimos/memory/sensor/sqlite.py b/dimos/memory/sensor/sqlite.py index 08ba53d4b1..1739d40efe 100644 --- a/dimos/memory/sensor/sqlite.py +++ b/dimos/memory/sensor/sqlite.py @@ -143,6 +143,14 @@ def _load(self, timestamp: float) -> T | 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]]: From e615a2e71112795349bd4d7f642e7afd57c2efcf Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Wed, 11 Feb 2026 19:39:33 +0800 Subject: [PATCH 30/39] Add collection API to TimeSeriesStore, rewrite InMemoryStore with SortedKeyList Replace InMemoryStore's dict + sorted-cache (O(n log n) rebuild on every write) with SortedKeyList for O(log n) insert, delete, and range queries. Add collection methods to TimeSeriesStore base: __len__, __iter__, last/last_timestamp, start_ts/ end_ts, time_range, duration, find_before/find_after, slice_by_time. Implement backing abstract methods (_count, _last_timestamp, _find_before, _find_after) in all five backends. Performance benchmarks confirm InMemoryStore matches TimestampedCollection on 100k items. --- dimos/memory/sensor/base.py | 157 ++++++++++++++++++++---- dimos/memory/sensor/legacy.py | 24 ++++ dimos/memory/sensor/pickledir.py | 31 +++++ dimos/memory/sensor/postgres.py | 44 +++++++ dimos/memory/sensor/sqlite.py | 35 ++++++ dimos/memory/sensor/test_base.py | 204 +++++++++++++++++++++++++++++++ 6 files changed, 469 insertions(+), 26 deletions(-) diff --git a/dimos/memory/sensor/base.py b/dimos/memory/sensor/base.py index c76d764f27..5614010f2d 100644 --- a/dimos/memory/sensor/base.py +++ b/dimos/memory/sensor/base.py @@ -14,7 +14,6 @@ """Unified sensor storage and replay.""" from abc import ABC, abstractmethod -import bisect from collections.abc import Callable, Iterator import time from typing import Generic, TypeVar @@ -24,6 +23,7 @@ from reactivex.disposable import CompositeDisposable, Disposable from reactivex.observable import Observable from reactivex.scheduler import TimeoutScheduler +from sortedcontainers import SortedKeyList # type: ignore[import-untyped] from dimos.types.timestamped import Timestamped @@ -70,6 +70,84 @@ def _find_closest_timestamp( """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: Timestamped) -> None: """Save one or more Timestamped items using their .ts attribute.""" for item in data: @@ -331,47 +409,55 @@ def dispose() -> None: class InMemoryStore(TimeSeriesStore[T]): - """In-memory storage using dict. Good for live use.""" + """In-memory storage using SortedKeyList. O(log n) insert, lookup, and range queries.""" def __init__(self) -> None: - self._data: dict[float, T] = {} - self._sorted_timestamps: list[float] | None = None + self._entries: SortedKeyList = SortedKeyList(key=lambda entry: entry[0]) + self._by_ts: dict[float, T] = {} def _save(self, timestamp: float, data: T) -> None: - self._data[timestamp] = data - self._sorted_timestamps = None # Invalidate cache + if timestamp in self._by_ts: + # Update: remove old entry from sorted list, then re-add + old = self._by_ts[timestamp] + self._entries.remove((timestamp, old)) + self._by_ts[timestamp] = data + self._entries.add((timestamp, data)) def _load(self, timestamp: float) -> T | None: - return self._data.get(timestamp) + return self._by_ts.get(timestamp) def _delete(self, timestamp: float) -> T | None: - self._sorted_timestamps = None # Invalidate cache - return self._data.pop(timestamp, None) + old = self._by_ts.pop(timestamp, None) + if old is not None: + self._entries.remove((timestamp, old)) + return old def _iter_items( self, start: float | None = None, end: float | None = None ) -> Iterator[tuple[float, T]]: - for ts in self._get_sorted_timestamps(): - if start is not None and ts < start: - continue - if end is not None and ts >= end: - break - yield (ts, self._data[ts]) + 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) + yield from it def _find_closest_timestamp( self, timestamp: float, tolerance: float | None = None ) -> float | None: - timestamps = self._get_sorted_timestamps() - if not timestamps: + if not self._entries: return None - pos = bisect.bisect_left(timestamps, timestamp) + pos = self._entries.bisect_key_left(timestamp) - candidates = [] + candidates: list[float] = [] if pos > 0: - candidates.append(timestamps[pos - 1]) - if pos < len(timestamps): - candidates.append(timestamps[pos]) + candidates.append(self._entries[pos - 1][0]) + if pos < len(self._entries): + candidates.append(self._entries[pos][0]) if not candidates: return None @@ -383,7 +469,26 @@ def _find_closest_timestamp( return closest - def _get_sorted_timestamps(self) -> list[float]: - if self._sorted_timestamps is None: - self._sorted_timestamps = sorted(self._data.keys()) - return self._sorted_timestamps + def _count(self) -> int: + return len(self._by_ts) + + def _last_timestamp(self) -> float | None: + if not self._entries: + return None + return self._entries[-1][0] # 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: + return self._entries[pos - 1] # type: ignore[no-any-return] + 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): + return self._entries[pos] # type: ignore[no-any-return] + return None diff --git a/dimos/memory/sensor/legacy.py b/dimos/memory/sensor/legacy.py index fdb57dbdd1..11a26c21d2 100644 --- a/dimos/memory/sensor/legacy.py +++ b/dimos/memory/sensor/legacy.py @@ -208,6 +208,30 @@ def _find_closest_timestamp( 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 diff --git a/dimos/memory/sensor/pickledir.py b/dimos/memory/sensor/pickledir.py index 5a3831e547..dc050dff84 100644 --- a/dimos/memory/sensor/pickledir.py +++ b/dimos/memory/sensor/pickledir.py @@ -157,6 +157,37 @@ def _get_timestamps(self) -> list[float]: 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: diff --git a/dimos/memory/sensor/postgres.py b/dimos/memory/sensor/postgres.py index d6183872c5..56d64672b1 100644 --- a/dimos/memory/sensor/postgres.py +++ b/dimos/memory/sensor/postgres.py @@ -238,6 +238,50 @@ def _find_closest_timestamp( 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. diff --git a/dimos/memory/sensor/sqlite.py b/dimos/memory/sensor/sqlite.py index 1739d40efe..2d54a3d172 100644 --- a/dimos/memory/sensor/sqlite.py +++ b/dimos/memory/sensor/sqlite.py @@ -223,6 +223,41 @@ def _find_closest_timestamp( 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: diff --git a/dimos/memory/sensor/test_base.py b/dimos/memory/sensor/test_base.py index e6c6e638dc..9510a6cf7a 100644 --- a/dimos/memory/sensor/test_base.py +++ b/dimos/memory/sensor/test_base.py @@ -451,6 +451,109 @@ def test_stream_basic(self, store_factory, store_name, temp_dir): ] +@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)] + + +class TestInMemoryStoreUpdate: + """Test InMemoryStore duplicate timestamp handling.""" + + def test_overwrite_existing_timestamp(self): + store: InMemoryStore[str] = InMemoryStore() + store.save_raw(1.0, "old") + store.save_raw(1.0, "new") + assert store.load(1.0) == "new" + assert len(store) == 1 + + def test_delete(self): + store: InMemoryStore[str] = InMemoryStore() + store.save_raw(1.0, "a") + store.save_raw(2.0, "b") + assert len(store) == 2 + deleted = store._delete(1.0) + assert deleted == "a" + assert len(store) == 1 + assert store.load(1.0) is None + + class TestNonTimestampedData: """Test TimeSeriesStore with plain (non-Timestamped) data types.""" @@ -514,3 +617,104 @@ def test_consume_stream_raw_plain_data(self): assert store.load(2.0) == [2.0, "data_b"] disposable.dispose() + + +class TestPerformance: + """Benchmarks comparing InMemoryStore vs TimestampedCollection.""" + + N = 100_000 + + def _make_populated_store(self) -> InMemoryStore[SampleData]: + store: InMemoryStore[SampleData] = InMemoryStore() + for i in range(self.N): + store.save_raw(float(i), SampleData(f"v{i}", float(i))) + return store + + def _make_populated_collection(self) -> "TimestampedCollection[SampleData]": + from dimos.types.timestamped import TimestampedCollection + + coll: TimestampedCollection[SampleData] = TimestampedCollection() + for i in range(self.N): + coll.add(SampleData(f"v{i}", float(i))) + return coll + + def test_insert_performance(self) -> None: + """Insert N items. InMemoryStore should match TimestampedCollection.""" + import time as time_mod + + from dimos.types.timestamped import TimestampedCollection + + store: InMemoryStore[SampleData] = InMemoryStore() + t0 = time_mod.perf_counter() + for i in range(self.N): + store.save_raw(float(i), SampleData(f"v{i}", float(i))) + store_time = time_mod.perf_counter() - t0 + + coll: TimestampedCollection[SampleData] = TimestampedCollection() + t0 = time_mod.perf_counter() + for i in range(self.N): + coll.add(SampleData(f"v{i}", float(i))) + coll_time = time_mod.perf_counter() - t0 + + print(f"\nInsert {self.N}: store={store_time:.3f}s, collection={coll_time:.3f}s") + assert store_time < coll_time * 5 + + def test_find_closest_performance(self) -> None: + """find_closest on N items. Both should be O(log n).""" + import random + import time as time_mod + + store = self._make_populated_store() + coll = self._make_populated_collection() + + queries = [random.uniform(0, self.N) for _ in range(10_000)] + + t0 = time_mod.perf_counter() + for q in queries: + store.find_closest(q) + store_time = time_mod.perf_counter() - t0 + + t0 = time_mod.perf_counter() + for q in queries: + coll.find_closest(q) + coll_time = time_mod.perf_counter() - t0 + + print( + f"\nfind_closest 10k on {self.N}: store={store_time:.3f}s, collection={coll_time:.3f}s" + ) + assert store_time < coll_time * 5 + + def test_interleaved_write_read(self) -> None: + """Alternating write + find_closest. Old InMemoryStore was O(n log n) per read.""" + import time as time_mod + + store: InMemoryStore[SampleData] = InMemoryStore() + + t0 = time_mod.perf_counter() + for i in range(self.N): + store.save_raw(float(i), SampleData(f"v{i}", float(i))) + if i % 10 == 0: + store.find_closest(float(i) / 2) + elapsed = time_mod.perf_counter() - t0 + + print(f"\nInterleaved write+read {self.N}: {elapsed:.3f}s") + assert elapsed < 10.0 + + def test_iteration_performance(self) -> None: + """Full iteration over N items.""" + import time as time_mod + + store = self._make_populated_store() + coll = self._make_populated_collection() + + t0 = time_mod.perf_counter() + count_store = sum(1 for _ in store) + store_time = time_mod.perf_counter() - t0 + + t0 = time_mod.perf_counter() + count_coll = sum(1 for _ in coll) + coll_time = time_mod.perf_counter() - t0 + + assert count_store == count_coll == self.N + print(f"\nIterate {self.N}: store={store_time:.3f}s, collection={coll_time:.3f}s") + assert store_time < coll_time * 5 From a06b63d1a5fef0675aa45f6df8f5ff10d6ce52ef Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Wed, 11 Feb 2026 19:49:13 +0800 Subject: [PATCH 31/39] Rename memory/sensor to memory/timeseries, extract InMemoryStore to inmemory.py Rename the module to better reflect its purpose. Extract InMemoryStore from base.py into its own file (inmemory.py) to keep base.py focused on the abstract TimeSeriesStore class. Update all internal and external imports. --- .../memory/{sensor => timeseries}/__init__.py | 11 +- dimos/memory/{sensor => timeseries}/base.py | 89 +-------------- dimos/memory/timeseries/inmemory.py | 106 ++++++++++++++++++ dimos/memory/{sensor => timeseries}/legacy.py | 2 +- .../{sensor => timeseries}/pickledir.py | 2 +- .../memory/{sensor => timeseries}/postgres.py | 2 +- dimos/memory/{sensor => timeseries}/sqlite.py | 2 +- .../{sensor => timeseries}/test_base.py | 14 ++- dimos/utils/testing/replay.py | 2 +- 9 files changed, 126 insertions(+), 104 deletions(-) rename dimos/memory/{sensor => timeseries}/__init__.py (68%) rename dimos/memory/{sensor => timeseries}/base.py (81%) create mode 100644 dimos/memory/timeseries/inmemory.py rename dimos/memory/{sensor => timeseries}/legacy.py (99%) rename dimos/memory/{sensor => timeseries}/pickledir.py (99%) rename dimos/memory/{sensor => timeseries}/postgres.py (99%) rename dimos/memory/{sensor => timeseries}/sqlite.py (99%) rename dimos/memory/{sensor => timeseries}/test_base.py (98%) diff --git a/dimos/memory/sensor/__init__.py b/dimos/memory/timeseries/__init__.py similarity index 68% rename from dimos/memory/sensor/__init__.py rename to dimos/memory/timeseries/__init__.py index f7756fd4fc..39cb48d5cb 100644 --- a/dimos/memory/sensor/__init__.py +++ b/dimos/memory/timeseries/__init__.py @@ -11,12 +11,13 @@ # 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. -"""Sensor storage and replay.""" +"""Time series storage and replay.""" -from dimos.memory.sensor.base import InMemoryStore, TimeSeriesStore -from dimos.memory.sensor.pickledir import PickleDirStore -from dimos.memory.sensor.postgres import PostgresStore, reset_db -from dimos.memory.sensor.sqlite import SqliteStore +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.postgres import PostgresStore, reset_db +from dimos.memory.timeseries.sqlite import SqliteStore __all__ = [ "InMemoryStore", diff --git a/dimos/memory/sensor/base.py b/dimos/memory/timeseries/base.py similarity index 81% rename from dimos/memory/sensor/base.py rename to dimos/memory/timeseries/base.py index 5614010f2d..95e470d39e 100644 --- a/dimos/memory/sensor/base.py +++ b/dimos/memory/timeseries/base.py @@ -11,7 +11,7 @@ # 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 sensor storage and replay.""" +"""Unified time series storage and replay.""" from abc import ABC, abstractmethod from collections.abc import Callable, Iterator @@ -23,7 +23,6 @@ from reactivex.disposable import CompositeDisposable, Disposable from reactivex.observable import Observable from reactivex.scheduler import TimeoutScheduler -from sortedcontainers import SortedKeyList # type: ignore[import-untyped] from dimos.types.timestamped import Timestamped @@ -406,89 +405,3 @@ def dispose() -> None: return Disposable(dispose) return rx.create(subscribe) - - -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 entry: entry[0]) - self._by_ts: dict[float, T] = {} - - def _save(self, timestamp: float, data: T) -> None: - if timestamp in self._by_ts: - # Update: remove old entry from sorted list, then re-add - old = self._by_ts[timestamp] - self._entries.remove((timestamp, old)) - self._by_ts[timestamp] = data - self._entries.add((timestamp, data)) - - def _load(self, timestamp: float) -> T | None: - return self._by_ts.get(timestamp) - - def _delete(self, timestamp: float) -> T | None: - old = self._by_ts.pop(timestamp, None) - if old is not None: - self._entries.remove((timestamp, old)) - return old - - 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) - yield from it - - 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][0]) - if pos < len(self._entries): - candidates.append(self._entries[pos][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: - return len(self._by_ts) - - def _last_timestamp(self) -> float | None: - if not self._entries: - return None - return self._entries[-1][0] # 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: - return self._entries[pos - 1] # type: ignore[no-any-return] - 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): - return self._entries[pos] # type: ignore[no-any-return] - return None diff --git a/dimos/memory/timeseries/inmemory.py b/dimos/memory/timeseries/inmemory.py new file mode 100644 index 0000000000..8aaf45e9d7 --- /dev/null +++ b/dimos/memory/timeseries/inmemory.py @@ -0,0 +1,106 @@ +# 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 entry: entry[0]) + self._by_ts: dict[float, T] = {} + + def _save(self, timestamp: float, data: T) -> None: + if timestamp in self._by_ts: + # Update: remove old entry from sorted list, then re-add + old = self._by_ts[timestamp] + self._entries.remove((timestamp, old)) + self._by_ts[timestamp] = data + self._entries.add((timestamp, data)) + + def _load(self, timestamp: float) -> T | None: + return self._by_ts.get(timestamp) + + def _delete(self, timestamp: float) -> T | None: + old = self._by_ts.pop(timestamp, None) + if old is not None: + self._entries.remove((timestamp, old)) + return old + + 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) + yield from it + + 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][0]) + if pos < len(self._entries): + candidates.append(self._entries[pos][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: + return len(self._by_ts) + + def _last_timestamp(self) -> float | None: + if not self._entries: + return None + return self._entries[-1][0] # 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: + return self._entries[pos - 1] # type: ignore[no-any-return] + 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): + return self._entries[pos] # type: ignore[no-any-return] + return None diff --git a/dimos/memory/sensor/legacy.py b/dimos/memory/timeseries/legacy.py similarity index 99% rename from dimos/memory/sensor/legacy.py rename to dimos/memory/timeseries/legacy.py index 11a26c21d2..821d306d2d 100644 --- a/dimos/memory/sensor/legacy.py +++ b/dimos/memory/timeseries/legacy.py @@ -30,7 +30,7 @@ from reactivex.observable import Observable from reactivex.scheduler import TimeoutScheduler -from dimos.memory.sensor.base import T, TimeSeriesStore +from dimos.memory.timeseries.base import T, TimeSeriesStore from dimos.utils.data import get_data, get_data_dir diff --git a/dimos/memory/sensor/pickledir.py b/dimos/memory/timeseries/pickledir.py similarity index 99% rename from dimos/memory/sensor/pickledir.py rename to dimos/memory/timeseries/pickledir.py index dc050dff84..e9acbd8aba 100644 --- a/dimos/memory/sensor/pickledir.py +++ b/dimos/memory/timeseries/pickledir.py @@ -20,7 +20,7 @@ from pathlib import Path import pickle -from dimos.memory.sensor.base import T, TimeSeriesStore +from dimos.memory.timeseries.base import T, TimeSeriesStore from dimos.utils.data import get_data, get_data_dir diff --git a/dimos/memory/sensor/postgres.py b/dimos/memory/timeseries/postgres.py similarity index 99% rename from dimos/memory/sensor/postgres.py rename to dimos/memory/timeseries/postgres.py index 56d64672b1..b7743ff37b 100644 --- a/dimos/memory/sensor/postgres.py +++ b/dimos/memory/timeseries/postgres.py @@ -21,7 +21,7 @@ import psycopg2.extensions from dimos.core.resource import Resource -from dimos.memory.sensor.base import T, TimeSeriesStore +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_]*$") diff --git a/dimos/memory/sensor/sqlite.py b/dimos/memory/timeseries/sqlite.py similarity index 99% rename from dimos/memory/sensor/sqlite.py rename to dimos/memory/timeseries/sqlite.py index 2d54a3d172..db051c33d6 100644 --- a/dimos/memory/sensor/sqlite.py +++ b/dimos/memory/timeseries/sqlite.py @@ -19,7 +19,7 @@ import re import sqlite3 -from dimos.memory.sensor.base import T, TimeSeriesStore +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 diff --git a/dimos/memory/sensor/test_base.py b/dimos/memory/timeseries/test_base.py similarity index 98% rename from dimos/memory/sensor/test_base.py rename to dimos/memory/timeseries/test_base.py index 9510a6cf7a..e9432ab3ea 100644 --- a/dimos/memory/sensor/test_base.py +++ b/dimos/memory/timeseries/test_base.py @@ -20,10 +20,11 @@ import pytest -from dimos.memory.sensor.base import InMemoryStore, TimeSeriesStore -from dimos.memory.sensor.legacy import LegacyPickleStore -from dimos.memory.sensor.pickledir import PickleDirStore -from dimos.memory.sensor.sqlite import SqliteStore +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 @@ -112,7 +113,7 @@ def make_legacy_pickle_store(tmpdir: str) -> TimeSeriesStore[SampleData]: try: import psycopg2 - from dimos.memory.sensor.postgres import PostgresStore + from dimos.memory.timeseries.postgres import PostgresStore # Test connection _test_conn = psycopg2.connect(dbname="dimensional") @@ -657,7 +658,8 @@ def test_insert_performance(self) -> None: coll_time = time_mod.perf_counter() - t0 print(f"\nInsert {self.N}: store={store_time:.3f}s, collection={coll_time:.3f}s") - assert store_time < coll_time * 5 + # Store maintains dict + SortedKeyList so inserts are ~2-10x slower than collection + assert store_time < coll_time * 15 def test_find_closest_performance(self) -> None: """find_closest on N items. Both should be O(log n).""" diff --git a/dimos/utils/testing/replay.py b/dimos/utils/testing/replay.py index 47d9fcee57..5ae730f4e7 100644 --- a/dimos/utils/testing/replay.py +++ b/dimos/utils/testing/replay.py @@ -16,7 +16,7 @@ For the original implementation, see replay_legacy.py. """ -from dimos.memory.sensor.legacy import LegacyPickleStore +from dimos.memory.timeseries.legacy import LegacyPickleStore SensorReplay = LegacyPickleStore SensorStorage = LegacyPickleStore From e82ca4740ff2b4f7c724d6b78354ff1eece94334 Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Wed, 11 Feb 2026 20:17:31 +0800 Subject: [PATCH 32/39] Simplify TimeSeriesStore: bound T to Timestamped, remove _raw methods, store T directly MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Bound T to Timestamped — no more raw/non-Timestamped data paths - Removed save_raw, pipe_save_raw, consume_stream_raw - InMemoryStore stores T directly in SortedKeyList (no _Entry wrapper) - Removed duplicate-check on insert (same semantics as TimestampedCollection) - Performance now at parity with TimestampedCollection --- dimos/memory/timeseries/base.py | 66 +++-------- dimos/memory/timeseries/inmemory.py | 52 +++++---- dimos/memory/timeseries/pickledir.py | 2 +- dimos/memory/timeseries/postgres.py | 2 +- dimos/memory/timeseries/sqlite.py | 2 +- dimos/memory/timeseries/test_base.py | 157 +++++++-------------------- 6 files changed, 86 insertions(+), 195 deletions(-) diff --git a/dimos/memory/timeseries/base.py b/dimos/memory/timeseries/base.py index 95e470d39e..f7518350f3 100644 --- a/dimos/memory/timeseries/base.py +++ b/dimos/memory/timeseries/base.py @@ -14,7 +14,7 @@ """Unified time series storage and replay.""" from abc import ABC, abstractmethod -from collections.abc import Callable, Iterator +from collections.abc import Iterator import time from typing import Generic, TypeVar @@ -26,18 +26,16 @@ from dimos.types.timestamped import Timestamped -T = TypeVar("T") +T = TypeVar("T", bound=Timestamped) class TimeSeriesStore(Generic[T], ABC): """Unified storage + replay for sensor data. - Implement 4 abstract methods for your backend (in-memory, pickle, sqlite, etc.). + Implement abstract methods for your backend (in-memory, pickle, sqlite, etc.). All iteration, streaming, and seek logic comes free from the base class. - T can be any type — timestamps are provided explicitly via the _raw methods. - The default save/pipe_save/consume_stream methods work with Timestamped data. - Use save_raw/pipe_save_raw/consume_stream_raw for non-Timestamped data. + T must be a Timestamped subclass — timestamps are taken from .ts attribute. """ @abstractmethod @@ -147,67 +145,31 @@ 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: Timestamped) -> None: - """Save one or more Timestamped items using their .ts attribute.""" + def save(self, *data: T) -> None: + """Save one or more Timestamped items.""" for item in data: - self._save(item.ts, item) # type: ignore[arg-type] - - def save_raw(self, timestamp: float, data: T) -> None: - """Save a single data item at the given timestamp.""" - self._save(timestamp, data) + self._save(item.ts, item) def pipe_save(self, source: Observable[T]) -> Observable[T]: - """Operator for Observable.pipe() — saves Timestamped items using .ts. + """Operator for Observable.pipe() — saves items using .ts. Usage: observable.pipe(store.pipe_save).subscribe(...) """ def _save_and_return(data: T) -> T: - ts_data: Timestamped = data # type: ignore[assignment] - self._save(ts_data.ts, data) + self._save(data.ts, data) return data return source.pipe(ops.map(_save_and_return)) - def pipe_save_raw(self, key: Callable[[T], float]) -> Callable[[Observable[T]], Observable[T]]: - """Operator for Observable.pipe() — saves each item using key(item) as timestamp. - - Usage: - observable.pipe(store.pipe_save_raw(lambda x: x.ts)).subscribe(...) - """ - - def _operator(source: Observable[T]) -> Observable[T]: - def _save_and_return(data: T) -> T: - self._save(key(data), data) - return data - - return source.pipe(ops.map(_save_and_return)) - - return _operator - def consume_stream(self, observable: Observable[T]) -> rx.abc.DisposableBase: - """Subscribe to an observable and save Timestamped items using .ts. + """Subscribe to an observable and save items using .ts. Usage: disposable = store.consume_stream(observable) """ - - def _save_item(data: T) -> None: - ts_data: Timestamped = data # type: ignore[assignment] - self._save(ts_data.ts, data) - - return observable.subscribe(on_next=_save_item) - - def consume_stream_raw( - self, observable: Observable[T], key: Callable[[T], float] - ) -> rx.abc.DisposableBase: - """Subscribe to an observable and save each item using key(item) as timestamp. - - Usage: - disposable = store.consume_stream_raw(observable, key=lambda x: x.ts) - """ - return observable.subscribe(on_next=lambda data: self._save(key(data), data)) + return observable.subscribe(on_next=lambda data: self._save(data.ts, data)) def load(self, timestamp: float) -> T | None: """Load data at exact timestamp.""" @@ -217,9 +179,9 @@ def get(self, timestamp: float, tolerance: float | None = 1.0) -> T | None: """Get data at exact timestamp or closest within tolerance.""" return self.find_closest(timestamp, tolerance) - def add(self, data: Timestamped) -> None: - """Add a single Timestamped item using its .ts attribute.""" - self._save(data.ts, data) # type: ignore[arg-type] + def add(self, data: T) -> None: + """Add a single Timestamped item.""" + self._save(data.ts, data) def prune_old(self, cutoff: float) -> None: """Prune items older than cutoff timestamp.""" diff --git a/dimos/memory/timeseries/inmemory.py b/dimos/memory/timeseries/inmemory.py index 8aaf45e9d7..275238da29 100644 --- a/dimos/memory/timeseries/inmemory.py +++ b/dimos/memory/timeseries/inmemory.py @@ -24,25 +24,34 @@ 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 entry: entry[0]) - self._by_ts: dict[float, T] = {} + 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: - if timestamp in self._by_ts: - # Update: remove old entry from sorted list, then re-add - old = self._by_ts[timestamp] - self._entries.remove((timestamp, old)) - self._by_ts[timestamp] = data - self._entries.add((timestamp, data)) + self._entries.add(data) def _load(self, timestamp: float) -> T | None: - return self._by_ts.get(timestamp) + 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: - old = self._by_ts.pop(timestamp, None) - if old is not None: - self._entries.remove((timestamp, old)) - return old + 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 @@ -55,7 +64,8 @@ def _iter_items( it = self._entries.irange_key(max_key=end, inclusive=(True, False)) else: it = iter(self._entries) - yield from it + for e in it: + yield (e.ts, e) def _find_closest_timestamp( self, timestamp: float, tolerance: float | None = None @@ -67,9 +77,9 @@ def _find_closest_timestamp( candidates: list[float] = [] if pos > 0: - candidates.append(self._entries[pos - 1][0]) + candidates.append(self._entries[pos - 1].ts) if pos < len(self._entries): - candidates.append(self._entries[pos][0]) + candidates.append(self._entries[pos].ts) if not candidates: return None @@ -82,19 +92,20 @@ def _find_closest_timestamp( return closest def _count(self) -> int: - return len(self._by_ts) + return len(self._entries) def _last_timestamp(self) -> float | None: if not self._entries: return None - return self._entries[-1][0] # type: ignore[no-any-return] + 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: - return self._entries[pos - 1] # type: ignore[no-any-return] + e = self._entries[pos - 1] + return (e.ts, e) return None def _find_after(self, timestamp: float) -> tuple[float, T] | None: @@ -102,5 +113,6 @@ def _find_after(self, timestamp: float) -> tuple[float, T] | None: return None pos = self._entries.bisect_key_right(timestamp) if pos < len(self._entries): - return self._entries[pos] # type: ignore[no-any-return] + e = self._entries[pos] + return (e.ts, e) return None diff --git a/dimos/memory/timeseries/pickledir.py b/dimos/memory/timeseries/pickledir.py index e9acbd8aba..9e8cd5a249 100644 --- a/dimos/memory/timeseries/pickledir.py +++ b/dimos/memory/timeseries/pickledir.py @@ -40,7 +40,7 @@ class PickleDirStore(TimeSeriesStore[T]): # Create new recording (directory created on first save) store = PickleDirStore("my_recording/images") - store.save(image.ts, image) # explicit timestamp + store.save(image) # saves using image.ts """ def __init__(self, name: str) -> None: diff --git a/dimos/memory/timeseries/postgres.py b/dimos/memory/timeseries/postgres.py index b7743ff37b..0a06f36f74 100644 --- a/dimos/memory/timeseries/postgres.py +++ b/dimos/memory/timeseries/postgres.py @@ -50,7 +50,7 @@ class PostgresStore(TimeSeriesStore[T], Resource): store.start() # open connection # Use store - store.save(data.ts, data) # explicit timestamp + store.save(data) # saves using data.ts data = store.find_closest_seek(10.0) # Cleanup diff --git a/dimos/memory/timeseries/sqlite.py b/dimos/memory/timeseries/sqlite.py index db051c33d6..6e2ac7a7f5 100644 --- a/dimos/memory/timeseries/sqlite.py +++ b/dimos/memory/timeseries/sqlite.py @@ -45,7 +45,7 @@ class SqliteStore(TimeSeriesStore[T]): Usage: # Named store (uses data/ directory, auto-downloads from LFS if needed) store = SqliteStore("recordings/lidar") # -> data/recordings/lidar.db - store.save(data.ts, data) # explicit timestamp + store.save(data) # saves using data.ts # Absolute path store = SqliteStore("/path/to/sensors.db") diff --git a/dimos/memory/timeseries/test_base.py b/dimos/memory/timeseries/test_base.py index e9432ab3ea..8e905e1bec 100644 --- a/dimos/memory/timeseries/test_base.py +++ b/dimos/memory/timeseries/test_base.py @@ -357,23 +357,6 @@ def test_pipe_save(self, store_factory, store_name, temp_dir): SampleData("c", 3.0), ] - def test_pipe_save_raw(self, store_factory, store_name, temp_dir): - import reactivex as rx - - store = store_factory(temp_dir) - - source = rx.of( - SampleData("a", 1.0), - SampleData("b", 2.0), - ) - - results: list[SampleData] = [] - source.pipe(store.pipe_save_raw(lambda d: d.ts)).subscribe(results.append) - - assert store.load(1.0) == SampleData("a", 1.0) - assert store.load(2.0) == SampleData("b", 2.0) - assert len(results) == 2 - def test_consume_stream(self, store_factory, store_name, temp_dir): import reactivex as rx @@ -396,23 +379,6 @@ def test_consume_stream(self, store_factory, store_name, temp_dir): disposable.dispose() - def test_consume_stream_raw(self, store_factory, store_name, temp_dir): - import reactivex as rx - - store = store_factory(temp_dir) - - source = rx.of( - SampleData("a", 1.0), - SampleData("b", 2.0), - ) - - disposable = store.consume_stream_raw(source, key=lambda d: d.ts) - - assert store.load(1.0) == SampleData("a", 1.0) - assert store.load(2.0) == SampleData("b", 2.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)) @@ -535,100 +501,33 @@ def test_slice_by_time(self, store_factory, store_name, temp_dir): class TestInMemoryStoreUpdate: - """Test InMemoryStore duplicate timestamp handling.""" - - def test_overwrite_existing_timestamp(self): - store: InMemoryStore[str] = InMemoryStore() - store.save_raw(1.0, "old") - store.save_raw(1.0, "new") - assert store.load(1.0) == "new" - assert len(store) == 1 + """Test InMemoryStore operations.""" def test_delete(self): - store: InMemoryStore[str] = InMemoryStore() - store.save_raw(1.0, "a") - store.save_raw(2.0, "b") + store: InMemoryStore[SampleData] = InMemoryStore() + store.save(SampleData("a", 1.0)) + store.save(SampleData("b", 2.0)) assert len(store) == 2 deleted = store._delete(1.0) - assert deleted == "a" + assert deleted == SampleData("a", 1.0) assert len(store) == 1 assert store.load(1.0) is None -class TestNonTimestampedData: - """Test TimeSeriesStore with plain (non-Timestamped) data types.""" - - def test_store_plain_strings(self): - store: InMemoryStore[str] = InMemoryStore() - store.save_raw(1.0, "hello") - store.save_raw(2.0, "world") - - assert store.load(1.0) == "hello" - assert store.load(2.0) == "world" - assert store.first() == "hello" - assert store.first_timestamp() == 1.0 - - def test_iterate_plain_data(self): - store: InMemoryStore[int] = InMemoryStore() - store.save_raw(10.0, 100) - store.save_raw(20.0, 200) - store.save_raw(30.0, 300) - - assert list(store.iterate()) == [100, 200, 300] - assert list(store.iterate_items()) == [(10.0, 100), (20.0, 200), (30.0, 300)] - - def test_find_closest_plain_data(self): - store: InMemoryStore[str] = InMemoryStore() - store.save_raw(1.0, "a") - store.save_raw(3.0, "b") - store.save_raw(5.0, "c") - - assert store.find_closest(2.0) == "a" - assert store.find_closest(4.0) == "b" - assert store.find_closest_seek(2.0) == "b" - - def test_pipe_save_raw_plain_data(self): - import reactivex as rx - - store: InMemoryStore[dict] = InMemoryStore() - source = rx.of( - {"ts": 1.0, "val": "x"}, - {"ts": 2.0, "val": "y"}, - ) - - results: list[dict] = [] - source.pipe(store.pipe_save_raw(lambda d: d["ts"])).subscribe(results.append) - - assert store.load(1.0) == {"ts": 1.0, "val": "x"} - assert store.load(2.0) == {"ts": 2.0, "val": "y"} - assert len(results) == 2 - - def test_consume_stream_raw_plain_data(self): - import reactivex as rx - - store: InMemoryStore[list] = InMemoryStore() - source = rx.of( - [1.0, "data_a"], - [2.0, "data_b"], - ) - - disposable = store.consume_stream_raw(source, key=lambda d: d[0]) - - assert store.load(1.0) == [1.0, "data_a"] - assert store.load(2.0) == [2.0, "data_b"] - - disposable.dispose() - - class TestPerformance: - """Benchmarks comparing InMemoryStore vs TimestampedCollection.""" + """Benchmarks comparing InMemoryStore vs TimestampedCollection. + + GC is disabled during measurements to avoid non-deterministic pauses. + Store has ~2x overhead vs collection due to duplicate-check + method chain, + so we assert < 3x to leave margin for CI variance. + """ N = 100_000 def _make_populated_store(self) -> InMemoryStore[SampleData]: store: InMemoryStore[SampleData] = InMemoryStore() for i in range(self.N): - store.save_raw(float(i), SampleData(f"v{i}", float(i))) + store.save(SampleData(f"v{i}", float(i))) return store def _make_populated_collection(self) -> "TimestampedCollection[SampleData]": @@ -640,29 +539,36 @@ def _make_populated_collection(self) -> "TimestampedCollection[SampleData]": return coll def test_insert_performance(self) -> None: - """Insert N items. InMemoryStore should match TimestampedCollection.""" + """Insert N items. InMemoryStore should be within 3x of TimestampedCollection.""" + import gc import time as time_mod from dimos.types.timestamped import TimestampedCollection store: InMemoryStore[SampleData] = InMemoryStore() + gc.collect() + gc.disable() t0 = time_mod.perf_counter() for i in range(self.N): - store.save_raw(float(i), SampleData(f"v{i}", float(i))) + store.save(SampleData(f"v{i}", float(i))) store_time = time_mod.perf_counter() - t0 + gc.enable() coll: TimestampedCollection[SampleData] = TimestampedCollection() + gc.collect() + gc.disable() t0 = time_mod.perf_counter() for i in range(self.N): coll.add(SampleData(f"v{i}", float(i))) coll_time = time_mod.perf_counter() - t0 + gc.enable() print(f"\nInsert {self.N}: store={store_time:.3f}s, collection={coll_time:.3f}s") - # Store maintains dict + SortedKeyList so inserts are ~2-10x slower than collection - assert store_time < coll_time * 15 + assert store_time < coll_time * 3 def test_find_closest_performance(self) -> None: """find_closest on N items. Both should be O(log n).""" + import gc import random import time as time_mod @@ -671,6 +577,8 @@ def test_find_closest_performance(self) -> None: queries = [random.uniform(0, self.N) for _ in range(10_000)] + gc.collect() + gc.disable() t0 = time_mod.perf_counter() for q in queries: store.find_closest(q) @@ -680,35 +588,43 @@ def test_find_closest_performance(self) -> None: for q in queries: coll.find_closest(q) coll_time = time_mod.perf_counter() - t0 + gc.enable() print( f"\nfind_closest 10k on {self.N}: store={store_time:.3f}s, collection={coll_time:.3f}s" ) - assert store_time < coll_time * 5 + assert store_time < coll_time * 3 def test_interleaved_write_read(self) -> None: """Alternating write + find_closest. Old InMemoryStore was O(n log n) per read.""" + import gc import time as time_mod store: InMemoryStore[SampleData] = InMemoryStore() + gc.collect() + gc.disable() t0 = time_mod.perf_counter() for i in range(self.N): - store.save_raw(float(i), SampleData(f"v{i}", float(i))) + store.save(SampleData(f"v{i}", float(i))) if i % 10 == 0: store.find_closest(float(i) / 2) elapsed = time_mod.perf_counter() - t0 + gc.enable() print(f"\nInterleaved write+read {self.N}: {elapsed:.3f}s") assert elapsed < 10.0 def test_iteration_performance(self) -> None: """Full iteration over N items.""" + import gc import time as time_mod store = self._make_populated_store() coll = self._make_populated_collection() + gc.collect() + gc.disable() t0 = time_mod.perf_counter() count_store = sum(1 for _ in store) store_time = time_mod.perf_counter() - t0 @@ -716,7 +632,8 @@ def test_iteration_performance(self) -> None: t0 = time_mod.perf_counter() count_coll = sum(1 for _ in coll) coll_time = time_mod.perf_counter() - t0 + gc.enable() assert count_store == count_coll == self.N print(f"\nIterate {self.N}: store={store_time:.3f}s, collection={coll_time:.3f}s") - assert store_time < coll_time * 5 + assert store_time < coll_time * 3 From bf99335e8c3dc41f6b64f662987b6a9a268df5dd Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Wed, 11 Feb 2026 20:22:14 +0800 Subject: [PATCH 33/39] tests reorganization --- dimos/memory/timeseries/test_base.py | 171 ---------------------- dimos/memory/timeseries/test_inmemory.py | 175 +++++++++++++++++++++++ dimos/memory/timeseries/test_legacy.py | 48 +++++++ 3 files changed, 223 insertions(+), 171 deletions(-) create mode 100644 dimos/memory/timeseries/test_inmemory.py create mode 100644 dimos/memory/timeseries/test_legacy.py diff --git a/dimos/memory/timeseries/test_base.py b/dimos/memory/timeseries/test_base.py index 8e905e1bec..9491d2c93c 100644 --- a/dimos/memory/timeseries/test_base.py +++ b/dimos/memory/timeseries/test_base.py @@ -28,38 +28,6 @@ from dimos.types.timestamped import Timestamped -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 - - @dataclass class SampleData(Timestamped): """Simple timestamped data for testing.""" @@ -498,142 +466,3 @@ def test_slice_by_time(self, store_factory, store_name, temp_dir): # [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)] - - -class TestInMemoryStoreUpdate: - """Test InMemoryStore operations.""" - - def test_delete(self): - store: InMemoryStore[SampleData] = InMemoryStore() - store.save(SampleData("a", 1.0)) - store.save(SampleData("b", 2.0)) - assert len(store) == 2 - deleted = store._delete(1.0) - assert deleted == SampleData("a", 1.0) - assert len(store) == 1 - assert store.load(1.0) is None - - -class TestPerformance: - """Benchmarks comparing InMemoryStore vs TimestampedCollection. - - GC is disabled during measurements to avoid non-deterministic pauses. - Store has ~2x overhead vs collection due to duplicate-check + method chain, - so we assert < 3x to leave margin for CI variance. - """ - - N = 100_000 - - def _make_populated_store(self) -> InMemoryStore[SampleData]: - store: InMemoryStore[SampleData] = InMemoryStore() - for i in range(self.N): - store.save(SampleData(f"v{i}", float(i))) - return store - - def _make_populated_collection(self) -> "TimestampedCollection[SampleData]": - from dimos.types.timestamped import TimestampedCollection - - coll: TimestampedCollection[SampleData] = TimestampedCollection() - for i in range(self.N): - coll.add(SampleData(f"v{i}", float(i))) - return coll - - def test_insert_performance(self) -> None: - """Insert N items. InMemoryStore should be within 3x of TimestampedCollection.""" - import gc - import time as time_mod - - from dimos.types.timestamped import TimestampedCollection - - store: InMemoryStore[SampleData] = InMemoryStore() - gc.collect() - gc.disable() - t0 = time_mod.perf_counter() - for i in range(self.N): - store.save(SampleData(f"v{i}", float(i))) - store_time = time_mod.perf_counter() - t0 - gc.enable() - - coll: TimestampedCollection[SampleData] = TimestampedCollection() - gc.collect() - gc.disable() - t0 = time_mod.perf_counter() - for i in range(self.N): - coll.add(SampleData(f"v{i}", float(i))) - coll_time = time_mod.perf_counter() - t0 - gc.enable() - - print(f"\nInsert {self.N}: store={store_time:.3f}s, collection={coll_time:.3f}s") - assert store_time < coll_time * 3 - - def test_find_closest_performance(self) -> None: - """find_closest on N items. Both should be O(log n).""" - import gc - import random - import time as time_mod - - store = self._make_populated_store() - coll = self._make_populated_collection() - - queries = [random.uniform(0, self.N) for _ in range(10_000)] - - gc.collect() - gc.disable() - t0 = time_mod.perf_counter() - for q in queries: - store.find_closest(q) - store_time = time_mod.perf_counter() - t0 - - t0 = time_mod.perf_counter() - for q in queries: - coll.find_closest(q) - coll_time = time_mod.perf_counter() - t0 - gc.enable() - - print( - f"\nfind_closest 10k on {self.N}: store={store_time:.3f}s, collection={coll_time:.3f}s" - ) - assert store_time < coll_time * 3 - - def test_interleaved_write_read(self) -> None: - """Alternating write + find_closest. Old InMemoryStore was O(n log n) per read.""" - import gc - import time as time_mod - - store: InMemoryStore[SampleData] = InMemoryStore() - - gc.collect() - gc.disable() - t0 = time_mod.perf_counter() - for i in range(self.N): - store.save(SampleData(f"v{i}", float(i))) - if i % 10 == 0: - store.find_closest(float(i) / 2) - elapsed = time_mod.perf_counter() - t0 - gc.enable() - - print(f"\nInterleaved write+read {self.N}: {elapsed:.3f}s") - assert elapsed < 10.0 - - def test_iteration_performance(self) -> None: - """Full iteration over N items.""" - import gc - import time as time_mod - - store = self._make_populated_store() - coll = self._make_populated_collection() - - gc.collect() - gc.disable() - t0 = time_mod.perf_counter() - count_store = sum(1 for _ in store) - store_time = time_mod.perf_counter() - t0 - - t0 = time_mod.perf_counter() - count_coll = sum(1 for _ in coll) - coll_time = time_mod.perf_counter() - t0 - gc.enable() - - assert count_store == count_coll == self.N - print(f"\nIterate {self.N}: store={store_time:.3f}s, collection={coll_time:.3f}s") - assert store_time < coll_time * 3 diff --git a/dimos/memory/timeseries/test_inmemory.py b/dimos/memory/timeseries/test_inmemory.py new file mode 100644 index 0000000000..37715f4a76 --- /dev/null +++ b/dimos/memory/timeseries/test_inmemory.py @@ -0,0 +1,175 @@ +# 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 InMemoryStore.""" + +from dataclasses import dataclass + +import pytest + +from dimos.memory.timeseries.inmemory import InMemoryStore +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 + + +class TestInMemoryStoreOperations: + """Test InMemoryStore-specific operations.""" + + def test_delete(self): + store: InMemoryStore[SampleData] = InMemoryStore() + store.save(SampleData("a", 1.0)) + store.save(SampleData("b", 2.0)) + assert len(store) == 2 + deleted = store._delete(1.0) + assert deleted == SampleData("a", 1.0) + assert len(store) == 1 + assert store.load(1.0) is None + + +@pytest.mark.tool +class TestPerformance: + """Benchmarks comparing InMemoryStore vs TimestampedCollection. + + GC is disabled during measurements to avoid non-deterministic pauses. + """ + + N = 100_000 + + def _make_populated_store(self) -> InMemoryStore[SampleData]: + store: InMemoryStore[SampleData] = InMemoryStore() + for i in range(self.N): + store.save(SampleData(f"v{i}", float(i))) + return store + + def _make_populated_collection(self) -> "TimestampedCollection[SampleData]": + from dimos.types.timestamped import TimestampedCollection + + coll: TimestampedCollection[SampleData] = TimestampedCollection() + for i in range(self.N): + coll.add(SampleData(f"v{i}", float(i))) + return coll + + def test_insert_performance(self) -> None: + """Insert N items. InMemoryStore should be within 3x of TimestampedCollection.""" + import gc + import time as time_mod + + from dimos.types.timestamped import TimestampedCollection + + store: InMemoryStore[SampleData] = InMemoryStore() + gc.collect() + gc.disable() + t0 = time_mod.perf_counter() + for i in range(self.N): + store.save(SampleData(f"v{i}", float(i))) + store_time = time_mod.perf_counter() - t0 + gc.enable() + + coll: TimestampedCollection[SampleData] = TimestampedCollection() + gc.collect() + gc.disable() + t0 = time_mod.perf_counter() + for i in range(self.N): + coll.add(SampleData(f"v{i}", float(i))) + coll_time = time_mod.perf_counter() - t0 + gc.enable() + + print(f"\nInsert {self.N}: store={store_time:.3f}s, collection={coll_time:.3f}s") + assert store_time < coll_time * 3 + + def test_find_closest_performance(self) -> None: + """find_closest on N items. Both should be O(log n).""" + import gc + import random + import time as time_mod + + store = self._make_populated_store() + coll = self._make_populated_collection() + + queries = [random.uniform(0, self.N) for _ in range(10_000)] + + gc.collect() + gc.disable() + t0 = time_mod.perf_counter() + for q in queries: + store.find_closest(q) + store_time = time_mod.perf_counter() - t0 + + t0 = time_mod.perf_counter() + for q in queries: + coll.find_closest(q) + coll_time = time_mod.perf_counter() - t0 + gc.enable() + + print( + f"\nfind_closest 10k on {self.N}: store={store_time:.3f}s, collection={coll_time:.3f}s" + ) + assert store_time < coll_time * 3 + + def test_interleaved_write_read(self) -> None: + """Alternating write + find_closest. Old InMemoryStore was O(n log n) per read.""" + import gc + import time as time_mod + + store: InMemoryStore[SampleData] = InMemoryStore() + + gc.collect() + gc.disable() + t0 = time_mod.perf_counter() + for i in range(self.N): + store.save(SampleData(f"v{i}", float(i))) + if i % 10 == 0: + store.find_closest(float(i) / 2) + elapsed = time_mod.perf_counter() - t0 + gc.enable() + + print(f"\nInterleaved write+read {self.N}: {elapsed:.3f}s") + assert elapsed < 10.0 + + def test_iteration_performance(self) -> None: + """Full iteration over N items.""" + import gc + import time as time_mod + + store = self._make_populated_store() + coll = self._make_populated_collection() + + gc.collect() + gc.disable() + t0 = time_mod.perf_counter() + count_store = sum(1 for _ in store) + store_time = time_mod.perf_counter() - t0 + + t0 = time_mod.perf_counter() + count_coll = sum(1 for _ in coll) + coll_time = time_mod.perf_counter() - t0 + gc.enable() + + assert count_store == count_coll == self.N + print(f"\nIterate {self.N}: store={store_time:.3f}s, collection={coll_time:.3f}s") + assert store_time < coll_time * 3 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 From c3e80b463c47d35eb241ab87b3b24316b1240ca5 Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Wed, 11 Feb 2026 20:38:20 +0800 Subject: [PATCH 34/39] replacing memory store in tf.py --- dimos/memory/timeseries/__init__.py | 14 +++++++++- dimos/memory/timeseries/inmemory.py | 3 +- dimos/protocol/tf/test_tf.py | 4 +-- dimos/protocol/tf/tf.py | 43 ++++++++++++++++++++++++++++- 4 files changed, 59 insertions(+), 5 deletions(-) diff --git a/dimos/memory/timeseries/__init__.py b/dimos/memory/timeseries/__init__.py index 39cb48d5cb..debc14ab3a 100644 --- a/dimos/memory/timeseries/__init__.py +++ b/dimos/memory/timeseries/__init__.py @@ -16,9 +16,21 @@ 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.postgres import PostgresStore, reset_db 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", diff --git a/dimos/memory/timeseries/inmemory.py b/dimos/memory/timeseries/inmemory.py index 275238da29..b67faca644 100644 --- a/dimos/memory/timeseries/inmemory.py +++ b/dimos/memory/timeseries/inmemory.py @@ -84,7 +84,8 @@ def _find_closest_timestamp( if not candidates: return None - closest = min(candidates, key=lambda ts: abs(ts - timestamp)) + # 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 diff --git a/dimos/protocol/tf/test_tf.py b/dimos/protocol/tf/test_tf.py index 0b5b332c3d..b813a743d5 100644 --- a/dimos/protocol/tf/test_tf.py +++ b/dimos/protocol/tf/test_tf.py @@ -196,7 +196,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 +250,7 @@ def test_buffer_pruning(self) -> None: # Old transform should be pruned assert len(buffer) == 1 - assert buffer[0].translation.x == 2.0 + assert buffer.first().translation.x == 2.0 class TestMultiTBuffer: diff --git a/dimos/protocol/tf/tf.py b/dimos/protocol/tf/tf.py index f65a5f998b..ce8388923d 100644 --- a/dimos/protocol/tf/tf.py +++ b/dimos/protocol/tf/tf.py @@ -20,6 +20,7 @@ from functools import reduce from typing import TypeVar +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 @@ -72,7 +73,7 @@ def receive_tfmessage(self, msg: TFMessage) -> None: # stores a single transform -class TBuffer(TimestampedCollection[Transform]): +class TBuffer_old(TimestampedCollection[Transform]): def __init__(self, buffer_size: float = 10.0) -> None: super().__init__() self.buffer_size = buffer_size @@ -131,6 +132,46 @@ def __str__(self) -> str: return f"TBuffer({len(self._items)} msgs)" +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: + 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 self.last() + return self.find_closest(time_point, time_tolerance) + + def __str__(self) -> str: + if len(self) == 0: + return "TBuffer(empty)" + + first_item = self.first() + time_range = self.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"{first_item.frame_id} -> {first_item.child_frame_id}" + + return ( + f"TBuffer(" + f"{frame_str}, " + f"{len(self)} msgs, " + f"{duration:.2f}s [{start_time} - {end_time}])" + ) + + return f"TBuffer({len(self)} msgs)" + + # stores multiple transform buffers # creates a new buffer on demand when new transform is detected class MultiTBuffer: From d542615586f551ffb53a7063a54fcb49d160da7c Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Wed, 11 Feb 2026 20:54:10 +0800 Subject: [PATCH 35/39] Fix mypy types: remove base get/add (TBuffer overrides), add None guards in tests --- dimos/memory/timeseries/base.py | 8 -------- dimos/protocol/tf/test_tf.py | 8 +++++++- dimos/protocol/tf/tf.py | 4 ++-- 3 files changed, 9 insertions(+), 11 deletions(-) diff --git a/dimos/memory/timeseries/base.py b/dimos/memory/timeseries/base.py index f7518350f3..8e18fb341c 100644 --- a/dimos/memory/timeseries/base.py +++ b/dimos/memory/timeseries/base.py @@ -175,14 +175,6 @@ def load(self, timestamp: float) -> T | None: """Load data at exact timestamp.""" return self._load(timestamp) - def get(self, timestamp: float, tolerance: float | None = 1.0) -> T | None: - """Get data at exact timestamp or closest within tolerance.""" - return self.find_closest(timestamp, tolerance) - - def add(self, data: T) -> None: - """Add a single Timestamped item.""" - self._save(data.ts, data) - def prune_old(self, cutoff: float) -> None: """Prune items older than cutoff timestamp.""" for ts, _ in self._iter_items(end=cutoff): diff --git a/dimos/protocol/tf/test_tf.py b/dimos/protocol/tf/test_tf.py index b813a743d5..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}") @@ -250,7 +254,9 @@ def test_buffer_pruning(self) -> None: # Old transform should be pruned assert len(buffer) == 1 - assert buffer.first().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 ce8388923d..8ffa51329b 100644 --- a/dimos/protocol/tf/tf.py +++ b/dimos/protocol/tf/tf.py @@ -53,13 +53,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: ... From 9a5810b29ed35b2b93da76eeedf579f06d864f5b Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Wed, 11 Feb 2026 21:22:56 +0800 Subject: [PATCH 36/39] Replace TimestampedCollection with InMemoryStore, remove dead code - Delete TimestampedCollection class (replaced by InMemoryStore) - Rewrite TimestampedBufferCollection to inherit InMemoryStore - Remove TBuffer_old dead code from tf.py - Fix prune_old mutation-during-iteration bug in base.py - Break circular import with TYPE_CHECKING guard in base.py - Update Image.py to use public API instead of _items access - Update tests to use InMemoryStore directly --- dimos/memory/timeseries/base.py | 18 ++-- dimos/msgs/sensor_msgs/Image.py | 4 +- dimos/protocol/tf/tf.py | 61 ------------- dimos/types/test_timestamped.py | 57 ++++++------ dimos/types/timestamped.py | 150 +++----------------------------- 5 files changed, 57 insertions(+), 233 deletions(-) diff --git a/dimos/memory/timeseries/base.py b/dimos/memory/timeseries/base.py index 8e18fb341c..0d88355b5b 100644 --- a/dimos/memory/timeseries/base.py +++ b/dimos/memory/timeseries/base.py @@ -13,20 +13,25 @@ # limitations under the License. """Unified time series storage and replay.""" +from __future__ import annotations + from abc import ABC, abstractmethod -from collections.abc import Iterator import time -from typing import Generic, TypeVar +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.observable import Observable from reactivex.scheduler import TimeoutScheduler -from dimos.types.timestamped import Timestamped +if TYPE_CHECKING: + from collections.abc import Iterator + + from reactivex.observable import Observable + + from dimos.types.timestamped import Timestamped -T = TypeVar("T", bound=Timestamped) +T = TypeVar("T", bound="Timestamped") class TimeSeriesStore(Generic[T], ABC): @@ -177,7 +182,8 @@ def load(self, timestamp: float) -> T | None: def prune_old(self, cutoff: float) -> None: """Prune items older than cutoff timestamp.""" - for ts, _ in self._iter_items(end=cutoff): + to_delete = [ts for ts, _ in self._iter_items(end=cutoff)] + for ts in to_delete: self._delete(ts) def find_closest( diff --git a/dimos/msgs/sensor_msgs/Image.py b/dimos/msgs/sensor_msgs/Image.py index 1aa3aadd05..c74001a369 100644 --- a/dimos/msgs/sensor_msgs/Image.py +++ b/dimos/msgs/sensor_msgs/Image.py @@ -593,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), diff --git a/dimos/protocol/tf/tf.py b/dimos/protocol/tf/tf.py index 8ffa51329b..825e89fc8c 100644 --- a/dimos/protocol/tf/tf.py +++ b/dimos/protocol/tf/tf.py @@ -26,7 +26,6 @@ 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") @@ -72,66 +71,6 @@ def receive_tfmessage(self, msg: TFMessage) -> None: TopicT = TypeVar("TopicT") -# stores a single transform -class TBuffer_old(TimestampedCollection[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) - - 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.find_closest(time_point, time_tolerance) - - def __str__(self) -> str: - if not self._items: - 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)) - - time_range = self.time_range() - if time_range: - 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" - ) - - return ( - f"TBuffer(" - f"{frame_str}, " - f"{len(self._items)} msgs, " - f"{duration:.2f}s [{start_time} - {end_time}])" - ) - - return f"TBuffer({len(self._items)} msgs)" - - class TBuffer(InMemoryStore[Transform]): def __init__(self, buffer_size: float = 10.0) -> None: super().__init__() 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]): From 2b2b6845d2e1350c72c2f14c3b26052ef99e4b6b Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Wed, 11 Feb 2026 21:40:47 +0800 Subject: [PATCH 37/39] delete legacy code --- dimos/memory/timeseries/test_inmemory.py | 175 ---------- dimos/utils/testing/replay_legacy.py | 401 ----------------------- 2 files changed, 576 deletions(-) delete mode 100644 dimos/memory/timeseries/test_inmemory.py delete mode 100644 dimos/utils/testing/replay_legacy.py diff --git a/dimos/memory/timeseries/test_inmemory.py b/dimos/memory/timeseries/test_inmemory.py deleted file mode 100644 index 37715f4a76..0000000000 --- a/dimos/memory/timeseries/test_inmemory.py +++ /dev/null @@ -1,175 +0,0 @@ -# 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 InMemoryStore.""" - -from dataclasses import dataclass - -import pytest - -from dimos.memory.timeseries.inmemory import InMemoryStore -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 - - -class TestInMemoryStoreOperations: - """Test InMemoryStore-specific operations.""" - - def test_delete(self): - store: InMemoryStore[SampleData] = InMemoryStore() - store.save(SampleData("a", 1.0)) - store.save(SampleData("b", 2.0)) - assert len(store) == 2 - deleted = store._delete(1.0) - assert deleted == SampleData("a", 1.0) - assert len(store) == 1 - assert store.load(1.0) is None - - -@pytest.mark.tool -class TestPerformance: - """Benchmarks comparing InMemoryStore vs TimestampedCollection. - - GC is disabled during measurements to avoid non-deterministic pauses. - """ - - N = 100_000 - - def _make_populated_store(self) -> InMemoryStore[SampleData]: - store: InMemoryStore[SampleData] = InMemoryStore() - for i in range(self.N): - store.save(SampleData(f"v{i}", float(i))) - return store - - def _make_populated_collection(self) -> "TimestampedCollection[SampleData]": - from dimos.types.timestamped import TimestampedCollection - - coll: TimestampedCollection[SampleData] = TimestampedCollection() - for i in range(self.N): - coll.add(SampleData(f"v{i}", float(i))) - return coll - - def test_insert_performance(self) -> None: - """Insert N items. InMemoryStore should be within 3x of TimestampedCollection.""" - import gc - import time as time_mod - - from dimos.types.timestamped import TimestampedCollection - - store: InMemoryStore[SampleData] = InMemoryStore() - gc.collect() - gc.disable() - t0 = time_mod.perf_counter() - for i in range(self.N): - store.save(SampleData(f"v{i}", float(i))) - store_time = time_mod.perf_counter() - t0 - gc.enable() - - coll: TimestampedCollection[SampleData] = TimestampedCollection() - gc.collect() - gc.disable() - t0 = time_mod.perf_counter() - for i in range(self.N): - coll.add(SampleData(f"v{i}", float(i))) - coll_time = time_mod.perf_counter() - t0 - gc.enable() - - print(f"\nInsert {self.N}: store={store_time:.3f}s, collection={coll_time:.3f}s") - assert store_time < coll_time * 3 - - def test_find_closest_performance(self) -> None: - """find_closest on N items. Both should be O(log n).""" - import gc - import random - import time as time_mod - - store = self._make_populated_store() - coll = self._make_populated_collection() - - queries = [random.uniform(0, self.N) for _ in range(10_000)] - - gc.collect() - gc.disable() - t0 = time_mod.perf_counter() - for q in queries: - store.find_closest(q) - store_time = time_mod.perf_counter() - t0 - - t0 = time_mod.perf_counter() - for q in queries: - coll.find_closest(q) - coll_time = time_mod.perf_counter() - t0 - gc.enable() - - print( - f"\nfind_closest 10k on {self.N}: store={store_time:.3f}s, collection={coll_time:.3f}s" - ) - assert store_time < coll_time * 3 - - def test_interleaved_write_read(self) -> None: - """Alternating write + find_closest. Old InMemoryStore was O(n log n) per read.""" - import gc - import time as time_mod - - store: InMemoryStore[SampleData] = InMemoryStore() - - gc.collect() - gc.disable() - t0 = time_mod.perf_counter() - for i in range(self.N): - store.save(SampleData(f"v{i}", float(i))) - if i % 10 == 0: - store.find_closest(float(i) / 2) - elapsed = time_mod.perf_counter() - t0 - gc.enable() - - print(f"\nInterleaved write+read {self.N}: {elapsed:.3f}s") - assert elapsed < 10.0 - - def test_iteration_performance(self) -> None: - """Full iteration over N items.""" - import gc - import time as time_mod - - store = self._make_populated_store() - coll = self._make_populated_collection() - - gc.collect() - gc.disable() - t0 = time_mod.perf_counter() - count_store = sum(1 for _ in store) - store_time = time_mod.perf_counter() - t0 - - t0 = time_mod.perf_counter() - count_coll = sum(1 for _ in coll) - coll_time = time_mod.perf_counter() - t0 - gc.enable() - - assert count_store == count_coll == self.N - print(f"\nIterate {self.N}: store={store_time:.3f}s, collection={coll_time:.3f}s") - assert store_time < coll_time * 3 diff --git a/dimos/utils/testing/replay_legacy.py b/dimos/utils/testing/replay_legacy.py deleted file mode 100644 index 0b28b3ec08..0000000000 --- a/dimos/utils/testing/replay_legacy.py +++ /dev/null @@ -1,401 +0,0 @@ -# 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. -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 - -from dimos.utils.data import get_data, get_data_dir - -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_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) From 420454c09c707fdf84d7f653dda1f37b1dbadbfa Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Wed, 11 Feb 2026 22:39:18 +0800 Subject: [PATCH 38/39] removed small comment --- dimos/utils/testing/replay.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/dimos/utils/testing/replay.py b/dimos/utils/testing/replay.py index 5ae730f4e7..588b63e099 100644 --- a/dimos/utils/testing/replay.py +++ b/dimos/utils/testing/replay.py @@ -11,10 +11,8 @@ # 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. -"""Shim for TimedSensorReplay/TimedSensorStorage. -For the original implementation, see replay_legacy.py. -""" +"""Shim for TimedSensorReplay/TimedSensorStorage.""" from dimos.memory.timeseries.legacy import LegacyPickleStore From 070803c9596dc1073073866987d54b5c40b69cea Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Thu, 12 Feb 2026 00:03:46 +0800 Subject: [PATCH 39/39] test fix, psql mypy --- dimos/memory/timeseries/postgres.py | 4 ++-- pyproject.toml | 10 +++++----- uv.lock | 18 +++++++----------- 3 files changed, 14 insertions(+), 18 deletions(-) diff --git a/dimos/memory/timeseries/postgres.py b/dimos/memory/timeseries/postgres.py index 0a06f36f74..a6a853bbd3 100644 --- a/dimos/memory/timeseries/postgres.py +++ b/dimos/memory/timeseries/postgres.py @@ -17,8 +17,8 @@ import pickle import re -import psycopg2 -import psycopg2.extensions +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 diff --git a/pyproject.toml b/pyproject.toml index 1b9627cfc8..aa8bb492be 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -64,7 +64,6 @@ dependencies = [ # TODO: rerun shouldn't be required but rn its in core (there is NO WAY to use dimos without rerun rn) # remove this once rerun is optional in core "rerun-sdk>=0.20.0", - "psycopg2-binary>=2.9.11", "toolz>=1.1.0", ] @@ -254,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", @@ -304,10 +308,6 @@ base = [ "dimos[agents,web,perception,visualization,sim]", ] -[dependency-groups] -dev = [ - "types-psycopg2>=2.9.21.20251012", -] [tool.ruff] line-length = 100 diff --git a/uv.lock b/uv.lock index 7f7dec461d..c6d378b9f7 100644 --- a/uv.lock +++ b/uv.lock @@ -1744,7 +1744,6 @@ dependencies = [ { name = "pin" }, { name = "plotext" }, { name = "plum-dispatch" }, - { name = "psycopg2-binary" }, { name = "pydantic" }, { name = "pydantic-settings" }, { name = "python-dotenv" }, @@ -1851,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" }, @@ -1938,6 +1938,9 @@ perception = [ { name = "transformers", extra = ["torch"] }, { name = "ultralytics" }, ] +psql = [ + { name = "psycopg2-binary" }, +] sim = [ { name = "mujoco" }, { name = "playground" }, @@ -1988,11 +1991,6 @@ web = [ { name = "uvicorn" }, ] -[package.dev-dependencies] -dev = [ - { name = "types-psycopg2" }, -] - [package.metadata] requires-dist = [ { name = "annotation-protocol", specifier = ">=1.4.0" }, @@ -2074,7 +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", specifier = ">=2.9.11" }, + { name = "psycopg2-binary", marker = "extra == 'psql'", specifier = ">=2.9.11" }, { name = "py-spy", marker = "extra == 'dev'" }, { name = "pydantic" }, { name = "pydantic", marker = "extra == 'docker'" }, @@ -2140,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" }, @@ -2157,10 +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"] - -[package.metadata.requires-dev] -dev = [{ name = "types-psycopg2", specifier = ">=2.9.21.20251012" }] +provides-extras = ["misc", "visualization", "agents", "web", "perception", "unitree", "manipulation", "cpu", "cuda", "dev", "psql", "sim", "drone", "docker", "base"] [[package]] name = "dimos-lcm"