SIENTIAPDE-994

Refactor activity methods and update requirements.txt to enhance functionality and remove deprecated filters. Added detailed docstrings for clarity and improved error handling in data processing workflows.

laborious/workflows/sub_workflows/format_and_export_prediction.py

@@ -11,7 +11,7 @@ class FormatAndExportPrediction():
     async def run(self, input_data: dict[str, Any]):
         path_flag = input_data['path_flag']
         data = input_data['data']
-        confidence = input_data['confidence']
+        prediction_confidence = input_data['prediction_confidence']
 
         if path_flag is None:
             # proceed with formatting and exporting
@@ -21,7 +21,7 @@ class FormatAndExportPrediction():
                     'data': data,
                     'timestamp': input_data['timestamp'],
                     'model_id': input_data['model_id'],
-                    'prediction_confidence': confidence,
+                    'prediction_confidence': prediction_confidence,
                 }
             )
 
@@ -32,7 +32,7 @@ class FormatAndExportPrediction():
                 {
                     'timestamp': input_data['timestamp'],
                     'model_id': input_data['model_id'],
-                    'prediction_confidence': confidence,
+                    'prediction_confidence': prediction_confidence,
                     'comment': input_data['comment']
                 }
             )

laborious/workflows/sub_workflows/prediction_process.py

@@ -13,8 +13,14 @@ class PredictionProcess():
     @workflow.run
     async def run(self, input_data: dict[str, Any]):
         data = input_data['data']
+        schema = input_data['schema']
+        table_name = input_data['table_name']
+        model = input_data['model']
+        filters = input_data['filters']
+        model_name = input_data['model_name']
+        model_retention = input_data['model_retention']
 
-        path_flag, _confidence = await workflow.execute_activity_method(
+        path_flag, confidence = await workflow.execute_activity_method(
             Gates.input_gate,
             {
                 'filters': input_data['filters'],
@@ -22,38 +28,102 @@ class PredictionProcess():
             }
         )
 
-        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']
-                }
-            )
+        if await self.path_flag_handler(
+            data, path_flag, confidence, schema, table_name, 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']
+                'model_name': model_name,
+                'model_retention': model_retention
             }
         )
 
         path_flag, confidence = await workflow.execute_activity_method(
             Gates.mlflow_gate,
             {
-                'filters': input_data['filters'],
+                'filters': filters,
                 'data': response_data,
                 'type': 'transform'
             }
         )
 
-        if path_flag == 'stop':
+        if await self.path_flag_handler(
+            data, path_flag, confidence, schema, table_name, model
+        ):
             return
+
+        response_data = await workflow.execute_activity_method(
+            MLFlow.request_predict,
+            {
+                'data': response_data,
+                'model_name': model_name,
+                'model_retention': model_retention
+            }
+        )
+
+        path_flag, confidence = await workflow.execute_activity_method(
+            Gates.mlflow_gate,
+            {
+                'filters': filters,
+                'data': response_data,
+                'type': 'predict'
+            }
+        )
+
+        if await self.path_flag_handler(
+            data, path_flag, confidence, schema, table_name, model
+        ):
+            return
+
+        await workflow.execute_child_workflow(
+            'format_and_export_prediction',
+            {
+                'path_flag': path_flag,
+                'data': response_data['content'],
+                'prediction_confidence': confidence,
+                'timestamp': response_data['timestamp'],
+                'model_id': model,
+                'model_name': model_name,
+                'model_retention': model_retention
+            }
+        )
+
+    async def path_flag_handler(self, data: dict[str, Any], path_flag: str,
+                                confidence: int, schema: str, table_name: str,
+                                model: str):
+        if path_flag == 'stop':
+            return True
+
+        elif path_flag == 'repeat':
+            # repeat last prediction
+            await workflow.execute_activity_method(
+                Postgres.repeat_last_prediction,
+                {
+                    'schema': schema,
+                    'table_name': table_name,
+                    'model': model
+                }
+            )
+            return True
+
+        elif path_flag == 'continue':
+            # call write workflow
+            workflow.execute_child_workflow(
+                'format_and_export_prediction',
+                {
+                    'path_flag': path_flag,
+                    'data': data,
+                    'prediction_confidence': confidence,
+                    'timestamp': last_timestamp,
+                    'model_id': model,
+                    'model_name': model_name,
+                    'model_retention': model_retention
+                }
+            )
+            return True
+
+        return False