sientia-dataops-scouter_temporal/scouter/activities/gates.py

from temporalio import workflow, activity

with workflow.unsafe.imports_passed_through():
    from sientia_do.notifications.models import NotificationLevel
    from sientia_do.temporal.activities.base import BaseActivity
    from scouter.utils.quality.filters import null_values_filter, out_of_bounds_filter
    from typing import Any
    import traceback
    from pandas import DataFrame

quality_gate_filters = {
    'NULL_VALUES_FILTER': null_values_filter,
    'OUT_OF_BOUNDS_FILTER': out_of_bounds_filter
}


class Gates(BaseActivity):

    def apply_aggregation(self, group: DataFrame, aggr_function: str,
                          metadata: dict[str, Any]) -> float | None | str:
        """
        Apply aggregation function to a group of data.

        Args:
            group (DataFrame): The group of data to apply the aggregation function to.
            aggr_function (str): The aggregation function to apply.

        Returns:
            float | None | str: The result of the aggregation function.
        """
        if len(group) == 1:
            return group['value'].item()

        # Apply aggregation function to value
        if aggr_function == 'lts':
            return group['value'].iloc[-1]
        else:
            group.dropna(inplace=True, subset=['value'])

            if group.empty:
                return None

            if aggr_function == 'avg':
                return group['value'].mean()
            elif aggr_function == 'mdn':
                return group['value'].median()
            elif aggr_function == 'max':
                return group['value'].max()
            elif aggr_function == 'min':
                return group['value'].min()
            else:
                self.send_notification(
                    metadata=metadata,
                    notification_id="AGGREGATION_ISSUES",
                    message=f"Invalid aggregation function: {aggr_function}",
                    block="aggregate_data",
                    level=NotificationLevel.ERROR,
                    attachment_content=traceback.format_exc()
                )
                return 'continue'

    @activity.defn(name="aggregate_data")
    async def aggregate_data(self, input_data: dict[str, Any]) -> dict[str, Any]:
        """
        Aggregates time series data by tag and name, applying specified
        aggregation functions and taking the latest timestamp.

        Args:
            input_data (dict[str, Any]): The data to aggregate. Contains:
                data (dict[str, Any]): The time series data.
                model_tags (dict[str, Any]): The tags configuration 
                containing aggregation functions.

        Returns:
            dict[str, Any]: The aggregated data.
        """

        metadata = input_data['metadata']

        try:
            # Convert input data to DataFrame
            df = DataFrame(input_data['data'])

            self.debug(
                f"Aggregating time series data: {df.to_string()}",
                metadata=metadata
            )

            # Initialize result dictionary
            result = {}

            # Group by tag and name
            grouped = df.groupby(['tag', 'name'])

            for (tag, name), group in grouped:
                # Get the aggregation function from model_tags
                aggr_function = input_data['model_tags'].get(
                    name, {}).get('aggr_func', 'lts')

                group.sort_values(by='timestamp', inplace=True)

                # Get the latest timestamp
                latest_timestamp = group['timestamp'].max()

                aggr_value = self.apply_aggregation(
                    group, aggr_function, metadata)

                if aggr_value == 'continue':
                    continue

                self.debug(
                    f"Aggregated data: {aggr_value}",
                    metadata=metadata
                )
                self.debug(
                    f"Latest timestamp: {latest_timestamp}",
                    metadata=metadata
                )
                self.debug(
                    f"Groups: {group.to_string()}",
                    metadata=metadata
                )
                self.debug(
                    f"group name: {name}",
                    metadata=metadata
                )
                self.debug(
                    f"group tag: {tag}",
                    metadata=metadata
                )

                # Store the result
                result[f"{tag}_{name}"] = {
                    'tag': tag,
                    'name': name,
                    'value': aggr_value,
                    'timestamp': latest_timestamp,
                    'aggregation_function': aggr_function
                }

            result_df = DataFrame(list(result.values()))
            self.debug(
                f"Aggregated data:\n {result_df.to_string()}",
                metadata=metadata
            )
            return result_df.to_dict()

        except Exception as e:
            trace = traceback.format_exc()

            self.send_notification(
                metadata=metadata,
                notification_id="AGGREGATION_ISSUES",
                message=f"Error aggregating data: {e}",
                block="aggregate_data",
                level=NotificationLevel.ERROR,
                attachment_content=trace
            )

            self.error(trace, metadata=metadata)
            raise e

    @activity.defn(name="data_quality_gate")
    async def data_quality_gate(self, input_data: dict[str, Any]) -> dict[str, Any]:
        """
        Data quality gate activity. for each selected filter,
        extracts filtered data, discards or keeps filtered data
        based on the filter.

        Args:
            input_data (dict[str, Any]): The data to validate. Contains:
                filters (dict[str, str]): The filters to apply. In format:
                    {filter_name: policy}.
                    filter_name: The name of the filter.
                    policy: The policy to apply. Can be "DISCARD" or "KEEP".
                data (dict[str, Any]): The data to validate.
                model_tags (dict[str, Any]): The tags of the model.
                    And it's respective configuration.

        Returns:
            dict[str, Any]: The data validated.
        """

        metadata = input_data['metadata']

        filters = input_data['filters']
        data = DataFrame(input_data['data'])
        model_tags = input_data['model_tags']

        self.debug(
            f"Applying quality gate to data: {data.to_string()}",
            metadata=metadata
        )

        for filter_name, policy in filters.items():
            if filter_name not in quality_gate_filters:
                self.warning(
                    f"Filter {filter_name} not found",
                    metadata=metadata
                )
                continue

            try:
                filtered_data = quality_gate_filters[filter_name](
                    data, model_tags)

            except Exception as e:
                trace = traceback.format_exc()
                self.send_notification(
                    metadata=metadata,
                    notification_id="DATA_QUALITY_GATE_ISSUES",
                    message=f"Error applying filter {filter_name}: {e}",
                    block="data_quality_gate",
                    level=NotificationLevel.ERROR,
                    attachment_content=trace
                )

                self.error(trace, metadata=metadata)

            else:
                if filtered_data.empty:
                    continue

                message = f"{len(filtered_data)} rows has quality issues: {filter_name}: {policy}"
                attachment = filtered_data.to_string()

                self.send_notification(
                    metadata=metadata,
                    notification_id=f"DATA_QUALITY_GATE_ISSUES__{filter_name}",
                    message=message,
                    block="data_quality_gate",
                    level=NotificationLevel.WARNING,
                    attachment_content=attachment
                )

                if policy == "DISCARD":
                    data = data[~data.index.isin(filtered_data.index)]

        self.debug(
            "Data quality gate applied",
            metadata=metadata
        )

        return data.to_dict()