SIENTIAPDE-994

Implement new get_last_timestamp method in Gates class, refactor MLFlow activity methods to return only transformed data, and update PredictionsBatch and PredictionProcess workflows to utilize Activities module. Add detailed docstrings for new methods and enhance test coverage for get_last_timestamp functionality.

laborious/workflows/predictions_batch.py

@@ -1,10 +1,7 @@
 from temporalio import workflow
 
 with workflow.unsafe.imports_passed_through():
-    from laborious.activities.postgres import Postgres
-    from laborious.activities.mlflow import MLFlow
-    from laborious.activities.gates import Gates
-    from laborious.activities.opc import OPC
+    from laborious.activities.activities import Activities
     from typing import Any
 
 
@@ -14,7 +11,7 @@ class PredictionsBatch():
     async def run(self, input_data: dict[str, Any]):
 
         await workflow.execute_activity_method(
-            Postgres.prepare_activity,
+            Activities.prepare_activity,
             {
                 'schedule_name': input_data['schedule_name'],
                 'model_name': input_data['model_name'],
@@ -23,63 +20,11 @@ class PredictionsBatch():
         )
 
         data = await workflow.execute_activity_method(
-            Postgres.load_custom_query,
+            Activities.load_custom_query,
             input_data['query']
         )
 
-        path_flag, confidence = await workflow.execute_activity_method(
-            Gates.input_gate,
-            {
-                'filters': input_data['filters'],
-                'data': data
-            }
-        )
+        input_data['data'] = data
 
-        if path_flag == 'stop':
-            return
-
-        if path_flag == 'continue':
-            # repeat last prediction
-            await workflow.execute_activity_method(
-                Postgres.repeat_last_prediction,
-                {
-                    'schema': input_data['schema'],
-                    'table_name': input_data['table_name'],
-                    'model': input_data['model']
-                }
-            )
-            return
-
-        response_data, last_timestamp = await workflow.execute_activity_method(
-            MLFlow.transform_data,
-            {
-                'data': data,
-                'model_name': input_data['model_name'],
-                'model_retention': input_data['model_retention']
-            }
-        )
-
-        path_flag, confidence = await workflow.execute_activity_method(
-            Gates.mlflow_gate,
-            {
-                'filters': input_data['filters'],
-                'data': response_data,
-                'type': 'transform'
-            }
-        )
-
-        if path_flag == 'stop':
-            return
-
-        if path_flag is None:
-            # procced with prediction
-            response_data = await workflow.execute_activity_method(
-                MLFlow.request_predict,
-                {
-                    'data': response_data,
-                    'model_name': input_data['model_name'],
-                    'model_retention': input_data['model_retention']
-                }
-            )
-
-            path_flag
+        await workflow.execute_child_workflow(
+            'prediction_process', input_data)

laborious/workflows/sub_workflows/format_and_export_prediction.py

@@ -9,6 +9,29 @@ with workflow.unsafe.imports_passed_through():
 class FormatAndExportPrediction():
     @workflow.run
     async def run(self, input_data: dict[str, Any]):
+        """
+        This workflow formats and exports predictions based on path_flag:
+        - If path_flag is None: formats prediction using input data, timestamp, model_id and confidence
+        - If path_flag exists: creates default prediction with timestamp, model_id, confidence and comment
+        Finally exports formatted prediction to postgres table
+        Args:
+            input_data(dict[str, Any]): The input data for the workflow. Contains the following keys:
+                - path_flag(str): The path flag to determine the type of prediction to format
+                - data(dict[str, Any]): The data to format
+                - prediction_confidence(float): The prediction confidence to be registered
+                - timestamp(str): The timestamp of the prediction, synchronized with the data
+                - model_id(str): The model id of the prediction
+                - model_name(str): The model name of the prediction
+                - model_retention(str): The model retention of the prediction
+                - comment(str): The comment to be registered
+                - schema(str): The schema of the prediction
+                - table_name(str): The table name of the prediction
+                - opc_servers(list[str]): The opc servers of the prediction
+                - opc_output_config(dict[str, Any]): The opc output config of the prediction
+
+        Returns:
+            bool: True if the workflow was successful, False otherwise.
+        """
         path_flag = input_data['path_flag']
         data = input_data['data']
         prediction_confidence = input_data['prediction_confidence']

laborious/workflows/sub_workflows/prediction_process.py

@@ -1,10 +1,7 @@
 from temporalio import workflow
 
 with workflow.unsafe.imports_passed_through():
-    from laborious.activities.postgres import Postgres
-    from laborious.activities.mlflow import MLFlow
-    from laborious.activities.gates import Gates
-    from laborious.activities.opc import OPC
+    from laborious.activities.activities import Activities
     from typing import Any
 
 
@@ -20,8 +17,15 @@ class PredictionProcess():
         model_name = input_data['model_name']
         model_retention = input_data['model_retention']
 
+        last_timestamp = await workflow.execute_activity_method(
+            Activities.get_last_timestamp,
+            {
+                'data': data
+            }
+        )
+
         path_flag, confidence = await workflow.execute_activity_method(
-            Gates.input_gate,
+            Activities.input_gate,
             {
                 'filters': input_data['filters'],
                 'data': data
@@ -29,12 +33,13 @@ class PredictionProcess():
         )
 
         if await self.path_flag_handler(
-            data, path_flag, confidence, schema, table_name, model
+            data, path_flag, confidence, schema, table_name,
+            model, last_timestamp, model_name, model_retention
         ):
             return
 
-        response_data, last_timestamp = await workflow.execute_activity_method(
-            MLFlow.transform_data,
+        response_data = await workflow.execute_activity_method(
+            Activities.request_transform,
             {
                 'data': data,
                 'model_name': model_name,
@@ -43,7 +48,7 @@ class PredictionProcess():
         )
 
         path_flag, confidence = await workflow.execute_activity_method(
-            Gates.mlflow_gate,
+            Activities.mlflow_response_gate,
             {
                 'filters': filters,
                 'data': response_data,
@@ -52,12 +57,28 @@ class PredictionProcess():
         )
 
         if await self.path_flag_handler(
-            data, path_flag, confidence, schema, table_name, model
+            data, path_flag, confidence, schema, table_name,
+            model, last_timestamp, model_name, model_retention
+        ):
+            return
+
+        path_flag, confidence = await workflow.execute_activity_method(
+            Activities.mlflow_content_gate,
+            {
+                'filters': filters,
+                'data': response_data,
+                'type': 'transform'
+            }
+        )
+
+        if await self.path_flag_handler(
+            data, path_flag, confidence, schema, table_name,
+            model, last_timestamp, model_name, model_retention
         ):
             return
 
         response_data = await workflow.execute_activity_method(
-            MLFlow.request_predict,
+            Activities.request_predict,
             {
                 'data': response_data,
                 'model_name': model_name,
@@ -66,7 +87,7 @@ class PredictionProcess():
         )
 
         path_flag, confidence = await workflow.execute_activity_method(
-            Gates.mlflow_gate,
+            Activities.mlflow_response_gate,
             {
                 'filters': filters,
                 'data': response_data,
@@ -75,7 +96,8 @@ class PredictionProcess():
         )
 
         if await self.path_flag_handler(
-            data, path_flag, confidence, schema, table_name, model
+            data, path_flag, confidence, schema, table_name,
+            model, last_timestamp, model_name, model_retention
         ):
             return
 
@@ -94,14 +116,32 @@ class PredictionProcess():
 
     async def path_flag_handler(self, data: dict[str, Any], path_flag: str,
                                 confidence: int, schema: str, table_name: str,
-                                model: str):
+                                model: str, last_timestamp: str, model_name: str,
+                                model_retention: str):
+        """
+        This function handles the path flag and the confidence of the prediction.
+        It returns True if the prediction should be stopped. If path_flag is 'repeat', it repeats the last prediction.
+        If path_flag is 'continue', it calls the write workflow. If path_flag is 'stop', it stops the prediction process.
+        Args:
+            data (dict[str, Any]): The data to be used for the prediction.
+            path_flag (str): The path flag to determine the type of prediction to format
+            confidence (int): The confidence of the prediction
+            schema (str): The schema of the prediction
+            table_name (str): The table name of the prediction
+            model (str): The model id of the prediction
+            last_timestamp (str): The timestamp of the last prediction
+            model_name (str): The model name of the prediction
+            model_retention (str): The model retention of the prediction
+        Returns:
+            bool: True if the prediction should be stopped, False otherwise.
+        """
         if path_flag == 'stop':
             return True
 
         elif path_flag == 'repeat':
             # repeat last prediction
             await workflow.execute_activity_method(
-                Postgres.repeat_last_prediction,
+                Activities.repeat_last_prediction,
                 {
                     'schema': schema,
                     'table_name': table_name,