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.repository.pi_web_api_client import PIWebAPIClient


PI_WEB_API_PREDICTION_ERROR_CONFIDENCE = 13


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

    This class provides Temporal activities for interacting with the PI Web API
    to write prediction and confidence values to industrial systems. It handles
    error scenarios gracefully by setting error confidence values and sending
    notifications when write operations fail.

    The class implements comprehensive error handling for both prediction and
    confidence value writes, ensuring that partial failures are properly
    reported and handled.
    """

    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 prediction and confidence data to PI Web API.

        This method writes prediction values and confidence scores to PI Web API
        using configured web IDs. It handles errors gracefully by setting error
        confidence values when prediction writes fail and sending notifications
        for both prediction and confidence write errors.

        Args:
            input_data (dict[str, Any]): The input data containing:
                - metadata (dict[str, Any]): Workflow execution metadata
                - pi_web_api_output_config (dict[str, Any]): PI Web API configuration with:
                    - endpoint (str): PI Web API endpoint URL
                    - prediction_tags (dict[str, str]): Mapping of tag names to web IDs for predictions
                    - confidence_tags (dict[str, str]): Mapping of tag names to web IDs for confidence
                - data (dict[str, Any]): Prediction data, its a dataframe converted to dict.
        Returns:
            dict[Any, Any]: Data dictionary with potentially modified confidence values
                If prediction write fails, prediction_confidence is set to error value (13)
        """
        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()