digitalkin 0.3.2.dev2__py3-none-any.whl → 0.3.2.dev3__py3-none-any.whl

digitalkin/models/module/setup_types.py

@@ -0,0 +1,463 @@
+"""Setup model types with dynamic schema resolution."""
+
+from __future__ import annotations
+
+import copy
+import types
+import typing
+from typing import TYPE_CHECKING, Any, Generic, TypeVar, cast, get_args, get_origin
+
+from pydantic import BaseModel, ConfigDict, create_model
+
+from digitalkin.logger import logger
+from digitalkin.models.module.tool_reference import ToolReference
+from digitalkin.utils.dynamic_schema import (
+    DynamicField,
+    get_fetchers,
+    has_dynamic,
+    resolve_safe,
+)
+
+if TYPE_CHECKING:
+    from pydantic.fields import FieldInfo
+
+    from digitalkin.services.registry import RegistryStrategy
+
+SetupModelT = TypeVar("SetupModelT", bound="SetupModel")
+
+
+class SetupModel(BaseModel, Generic[SetupModelT]):
+    """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 = field_info.json_schema_extra or {}
+            is_config = bool(extra.get("config", False)) if isinstance(extra, dict) else False
+            is_hidden = bool(extra.get("hidden", False)) if isinstance(extra, dict) else 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)
+                        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)
+                    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 = field_info.json_schema_extra or {}
+        new_extra = {**extra, **result.values} if isinstance(extra, dict) else result.values
+
+        # Create a deep copy of the FieldInfo to avoid shared mutable state
+        new_field_info = copy.deepcopy(field_info)
+        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)]
+        new_field_info.metadata = new_metadata
+
+        logger.debug(
+            "Refreshed '%s' with dynamic values: %s",
+            field_name,
+            list(result.values.keys()),
+        )
+
+        return new_field_info
+
+    def resolve_tool_references(self, registry: RegistryStrategy) -> None:
+        """Resolve all ToolReference fields in this setup instance.
+
+        Recursively walks through all fields, including nested BaseModel instances,
+        and resolves any ToolReference fields using the provided registry.
+
+        Args:
+            registry: Registry service to use for resolution.
+        """
+        self._resolve_tool_references_recursive(self, registry)
+
+    @classmethod
+    def _resolve_tool_references_recursive(
+        cls,
+        model_instance: BaseModel,
+        registry: RegistryStrategy,
+    ) -> None:
+        """Recursively resolve ToolReference fields in a model instance.
+
+        Args:
+            model_instance: The model instance to process.
+            registry: Registry service to use for resolution.
+        """
+        for field_name, field_value in model_instance.__dict__.items():
+            if field_value is None:
+                continue
+
+            cls._resolve_field_value(field_name, field_value, registry)
+
+    @classmethod
+    def _resolve_field_value(
+        cls,
+        field_name: str,
+        field_value: BaseModel | ToolReference | list | dict,
+        registry: RegistryStrategy,
+    ) -> None:
+        """Resolve a single field value, handling different types.
+
+        Args:
+            field_name: Name of the field for logging.
+            field_value: The value to process.
+            registry: Registry service to use for resolution.
+        """
+        if isinstance(field_value, ToolReference):
+            cls._resolve_single_tool_reference(field_name, field_value, registry)
+        elif isinstance(field_value, BaseModel):
+            cls._resolve_tool_references_recursive(field_value, registry)
+        elif isinstance(field_value, list):
+            cls._resolve_list_items(field_value, registry)
+        elif isinstance(field_value, dict):
+            cls._resolve_dict_values(field_value, registry)
+
+    @classmethod
+    def _resolve_single_tool_reference(
+        cls,
+        field_name: str,
+        tool_ref: ToolReference,
+        registry: RegistryStrategy,
+    ) -> None:
+        """Resolve a single ToolReference instance.
+
+        Args:
+            field_name: Name of the field for logging.
+            tool_ref: The ToolReference instance.
+            registry: Registry service to use for resolution.
+        """
+        try:
+            tool_ref.resolve(registry)
+            logger.debug(
+                "Resolved ToolReference field '%s'",
+                field_name,
+                extra={"field_name": field_name, "mode": tool_ref.config.mode.value},
+            )
+        except Exception:
+            logger.exception(
+                "Failed to resolve ToolReference field '%s'",
+                field_name,
+                extra={"field_name": field_name, "config": tool_ref.config.model_dump()},
+            )
+
+    @classmethod
+    def _resolve_list_items(
+        cls,
+        items: list,
+        registry: RegistryStrategy,
+    ) -> None:
+        """Resolve ToolReference instances in a list.
+
+        Args:
+            items: List of items to process.
+            registry: Registry service to use for resolution.
+        """
+        for item in items:
+            if isinstance(item, ToolReference):
+                cls._resolve_single_tool_reference("list_item", item, registry)
+            elif isinstance(item, BaseModel):
+                cls._resolve_tool_references_recursive(item, registry)
+
+    @classmethod
+    def _resolve_dict_values(
+        cls,
+        mapping: dict,
+        registry: RegistryStrategy,
+    ) -> None:
+        """Resolve ToolReference instances in a dict's values.
+
+        Args:
+            mapping: Dict to process.
+            registry: Registry service to use for resolution.
+        """
+        for item in mapping.values():
+            if isinstance(item, ToolReference):
+                cls._resolve_single_tool_reference("dict_value", item, registry)
+            elif isinstance(item, BaseModel):
+                cls._resolve_tool_references_recursive(item, registry)

digitalkin/models/module/tool_reference.py

@@ -0,0 +1,105 @@
+"""Tool reference types for archetype module configuration."""
+
+from enum import Enum
+from typing import TYPE_CHECKING
+
+from pydantic import BaseModel, Field, PrivateAttr, model_validator
+
+from digitalkin.models.services.registry import ModuleInfo
+
+if TYPE_CHECKING:
+    from digitalkin.services.registry import RegistryStrategy
+
+
+class ToolSelectionMode(str, Enum):
+    """Mode for tool selection in archetype setup."""
+
+    FIXED = "fixed"
+    TAG = "tag"
+    DISCOVERABLE = "discoverable"
+
+
+class ToolReferenceConfig(BaseModel):
+    """Configuration for how a tool should be selected."""
+
+    mode: ToolSelectionMode = Field(default=ToolSelectionMode.FIXED)
+    fixed_id: str | None = Field(default=None, description="Module ID for FIXED mode")
+    tag: str | None = Field(default=None, description="Search tag for TAG mode")
+    organization_id: str | None = Field(default=None, description="Filter by organization")
+
+    @model_validator(mode="after")
+    def validate_config(self) -> "ToolReferenceConfig":
+        """Validate configuration based on mode.
+
+        Returns:
+            Self if validation passes.
+
+        Raises:
+            ValueError: If required field is missing for the mode.
+        """
+        if self.mode == ToolSelectionMode.FIXED and not self.fixed_id:
+            msg = "fixed_id required when mode is FIXED"
+            raise ValueError(msg)
+        if self.mode == ToolSelectionMode.TAG and not self.tag:
+            msg = "tag required when mode is TAG"
+            raise ValueError(msg)
+        return self
+
+
+class ToolReference(BaseModel):
+    """Reference to a tool module for archetype configuration.
+
+    Frontend sets config, backend resolves to actual ModuleInfo.
+    The resolved module_id is persisted in selected_module_id for
+    subsequent StartModule calls without re-resolution.
+    """
+
+    config: ToolReferenceConfig
+    selected_module_id: str | None = Field(default=None, description="Resolved module ID after resolution")
+    _cached_info: ModuleInfo | None = PrivateAttr(default=None)
+
+    @property
+    def module_info(self) -> ModuleInfo | None:
+        """Get cached ModuleInfo if resolved."""
+        return self._cached_info
+
+    @property
+    def is_resolved(self) -> bool:
+        """Check if this reference has been resolved."""
+        return self._cached_info is not None or self.selected_module_id is not None
+
+    def resolve(self, registry: "RegistryStrategy") -> ModuleInfo | None:
+        """Resolve this reference using the provided registry.
+
+        For FIXED mode, looks up by fixed_id.
+        For TAG mode, searches by tag and takes first result.
+        For DISCOVERABLE mode, returns None (LLM handles at runtime).
+
+        Args:
+            registry: Registry service to use for resolution.
+
+        Returns:
+            ModuleInfo if resolved, None if not resolvable or DISCOVERABLE mode.
+        """
+        if self.config.mode == ToolSelectionMode.DISCOVERABLE:
+            return None
+
+        if self.config.mode == ToolSelectionMode.FIXED and self.config.fixed_id:
+            info = registry.discover_by_id(self.config.fixed_id)
+            if info:
+                self._cached_info = info
+                self.selected_module_id = self.config.fixed_id
+            return info
+
+        if self.config.mode == ToolSelectionMode.TAG and self.config.tag:
+            results = registry.search(
+                name=self.config.tag,
+                module_type="tool",
+                organization_id=self.config.organization_id,
+            )
+            if results:
+                self._cached_info = results[0]
+                self.selected_module_id = results[0].module_id
+                return results[0]
+
+        return None

digitalkin/models/module/utility.py

@@ -4,11 +4,12 @@ These protocols are automatically available to all modules and don't need to be
 explicitly included in module output unions.
 """
 
