digitalkin 0.3.2.dev2__py3-none-any.whl

modules/cpu_intensive_module.py

@@ -0,0 +1,280 @@
+"""Simple module calling an LLM."""
+
+import logging
+from collections.abc import Callable
+from typing import Any, ClassVar, Literal
+
+from digitalkin.grpc_servers.utils.models import ClientConfig, SecurityMode, ServerConfig, ServerMode
+from pydantic import BaseModel, Field
+
+from digitalkin.modules._base_module import BaseModule
+from digitalkin.services.services_models import ServicesStrategy
+from digitalkin.services.setup.setup_strategy import SetupData
+
+# Configure logging with clear formatting
+logging.basicConfig(
+    level=logging.DEBUG,
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+)
+logger = logging.getLogger(__name__)
+
+
+class MessageInputPayload(BaseModel):
+    """Message trigger model for the CPU Archetype module."""
+
+    payload_type: Literal["message"] = "message"
+    user_prompt: str = Field(
+        ...,
+        title="User Prompt",
+        description="The prompt provided by the user for processing.",
+    )
+
+
+class InputFile(BaseModel):
+    """File model for the CPU Archetype module."""
+
+    name: str = Field(
+        ...,
+        title="File Name",
+        description="The name of the file to be processed.",
+    )
+    content: bytes = Field(
+        ...,
+        title="File Content",
+        description="The content of the file to be processed.",
+    )
+
+    file_type: str = Field(
+        ...,
+        title="File Type",
+        description="The type of the file to be processed.",
+    )
+
+
+class FileInputPayload(BaseModel):
+    """File input model for the CPU Archetype module."""
+
+    payload_type: Literal["file"] = "file"
+    files: list[InputFile] = Field(
+        ...,
+        title="Files",
+        description="List of files to be processed.",
+    )
+
+
+class CPUInput(BaseModel):
+    """Input model defining what data the module expects."""
+
+    payload: MessageInputPayload | FileInputPayload = Field(
+        ...,
+        discriminator="payload_type",
+        title="Payload",
+        description="Either a message or list of file input.",
+    )
+
+
+class MessageOutputPayload(BaseModel):
+    """Message output model for the CPU Archetype module."""
+
+    payload_type: Literal["message"] = "message"
+    user_response: str = Field(
+        ...,
+        title="User Response",
+        description="The response generated by the assistant based on the user prompt.",
+    )
+
+
+class OutputFile(BaseModel):
+    """File model for the CPU Archetype module."""
+
+    name: str = Field(
+        ...,
+        title="File Name",
+        description="The name of the file to be processed.",
+    )
+    url: str | None = Field(
+        ...,
+        title="File URL",
+        description="The URL of the file to be processed.",
+    )
+
+    message: str | None = Field(
+        None,
+        title="Message",
+        description="Optional message associated with the file.",
+    )
+
+
+class FileOutputPayload(BaseModel):
+    """File output model for the CPU Archetype module."""
+
+    payload_type: Literal["file"] = "file"
+    files: list[OutputFile] = Field(
+        ...,
+        title="Files",
+        description="List of files generated by the assistant.",
+    )
+
+
+class CPUOutput(BaseModel):
+    """Output model defining what data the module produces."""
+
+    payload: MessageOutputPayload | FileOutputPayload = Field(
+        ...,
+        discriminator="payload_type",
+        title="Payload",
+        description="Either a message or file response.",
+    )
+
+
+class CPUConfigSetup(BaseModel):
+    """Config Setup model definining data that will be pre-computed for each setup and module instance."""
+
+    files: list[str] = Field(
+        ...,
+        title="Files to embed",
+        description="List of files to embed in the setup lifecycle.",
+    )
+
+
+class CPUSetup(BaseModel):
+    """Setup model defining module configuration parameters."""
+
+    model_name: str = Field(
+        ...,
+        title="Model Name",
+        description="The name of the CPU model to use for processing.",
+    )
+    developer_prompt: str = Field(
+        ...,
+        title="Developer Prompt",
+        description="The developer prompt new versions of system prompt, it defines the behavior of the assistant.",
+    )
+    temperature: float = Field(
+        0.7,
+        title="Temperature",
+        description="Controls the randomness of the model's output. Higher values make output more random.",
+    )
+    max_tokens: int = Field(
+        100,
+        title="Max Tokens",
+        description="The maximum number of tokens to generate in the response.",
+    )
+
+
+class CPUToolSecret(BaseModel):
+    """Secret model defining module configuration parameters."""
+
+
+server_config = ServerConfig(
+    host="[::]",
+    port=50151,
+    mode=ServerMode.ASYNC,
+    security=SecurityMode.INSECURE,
+    max_workers=10,
+    credentials=None,
+)
+
+
+client_config = ClientConfig(
+    host="[::]",
+    port=50151,
+    mode=ServerMode.ASYNC,
+    security=SecurityMode.INSECURE,
+    credentials=None,
+)
+
+
+class CPUIntensiveModule(BaseModule[CPUInput, CPUOutput, CPUSetup, CPUToolSecret, None]):
+    """A CPU endpoint tool module module."""
+
+    name = "CPUIntensiveModule"
+    description = "A module that interacts with CPU API to process text"
+
+    # Define the schema formats for the module
+    input_format = CPUInput
+    output_format = CPUOutput
+    setup_format = CPUSetup
+    secret_format = CPUToolSecret
+
+    # Define module metadata for discovery
+    metadata: ClassVar[dict[str, Any]] = {
+        "name": "CPUIntensiveModule",
+        "description": "Transforms input text using a streaming LLM response.",
+        "version": "1.0.0",
+        "tags": ["text", "transformation", "encryption", "streaming"],
+    }
+    # Define services_config_params with default values
+    services_config_strategies: ClassVar[dict[str, ServicesStrategy | None]] = {}
+    services_config_params: ClassVar[dict[str, dict[str, Any | None] | None]] = {
+        "storage": {
+            "config": {"chat_history": None},
+            "client_config": client_config,
+        },
+        "filesystem": {
+            "client_config": client_config,
+        },
+        "cost": {
+            "config": {},
+            "client_config": client_config,
+        },
+    }
+
+    async def initialize(self, setup_data: SetupData) -> None:
+        """Initialize the module capabilities.
+
+        This method is called when the module is loaded by the server.
+        Use it to set up module-specific resources or configurations.
+        """
+
+    async def run(
+        self,
+        input_data: CPUInput,
+        setup_data: CPUSetup,
+        callback: Callable,
+    ) -> None:
+        """Run the module.
+
+        Args:
+            input_data: Input data for the module
+            setup_data: Setup data for the module
+            callback: Callback function to report progress
+
+        Raises:
+            ValueError: If the payload type is unknown
+        """
+        # Validate the input data
+        input_model = self.input_format.model_validate(input_data)
+        self.setup_format.model_validate(setup_data)
+        logger.debug("Running with input data: %s", input_model)
+
+        if not hasattr(input_model, "payload"):
+            error_msg = "Input data is missing 'payload' field"
+            raise ValueError(error_msg)
+
+        if not hasattr(input_model.payload, "payload_type"):
+            error_msg = "Input payload is missing 'type' field"
+            raise ValueError(error_msg)
+
+        total = 0
+        input = MessageInputPayload.model_validate(input_model.payload).user_prompt
+
+        for i in range(int(input)):
+            total += i * i
+            if i % 100 == 0 or i == int(input) - 1:
+                message_output_payload = MessageOutputPayload(
+                    payload_type="message",
+                    user_response=f"result iteration {i}: {total}",
+                )
+                output_model = self.output_format.model_validate({"payload": message_output_payload})
+                await callback(output_data=output_model)
+        logger.info("Job %s completed", self.job_id)
+
+    async def cleanup(self) -> None:
+        """Clean up any resources when the module is stopped.
+
+        This method is called when the module is being shut down.
+        Use it to close connections, free resources, etc.
+        """
+        logger.info("Cleaning up module %s", self.metadata["name"])
+        # Release any resources here if needed.

