digitalkin 0.2.12__py3-none-any.whl → 0.2.14__py3-none-any.whl

modules/cpu_intensive_module.py

@@ -0,0 +1,281 @@
+"""Simple module calling an LLM."""
+
+import logging
+from collections.abc import Callable
+from typing import Any, ClassVar, Literal
+
+from pydantic import BaseModel, Field
+
+from digitalkin.grpc_servers.utils.models import ClientConfig, SecurityMode, ServerConfig, ServerMode
+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": {
+            "config": {},
+            "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/minimal_llm_module.py

@@ -1,44 +1,166 @@
 """Simple module calling an LLM."""
 
 import logging
+import os
 from collections.abc import Callable
-from typing import Any, ClassVar
+from typing import Any, ClassVar, Literal
 
-import grpc
 import openai
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
-from digitalkin.grpc_servers.utils.models import SecurityMode, ClientConfig, ServerMode
+from digitalkin.grpc_servers.utils.models import ClientConfig, SecurityMode, ServerMode
 from digitalkin.modules._base_module import BaseModule
-from digitalkin.services.setup.setup_strategy import SetupData
+from digitalkin.services.services_models import ServicesStrategy
 
 # Configure logging with clear formatting
 logging.basicConfig(
-    level=logging.INFO,
+    level=logging.DEBUG,
     format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
 )
 logger = logging.getLogger(__name__)
 
 
-# Define schema models using Pydantic
-class OpenAIToolInput(BaseModel):
+class MessageInputPayload(BaseModel):
+    """Message trigger model for the OpenAI 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 OpenAI 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 OpenAI Archetype module."""
+
+    payload_type: Literal["file"] = "file"
+    files: list[InputFile] = Field(
+        ...,
+        title="Files",
+        description="List of files to be processed.",
+    )
+
+
+class OpenAIInput(BaseModel):
     """Input model defining what data the module expects."""
 
-    prompt: str
+    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 OpenAI 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 OpenAIToolOutput(BaseModel):
+class OutputFile(BaseModel):
+    """File model for the OpenAI 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 OpenAI Archetype module."""
+
+    payload_type: Literal["file"] = "file"
+    files: list[OutputFile] = Field(
+        ...,
+        title="Files",
+        description="List of files generated by the assistant.",
+    )
+
+
+class OpenAIOutput(BaseModel):
     """Output model defining what data the module produces."""
 
-    response: str
+    payload: MessageOutputPayload | FileOutputPayload = Field(
+        ...,
+        discriminator="payload_type",
+        title="Payload",
+        description="Either a message or file response.",
+    )
 
 
-class OpenAIToolSetup(BaseModel):
+class OpenAISetup(BaseModel):
     """Setup model defining module configuration parameters."""
 