+from datetime import datetime, timezone
 from typing import Any, ClassVar, Literal
 
 from pydantic import BaseModel, Field
 
-from digitalkin.models.module.module_types import DataTrigger
+from digitalkin.models.module.base_types import DataTrigger
 
 
 class UtilityProtocol(DataTrigger):
@@ -27,6 +28,26 @@ class EndOfStreamOutput(UtilityProtocol):
     protocol: Literal["end_of_stream"] = "end_of_stream"  # type: ignore
 
 
+class ModuleStartInfoOutput(UtilityProtocol):
+    """Output sent when module starts with execution context.
+
+    This protocol is sent as the first message when a module starts,
+    providing the client with essential execution context information.
+    """
+
+    protocol: Literal["module_start_info"] = "module_start_info"  # type: ignore
+    job_id: str = Field(..., description="Unique job identifier")
+    mission_id: str = Field(..., description="Mission identifier")
+    setup_id: str = Field(..., description="Setup identifier")
+    setup_version_id: str = Field(..., description="Setup version identifier")
+    module_id: str = Field(..., description="Module identifier")
+    module_name: str = Field(..., description="Human-readable module name")
+    started_at: str = Field(
+        default_factory=lambda: datetime.now(tz=timezone.utc).isoformat(),
+        description="ISO timestamp when module started",
+    )
+
+
 class HealthcheckPingInput(UtilityProtocol):
     """Input for healthcheck ping request."""
 

