sientia-dataops-laborious_temporal/laborious/activities/api.py

from temporalio import activity, workflow

with workflow.unsafe.imports_passed_through():
    import traceback
    from typing import Any

    from pandas import DataFrame

    from sientia_do.notifications.handlers import NotificationHandler
    from sientia_do.notifications.models import NotificationLevel
    from sientia_do.observability.logger import Logger
    from sientia_do.observability.metrics_controller import MetricsController
    from sientia_do.observability.sientia_monitoring import SientiaMonitoring
    from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
    from sientia_do.repository.pi_web_api_client import PIWebAPIClient


PI_WEB_API_PREDICTION_ERROR_CONFIDENCE = 13

class API(SientiaMonitoring):
    """
    PI Web API operations for writing data to PI Web API.

    This class provides Temporal activities for interacting with the PI Web API
    to write data to PI Web API.
    """

    def __init__(
        self,
        base_url: str,
        auth_type: str,
        auth_token: str,
        logger: Logger,
        notification_handler: NotificationHandler,
        metrics_controller: MetricsController,
    ) -> None:
        """
        Initialize API activity with PI Web API client.

        Args:
            base_url (str): Base URL of the PI Web API server
            auth_type (str): Authentication type ('basic' or 'bearer')
            auth_token (str): Authentication token
            logger (Logger): Logger instance for operation logging
            notification_handler (NotificationHandler): Handler for system notifications
            metrics_controller (MetricsController): Controller for metrics collection
        """
        SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
        self.pi_web_api_client = PIWebAPIClient(
            base_url=base_url,
            auth_config={
                'type': auth_type,
                'token': auth_token,
            },
            logger=logger,
            notification_handler=notification_handler,
            metrics_controller=metrics_controller,
        )

    def close(self) -> None:
        """
        Close the PI Web API client and shutdown monitoring services.
        """
        self.pi_web_api_client.close()
        SientiaMonitoring.shutdown(self)

    @activity.defn(name='write_pi_web_api_data')
    async def write_pi_web_api_data(self, input_data: dict[str, Any]) -> dict[Any, Any]:
        """
        Write data to PI Web API.

        Args:
            input_data (dict[str, Any]): The input data. Containing:
                - metadata (dict[str, Any]): The metadata.
                - pi_web_api_output_config (dict[str, Any]): The PI Web API output configuration.
                - data (dict[str, Any]): The data to write.
        """
        metadata = input_data['metadata']
        data = DataFrame(input_data['data'])
        pi_web_api_output_config = input_data['pi_web_api_output_config']
        
        self.info('Writing data to PI Web API...', metadata)

        endpoint = pi_web_api_output_config['endpoint']

        raw_prediction_tags = pi_web_api_output_config['prediction_tags']
        raw_confidence_tags = pi_web_api_output_config['confidence_tags']
        prediction_tags = list[str](raw_prediction_tags.values())
        confidence_tags = list[str](raw_confidence_tags.values())

        prediction_value = data.head(1)['prediction'].values[0]
        confidence_value = data.head(1)['prediction_confidence'].values[0]

        try:

            await self.pi_web_api_client.write_value(
                web_ids=prediction_tags,
                value={
                    'Timestamp': data.head(1)['timestamp'].values[0],
                    'Value': prediction_value,
                },
                endpoint=endpoint,
                metadata=metadata,
            )

        except Exception as e:
            trace = traceback.format_exc()
            await self.send_notification_async(
                metadata=metadata,
                notification_id='WRITE_PI_WEB_API_PREDICTION_ERROR',
                message=f'Error writing prediction data to PI Web API: {e}\n Tags: {raw_prediction_tags}',
                block='write_pi_web_api_data',
                level=NotificationLevel.ERROR,
                attachment_content=trace,
            )

            data['prediction_confidence'] = PI_WEB_API_PREDICTION_ERROR_CONFIDENCE

            return data.to_dict()

        try:
            await self.pi_web_api_client.write_value(
                web_ids=confidence_tags,
                value={
                    'Timestamp': data.head(1)['timestamp'].values[0],
                    'Value': confidence_value,
                },
                endpoint=endpoint,
                metadata=metadata,
            )
        except Exception as e:
            trace = traceback.format_exc()
            await self.send_notification_async(
                metadata=metadata,
                notification_id='WRITE_PI_WEB_API_CONFIDENCE_ERROR',
                message=f'Error writing confidence data to PI Web API: {e}\n Tags: {raw_confidence_tags}',
                block='write_pi_web_api_data',
                level=NotificationLevel.ERROR,
                attachment_content=trace,
            )
        
        return data.to_dict()