digitalkin 0.3.2.dev2__py3-none-any.whl

digitalkin/models/module/module_context.py

@@ -0,0 +1,149 @@
+"""Define the module context used in the triggers."""
+
+import os
+from datetime import tzinfo
+from types import SimpleNamespace
+from typing import Any
+from zoneinfo import ZoneInfo
+
+from digitalkin.services.agent.agent_strategy import AgentStrategy
+from digitalkin.services.communication.communication_strategy import CommunicationStrategy
+from digitalkin.services.cost.cost_strategy import CostStrategy
+from digitalkin.services.filesystem.filesystem_strategy import FilesystemStrategy
+from digitalkin.services.identity.identity_strategy import IdentityStrategy
+from digitalkin.services.registry.registry_strategy import RegistryStrategy
+from digitalkin.services.snapshot.snapshot_strategy import SnapshotStrategy
+from digitalkin.services.storage.storage_strategy import StorageStrategy
+from digitalkin.services.user_profile.user_profile_strategy import UserProfileStrategy
+
+
+class Session(SimpleNamespace):
+    """Session data container with mandatory setup_id and mission_id."""
+
+    job_id: str
+    mission_id: str
+    setup_id: str
+    setup_version_id: str
+    timezone: tzinfo
+
+    def __init__(
+        self,
+        job_id: str,
+        mission_id: str,
+        setup_id: str,
+        setup_version_id: str,
+        timezone: tzinfo | None = None,
+        **kwargs: dict[str, Any],
+    ) -> None:
+        """Init Module Session.
+
+        Raises:
+            ValueError: If mandatory args are missing.
+        """
+        if not setup_id:
+            msg = "setup_id is mandatory"
+            raise ValueError(msg)
+        if not setup_version_id:
+            msg = "setup_version_id is mandatory"
+            raise ValueError(msg)
+        if not mission_id:
+            msg = "mission_id is mandatory"
+            raise ValueError(msg)
+        if not job_id:
+            msg = "job_id is mandatory"
+            raise ValueError(msg)
+
+        self.job_id = job_id
+        self.mission_id = mission_id
+        self.setup_id = setup_id
+        self.setup_version_id = setup_version_id
+        self.timezone = timezone or ZoneInfo(os.environ.get("DIGITALKIN_TIMEZONE", "Europe/Paris"))
+
+        super().__init__(**kwargs)
+
+    def current_ids(self) -> dict[str, str]:
+        """Return current session ids as a dictionary.
+
+        Returns:
+            A dictionary containing the current session ids.
+        """
+        return {
+            "job_id": self.job_id,
+            "mission_id": self.mission_id,
+            "setup_id": self.setup_id,
+            "setup_version_id": self.setup_version_id,
+        }
+
+
+class ModuleContext:
+    """ModuleContext provides a container for strategies and resources used by a module.
+
+    This context object is designed to be passed to module components, providing them with
+    access to shared strategies and resources. Additional attributes may be set dynamically.
+    """
+
+    # services list
+    agent: AgentStrategy
+    communication: CommunicationStrategy
+    cost: CostStrategy
+    filesystem: FilesystemStrategy
+    identity: IdentityStrategy
+    registry: RegistryStrategy
+    snapshot: SnapshotStrategy
+    storage: StorageStrategy
+    user_profile: UserProfileStrategy
+
+    session: Session
+    callbacks: SimpleNamespace
+    metadata: SimpleNamespace
+    helpers: SimpleNamespace
+    state: SimpleNamespace = SimpleNamespace()
+
+    def __init__(  # noqa: PLR0913, PLR0917
+        self,
+        agent: AgentStrategy,
+        communication: CommunicationStrategy,
+        cost: CostStrategy,
+        filesystem: FilesystemStrategy,
+        identity: IdentityStrategy,
+        registry: RegistryStrategy,
+        snapshot: SnapshotStrategy,
+        storage: StorageStrategy,
+        user_profile: UserProfileStrategy,
+        session: dict[str, Any],
+        metadata: dict[str, Any] = {},
+        helpers: dict[str, Any] = {},
+        callbacks: dict[str, Any] = {},
+    ) -> None:
+        """Register mandatory services, session, metadata and callbacks.
+
+        Args:
+            agent: AgentStrategy.
+            communication: CommunicationStrategy.
+            cost: CostStrategy.
+            filesystem: FilesystemStrategy.
+            identity: IdentityStrategy.
+            registry: RegistryStrategy.
+            snapshot: SnapshotStrategy.
+            storage: StorageStrategy.
+            user_profile: UserProfileStrategy.
+            metadata: dict defining differents Module metadata.
+            helpers: dict different user defined helpers.
+            session: dict referring the session IDs or informations.
+            callbacks: Functions allowing user to agent interaction.
+        """
+        # Core services
+        self.agent = agent
+        self.communication = communication
+        self.cost = cost
+        self.filesystem = filesystem
+        self.identity = identity
+        self.registry = registry
+        self.snapshot = snapshot
+        self.storage = storage
+        self.user_profile = user_profile
+
+        self.metadata = SimpleNamespace(**metadata)
+        self.session = Session(**session)
+        self.helpers = SimpleNamespace(**helpers)
+        self.callbacks = SimpleNamespace(**callbacks)

