digitalkin 0.2.7__tar.gz → 0.2.9__tar.gz

{digitalkin-0.2.7 → digitalkin-0.2.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: digitalkin
-Version: 0.2.7
+Version: 0.2.9
 Summary: SDK to build kin used in DigitalKin
 Author-email: "DigitalKin.ai" <contact@digitalkin.ai>
 License: Attribution-NonCommercial-ShareAlike 4.0 International
@@ -452,20 +452,20 @@ Classifier: License :: Other/Proprietary License
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: digitalkin-proto>=0.1.7
+Requires-Dist: digitalkin-proto>=0.1.8
 Requires-Dist: grpcio-health-checking>=1.71.0
 Requires-Dist: grpcio-reflection>=1.71.0
 Requires-Dist: grpcio-status>=1.71.0
-Requires-Dist: openai>=1.70.0
-Requires-Dist: pydantic>=2.11.3
+Requires-Dist: openai>=1.76.2
+Requires-Dist: pydantic>=2.11.4
 Provides-Extra: dev
 Requires-Dist: pytest>=8.3.4; extra == "dev"
 Requires-Dist: pytest-asyncio>=0.26.0; extra == "dev"
 Requires-Dist: pytest-cov>=6.1.0; extra == "dev"
-Requires-Dist: typos>=1.31.1; extra == "dev"
-Requires-Dist: ruff>=0.11.2; extra == "dev"
+Requires-Dist: typos>=1.31.2; extra == "dev"
+Requires-Dist: ruff>=0.11.7; extra == "dev"
 Requires-Dist: mypy>=1.15.0; extra == "dev"
-Requires-Dist: pyright>=1.1.398; extra == "dev"
+Requires-Dist: pyright>=1.1.400; extra == "dev"
 Requires-Dist: pre-commit>=4.2.0; extra == "dev"
 Requires-Dist: bump2version>=1.0.1; extra == "dev"
 Requires-Dist: build>=1.2.2; extra == "dev"

{digitalkin-0.2.7 → digitalkin-0.2.9}/examples/modules/storage_module.py

@@ -119,7 +119,12 @@ class ExampleModule(ArchetypeModule[ExampleInput, ExampleOutput, ExampleSetup, E
         )
 
         # Store the output data in storage
-        storage_id = self.storage.store("example_outputs", output_data.model_dump(), data_type="OUTPUT")
+        storage_id = self.storage.store(
+            collection="example",
+            record_id=f"example_outputs",
+            data=output_data.model_dump(),
+            data_type="OUTPUT"
+        )
 
         logger.info("Stored output data with ID: %s", storage_id)
 
@@ -159,7 +164,7 @@ async def test_module() -> None:
 
     # Check the storage
     if module.status == ModuleStatus.STOPPED:
-        result: StorageRecord = module.storage.read("example_outputs")
+        result: StorageRecord = module.storage.read("example", "example_outputs")
         if result:
             pass
 
@@ -170,10 +175,10 @@ def test_storage_directly() -> None:
     storage = ServicesConfig().storage(mission_id="test-mission", config={"test_table": ExampleStorage})
 
     # Create a test record
-    storage.store("test_table", {"test_key": "test_value"}, "OUTPUT")
+    storage.store("example", "test_table", {"test_key": "test_value"}, "OUTPUT")
 
     # Retrieve the record
-    retrieved = storage.read("test_table")
+    retrieved = storage.read("example", "test_table")
 
     if retrieved:
         pass

{digitalkin-0.2.7 → digitalkin-0.2.9}/pyproject.toml

@@ -12,7 +12,7 @@
 
     keywords = [ "digitalkin", "kin", "agent", "gprc", "sdk" ]
     # Version of the package automatically updated by bump2version (that is why it is separated)
-    version = "0.2.7"
+    version = "0.2.9"
 
     classifiers = [
         "Development Status :: 3 - Alpha",
@@ -29,12 +29,12 @@
     ]
 
     dependencies = [
-        "digitalkin-proto>=0.1.7",
+        "digitalkin-proto>=0.1.8",
         "grpcio-health-checking>=1.71.0",
         "grpcio-reflection>=1.71.0",
         "grpcio-status>=1.71.0",
-        "openai>=1.70.0",
-        "pydantic>=2.11.3",
+        "openai>=1.76.2",
+        "pydantic>=2.11.4",
     ]
 
     [project.optional-dependencies]
@@ -42,10 +42,10 @@
             "pytest>=8.3.4",
             "pytest-asyncio>=0.26.0",
             "pytest-cov>=6.1.0",
-            "typos>=1.31.1",
-            "ruff>=0.11.2",
+            "typos>=1.31.2",
+            "ruff>=0.11.7",
             "mypy>=1.15.0",
-            "pyright>=1.1.398",
+            "pyright>=1.1.400",
             "pre-commit>=4.2.0",
             "bump2version>=1.0.1",
             "build>=1.2.2",

{digitalkin-0.2.7 → digitalkin-0.2.9}/src/digitalkin/__version__.py

@@ -5,4 +5,4 @@ from importlib.metadata import PackageNotFoundError, version
 try:
     __version__ = version("digitalkin")
 except PackageNotFoundError:
-    __version__ = "0.2.7"
+    __version__ = "0.2.9"

digitalkin-0.2.9/src/digitalkin/services/storage/default_storage.py

@@ -0,0 +1,220 @@
+"""This module implements the default storage strategy."""
+
+import datetime
+import json
+import logging
+import tempfile
+from pathlib import Path
+from typing import Any
+
+from pydantic import BaseModel
+
+from digitalkin.services.storage.storage_strategy import DataType, StorageRecord, StorageStrategy
+
+logger = logging.getLogger(__name__)
+
+
+class DefaultStorage(StorageStrategy):
+    """Persist records in a local JSON file for quick local development.
+
+    File format: a JSON object of
+      { "<collection>:<record_id>": { ... StorageRecord fields ... },
+    """
+
+    @staticmethod
+    def _json_default(o: Any) -> str:  # noqa: ANN401
+        """JSON serializer for non-standard types (datetime → ISO).
+
+        Args:
+            o: The object to serialize
+
+        Returns:
+            str: The serialized object
+
+        Raises:
+            TypeError: If the object is not serializable
+        """
+        if isinstance(o, datetime.datetime):
+            return o.isoformat()
+        msg = f"Type {o.__class__.__name__} not serializable"
+        raise TypeError(msg)
+
+    def _load_from_file(self) -> dict[str, StorageRecord]:
+        """Load storage data from the file.
+
+        Returns:
+            A dictionary containing the loaded storage records
+        """
+        if not self.storage_file.exists():
+            return {}
+
+        try:
+            raw = json.loads(self.storage_file.read_text(encoding="utf-8"))
+            out: dict[str, StorageRecord] = {}
+
+            for key, rd in raw.items():
+                # rd is a dict with the StorageRecord fields
+                model_cls = self.config.get(rd["collection"])
+                if not model_cls:
+                    logger.warning("No model for collection '%s'", rd["collection"])
+                    continue
+                data_model = model_cls.model_validate(rd["data"])
+                rec = StorageRecord(
+                    mission_id=rd["mission_id"],
+                    collection=rd["collection"],
+                    record_id=rd["record_id"],
+                    data=data_model,
+                    data_type=DataType[rd["data_type"]],
+                    creation_date=datetime.datetime.fromisoformat(rd["creation_date"])
+                    if rd.get("creation_date")
+                    else None,
+                    update_date=datetime.datetime.fromisoformat(rd["update_date"]) if rd.get("update_date") else None,
+                )
+                out[key] = rec
+        except Exception:
+            logger.exception("Failed to load default storage file")
+            return {}
+        return out
+
+    def _save_to_file(self) -> None:
+        """Atomically write `self.storage` back to disk as JSON."""
+        self.storage_file.parent.mkdir(parents=True, exist_ok=True)
+        with tempfile.NamedTemporaryFile(
+            mode="w", encoding="utf-8", delete=False, dir=str(self.storage_file.parent), suffix=".tmp"
+        ) as temp:
+            try:
+                # Convert storage to a serializable format
+                serial: dict[str, dict] = {}
+                for key, record in self.storage.items():
+                    serial[key] = {
+                        "mission_id": record.mission_id,
+                        "collection": record.collection,
+                        "record_id": record.record_id,
+                        "data_type": record.data_type.name,
+                        "data": record.data.model_dump(),
+                        "creation_date": record.creation_date.isoformat() if record.creation_date else None,
+                        "update_date": record.update_date.isoformat() if record.update_date else None,
+                    }
+                json.dump(serial, temp, indent=2, default=self._json_default)
+                temp.flush()
+                Path(temp.name).replace(self.storage_file)
+            except Exception:
+                logger.exception("Unexpected error saving storage")
+
+    def _store(self, record: StorageRecord) -> StorageRecord:
+        """Store a new record in the database and persist to file.
+
+        Args:
+            record: The record to store
+
+        Returns:
+            str: The ID of the new record
+
+        Raises:
+            ValueError: If the record already exists
+        """
+        key = f"{record.collection}:{record.record_id}"
+        if key in self.storage:
+            msg = f"Document {key!r} already exists"
+            raise ValueError(msg)
+        now = datetime.datetime.now(datetime.timezone.utc)
+        record.creation_date = now
+        record.update_date = now
+        self.storage[key] = record
+        self._save_to_file()
+        logger.debug("Created %s", key)
+        return record
+
+    def _read(self, collection: str, record_id: str) -> StorageRecord | None:
+        """Get records from the database.
+
+        Args:
+            collection: The unique name to retrieve data for
+            record_id: The unique ID of the record
+
+        Returns:
+            StorageRecord: The corresponding record
+        """
+        key = f"{collection}:{record_id}"
+        return self.storage.get(key)
+
+    def _update(self, collection: str, record_id: str, data: BaseModel) -> StorageRecord | None:
+        """Update records in the database and persist to file.
+
+        Args:
+            collection: The unique name to retrieve data for
+            record_id: The unique ID of the record
+            data: The data to modify
+
+        Returns:
+            StorageRecord: The modified record
+        """
+        key = f"{collection}:{record_id}"
+        rec = self.storage.get(key)
+        if not rec:
+            return None
+        rec.data = data
+        rec.update_date = datetime.datetime.now(datetime.timezone.utc)
+        self._save_to_file()
+        logger.debug("Modified %s", key)
+        return rec
+
+    def _remove(self, collection: str, record_id: str) -> bool:
+        """Delete records from the database and update file.
+
+        Args:
+            collection: The unique name to retrieve data for
+            record_id: The unique ID of the record
+
+        Returns:
+            bool: True if the record was removed, False otherwise
+        """
+        key = f"{collection}:{record_id}"
+        if key not in self.storage:
+            return False
+        del self.storage[key]
+        self._save_to_file()
+        logger.debug("Removed %s", key)
+        return True
+
+    def _list(self, collection: str) -> list[StorageRecord]:
+        """Implements StorageStrategy._list.
+
+        Args:
+            collection: The unique name to retrieve data for
+
+        Returns:
+            A list of storage records
+        """
+        prefix = f"{collection}:"
+        return [r for k, r in self.storage.items() if k.startswith(prefix)]
+
+    def _remove_collection(self, collection: str) -> bool:
+        """Implements StorageStrategy._remove_collection.
+
+        Args:
+            collection: The unique name to retrieve data for
+
+        Returns:
+            bool: True if the collection was removed, False otherwise
+        """
+        prefix = f"{collection}:"
+        to_delete = [k for k in self.storage if k.startswith(prefix)]
+        for k in to_delete:
+            del self.storage[k]
+        self._save_to_file()
+        logger.debug("Removed collection %s (%d docs)", collection, len(to_delete))
+        return True
+
+    def __init__(
+        self,
+        mission_id: str,
+        config: dict[str, type[BaseModel]],
+        storage_file_path: str = "local_storage",
+        **kwargs,  # noqa: ANN003, ARG002
+    ) -> None:
+        """Initialize the storage."""
+        super().__init__(mission_id=mission_id, config=config)
+        self.storage_file_path = f"{self.mission_id}_{storage_file_path}.json"
+        self.storage_file = Path(self.storage_file_path)
+        self.storage = self._load_from_file()

digitalkin-0.2.9/src/digitalkin/services/storage/grpc_storage.py

@@ -0,0 +1,197 @@
+"""This module implements the default storage strategy."""
+
+import logging
+
+from digitalkin_proto.digitalkin.storage.v2 import data_pb2, storage_service_pb2_grpc
+from google.protobuf import json_format
+from google.protobuf.struct_pb2 import Struct
+from pydantic import BaseModel
+
+from digitalkin.grpc_servers.utils.grpc_client_wrapper import GrpcClientWrapper
+from digitalkin.grpc_servers.utils.models import ClientConfig
+from digitalkin.services.storage.storage_strategy import DataType, StorageRecord, StorageServiceError, StorageStrategy
+
+logger = logging.getLogger(__name__)
+
+
+class GrpcStorage(StorageStrategy, GrpcClientWrapper):
+    """This class implements the default storage strategy."""
+
+    def _build_record_from_proto(self, proto: data_pb2.StorageRecord) -> StorageRecord:
+        """Convert a protobuf StorageRecord message into our Pydantic model.
+
+        Args:
+            proto: gRPC StorageRecord
+
+        Returns:
+            A fully validated StorageRecord.
+        """
+        raw = json_format.MessageToDict(
+            proto,
+            preserving_proto_field_name=True,
+            always_print_fields_with_no_presence=True,
+        )
+        mission = raw["mission_id"]
+        coll = raw["collection"]
+        rid = raw["record_id"]
+        dtype = DataType[raw["data_type"]]
+        payload = raw.get("data", {})
+
+        validated = self._validate_data(rid, payload)
+        return StorageRecord(
+            mission_id=mission,
+            collection=coll,
+            record_id=rid,
+            data=validated,
+            data_type=dtype,
+            creation_date=raw.get("creation_date"),
+            update_date=raw.get("update_date"),
+        )
+
+    def _store(self, record: StorageRecord) -> StorageRecord:
+        """Create a new record in the database.
+
+        Parameters:
+            record: The record to store
+
+        Returns:
+            StorageRecord: The corresponding record
+
+        Raises:
+            StorageServiceError: If there is an error while storing the record
+        """
+        try:
+            data_struct = Struct()
+            data_struct.update(record.data.model_dump())
+            req = data_pb2.StoreRecordRequest(
+                data=data_struct,
+                mission_id=record.mission_id,
+                collection=record.collection,
+                record_id=record.record_id,
+                data_type=record.data_type.name,
+            )
+            resp = self.exec_grpc_query("StoreRecord", req)
+            return self._build_record_from_proto(resp.stored_data)
+        except Exception as e:
+            logger.exception("gRPC StoreRecord failed for %s:%s", record.collection, record.record_id)
+            raise StorageServiceError(str(e)) from e
+
+    def _read(self, collection: str, record_id: str) -> StorageRecord | None:
+        """Fetch a single document by collection + record_id.
+
+        Returns:
+            StorageData: The record
+        """
+        try:
+            req = data_pb2.ReadRecordRequest(
+                mission_id=self.mission_id,
+                collection=collection,
+                record_id=record_id,
+            )
+            resp = self.exec_grpc_query("ReadRecord", req)
+            return self._build_record_from_proto(resp.stored_data)
+        except Exception:
+            logger.exception("gRPC ReadRecord failed for %s:%s", collection, record_id)
+            return None
+
+    def _update(self, collection: str, record_id: str, data: BaseModel) -> StorageRecord | None:
+        """Overwrite a document via gRPC.
+
+        Args:
+            collection: The unique name for the record type
+            record_id: The unique ID for the record
+            data: The validated data model
+
+        Returns:
+            StorageRecord: The updated record
+        """
+        try:
+            struct = Struct()
+            struct.update(data.model_dump())
+            req = data_pb2.UpdateRecordRequest(
+                data=struct,
+                mission_id=self.mission_id,
+                collection=collection,
+                record_id=record_id,
+            )
+            resp = self.exec_grpc_query("ModifyRecord", req)
+            return self._build_record_from_proto(resp.stored_data)
+        except Exception:
+            logger.exception("gRPC ModifyRecord failed for %s:%s", collection, record_id)
+            return None
+
+    def _remove(self, collection: str, record_id: str) -> bool:
+        """Delete a document via gRPC.
+
+        Args:
+            collection: The unique name for the record type
+            record_id: The unique ID for the record
+
+        Returns:
+            bool: True if the record was deleted, False otherwise
+        """
+        try:
+            req = data_pb2.RemoveRecordRequest(
+                mission_id=self.mission_id,
+                collection=collection,
+                record_id=record_id,
+            )
+            self.exec_grpc_query("RemoveRecord", req)
+        except Exception:
+            logger.exception("gRPC RemoveRecord failed for %s:%s", collection, record_id)
+            return False
+        return True
+
+    def _list(self, collection: str) -> list[StorageRecord]:
+        """List all documents in a collection via gRPC.
+
+        Args:
+            collection: The unique name for the record type
+
+        Returns:
+            list[StorageRecord]: A list of storage records
+        """
+        try:
+            req = data_pb2.ListRecordsRequest(
+                mission_id=self.mission_id,
+                collection=collection,
+            )
+            resp = self.exec_grpc_query("ListRecords", req)
+            return [self._build_record_from_proto(r) for r in resp.records]
+        except Exception:
+            logger.exception("gRPC ListRecords failed for %s", collection)
+            return []
+
+    def _remove_collection(self, collection: str) -> bool:
+        """Delete an entire collection via gRPC.
+
+        Args:
+            collection: The unique name for the record type
+
+        Returns:
+            bool: True if the collection was deleted, False otherwise
+        """
+        try:
+            req = data_pb2.RemoveCollectionRequest(
+                mission_id=self.mission_id,
+                collection=collection,
+            )
+            self.exec_grpc_query("RemoveCollection", req)
+        except Exception:
+            logger.exception("gRPC RemoveCollection failed for %s", collection)
+            return False
+        return True
+
+    def __init__(
+        self,
+        mission_id: str,
+        config: dict[str, type[BaseModel]],
+        client_config: ClientConfig,
+        **kwargs,  # noqa: ANN003, ARG002
+    ) -> None:
+        """Initialize the storage."""
+        super().__init__(mission_id=mission_id, config=config)
+
+        channel = self._init_channel(client_config)
+        self.stub = storage_service_pb2_grpc.StorageServiceStub(channel)
+        logger.debug("Channel client 'storage' initialized succesfully")