digitalkin 0.3.2.dev2__py3-none-any.whl

modules/minimal_llm_module.py

@@ -0,0 +1,347 @@
+"""Simple module calling an LLM."""
+
+import logging
+import os
+from collections.abc import Callable
+from typing import Any, ClassVar, Literal
+
+import openai
+from digitalkin.grpc_servers.utils.models import ClientConfig, SecurityMode, ServerMode
+from pydantic import BaseModel, Field
+
+from digitalkin.modules._base_module import BaseModule
+from digitalkin.services.services_models import ServicesStrategy
+
+# 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 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."""
+
+    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 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."""
+
+    payload: MessageOutputPayload | FileOutputPayload = Field(
+        ...,
+        discriminator="payload_type",
+        title="Payload",
+        description="Either a message or file response.",
+    )
+
+
+class OpenAISetup(BaseModel):
+    """Setup model defining module configuration parameters."""
+
+    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):
+    """Secret model defining module configuration parameters."""
+
+
+client_config = ClientConfig(
+    host="[::]",
+    port=50151,
+    mode=ServerMode.ASYNC,
+    security=SecurityMode.INSECURE,
+    credentials=None,
+)
+
+
+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
+    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": "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: ClassVar[dict[str, ServicesStrategy | None]] = {}
+    services_config_params: ClassVar[dict[str, dict[str, Any | None] | None]] = {
+        "storage": {
+            "config": {"setups": OpenAISetup},
+            "client_config": client_config,
+        },
+        "filesystem": {
+            "config": {},
+            "client_config": client_config,
+        },
+        "cost": {
+            "config": {},
+            "client_config": client_config,
+        },
+    }
+
+    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.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(
+            "Module %s initialized with capabilities: %s",
+            self.metadata["name"],
+            self.capabilities,
+        )
+
+    async def run_message(
+        self,
+        input_model: MessageInputPayload,
+        setup_model: OpenAISetup,
+        callback: Callable,
+    ) -> None:
+        """Process input text and stream LLM responses.
+
+        Args:
+            input_data: Contains the text to process.
+            setup_data: Contains model configuration and development prompt.
+            callback: Function to send output data back to the client.
+
+        Raises:
+            grpc.RpcError: If gRPC communication fails.
+            openai.AuthenticationError: If authentication with OpenAI fails.
+            openai.APIConnectionError: If an API connection error occurs.
+            Exception: For any unexpected runtime errors.
+        """
+        # 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,
+        )
+        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:
+        """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/text_transform_module.py