digitalkin/models/module/module_types.py

@@ -0,0 +1,393 @@
+"""Types for module models."""
+
+from __future__ import annotations
+
+import copy
+import types
+import typing
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar, cast, get_args, get_origin
+
+from pydantic import BaseModel, ConfigDict, Field, create_model
+
+from digitalkin.logger import logger
+from digitalkin.utils.dynamic_schema import (
+    DynamicField,
+    get_fetchers,
+    has_dynamic,
+    resolve_safe,
+)
+
+if TYPE_CHECKING:
+    from pydantic.fields import FieldInfo
+
+
+class DataTrigger(BaseModel):
+    """Defines the root input/output model exposing the protocol.
+
+    The mandatory protocol is important to define the module beahvior following the user or agent input/output.
+
+    Example:
+        class MyInput(DataModel):
+            root: DataTrigger
+            user_define_data: Any
+
+        # Usage
+        my_input = MyInput(root=DataTrigger(protocol="message"))
+        print(my_input.root.protocol)  # Output: message
+    """
+
+    protocol: ClassVar[str]
+    created_at: str = Field(
+        default_factory=lambda: datetime.now(tz=timezone.utc).isoformat(),
+        title="Created At",
+        description="Timestamp when the payload was created.",
+    )
+
+
+DataTriggerT = TypeVar("DataTriggerT", bound=DataTrigger)
+
+
+class DataModel(BaseModel, Generic[DataTriggerT]):
+    """Base definition of input/output model showing mandatory root fields.
+
+    The Model define the Module Input/output, usually referring to multiple input/output type defined by an union.
+
+    Example:
+        class ModuleInput(DataModel):
+            root: FileInput | MessageInput
+    """
+
+    root: DataTriggerT
+    annotations: dict[str, str] = Field(
+        default={},
+        title="Annotations",
+        description="Additional metadata or annotations related to the output. ex {'role': 'user'}",
+    )
+
+
+InputModelT = TypeVar("InputModelT", bound=DataModel)
+OutputModelT = TypeVar("OutputModelT", bound=DataModel)
+SecretModelT = TypeVar("SecretModelT", bound=BaseModel)
+SetupModelT = TypeVar("SetupModelT", bound="SetupModel")
+
+
+class SetupModel(BaseModel):
+    """Base definition of setup model showing mandatory root fields.
+
+    Optionally, the setup model can define a config option in json_schema_extra
+    to be used to initialize the Kin. Supports dynamic schema providers for
+    runtime value generation.
+
+    Attributes:
+        model_fields: Inherited from Pydantic BaseModel, contains field definitions.
+
+    See Also:
+        - Documentation: docs/api/dynamic_schema.md
+        - Tests: tests/modules/test_setup_model.py
+    """
+
+    @classmethod
+    async def get_clean_model(
+        cls,
+        *,
+        config_fields: bool,
+        hidden_fields: bool,
+        force: bool = False,
+    ) -> type[SetupModelT]:
+        """Dynamically builds and returns a new BaseModel subclass with filtered fields.
+
+        This method filters fields based on their `json_schema_extra` metadata:
+        - Fields with `{"config": True}` are included only when `config_fields=True`
+        - Fields with `{"hidden": True}` are included only when `hidden_fields=True`
+
+        When `force=True`, fields with dynamic schema providers will have their
+        providers called to fetch fresh values for schema metadata like enums.
+        This includes recursively processing nested BaseModel fields.
+
+        Args:
+            config_fields: If True, include fields marked with `{"config": True}`.
+                These are typically initial configuration fields.
+            hidden_fields: If True, include fields marked with `{"hidden": True}`.
+                These are typically runtime-only fields not shown in initial config.
+            force: If True, refresh dynamic schema fields by calling their providers.
+                Use this when you need up-to-date values from external sources like
+                databases or APIs. Default is False for performance.
+
+        Returns:
+            A new BaseModel subclass with filtered fields.
+        """
+        clean_fields: dict[str, Any] = {}
+
+        for name, field_info in cls.model_fields.items():
+            extra = getattr(field_info, "json_schema_extra", {}) or {}
+            is_config = bool(extra.get("config", False))
+            is_hidden = bool(extra.get("hidden", False))
+
+            # Skip config unless explicitly included
+            if is_config and not config_fields:
+                logger.debug("Skipping '%s' (config-only)", name)
+                continue
+
+            # Skip hidden unless explicitly included
+            if is_hidden and not hidden_fields:
+                logger.debug("Skipping '%s' (hidden-only)", name)
+                continue
+
+            # Refresh dynamic schema fields when force=True
+            current_field_info = field_info
+            current_annotation = field_info.annotation
+
+            if force:
+                # Check if this field has DynamicField metadata
+                if has_dynamic(field_info):
+                    current_field_info = await cls._refresh_field_schema(name, field_info)
+
+                # Check if the annotation is a nested BaseModel that might have dynamic fields
+                nested_model = cls._get_base_model_type(current_annotation)
+                if nested_model is not None:
+                    refreshed_nested = await cls._refresh_nested_model(nested_model)
+                    if refreshed_nested is not nested_model:
+                        # Update annotation to use refreshed nested model
+                        current_annotation = refreshed_nested
+                        # Create new field_info with updated annotation (deep copy for safety)
+                        current_field_info = copy.deepcopy(current_field_info)
+                        setattr(current_field_info, "annotation", current_annotation)
+
+            clean_fields[name] = (current_annotation, current_field_info)
+
+        # Dynamically create a model e.g. "SetupModel"
+        m = create_model(
+            f"{cls.__name__}",
+            __base__=BaseModel,
+            __config__=ConfigDict(arbitrary_types_allowed=True),
+            **clean_fields,
+        )
+        return cast("type[SetupModelT]", m)
+
+    @classmethod
+    def _get_base_model_type(cls, annotation: type | None) -> type[BaseModel] | None:
+        """Extract BaseModel type from an annotation.
+
+        Handles direct types, Optional, Union, list, dict, set, tuple, and other generics.
+
+        Args:
+            annotation: The type annotation to inspect.
+
+        Returns:
+            The BaseModel subclass if found, None otherwise.
+        """
+        if annotation is None:
+            return None
+
+        # Direct BaseModel subclass check
+        if isinstance(annotation, type) and issubclass(annotation, BaseModel):
+            return annotation
+
+        origin = get_origin(annotation)
+        if origin is None:
+            return None
+
+        args = get_args(annotation)
+        return cls._extract_base_model_from_args(origin, args)
+
+    @classmethod
+    def _extract_base_model_from_args(
+        cls,
+        origin: type,
+        args: tuple[type, ...],
+    ) -> type[BaseModel] | None:
+        """Extract BaseModel from generic type arguments.
+
+        Args:
+            origin: The generic origin type (list, dict, Union, etc.).
+            args: The type arguments.
+
+        Returns:
+            The BaseModel subclass if found, None otherwise.
+        """
+        # Union/Optional: check each arg (supports both typing.Union and types.UnionType)
+        # Python 3.10+ uses types.UnionType for X | Y syntax
+        if origin is typing.Union or origin is types.UnionType:
+            return cls._find_base_model_in_args(args)
+
+        # list, set, frozenset: check first arg
+        if origin in {list, set, frozenset} and args:
+            return cls._check_base_model(args[0])
+
+        # dict: check value type (second arg)
+        dict_value_index = 1
+        if origin is dict and len(args) > dict_value_index:
+            return cls._check_base_model(args[dict_value_index])
+
+        # tuple: check first non-ellipsis arg
+        if origin is tuple:
+            return cls._find_base_model_in_args(args, skip_ellipsis=True)
+
+        return None
+
+    @classmethod
+    def _check_base_model(cls, arg: type) -> type[BaseModel] | None:
+        """Check if arg is a BaseModel subclass.
+
+        Returns:
+            The BaseModel subclass if arg is one, None otherwise.
+        """
+        if isinstance(arg, type) and issubclass(arg, BaseModel):
+            return arg
+        return None
+
+    @classmethod
+    def _find_base_model_in_args(
+        cls,
+        args: tuple[type, ...],
+        *,
+        skip_ellipsis: bool = False,
+    ) -> type[BaseModel] | None:
+        """Find first BaseModel in args.
+
+        Returns:
+            The first BaseModel subclass found, None otherwise.
+        """
+        for arg in args:
+            if arg is type(None):
+                continue
+            if skip_ellipsis and arg is ...:
+                continue
+            result = cls._check_base_model(arg)
+            if result is not None:
+                return result
+        return None
+
+    @classmethod
+    async def _refresh_nested_model(cls, model_cls: type[BaseModel]) -> type[BaseModel]:
+        """Refresh dynamic fields in a nested BaseModel.
+
+        Creates a new model class with all DynamicField metadata resolved.
+
+        Args:
+            model_cls: The nested model class to refresh.
+
+        Returns:
+            A new model class with refreshed fields, or the original if no changes.
+        """
+        has_changes = False
+        clean_fields: dict[str, Any] = {}
+
+        for name, field_info in model_cls.model_fields.items():
+            current_field_info = field_info
+            current_annotation = field_info.annotation
+
+            # Check if field has DynamicField metadata
+            if has_dynamic(field_info):
+                current_field_info = await cls._refresh_field_schema(name, field_info)
+                has_changes = True
+
+            # Recursively check nested models
+            nested_model = cls._get_base_model_type(current_annotation)
+            if nested_model is not None:
+                refreshed_nested = await cls._refresh_nested_model(nested_model)
+                if refreshed_nested is not nested_model:
+                    current_annotation = refreshed_nested
+                    current_field_info = copy.deepcopy(current_field_info)
+                    setattr(current_field_info, "annotation", current_annotation)
+                    has_changes = True
+
+            clean_fields[name] = (current_annotation, current_field_info)
+
+        if not has_changes:
+            return model_cls
+
+        # Create new model with refreshed fields
+        logger.debug("Creating refreshed nested model for '%s'", model_cls.__name__)
+        return create_model(
+            model_cls.__name__,
+            __base__=BaseModel,
+            __config__=ConfigDict(arbitrary_types_allowed=True),
+            **clean_fields,
+        )
+
+    @classmethod
+    async def _refresh_field_schema(cls, field_name: str, field_info: FieldInfo) -> FieldInfo:
+        """Refresh a field's json_schema_extra with fresh values from dynamic providers.
+
+        This method calls all dynamic providers registered for a field (via Annotated
+        metadata) and creates a new FieldInfo with the resolved values. The original
+        field_info is not modified.
+
+        Uses `resolve_safe()` for structured error handling, allowing partial success
+        when some fetchers fail. Successfully resolved values are still applied.
+
+        Args:
+            field_name: The name of the field being refreshed (used for logging).
+            field_info: The original FieldInfo object containing the dynamic providers.
+
+        Returns:
+            A new FieldInfo object with the same attributes as the original, but with
+            `json_schema_extra` containing resolved values and Dynamic metadata removed.
+
+        Note:
+            If all fetchers fail, the original field_info is returned unchanged.
+            If some fetchers fail, successfully resolved values are still applied.
+        """
+        fetchers = get_fetchers(field_info)
+
+        if not fetchers:
+            return field_info
+
+        fetcher_keys = list(fetchers.keys())
+        logger.debug(
+            "Refreshing dynamic schema for field '%s' with fetchers: %s",
+            field_name,
+            fetcher_keys,
+            extra={"field_name": field_name, "fetcher_keys": fetcher_keys},
+        )
+
+        # Resolve all fetchers with structured error handling
+        result = await resolve_safe(fetchers)
+
+        # Log any errors that occurred with full details
+        if result.errors:
+            for key, error in result.errors.items():
+                logger.warning(
+                    "Failed to resolve '%s' for field '%s': %s: %s",
+                    key,
+                    field_name,
+                    type(error).__name__,
+                    str(error) or "(no message)",
+                    extra={
+                        "field_name": field_name,
+                        "fetcher_key": key,
+                        "error_type": type(error).__name__,
+                        "error_message": str(error),
+                        "error_repr": repr(error),
+                    },
+                )
+
+        # If no values were resolved, return original field_info
+        if not result.values:
+            logger.warning(
+                "All fetchers failed for field '%s', keeping original",
+                field_name,
+            )
+            return field_info
+
+        # Build new json_schema_extra with resolved values merged
+        extra = getattr(field_info, "json_schema_extra", {}) or {}
+        new_extra = {**extra, **result.values}
+
+        # Create a deep copy of the FieldInfo to avoid shared mutable state
+        new_field_info = copy.deepcopy(field_info)
+        setattr(new_field_info, "json_schema_extra", new_extra)
+
+        # Remove Dynamic from metadata (it's been resolved)
+        new_metadata = [m for m in new_field_info.metadata if not isinstance(m, DynamicField)]
+        setattr(new_field_info, "metadata", new_metadata)
+
+        logger.debug(
+            "Refreshed '%s' with dynamic values: %s",
+            field_name,
+            list(result.values.keys()),
+        )
+
+        return new_field_info

