digitalkin 0.3.2.dev2__py3-none-any.whl

digitalkin/grpc_servers/utils/utility_schema_extender.py

@@ -0,0 +1,97 @@
+"""Utility schema extender for gRPC API responses.
+
+This module extends module schemas with SDK utility protocols for API responses.
+"""
+
+from typing import Annotated, Union, get_args, get_origin
+
+from pydantic import Field, create_model
+
+from digitalkin.models.module.module_types import DataModel
+from digitalkin.models.module.utility import (
+    EndOfStreamOutput,
+    HealthcheckPingInput,
+    HealthcheckPingOutput,
+    HealthcheckServicesInput,
+    HealthcheckServicesOutput,
+    HealthcheckStatusInput,
+    HealthcheckStatusOutput,
+)
+
+
+class UtilitySchemaExtender:
+    """Extends module schemas with SDK utility protocols for API responses.
+
+    This class provides methods to create extended Pydantic models that include
+    both user-defined protocols and SDK utility protocols in their schemas.
+    """
+
+    _output_protocols = (
+        EndOfStreamOutput,
+        HealthcheckPingOutput,
+        HealthcheckServicesOutput,
+        HealthcheckStatusOutput,
+    )
+
+    _input_protocols = (
+        HealthcheckPingInput,
+        HealthcheckServicesInput,
+        HealthcheckStatusInput,
+    )
+
+    @classmethod
+    def _extract_union_types(cls, annotation: type) -> tuple:
+        """Extract individual types from a Union or Annotated[Union, ...] annotation.
+
+        Returns:
+            A tuple of individual types contained in the Union.
+        """
+        if get_origin(annotation) is Annotated:
+            inner_args = get_args(annotation)
+            if inner_args:
+                return cls._extract_union_types(inner_args[0])
+        if get_origin(annotation) is Union:
+            return get_args(annotation)
+        return (annotation,)
+
+    @classmethod
+    def create_extended_output_model(cls, base_model: type[DataModel]) -> type[DataModel]:
+        """Create an extended output model that includes utility output protocols.
+
+        Args:
+            base_model: The module's output_format class (a DataModel subclass).
+
+        Returns:
+            A new DataModel subclass with root typed as Union[original_types, utility_types].
+        """
+        original_annotation = base_model.model_fields["root"].annotation
+        original_types = cls._extract_union_types(original_annotation)
+        extended_types = (*original_types, *cls._output_protocols)
+        extended_root = Annotated[extended_types, Field(discriminator="protocol")]  # type: ignore[valid-type]
+        return create_model(
+            f"{base_model.__name__}Utilities",
+            __base__=DataModel,
+            root=(extended_root, ...),
+            annotations=(dict[str, str], Field(default={})),
+        )
+
+    @classmethod
+    def create_extended_input_model(cls, base_model: type[DataModel]) -> type[DataModel]:
+        """Create an extended input model that includes utility input protocols.
+
+        Args:
+            base_model: The module's input_format class (a DataModel subclass).
+
+        Returns:
+            A new DataModel subclass with root typed as Union[original_types, utility_types].
+        """
+        original_annotation = base_model.model_fields["root"].annotation
+        original_types = cls._extract_union_types(original_annotation)
+        extended_types = (*original_types, *cls._input_protocols)
+        extended_root = Annotated[extended_types, Field(discriminator="protocol")]  # type: ignore[valid-type]
+        return create_model(
+            f"{base_model.__name__}Utilities",
+            __base__=DataModel,
+            root=(extended_root, ...),
+            annotations=(dict[str, str], Field(default={})),
+        )

digitalkin/logger.py

