digitalkin 0.3.1.dev1__py3-none-any.whl → 0.3.2a2__py3-none-any.whl

digitalkin/models/module/tool_cache.py

@@ -0,0 +1,230 @@
+"""Tool cache for resolved tool references."""
+
+from typing import TYPE_CHECKING, Any
+
+from pydantic import BaseModel, Field
+
+from digitalkin.logger import logger
+from digitalkin.models.services.registry import ModuleInfo
+
+
+class SelectedTool(BaseModel):
+    """Selected tool information."""
+
+    setup_id: str = ""
+    module_id: str = ""
+    slug: str = ""
+    name: str = ""
+
+
+if TYPE_CHECKING:
+    from digitalkin.services.communication import CommunicationStrategy
+
+
+class ToolParameter(BaseModel):
+    """Definition of a single tool parameter.
+
+    Attributes:
+        name: Parameter name.
+        type: JSON Schema type (string, integer, number, boolean, array, object).
+        description: Parameter description for the LLM.
+        required: Whether this parameter is required.
+        enum: Optional list of allowed values.
+        items: Optional schema for array item types.
+        properties: Optional schema for object properties.
+    """
+
+    name: str
+    type: str
+    description: str
+    required: bool = True
+    enum: list[str] | None = None
+    items: dict[str, Any] | None = None
+    properties: dict[str, Any] | None = None
+
+
+class ToolDefinition(BaseModel):
+    """Complete definition of an LLM tool with grouped parameters.
+
+    Attributes:
+        name: Tool name (from protocol const or trigger class name).
+        description: Tool description (from trigger docstring).
+        parameters: List of parameter definitions.
+    """
+
+    name: str
+    description: str
+    parameters: list[ToolParameter] = Field(default_factory=list)
+
+
+class ToolModuleInfo(ModuleInfo):
+    """Module info for tool modules."""
+
+    tools: list[ToolDefinition] = Field(default_factory=list)
+    tool_name: str = ""
+    setup_id: str = ""
+    cost_config: dict[str, Any] = Field(default_factory=dict)
+
+    @property
+    def slug(self) -> str:
+        """Module ID."""
+        return self.setup_id + "_" + self.tool_name
+
+
+class ToolCache(BaseModel):
+    """Registry cache storing resolved tool references by setup field name."""
+
+    entries: dict[str, ToolModuleInfo] = Field(default_factory=dict)
+
+    def add(self, tool_module_info: ToolModuleInfo) -> None:
+        """Add a tool to the cache.
+
+        Args:
+            tool_module_info: Resolved tool module information.
+        """
+        self.entries[tool_module_info.slug] = tool_module_info
+        logger.debug(
+            "Tool cached",
+            extra={
+                "slug": tool_module_info.slug,
+                "module_id": tool_module_info.module_id,
+                "setup_id": tool_module_info.setup_id,
+            },
+        )
+
+    def get(
+        self,
+        slug: str,
+    ) -> ToolModuleInfo | None:
+        """Get a tool from cache, optionally querying registry on miss.
+
+        Args:
+            slug: Field name to look up.
+
+        Returns:
+            ToolModuleInfo if found, None otherwise.
+        """
+        return self.entries.get(slug)
+
+    def clear(self) -> None:
+        """Clear all cache entries."""
+        self.entries.clear()
+
+    def list_tools(self) -> list[str]:
+        """List all cached tool names.
+
+        Returns:
+            List of setup field names in cache.
+        """
+        return list(self.entries.keys())
+
+
+async def module_info_to_tool_module_info(
+    module_info: ModuleInfo,
+    setup_id: str,
+    tool_name: str,
+    communication: "CommunicationStrategy",
+    *,
+    llm_format: bool = True,
+) -> ToolModuleInfo:
+    """Convert ModuleInfo to ToolModuleInfo by fetching schemas via gRPC.
+
+    Fetches the module's input schema and extracts tool definitions from
+    the discriminated union structure.
+
+    Args:
+        module_info: Module info from registry.
+        setup_id: Setup ID of the selected tool.
+        tool_name: Name of the tool.
+        communication: Communication strategy for gRPC calls.
+        llm_format: Use LLM-friendly schema format.
+
+    Returns:
+        ToolModuleInfo with tools extracted from input schema.
+    """
+    schemas = await communication.get_module_schemas(
+        module_info.address,
+        module_info.port,
+        llm_format=llm_format,
+    )
+
+    input_schema = schemas.get("input", {})
+    if llm_format:
+        input_schema = input_schema.get("json_schema", input_schema)
+
+    tools = _extract_tools_from_schema(input_schema)
+    cost_config = schemas.get("cost", {})
+
+    return ToolModuleInfo(
+        module_id=module_info.module_id,
+        module_type=module_info.module_type,
+        address=module_info.address,
+        port=module_info.port,
+        version=module_info.version,
+        module_name=module_info.module_name,
+        documentation=module_info.documentation,
+        status=module_info.status,
+        tools=tools,
+        setup_id=setup_id,
+        tool_name=tool_name,
+        cost_config=cost_config,
+    )
+
+
+def _extract_tools_from_schema(schema: dict[str, Any]) -> list[ToolDefinition]:
+    """Extract tool definitions from a discriminated union input schema.
+
+    Args:
+        schema: JSON schema with $defs containing protocol-based types.
+
+    Returns:
+        List of ToolDefinition with parameters grouped by trigger.
+    """
+    tools: list[ToolDefinition] = []
+    defs = schema.get("$defs", {})
+
+    # Skip SDK utility protocols
+    utility_protocols = {"HealthcheckPingInput", "HealthcheckServicesInput", "HealthcheckStatusInput"}
+
+    for def_name, def_schema in defs.items():
+        if def_name in utility_protocols:
+            continue
+
+        properties = def_schema.get("properties", {})
+        protocol_prop = properties.get("protocol", {})
+
+        # Skip if no protocol const (not a tool input type)
+        if "const" not in protocol_prop:
+            continue
+
+        # Extract tool-level info from trigger
+        tool_name = protocol_prop.get("const", def_name)
+        tool_description = def_schema.get("description", "")
+
+        required_fields = set(def_schema.get("required", []))
+        parameters: list[ToolParameter] = []
+
+        for prop_name, prop_info in properties.items():
+            if prop_name in {"protocol", "created_at"}:
+                continue
+
+            param = ToolParameter(
+                name=prop_name,
+                type=prop_info.get("type", "string"),
+                description=prop_info.get("description", ""),
+                required=prop_name in required_fields,
+                enum=prop_info.get("enum"),
+                items=prop_info.get("items"),
+                properties=prop_info.get("properties"),
+            )
+            parameters.append(param)
+
+        tools.append(
+            ToolDefinition(
+                name=tool_name,
+                description=tool_description,
+                parameters=parameters,
+            )
+        )
+
+    return tools