-    openai_key: str
-    model_name: str
-    dev_prompt: str
+    model_name: str = Field(
+        ...,
+        title="Model Name",
+        description="The name of the OpenAI 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 OpenAIConfigSetup(BaseModel):
+    """Setup model defining module configuration parameters."""
+
+    rag_files: list[bytes] = Field(
+        ...,
+        title="RAG Files",
+        description="Files used for retrieval-augmented generation (RAG) with the OpenAI module.",
+    )
 
 
 class OpenAIToolSecret(BaseModel):
@@ -54,47 +176,80 @@ client_config = ClientConfig(
 )
 
 
-class OpenAIToolModule(BaseModule[OpenAIToolInput, OpenAIToolOutput, OpenAIToolSetup, OpenAIToolSecret]):
+class OpenAIToolModule(
+    BaseModule[
+        OpenAIInput,
+        OpenAIOutput,
+        OpenAISetup,
+        OpenAIToolSecret,
+        OpenAIConfigSetup,
+    ]
+):
     """A openAI endpoint tool module module."""
 
     name = "OpenAIToolModule"
     description = "A module that interacts with OpenAI API to process text"
 
     # Define the schema formats for the module
-    input_format = OpenAIToolInput
-    output_format = OpenAIToolOutput
-    setup_format = OpenAIToolSetup
+    config_setup_format = OpenAIConfigSetup
+    input_format = OpenAIInput
+    output_format = OpenAIOutput
+    setup_format = OpenAISetup
     secret_format = OpenAIToolSecret
 
     openai_client: openai.OpenAI
 
     # Define module metadata for discovery
     metadata: ClassVar[dict[str, Any]] = {
-        "name": "Minimal_LLM_Tool",
+        "name": "OpenAIToolModule",
         "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 = {}
-    services_config_params = {
+    services_config_strategies: ClassVar[dict[str, ServicesStrategy | None]] = {}
+    services_config_params: ClassVar[dict[str, dict[str, Any | None] | None]] = {
         "storage": {
-            "config": {"setups": OpenAIToolSetup},
+            "config": {"setups": OpenAISetup},
             "client_config": client_config,
         },
         "filesystem": {
             "config": {},
             "client_config": client_config,
         },
+        "cost": {
+            "config": {},
+            "client_config": client_config,
+        },
     }
 
-    async def initialize(self, setup_data: SetupData) -> None:
+    async def run_config_setup(
+        self,
+        config_setup_data: OpenAIConfigSetup,
+        setup_data: OpenAISetup,
+        callback: Callable,
+    ) -> None:
+        """Configure the module with additional setup data.
+
+        Args:
+            config_setup_data: Additional configuration content.
+            setup_data: Initial setup data for the module.
+            callback: Function to send output data back to the client.
+        """
+        logger.info("Configuring OpenAIToolModule with additional setup data. %s", config_setup_data)
+
+        # Here you can process config_content and update setup_data as needed
+        # For now, we just return the original setup_data
+        setup_data.developer_prompt = "| + |".join(f.decode("utf-8") for f in config_setup_data.rag_files)
+        await callback(setup_data)
+
+    async def initialize(self, setup_data: OpenAISetup) -> 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.
         """
-        self.openai_client = openai.OpenAI(api_key=setup_data.current_setup_version.content["openai_key"])
+        self.client: openai.AsyncOpenAI = openai.AsyncOpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
         # Define what capabilities this module provides
         self.capabilities = ["text-processing", "streaming", "transformation"]
         logger.info(
@@ -103,10 +258,10 @@ class OpenAIToolModule(BaseModule[OpenAIToolInput, OpenAIToolOutput, OpenAIToolS
             self.capabilities,
         )
 
-    async def run(
+    async def run_message(
         self,
-        input_data: dict[str, Any],
-        setup_data: SetupData,
+        input_model: MessageInputPayload,
+        setup_model: OpenAISetup,
         callback: Callable,
     ) -> None:
         """Process input text and stream LLM responses.
@@ -122,37 +277,64 @@ class OpenAIToolModule(BaseModule[OpenAIToolInput, OpenAIToolOutput, OpenAIToolS
             openai.APIConnectionError: If an API connection error occurs.
             Exception: For any unexpected runtime errors.
         """
-        logger.info(
-            "Running job %s with prompt: '%s' on model: %s",
-            self.job_id,
-            input_data["prompt"],
-            setup_data.current_setup_version.content["model_name"],
+        # response = await self.client.responses.create(
+        #     model=setup_model.model_name,
+        #     instructions=setup_model.developer_prompt,
+        #     temperature=setup_model.temperature,
+        #     max_output_tokens=setup_model.max_tokens,
+        #     input=input_model.user_prompt,
+        # )
+        # logger.info("Recieved answer from OpenAI: %s", response)
+
+        # Get and save the output data
+        message_output_payload = MessageOutputPayload(
+            payload_type="message",
+            user_response="Mock data",
+            # user_response=response.output_text,
         )
-        try:
-            response = self.openai_client.responses.create(
-                model=setup_data.current_setup_version.content["model_name"],
-                tools=[{"type": "web_search_preview"}],
-                instructions=setup_data.current_setup_version.content["dev_prompt"],
-                input=input_data["prompt"],
-            )
-            if not response.output_text:
-                raise openai.APIConnectionError
-            output_data = OpenAIToolOutput(response=response.output_text).model_dump()
-
-        except openai.AuthenticationError as _:
-            message = "Authentication Error, OPENAI auth token was never set."
-            logger.exception(message)
-            output_data = {
-                "error": {
-                    "code": grpc.StatusCode.UNAUTHENTICATED,
-                    "error_message": message,
-                }
-            }
-        except openai.APIConnectionError as _:
-            message = "API Error, please try again."
-            logger.exception(message)
-            output_data = {"error": {"code": grpc.StatusCode.UNAVAILABLE, "error_message": message}}
-        await callback(job_id=self.job_id, output_data=output_data)
+        output_model = self.output_format.model_validate({"payload": message_output_payload})
+        await callback(output_data=output_model)
+
+    async def run(
+        self,
+        input_data: OpenAIInput,
+        setup_data: OpenAISetup,
+        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)
+        setup_model = 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)
+
+        if input_model.payload.payload_type == "message":
+            # Validate against MessageInputPayload
+            message_payload = MessageInputPayload.model_validate(input_model.payload)
+            await self.run_message(message_payload, setup_model, callback)
+        elif input_model.payload.payload_type == "file":
+            # Validate against FileInputPayload
+            file_payload = FileInputPayload.model_validate(input_model.payload)
+            await self.run_file(file_payload, setup_model, callback)
+        else:
+            error_msg = f"Unknown input type '{input_model.payload.payload_type}'. Expected 'message' or 'file'."
+            raise ValueError(error_msg)
         logger.info("Job %s completed", self.job_id)
 
     async def cleanup(self) -> None:

modules/storage_module.py

@@ -3,7 +3,7 @@
 import asyncio
 import datetime
 from collections.abc import Callable
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from pydantic import BaseModel, Field
 
@@ -12,7 +12,9 @@ from digitalkin.models.module import ModuleStatus
 from digitalkin.modules.archetype_module import ArchetypeModule
 from digitalkin.services.services_config import ServicesConfig
 from digitalkin.services.services_models import ServicesMode
-from digitalkin.services.storage.storage_strategy import StorageRecord
+
+if TYPE_CHECKING:
+    from digitalkin.services.storage.storage_strategy import StorageRecord
 
 
 class ExampleInput(BaseModel):
@@ -120,10 +122,7 @@ class ExampleModule(ArchetypeModule[ExampleInput, ExampleOutput, ExampleSetup, E
 
         # Store the output data in storage
         storage_id = self.storage.store(
-            collection="example",
-            record_id=f"example_outputs",
-            data=output_data.model_dump(),
-            data_type="OUTPUT"
+            collection="example", record_id="example_outputs", data=output_data.model_dump(), data_type="OUTPUT"
         )
 
         logger.info("Stored output data with ID: %s", storage_id)