digitalkin/models/module/utility.py

@@ -0,0 +1,146 @@
+"""Utility protocols for SDK-provided functionality.
+
+These protocols are automatically available to all modules and don't need to be
+explicitly included in module output unions.
+"""
+
+from typing import Any, ClassVar, Literal
+
+from pydantic import BaseModel, Field
+
+from digitalkin.models.module.module_types import DataTrigger
+
+
+class UtilityProtocol(DataTrigger):
+    """Base class for SDK-provided utility protocols.
+
+    All SDK utility protocols inherit from this class to enable:
+    - Easy identification of SDK vs user-defined protocols
+    - Auto-injection capability
+    - Consistent behavior across the SDK
+    """
+
+
+class EndOfStreamOutput(UtilityProtocol):
+    """Signal that the stream has ended."""
+
+    protocol: Literal["end_of_stream"] = "end_of_stream"  # type: ignore
+
+
+class HealthcheckPingInput(UtilityProtocol):
+    """Input for healthcheck ping request."""
+
+    protocol: Literal["healthcheck_ping"] = "healthcheck_ping"  # type: ignore
+
+
+class HealthcheckPingOutput(UtilityProtocol):
+    """Output for healthcheck ping response.
+
+    Simple alive check that returns "pong" status.
+    """
+
+    protocol: Literal["healthcheck_ping"] = "healthcheck_ping"  # type: ignore
+    status: Literal["pong"] = "pong"
+    latency_ms: float | None = Field(
+        default=None,
+        description="Round-trip latency in milliseconds",
+    )
+
+
+class ServiceHealthStatus(BaseModel):
+    """Health status of a single service."""
+
+    name: str = Field(..., description="Name of the service")
+    status: Literal["healthy", "unhealthy", "unknown"] = Field(
+        ...,
+        description="Health status of the service",
+    )
+    message: str | None = Field(
+        default=None,
+        description="Optional message about the service status",
+    )
+
+
+class HealthcheckServicesInput(UtilityProtocol):
+    """Input for healthcheck services request."""
+
+    protocol: Literal["healthcheck_services"] = "healthcheck_services"  # type: ignore
+
+
+class HealthcheckServicesOutput(UtilityProtocol):
+    """Output for healthcheck services response.
+
+    Reports the health status of all configured services.
+    """
+
+    protocol: Literal["healthcheck_services"] = "healthcheck_services"  # type: ignore
+    services: list[ServiceHealthStatus] = Field(
+        ...,
+        description="List of service health statuses",
+    )
+    overall_status: Literal["healthy", "degraded", "unhealthy"] = Field(
+        ...,
+        description="Overall health status based on all services",
+    )
+
+
+class HealthcheckStatusInput(UtilityProtocol):
+    """Input for healthcheck status request."""
+
+    protocol: Literal["healthcheck_status"] = "healthcheck_status"  # type: ignore
+
+
+class HealthcheckStatusOutput(UtilityProtocol):
+    """Output for healthcheck status response.
+
+    Comprehensive module status including uptime, active jobs, and metadata.
+    """
+
+    protocol: Literal["healthcheck_status"] = "healthcheck_status"  # type: ignore
+    module_name: str = Field(..., description="Name of the module")
+    module_status: str = Field(..., description="Current status of the module")
+    uptime_seconds: float | None = Field(
+        default=None,
+        description="Module uptime in seconds",
+    )
+    active_jobs: int = Field(
+        default=0,
+        description="Number of currently active jobs",
+    )
+    metadata: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Additional metadata about the module",
+    )
+
+
+class UtilityRegistry:
+    """Registry for SDK-provided built-in triggers.
+
+    Example:
+        builtin_triggers = UtilityRegistry.get_builtin_triggers()
+    """
+
+    _builtin_triggers: ClassVar[tuple | None] = None
+
+    @classmethod
+    def get_builtin_triggers(cls) -> tuple:
+        """Get all SDK-provided built-in trigger handlers.
+
+        Uses lazy loading to avoid circular imports with the modules package.
+
+        Returns:
+            Tuple of TriggerHandler subclasses for built-in functionality.
+        """
+        if cls._builtin_triggers is None:
+            from digitalkin.modules.triggers.healthcheck_ping_trigger import HealthcheckPingTrigger  # noqa: PLC0415
+            from digitalkin.modules.triggers.healthcheck_services_trigger import (  # noqa: PLC0415
+                HealthcheckServicesTrigger,
+            )
+            from digitalkin.modules.triggers.healthcheck_status_trigger import HealthcheckStatusTrigger  # noqa: PLC0415
+
+            cls._builtin_triggers = (
+                HealthcheckPingTrigger,
+                HealthcheckServicesTrigger,
+                HealthcheckStatusTrigger,
+            )
+        return cls._builtin_triggers

digitalkin/models/services/__init__.py

@@ -0,0 +1,10 @@
+"""This module contains the models for the services."""
+
+from digitalkin.models.services.storage import BaseMessage, BaseRole, ChatHistory, Role
+
+__all__ = [
+    "BaseMessage",
+    "BaseRole",
+    "ChatHistory",
+    "Role",
+]