SIENTIAPDE-1102

Refactor OPC activity to enhance data writing and confidence processing; update tests accordingly

laborious/activities/opc.py

@@ -11,6 +11,8 @@ with workflow.unsafe.imports_passed_through():
     import traceback
     from pandas import DataFrame
 
+OPC_WRITTING_ERROR_CONFIDENCE = 12
+
 
 class OPC(BaseActivity):
     def __init__(self, opc_servers: dict[str, dict[str, Any]],
@@ -38,7 +40,7 @@ class OPC(BaseActivity):
         BaseActivity.__init__(self, logger, notification_handler)
 
     def write_data(self, server: str, tag: str, data: Any,
-                   data_type: str, tag_type: str):
+                   data_type: str, tag_type: str) -> bool:
         """
         Write data to OPC server.
 
@@ -48,10 +50,13 @@ class OPC(BaseActivity):
             - data (Any): The data to write.
             - data_type (str): The data type.
             - tag_type (str): The tag type.
+
+        Returns:
+            - bool: True if the data was written successfully, False otherwise.
         """
 
         try:
-            self.opc_repository[server].write_data(
+            return self.opc_repository[server].write_data(
                 tag, data, data_type)
             self.logger.debug(f"Wrote {tag_type} to {tag}")
         except Exception as e:
@@ -66,7 +71,7 @@ class OPC(BaseActivity):
             raise e
 
     @activity.defn(name='write_opc_data')
-    async def write_opc_data(self, input_data: dict[str, Any]):
+    async def write_opc_data(self, input_data: dict[str, Any]) -> dict[Any, Any]:
         """
         Write prediction and confidence data to OPC servers. The two writing
         operations are optional and independent of each other.
@@ -80,12 +85,17 @@ class OPC(BaseActivity):
                     - prediction_tags(dict[str, Any]): The tags to write to the OPC servers.
                     - confidence_tags(dict[str, Any]): The tags to write to the OPC servers.
 
+        Returns:
+            - dict[Any, Any]: The data that was written to the OPC servers.
+
         """
         self.logger.debug("Writing data to OPC servers...")
         data = DataFrame(input_data['data'])
         opc_output_config = input_data['opc_output_config']
         self.logger.debug(data)
 
+        success = True
+
         for server, config in opc_output_config.items():
             if self.opc_repository.get(server) is None:
                 self.logger.error(f"OPC server {server} not found")
@@ -93,13 +103,14 @@ class OPC(BaseActivity):
 
             if 'prediction_tags' in config:
                 for tag, tag_config in config['prediction_tags'].items():
-                    self.write_data(
+                    success = success and self.write_data(
                         server=server,
                         tag=tag,
                         data=data.head(1)['prediction'].values[0],
                         data_type=tag_config['data_type'],
                         tag_type='prediction'
                     )
+
             if 'confidence_tags' in config:
                 for tag, tag_config in config['confidence_tags'].items():
                     self.write_data(
@@ -110,6 +121,34 @@ class OPC(BaseActivity):
                         tag_type='confidence'
                     )
 
+        return self.process_confidence(data, success)
+
+    def process_confidence(self, data: DataFrame, success: bool) -> dict[Any, Any]:
+        """
+        Processes the confidence of OPC server write operations and updates the DataFrame accordingly.
+
+        If the write operation was not successful, sets the 'prediction_confidence' column in the DataFrame
+        to a predefined error confidence value and logs a debug message. Otherwise, logs a success message.
+
+        Args:
+            data (DataFrame): The DataFrame containing the data to be processed.
+            success (bool): Indicates whether the data was successfully written to the OPC servers.
+
+        Returns:
+            dict[Any, Any]: The processed data as a dictionary.
+        """
+        if not success:
+            data['prediction_confidence'] = OPC_WRITTING_ERROR_CONFIDENCE
+            self.logger.debug(
+                "Some data could not be written to OPC servers, setting confidence to "
+                f"{OPC_WRITTING_ERROR_CONFIDENCE}."
+            )
+
+        else:
+            self.logger.info("Data written to OPC servers successfully.")
+
+        return data.to_dict()
+
     def shutdown(self):
         for opc in self.opc_repository.values():
             opc.disconnect()

laborious/utils/repository/opc_repository.py

@@ -2,9 +2,11 @@ import traceback
 from logging import Logger
 from datetime import datetime
 from pathlib import Path
+from typing import Any
 from asyncua.sync import Client
 from asyncua.crypto.security_policies import SecurityPolicyBasic256
 from asyncua.ua import DataValue, Variant, VariantType
+from regex import F
 from sientia_do.notifications.handlers import NotificationHandler
 from sientia_do.notifications.models import NotificationLevel
 
@@ -190,7 +192,7 @@ class OpcRepository():
 
         return True
 
-    def write_data(self, node, value, data_type):
+    def write_data(self, node: str, value: Any, data_type: str) -> bool:
         """
         Writes data to the OPC server.
         If the connection is not established, it attempts to reconnect.
@@ -202,7 +204,7 @@ class OpcRepository():
         If the client is connected, it returns True.
         """
         if not self.validate_connection():
-            return
+            return False
         try:
             node = self.client.get_node(node)
         except Exception as e:
@@ -216,7 +218,7 @@ class OpcRepository():
             )
             self.logger.error(trace)
             self.error_count += 1
-            return
+            return False
 
         if data_type not in data_type_map:
             self.notification_handler.build_and_send_notification(
@@ -225,7 +227,7 @@ class OpcRepository():
                 block="opc_repository",
                 level=NotificationLevel.ERROR
             )
-            return
+            return False
 
         data = data_type_map[data_type]['converter'](value)
         self.logger.info(f'Writing {data} - {type(data)} to {node}')
@@ -245,5 +247,7 @@ class OpcRepository():
             )
             self.logger.error(trace)
             self.error_count += 1
-            return
+            return False
         self.error_count = 0
+
+        return True

laborious/workflows/sub_workflows/format_and_export_prediction.py

@@ -68,20 +68,8 @@ class FormatAndExportPrediction():
                 start_to_close_timeout=timedelta(seconds=60)
             )
 
-        # write to postgres
-        postgres_holder = workflow.execute_activity_method(
-            Activities.export_data_to_postgres,
-            {
-                'schema': input_data['schema'],
-                'table_name': input_data['table_name'],
-                'data': prediction
-            },
-            retry_policy=retry_policy,
-            start_to_close_timeout=timedelta(seconds=60)
-        )
-
         # write to opc
-        opc_holder = workflow.execute_activity_method(
+        prediction = await workflow.execute_activity_method(
             Activities.write_opc_data,
             {
                 'opc_output_config': input_data['opc_output_config'],
@@ -91,5 +79,14 @@ class FormatAndExportPrediction():
             start_to_close_timeout=timedelta(seconds=60)
         )
 
-        await postgres_holder
-        await opc_holder
+        # write to postgres
+        await workflow.execute_activity_method(
+            Activities.export_data_to_postgres,
+            {
+                'schema': input_data['schema'],
+                'table_name': input_data['table_name'],
+                'data': prediction,
+            },
+            retry_policy=retry_policy,
+            start_to_close_timeout=timedelta(seconds=60)
+        )