@@ -0,0 +1,157 @@
+"""This module sets up a logger."""
+
+import json
+import logging
+import os
+import sys
+from datetime import datetime, timezone
+from typing import Any, ClassVar
+
+
+class ColorJSONFormatter(logging.Formatter):
+    """Color JSON formatter for development (pretty-printed with colors)."""
+
+    def __init__(self, *, is_production: bool = False) -> None:
+        """Initialize the formatter.
+
+        Args:
+            is_production: Whether the application is running in production.
+        """
+        self.is_production = is_production
+        super().__init__()
+
+    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"
+
+    COLORS: ClassVar[dict[int, str]] = {
+        logging.DEBUG: grey,
+        logging.INFO: green,
+        logging.WARNING: yellow,
+        logging.ERROR: red,
+        logging.CRITICAL: bold_red,
+    }
+
+    def format(self, record: logging.LogRecord) -> str:
+        """Format the log record as colored JSON for development.
+
+        Args:
+            record: The log record to format.
+
+        Returns:
+            str: The colored JSON formatted log record.
+        """
+        log_obj: dict[str, Any] = {
+            "timestamp": datetime.fromtimestamp(record.created, tz=timezone.utc).isoformat(),
+            "level": record.levelname.lower(),
+            "message": record.getMessage(),
+            "module": record.module,
+            "location": f"{record.pathname}:{record.lineno}:{record.funcName}",
+        }
+        # Add exception info if present
+        if record.exc_info:
+            log_obj["exception"] = self.formatException(record.exc_info)
+
+        # Add any extra fields
+        skip_attrs = {
+            "name",
+            "msg",
+            "args",
+            "created",
+            "filename",
+            "funcName",
+            "levelname",
+            "levelno",
+            "lineno",
+            "module",
+            "msecs",
+            "message",
+            "pathname",
+            "process",
+            "processName",
+            "relativeCreated",
+            "thread",
+            "threadName",
+            "exc_info",
+            "exc_text",
+            "stack_info",
+        }
+
+        extras = {key: value for key, value in record.__dict__.items() if key not in skip_attrs}
+
+        if extras:
+            log_obj["extra"] = extras
+
+        # Pretty print with color
+        color = self.COLORS.get(record.levelno, self.grey)
+        if self.is_production:
+            log_obj["message"] = f"{color}{log_obj.get('message', '')}{self.reset}"
+            return json.dumps(log_obj, default=str, separators=(",", ":"))
+        json_str = json.dumps(log_obj, indent=2, default=str)
+        json_str = json_str.replace("\\n", "\n")
+        return f"{color}{json_str}{self.reset}"
+
+
+def setup_logger(
+    name: str,
+    level: int = logging.INFO,
+    additional_loggers: dict[str, int] | None = None,
+    *,
+    is_production: bool | None = None,
+    configure_root: bool = True,
+) -> logging.Logger:
+    """Set up a logger with the ColorJSONFormatter.
+
+    Args:
+        name: Name of the logger to create
+        level: Logging level (default: logging.INFO)
+        is_production: Whether running in production. If None, checks RAILWAY_SERVICE_NAME env var
+        configure_root: Whether to configure root logger (default: True)
+        additional_loggers: Dict of additional logger names and their levels to configure
+
+    Returns:
+        logging.Logger: Configured logger instance
+    """
+    # Determine if we're in production
+    if is_production is None:
+        is_production = os.getenv("RAILWAY_SERVICE_NAME") is not None
+
+    # Configure root logger if requested
+    if configure_root:
+        logging.basicConfig(
+            level=logging.DEBUG,
+            stream=sys.stdout,
+            datefmt="%Y-%m-%d %H:%M:%S",
+        )
+
+    # Configure additional loggers
+    if additional_loggers:
+        for logger_name, logger_level in additional_loggers.items():
+            logging.getLogger(logger_name).setLevel(logger_level)
+
+    # Create and configure the main logger
+    logger = logging.getLogger(name)
+    logger.setLevel(level)
+    # Only add handler if not already configured
+    if not logger.handlers:
+        ch = logging.StreamHandler()
+        ch.setLevel(level)
+        ch.setFormatter(ColorJSONFormatter(is_production=is_production))
+        logger.addHandler(ch)
+        logger.propagate = False
+
+    return logger
+
+
+logger = setup_logger(
+    "digitalkin",
+    level=logging.INFO,
+    additional_loggers={
+        "grpc": logging.DEBUG,
+        "asyncio": logging.DEBUG,
+    },
+)

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,110 @@
+"""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
+        """
+        # TODO: define mission-specific chat history key not dependant on mission_id
+        # or need customization by user
+        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=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,93 @@
+"""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
+        """
+        # TODO: define mission-specific chat history key not dependant on mission_id
+        # or need customization by user
+        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
+            files: list of files model
+
+        Raises:
+            StorageServiceError: If history update fails
+        """
+        history_key = self._get_history_key(context)
+        file_history = self.load_file_history(context)
+
+        file_history.files.extend(files)
+        if len(file_history.files) == len(files):
+            # 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,
+                file_history.model_dump(),
+                data_type="OUTPUT",
+            )
+        else:
+            self.log_debug(context, f"Updating file history for session: {history_key}")
+            self.update_storage(
+                context,
+                self.FILE_HISTORY_COLLECTION,
+                history_key,
+                file_history.model_dump(),
+            )

digitalkin/mixins/filesystem_mixin.py

@@ -0,0 +1,46 @@
+"""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) -> FilesystemRecord:
+        """Retrieve a file by ID with the content.
+
+        Args:
+            context: Module context containing the filesystem strategy
+            file_id: Unique identifier for the file
+
+        Returns:
+            File object with metadata and optionally content
+
+        Raises:
+            FilesystemServiceError: If file retrieval fails
+        """
+        return context.filesystem.get_file(file_id, include_content=True)

digitalkin/mixins/logger_mixin.py

@@ -0,0 +1,51 @@
+"""Logger Mixin to ease and merge every logs."""
+
+from digitalkin.models.module.module_context import ModuleContext
+
+
+class LoggerMixin:
+    """Mixin providing callback operations through the callbacks strategy.
+
+    This mixin wraps callback strategy calls to provide a cleaner API
+    for logging and messaging in trigger handlers.
+    """
+
+    @staticmethod
+    def log_debug(context: ModuleContext, message: str) -> None:
+        """Log debug message using the callbacks strategy.
+
+        Args:
+            context: Module context containing the callbacks strategy
+            message: Debug message to log
+        """
+        return context.callbacks.logger.debug(message, extra=context.session.current_ids())
+
+    @staticmethod
+    def log_info(context: ModuleContext, message: str) -> None:
+        """Log info message using the callbacks strategy.
+
+        Args:
+            context: Module context containing the callbacks strategy
+            message: Info message to log
+        """
+        return context.callbacks.logger.info(message, extra=context.session.current_ids())
+
+    @staticmethod
+    def log_warning(context: ModuleContext, message: str) -> None:
+        """Log warning message using the callbacks strategy.
+
+        Args:
+            context: Module context containing the callbacks strategy
+            message: Warning message to log
+        """
+        return context.callbacks.logger.warning(message, extra=context.session.current_ids())
+
+    @staticmethod
+    def log_error(context: ModuleContext, message: str) -> None:
+        """Log error message using the callbacks strategy.
+
+        Args:
+            context: Module context containing the callbacks strategy
+            message: Error message to log
+        """
+        return context.callbacks.logger.error(message, extra=context.session.current_ids())