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, config in filters.items(): policy = config['policy'] 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()