digitalkin 0.1.1__py3-none-any.whl → 0.2.0__py3-none-any.whl

digitalkin/services/storage/storage_strategy.py

@@ -1,42 +1,213 @@
 """This module contains the abstract base class for storage strategies."""
 
+import datetime
 from abc import ABC, abstractmethod
-from typing import Any
+from enum import Enum
+from typing import Any, Literal, TypeGuard
 
+from pydantic import BaseModel, Field
 
-class StorageStrategy(ABC):
-    """Abstract base class for storage strategies."""
+from digitalkin.services.base_strategy import BaseStrategy
 
-    def __init__(self) -> None:
-        """Initialize the storage strategy."""
 
-    def __post_init__(self, *args, **kwargs) -> None:  # noqa: ANN002, ANN003
-        """Initialize the storage strategy."""
+class StorageServiceError(Exception):
+    """Base exception for Setup service errors."""
 
-    @abstractmethod
-    def connect(self) -> bool:
-        """Establish connection to the database."""
 
-    @abstractmethod
-    def disconnect(self) -> bool:
-        """Close connection to the database."""
+class DataType(Enum):
+    """Enum defining the types of data that can be stored."""
 
-    @abstractmethod
-    def create(self, table: str, data: dict[str, Any]) -> str:
-        """Create a new record in the database."""
+    OUTPUT = "OUTPUT"
+    VIEW = "VIEW"
+    LOGS = "LOGS"
+    OTHER = "OTHER"
+
+
+class StorageRecord(BaseModel):
+    """Container for stored records with metadata."""
+
+    # Metadata
+    mission_id: str = Field(description="The ID of the mission this record is associated with")
+    name: str = Field(description="The name of the record")
+    creation_date: datetime.datetime | None = Field(default=None, description="The date the record was created")
+    update_date: datetime.datetime | None = Field(default=None, description="The date the record was last updated")
+    data_type: DataType = Field(default=DataType.OUTPUT, description="The type of data stored")
+    # Actual data payload
+    data: BaseModel = Field(description="The data stored in the record")
+
+
+class StorageStrategy(BaseStrategy, ABC):
+    """Abstract base class for storage strategies.
+
+    This strategy defines how data is stored and retrieved, with
+    type validation through registered Pydantic models.
+    """
+
+    def __init__(self, mission_id: str, config: dict[str, type[BaseModel]]) -> None:
+        """Initialize the storage strategy.
+
+        Args:
+            mission_id: The ID of the mission this strategy is associated with
+            config: A dictionary mapping names to Pydantic model classes
+        """
+        super().__init__(mission_id)
+        # Schema configuration mapping keys to model classes
+        self.config: dict[str, type[BaseModel]] = config
+
+    @staticmethod
+    def _is_valid_data_type_name(value: str) -> TypeGuard[str]:
+        return value in DataType.__members__
 
     @abstractmethod
-    def get(self, table: str, data: dict[str, Any]) -> list[dict[str, Any]]:
-        """Get records from the database."""
+    def _store(self, record: StorageRecord) -> StorageRecord:
+        """Store a new record in the storage.
+
+        Args:
+            record: The record to store
+
+        Returns:
+            The ID of the created record
+        """
+
+    def store(
+        self,
+        name: str,
+        data: dict[str, Any],
+        data_type: Literal["OUTPUT", "VIEW", "LOGS", "OTHER"] = "OUTPUT",
+    ) -> StorageRecord:
+        """Store a new record in the storage.
+
+        Args:
+            name: The unique name to store the data under
+            data: The data to store
+            data_type: The type of data being stored (default: OUTPUT)
+
+        Returns:
+            The ID of the created record
+
+        Raises:
+            ValueError: If the data type is invalid or if validation fails
+        """
+        if not self._is_valid_data_type_name(data_type):
+            msg = f"Invalid data type '{data_type}'. Must be one of {list(DataType.__members__.keys())}"
+            raise ValueError(msg)
+        data_type_enum = DataType[data_type]
+        validated_data = self._validate_data(name, {**data, "mission_id": self.mission_id})
+        record = self._create_storage_record(name, validated_data, data_type_enum)
+        return self._store(record)
 
     @abstractmethod