digitalkin/models/module/tool_reference.py

@@ -0,0 +1,160 @@
+"""Tool reference types for module configuration."""
+
+from typing import Annotated
+
+from pydantic import AfterValidator, BaseModel, BeforeValidator, Field, PlainSerializer
+from pydantic.annotated_handlers import GetJsonSchemaHandler
+from pydantic.json_schema import JsonSchemaValue
+from pydantic_core import CoreSchema
+
+from digitalkin.models.module.tool_cache import ToolModuleInfo, module_info_to_tool_module_info
+from digitalkin.services.communication.communication_strategy import CommunicationStrategy
+from digitalkin.services.registry import RegistryStrategy
+
+
+class ToolReference(BaseModel):
+    """Tool selection containing setup IDs."""
+
+    selected_tools: list[str] = Field(default=[], description="Setup IDs of selected tools.")
+
+    async def resolve(self, registry: RegistryStrategy, communication: CommunicationStrategy) -> list[ToolModuleInfo]:
+        """Resolve selected tools using the registry.
+
+        Args:
+            registry: Registry service for module discovery.
+            communication: Communication service for module schemas.
+
+        Returns:
+            List of ToolModuleInfo for resolved tools.
+        """
+        resolved: list[ToolModuleInfo] = []
+        for setup_id in self.selected_tools:
+            setup = registry.get_setup(setup_id)
+            if setup and setup.module_id:
+                info = registry.discover_by_id(setup.module_id)
+                if info:
+                    resolved.append(await module_info_to_tool_module_info(info, setup_id, setup.name, communication))
+        return resolved
+
+
+class _ToolReferenceInputSchema:
+    """Custom JSON schema generator with configurable maxItems and ui:options."""
+
+    def __init__(
+        self,
+        setup_ids: list[str],
+        module_ids: list[str] | None,
+        tag_ids: list[str],
+        categories: list[str],
+        max_tools: int = 0,
+        min_tools: int = 0,
+    ) -> None:
+        self.setup_ids = setup_ids
+        self.module_ids = module_ids
+        self.tag_ids = tag_ids
+        self.max_tools = max_tools
+        self.min_tools = min_tools
+        self.categories = categories
+
+    def __get_pydantic_json_schema__(  # noqa: PLW3201
+        self,
+        _schema: CoreSchema,
+        _handler: GetJsonSchemaHandler,
+    ) -> JsonSchemaValue:
+        """Generate JSON schema as array for UI, hiding ToolReference complexity.
+
+        Returns:
+            JSON schema as array with ui:widget toolSelect.
+        """
+        json_schema: dict[str, object] = {
+            "type": "array",
+            "items": {"type": "string"},
+        }
+        if self.max_tools > 0:
+            json_schema["maxItems"] = self.max_tools
+        if self.min_tools > 0:
+            json_schema["minItems"] = self.min_tools
+        json_schema["ui:widget"] = "toolSelect"
+        json_schema["ui:options"] = {
+            "setupIds": self.setup_ids,
+            "tagIds": self.tag_ids,
+            "categories": self.categories,
+            "moduleIds": self.module_ids or [],
+            "showModules": self.module_ids is not None,
+        }
+        return json_schema
+
+
+def tool_reference_input(
+    setup_ids: list[str] = [],
+    module_ids: list[str] | None = [],
+    tag_ids: list[str] = [],
+    categories: list[str] = [],
+    max_tools: int = 0,
+    min_tools: int = 0,
+) -> type[ToolReference]:
+    """Create ToolReferenceInput type with schema options and validation.
+
+    Args:
+        setup_ids: Setup IDs for the user to choose from.
+        module_ids: Module IDs for the user to choose from.
+        tag_ids: Tag IDs for the user to choose from.
+        categories: Categories for the user to choose from.
+        max_tools: Maximum tools allowed. 0 for unlimited.
+        min_tools: Minimum tools required. 0 for no minimum.
+
+    Returns:
+        Annotated type for use in Pydantic models.
+    """
+
+    def convert_to_tool_reference(v: object) -> ToolReference | object:
+        """Convert list of setup IDs to ToolReference.
+
+        Returns:
+            ToolReference if input is list, otherwise original value.
+        """
+        if isinstance(v, list):
+            return ToolReference(selected_tools=v)
+        return v
+
+    def validate_tools_count(v: ToolReference) -> ToolReference:
+        """Validate selected_tools count against min/max constraints.
+
+        Returns:
+            The validated ToolReference.
+
+        Raises:
+            ValueError: If count is below min_tools or above max_tools.
+        """
+        count = len(v.selected_tools)
+        if min_tools > 0 and count < min_tools:
+            msg = f"At least {min_tools} tools required, got {count}"
+            raise ValueError(msg)
+        if max_tools > 0 and count > max_tools:
+            msg = f"At most {max_tools} tools allowed, got {count}"
+            raise ValueError(msg)
+        return v
+
+    def serialize_to_list(v: ToolReference) -> list[str]:
+        """Serialize ToolReference as plain list for frontend compatibility.
+
+        Returns:
+            List of setup IDs.
+        """
+        return v.selected_tools
+
+    return Annotated[  # type: ignore[return-value]
+        ToolReference,
+        BeforeValidator(convert_to_tool_reference),
+        AfterValidator(validate_tools_count),
+        PlainSerializer(serialize_to_list, return_type=list[str]),
+        _ToolReferenceInputSchema(
+            setup_ids=setup_ids,
+            module_ids=module_ids,
+            tag_ids=tag_ids,
+            categories=categories,
+            max_tools=max_tools,
+            min_tools=min_tools,
+        ),
+        Field(default=[]),
+    ]