modules/dynamic_setup_module.py

@@ -0,0 +1,338 @@
+"""Example module demonstrating dynamic schema fields in SetupModel.
+
+This example shows how to use the Dynamic metadata class with async fetchers
+to populate field schemas (like enums) at runtime. This is useful when the
+available options come from external sources like databases or APIs.
+
+Usage:
+    # Start the module server
+    python examples/modules/dynamic_setup_module.py
+
+    # Or import and use in your own code
+    from examples.modules.dynamic_setup_module import DynamicSetupModule
+"""
+
+import asyncio
+import logging
+from typing import Annotated, Any, ClassVar
+
+from pydantic import BaseModel, Field
+
+from digitalkin.models.module.module_context import ModuleContext
+from digitalkin.models.module.module_types import DataModel, DataTrigger, SetupModel
+from digitalkin.modules._base_module import BaseModule
+from digitalkin.services.services_models import ServicesStrategy
+from digitalkin.utils import Dynamic
+
+# Configure logging
+logging.basicConfig(
+    level=logging.DEBUG,
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+)
+logger = logging.getLogger(__name__)
+
+
+# =============================================================================
+# Simulated External Services (replace with real implementations)
+# =============================================================================
+
+
+class MockModelRegistry:
+    """Simulates an external model registry service.
+
+    In a real application, this would be a connection to a database,
+    API service, or configuration management system.
+    """
+
+    _models: ClassVar[list[str]] = ["gpt-4", "gpt-4-turbo", "gpt-3.5-turbo"]
+    _languages: ClassVar[list[str]] = ["en", "fr", "de", "es", "it", "pt"]
+
+    @classmethod
+    async def fetch_available_models(cls) -> list[str]:
+        """Fetch available models from the registry.
+
+        Simulates an async API call with a small delay.
+        """
+        await asyncio.sleep(0.1)  # Simulate network latency
+        logger.info("Fetched %d models from registry", len(cls._models))
+        return cls._models.copy()
+
+    @classmethod
+    async def fetch_supported_languages(cls) -> list[str]:
+        """Fetch supported languages from the registry."""
+        await asyncio.sleep(0.05)  # Simulate network latency
+        logger.info("Fetched %d languages from registry", len(cls._languages))
+        return cls._languages.copy()
+
+    @classmethod
+    def get_default_model(cls) -> str:
+        """Get the default model (sync fetcher example)."""
+        return cls._models[0] if cls._models else "gpt-4"
+
+
+# =============================================================================
+# Dynamic Fetcher Functions
+# =============================================================================
+
+
+async def fetch_models() -> list[str]:
+    """Async fetcher for available model names.
+
+    This function is called when SetupModel.get_clean_model(force=True)
+    is invoked, typically during module initialization or schema refresh.
+    """
+    return await MockModelRegistry.fetch_available_models()
+
+
+async def fetch_languages() -> list[str]:
+    """Async fetcher for supported languages."""
+    return await MockModelRegistry.fetch_supported_languages()
+
+
+def get_temperature_range() -> dict[str, float]:
+    """Sync fetcher example returning min/max for temperature.
+
+    Demonstrates that fetchers can return any JSON-serializable value,
+    not just lists for enums.
+    """
+    return {"minimum": 0.0, "maximum": 2.0}
+
+
+# =============================================================================
+# Setup Model with Dynamic Fields
+# =============================================================================
+
+
+class DynamicAgentSetup(SetupModel):
+    """Setup model demonstrating dynamic schema fields.
+
+    Fields marked with Dynamic(...) will have their schema values
+    refreshed at runtime when get_clean_model(force=True) is called.
+
+    Attributes:
+        model_name: The LLM model to use. Enum values fetched from registry.
+        language: Output language. Enum values fetched dynamically.
+        temperature: Sampling temperature. Static field for comparison.
+        max_tokens: Maximum tokens to generate.
+        system_prompt: The system prompt for the model.
+    """
+
+    # Dynamic field: enum values fetched asynchronously from model registry
+    model_name: Annotated[str, Dynamic(enum=fetch_models)] = Field(
+        default="gpt-4",
+        title="Model Name",
+        description="The LLM model to use for generation.",
+        json_schema_extra={
+            "config": True,  # Shown in initial configuration
+            "ui:widget": "select",
+        },
+    )
+
+    # Dynamic field: language options fetched asynchronously
+    language: Annotated[str, Dynamic(enum=fetch_languages)] = Field(
+        default="en",
+        title="Output Language",
+        description="The language for generated responses.",
+        json_schema_extra={
+            "config": True,
+            "ui:widget": "select",
+        },
+    )
+
+    # Static field: no dynamic fetcher, values defined at class time
+    temperature: float = Field(
+        default=0.7,
+        ge=0.0,
+        le=2.0,
+        title="Temperature",
+        description="Controls randomness. Higher values = more creative.",
+        json_schema_extra={"config": True},
+    )
+
+    # Static field with hidden flag (runtime-only, not in initial config)
+    max_tokens: int = Field(
+        default=1024,
+        ge=1,
+        le=4096,
+        title="Max Tokens",
+        description="Maximum tokens in the response.",
+        json_schema_extra={"hidden": True},
+    )
+
+    # Static field without any special flags
+    system_prompt: str = Field(
+        default="You are a helpful assistant.",
+        title="System Prompt",
+        description="The system prompt defining assistant behavior.",
+    )
+
+
+# =============================================================================
+# Input/Output Models (Using DataModel/DataTrigger pattern)
+# =============================================================================
+
+
+class MessageInputTrigger(DataTrigger):
+    """Message input trigger following DigitalKin DataTrigger pattern.
+
+    The protocol field determines which trigger handler processes this input.
+    """
+
+    protocol: str = "message"
+    content: str = Field(default="", description="The user message content.")
+
+
+class DynamicModuleInput(DataModel[MessageInputTrigger]):
+    """Input model following DigitalKin DataModel pattern.
+
+    Wraps the trigger in a root field with optional annotations.
+    """
+
+    root: MessageInputTrigger = Field(default_factory=MessageInputTrigger)
+
+
+class MessageOutputTrigger(DataTrigger):
+    """Message output trigger following DigitalKin DataTrigger pattern."""
+
+    protocol: str = "message"
+    content: str = Field(default="", description="The generated response.")
+    model_used: str = Field(default="", description="The model that generated this response.")
+    language: str = Field(default="", description="The output language.")
+
+
+class DynamicModuleOutput(DataModel[MessageOutputTrigger]):
+    """Output model following DigitalKin DataModel pattern."""
+
+    root: MessageOutputTrigger = Field(default_factory=MessageOutputTrigger)
+
+
+class DynamicModuleSecret(BaseModel):
+    """Secret model (empty for this example)."""
+
+
+# =============================================================================
+# Module Implementation
+# =============================================================================
+
+
+class DynamicSetupModule(
+    BaseModule[
+        DynamicModuleInput,
+        DynamicModuleOutput,
+        DynamicAgentSetup,
+        DynamicModuleSecret,
+    ]
+):
+    """Example module demonstrating dynamic schema in SetupModel.
+
+    This module shows how to:
+    1. Define setup fields with Dynamic() metadata for runtime enum fetching
+    2. Mix static and dynamic fields in the same SetupModel
+    3. Use async fetchers that simulate external service calls
+    4. Follow DigitalKin's DataModel/DataTrigger pattern for I/O
+
+    The key integration point is in the gRPC servicer, which calls
+    SetupModel.get_clean_model(force=True) to refresh dynamic values
+    before returning schema information to clients.
+    """
+
+    name = "DynamicSetupModule"
+    description = "Demonstrates dynamic schema fields in module setup"
+
+    # Schema format definitions
+    input_format = DynamicModuleInput
+    output_format = DynamicModuleOutput
+    setup_format = DynamicAgentSetup
+    secret_format = DynamicModuleSecret
+
+    # Module metadata
+    metadata: ClassVar[dict[str, Any]] = {
+        "name": "DynamicSetupModule",
+        "description": "Example module with dynamic setup schema",
+        "version": "1.0.0",
+        "tags": ["example", "dynamic-schema"],
+    }
+
+    # Services configuration (empty for this example)
+    services_config_strategies: ClassVar[dict[str, ServicesStrategy | None]] = {}
+    services_config_params: ClassVar[dict[str, dict[str, Any | None] | None]] = {}
+
+    async def initialize(self, context: ModuleContext, setup_data: DynamicAgentSetup) -> None:
+        """Initialize the module with setup data.
+
+        Args:
+            context: The module context with services and session info.
+            setup_data: The validated setup configuration.
+        """
+        logger.info(
+            "Initializing DynamicSetupModule with model=%s, language=%s",
+            setup_data.model_name,
+            setup_data.language,
+        )
+        self.setup = setup_data
+
+    async def cleanup(self) -> None:
+        """Clean up resources."""
+        logger.info("Cleaning up DynamicSetupModule")
+
+
+# =============================================================================
+# Demonstration Script
+# =============================================================================
+
+
+async def demonstrate_dynamic_schema() -> None:
+    """Demonstrate the dynamic schema functionality."""
+    # 1. Show schema WITHOUT force (dynamic fields not resolved)
+
+    model_no_force = await DynamicAgentSetup.get_clean_model(
+        config_fields=True,
+        hidden_fields=False,
+        force=False,
+    )
+    schema_no_force = model_no_force.model_json_schema()
+
+    # Check if enum is present
+    model_name_schema = schema_no_force.get("properties", {}).get("model_name", {})
+    if "enum" in model_name_schema:
+        pass
+
+    # 2. Show schema WITH force (dynamic fields resolved)
+
+    model_with_force = await DynamicAgentSetup.get_clean_model(
+        config_fields=True,
+        hidden_fields=False,
+        force=True,
+    )
+    schema_with_force = model_with_force.model_json_schema()
+
+    # Check enum values after force
+    model_name_schema = schema_with_force.get("properties", {}).get("model_name", {})
+    if "enum" in model_name_schema:
+        pass
+
+    language_schema = schema_with_force.get("properties", {}).get("language", {})
+    if "enum" in language_schema:
+        pass
+
+    # 3. Show that static json_schema_extra is preserved
+
+    # 4. Show field filtering
+
+    # Config fields only (hidden excluded)
+    await DynamicAgentSetup.get_clean_model(
+        config_fields=True,
+        hidden_fields=False,
+        force=False,
+    )
+
+    # All fields including hidden
+    await DynamicAgentSetup.get_clean_model(
+        config_fields=True,
+        hidden_fields=True,
+        force=False,
+    )
+
+
+if __name__ == "__main__":
+    asyncio.run(demonstrate_dynamic_schema())