-    def update(self, table: str, data: dict[str, Any]) -> int:
-        """Update records in the database."""
+    def _read(self, name: str) -> StorageRecord | None:
+        """Get records from storage by key.
+
+        Args:
+            name: The unique name to retrieve data for
+
+        Returns:
+            A storage record with validated data
+        """
+
+    def read(self, name: str) -> StorageRecord | None:
+        """Get records from storage by key.
+
+        Args:
+            name: The unique name to retrieve data for
+
+        Returns:
+            A storage record with validated data
+        """
+        return self._read(name)
 
     @abstractmethod
-    def delete(self, table: str, data: dict[str, Any]) -> int:
-        """Delete records from the database."""
+    def _modify(self, name: str, data: BaseModel) -> StorageRecord | None:
+        """Update a record in the storage.
+
+        Args:
+            name: The unique name for the record type
+            data: The new data to store
+
+        Returns:
+            StorageRecord: The modified record
+        """
+
+    def modify(self, name: str, data: dict[str, Any]) -> StorageRecord | None:
+        """Update a record in the storage (overwrite all the data).
+
+        Args:
+            name: The unique name for the record type
+            data: The new data to store
+
+        Returns:
+            StorageRecord: The modified record
+        """
+        validated_data = self._validate_data(name, data)
+        return self._modify(name, validated_data)
 
     @abstractmethod
-    def get_all(self) -> dict[str, list[dict[str, Any]]]:
-        """Get all records from the database."""
+    def _remove(self, name: str) -> bool:
+        """Delete a record from the storage.
+
+        Args:
+            name: The unique name for the record type
+
+        Returns:
+            True if the deletion was successful, False otherwise
+        """
+
+    def remove(self, name: str) -> bool:
+        """Delete a record from the storage.
+
+        Args:
+            name: The unique name for the record type
+
+        Returns:
+            True if the deletion was successful, False otherwise
+        """
+        return self._remove(name)
+
+    def _validate_data(self, name: str, data: dict[str, Any]) -> BaseModel:
+        """Validate data against the model schema for the given key.
+
+        Args:
+            name: The unique name to get the model type for
+            data: The data to validate
+
+        Returns:
+            A validated model instance
+
+        Raises:
+            ValueError: If the key has no associated model or validation fails
+        """
+        model_cls = self.config.get(name)
+        if not model_cls:
+            msg = f"No model schema defined for name: {name}"
+            raise ValueError(msg)
+
+        try:
+            return model_cls.model_validate(data)
+        except Exception as e:
+            msg = f"Data validation failed for key '{name}': {e!s}"
+            raise ValueError(msg) from e
+
+    def _create_storage_record(
+        self,
+        name: str,
+        validated_data: BaseModel,
+        data_type: DataType,
+    ) -> StorageRecord:
+        """Create a storage record with metadata.
+
+        Args:
+            name: The unique name for the record
+            validated_data: The validated data model
+            data_type: The type of data
+
+        Returns:
+            A complete storage record with metadata
+        """
+        return StorageRecord(
+            mission_id=self.mission_id,
+            name=name,
+            data=validated_data,
+            data_type=data_type,
+        )

digitalkin/utils/arg_parser.py

@@ -6,7 +6,7 @@ from argparse import Action, ArgumentParser, Namespace, _HelpAction, _SubParsers
 from collections.abc import Sequence
 from typing import Any
 
-from digitalkin.services.service_provider import ServiceProvider
+from digitalkin.services.services_models import ServicesMode
 
 logging.getLogger().setLevel(logging.INFO)
 logger = logging.getLogger(__name__)
@@ -99,38 +99,37 @@ class ArgParser:
 class DevelopmentModeMappingAction(Action):
     """."""
 
-    default: ServiceProvider | None
-    class_mapping: dict[str, ServiceProvider]
-
     def __init__(
         self,
         env_var: str,
-        class_mapping: dict[str, ServiceProvider],
         required: bool = True,  # noqa: FBT001, FBT002
         default: str | None = None,
         **kwargs: dict[str, Any],
     ) -> None:
         """."""
-        default = class_mapping.get(os.environ.get(env_var, default), None)  # type: ignore
-        if default is None:
-            logger.error("Invalid default value: %s, for the Service Provider in the module", default)
+        default = ServicesMode(os.environ.get(env_var, default))
 
         if required and default:
             required = False