digitalkin/models/module/utility.py

@@ -0,0 +1,167 @@
+"""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 datetime import datetime, timezone
+from typing import Any, ClassVar, Literal
+
+from pydantic import BaseModel, Field
+
+from digitalkin.models.module.base_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 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."""
+
+    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/cost.py

@@ -2,7 +2,7 @@
 
 from datetime import datetime, timezone
 from enum import Enum
-from typing import Any
+from typing import Annotated, Any, Literal
 
 from pydantic import BaseModel, Field
 
@@ -35,6 +35,27 @@ class CostConfig(BaseModel):
     rate: float
 
 
+class QuantityLimit(BaseModel):
+    """Cost limit based on quantity (e.g., max 10000 tokens)."""
+
+    limit_type: Literal["quantity"] = "quantity"
+    name: str
+    type: CostTypeEnum
+    max_value: float
+
+
+class AmountLimit(BaseModel):
+    """Cost limit based on cost amount in dollars (e.g., max $1.00)."""
+
+    limit_type: Literal["amount"] = "amount"
+    name: str
+    type: CostTypeEnum
+    max_value: float
+
+
+CostLimit = Annotated[QuantityLimit | AmountLimit, Field(discriminator="limit_type")]
+
+
 class CostEvent(BaseModel):
     """Pydantic model that represents a cost event registered during service execution.
 

digitalkin/models/services/registry.py

@@ -0,0 +1,77 @@
+"""Registry data models."""
+
+from enum import Enum
+from typing import Any
+
+from pydantic import BaseModel
+
+
+class RegistryModuleStatus(str, Enum):
+    """Module status in the registry."""
+
+    UNSPECIFIED = "unspecified"
+    READY = "ready"
+    ACTIVE = "active"
+    ARCHIVED = "archived"
+
+
+class RegistryModuleType(str, Enum):
+    """Module type in the registry."""
+
+    UNSPECIFIED = "unspecified"
+    ARCHETYPE = "archetype"
+    TOOL = "tool"
+
+
+class ModuleInfo(BaseModel):
+    """Module information from registry."""
+
+    module_id: str = ""
+    module_type: RegistryModuleType = RegistryModuleType.UNSPECIFIED
+    address: str = ""
+    port: int = 0
+    version: str = ""
+    module_name: str = ""
+    documentation: str | None = None
+    status: RegistryModuleStatus | None = None
+
+
+class RegistrySetupStatus(str, Enum):
+    """Setup status in the registry."""
+
+    UNSPECIFIED = "unspecified"
+    DRAFT = "draft"
+    WAITING_FOR_APPROVAL = "waiting_for_approval"
+    READY = "ready"
+    PAUSED = "paused"
+    FAILED = "failed"
+    ARCHIVED = "archived"
+    NEEDS_CONFIGURATION = "needs_configuration"
+    CONFIGURATION_FAILED = "configuration_failed"
+    CONFIGURATION_SUCCEEDED = "configuration_succeeded"
+
+
+class RegistryVisibility(str, Enum):
+    """Visibility in the registry."""
+
+    UNSPECIFIED = "unspecified"
+    PUBLIC = "public"
+    PRIVATE = "private"
+    INTERNAL = "internal"
+
+
+class SetupInfo(BaseModel):
+    """Setup information from registry."""
+
+    setup_id: str
+    name: str
+    documentation: str | None = None
+    status: RegistrySetupStatus | None = None
+    visibility: RegistryVisibility | None = None
+    organization_id: str | None = None
+    owner_id: str | None = None
+    card_id: str | None = None
+    module_id: str | None = None
+    setup_version_id: str | None = None
+    setup_version: str | None = None
+    config: dict[str, Any] | None = None

digitalkin/modules/__init__.py

@@ -4,4 +4,8 @@ from digitalkin.modules.archetype_module import ArchetypeModule
 from digitalkin.modules.tool_module import ToolModule
 from digitalkin.modules.trigger_handler import TriggerHandler
 
-__all__ = ["ArchetypeModule", "ToolModule", "TriggerHandler"]
+__all__ = [
+    "ArchetypeModule",
+    "ToolModule",
+    "TriggerHandler",
+]