digitalkin 0.2.26__py3-none-any.whl → 0.3.0.dev1__py3-none-any.whl

digitalkin/__version__.py

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

digitalkin/grpc_servers/module_server.py

@@ -1,20 +1,9 @@
 """Module gRPC server implementation for DigitalKin."""
 
-from pathlib import Path
 import uuid
+from pathlib import Path
 
 import grpc
-
-from digitalkin.grpc_servers._base_server import BaseServer
-from digitalkin.grpc_servers.module_servicer import ModuleServicer
-from digitalkin.grpc_servers.utils.exceptions import ServerError
-from digitalkin.grpc_servers.utils.models import (
-    ClientConfig,
-    ModuleServerConfig,
-    SecurityMode,
-)
-from digitalkin.modules._base_module import BaseModule
-
 from digitalkin_proto.digitalkin.module.v2 import (
     module_service_pb2,
     module_service_pb2_grpc,
@@ -24,7 +13,17 @@ from digitalkin_proto.digitalkin.module_registry.v2 import (
     module_registry_service_pb2_grpc,
     registration_pb2,
 )
+
+from digitalkin.grpc_servers._base_server import BaseServer
+from digitalkin.grpc_servers.module_servicer import ModuleServicer
+from digitalkin.grpc_servers.utils.exceptions import ServerError
+from digitalkin.grpc_servers.utils.models import (
+    ClientConfig,
+    ModuleServerConfig,
+    SecurityMode,
+)
 from digitalkin.logger import logger
+from digitalkin.modules._base_module import BaseModule
 
 
 class ModuleServer(BaseServer):
@@ -50,7 +49,8 @@ class ModuleServer(BaseServer):
 
         Args:
             module_class: The module instance to be served.
-            config: Server configuration including registry address if auto-registration is desired.
+            server_config: Server configuration including registry address if auto-registration is desired.
+            client_config: Client configuration used by services.
         """
         super().__init__(server_config)
         self.module_class = module_class
@@ -79,10 +79,9 @@ class ModuleServer(BaseServer):
 
     def start(self) -> None:
         """Start the module server and register with the registry if configured."""
-        logger.info("Starting module server",extra={"server_config": self.server_config})
+        logger.info("Starting module server", extra={"server_config": self.server_config})
         super().start()
 
-        logger.debug("Starting module server",extra={"server_config": self.server_config})
         # If a registry address is provided, register the module
         if self.server_config.registry_address:
             try:
@@ -91,14 +90,12 @@ class ModuleServer(BaseServer):
                 logger.exception("Failed to register with registry")
 
         if self.module_servicer is not None:
-            logger.debug(
-                "Setup post init started",extra={"client_config": self.client_config}
-            )
+            logger.debug("Setup post init started", extra={"client_config": 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.info("Starting module server",extra={"server_config": self.server_config})
+        logger.info("Starting module server", extra={"server_config": self.server_config})
         await super().start_async()
         # If a registry address is provided, register the module
         if self.server_config.registry_address:
@@ -108,10 +105,8 @@ class ModuleServer(BaseServer):
                 logger.exception("Failed to register with registry")
 
         if self.module_servicer is not None:
-            logger.info(
-                "Setup post init started",extra={"client_config": self.client_config}
-            )
-            await self.module_servicer.job_manager._start()
+            logger.info("Setup post init started", extra={"client_config": self.client_config})
+            await self.module_servicer.job_manager.start()
             self.module_servicer.setup.__post_init__(self.client_config)
 
     def stop(self, grace: float | None = None) -> None:
@@ -133,7 +128,8 @@ class ModuleServer(BaseServer):
         """
         logger.debug(
             "Registering module with registry at %s",
-            self.server_config.registry_address, extra={"server_config": self.server_config}
+            self.server_config.registry_address,
+            extra={"server_config": self.server_config},
         )
 
         # Create appropriate channel based on security mode
@@ -148,16 +144,11 @@ class ModuleServer(BaseServer):
 
             metadata = metadata_pb2.Metadata(
                 name=self.module_class.metadata["name"],
-                tags=[
-                    metadata_pb2.Tag(tag=tag)
-                    for tag in self.module_class.metadata["tags"]
-                ],
+                tags=[metadata_pb2.Tag(tag=tag) for tag in self.module_class.metadata["tags"]],
                 description=self.module_class.metadata["description"],
             )
 
-            self.module_class.metadata["module_id"] = (
-                f"{self.module_class.metadata['name']}:{uuid.uuid4()}"
-            )
+            self.module_class.metadata["module_id"] = f"{self.module_class.metadata['name']}:{uuid.uuid4()}"
             # Create registration request
             request = registration_pb2.RegisterRequest(
                 module_id=self.module_class.metadata["module_id"],
@@ -173,7 +164,7 @@ class ModuleServer(BaseServer):
                     "Request sent to registry for module: %s:%s",
                     self.module_class.metadata["name"],
                     self.module_class.metadata["module_id"],
-                    extra={"module_info": self.module_class.metadata}
+                    extra={"module_info": self.module_class.metadata},
                 )
                 response = stub.RegisterModule(request)
 
@@ -235,9 +226,7 @@ class ModuleServer(BaseServer):
         ):
             # Secure channel
             # Secure channel