-        self.class_mapping = class_mapping
         super().__init__(default=default, required=required, **kwargs)  # type: ignore
 
     def __call__(
         self,
-        parser: ArgumentParser,
+        parser: ArgumentParser,  # noqa: ARG002
         namespace: Namespace,
         values: str | Sequence[Any] | None,
         option_string: str | None = None,  # noqa: ARG002
     ) -> None:
-        """Set the attribute to the corresponding class."""
-        if values not in self.class_mapping:
-            logger.error("Invalid mode: %s, dest: %s not set!", values, self.dest)
-            parser.error(f"Invalid mode: {values}")
-
-        values = self.class_mapping[values]  # type: ignore
-        setattr(namespace, self.dest, values)
+        """Set the attribute to the corresponding class.
+
+        Raises:
+            TypeError: if the value is not a string.
+        """
+        # Check if the value is a string and convert it to lowercase
+        if isinstance(values, str):
+            values = values.lower()
+        else:
+            msg = "values must be a string"
+            raise TypeError(msg)
+        mode = ServicesMode(values)
+        setattr(namespace, self.dest, mode)

{digitalkin-0.1.1.dist-info → digitalkin-0.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: digitalkin
-Version: 0.1.1
+Version: 0.2.0
 Summary: SDK to build kin used in DigitalKin
 Author-email: "DigitalKin.ai" <contact@digitalkin.ai>
 License: Attribution-NonCommercial-ShareAlike 4.0 International
@@ -452,25 +452,26 @@ Classifier: License :: Other/Proprietary License
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: digitalkin-proto>=0.0.8
+Requires-Dist: digitalkin-proto>=0.1.5
 Requires-Dist: grpcio-health-checking>=1.71.0
 Requires-Dist: grpcio-reflection>=1.71.0
 Requires-Dist: grpcio-status>=1.71.0
-Requires-Dist: pydantic>=2.10.6
+Requires-Dist: pydantic>=2.11.3
 Provides-Extra: dev
 Requires-Dist: pytest>=8.3.4; extra == "dev"
-Requires-Dist: pytest-asyncio>=0.25.3; extra == "dev"
-Requires-Dist: pytest-cov>=6.0.0; extra == "dev"
-Requires-Dist: typos>=1.30.2; 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: mypy>=1.15.0; extra == "dev"
-Requires-Dist: pyright>=1.1.397; extra == "dev"
+Requires-Dist: pyright>=1.1.398; 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"
 Requires-Dist: twine>=6.1.0; extra == "dev"
 Requires-Dist: cryptography>=44.0.2; extra == "dev"
 Requires-Dist: grpcio-testing>=1.71.0; extra == "dev"
+Requires-Dist: freezegun>=1.5.1; extra == "dev"
 Provides-Extra: examples
 Requires-Dist: openai>=1.66.3; extra == "examples"
 Dynamic: license-file

digitalkin-0.2.0.dist-info/RECORD

@@ -0,0 +1,78 @@
+base_server/__init__.py,sha256=gs8t9Dg0dNVHRdYYbEQY8bn8tzEUv2zE6eBhKNPG3kU,88
+base_server/server_async_insecure.py,sha256=Rvj5Xj8tDMglAj6aOGlQdMeD-FL97_OBEhHYHxxQrVo,3887
+base_server/server_async_secure.py,sha256=SV_CqEmD6YNHHP2Y369T2PLPJ-9JG3bvXytchX_Ensk,4684
+base_server/server_sync_insecure.py,sha256=VgSH8YghagK3fiwHhc7d6__zjN6lj_FgBUponFbOoxM,3127
+base_server/server_sync_secure.py,sha256=MMkq67vAZNiDLJSySjoXtHIEIK1pAlRz57n75Egnndk,3939
+base_server/mock/__init__.py,sha256=YZFT-F1l_TpvJYuIPX-7kTeE1CfOjhx9YmNRXVoi-jQ,143
+base_server/mock/mock_pb2.py,sha256=sETakcS3PAAm4E-hTCV1jIVaQTPEAIoVVHupB8Z_k7Y,1843
+base_server/mock/mock_pb2_grpc.py,sha256=BbOT70H6q3laKgkHfOx1QdfmCS_HxCY4wCOX84YAdG4,3180
+digitalkin/__init__.py,sha256=7LLBAba0th-3SGqcpqFO-lopWdUkVLKzLZiMtB-mW3M,162
+digitalkin/__version__.py,sha256=c4f1Pc-TEf7O1UkIJonXwVhP7WicVMKecmjTFvgmDLQ,190
+digitalkin/logger.py,sha256=9cDgyJV2QXXT8F--xRODFlZyDgjuTTXNdpCU3GdqCsk,382
+digitalkin/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+digitalkin/grpc_servers/__init__.py,sha256=0cJBlwipSmFdXkyH3T0i6OJ1WpAtNsZgYX7JaSnkbtg,804
+digitalkin/grpc_servers/_base_server.py,sha256=ec4xmgAuOMVg45a63O_PEa2T7mI4tJ6boxcXauFyZ5g,18649
+digitalkin/grpc_servers/module_server.py,sha256=0HVOygMGP82WDU8I_04ZrFQ1Tqme0bM5nyqMVmoOL1Y,9072
+digitalkin/grpc_servers/module_servicer.py,sha256=togOyvogAYkBv353S-9X_pP5yI8MvxupOIBSvq39jA4,11816
+digitalkin/grpc_servers/registry_server.py,sha256=AjN-OXLGsJtvddqLKeHosJVdSDqvjfpkQM3JcPl5PSQ,2120
+digitalkin/grpc_servers/registry_servicer.py,sha256=mCAjNhdMq5DozZMEPsJK__DIxePEYxSWV-gAq-Xctk4,16469
+digitalkin/grpc_servers/utils/exceptions.py,sha256=I00OM8p8up20He4dU1fiHsvdLj1DymjR_UmoeUm2MSA,785
+digitalkin/grpc_servers/utils/factory.py,sha256=jm6rFjiqmtSv7BIHNAOxsG9xXtSvWpx9TfzSQiX97MQ,5899
+digitalkin/grpc_servers/utils/grpc_client_wrapper.py,sha256=y9zGzmeESAhrzzQcekA8rE_P9gyPvOjxWwnZAc3fpx4,2435
+digitalkin/grpc_servers/utils/models.py,sha256=AXv4FQ3RRmcr5sZaaZc72ED1LtGJEaxmFIZP0Xlo-PY,5483
+digitalkin/grpc_servers/utils/types.py,sha256=rQ78s4nAet2jy-NIDj_PUWriT0kuGHr_w6ELjmjgBao,539
+digitalkin/models/__init__.py,sha256=hDHtUfswaNh8wo4NZaBItg9JqC0uNSRqXArNWSrGynY,163
+digitalkin/models/module/__init__.py,sha256=ihvRlemJuFvU4armZIL-Vq-zaJx9UrjDoJAVPbEG8jw,345
+digitalkin/models/module/module.py,sha256=vlIeWmH61NVEiJpRx0Crk4iB8hAUV4mG30likSyRFP4,686
+digitalkin/models/module/module_types.py,sha256=SxwzfDDnkyuQe50glVrxN6VPuJCkbB313FRyVZugIMw,312
+digitalkin/models/services/__init__.py,sha256=HsW7MUGFPvH7Ri28WN4BHHBfEQk5dzU_9FOWAc-0lSE,56
+digitalkin/models/services/cost.py,sha256=QTEuFD6xz62nob0z4ksE-INJWcZ-iFiuNW5mvXhpFes,1599
+digitalkin/models/services/storage.py,sha256=cYTVIriGKiprF9OerhSxmc_jM6fUTVwmeon1yQCinkE,143
+digitalkin/modules/__init__.py,sha256=ppYARmhvdVi55ofC0QZerIempSlcJYDeCXhcl4qXObw,278
+digitalkin/modules/_base_module.py,sha256=f4AEZFyl0meT_V6esnEFyUtuN7ZXvhHgIYa1ZOMe7tM,7651
+digitalkin/modules/archetype_module.py,sha256=T2Ehj7EpAC2MO9WQbJv39hqRw7rh3exhVZTEL3JPM8U,421
+digitalkin/modules/job_manager.py,sha256=q48ZJwWmOZ4xJKV2QApI6x_Nfgviej-AoIV8BJQ9B-o,6110
+digitalkin/modules/tool_module.py,sha256=86g0M1wHZ1ReIc7AkKfyjnlGN2QYJBGxrEQpKVlyrZI,421
+digitalkin/modules/trigger_module.py,sha256=kVoI4Gdkw7WWUP5T6hSCNqw5FxibTxL6Tpq9KP7gg78,379
+digitalkin/services/__init__.py,sha256=LqGk_5DJy8Bzz62ajIq9jCeYNKQUIgtSCpafZk15FLc,910
+digitalkin/services/base_strategy.py,sha256=U7J09LkNvTaj4bkX-gF86iU_FcUeHK6AKp4zyUPx6eI,470
+digitalkin/services/services_config.py,sha256=6mNYtw07ZNZbk2b4e3xKEpYfozD_NUZyUY3iiein_10,7256
+digitalkin/services/services_models.py,sha256=5zXkWcfKnXGwQi9sN4OAL3XrgqOcmsTl8ai5Mi4RPsw,1668
+digitalkin/services/agent/__init__.py,sha256=vJc8JN0pdtA8ecypLBeHrwAUIW6H2C8NyW-dk24rTpk,244
+digitalkin/services/agent/agent_strategy.py,sha256=42Q9RciHX6tg3CgDQkbrlIx4h_TX0WIuSpLmCjitVmA,492
+digitalkin/services/agent/default_agent.py,sha256=4N_E_eQxJGOx1KVUUg5jNOje-3ncMxF3ePB-uDuGrJc,345
+digitalkin/services/cost/__init__.py,sha256=Wi9ZB4LSXFsUYgkX-V1UJQkVXYDNDpp8q2dXccR2uRM,303
+digitalkin/services/cost/cost_strategy.py,sha256=KkVLT8X2EY58ul-27Gj6lFx7dsk1SZgg-UunBC0D8Ks,986
+digitalkin/services/cost/default_cost.py,sha256=LaTOBIdGOwh4-9vDlnefm8eKDvSWglJOS7BFLOT_G9Q,811
+digitalkin/services/cost/grpc_cost.py,sha256=k0Df_41dwrowBjtz11eqYQyWCMIjY-wEPEytbo2j4gE,3215
+digitalkin/services/filesystem/__init__.py,sha256=BhwMl_BUvM0d65fmglkp0SVwn3RfYiUOKJgIMnOCaGM,381
+digitalkin/services/filesystem/default_filesystem.py,sha256=tTOsBt0MSjexKgWm657y8BLppMYJtmJIVdqzsSRRaYk,7264
+digitalkin/services/filesystem/filesystem_strategy.py,sha256=VH93jkijFtyVA6XaH_rsfdSyeEpkqea54NAvHFLJDjM,2216
+digitalkin/services/filesystem/grpc_filesystem.py,sha256=rtVfutPoggB1AG6vcxrSMnUOu20h9tlWFfwoY1PMliw,8180
+digitalkin/services/identity/__init__.py,sha256=InkeyLgFYYwItx8mePA8HpfacOMWZwwuc0G4pWtKq9s,270
+digitalkin/services/identity/default_identity.py,sha256=Y2auZHrGSZTIN5D8HyjLvLcNbYFM1CNUE23x7p5VIGw,386
+digitalkin/services/identity/identity_strategy.py,sha256=skappBbds1_qa0Gr24FGrNX1N0_OYhYT1Lh7dUaAirE,429
+digitalkin/services/registry/__init__.py,sha256=Zl4QAkCe9tOmmKGBWVuLQVFepdZiL0ec3VDj27IeyYM,270
+digitalkin/services/registry/default_registry.py,sha256=VnWkF6nHpFxUKuUbZLPqzXqdA6oXmyV_ySpeuOCf_ko,277
+digitalkin/services/registry/registry_strategy.py,sha256=uBXgZIv25jeXbeVO8vWvlNPxxNYu7_KiCw2PoE6AWr8,423
+digitalkin/services/setup/default_setup.py,sha256=x1J6trXhyLC7V2OTe5pRY5mIFkQ4oyi3-aG50a1G1U4,8253
+digitalkin/services/setup/grpc_setup.py,sha256=Bo84gLZSK8DKqHAVSVFOQdtZQkfpp78eNgjJnGhXQfU,12449
+digitalkin/services/setup/setup_strategy.py,sha256=ZnJ_HwWCkHCPrqKekSD5L9y3p8wMwfjQ8sj2hLZq6go,4004
+digitalkin/services/snapshot/__init__.py,sha256=Uzlnzo0CYlSpVsdiI37hW7xQk8hu3YA1fOI6O6MSzB0,270
+digitalkin/services/snapshot/default_snapshot.py,sha256=Mb8QwWRsHh9I_tN0ln_ZiFa1QCZxOVWmuVLemQOTWpc,1058
+digitalkin/services/snapshot/snapshot_strategy.py,sha256=B1TU3V_k9A-OdqBkdyc41-ihnrW5Btcwd1KyQdHT46A,898
+digitalkin/services/storage/__init__.py,sha256=T-ocYLLphudkQgzvG47jBOm5GQsRFRIGA88y7Ur4akg,341
+digitalkin/services/storage/default_storage.py,sha256=bHNPm8nLvytKqKP2ntLkikvqH1qDKOwofrStVJH6PJg,7765
+digitalkin/services/storage/grpc_storage.py,sha256=b9z-2Q7mui8H0yn9GNb_RLu5RoLNFc55Dw2Xe_YklvY,4042
+digitalkin/services/storage/storage_strategy.py,sha256=vGo4aYkEp_GZV11m7vd-xY_Z3gVa5K0gMTzbj2Au_3o,6600
+digitalkin/utils/__init__.py,sha256=sJnY-ZUgsjMfojAjONC1VN14mhgIDnzyOlGkw21rRnM,28
+digitalkin/utils/arg_parser.py,sha256=3YyI6oZhhrlTmPTrzlwpQzbCNWDFAT3pggcLxNtJoc0,4388
+digitalkin-0.2.0.dist-info/licenses/LICENSE,sha256=Ies4HFv2r2hzDRakJYxk3Y60uDFLiG-orIgeTpstnIo,20327
+modules/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+modules/minimal_llm_module.py,sha256=JrRPTx7DCYrbmnfkE98xWBU_6x5lr7i8MFP16ATKkfo,5531
+modules/storage_module.py,sha256=HktIy4j8-0w7ugMA7HqcNZymZIGruhchsWKXje9qR4o,6264
+modules/text_transform_module.py,sha256=fAC6r_Ujca1Tz1qdWL4hTPZFn3gFWIVNj5-rytQMObE,7191
+digitalkin-0.2.0.dist-info/METADATA,sha256=692wgNWfg0bqDK5y2cA-UxS3QvpYUEAtBP1XeTno3QE,29095
+digitalkin-0.2.0.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+digitalkin-0.2.0.dist-info/top_level.txt,sha256=5_5e35inSM5YfWNZE21p5wGBojiVtQQML_WzbEk4BRU,31
+digitalkin-0.2.0.dist-info/RECORD,,

{digitalkin-0.1.1.dist-info → digitalkin-0.2.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (78.0.2)
+Generator: setuptools (78.1.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

digitalkin-0.2.0.dist-info/top_level.txt

@@ -0,0 +1,3 @@
+base_server
+digitalkin
+modules

modules/minimal_llm_module.py

@@ -0,0 +1,162 @@
+"""Simple module calling an LLM."""
+
+import logging
+from collections.abc import Callable
+from typing import Any, ClassVar
+
+import grpc
+import openai
+from pydantic import BaseModel
+
+from digitalkin.grpc_servers.utils.models import SecurityMode, ServerConfig, ServerMode
+from digitalkin.modules._base_module import BaseModule
+from digitalkin.services.setup.setup_strategy import SetupData
+
+# Configure logging with clear formatting
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+)
+logger = logging.getLogger(__name__)
+
+
+# Define schema models using Pydantic
+class OpenAIToolInput(BaseModel):
+    """Input model defining what data the module expects."""
+
+    prompt: str
+
+
+class OpenAIToolOutput(BaseModel):
+    """Output model defining what data the module produces."""
+
+    response: str
+
+
+class OpenAIToolSetup(BaseModel):
+    """Setup model defining module configuration parameters."""
+
+    openai_key: str
+    model_name: str
+    dev_prompt: str
+
+
+class OpenAIToolSecret(BaseModel):
+    """Secret model defining module configuration parameters."""
+
+
+server_config = ServerConfig(
+    host="[::]",
+    port=50151,
+    mode=ServerMode.ASYNC,
+    security=SecurityMode.INSECURE,
+    max_workers=10,
+    credentials=None,
+)
+
+
+class OpenAIToolModule(BaseModule[OpenAIToolInput, OpenAIToolOutput, OpenAIToolSetup, OpenAIToolSecret]):
+    """A openAI endpoint tool module module."""
+
+    name = "OpenAIToolModule"
+    description = "A module that interacts with OpenAI API to process text"
+
+    # Define the schema formats for the module
+    input_format = OpenAIToolInput
+    output_format = OpenAIToolOutput
+    setup_format = OpenAIToolSetup
+    secret_format = OpenAIToolSecret
+
+    openai_client: openai.OpenAI
+
+    # Define module metadata for discovery
+    metadata: ClassVar[dict[str, Any]] = {
+        "name": "Minimal_LLM_Tool",
+        "description": "Transforms input text using a streaming LLM response.",
+        "version": "1.0.0",
+        "tags": ["text", "transformation", "encryption", "streaming"],
+    }
+    # Define services_config_params with default values
+    services_config_strategies = {}
+    services_config_params = {
+        "storage": {
+            "config": {"setups": OpenAIToolSetup},
+            "server_config": server_config,
+        }
+    }
+
+    async def initialize(self, setup_data: SetupData) -> None:
+        """Initialize the module capabilities.
+
+        This method is called when the module is loaded by the server.
+        Use it to set up module-specific resources or configurations.
+        """
+        self.openai_client = openai.OpenAI(api_key=setup_data.current_setup_version.content["openai_key"])
+        # Define what capabilities this module provides
+        self.capabilities = ["text-processing", "streaming", "transformation"]
+        logger.info(
+            "Module %s initialized with capabilities: %s",
+            self.metadata["name"],
+            self.capabilities,
+        )
+
+    async def run(
+        self,
+        input_data: dict[str, Any],
+        setup_data: SetupData,
+        callback: Callable,
+    ) -> None:
+        """Process input text and stream LLM responses.
+
+        Args:
+            input_data: Contains the text to process.
+            setup_data: Contains model configuration and development prompt.
+            callback: Function to send output data back to the client.
+
+        Raises:
+            grpc.RpcError: If gRPC communication fails.
+            openai.AuthenticationError: If authentication with OpenAI fails.
+            openai.APIConnectionError: If an API connection error occurs.
+            Exception: For any unexpected runtime errors.
+        """
+        logger.info(
+            "Running job %s with prompt: '%s' on model: %s",
+            self.job_id,
+            input_data["prompt"],
+            setup_data.current_setup_version.content["model_name"],
+        )
+        try:
+            response = self.openai_client.responses.create(
+                model=setup_data.current_setup_version.content["model_name"],
+                tools=[{"type": "web_search_preview"}],
+                instructions=setup_data.current_setup_version.content["dev_prompt"],
+                input=input_data["prompt"],
+            )
+            if not response.output_text:
+                raise openai.APIConnectionError
+            output_data = OpenAIToolOutput(response=response.output_text).model_dump()
+
+        except openai.AuthenticationError as _:
+            message = "Authentication Error, OPENAI auth token was never set."
+            logger.exception(message)
+            output_data = {
+                "error": {
+                    "code": grpc.StatusCode.UNAUTHENTICATED,
+                    "error_message": message,
+                }
+            }
+        except openai.APIConnectionError as _:
+            message = "API Error, please try again."
+            logger.exception(message)
+            output_data = {"error": {"code": grpc.StatusCode.UNAVAILABLE, "error_message": message}}
+        await callback(job_id=self.job_id, output_data=output_data)
+        logger.info("Job %s completed", self.job_id)
+
+    async def cleanup(self) -> None:
+        """Clean up any resources when the module is stopped.
+
+        This method is called when the module is being shut down.
+        Use it to close connections, free resources, etc.
+        """
+        logger.info("Cleaning up module %s", self.metadata["name"])
+        # Release any resources here if needed.