digitalkin 0.2.23__py3-none-any.whl → 0.3.1.dev2__py3-none-any.whl

digitalkin/grpc_servers/module_servicer.py

@@ -5,7 +5,7 @@ from collections.abc import AsyncGenerator
 from typing import Any
 
 import grpc
-from digitalkin_proto.digitalkin.module.v2 import (
+from digitalkin_proto.agentic_mesh_protocol.module.v1 import (
     information_pb2,
     lifecycle_pb2,
     module_service_pb2_grpc,
@@ -13,12 +13,12 @@ from digitalkin_proto.digitalkin.module.v2 import (
 )
 from google.protobuf import json_format, struct_pb2
 
+from digitalkin.core.job_manager.base_job_manager import BaseJobManager
 from digitalkin.grpc_servers.utils.exceptions import ServicerError
 from digitalkin.logger import logger
+from digitalkin.models.core.job_manager_models import JobManagerMode
 from digitalkin.models.module.module import ModuleStatus
 from digitalkin.modules._base_module import BaseModule
-from digitalkin.modules.job_manager.base_job_manager import BaseJobManager
-from digitalkin.modules.job_manager.job_manager_models import JobManagerMode
 from digitalkin.services.services_models import ServicesMode
 from digitalkin.services.setup.default_setup import DefaultSetup
 from digitalkin.services.setup.grpc_setup import GrpcSetup
@@ -76,9 +76,9 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         self.job_manager = job_manager_class(module_class, self.args.services_mode)
 
         logger.debug(
-            "ModuleServicer initialized with job manager: %s | %s",
+            "ModuleServicer initialized with job manager: %s",
             self.args.job_manager_mode,
-            self.job_manager,
+            extra={"job_manager": self.job_manager},
         )
         self.setup = GrpcSetup() if self.args.services_mode == ServicesMode.REMOTE else DefaultSetup()
 
@@ -112,7 +112,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         # TODO: Secret should be used here as well
         setup_version = request.setup_version
         config_setup_data = self.module_class.create_config_setup_model(json_format.MessageToDict(request.content))
-        setup_version_data = self.module_class.create_setup_model(
+        setup_version_data = await self.module_class.create_setup_model(
             json_format.MessageToDict(request.setup_version.content),
             config_fields=True,
         )
@@ -172,7 +172,8 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         )
         # Process the module input
         # TODO: Check failure of input data format
-        input_data = self.module_class.create_input_model(dict(request.input.items()))
+        input_data = self.module_class.create_input_model(json_format.MessageToDict(request.input))
+
         setup_data_class = self.setup.get_setup(
             setup_dict={
                 "setup_id": request.setup_id,
@@ -184,7 +185,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
             msg = "No setup data returned."
             raise ServicerError(msg)
 
-        setup_data = self.module_class.create_setup_model(setup_data_class.current_setup_version.content)
+        setup_data = await self.module_class.create_setup_model(setup_data_class.current_setup_version.content)
 
         # create a task to run the module in background
         job_id = await self.job_manager.create_module_instance_job(
@@ -201,27 +202,37 @@ 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:
-                    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("Error in output_data")
-                    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__":
+                        logger.info(
+                            "End of stream via __END_OF_STREAM__",
+                            extra={"job_id": job_id, "mission_id": request.mission_id},
+                        )
+                        break
+
+                    logger.info("Yielding message from job %s: %s", job_id, message)
+                    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.wait_for_completion(job_id)
+            await self.job_manager.clean_session(job_id, mission_id=request.mission_id)
+
         logger.info("Job %s finished", job_id)
 
     async def StopModule(  # noqa: N802
@@ -248,7 +259,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
             context.set_details(message)
             return lifecycle_pb2.StopModuleResponse(success=False)
 
-        logger.debug("Job %s stopped successfully", request.job_id)
+        logger.debug("Job %s stopped successfully", request.job_id, extra={"job_id": request.job_id})
         return lifecycle_pb2.StopModuleResponse(success=True)
 
     async def GetModuleStatus(  # noqa: N802
@@ -339,7 +350,9 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         # Get input schema if available
         try:
             # Convert schema to proto format
-            input_schema_proto = self.module_class.get_input_format(llm_format=request.llm_format)
+            input_schema_proto = await self.module_class.get_input_format(
+                llm_format=request.llm_format,
+            )
             input_format_struct = json_format.Parse(
                 text=input_schema_proto,
                 message=struct_pb2.Struct(),  # pylint: disable=no-member
@@ -375,7 +388,9 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         # Get output schema if available
         try:
             # Convert schema to proto format
-            output_schema_proto = self.module_class.get_output_format(llm_format=request.llm_format)
+            output_schema_proto = await self.module_class.get_output_format(
+                llm_format=request.llm_format,
+            )
             output_format_struct = json_format.Parse(
                 text=output_schema_proto,
                 message=struct_pb2.Struct(),  # pylint: disable=no-member
@@ -411,7 +426,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         # Get setup schema if available
         try:
             # Convert schema to proto format
-            setup_schema_proto = self.module_class.get_setup_format(llm_format=request.llm_format)
+            setup_schema_proto = await self.module_class.get_setup_format(llm_format=request.llm_format)
             setup_format_struct = json_format.Parse(
                 text=setup_schema_proto,
                 message=struct_pb2.Struct(),  # pylint: disable=no-member
@@ -428,7 +443,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
             setup_schema=setup_format_struct,
         )
 
-    def GetModuleSecret(  # noqa: N802
+    async def GetModuleSecret(  # noqa: N802
         self,
         request: information_pb2.GetModuleSecretRequest,
         context: grpc.ServicerContext,
@@ -447,7 +462,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         # Get secret schema if available
         try:
             # Convert schema to proto format
-            secret_schema_proto = self.module_class.get_secret_format(llm_format=request.llm_format)
+            secret_schema_proto = await self.module_class.get_secret_format(llm_format=request.llm_format)
             secret_format_struct = json_format.Parse(
                 text=secret_schema_proto,
                 message=struct_pb2.Struct(),  # pylint: disable=no-member
@@ -483,7 +498,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         # Get setup schema if available
         try:
             # Convert schema to proto format
-            config_setup_schema_proto = self.module_class.get_config_setup_format(llm_format=request.llm_format)
+            config_setup_schema_proto = await self.module_class.get_config_setup_format(llm_format=request.llm_format)
             config_setup_format_struct = json_format.Parse(
                 text=config_setup_schema_proto,
                 message=struct_pb2.Struct(),  # pylint: disable=no-member

digitalkin/grpc_servers/registry_server.py

@@ -1,14 +1,14 @@
 """Registry gRPC server implementation for DigitalKin."""
 
-from digitalkin_proto.digitalkin.module_registry.v2 import (
+from digitalkin_proto.agentic_mesh_protocol.module_registry.v1 import (
     module_registry_service_pb2,
     module_registry_service_pb2_grpc,
 )
 
 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
 from digitalkin.logger import logger
+from digitalkin.models.grpc_servers.models import RegistryServerConfig
 
 
 class RegistryServer(BaseServer):

digitalkin/grpc_servers/registry_servicer.py

@@ -9,7 +9,7 @@ from collections.abc import Iterator
 from enum import Enum
 
 import grpc
-from digitalkin_proto.digitalkin.module_registry.v2 import (
+from digitalkin_proto.agentic_mesh_protocol.module_registry.v1 import (
     discover_pb2,
     metadata_pb2,
     module_registry_service_pb2_grpc,
@@ -344,7 +344,7 @@ class RegistryServicer(module_registry_service_pb2_grpc.ModuleRegistryServiceSer
             return status_pb2.ModuleStatusResponse()
 
         module = self.registered_modules[request.module_id]
-        return status_pb2.ModuleStatusResponse(module_id=module.module_id, status=module.status.name)
+        return status_pb2.ModuleStatusResponse(module_id=module.module_id, status=module.status.value)
 
     def ListModuleStatus(  # noqa: N802
         self,
@@ -379,7 +379,7 @@ class RegistryServicer(module_registry_service_pb2_grpc.ModuleRegistryServiceSer
             list_size = len(self.registered_modules)
 
         modules_statuses = [
-            status_pb2.ModuleStatusResponse(module_id=module.module_id, status=module.status.name)
+            status_pb2.ModuleStatusResponse(module_id=module.module_id, status=module.status.value)
             for module in list(self.registered_modules.values())[request.offset : request.offset + list_size]
         ]
 
@@ -409,7 +409,7 @@ class RegistryServicer(module_registry_service_pb2_grpc.ModuleRegistryServiceSer
         for module in self.registered_modules.values():
             yield status_pb2.ModuleStatusResponse(
                 module_id=module.module_id,
-                status=module.status.name,
+                status=module.status.value,
             )
 
     def UpdateModuleStatus(  # noqa: N802

digitalkin/grpc_servers/utils/__init__.py

@@ -0,0 +1 @@
+"""gRPC servers utilities package."""

digitalkin/grpc_servers/utils/exceptions.py

@@ -27,11 +27,3 @@ class ServerStateError(ServerError):
 
 class ReflectionError(ServerError):
     """Error related to gRPC reflection service."""
-
-
-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

@@ -6,8 +6,8 @@ from typing import Any
 import grpc
 
 from digitalkin.grpc_servers.utils.exceptions import ServerError
-from digitalkin.grpc_servers.utils.models import ClientConfig, SecurityMode
 from digitalkin.logger import logger
+from digitalkin.models.grpc_servers.models import ClientConfig, SecurityMode
 
 
 class GrpcClientWrapper:
@@ -62,11 +62,11 @@ class GrpcClientWrapper:
         """
         try:
             # Call the register method
-            logger.debug("send request to %s", query_endpoint)
+            logger.debug("send request to %s", query_endpoint, extra={"request": request})
             response = getattr(self.stub, query_endpoint)(request)
-            logger.debug("receive response from request to registry: %s", response)
+            logger.debug("receive response from request to %s", query_endpoint, extra={"response": response})
         except grpc.RpcError as e:
-            logger.exception("RPC error during %s: %s", query_endpoint, e.details())
+            logger.exception("RPC error during %s", query_endpoint, extra={"error": e.details()})
             raise ServerError
         else:
             return response

digitalkin/grpc_servers/utils/grpc_error_handler.py

@@ -0,0 +1,53 @@
+"""Shared error handling utilities for gRPC services."""
+
+from collections.abc import Generator
+from contextlib import contextmanager
+from typing import Any
+
+from digitalkin.grpc_servers.utils.exceptions import ServerError
+from digitalkin.logger import logger
+
+
+class GrpcErrorHandlerMixin:
+    """Mixin class providing common gRPC error handling functionality."""
+
+    @contextmanager
+    def handle_grpc_errors(  # noqa: PLR6301
+        self,
+        operation: str,
+        service_error_class: type[Exception] | None = None,
+    ) -> Generator[Any, Any, Any]:
+        """Handle gRPC errors for the given operation.
+
+        Args:
+            operation: Name of the operation being performed.
+            service_error_class: Optional specific service exception class to raise.
+                                If not provided, uses the generic ServerError.
+
+        Yields:
+            Context for the operation.
+
+        Raises:
+            ServerError: For gRPC-related errors.
+            service_error_class: For service-specific errors if provided.
+        """
+        if service_error_class is None:
+            service_error_class = ServerError
+
+        try:
+            yield
+        except service_error_class as e:
+            # Re-raise service-specific errors as-is
+            msg = f"{service_error_class.__name__} in {operation}: {e}"
+            logger.exception(msg)
+            raise service_error_class(msg) from e
+        except ServerError as e:
+            # Handle gRPC server errors
+            msg = f"gRPC {operation} failed: {e}"
+            logger.exception(msg)
+            raise ServerError(msg) from e
+        except Exception as e:
+            # Handle unexpected errors
+            msg = f"Unexpected error in {operation}: {e}"
+            logger.exception(msg)
+            raise service_error_class(msg) from e

digitalkin/logger.py

@@ -11,6 +11,15 @@ 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"
@@ -39,12 +48,10 @@ 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(),
-            "location": f"{record.filename}:{record.lineno}",
-            "function": record.funcName,
+            "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)
@@ -81,28 +88,70 @@ class ColorJSONFormatter(logging.Formatter):
 
         # 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)
-        if os.getenv("RAILWAY_SERVICE_NAME"):
-            json_str.replace("\\n", "\n")
+        json_str = json_str.replace("\\n", "\n")
         return f"{color}{json_str}{self.reset}"
 
 
-logging.basicConfig(
-    level=logging.DEBUG,
-    stream=sys.stdout,
-    datefmt="%Y-%m-%d %H:%M:%S",
+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,
+    },
 )
-
-logging.getLogger("grpc").setLevel(logging.DEBUG)
-logging.getLogger("asyncio").setLevel(logging.DEBUG)
-
-
-logger = logging.getLogger("digitalkin")
-
-if not logger.handlers:
-    ch = logging.StreamHandler()
-    ch.setLevel(logging.INFO)
-    ch.setFormatter(ColorJSONFormatter())
-
-    logger.addHandler(ch)
-    logger.propagate = False

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)