@@ -0,0 +1,203 @@
+"""Simple module example transforming a text."""
+
+import logging
+from collections.abc import Callable
+from typing import Any, ClassVar
+
+from digitalkin.grpc_servers.utils.models import ClientConfig, SecurityMode, ServerMode
+from pydantic import BaseModel
+
+from digitalkin.modules._base_module import BaseModule
+from digitalkin.services.setup.setup_strategy import SetupData
+from digitalkin.services.storage.storage_strategy import DataType, StorageRecord
+
+# Configure logging with clear formatting
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+)
+logger = logging.getLogger(__name__)
+
+
+# Define schema models using Pydantic
+class TextTransformInput(BaseModel):
+    """Input model defining what data the module expects."""
+
+    text: str
+    transform_count: int = 1  # Default to 1 transformation
+
+
+class TextTransformOutput(BaseModel):
+    """Output model defining what data the module produces."""
+
+    transformed_text: str
+    iteration: int  # Tracks which transformation this is
+
+
+class TextTransformSetup(BaseModel):
+    """Setup model defining module configuration parameters."""
+
+    shift_amount: int = 1  # Default Caesar shift by 1
+    uppercase: bool = False  # Whether to convert to uppercase
+
+
+class TextTransformSecret(BaseModel):
+    """Secret model defining module configuration parameters."""
+
+
+class TextTransformStorage(BaseModel):
+    """Secret model defining module configuration parameters."""
+
+    module: str = "Text_Transform_Module"
+    user: str = "user"
+    consumption: int = 0
+    ended: bool = False
+
+
+client_config = ClientConfig(
+    host="[::]",
+    port=50151,
+    mode=ServerMode.ASYNC,
+    security=SecurityMode.INSECURE,
+    credentials=None,
+)
+
+
+class TextTransformModule(BaseModule[TextTransformInput, TextTransformOutput, TextTransformSetup, TextTransformSecret]):
+    """A text transformation module that demonstrates streaming capabilities.
+
+    This module takes text input and performs multiple transformations on it,
+    sending back each transformation as a separate output message.
+    """
+
+    # Define the schema formats for the module
+    name = "Text_Transform_Module"
+    input_format = TextTransformInput
+    output_format = TextTransformOutput
+    setup_format = TextTransformSetup
+    secret_format = TextTransformSecret
+
+    # Define module metadata for discovery
+    metadata: ClassVar[dict[str, Any]] = {
+        "name": "Text_Transform_Module",
+        "description": "Transforms input text using Caesar cipher with streaming output",
+        "version": "1.0.0",
+        "tags": ["text", "transformation", "encryption", "streaming"],
+    }
+
+    # Define services_config_params with default values
+    services_config_strategies = {}
+    services_config_params = {
+        "storage": {
+            "config": {"monitor": TextTransformStorage, "setups": TextTransformStorage},
+            "client_config": client_config,
+        },
+        "filesystem": {
+            "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.
+
+        Raises:
+            Exception: If initialization fails.
+        """
+        # Define what capabilities this module provides
+        self.capabilities = ["text-processing", "streaming", "transformation"]
+        logger.info(
+            "Module %s initialized with capabilities: %s",
+            self.metadata["name"],
+            self.capabilities,
+        )
+
+        self.db_id = self.storage.store(
+            "monitor",
+            {
+                "module": self.metadata["name"],
+                "user": f"xxxx+{self.job_id}",
+                "consumption": 0,
+                "ended": False,
+            },
+            data_type=DataType.VIEW,
+        )
+
+    async def run(
+        self,
+        input_data: dict[str, Any],
+        setup_data: SetupData,
+        callback: Callable,
+    ) -> None:
+        """Process input text and stream transformation results.
+
+        This method implements a Caesar cipher transformation on input text.
+        It demonstrates streaming capability by sending multiple outputs through
+        the callback for each transformation iteration.
+
+        Args:
+            input_data: Contains the text to transform and number of iterations.
+            setup_data: Contains shift amount and uppercase flags.
+            callback: Function to send output data back to the client.
+        """
+        text = input_data["text"]
+        transform_count = int(input_data["transform_count"])
+        logger.info("%s  | %s", setup_data, type(setup_data))
+        shift_amount = int(setup_data.current_setup_version.content["shift_amount"])
+        uppercase = setup_data.current_setup_version.content["uppercase"]
+
+        logger.info(
+            "Running job %s with text: '%s', iterations: %s",
+            self.job_id,
+            text,
+            transform_count,
+        )
+
+        # Process the text for each iteration
+        for i in range(transform_count):
+            # Apply Caesar cipher (shift each character by specified amount)
+            transformed = "".join([chr(ord(char) + shift_amount) if char.isalpha() else char for char in text])
+
+            # Apply uppercase transformation if configured
+            if uppercase:
+                transformed = transformed.upper()
+
+            output_data = TextTransformOutput(transformed_text=transformed, iteration=i + 1)
+
+            logger.info(
+                "Sending transformation %s/%s: '%s'",
+                i + 1,
+                transform_count,
+                transformed,
+            )
+
+            monitor_obj: StorageRecord | None = self.storage.read("monitor")
+            if monitor_obj is None:
+                logger.error("Monitor object not found in storage.")
+                break
+            monitor_obj.data.consumption += 1
+            updated_monitor_obj: StorageRecord | None = self.storage.modify("monitor", monitor_obj.data.model_dump())
+            self.db_id = updated_monitor_obj.name if updated_monitor_obj else "monitor"
+
+            # Send results through callback and wait for acknowledgment
+            await callback(job_id=self.job_id, output_data=output_data.model_dump())
+            text = transformed
+
+        logger.info("Job %s completed with %s transformations", self.job_id, transform_count)
+
+    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(f"Cleaning up module {self.metadata['name']}")
+        monitor_obj = self.storage.read("monitor")
+        if monitor_obj is None:
+            logger.error("Monitor object not found in storage.")
+            return
+        monitor_obj.data.ended = True
+        updated_monitor_obj: StorageRecord | None = self.storage.modify("monitor", monitor_obj.data.model_dump())
+        self.db_id = updated_monitor_obj.name if updated_monitor_obj else "monitor"