-            root_certificates = Path(
-                self.client_config.credentials.root_cert_path
-            ).read_bytes()
+            root_certificates = Path(self.client_config.credentials.root_cert_path).read_bytes()
 
             # mTLS channel
             private_key = None
@@ -246,12 +235,8 @@ class ModuleServer(BaseServer):
                 self.client_config.credentials.client_cert_path is not None
                 and self.client_config.credentials.client_key_path is not None
             ):
-                private_key = Path(
-                    self.client_config.credentials.client_key_path
-                ).read_bytes()
-                certificate_chain = Path(
-                    self.client_config.credentials.client_cert_path
-                ).read_bytes()
+                private_key = Path(self.client_config.credentials.client_key_path).read_bytes()
+                certificate_chain = Path(self.client_config.credentials.client_cert_path).read_bytes()
 
             # Create channel credentials
             channel_credentials = grpc.ssl_channel_credentials(
@@ -259,9 +244,7 @@ class ModuleServer(BaseServer):
                 certificate_chain=certificate_chain,
                 private_key=private_key,
             )
-            return grpc.secure_channel(
-                self.server_config.registry_address, channel_credentials
-            )
+            return grpc.secure_channel(self.server_config.registry_address, channel_credentials)
         # Insecure channel
         return grpc.insecure_channel(self.server_config.registry_address)
 

digitalkin/grpc_servers/module_servicer.py

@@ -201,28 +201,33 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
             yield lifecycle_pb2.StartModuleResponse(success=False)
             return
 
