digitalkin 0.2.11__py3-none-any.whl → 0.2.13__py3-none-any.whl

modules/minimal_llm_module.py

@@ -1,44 +1,157 @@
 """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.services_models import ServicesStrategy
 from digitalkin.services.setup.setup_strategy import SetupData
 
 # 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 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 OpenAIToolOutput(BaseModel):
+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 OpenAIToolSecret(BaseModel):
@@ -54,38 +167,42 @@ client_config = ClientConfig(
 )
 
 
-class OpenAIToolModule(BaseModule[OpenAIToolInput, OpenAIToolOutput, OpenAIToolSetup, OpenAIToolSecret]):
+class OpenAIToolModule(BaseModule[OpenAIInput, OpenAIOutput, OpenAISetup, OpenAIToolSecret]):
     """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
+    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:
@@ -94,7 +211,7 @@ class OpenAIToolModule(BaseModule[OpenAIToolInput, OpenAIToolOutput, OpenAIToolS
         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 +220,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 +239,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)

modules/text_transform_module.py

@@ -6,7 +6,7 @@ from typing import Any, ClassVar
 
 from pydantic import BaseModel
 
-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.storage.storage_strategy import DataType, StorageRecord

digitalkin/modules/job_manager.py

@@ -1,176 +0,0 @@
-"""Background module manager."""
-
-import asyncio
-import uuid
-from argparse import ArgumentParser, Namespace
-from collections.abc import Callable, Coroutine
-from typing import Any
-
-from digitalkin.logger import logger
-from digitalkin.models import ModuleStatus
-from digitalkin.models.module import InputModelT, OutputModelT, SecretModelT, SetupModelT
-from digitalkin.modules._base_module import BaseModule
-from digitalkin.services.services_config import ServicesConfig
-from digitalkin.services.services_models import ServicesMode
-from digitalkin.utils.arg_parser import ArgParser, DevelopmentModeMappingAction
-
-
-class JobManager(ArgParser):
-    """Background module manager."""
-
-    args: Namespace
-
-    @staticmethod
-    async def _job_specific_callback(
-        callback: Callable[[str, OutputModelT], Coroutine[Any, Any, None]], job_id: str
-    ) -> Callable[[OutputModelT], Coroutine[Any, Any, None]]:
-        """Return a callback function for the job.
-
-        Args:
-            callback: Callback function to be called when the job is done
-            job_id: Identifiant du module
-
-        Returns:
-            Callable: Callback function
-        """
-
-        def callback_wrapper(output_data: OutputModelT) -> Coroutine[Any, Any, None]:
-            """Wrapper for the callback function.
-
-            Args:
-                output_data: Output data of the job
-
-            Returns:
-                Coroutine: Callback function
-            """
-            return callback(job_id, output_data)
-
-        return callback_wrapper
-
-    def _add_parser_args(self, parser: ArgumentParser) -> None:
-        super()._add_parser_args(parser)
-        parser.add_argument(
-            "-d",
-            "--dev-mode",
-            env_var="SERVICE_MODE",
-            choices=ServicesMode.__members__,
-            default="local",
-            action=DevelopmentModeMappingAction,
-            dest="services_mode",
-            help="Define Module Service configurations for endpoints",
-        )
-
-    def __init__(self, module_class: type[BaseModule]) -> None:
-        """Initialize the job manager."""
-        self.module_class = module_class
-        self.modules: dict[str, BaseModule] = {}
-        self._lock = asyncio.Lock()
-        super().__init__()
-
-        services_config = ServicesConfig(
-            services_config_strategies=self.module_class.services_config_strategies,
-            services_config_params=self.module_class.services_config_params,
-            mode=self.args.services_mode,
-        )
-        setattr(self.module_class, "services_config", services_config)
-
-    async def create_job(  # noqa: D417
-        self,
-        input_data: InputModelT,
-        setup_data: SetupModelT,
-        mission_id: str,
-        callback: Callable[[str, OutputModelT], Coroutine[Any, Any, None]],
-    ) -> tuple[str, BaseModule[InputModelT, OutputModelT, SetupModelT, SecretModelT]]:  # type: ignore
-        """Start new module job in background (asyncio).
-
-        Args:
-            module_class: Classe du module à instancier
-            *args: Arguments à passer au constructeur du module
-            **kwargs: Arguments à passer au constructeur du module
-
-        Returns:
-            str: job_id of the module entity
-        """
-        job_id = str(uuid.uuid4())
-        """TODO: check uniqueness of the job_id"""
-        # Création et démarrage du module
-        module = self.module_class(job_id, mission_id=mission_id)
-        self.modules[job_id] = module
-        try:
-            await module.start(input_data, setup_data, await JobManager._job_specific_callback(callback, job_id))
-            logger.info("Module %s (%s) started successfully", job_id, module.name)
-        except Exception:
-            # En cas d'erreur, supprimer le module du gestionnaire
-            del self.modules[job_id]
-            logger.exception("Échec du démarrage du module %s: %s", job_id)
-            raise
-        else:
-            return job_id, module
-
-    async def stop_module(self, job_id: str) -> bool:
-        """Arrête un module en cours d'exécution.
-
-        Args:
-            job_id: Identifiant du module à arrêter
-
-        Returns:
-            True si le module a été arrêté, False s'il n'existe pas.
-        """
-        async with self._lock:
-            module = self.modules.get(job_id)
-            if not module:
-                logger.warning(f"Module {job_id} introuvable")
-                return False
-            try:
-                await module.stop()
-                logger.info(f"Module {job_id} ({module.name}) arrêté avec succès")
-            except Exception as e:
-                logger.error(f"Erreur lors de l'arrêt du module {job_id}: {e}")
-                raise
-            else:
-                return True
-
-    def get_module_status(self, job_id: str) -> ModuleStatus | None:
-        """Obtient le statut d'un module.
-
-        Args:
-            job_id: Identifiant du module
-
-        Returns:
-            Le statut du module ou None si le module n'existe pas.
-        """
-        module = self.modules.get(job_id)
-        return module.status if module else None
-
-    def get_module(self, job_id: str) -> BaseModule | None:
-        """Récupère une référence au module.
-
-        Args:
-            job_id: Identifiant du module
-
-        Returns:
-            Le module ou None s'il n'existe pas.
-        """
-        return self.modules.get(job_id)
-
-    async def stop_all_modules(self) -> None:
-        """Arrête tous les modules en cours d'exécution."""
-        async with self._lock:
-            stop_tasks = [self.stop_module(job_id) for job_id in list(self.modules.keys())]
-            if stop_tasks:
-                await asyncio.gather(*stop_tasks, return_exceptions=True)
-
-    def list_modules(self) -> dict[str, dict[str, Any]]:
-        """Liste tous les modules avec leur statut.
-
-        Returns:
-            Dictionnaire des modules avec leurs informations.
-        """
-        return {
-            job_id: {
-                "name": module.name,
-                "status": module.status,
-                "class": module.__class__.__name__,
-            }
-            for job_id, module in self.modules.items()
-        }