SIENTIAPDE-994

Update requirements.txt with new dependencies and refactor activity methods for improved functionality and error handling

laborious/activities/gates.py

@@ -1,19 +1,31 @@
 from temporalio import activity, workflow
 
+
 with workflow.unsafe.imports_passed_through():
+    import traceback
     from logging import Logger
     from sientia_do.notifications.handlers import NotificationHandler
     from laborious.activities.base import BaseActivity
     from typing import Any
     from laborious.utils.filters.conditional_filters import filter_empty_data, filter_specific_variables_null_values
     from pandas import DataFrame
+    from sientia_do.notifications.models import NotificationLevel
+    from laborious.utils.filters.mlflow_filters import nan_values_filter, api_error_filter
 
-
-filter_functions = {
+input_filter_functions = {
     'SPECIFIC_VARIABLES_NULL_VALUES': filter_specific_variables_null_values,
     'EMPTY_DATA': filter_empty_data
 }
 
+transform_filter_functions = {
+    'response_filter': {
+        'API_ERROR': api_error_filter,
+    },
+    'content_filter': {
+        'NAN_VALUES': nan_values_filter,
+    }
+}
+
 
 class Gates(BaseActivity):
     def __init__(self, logger: Logger, notification_handler: NotificationHandler):
@@ -25,7 +37,9 @@ class Gates(BaseActivity):
         Filters the data based on the filters. The return value is a tuple with the first element
         being the policy and the second element being the confidence status.
         Args:
-            input_data (dict): The input data.
+            input_data (dict): The input data. Contains:
+                filters (dict): The filters to apply.
+                data (dict[str, Any]): The data to filter.
         Returns:
             tuple[str, int]: ('stop', -1) if some filter policy is 'stop', ('continue', 2)
             if no filter policy is 'stop' and some filter policy is 'continue',
@@ -36,8 +50,18 @@ class Gates(BaseActivity):
 
         filter_output = []
         for fil, config in filters.items():
-            if filter_functions[fil](data, config):
-                filter_output.append(config['POLICY'])
+            try:
+                if input_filter_functions[fil](data, config):
+                    filter_output.append(config['POLICY'])
+            except Exception as e:
+                trace = traceback.format_exc()
+                self.notification_handler.build_and_send_notification(
+                    notification_id=f"INTPUT_GATE_ERROR__{fil}",
+                    message=f"Error in filter {fil}:{config}: \n {e}",
+                    block="input_gate",
+                    level=NotificationLevel.ERROR,
+                    attachment_content=trace
+                )
 
         if 'stop' in filter_output:
             return 'stop', -1
@@ -45,3 +69,75 @@ class Gates(BaseActivity):
             return 'continue', 2
 
         return None, 0
+
+    @activity.defn(name="mlflow_gate")
+    async def mlflow_gate(self, input_data: dict[str, Any]) -> tuple[str, int]:
+
+        filters = input_data['filters']
+        data = DataFrame(input_data['data'])
+        gate_type = input_data['type']
+
+        filter_output = []
+        for fil, config in filters.items():
+            if transform_filter_functions['response_filter'][fil](data, config):
+                filter_output.append(config['POLICY'])
+                self.notification_handler.build_and_send_notification(
+                    notification_id=f"{gate_type.upper()}_GATE_RESPONSE_FILTER__{fil}",
+                    message=data['content']['message'],
+                    block="mlflow_gate",
+                    level=NotificationLevel.WARNING,
+                    attachment_content=data['content']['traceback']
+                )
+
+        if 'stop' in filter_output:
+            return 'stop', -1
+        elif 'continue' in filter_output:
+            return 'continue', 10
+
+        if gate_type == 'predict':
+            return None, 0
+
+        data = DataFrame(data['content'])
+
+        for fil, config in filters.items():
+            if transform_filter_functions['content_filter'][fil](data, config):
+                filter_output.append(config['POLICY'])
+                self.notification_handler.build_and_send_notification(
+                    notification_id=f"{gate_type.upper()}_GATE_CONTENT_FILTER__{fil}",
+                    message=f"Data not passed the content filter {fil}:{config}",
+                    block="mlflow_gate",
+                    level=NotificationLevel.WARNING,
+                    attachment_content=data.to_string()
+                )
+
+        if 'stop' in filter_output:
+            return 'stop', -1
+        elif 'continue' in filter_output:
+            return 'continue', 18
+
+        return None, 0
+
+    @activity.defn(name="format_prediction")
+    async def format_prediction(self, input_data: dict[str, Any]) -> str:
+        data = DataFrame(input_data['data'])
+
+        data['timestamp'] = input_data['timestamp']
+        data['model_id'] = input_data['model_id']
+        data['prediction_confidence'] = input_data['prediction_confidence']
+        data['prediction_status'] = 'Good'
+        data['comment'] = ""
+        data.sort_values(by='timestamp', inplace=True)
+
+        return data.to_dict()
+
+    @activity.defn(name="format_default_prediction")
+    async def format_default_prediction(self, input_data: dict[str, Any]) -> str:
+        return DataFrame({
+            'prediction': [0],
+            'response_time': [0],
+            'timestamp': [input_data['timestamp']],
+            'model_id': [input_data['model_id']],
+            'prediction_confidence': [input_data['prediction_confidence']],
+            'prediction_status': ['Bad'],
+            'comment': [input_data['comment']]
+        }).to_dict()