digitalkin 0.2.15__tar.gz → 0.2.17__tar.gz

{digitalkin-0.2.15 → digitalkin-0.2.17}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: digitalkin
-Version: 0.2.15
+Version: 0.2.17
 Summary: SDK to build kin used in DigitalKin
 Author-email: "DigitalKin.ai" <contact@digitalkin.ai>
 License: Attribution-NonCommercial-ShareAlike 4.0 International

{digitalkin-0.2.15 → digitalkin-0.2.17}/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.15"
+    version = "0.2.17"
 
     classifiers = [
         "Development Status :: 3 - Alpha",
@@ -212,7 +212,7 @@
                 "ANN401",
                 "PLR0912",
                 "ANN201",
-                "D100"
+                "D100",
             ]
 
         [tool.ruff.lint.pylint]

{digitalkin-0.2.15 → digitalkin-0.2.17}/src/digitalkin/__version__.py

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

{digitalkin-0.2.15 → digitalkin-0.2.17}/src/digitalkin/grpc_servers/module_server.py

@@ -79,10 +79,10 @@ class ModuleServer(BaseServer):
 
     def start(self) -> None:
         """Start the module server and register with the registry if configured."""
-        logger.critical(self.server_config)
+        logger.info(self.server_config)
         super().start()
 
-        logger.critical(self.server_config)
+        logger.info(self.server_config)
         # If a registry address is provided, register the module
         if self.server_config.registry_address:
             try:
@@ -91,17 +91,17 @@ class ModuleServer(BaseServer):
                 logger.exception("Failed to register with registry")
 
         if self.module_servicer is not None:
-            logger.critical(
+            logger.info(
                 "Setup post init started with config: %s", self.client_config
             )
             self.module_servicer.setup.__post_init__(self.client_config)
 
     async def start_async(self) -> None:
         """Start the module server and register with the registry if configured."""
-        logger.critical(self.server_config)
+        logger.info(self.server_config)
         await super().start_async()
 
-        logger.critical(self.server_config)
+        logger.info(self.server_config)
         # If a registry address is provided, register the module
         if self.server_config.registry_address:
             try:
@@ -110,7 +110,7 @@ class ModuleServer(BaseServer):
                 logger.exception("Failed to register with registry")
 
         if self.module_servicer is not None:
-            logger.critical(
+            logger.info(
                 "Setup post init started with config: %s", self.client_config
             )
             await self.module_servicer.job_manager._start()

{digitalkin-0.2.15 → digitalkin-0.2.17}/src/digitalkin/grpc_servers/module_servicer.py

@@ -70,6 +70,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
             module_class: The module type to serve.
         """
         super().__init__()
+        module_class.discover()
         self.module_class = module_class
         job_manager_class = self.args.job_manager_mode.get_manager_class()
         self.job_manager = job_manager_class(module_class, self.args.services_mode)

digitalkin-0.2.17/src/digitalkin/models/module/__init__.py

@@ -0,0 +1,26 @@
+"""This module contains the models for the modules."""
+
+from digitalkin.models.module.module import Module, ModuleStatus
+from digitalkin.models.module.module_context import ModuleContext
+from digitalkin.models.module.module_types import (
+    ConfigSetupModelT,
+    InputModel,
+    InputModelT,
+    InputTrigger,
+    OutputModelT,
+    SecretModelT,
+    SetupModelT,
+)
+
+__all__ = [
+    "ConfigSetupModelT",
+    "InputModel",
+    "InputModelT",
+    "InputTrigger",
+    "Module",
+    "ModuleContext",
+    "ModuleStatus",
+    "OutputModelT",
+    "SecretModelT",
+    "SetupModelT",
+]

digitalkin-0.2.17/src/digitalkin/models/module/module_context.py

@@ -0,0 +1,24 @@
+"""Define the module context used in the triggers."""
+
+from types import SimpleNamespace
+
+from digitalkin.services.cost.cost_strategy import CostStrategy
+from digitalkin.services.filesystem.filesystem_strategy import FilesystemStrategy
+from digitalkin.services.storage.storage_strategy import StorageStrategy
+
+
+class ModuleContext(SimpleNamespace):
+    """ModuleContext provides a container for strategies and resources used by a module.
+
+    Attributes:
+        cost (CostStrategy): The strategy used to calculate or manage costs within the module.
+        filesystem (FilesystemStrategy): The strategy for interacting with the filesystem.
+        storage (StorageStrategy): The strategy for handling storage operations.
+
+    This context object is designed to be passed to module components, providing them with
+    access to shared strategies and resources. Additional attributes may be set dynamically.
+    """
+
+    cost: CostStrategy
+    filesystem: FilesystemStrategy
+    storage: StorageStrategy

digitalkin-0.2.17/src/digitalkin/models/module/module_types.py

@@ -0,0 +1,43 @@
+"""Types for module models."""
+
+from typing import TypeVar
+
+from pydantic import BaseModel
+
+
+class InputTrigger(BaseModel):
+    """Defines the root input model exposing the protocol.
+
+    The mandatory protocol is important to define the module beahvior following the user or agent input.
+
+    Example:
+        class MyInput(InputModel):
+            root: InputTrigger
+            user_define_data: Any
+
+        # Usage
+        my_input = MyInput(root=InputTrigger(protocol="message"))
+        print(my_input.root.protocol)  # Output: message
+    """
+
+    protocol: str
+
+
+class InputModel(BaseModel):
+    """Base definition of input model showing mandatory root fields.
+
+    The Model define the Module Input, usually referring to multiple input type defined by an union.
+
+    Example:
+        class ModuleInput(InputModel):
+            root: FileInput | MessageInput
+    """
+
+    root: InputTrigger
+
+
+ConfigSetupModelT = TypeVar("ConfigSetupModelT", bound=BaseModel | None)
+InputModelT = TypeVar("InputModelT", bound=InputModel)
+OutputModelT = TypeVar("OutputModelT", bound=BaseModel)
+SetupModelT = TypeVar("SetupModelT", bound=BaseModel)
+SecretModelT = TypeVar("SecretModelT", bound=BaseModel)

{digitalkin-0.2.15 → digitalkin-0.2.17}/src/digitalkin/modules/__init__.py

@@ -2,6 +2,6 @@
 
 from digitalkin.modules.archetype_module import ArchetypeModule
 from digitalkin.modules.tool_module import ToolModule
-from digitalkin.modules.trigger_module import TriggerModule
+from digitalkin.modules.trigger_handler import TriggerHandler
 
-__all__ = ["ArchetypeModule", "ToolModule", "TriggerModule"]
+__all__ = ["ArchetypeModule", "ToolModule", "TriggerHandler"]

{digitalkin-0.2.15 → digitalkin-0.2.17}/src/digitalkin/modules/_base_module.py

@@ -19,6 +19,8 @@ from digitalkin.models.module import (
     SecretModelT,
     SetupModelT,
 )
+from digitalkin.models.module.module_context import ModuleContext
+from digitalkin.modules.trigger_handler import TriggerHandler
 from digitalkin.services.agent.agent_strategy import AgentStrategy
 from digitalkin.services.cost.cost_strategy import CostStrategy
 from digitalkin.services.filesystem.filesystem_strategy import FilesystemStrategy
@@ -28,17 +30,18 @@ from digitalkin.services.services_config import ServicesConfig, ServicesStrategy
 from digitalkin.services.snapshot.snapshot_strategy import SnapshotStrategy
 from digitalkin.services.storage.storage_strategy import StorageStrategy
 from digitalkin.utils.llm_ready_schema import llm_ready_schema
+from digitalkin.utils.package_discover import ModuleDiscoverer
 
 
 class ModuleErrorModel(BaseModel):
-    """Typed error/code model."""
+    """typed error/code model."""
 
     code: str
     exception: str
     short_description: str
 
 
-class BaseModule(
+class BaseModule(  # noqa: PLR0904
     ABC,
     Generic[
         InputModelT,
@@ -60,6 +63,9 @@ class BaseModule(
     secret_format: type[SecretModelT]
     metadata: ClassVar[dict[str, Any]]
 
+    context: ModuleContext
+    triggers_discoverer: ClassVar[ModuleDiscoverer]
+
     # service config params
     services_config_strategies: ClassVar[dict[str, ServicesStrategy | None]]
     services_config_params: ClassVar[dict[str, dict[str, Any | None] | None]]
@@ -95,6 +101,13 @@ class BaseModule(
         # Initialize services configuration
         self._init_strategies()
 
+        # Initialize minimum context
+        self.context = ModuleContext(
+            storage=self.storage,
+            cost=self.cost,
+            filesystem=self.filesystem,
+        )
+
     @property
     def status(self) -> ModuleStatus:
         """Get the module status.
@@ -165,10 +178,12 @@ class BaseModule(
         Returns:
             The JSON schema of the config setup format as a string.
         """
-        if cls.config_setup_format is not None:
+        config_setup_format = getattr(cls, "config_setup_format", None)
+
+        if config_setup_format is not None:
             if llm_format:
-                return json.dumps(llm_ready_schema(cls.config_setup_format), indent=2)
-            return json.dumps(cls.config_setup_format.model_json_schema(), indent=2)
+                return json.dumps(llm_ready_schema(config_setup_format), indent=2)
+            return json.dumps(config_setup_format.model_json_schema(), indent=2)
         msg = "'%s' class does not define an 'config_setup_format'."
         raise OptionalFeatureNotImplementedError(msg)
 
@@ -249,6 +264,32 @@ class BaseModule(
         """
         return cls.output_format(**output_data)
 
+    @classmethod
+    def discover(cls) -> None:
+        """Discover and register all TriggerHandler subclasses in the specified package or current directory.
+
+        Dynamically import all Python modules in the specified package or current directory,
+        triggering class registrations for subclasses of TriggerHandler whose names end with 'Trigger'.
+
+        If a package is provided, all .py files within its path are imported; otherwise, the current
+        working directory is searched. For each imported module, any class matching the criteria is
+        registered via cls.register(). Errors during import are logged at debug level.
+        """
+        cls.triggers_discoverer.discover_modules()
+        logger.debug("discovered: %s", cls.triggers_discoverer)
+
+    @classmethod
+    def register(cls, handler_cls: type[TriggerHandler]) -> type[TriggerHandler]:
+        """Dynamically register the trigger class.
+
+        Args:
+            handler_cls: type of the trigger handler to register.
+
+        Returns:
+            type of the trigger handler.
+        """
+        return cls.triggers_discoverer.register_trigger(handler_cls)
+
     @abstractmethod
     async def run_config_setup(
         self,
@@ -269,15 +310,42 @@ class BaseModule(
         """Initialize the module."""
         raise NotImplementedError
 
-    @abstractmethod
     async def run(
         self,
         input_data: InputModelT,
         setup_data: SetupModelT,
         callback: Callable[[OutputModelT], Coroutine[Any, Any, None]],
     ) -> None:
-        """Run the module."""
-        raise NotImplementedError
+        """Run the module with the given input and setup data.
+
+        This method validates the input data, determines the protocol from the input,
+        and dispatches the request to the corresponding trigger handler. The trigger handler
+        is responsible for processing the input and invoking the callback with the result.
+
+        Triggers:
+            - The method is triggered when a module run is requested with specific input and setup data.
+            - The protocol specified in the input determines which trigger handler is invoked.
+
+        Args:
+            input_data (InputModelT): The input data to be processed by the module.
+            setup_data (SetupModelT): The setup or configuration data required for the module.
+            callback (Callable[[OutputModelT], Coroutine[Any, Any, None]]): callback to be invoked to stream any result.
+
+        Raises:
+            ValueError: If no handler for the protocol is found.
+        """
+        input_instance = self.input_format.model_validate(input_data)
+        handler_instance = self.triggers_discoverer.get_trigger(
+            input_instance.root.protocol,
+            input_instance.root,
+        )
+
+        await handler_instance.handle(
+            input_instance.root,
+            setup_data,
+            callback,
+            self.context,
+        )
 
     @abstractmethod
     async def cleanup(self) -> None:
@@ -338,6 +406,8 @@ class BaseModule(
             return
 
         try:
+            logger.info("Init the discod input handlers.")
+            self.triggers_discoverer.init_handlers(self.context)
             logger.info("Run lifecycle")
             self._status = ModuleStatus.RUNNING
             self._task = asyncio.create_task(

digitalkin-0.2.17/src/digitalkin/modules/trigger_handler.py

@@ -0,0 +1,47 @@
+"""Definition of the Trigger type."""
+
+from abc import ABC, abstractmethod
+from collections.abc import Callable, Coroutine
+from typing import Any, ClassVar, Generic
+
+from digitalkin.models.module.module_types import InputModelT, OutputModelT, SetupModelT
+from digitalkin.modules._base_module import ModuleContext
+
+
+class TriggerHandler(ABC, Generic[InputModelT, SetupModelT, OutputModelT]):
+    """Base class for all input-trigger handlers.
+
+    Each handler declares:
+      - protocol_key: the Literal value this handler processes
+      - handle(): logic to process the validated payload
+    """
+
+    protocol: ClassVar[str]
+    input_format: type[InputModelT]
+    output_format: type[OutputModelT]
+
+    def __init__(self, context: ModuleContext) -> None:
+        """Initialize the TriggerHandler with the given context."""
+
+    @abstractmethod
+    async def handle(
+        self,
+        input_data: InputModelT,
+        setup_data: SetupModelT,
+        callback: Callable[[Any], Coroutine[Any, Any, None]],
+        context: ModuleContext,
+    ) -> None:
+        """Asynchronously processes the input data specific to Handler and streams results via the provided callback.
+
+        Args:
+            input_data (InputModelT): The input data to be processed by the handler.
+            setup_data (SetupModelT): The setup or configuration data required for processing.
+            callback (Callable[[Any], Coroutine[Any, Any, None]]): callback that stream results.
+            context (ModuleContext): The context object containing module-specific information and resources.
+
+        Returns:
+            Any: The result of the processing, if applicable.
+
+        Note:
+            The callback must be awaited to ensure results are streamed correctly during processing.
+        """

{digitalkin-0.2.15 → digitalkin-0.2.17}/src/digitalkin/services/filesystem/default_filesystem.py

@@ -5,7 +5,7 @@ import os
 import tempfile
 import uuid
 from pathlib import Path
-from typing import Any
+from typing import Any, Literal
 
 from digitalkin.logger import logger
 from digitalkin.services.filesystem.filesystem_strategy import (
@@ -259,7 +259,17 @@ class DefaultFilesystem(FilesystemStrategy):
         self,
         file_id: str,
         content: bytes | None = None,
-        file_type: str | None = None,
+        file_type: Literal[
+            "UNSPECIFIED",
+            "DOCUMENT",
+            "IMAGE",
+            "VIDEO",
+            "AUDIO",
+            "ARCHIVE",
+            "CODE",
+            "OTHER",
+        ]
+        | None = None,
         content_type: str | None = None,
         metadata: dict[str, Any] | None = None,
         new_name: str | None = None,

{digitalkin-0.2.15 → digitalkin-0.2.17}/src/digitalkin/services/filesystem/filesystem_strategy.py

@@ -2,7 +2,7 @@
 
 from abc import ABC, abstractmethod
 from datetime import datetime
-from typing import Any
+from typing import Any, Literal
 
 from pydantic import BaseModel, Field
 
@@ -34,7 +34,21 @@ class FileFilter(BaseModel):
 
     names: list[str] | None = Field(default=None, description="Filter by file names (exact matches)")
     file_ids: list[str] | None = Field(default=None, description="Filter by file IDs")
-    file_types: list[str] | None = Field(default=None, description="Filter by file types")
+    file_types: (
+        list[
+            Literal[
+                "UNSPECIFIED",
+                "DOCUMENT",
+                "IMAGE",
+                "AUDIO",
+                "VIDEO",
+                "ARCHIVE",
+                "CODE",
+                "OTHER",
+            ]
+        ]
+        | None
+    ) = Field(default=None, description="Filter by file types")
     created_after: datetime | None = Field(default=None, description="Filter files created after this timestamp")
     created_before: datetime | None = Field(default=None, description="Filter files created before this timestamp")
     updated_after: datetime | None = Field(default=None, description="Filter files updated after this timestamp")
@@ -52,7 +66,16 @@ class UploadFileData(BaseModel):
 
     content: bytes = Field(description="The content of the file")
     name: str = Field(description="The name of the file")
-    file_type: str = Field(description="The type of the file")
+    file_type: Literal[
+        "UNSPECIFIED",
+        "DOCUMENT",
+        "IMAGE",
+        "AUDIO",
+        "VIDEO",
+        "ARCHIVE",
+        "CODE",
+        "OTHER",
+    ] = Field(description="The type of the file")
     content_type: str | None = Field(default=None, description="The content type of the file")
     metadata: dict[str, Any] | None = Field(default=None, description="The metadata of the file")
     replace_if_exists: bool = Field(default=False, description="Whether to replace the file if it already exists")
@@ -153,7 +176,17 @@ class FilesystemStrategy(BaseStrategy, ABC):
         self,
         file_id: str,
         content: bytes | None = None,
-        file_type: str | None = None,
+        file_type: Literal[
+            "UNSPECIFIED",
+            "DOCUMENT",
+            "IMAGE",
+            "VIDEO",
+            "AUDIO",
+            "ARCHIVE",
+            "CODE",
+            "OTHER",
+        ]
+        | None = None,
         content_type: str | None = None,
         metadata: dict[str, Any] | None = None,
         new_name: str | None = None,

{digitalkin-0.2.15 → digitalkin-0.2.17}/src/digitalkin/services/filesystem/grpc_filesystem.py

@@ -2,7 +2,7 @@
 
 from collections.abc import Generator
 from contextlib import contextmanager
-from typing import Any
+from typing import Any, Literal
 
 from digitalkin_proto.digitalkin.filesystem.v1 import filesystem_pb2, filesystem_service_pb2_grpc
 from google.protobuf import struct_pb2
@@ -103,12 +103,11 @@ class GrpcFilesystem(FilesystemStrategy, GrpcClientWrapper):
             status=self._file_status_to_enum(filters.status) if filters.status else None,
         )
 
-    def _file_proto_to_data(self, file: filesystem_pb2.File, content: bytes | None = None) -> FilesystemRecord:
+    def _file_proto_to_data(self, file: filesystem_pb2.File) -> FilesystemRecord:
         """Convert a File proto message to FilesystemRecord.
 
         Args:
             file: The File proto message to convert
-            content: The content of the file
 
         Returns:
             FilesystemRecord: The converted data
@@ -124,7 +123,7 @@ class GrpcFilesystem(FilesystemStrategy, GrpcClientWrapper):
             metadata=MessageToDict(file.metadata),
             storage_url=file.storage_url,
             status=filesystem_pb2.FileStatus.Name(file.status),
-            content=content,
+            content=file.content,
         )
 
     def __init__(
@@ -213,13 +212,23 @@ class GrpcFilesystem(FilesystemStrategy, GrpcClientWrapper):
 
             response: filesystem_pb2.GetFileResponse = self.exec_grpc_query("GetFile", request)
 
-            return self._file_proto_to_data(response.file, response.content)
+            return self._file_proto_to_data(response.file)
 
     def update_file(
         self,
         file_id: str,
         content: bytes | None = None,
-        file_type: str | None = None,
+        file_type: Literal[
+            "UNSPECIFIED",
+            "DOCUMENT",
+            "IMAGE",
+            "VIDEO",
+            "AUDIO",
+            "ARCHIVE",
+            "CODE",
+            "OTHER",
+        ]
+        | None = None,
         content_type: str | None = None,
         metadata: dict[str, Any] | None = None,
         new_name: str | None = None,