digitalkin 0.2.12__py3-none-any.whl → 0.2.14__py3-none-any.whl

digitalkin/grpc_servers/registry_server.py

@@ -1,7 +1,5 @@
 """Registry gRPC server implementation for DigitalKin."""
 
-import logging
-
 from digitalkin_proto.digitalkin.module_registry.v2 import (
     module_registry_service_pb2,
     module_registry_service_pb2_grpc,
@@ -10,8 +8,7 @@ from digitalkin_proto.digitalkin.module_registry.v2 import (
 from digitalkin.grpc_servers._base_server import BaseServer
 from digitalkin.grpc_servers.registry_servicer import RegistryModule, RegistryServicer
 from digitalkin.grpc_servers.utils.models import RegistryServerConfig
-
-logger = logging.getLogger(__name__)
+from digitalkin.logger import logger
 
 
 class RegistryServer(BaseServer):
@@ -48,14 +45,14 @@ class RegistryServer(BaseServer):
             msg = "Server must be created before registering servicers"
             raise RuntimeError(msg)
 
-        logger.info("Registering registry servicer")
+        logger.debug("Registering registry servicer")
         self.registry_servicer = RegistryServicer()
         self.register_servicer(
             self.registry_servicer,
             module_registry_service_pb2_grpc.add_ModuleRegistryServiceServicer_to_server,
             service_descriptor=module_registry_service_pb2.DESCRIPTOR,
         )
-        logger.info("Registered registry servicer")
+        logger.debug("Registered registry servicer")
 
     def get_registered_modules(self) -> list[RegistryModule]:
         """Get a list of all registered modules.

digitalkin/grpc_servers/registry_servicer.py

@@ -5,7 +5,6 @@ which handles registration, deregistration, discovery, and status management
 of DigitalKin modules.
 """
 
-import logging
 from collections.abc import Iterator
 from enum import Enum
 
@@ -20,7 +19,7 @@ from digitalkin_proto.digitalkin.module_registry.v2 import (
 from pydantic import BaseModel
 from typing_extensions import Self
 
-logger = logging.getLogger(__name__)
+from digitalkin.logger import logger
 
 
 class ExtendedEnum(Enum):
@@ -184,7 +183,7 @@ class RegistryServicer(module_registry_service_pb2_grpc.ModuleRegistryServiceSer
             registration_pb2.RegisterResponse: A response indicating success or failure.
         """
         module_id = request.module_id
-        logger.info("Registering module: %s", module_id)
+        logger.debug("Registering module: %s", module_id)
 
         # Check if module is already registered
         if module_id in self.registered_modules:
@@ -207,7 +206,7 @@ class RegistryServicer(module_registry_service_pb2_grpc.ModuleRegistryServiceSer
             message=None,
         )
 
-        logger.info("Module %s registered at %s:%d", module_id, request.address, request.port)
+        logger.debug("Module %s registered at %s:%d", module_id, request.address, request.port)
         return registration_pb2.RegisterResponse(success=True)
 
     def DeregisterModule(  # noqa: N802
@@ -228,7 +227,7 @@ class RegistryServicer(module_registry_service_pb2_grpc.ModuleRegistryServiceSer
             registration_pb2.DeregisterResponse: A response indicating success or failure.
         """
         module_id = request.module_id
-        logger.info("Deregistering module: %s", module_id)
+        logger.debug("Deregistering module: %s", module_id)
 
         # Check if module exists in registry
         if module_id not in self.registered_modules:
@@ -242,7 +241,7 @@ class RegistryServicer(module_registry_service_pb2_grpc.ModuleRegistryServiceSer
         # Remove the module
         del self.registered_modules[module_id]
 
-        logger.info("Module %s deregistered", module_id)
+        logger.debug("Module %s deregistered", module_id)
         return registration_pb2.DeregisterResponse(success=True)
 
     def DiscoverInfoModule(  # noqa: N802
@@ -261,7 +260,7 @@ class RegistryServicer(module_registry_service_pb2_grpc.ModuleRegistryServiceSer
         Returns:
             discover_pb2.DiscoverInfoResponse: A response containing the module's information.
         """
-        logger.info("Discovering module: %s", request.module_id)
+        logger.debug("Discovering module: %s", request.module_id)
 
         # Check if module exists in registry
         if request.module_id not in self.registered_modules:
@@ -289,24 +288,24 @@ class RegistryServicer(module_registry_service_pb2_grpc.ModuleRegistryServiceSer
         Returns:
             discover_pb2.DiscoverSearchResponse: A response containing matching modules.
         """
-        logger.info("Discovering modules with criteria:")
+        logger.debug("Discovering modules with criteria:")
 
         # Start with all modules
         results = list(self.registered_modules.values())
-        logger.info("%s", list(results))
+        logger.debug("%s", list(results))
         # Filter by name if specified
         if request.name:
-            logger.info("\tname %s", request.name)
+            logger.debug("\tname %s", request.name)
             results = [m for m in results if request.name in m.metadata.name]
 
         # Filter by type if specified
         if request.module_type:
-            logger.info("\tmodule_type %s", request.module_type)
+            logger.debug("\tmodule_type %s", request.module_type)
             results = [m for m in results if m.module_type == request.module_type]
 
         # Filter by tags if specified
         if request.tags:
-            logger.info("\ttags %s", request.tags)
+            logger.debug("\ttags %s", request.tags)
             results = [m for m in results if any(tag in m.metadata.tags for tag in request.tags)]
 
         # Filter by description if specified
@@ -315,7 +314,7 @@ class RegistryServicer(module_registry_service_pb2_grpc.ModuleRegistryServiceSer
             results = [m for m in results if request.description in m.metadata.description]
         """
 
-        logger.info("Found %d matching modules", len(results))
+        logger.debug("Found %d matching modules", len(results))
         return discover_pb2.DiscoverSearchResponse(modules=[r.to_proto() for r in results])
 
     def GetModuleStatus(  # noqa: N802
@@ -334,7 +333,7 @@ class RegistryServicer(module_registry_service_pb2_grpc.ModuleRegistryServiceSer
         Returns:
             status_pb2.ModuleStatusResponse: A response containing the module's status.
         """
-        logger.info("Getting status for module: %s", request.module_id)
+        logger.debug("Getting status for module: %s", request.module_id)
 
         # Check if module exists in registry
         if request.module_id not in self.registered_modules:
@@ -363,7 +362,7 @@ class RegistryServicer(module_registry_service_pb2_grpc.ModuleRegistryServiceSer
         Returns:
             status_pb2.ListModulesStatusResponse: A response containing a list of module statuses.
         """
-        logger.info(
+        logger.debug(
             "Getting registered modules with offset %d and limit %d",
             request.offset,
             request.list_size,
@@ -384,7 +383,7 @@ class RegistryServicer(module_registry_service_pb2_grpc.ModuleRegistryServiceSer
             for module in list(self.registered_modules.values())[request.offset : request.offset + list_size]
         ]
 
-        logger.info("Found %d registered modules", len(modules_statuses))
+        logger.debug("Found %d registered modules", len(modules_statuses))
         return status_pb2.ListModulesStatusResponse(
             list_size=len(modules_statuses),
             modules_statuses=modules_statuses,
@@ -406,7 +405,7 @@ class RegistryServicer(module_registry_service_pb2_grpc.ModuleRegistryServiceSer
         Yields:
             status_pb2.ModuleStatusResponse: Responses containing individual module statuses.
         """
-        logger.info("Streaming all %d registered modules", len(self.registered_modules))
+        logger.debug("Streaming all %d registered modules", len(self.registered_modules))
         for module in self.registered_modules.values():
             yield status_pb2.ModuleStatusResponse(
                 module_id=module.module_id,
@@ -431,7 +430,7 @@ class RegistryServicer(module_registry_service_pb2_grpc.ModuleRegistryServiceSer
             status_pb2.UpdateStatusResponse: A response indicating success or failure.
         """
         module_id = request.module_id
-        logger.info("Updating status for module: %s to %s", module_id, request.status)
+        logger.debug("Updating status for module: %s to %s", module_id, request.status)
 
         # Check if module exists in registry
         if request.module_id not in self.registered_modules:
@@ -453,5 +452,5 @@ class RegistryServicer(module_registry_service_pb2_grpc.ModuleRegistryServiceSer
         module_info = self.registered_modules[module_id]
         module_info.status = ModuleStatus(request.status)
 
-        logger.info("Status for module %s updated to %s", module_id, request.status)
+        logger.debug("Status for module %s updated to %s", module_id, request.status)
         return status_pb2.UpdateStatusResponse(success=True)

digitalkin/grpc_servers/utils/exceptions.py

@@ -31,3 +31,7 @@ class ReflectionError(ServerError):
 
 class HealthCheckError(ServerError):
     """Error related to gRPC health check service."""
+
+
+class OptionalFeatureNotImplementedError(NotImplementedError):
+    """Raised when an optional feature is not implemented, but was requested."""

digitalkin/grpc_servers/utils/grpc_client_wrapper.py

@@ -1,6 +1,5 @@
 """Client wrapper to ease channel creation with specific ServerConfig."""
 
-import logging
 from pathlib import Path
 from typing import Any
 
@@ -8,8 +7,7 @@ import grpc
 
 from digitalkin.grpc_servers.utils.exceptions import ServerError
 from digitalkin.grpc_servers.utils.models import ClientConfig, SecurityMode
-
-logger = logging.getLogger(__name__)
+from digitalkin.logger import logger
 
 
 class GrpcClientWrapper:
@@ -64,9 +62,9 @@ class GrpcClientWrapper:
         """
         try:
             # Call the register method
-            logger.warning("send request to %s", query_endpoint)
+            logger.debug("send request to %s", query_endpoint)
             response = getattr(self.stub, query_endpoint)(request)
-            logger.warning("recive response from request to registry: %s", response)
+            logger.debug("receive response from request to registry: %s", response)
         except grpc.RpcError:
             logger.exception("RPC error during registration:")
             raise ServerError

digitalkin/logger.py

@@ -2,11 +2,46 @@
 
 import logging
 import sys
+from typing import ClassVar
+
+
+class ColorFormatter(logging.Formatter):
+    """Color formatter for logging."""
+
+    grey = "\x1b[38;20m"
+    green = "\x1b[32;20m"
+    blue = "\x1b[34;20m"
+    yellow = "\x1b[33;20m"
+    red = "\x1b[31;20m"
+    bold_red = "\x1b[31;1m"
+    reset = "\x1b[0m"
+    format: str = "%(asctime)s - %(name)s - %(levelname)s - %(message)s (%(filename)s:%(lineno)d)"  # type: ignore
+
+    FORMATS: ClassVar[dict[int, str]] = {
+        logging.DEBUG: grey + format + reset + "\n",  # type: ignore
+        logging.INFO: green + format + reset + "\n",  # type: ignore
+        logging.WARNING: yellow + format + reset + "\n",  # type: ignore
+        logging.ERROR: red + format + reset + "\n",  # type: ignore
+        logging.CRITICAL: bold_red + format + reset + "\n",  # type: ignore
+    }
+
+    def format(self, record: logging.LogRecord) -> str:  # type: ignore
+        """Format the log record.
+
+        Args:
+            record: The log record to format.
+
+        Returns:
+            str: The formatted log record.
+        """
+        log_fmt = self.FORMATS.get(record.levelno)
+        formatter = logging.Formatter(log_fmt)
+        return formatter.format(record)
+
 
 logging.basicConfig(
     level=logging.DEBUG,
     stream=sys.stdout,
-    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
     datefmt="%Y-%m-%d %H:%M:%S",
 )
 
@@ -15,3 +50,12 @@ logging.getLogger("asyncio").setLevel(logging.DEBUG)
 
 
 logger = logging.getLogger("digitalkin")
+
+if not logger.handlers:
+    ch = logging.StreamHandler()
+    ch.setLevel(logging.INFO)
+
+    ch.setFormatter(ColorFormatter())
+
+    logger.addHandler(ch)
+    logger.propagate = False

digitalkin/models/module/__init__.py

@@ -2,10 +2,11 @@
 
 from digitalkin.models.module.module import Module, ModuleStatus
 from digitalkin.models.module.module_types import (
+    ConfigSetupModelT,
     InputModelT,
     OutputModelT,
     SecretModelT,
     SetupModelT,
 )
 
-__all__ = ["InputModelT", "Module", "ModuleStatus", "OutputModelT", "SecretModelT", "SetupModelT"]
+__all__ = ["ConfigSetupModelT", "InputModelT", "Module", "ModuleStatus", "OutputModelT", "SecretModelT", "SetupModelT"]

digitalkin/models/module/module.py

@@ -14,6 +14,7 @@ class ModuleStatus(Enum):
     STOPPING = auto()  # Module is stopping
     STOPPED = auto()  # Module stop successfuly
     FAILED = auto()  # Module stopped due to internal error
+    CANCELLED = auto()  # Module stopped due to internal error
     NOT_FOUND = auto()
 
 

digitalkin/models/module/module_types.py

@@ -4,6 +4,7 @@ from typing import TypeVar
 
 from pydantic import BaseModel
 
+ConfigSetupModelT = TypeVar("ConfigSetupModelT", bound=BaseModel | None)
 InputModelT = TypeVar("InputModelT", bound=BaseModel)
 OutputModelT = TypeVar("OutputModelT", bound=BaseModel)
 SetupModelT = TypeVar("SetupModelT", bound=BaseModel)

digitalkin/modules/_base_module.py

@@ -7,8 +7,18 @@ from abc import ABC, abstractmethod
 from collections.abc import Callable, Coroutine
 from typing import Any, ClassVar, Generic
 
+from pydantic import BaseModel
+
+from digitalkin.grpc_servers.utils.exceptions import OptionalFeatureNotImplementedError
 from digitalkin.logger import logger
-from digitalkin.models.module import InputModelT, ModuleStatus, OutputModelT, SecretModelT, SetupModelT
+from digitalkin.models.module import (
+    ConfigSetupModelT,
+    InputModelT,
+    ModuleStatus,
+    OutputModelT,
+    SecretModelT,
+    SetupModelT,
+)
 from digitalkin.services.agent.agent_strategy import AgentStrategy
 from digitalkin.services.cost.cost_strategy import CostStrategy
 from digitalkin.services.filesystem.filesystem_strategy import FilesystemStrategy
@@ -20,11 +30,30 @@ from digitalkin.services.storage.storage_strategy import StorageStrategy
 from digitalkin.utils.llm_ready_schema import llm_ready_schema
 
 
-class BaseModule(ABC, Generic[InputModelT, OutputModelT, SetupModelT, SecretModelT]):
+class ModuleErrorModel(BaseModel):
+    """Typed error/code model."""
+
+    code: str
+    exception: str
+    short_description: str
+
+
+class BaseModule(
+    ABC,
+    Generic[
+        InputModelT,
+        OutputModelT,
+        SetupModelT,
+        SecretModelT,
+        ConfigSetupModelT,
+    ],
+):
     """BaseModule is the abstract base for all modules in the DigitalKin SDK."""
 
     name: str
     description: str
+
+    config_setup_format: type[ConfigSetupModelT]
     input_format: type[InputModelT]
     output_format: type[OutputModelT]
     setup_format: type[SetupModelT]
@@ -126,6 +155,23 @@ class BaseModule(ABC, Generic[InputModelT, OutputModelT, SetupModelT, SecretMode
         msg = "'%s' class does not define an 'output_format'."
         raise NotImplementedError(msg)
 
+    @classmethod
+    def get_config_setup_format(cls, *, llm_format: bool) -> str:
+        """Gets the JSON schema of the config setup format model.
+
+        Raises:
+            OptionalFeatureNotImplementedError: If the `config_setup_format` is not defined.
+
+        Returns:
+            The JSON schema of the config setup format as a string.
+        """
+        if cls.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)
+        msg = "'%s' class does not define an 'config_setup_format'."
+        raise OptionalFeatureNotImplementedError(msg)
+
     @classmethod
     def get_setup_format(cls, *, llm_format: bool) -> str:
         """Gets the JSON schema of the setup format model.
@@ -143,6 +189,18 @@ class BaseModule(ABC, Generic[InputModelT, OutputModelT, SetupModelT, SecretMode
         msg = "'%s' class does not define an 'setup_format'."
         raise NotImplementedError(msg)
 
+    @classmethod
+    def create_config_setup_model(cls, config_setup_data: dict[str, Any]) -> ConfigSetupModelT:
+        """Create the setup model from the setup data.
+
+        Args:
+            config_setup_data: The setup data to create the model from.
+
+        Returns:
+            The setup model.
+        """
+        return cls.config_setup_format(**config_setup_data)
+
     @classmethod
     def create_input_model(cls, input_data: dict[str, Any]) -> InputModelT:
         """Create the input model from the input data.
@@ -191,6 +249,21 @@ class BaseModule(ABC, Generic[InputModelT, OutputModelT, SetupModelT, SecretMode
         """
         return cls.output_format(**output_data)
 
+    @abstractmethod
+    async def run_config_setup(
+        self,
+        config_setup_data: ConfigSetupModelT,
+        setup_data: SetupModelT,
+        callback: Callable,
+    ) -> None:
+        """Run config setup the module.
+
+        Raises:
+            OptionalFeatureNotImplementedError: If the config setup feature is not implemented.
+        """
+        msg = f"'{self}' class does not define an optional 'run_config_setup' attribute."
+        raise OptionalFeatureNotImplementedError(msg)
+
     @abstractmethod
     async def initialize(self, setup_data: SetupModelT) -> None:
         """Initialize the module."""
@@ -223,30 +296,59 @@ class BaseModule(ABC, Generic[InputModelT, OutputModelT, SetupModelT, SecretMode
             asyncio.CancelledError: If the module is cancelled
         """
         try:
+            logger.warning("Starting module %s", self.name)
             await self.run(input_data, setup_data, callback)
-            await self.stop()
+            logger.warning("Module %s finished", self.name)
         except asyncio.CancelledError:
-            logger.info(f"Module {self.name} cancelled")
+            self._status = ModuleStatus.CANCELLED
+            logger.error(f"Module {self.name} cancelled")
         except Exception:
             self._status = ModuleStatus.FAILED
             logger.exception("Error inside module %s", self.name)
         else:
             self._status = ModuleStatus.STOPPED
+        finally:
+            await self.stop()
 
     async def start(
         self,
         input_data: InputModelT,
         setup_data: SetupModelT,
-        callback: Callable[[OutputModelT], Coroutine[Any, Any, None]],
+        callback: Callable[[OutputModelT | ModuleErrorModel], Coroutine[Any, Any, None]],
+        done_callback: Callable | None = None,
     ) -> None:
         """Start the module."""
         try:
+            logger.info("Inititalize module")
             await self.initialize(setup_data=setup_data)
+        except Exception as e:
+            self._status = ModuleStatus.FAILED
+            short_description = "Error initializing module"
+            logger.exception("%s: %s", short_description, e)
+            await callback(
+                ModuleErrorModel(
+                    code=str(self._status),
+                    short_description=short_description,
+                    exception=str(e),
+                )
+            )
+            if done_callback is not None:
+                await done_callback(None)
+            await self.stop()
+            return
+
+        try:
+            logger.info("Run lifecycle")
             self._status = ModuleStatus.RUNNING
-            self._task = asyncio.create_task(self._run_lifecycle(input_data, setup_data, callback))
+            self._task = asyncio.create_task(
+                self._run_lifecycle(input_data, setup_data, callback),
+                name="module_lifecycle",
+            )
+            if done_callback is not None:
+                self._task.add_done_callback(done_callback)
         except Exception:
             self._status = ModuleStatus.FAILED
-            logger.exception("Error starting module")
+            logger.exception("Error during module lifecyle")
 
     async def stop(self) -> None:
         """Stop the module."""
@@ -263,3 +365,18 @@ class BaseModule(ABC, Generic[InputModelT, OutputModelT, SetupModelT, SecretMode
         except Exception:
             self._status = ModuleStatus.FAILED
             logger.exception("Error stopping module")
+
+    async def start_config_setup(
+        self,
+        config_setup_data: ConfigSetupModelT,
+        setup_data: SetupModelT,
+        callback: Callable[[OutputModelT | ModuleErrorModel], Coroutine[Any, Any, None]],
+    ) -> None:
+        """Start the module."""
+        try:
+            logger.info("Run Config Setup lifecycle")
+            self._status = ModuleStatus.RUNNING
+            await self.run_config_setup(config_setup_data, setup_data, callback)
+        except Exception:
+            self._status = ModuleStatus.FAILED
+            logger.exception("Error during module lifecyle")

digitalkin/modules/archetype_module.py

@@ -3,8 +3,18 @@
 from abc import ABC
 
 from digitalkin.models.module import InputModelT, OutputModelT, SecretModelT, SetupModelT
+from digitalkin.models.module.module_types import ConfigSetupModelT
 from digitalkin.modules._base_module import BaseModule
 
 
-class ArchetypeModule(BaseModule[InputModelT, OutputModelT, SetupModelT, SecretModelT], ABC):
+class ArchetypeModule(
+    BaseModule[
+        InputModelT,
+        OutputModelT,
+        SetupModelT,
+        SecretModelT,
+        ConfigSetupModelT,
+    ],
+    ABC,
+):
     """ArchetypeModule extends BaseModule to implement specific module types."""