-        async with self.job_manager.generate_stream_consumer(job_id) as stream:  # type: ignore
-            async for message in stream:
-                if message.get("error", None) is not None:
-                    logger.error("Error in output_data", extra={"message": message})
-                    context.set_code(message["error"]["code"])
-                    context.set_details(message["error"]["error_message"])
-                    yield lifecycle_pb2.StartModuleResponse(success=False, job_id=job_id)
-                    break
-
-                if message.get("exception", None) is not None:
-                    logger.error("Exception in output_data", extra={"message": message})
-                    context.set_code(message["short_description"])
-                    context.set_details(message["exception"])
-                    yield lifecycle_pb2.StartModuleResponse(success=False, job_id=job_id)
-                    break
-
-                if message.get("code", None) is not None and message.get("code") == "__END_OF_STREAM__":
-                    yield lifecycle_pb2.StartModuleResponse(success=True, job_id=job_id)
-                    break
-
-                proto = json_format.ParseDict(message, struct_pb2.Struct(), ignore_unknown_fields=True)
-                yield lifecycle_pb2.StartModuleResponse(success=True, output=proto, job_id=job_id)
+        try:
+            async with self.job_manager.generate_stream_consumer(job_id) as stream:  # type: ignore
+                async for message in stream:
+                    if message.get("error", None) is not None:
+                        logger.error("Error in output_data", extra={"message": message})
+                        context.set_code(message["error"]["code"])
+                        context.set_details(message["error"]["error_message"])
+                        yield lifecycle_pb2.StartModuleResponse(success=False, job_id=job_id)
+                        break
+
+                    if message.get("exception", None) is not None:
+                        logger.error("Exception in output_data", extra={"message": message})
+                        context.set_code(message["short_description"])
+                        context.set_details(message["exception"])
+                        yield lifecycle_pb2.StartModuleResponse(success=False, job_id=job_id)
+                        break
+
+                    if message.get("code", None) is not None and message.get("code") == "__END_OF_STREAM__":
+                        yield lifecycle_pb2.StartModuleResponse(success=True, job_id=job_id)
+                        break
+
+                    proto = json_format.ParseDict(message, struct_pb2.Struct(), ignore_unknown_fields=True)
+                    yield lifecycle_pb2.StartModuleResponse(success=True, output=proto, job_id=job_id)
+        finally:
+            await self.job_manager.tasks[job_id]
+            await self.job_manager.clean_session(job_id)
+
         logger.info("Job %s finished", job_id)
 
     async def StopModule(  # noqa: N802

digitalkin/grpc_servers/utils/models.py

@@ -262,7 +262,7 @@ class ModuleServerConfig(ServerConfig):
         registry_address: Address of the registry server
     """
 
-    registry_address: str | None = Field(None, description="Address of the registry server")
+    registry_address: str = Field(..., description="Address of the registry server")
 
 
 class RegistryServerConfig(ServerConfig):

digitalkin/logger.py

@@ -48,17 +48,9 @@ class ColorJSONFormatter(logging.Formatter):
         log_obj: dict[str, Any] = {
             "timestamp": datetime.fromtimestamp(record.created, tz=timezone.utc).isoformat(),
             "level": record.levelname.lower(),
-            "logger": record.name,
             "message": record.getMessage(),
             "module": record.module,
-            "location": f"{record.filename}:{record.lineno}",
-            "function": record.funcName,
-            "pathname": record.pathname,
-            "process": record.process,
-            "processName": record.processName,
-            "relativeCreated": record.relativeCreated,
-            "thread": record.thread,
-            "threadName": record.threadName,
+            "location": f"{record.pathname}:{record.lineno}:{record.funcName}",
         }
         # Add exception info if present
         if record.exc_info:

digitalkin/mixins/__init__.py

@@ -0,0 +1,19 @@
+"""Mixin definitions."""
+
+from digitalkin.mixins.base_mixin import BaseMixin
+from digitalkin.mixins.callback_mixin import UserMessageMixin
+from digitalkin.mixins.chat_history_mixin import ChatHistoryMixin
+from digitalkin.mixins.cost_mixin import CostMixin
+from digitalkin.mixins.filesystem_mixin import FilesystemMixin
+from digitalkin.mixins.logger_mixin import LoggerMixin
+from digitalkin.mixins.storage_mixin import StorageMixin
+
+__all__ = [
+    "BaseMixin",
+    "ChatHistoryMixin",
+    "CostMixin",
+    "FilesystemMixin",
+    "LoggerMixin",
+    "StorageMixin",
+    "UserMessageMixin",
+]

digitalkin/mixins/base_mixin.py

@@ -0,0 +1,10 @@
+"""Simple toolkit class with basic and simple API access in the Triggers."""
+
+from digitalkin.mixins.chat_history_mixin import ChatHistoryMixin
+from digitalkin.mixins.cost_mixin import CostMixin
+from digitalkin.mixins.file_history_mixin import FileHistoryMixin
+from digitalkin.mixins.logger_mixin import LoggerMixin
+
+
+class BaseMixin(CostMixin, ChatHistoryMixin, FileHistoryMixin, LoggerMixin):
+    """Base Mixin to access to minimum Module Context functionnalities in the Triggers."""

digitalkin/mixins/callback_mixin.py

@@ -0,0 +1,24 @@
+"""User callback to send a message from the Trigger."""
+
+from typing import Generic
+
+from digitalkin.models.module.module_context import ModuleContext
+from digitalkin.models.module.module_types import OutputModelT
+
+
+class UserMessageMixin(Generic[OutputModelT]):
+    """Mixin providing callback operations through the callbacks .
+
+    This mixin wraps callback strategy calls to provide a cleaner API
+    for direct messaging in trigger handlers.
+    """
+
+    @staticmethod
+    async def send_message(context: ModuleContext, output: OutputModelT) -> None:
+        """Send a message using the callbacks strategy.
+
+        Args:
+            context: Module context containing the callbacks strategy.
+            output: Message to send with the Module defined output Type.
+        """
+        await context.callbacks.send_message(output)

digitalkin/mixins/chat_history_mixin.py

@@ -0,0 +1,108 @@
+"""Context mixins providing ergonomic access to service strategies.
+
+This module provides mixins that wrap service strategy calls with cleaner APIs,
+following Django/FastAPI patterns where context is passed explicitly to each method.
+"""
+
+from typing import Any, Generic
+
+from digitalkin.mixins.callback_mixin import UserMessageMixin
+from digitalkin.mixins.logger_mixin import LoggerMixin
+from digitalkin.mixins.storage_mixin import StorageMixin
+from digitalkin.models.module.module_context import ModuleContext
+from digitalkin.models.module.module_types import InputModelT, OutputModelT
+from digitalkin.models.services.storage import BaseMessage, ChatHistory, Role
+
+
+class ChatHistoryMixin(UserMessageMixin, StorageMixin, LoggerMixin, Generic[InputModelT, OutputModelT]):
+    """Mixin providing chat history operations through storage strategy.
+
+    This mixin provides a higher-level API for managing chat history,
+    using the storage strategy as the underlying persistence mechanism.
+    """
+
+    CHAT_HISTORY_COLLECTION = "chat_history"
+    CHAT_HISTORY_RECORD_ID = "full_chat_history"
+
+    def _get_history_key(self, context: ModuleContext) -> str:
+        """Get session-specific history key.
+
+        Args:
+            context: Module context containing session information
+
+        Returns:
+            Unique history key for the current session
+        """
+        mission_id = getattr(context.session, "mission_id", None) or "default"
+        return f"{self.CHAT_HISTORY_RECORD_ID}_{mission_id}"
+
+    def load_chat_history(self, context: ModuleContext) -> ChatHistory:
+        """Load chat history for the current session.
+
+        Args:
+            context: Module context containing storage strategy
+
+        Returns:
+            Chat history object, empty if none exists or loading fails
+        """
+        history_key = self._get_history_key(context)
+
+        if (raw_history := self.read_storage(context, self.CHAT_HISTORY_COLLECTION, history_key)) is not None:
+            return ChatHistory.model_validate(raw_history.data)
+        return ChatHistory(messages=[])
+
+    def append_chat_history_message(
+        self,
+        context: ModuleContext,
+        role: Role,
+        content: Any,  # noqa: ANN401
+    ) -> None:
+        """Append a message to chat history.
+
+        Args:
+            context: Module context containing storage strategy
+            role: Message role (user, assistant, system)
+            content: Message content
+
+        Raises:
+            StorageServiceError: If history update fails
+        """
+        history_key = self._get_history_key(context)
+        chat_history = self.load_chat_history(context)
+
+        chat_history.messages.append(BaseMessage(role=role, content=content))
+        if len(chat_history.messages) == 1:
+            # Create new record
+            self.log_debug(context, f"Creating new chat history for session: {history_key}")
+            self.store_storage(
+                context,
+                self.CHAT_HISTORY_COLLECTION,
+                history_key,
+                chat_history.model_dump(),
+                data_type="OUTPUT",
+            )
+        else:
+            self.log_debug(context, f"Updating chat history for session: {history_key}")
+            self.update_storage(
+                context,
+                self.CHAT_HISTORY_COLLECTION,
+                history_key,
+                chat_history.model_dump(),
+            )
+
+    async def save_send_message(
+        self,
+        context: ModuleContext,
+        output: OutputModelT,
+        role: Role,
+    ) -> None:
+        """Save the output message to the chat history and send a response to the Module request.
+
+        Args:
+            context: Module context containing storage strategy
+            role: Message role (user, assistant, system)
+            output: Message content as Pydantic Class
+        """
+        # TO-DO: we should define a default output message type to ease user experience
+        self.append_chat_history_message(context=context, role=role, content=str(output.root))
+        await self.send_message(context=context, output=output)

digitalkin/mixins/cost_mixin.py

@@ -0,0 +1,76 @@
+"""Cost Mixin to ease trigger deveolpment."""
+
+from typing import Literal
+
+from digitalkin.models.module.module_context import ModuleContext
+from digitalkin.services.cost.cost_strategy import CostData
+
+
+class CostMixin:
+    """Mixin providing cost tracking operations through the cost strategy.
+
+    This mixin wraps cost strategy calls to provide a cleaner API
+    for cost tracking in trigger handlers.
+    """
+
+    @staticmethod
+    def add_cost(context: ModuleContext, name: str, cost_config_name: str, quantity: float) -> None:
+        """Add a cost entry using the cost strategy.
+
+        Args:
+            context: Module context containing the cost strategy
+            name: Name/identifier for this cost entry
+            cost_config_name: Name of the cost configuration to use
+            quantity: Quantity of units consumed
+
+        Raises:
+            CostServiceError: If cost addition fails
+        """
+        return context.cost.add(name, cost_config_name, quantity)
+
+    @staticmethod
+    def get_cost(context: ModuleContext, name: str) -> list[CostData]:
+        """Get cost entries for a specific name.
+
+        Args:
+            context: Module context containing the cost strategy
+            name: Name/identifier to get costs for
+
+        Returns:
+            List of cost data entries
+
+        Raises:
+            CostServiceError: If cost retrieval fails
+        """
+        return context.cost.get(name)
+
+    @staticmethod
+    def get_costs(
+        context: ModuleContext,
+        names: list[str] | None = None,
+        cost_types: list[
+            Literal[
+                "TOKEN_INPUT",
+                "TOKEN_OUTPUT",
+                "API_CALL",
+                "STORAGE",
+                "TIME",
+                "OTHER",
+            ]
+        ]
+        | None = None,
+    ) -> list[CostData]:
+        """Get filtered cost entries.
+
+        Args:
+            context: Module context containing the cost strategy
+            names: Optional list of names to filter by
+            cost_types: Optional list of cost types to filter by
+
+        Returns:
+            List of filtered cost data entries
+
+        Raises:
+            CostServiceError: If cost retrieval fails
+        """
+        return context.cost.get_filtered(names, cost_types)

digitalkin/mixins/file_history_mixin.py

@@ -0,0 +1,99 @@
+"""Context mixins providing ergonomic access to service strategies.
+
+This module provides mixins that wrap service strategy calls with cleaner APIs,
+following Django/FastAPI patterns where context is passed explicitly to each method.
+"""
+
+from digitalkin.mixins.logger_mixin import LoggerMixin
+from digitalkin.mixins.storage_mixin import StorageMixin
+from digitalkin.models.module.module_context import ModuleContext
+from digitalkin.models.services.storage import FileHistory, FileModel
+
+
+class FileHistoryMixin(StorageMixin, LoggerMixin):
+    """Mixin providing File history operations through storage strategy.
+
+    This mixin provides a higher-level API for managing File history,
+    using the storage strategy as the underlying persistence mechanism.
+    """
+
+    file_history_front: FileHistory = FileHistory(files=[])
+    FILE_HISTORY_COLLECTION = "file_history"
+    FILE_HISTORY_RECORD_ID = "full_file_history"
+
+    def _get_history_key(self, context: ModuleContext) -> str:
+        """Get session-specific history key.
+
+        Args:
+            context: Module context containing session information
+
+        Returns:
+            Unique history key for the current session
+        """
+        mission_id = getattr(context.session, "mission_id", None) or "default"
+        return f"{self.FILE_HISTORY_RECORD_ID}_{mission_id}"
+
+    def load_file_history(self, context: ModuleContext) -> FileHistory:
+        """Load File history for the current session.
+
+        Args:
+            context: Module context containing storage strategy
+
+        Returns:
+            File history object, empty if none exists or loading fails
+        """
+        history_key = self._get_history_key(context)
+
+        if self.file_history_front is None:
+            try:
+                record = self.read_storage(
+                    context,
+                    self.FILE_HISTORY_COLLECTION,
+                    history_key,
+                )
+                if record and record.data:
+                    return FileHistory.model_validate(record.data)
+            except Exception as e:
+                self.log_warning(context, f"Failed to load File history: {e}")
+        return self.file_history_front
+
+    def append_files_history(self, context: ModuleContext, files: list[FileModel]) -> None:
+        """Append a message to File history.
+
+        Args:
+            context: Module context containing storage strategy
+            role: Message role (user, assistant, system)
+            files: list of files model
+
+        Raises:
+            StorageServiceError: If history update fails
+        """
+        history_key = self._get_history_key(context)
+
+        try:
+            if not self.file_history_front.files:
+                self.file_history_front = self.load_file_history(context)
+            self.file_history_front.files.extend(files)
+        except Exception as e:
+            self.log_error(context, f"Failed to append message to File history: {e}")
+
+        try:
+            self.log_debug(context, f"Updating File history for session: {history_key}")
+            self.update_storage(
+                context,
+                self.FILE_HISTORY_COLLECTION,
+                history_key,
+                self.file_history_front.model_dump(),
+            )
+        except Exception as e:
+            self.log_error(context, f"Updating File history for session: {history_key} with error: {e}")
+
+            # Create new record
+            self.log_debug(context, f"Creating new File history for session: {history_key}")
+            self.store_storage(
+                context,
+                self.FILE_HISTORY_COLLECTION,
+                history_key,
+                self.file_history_front.model_dump(),
+                data_type="OUTPUT",
+            )

digitalkin/mixins/filesystem_mixin.py

@@ -0,0 +1,47 @@
+"""Filesystem Mixin to ease filesystem use."""
+
+from typing import Any
+
+from digitalkin.models.module.module_context import ModuleContext
+from digitalkin.services.filesystem.filesystem_strategy import FilesystemRecord
+
+
+class FilesystemMixin:
+    """Mixin providing filesystem operations through the filesystem strategy.
+
+    This mixin wraps filesystem strategy calls to provide a cleaner API
+    for file operations in trigger handlers.
+    """
+
+    @staticmethod
+    def upload_files(context: ModuleContext, files: list[Any]) -> tuple[list[FilesystemRecord], int, int]:
+        """Upload files using the filesystem strategy.
+
+        Args:
+            context: Module context containing the filesystem strategy
+            files: List of files to upload
+
+        Returns:
+            Tuple of (all_files, succeeded_files, failed_files)
+
+        Raises:
+            FilesystemServiceError: If upload operation fails
+        """
+        return context.filesystem.upload_files(files)
+
+    @staticmethod
+    def get_file(context: ModuleContext, file_id: str) -> tuple[FilesystemRecord, bytes | None]:
+        """Retrieve a file by ID with the content.
+
+        Args:
+            context: Module context containing the filesystem strategy
+            file_id: Unique identifier for the file
+            include_content: Whether to include file content in response
+
+        Returns:
+            File object with metadata and optionally content
+
+        Raises:
+            FilesystemServiceError: If file retrieval fails
+        """
+        return context.filesystem.get_file(file_id, include_content=True)