digitalkin/modules/_base_module.py

@@ -18,7 +18,9 @@ from digitalkin.models.module.module_types import (
     SecretModelT,
     SetupModelT,
 )
-from digitalkin.models.module.utility import EndOfStreamOutput
+from digitalkin.models.module.tool_reference import ToolReference
+from digitalkin.models.module.utility import EndOfStreamOutput, ModuleStartInfoOutput
+from digitalkin.models.services.registry import ModuleInfo
 from digitalkin.models.services.storage import BaseRole
 from digitalkin.modules.trigger_handler import TriggerHandler
 from digitalkin.services.services_config import ServicesConfig, ServicesStrategy
@@ -76,6 +78,7 @@ class BaseModule(  # noqa: PLR0904
                 registry: RegistryStrategy
                 snapshot: SnapshotStrategy
                 storage: StorageStrategy
+                user_profile: UserProfileStrategy
         """
         logger.debug("Service initialisation: %s", self.services_config_strategies.keys())
         return {
@@ -436,6 +439,22 @@ class BaseModule(  # noqa: PLR0904
         else:
             self._status = ModuleStatus.STOPPING
 
+    @staticmethod
+    def _extract_tool_cache(setup_data: SetupModelT) -> dict[str, ModuleInfo]:
+        """Extract resolved ToolReference info from setup data.
+
+        Args:
+            setup_data: The setup model containing potential ToolReference fields.
+
+        Returns:
+            Dict mapping field names to resolved ModuleInfo.
+        """
+        cache: dict[str, ModuleInfo] = {}
+        for name, value in setup_data.__dict__.items():
+            if isinstance(value, ToolReference) and value.module_info:
+                cache[name] = value.module_info
+        return cache
+
     async def start(
         self,
         input_data: InputModelT,
@@ -446,7 +465,26 @@ class BaseModule(  # noqa: PLR0904
         """Start the module."""
         try:
             self.context.callbacks.send_message = callback
-            logger.info(f"Inititalize module {self.context.session.job_id}")
+
+            # Populate tool cache from resolved ToolReference fields
+            self.context.tool_cache = self._extract_tool_cache(setup_data)
+
+            # Send module start info as first message
+            await callback(
+                DataModel(
+                    root=ModuleStartInfoOutput(
+                        job_id=self.context.session.job_id,
+                        mission_id=self.context.session.mission_id,
+                        setup_id=self.context.session.setup_id,
+                        setup_version_id=self.context.session.setup_version_id,
+                        module_id=self.get_module_id(),
+                        module_name=self.name,
+                    ),
+                    annotations={"role": BaseRole.SYSTEM},
+                )
+            )
+
+            logger.info("Initialize module %s", self.context.session.job_id)
             await self.initialize(self.context, setup_data)
         except Exception as e:
             self._status = ModuleStatus.FAILED
@@ -467,7 +505,7 @@ class BaseModule(  # noqa: PLR0904
         try:
             logger.debug("Init the discovered input handlers.")
             self.triggers_discoverer.init_handlers(self.context)
-            logger.debug(f"Run lifecycle {self.context.session.job_id}")
+            logger.debug("Run lifecycle %s", self.context.session.job_id)
             await self._run_lifecycle(input_data, setup_data)
         except Exception:
             self._status = ModuleStatus.FAILED

digitalkin/modules/triggers/__init__.py

@@ -6,7 +6,3 @@ They provide standard functionality available to all modules.
 Note: These are internal triggers. External code should not import them directly.
 Use UtilityRegistry.get_builtin_triggers() to access the trigger classes.
 """
-
-# No public exports - all triggers are internal
-# Access via: UtilityRegistry.get_builtin_triggers()
-__all__: list[str] = []