Merge pull request #30 from Aignosi/feature/SIENTIAPDE-1273

SIENTIAPDE-1273: Refactor Data Handling, Enhance Security Analysis, and Update Dependencies

README.md

@@ -130,8 +130,14 @@ Laborious uses a Temporal-based architecture with strong separation of concerns
 - `minimal_retrain.py`: Automated model retraining and production update
 
 #### **Activities (`laborious/activities/`)**
-- `gates.py`: Data quality validation and filtering
+- `gates.py`: Data quality validation, filtering, and data formatting operations
-- `mlflow.py`: Transform and predict operations
+  - Input/response/content gates for quality validation
+  - Prediction and transformed data formatting
+  - Retrain report formatting and metrics recording
+- `mlflow.py`: Transform, predict, and model management operations
+  - MLFlow model transformation and prediction
+  - Model retraining and production updates
+  - Reference data retrieval from MLflow Model Registry
 - `opc.py`: OPC UA export to industrial systems (optional)
 - `activities.py`: Aggregates activity interfaces
 
@@ -146,7 +152,9 @@ Laborious uses a Temporal-based architecture with strong separation of concerns
 #### **1. Batch Prediction Pipeline**
 ```
 Input Data (PostgreSQL) → Data Quality Gates → MLFlow Transform →
-MLFlow Prediction → Response Validation → Export (PostgreSQL [+ OPC])
+MLFlow Prediction → Response Validation → Format & Export
+  ├─→ Predictions → PostgreSQL [+ OPC]
+  └─→ Transformed Data → PostgreSQL (optional)
 ```
 
 #### **2. Model Retraining Pipeline**
@@ -318,27 +326,37 @@ The **FormatAndExportPrediction** workflow handles prediction data formatting an
 #### Execution Flow
 1. **Path Decision**: Determines formatting path based on configuration
 2. **Data Formatting**: Formats prediction data for specific output requirements
-3. **PostgreSQL Export**: Writes formatted predictions to database
+3. **Transformed Data Processing**: Optionally formats and exports transformed data separately
-4. **OPC Export**: Writes predictions to OPC servers
+4. **PostgreSQL Export**: Writes formatted predictions to database
-5. **Metrics Recording**: Records export performance and success metrics
+5. **OPC Export**: Writes predictions to OPC servers
+6. **Metrics Recording**: Records export performance and success metrics
 
 #### Key Features
 - **Flexible Formatting**: Configurable output formats for different destinations
 - **Multi-Destination Export**: PostgreSQL and OPC server integration
+- **Transformed Data Export**: Optional separate export of MLFlow transformed data
 - **Performance Monitoring**: Comprehensive metrics for export operations
 - **Error Handling**: Robust error handling with notification integration
 
 #### Architecture Diagram
 ```mermaid
 flowchart LR
-    A[1. format_prediction/format_default_prediction] --> B[2. write_opc_data] --> C[3. export_data_to_postgres] --> D[4. write_metrics]
+    A[1. format_prediction/format_default_prediction] --> B[2. format_transformed_data] --> C[3. write_opc_data] --> D[4. export_data_to_postgres] --> E[5. write_metrics]
     
     A -.-> Format[Data Formatting]
-    B -.-> OPC[OPC Servers]
+    B -.-> Transform[Transformed Data]
-    C -.-> PostgreSQL[(PostgreSQL)]
+    C -.-> OPC[OPC Servers]
-    D -.-> Prometheus[Prometheus]
+    D -.-> PostgreSQL[(PostgreSQL)]
+    E -.-> Prometheus[Prometheus]
 ```
 
+#### Transformed Data Export
+When `transformed_data` is provided in the input, the workflow will:
+- Format the transformed data using `format_transformed_data` activity
+- Export it to a separate table (`transform_table_name`) asynchronously
+- Wait for both prediction and transformed data exports to complete
+- This enables separate tracking of model transformations for analysis and debugging
+
 ### 4. Minimal Retrain Workflow (`minimal_retrain.py`)
 
 The **MinimalRetrain** workflow handles automated model retraining and production model updates.
@@ -579,11 +597,24 @@ The workflow at `.github/workflows/quality-gate.yml` executes validations on eac
 ```
 tests/
 ├── activities/           # Activity implementation tests
-├── workflow/            # Workflow orchestration tests
+│   ├── test_gates.py    # Data quality gates and formatting tests
+│   ├── test_mlflow.py   # MLFlow operations and reference data tests
+│   └── ...              # Other activity tests
+├── workflows/           # Workflow orchestration tests
+│   └── subworkflows/    # Sub-workflow tests
+│       └── test_format_and_export_prediction.py  # Export workflow tests
 ├── utils/               # Utility function tests
 └── integration/         # End-to-end workflow tests
 ```
 
+### Test Coverage
+The test suite provides comprehensive coverage for:
+- **Data Quality Gates**: Input, response, and content validation filters
+- **Data Formatting**: Prediction, transformed data, and retrain report formatting
+- **MLFlow Operations**: Transform, predict, retrain, and reference data retrieval
+- **Workflow Orchestration**: Complete workflow execution paths and error handling
+- **Metrics Recording**: Performance monitoring and OPC export metrics
+
 ### Test Execution
 ```bash
 # Install test dependencies

laborious/activities/activities.py

@@ -9,11 +9,12 @@ with workflow.unsafe.imports_passed_through():
 
     from laborious.activities.gates import Gates
     from laborious.activities.mlflow import MLFlow
+    from laborious.activities.model_metrics import ModelMetrics
     from laborious.activities.opc import OPC
     from laborious.activities.storage import Storage
 
 
-class Activities(Storage, MLFlow, Gates, OPC):
+class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics):
     """
     Main activities orchestrator for the Laborious system.
 
@@ -108,6 +109,13 @@ class Activities(Storage, MLFlow, Gates, OPC):
             metrics_controller=metrics_controller,
         )
 
+        ModelMetrics.__init__(
+            self,
+            logger=logger,
+            notification_handler=notification_handler,
+            metrics_controller=metrics_controller,
+        )
+
     async def shutdown(self):
         """
         Gracefully shutdown all activities and clean up resources.
@@ -124,3 +132,4 @@ class Activities(Storage, MLFlow, Gates, OPC):
         MLFlow.close(self)
         Gates.close(self)
         await OPC.close(self)
+        ModelMetrics.close(self)

laborious/activities/gates.py

@@ -402,8 +402,54 @@ class Gates(SientiaMonitoring):
 
         return policy_type, int(policy_value)
 
+    @activity.defn(name='format_transformed_data')
+    async def format_transformed_data(self, input_data: dict[str, Any]) -> dict:
+        """
+        Format transformed data for storage and export operations.
+
+        This method formats transformed data from MLFlow model transformations
+        into a standardized format suitable for database storage. It converts
+        wide-format data (columns as variables) into long-format (melted)
+        with proper timestamp handling and model identification.
+
+        The formatting process includes:
+        1. Converting input data dictionary to DataFrame
+        2. Extracting timestamps from DataFrame index
+        3. Resetting index to create sequential row numbers
+        4. Melting data from wide format to long format (variable-value pairs)
+        5. Adding model_id for data lineage tracking
+
+        Args:
+            input_data (dict): Input data containing:
+                - metadata (dict): Workflow execution metadata
+                - data (dict[str, Any]): Transformed data to format (DataFrame-compatible dict)
+                - model_id (str): Unique identifier for the ML model
+
+        Returns:
+            dict: Formatted data dictionary with keys:
+                - timestamp (dict): Timestamp values indexed by row number
+                - variable (dict): Variable names indexed by row number
+                - value (dict): Variable values indexed by row number
+                - model_id (dict): Model identifiers indexed by row number
+        """
+        metadata = input_data['metadata']
+
+        model_id = input_data['model_id']
+
+        self.info('Formatting transformed data...', metadata)
+
+        data = DataFrame(input_data['data'])
+
+        data['timestamp'] = data.index
+        data = data.reset_index(drop=True)
+
+        data = data.melt(id_vars='timestamp', var_name='variable', value_name='value')
+        data['model_id'] = model_id
+
+        return data.to_dict()
+
     @activity.defn(name='format_prediction')
-    async def format_prediction(self, input_data: dict[str, Any]) -> dict[Any, Any]:
+    async def format_prediction(self, input_data: dict[str, Any]) -> dict:
         """
         Format prediction data according to configured storage policies.
 
@@ -476,7 +522,7 @@ class Gates(SientiaMonitoring):
         return data.to_dict()
 
     @activity.defn(name='format_default_prediction')
-    async def format_default_prediction(self, input_data: dict[str, Any]) -> dict[Any, Any]:
+    async def format_default_prediction(self, input_data: dict[str, Any]) -> dict:
         """
         Create and format default prediction data for error conditions.
 
@@ -521,9 +567,47 @@ class Gates(SientiaMonitoring):
         return data.to_dict()
 
     @activity.defn(name='format_retrain_report')
-    async def format_retrain_report(self, input_data: dict[str, Any]) -> dict[Any, Any]:
+    async def format_retrain_report(self, input_data: dict[str, Any]) -> dict:
         """
-        Format retrain report data according to configured storage policies.
+        Format retrain report data for storage and audit trail maintenance.
+
+        This method formats model retraining operation results into a standardized
+        report format suitable for database storage and operational monitoring.
+        It captures retraining status, timestamps, and model version information
+        for comprehensive audit trails and operational visibility.
+
+        The formatting process includes:
+        1. Extracting retraining experiment response data
+        2. Capturing model update report information (version, MLflow IDs)
+        3. Formatting timestamps and status information
+        4. Conditionally including version information for successful retrains
+
+        Args:
+            input_data (dict): Input data containing:
+                - metadata (dict): Workflow execution metadata
+                - experiment_response (dict): Retraining experiment response containing:
+                    - success (bool): Retraining operation success status
+                    - timestamp (str): Timestamp of the retraining operation
+                    - message (str): Status message or error description
+                - update_report (dict): Model update report containing:
+                    - version (str): New model version identifier
+                    - mlflow_run_id (str): MLflow run identifier
+                    - mlflow_experiment_id (str): MLflow experiment identifier
+                - model_id (str): Unique identifier for the ML model
+                - model_name (str): Name of the ML model
+
+        Returns:
+            dict: Formatted retrain report dictionary with keys:
+                - model_id (dict): Model identifiers indexed by row number
+                - model_name (dict): Model names indexed by row number
+                - timestamp (dict): Retraining timestamps indexed by row number
+                - status (dict): Retraining status messages indexed by row number
+                - version (dict, optional): Model versions indexed by row number
+                    Only included if experiment_response['success'] is True
+                - mlflow_run_id (dict, optional): MLflow run IDs indexed by row number
+                    Only included if experiment_response['success'] is True
+                - mlflow_experiment_id (dict, optional): MLflow experiment IDs indexed by row number
+                    Only included if experiment_response['success'] is True
         """
         metadata = input_data['metadata']
         self.info('Formatting retrain report...', metadata)

laborious/activities/mlflow.py

@@ -418,3 +418,51 @@ class MLFlow(SientiaMonitoring):
             )
             self.error(trace, metadata=metadata)
             raise e
+
+    @activity.defn(name='get_reference_data')
+    async def get_reference_data(self, input_data: dict[str, Any]) -> list[dict] | None:
+        """
+        Get reference data from the MLflow Model Registry.
+
+        This method retrieves evaluation reference data stored as artifacts in the
+        MLflow Model Registry. The reference data is typically used for model
+        drift detection, performance comparison, and quality validation. The method
+        loads the data from a CSV artifact file and formats timestamps for
+        consistent processing.
+
+        The method handles:
+        1. Loading evaluation data artifact from MLflow Model Registry
+        2. Timestamp parsing and formatting for consistency
+        3. Data conversion to dictionary format for workflow consumption
+        4. Graceful handling of missing reference data
+
+        Args:
+            input_data (dict): Input data containing:
+                - metadata (dict): Workflow execution metadata
+                - model_name (str): Name of the MLFlow model to get reference data from
+
+        Returns:
+            list[dict[Hashable, Any]] | None: Reference data from the MLflow Model Registry
+                as a list of dictionaries. Returns None if reference data is not found
+                or if the artifact does not exist.
+
+        Raises:
+            Exception: If artifact loading fails or encounters errors during processing
+        """
+
+        metadata = input_data['metadata']
+        model_name = input_data['model_name']
+        artifact = 'evaluation_data.csv'
+
+        reference_data = await self.model_monitoring_repository.load_artifact_dataframe(
+            model_name=model_name, artifact_path=artifact, metadata=metadata
+        )
+
+        if reference_data is None:
+            self.warning(f'Reference data not found for model {model_name}', metadata)
+            return None
+
+        reference_data['timestamp'] = to_datetime(reference_data['timestamp'])
+        reference_data['timestamp'] = reference_data['timestamp'].dt.strftime(DATETIME_FORMAT)
+
+        return reference_data.to_dict(orient='records')

laborious/activities/model_metrics.py

@@ -0,0 +1,354 @@
+from temporalio import activity, workflow
+
+with workflow.unsafe.imports_passed_through():
+    import time
+    import traceback
+    import warnings
+    from typing import Any
+
+    import numpy as np
+    from pandas import DataFrame, Index, to_datetime
+    from sientia.ModelAnalysis import ModelAnalysis
+    from sientia_do.notifications.handlers import CoreNotificationHandler as 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, DATETIME_FORMAT_WITH_TZ
+
+    from laborious import metrics
+
+warnings.filterwarnings('ignore', category=RuntimeWarning, message='Degrees of freedom <= 0')
+warnings.filterwarnings(
+    'ignore', category=RuntimeWarning, message='invalid value encountered in scalar divide'
+)
+
+
+class ModelMetrics(SientiaMonitoring):
+    """
+    Metrics activities for the Laborious system.
+
+    This class provides activities for writing metrics to the Prometheus monitoring system.
+    """
+
+    def __init__(
+        self,
+        logger: Logger,
+        notification_handler: NotificationHandler,
+        metrics_controller: MetricsController,
+    ):
+        SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
+
+    def close(self) -> None:
+        """
+        Close the model metrics activity and clean up resources.
+        """
+        SientiaMonitoring.shutdown(self)
+
+    def __del__(self):
+        self.close()
+
+    async def get_drift_metrics(
+        self,
+        reference_data: DataFrame,
+        target_data: DataFrame,
+        target_name: str,
+        reference_columns: Index,
+        drift_metrics: list[str],
+        chunk_period: str,
+        metadata: dict[str, Any],
+    ) -> DataFrame:
+        """
+        Calculate univariate drift metrics for a model.
+        Args:
+            model_analysis (ModelAnalysis): Model analysis object
+            reference_data (DataFrame): Reference data
+            target_data (DataFrame): Target data
+            reference_columns (list[str]): Reference columns
+            drift_metrics (list[str]): Drift metrics
+            metadata (dict[str, Any]): Workflow execution metadata
+        """
+
+        config = {
+            'target': target_name,
+            'prediction': 'prediction',
+            'timestamp': 'timestamp',
+            'features': reference_columns,
+        }
+
+        model_analysis = ModelAnalysis(config=config)
+
+        self.debug(
+            f'Reference data: Size {reference_data.shape} \n{reference_data.head(5).to_string()}',
+            metadata,
+        )
+
+        self.debug(
+            f'Target data: Size {target_data.shape} \n{target_data.head(5).to_string()}', metadata
+        )
+
+        core_labels = self.get_core_labels(metadata, operation_type='detect_univariate_drift')
+        start_time = time.time()
+        try:
+            univariate_drift = model_analysis.detect_univariate_drift(
+                reference_df=reference_data,
+                analysis_df=target_data,
+                features=reference_columns,
+                timestamp_col=config['timestamp'],
+                methods=drift_metrics,
+                chunk_period=chunk_period,
+            )
+        except Exception as e:
+            self.error(f'Error detecting univariate drift: {e}', metadata)
+            await self.emit_metric(
+                metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels
+            )
+            raise e
+        await self.observe_lag(start_time, metrics.MODEL_ANALYZE_LAG, core_labels)
+        await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_COUNT, tags=core_labels)
+
+        core_labels = self.get_core_labels(metadata, operation_type='detect_multivariate_drift')
+        start_time = time.time()
+        try:
+            multivariate_drift = model_analysis.detect_multivariate_drift(
+                reference_df=reference_data,
+                analysis_df=target_data,
+                features=reference_columns,
+                timestamp_col=config['timestamp'],
+                chunk_period=chunk_period,
+            )
+        except Exception as e:
+            self.error(f'Error detecting multivariate drift: {e}', metadata)
+            await self.emit_metric(
+                metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels
+            )
+            raise e
+        await self.observe_lag(start_time, metrics.MODEL_ANALYZE_LAG, core_labels)
+        await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_COUNT, tags=core_labels)
+
+        start_time = time.time()
+        core_labels = self.get_core_labels(metadata, operation_type='get_drift_metrics_dataframe')
+        try:
+            drift_df = model_analysis.get_drift_metrics_dataframe(
+                univariate_drift=univariate_drift,
+                multivariate_drift=multivariate_drift,
+            )
+        except Exception as e:
+            self.error(f'Error getting drift metrics: {e}', metadata)
+            await self.emit_metric(
+                metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels
+            )
+            raise e
+        await self.observe_lag(start_time, metrics.MODEL_ANALYZE_LAG, core_labels)
+        await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_COUNT, tags=core_labels)
+
+        self.debug(
+            f'Drift dataframe: Size {drift_df.shape} \n{drift_df.head(5).to_string()}', metadata
+        )
+
+        return drift_df
+
+    @activity.defn(name='calculate_drift')
+    async def calculate_drift(self, input_data: dict[str, Any]) -> list[dict]:
+        """
+        Calculate drift metrics for a model.
+
+        Args:
+            input_data (dict[str, Any]): Input data containing:
+                - metadata (dict): Workflow execution metadata
+                - model_name (str): Name of the MLFlow model to calculate drift for
+                - reference_data (pd.DataFrame): Reference data for the model
+                - target_data (pd.DataFrame): Target data for calculating drift
+                - target_name (str): Name of the target column
+                - drift_metrics (list[str]): List of drift metrics to calculate
+        """
+        metadata = input_data['metadata']
+        model_name = input_data['model_name']
+        model_id = input_data['model_id']
+        reference_raw_data = input_data['reference_data']
+        target_data = DataFrame(input_data['target_data'])
+        target_name = input_data['target_name']
+        drift_metrics = input_data['drift_metrics']
+        chunk_period = input_data['chunk_period']
+
+        if chunk_period not in ['min', 's']:
+            self.error(f'Invalid chunk period: {chunk_period}', metadata)
+            raise ValueError(f'Invalid chunk period: {chunk_period}, must be "min" or "s"')
+
+        self.info(f'Calculating drift for model {model_name}', metadata)
+
+        target_data = target_data.pivot(index='timestamp', columns='variable', values='value')
+        target_data['timestamp'] = target_data.index
+        target_data['timestamp'] = to_datetime(target_data['timestamp'])
+        target_data['timestamp'] = target_data['timestamp'].dt.strftime(DATETIME_FORMAT)
+        target_data = target_data.reset_index(drop=True)
+        target_data.dropna(inplace=True)
+
+        if reference_raw_data is not None:
+            self.info('Using reference data', metadata)
+            reference_data = DataFrame(reference_raw_data)
+            accurate = True
+        else:
+            # Get 30% first rows of target_data
+            self.warning('Using 30% first rows of target data as reference data', metadata)
+            target_data.sort_values(by='timestamp', ascending=True, inplace=True)
+            reference_data = target_data.head(int(len(target_data) * 0.3))
+            accurate = False
+
+            await self.send_notification_async(
+                metadata=metadata,
+                notification_id='MODEL_METRICS_REFERENCE_DATA_WARNING',
+                message='Using 30% first rows of target data as reference data',
+                block='model_metrics',
+                level=NotificationLevel.WARNING,
+                attachment_content=reference_data.to_csv(),
+            )
+
+        reference_columns = reference_data.drop(
+            columns=[target_name, 'timestamp', 'target', 'prediction'], errors='ignore'
+        ).columns
+
+        try:
+            drift_df = await self.get_drift_metrics(
+                reference_data=reference_data,
+                target_data=target_data,
+                target_name=target_name,
+                reference_columns=reference_columns,
+                drift_metrics=drift_metrics,
+                chunk_period=chunk_period,
+                metadata=metadata,
+            )
+        except Exception as e:
+            self.error(f'Error getting drift metrics: {e}', metadata)
+            await self.send_notification_async(
+                metadata=metadata,
+                notification_id='MODEL_METRICS_GET_DRIFT_METRICS_ERROR',
+                message=f'Error getting drift metrics: {e}',
+                block='model_metrics',
+                level=NotificationLevel.ERROR,
+                attachment_content=traceback.format_exc(),
+            )
+            return []
+
+        if drift_df.empty:
+            self.warning('No drift metrics found', metadata)
+            return []
+
+        # Drop unnecessary columns
+        drift_df.drop(columns=['p_value'], inplace=True)
+
+        # Extract timestamps only until minutes
+        if chunk_period == 'min':
+            target_timestamps = target_data['timestamp'].apply(lambda x: x[:16])
+        else:
+            target_timestamps = target_data['timestamp']
+
+        # Drop rows where timestamp is not in target data, to avoid save drift from reference
+        drift_df = drift_df[drift_df['timestamp'].isin(target_timestamps)]
+
+        if drift_df.empty:
+            self.warning(
+                'No drift metrics found after dropping rows where timestamp is not in target data',
+                metadata,
+            )
+            return []
+
+        # Rename columns to match database columns
+        drift_df.rename(
+            columns={
+                'metric': 'method',
+                'statistic': 'value',
+            },
+            inplace=True,
+        )
+
+        # Drop duplicates
+        drift_df.drop_duplicates(
+            subset=['timestamp', 'method', 'feature'], keep='first', inplace=True
+        )
+
+        drift_df['model_id'] = model_id
+        drift_df['accurate'] = accurate
+
+        drift_df['timestamp'] = to_datetime(drift_df['timestamp'])
+        drift_df['timestamp'] = drift_df['timestamp'].dt.tz_localize('UTC')
+        drift_df['timestamp'] = drift_df['timestamp'].dt.strftime(DATETIME_FORMAT_WITH_TZ)
+
+        self.debug(
+            f'Drift dataframe: Size {drift_df.shape} \n{drift_df.head(5).to_string()}', metadata
+        )
+
+        self.debug(f'Drift dataframe: {drift_df.head(5).to_string()}', metadata)
+
+        return drift_df.to_dict(orient='records')
+
+    @activity.defn(name='calculate_simple_metrics')
+    async def calculate_simple_metrics(self, input_data: dict[str, Any]) -> list[dict]:
+        """
+        Calculate simple metrics for a model. Metrics available are:
+            - rmse
+            - mse
+            - mae
+            - r2
+            - accuracy
+            - precision
+            - recall
+            - f1
+        Args:
+            input_data (dict[str, Any]): Input data containing:
+                - metadata (dict): Workflow execution metadata
+                - model_id (str): ID of the MLFlow model
+                - target_data (pd.DataFrame): Target data for calculating metrics, containing target and prediction columns
+                - metrics (list[str]): List of metrics to calculate
+        Returns:
+            dict[Hashable, Any]: Dictionary containing the calculated metrics
+        """
+
+        metadata = input_data['metadata']
+        model_id = input_data['model_id']
+        target_data = DataFrame(input_data['target_data'])
+        metrics = input_data['metrics']
+        interval_minutes = input_data['interval_minutes']
+
+        data_size = target_data.shape[0]
+
+        output_data = []
+
+        diff = target_data['target'] - target_data['prediction']
+        diff_squared = diff**2
+
+        self.info(f'Calculating simple metrics for model {model_id}: {metrics}', metadata)
+
+        for metric in metrics:
+            if metric == 'rmse':
+                output_data.append({'metric': 'rmse', 'value': np.sqrt(np.mean(diff_squared))})
+            elif metric == 'mse':
+                output_data.append({'metric': 'mse', 'value': np.mean(diff_squared)})
+            elif metric == 'mae':
+                output_data.append({'metric': 'mae', 'value': np.mean(np.abs(diff))})
+            elif metric == 'r2':
+                y_true = target_data['target']
+                y_mean = np.mean(y_true)
+
+                ss_res = np.sum(diff_squared)
+                ss_tot = np.sum((y_true - y_mean) ** 2)
+
+                # Evita divisão por zero
+                if ss_tot == 0:
+                    r2_score = 0.0
+                else:
+                    r2_score = 1 - (ss_res / ss_tot)
+
+                output_data.append({'metric': 'r2', 'value': r2_score})
+
+        data = DataFrame(output_data)
+        data['model_id'] = model_id
+        data['timestamp'] = target_data['timestamp'].max()
+        data['data_size'] = data_size
+        data['interval_minutes'] = interval_minutes
+
+        self.debug(
+            f'Simple metrics dataframe: Size {data.shape} \n{data.head(5).to_string()}', metadata
+        )
+
+        return data.to_dict(orient='records')

laborious/activities/opc.py

@@ -2,6 +2,7 @@ from temporalio import activity, workflow
 
 with workflow.unsafe.imports_passed_through():
     import traceback
+    from collections.abc import Hashable
     from typing import Any
 
     from pandas import DataFrame
@@ -275,7 +276,7 @@ class OPC(SientiaMonitoring):
     @activity.defn(name='write_opc_data')
     async def write_opc_data(
         self, input_data: dict[str, Any]
-    ) -> tuple[dict[Any, Any], dict[str, dict[str, float | None]]]:
+    ) -> tuple[dict[Hashable, Any], dict[str, dict[str, float | None]]]:
         """
         Write prediction and confidence data to OPC servers. The two writing
         operations are optional and independent of each other.
@@ -324,7 +325,7 @@ class OPC(SientiaMonitoring):
 
     def process_confidence(
         self, data: DataFrame, success: bool, metadata: dict[str, Any]
-    ) -> dict[Any, Any]:
+    ) -> dict[Hashable, Any]:
         """
         Process prediction confidence based on OPC write operation success.
 

laborious/metrics.py

@@ -169,3 +169,22 @@ MODEL_WRITE_ERROR_COUNT = Counter(
     'Number of errors writing to the model',
     SIENTIA_CORE_LABELS,
 )
+
+MODEL_ANALYZE_LAG = Histogram(
+    'laborious_model_analyze_lag',
+    'Lag between the start and end of analyze operations',
+    SIENTIA_CORE_LABELS,
+    buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0],
+)
+
+MODEL_ANALYZE_COUNT = Counter(
+    'laborious_model_analyze_count',
+    'Number of analyze operations',
+    SIENTIA_CORE_LABELS,
+)
+
+MODEL_ANALYZE_ERROR_COUNT = Counter(
+    'laborious_model_analyze_error_count',
+    'Number of errors during analyze operations',
+    SIENTIA_CORE_LABELS,
+)

laborious/utils/repository/model_repository.py

@@ -20,6 +20,7 @@ import threading
 import time
 import traceback
 from datetime import datetime, timedelta
+from io import StringIO
 from os import environ, makedirs, path
 from shutil import rmtree
 from typing import Any, Literal, overload
@@ -215,6 +216,26 @@ class MLFlowRepository(SientiaMonitoring):
         run_info = mlflow.get_run(run_id)
         return run_info.data.params
 
+    def check_artifact_exists(
+        self, run_id: str, artifact_path: str, metadata: dict[str, Any]
+    ) -> bool:
+        """
+        Check if an artifact exists in the MLflow Model Registry.
+
+        Args:
+            run_id (str): Run identifier to inspect.
+            artifact_path (str): Path to the artifact to check.
+
+        Returns:
+            bool: True if the artifact exists, False otherwise.
+        """
+        artifacts = self.client.list_artifacts(run_id)
+
+        self.debug(f'Artifacts of {run_id}: \n{artifacts}', metadata)
+        self.debug(f'Looking for artifact {artifact_path} in {run_id}', metadata)
+
+        return any(artifact.path == artifact_path for artifact in artifacts)
+
     """
     Functions related to download and load models
     """
@@ -259,6 +280,45 @@ class MLFlowRepository(SientiaMonitoring):
 
         return artifacts
 
+    async def load_artifact_dataframe(
+        self, model_name: str, artifact_path: str, metadata: dict[str, Any]
+    ) -> pd.DataFrame | None:
+        """
+        Load the dataframe content of an artifact from the MLflow Model Registry.
+
+        Args:
+            model_name (str): The name of the model to download from the registry.
+            artifact_path (str): The path to the artifact to load.
+            metadata (dict[str, Any]): Metadata used for structured logging.
+
+        Returns:
+            pd.DataFrame: The dataframe content of the artifact.
+        """
+        run_id = self.get_model_run_id(model_name=model_name, stage='Production')
+        core_labels = self.get_core_labels(metadata, operation_type='load_text')
+
+        if not self.check_artifact_exists(run_id, artifact_path, metadata):
+            return None
+
+        artifact_path = path.join('runs:/', run_id, artifact_path)
+
+        start_time = time.time()
+        try:
+            content = mlflow.artifacts.load_text(artifact_path)
+        except Exception as e:
+            await self.emit_metric(metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=core_labels)
+            raise e
+
+        await self.observe_lag(start_time, metrics.MODEL_READ_LAG, core_labels)
+        await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=core_labels)
+
+        self.debug(f'Content of {run_id}/{artifact_path}: \n{content}', metadata)
+
+        dataframe = pd.read_csv(StringIO(content))
+
+        self.info(f'Loaded dataframe from {model_name}:{artifact_path}', metadata)
+        return dataframe
+
     async def load_predict_model(
         self, model_name: str, metadata: dict[str, Any], flavor: str = 'sklearn'
     ) -> Any:
@@ -644,6 +704,47 @@ class MLFlowRepository(SientiaMonitoring):
     Functions related to model retraining
     """
 
+    def get_prediction_data(
+        self,
+        prediction_model: Any,
+        retrain_dataset: pd.DataFrame,
+        target_name: str,
+        predict_flavor: str,
+    ) -> pd.DataFrame:
+        """
+        Get prediction data from prediction model.
+        """
+        input_index = retrain_dataset.index
+
+        if predict_flavor == 'pyfunc':
+            prediction_data = prediction_model.predict({}, retrain_dataset)
+        else:
+            prediction_data = prediction_model.predict(retrain_dataset)
+
+        if isinstance(prediction_data, pd.DataFrame):
+            prediction_data.columns = pd.Index(['prediction'])
+
+        else:
+            prediction_data = pd.DataFrame(prediction_data, columns=['prediction'])
+
+        prediction_data.index = input_index
+
+        # Merge prediction data with retrain_dataset on index
+        prediction_data = pd.merge(  # NOSONAR
+            retrain_dataset, prediction_data, left_index=True, right_index=True, how='left'
+        )
+
+        # Rename column "target_name" to "target"
+        prediction_data.rename(columns={target_name: 'target'}, inplace=True)
+
+        prediction_data['timestamp'] = prediction_data.index
+
+        prediction_data.reset_index(drop=True, inplace=True)
+
+        prediction_data = prediction_data.sort_values(by='timestamp', ascending=True)
+
+        return prediction_data
+
     async def fit_models(
         self,
         model_name: str,
@@ -653,7 +754,7 @@ class MLFlowRepository(SientiaMonitoring):
         transform_flavor: str = 'sklearn',
         predict_flavor: str = 'sklearn',
         target_name: str | None = None,
-    ) -> dict[str, dict[str, Any]]:
+    ) -> dict[str, Any]:
         """
         Prepare models and data for a retraining run.
 
@@ -763,6 +864,11 @@ class MLFlowRepository(SientiaMonitoring):
 
         prediction_model.fit(retrain_dataset)
 
+        # get prediction data
+        prediction_data = self.get_prediction_data(
+            prediction_model, retrain_dataset, target_name, predict_flavor
+        )
+
         self.info(f'Model experiment creation completed successfully for {model_name}', metadata)
 
         retrain_data = {
@@ -771,6 +877,7 @@ class MLFlowRepository(SientiaMonitoring):
                 'artifact_path': prediction_artifact_path,
             },
             'data_model': {'model': data_model, 'artifact_path': data_artifact_path},
+            'prediction_data': prediction_data,
         }
         return retrain_data
 
@@ -849,6 +956,7 @@ class MLFlowRepository(SientiaMonitoring):
 
         prediction_model = retrain_data['prediction_model']
         data_model = retrain_data['data_model']
+        prediction_data = retrain_data['prediction_data']
 
         model_temp_path = path.join(ARTIFACTS_PATH, model_name)
 
@@ -872,10 +980,12 @@ class MLFlowRepository(SientiaMonitoring):
         self.debug(f'Attributes: {retrain_params}', metadata)
 
         data_path = f'{model_temp_path}/retrain_data.csv'
+        prediction_data_path = f'{model_temp_path}/evaluation_data.csv'
 
         makedirs(model_temp_path, exist_ok=True)
 
-        data.to_csv(data_path, index=True)
+        data.to_csv(data_path, index=False)
+        prediction_data.to_csv(prediction_data_path, index=False)
 
         self.info(
             f'Starting model upload for {experiment_name} with run name {current_run_name}',
@@ -909,6 +1019,7 @@ class MLFlowRepository(SientiaMonitoring):
 
                 # log the data raw
                 mlflow.log_artifact(data_path)
+                mlflow.log_artifact(prediction_data_path)
 
         except Exception as e:
             await self.emit_metric(metric_object=metrics.MODEL_WRITE_ERROR_COUNT, tags=core_labels)

laborious/worker/worker.py

@@ -5,12 +5,15 @@ This module provides the main worker implementation for the Sientia DataOps Labo
 It orchestrates Temporal workers, manages task queues, and handles the lifecycle of
 prediction and retraining workflows.
 
-The worker supports two main task queues:
+The worker supports multiple task queues:
-- predictions_batch-queue: Handles batch prediction workflows
+- predictions_batch-queue: Handles batch prediction workflows (heavy workload)
 - minimal_retrain-queue: Handles model retraining workflows
+- drift-queue: Handles drift detection workflows
+- simple_metrics-queue: Handles simple metrics calculation workflows
 
 Key Features:
-- Automatic scaling with PollerBehaviorAutoscaling
+- Resource-based scaling with WorkerTuner (CPU and memory aware)
+- Automatic polling scaling with PollerBehaviorAutoscaling
 - Prometheus metrics integration
 - Comprehensive error handling and logging
 - Graceful shutdown with cleanup
@@ -23,16 +26,37 @@ Environment Variables:
 - HTTP_METRICS_PORT: Prometheus metrics server port (default: 9090)
 - HTTP_SDK_METRICS_PORT: Temporal SDK metrics port (default: 9091)
 - PROJECT_NAME: Project name for notifications (default: laborious)
+
+Tuner Configuration (Resource-based scaling):
+- TUNER_TARGET_MEMORY_USAGE: Target memory usage (0.0-1.0, default: 0.75)
+- TUNER_TARGET_CPU_USAGE: Target CPU usage (0.0-1.0, default: 0.80)
+- TUNER_WORKFLOW_MIN_SLOTS: Minimum workflow slots (default: 5)
+- TUNER_WORKFLOW_MAX_SLOTS: Maximum workflow slots (default: 50)
+- TUNER_ACTIVITY_MIN_SLOTS: Minimum activity slots (default: 5)
+- TUNER_ACTIVITY_MAX_SLOTS: Maximum activity slots (default: 50)
+- TUNER_WORKFLOW_RAMP_THROTTLE_MS: Workflow ramp throttle in ms (default: 100)
+- TUNER_ACTIVITY_RAMP_THROTTLE_MS: Activity ramp throttle in ms (default: 50)
+
+Poller Configuration:
+- POLLER_MINIMUM: Minimum number of pollers (default: 1)
+- POLLER_MAXIMUM: Maximum number of pollers (default: 10)
+- POLLER_INITIAL: Initial number of pollers (default: 2)
 """
 
+
 from temporalio import client, workflow
 from temporalio.runtime import PrometheusConfig, Runtime, TelemetryConfig
-from temporalio.worker import PollerBehaviorAutoscaling, Worker
+from temporalio.worker import (
+    PollerBehaviorAutoscaling,
+    ResourceBasedSlotConfig,
+    Worker,
+    WorkerTuner,
+)
 
 with workflow.unsafe.imports_passed_through():
     import asyncio
-    import os
     import sys
+    from datetime import timedelta
 
     from prometheus_client import start_http_server
     from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
@@ -47,17 +71,63 @@ with workflow.unsafe.imports_passed_through():
         build_opc_config,
         build_postgres_config,
     )
+    from laborious.workflows.drift import Drift
     from laborious.workflows.minimal_retrain import MinimalRetrain
     from laborious.workflows.predictions_batch import PredictionsBatch
+    from laborious.workflows.simple_metrics import SimpleMetrics
     from laborious.workflows.sub_workflows.format_and_export_prediction import (
         FormatAndExportPrediction,
     )
     from laborious.workflows.sub_workflows.prediction_process import PredictionProcess
+    import os
 
 POD_ID = os.getenv('POD_ID')
 SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', '9091'))
 
 
+def create_resource_tuner() -> WorkerTuner:
+    """Create a resource-based tuner from environment variables."""
+    target_memory = float(os.getenv('TUNER_TARGET_MEMORY_USAGE', '0.75'))
+    target_cpu = float(os.getenv('TUNER_TARGET_CPU_USAGE', '0.50'))
+    workflow_min = int(os.getenv('TUNER_WORKFLOW_MIN_SLOTS', '5'))
+    workflow_max = int(os.getenv('TUNER_WORKFLOW_MAX_SLOTS', '50'))
+    activity_min = int(os.getenv('TUNER_ACTIVITY_MIN_SLOTS', '5'))
+    activity_max = int(os.getenv('TUNER_ACTIVITY_MAX_SLOTS', '50'))
+    local_activity_min = int(os.getenv('TUNER_LOCAL_ACTIVITY_MIN_SLOTS', '1'))
+    local_activity_max = int(os.getenv('TUNER_LOCAL_ACTIVITY_MAX_SLOTS', '30'))
+    workflow_ramp = int(os.getenv('TUNER_WORKFLOW_RAMP_THROTTLE_MS', '100'))
+    activity_ramp = int(os.getenv('TUNER_ACTIVITY_RAMP_THROTTLE_MS', '50'))
+    local_activity_ramp = int(os.getenv('TUNER_LOCAL_ACTIVITY_RAMP_THROTTLE_MS', '50'))
+
+    return WorkerTuner.create_resource_based(
+        target_memory_usage=target_memory,
+        target_cpu_usage=target_cpu,
+        workflow_config=ResourceBasedSlotConfig(
+            minimum_slots=workflow_min,
+            maximum_slots=workflow_max,
+            ramp_throttle=timedelta(milliseconds=workflow_ramp),
+        ),
+        activity_config=ResourceBasedSlotConfig(
+            minimum_slots=activity_min,
+            maximum_slots=activity_max,
+            ramp_throttle=timedelta(milliseconds=activity_ramp),
+        ),
+        local_activity_config=ResourceBasedSlotConfig(
+            minimum_slots=local_activity_min,
+            maximum_slots=local_activity_max,
+            ramp_throttle=timedelta(milliseconds=local_activity_ramp),
+        ),
+    )
+
+
+def create_poller_behavior() -> PollerBehaviorAutoscaling:
+    """Create poller behavior from environment variables."""
+    minimum = int(os.getenv('POLLER_MINIMUM', '1'))
+    maximum = int(os.getenv('POLLER_MAXIMUM', '10'))
+    initial = int(os.getenv('POLLER_INITIAL', '2'))
+    return PollerBehaviorAutoscaling(minimum=minimum, maximum=maximum, initial=initial)
+
+
 async def main():
     """
     Main entry point for the Laborious worker application.
@@ -136,6 +206,9 @@ async def main():
 
     logger.custom_info('Starting Workers...', metadata)
 
+    tuner = create_resource_tuner()
+    poller = create_poller_behavior()
+
     workers = [
         Worker(
             temporal_client,
@@ -149,12 +222,39 @@ async def main():
                 activities.format_retrain_report,
                 activities.export_data_to_postgres,
             ],
-            max_concurrent_workflow_tasks=50,
+            tuner=tuner,
-            max_concurrent_activities=50,
-            max_concurrent_local_activities=50,
             max_cached_workflows=2,
-            workflow_task_poller_behavior=PollerBehaviorAutoscaling(),
+            workflow_task_poller_behavior=poller,
-            activity_task_poller_behavior=PollerBehaviorAutoscaling(),
+            activity_task_poller_behavior=poller,
+        ),
+        Worker(
+            temporal_client,
+            task_queue='drift-queue',
+            workflows=[Drift],
+            activities=[
+                activities.load_custom_query,
+                activities.get_reference_data,
+                activities.calculate_drift,
+                activities.export_data_to_postgres,
+            ],
+            tuner=tuner,
+            max_cached_workflows=2,
+            workflow_task_poller_behavior=poller,
+            activity_task_poller_behavior=poller,
+        ),
+        Worker(
+            temporal_client,
+            task_queue='simple_metrics-queue',
+            workflows=[SimpleMetrics],
+            activities=[
+                activities.load_custom_query,
+                activities.calculate_simple_metrics,
+                activities.export_data_to_postgres,
+            ],
+            tuner=tuner,
+            max_cached_workflows=2,
+            workflow_task_poller_behavior=poller,
+            activity_task_poller_behavior=poller,
         ),
         Worker(
             temporal_client,
@@ -168,6 +268,7 @@ async def main():
                 activities.input_gate,
                 activities.mlflow_response_gate,
                 activities.mlflow_content_gate,
+                activities.format_transformed_data,
                 activities.format_prediction,
                 activities.format_default_prediction,
                 activities.get_last_timestamp,
@@ -179,12 +280,10 @@ async def main():
                 activities.export_data_to_postgres,
                 activities.write_metrics,
             ],
-            max_concurrent_workflow_tasks=50,
+            tuner=tuner,
-            max_concurrent_activities=50,
-            max_concurrent_local_activities=50,
             max_cached_workflows=200,
-            workflow_task_poller_behavior=PollerBehaviorAutoscaling(),
+            workflow_task_poller_behavior=poller,
-            activity_task_poller_behavior=PollerBehaviorAutoscaling(),
+            activity_task_poller_behavior=poller,
         ),
     ]
 

laborious/workflows/drift.py

@@ -0,0 +1,107 @@
+from temporalio import workflow
+
+with workflow.unsafe.imports_passed_through():
+    from datetime import timedelta
+    from typing import Any
+
+    from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
+    from sientia_do.temporal.policies import retry_policy
+
+    from laborious.activities.activities import Activities
+
+
+@workflow.defn(name='drift')
+class Drift:
+    @workflow.run
+    async def run(self, input_data: dict[str, Any]):
+        """
+        Execute the drift workflow.
+
+        This method orchestrates the complete drift process by:
+        1. Loading data using the provided custom SQL query
+        2. Preparing prediction configuration and filters
+        3. Delegating to the PredictionProcess workflow for ML operations
+        """
+
+        metadata = {
+            'metadata': {
+                'schedule_name': input_data['schedule_name'],
+                'model_name': input_data['model_name'],
+                'model_id': input_data['model_id'],
+                'workflow_name': 'drift',
+            }
+        }
+
+        print(f'Input data: {input_data}', metadata)
+
+        model_config = input_data['model_config']
+        target_name = model_config['target']
+
+        gathering_query = f"""
+            SELECT *
+            FROM "{input_data['schema']}"."{input_data['source_table_name']}"
+            WHERE
+                model_id = '{input_data['model_id']}' AND
+                timestamp > NOW() - INTERVAL '{input_data['interval']} minutes'
+            ORDER BY timestamp ASC
+        """
+
+        target_data_handler = workflow.start_local_activity_method(
+            Activities.load_custom_query,
+            {
+                **metadata,
+                'query': gathering_query,
+                'datetime_columns': ['timestamp', 'created_at'],
+                'orient': 'records',
+            },
+            retry_policy=retry_policy,
+            start_to_close_timeout=timedelta(seconds=300),
+        )
+
+        reference_data_handler = workflow.start_local_activity_method(
+            Activities.get_reference_data,
+            {**metadata, 'model_name': input_data['model_name']},
+            retry_policy=retry_policy,
+            start_to_close_timeout=timedelta(seconds=300),
+        )
+
+        target_data = await target_data_handler
+        reference_data = await reference_data_handler
+
+        if not target_data:
+            return
+
+        drift_data = await workflow.execute_local_activity_method(
+            Activities.calculate_drift,
+            {
+                **metadata,
+                'target_data': target_data,
+                'reference_data': reference_data,
+                'model_name': input_data['model_name'],
+                'model_id': input_data['model_id'],
+                'target_name': target_name,
+                'drift_metrics': input_data.get(
+                    'drift_metrics', ['kolmogorov_smirnov', 'jensen_shannon', 'wasserstein']
+                ),
+                'chunk_period': input_data.get('chunk_period', 'min'),
+            },
+            retry_policy=retry_policy,
+            start_to_close_timeout=timedelta(seconds=300),
+        )
+
+        if drift_data:
+            await workflow.execute_activity_method(
+                Activities.export_data_to_postgres,
+                {
+                    **metadata,
+                    'data': drift_data,
+                    'schema': input_data['schema'],
+                    'table_name': input_data['target_table_name'],
+                    'timestamp_conversion': {
+                        'column': 'timestamp',
+                        'format': DATETIME_FORMAT_WITH_TZ,
+                    },
+                },
+                retry_policy=retry_policy,
+                start_to_close_timeout=timedelta(seconds=300),
+            )

laborious/workflows/predictions_batch.py

@@ -98,6 +98,7 @@ class PredictionsBatch:
             'data': data,
             'schema': input_data['schema'],
             'table_name': input_data['table_name'],
+            'transform_table_name': input_data['transform_table_name'],
             'model_id': input_data['model_id'],
             'model_name': input_data['model_name'],
             'input_filters': input_data.get('input_filters', {'EMPTY_DATA': {'POLICY': 'STOP'}}),

laborious/workflows/simple_metrics.py

@@ -0,0 +1,95 @@
+from temporalio import workflow
+
+with workflow.unsafe.imports_passed_through():
+    from datetime import timedelta
+    from typing import Any
+
+    from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
+    from sientia_do.temporal.policies import retry_policy
+
+    from laborious.activities.activities import Activities
+
+
+@workflow.defn(name='simple_metrics')
+class SimpleMetrics:
+    @workflow.run
+    async def run(self, input_data: dict[str, Any]):
+        """
+        Execute the simple metrics workflow.
+        """
+        metadata = {
+            'metadata': {
+                'model_id': input_data['model_id'],
+                'model_name': input_data['model_name'],
+                'workflow_name': 'simple_metrics',
+                'schedule_name': input_data['schedule_name'],
+            }
+        }
+
+        model_id = input_data['model_id']
+        interval_minutes = input_data['interval_minutes']
+
+        model_config = input_data['model_config']
+        target_name = model_config['target']
+
+        query = f"""
+            select p."timestamp", p.prediction, ld.value as "target"
+            from "{input_data['schema']}"."{input_data['predictions_table_name']}" p
+            inner join "{input_data['schema']}"."{input_data['data_table_name']}" ld
+                on p."timestamp" = ld."timestamp"
+            where
+                p.model_id = '{model_id}' and
+                p.prediction is not null and
+                ld.variable = '{target_name}' and
+                ld.value is not null and
+                p."timestamp" >= NOW() - INTERVAL '{interval_minutes} minutes'
+            order by
+                p."timestamp" desc;
+        """
+
+        target_data = await workflow.execute_local_activity_method(
+            Activities.load_custom_query,
+            {
+                **metadata,
+                'query': query,
+                'datetime_columns': ['timestamp'],
+                'orient': 'records',
+            },
+            retry_policy=retry_policy,
+            start_to_close_timeout=timedelta(seconds=300),
+        )
+
+        if not target_data:
+            return
+
+        simple_metrics = await workflow.execute_local_activity_method(
+            Activities.calculate_simple_metrics,
+            {
+                **metadata,
+                'model_id': model_id,
+                'target_data': target_data,
+                'metrics': input_data.get('metrics', ['rmse', 'mse', 'mae', 'r2']),
+                'interval_minutes': interval_minutes,
+            },
+            retry_policy=retry_policy,
+            start_to_close_timeout=timedelta(seconds=300),
+        )
+
+        if not simple_metrics:
+            return
+
+        await workflow.execute_activity_method(
+            Activities.export_data_to_postgres,
+            {
+                **metadata,
+                'data': simple_metrics,
+                'schema': input_data['schema'],
+                'table_name': input_data['target_table_name'],
+                'timestamp_conversion': {
+                    'column': 'timestamp',
+                    'format': DATETIME_FORMAT_WITH_TZ,
+                },
+            },
+            retry_policy=retry_policy,
+            start_to_close_timeout=timedelta(seconds=300),
+        )

laborious/workflows/sub_workflows/format_and_export_prediction.py

@@ -50,29 +50,45 @@ class FormatAndExportPrediction:
         Args:
             input_data: Complete configuration for the export workflow
                 Required keys:
+                    - metadata (dict): Workflow execution metadata
                     - path_flag (str | None): Decision path flag for formatting strategy
+                        - None: Normal prediction path with full formatting
+                        - Any other value: Default prediction path for error conditions
                     - data (dict[str, Any]): Prediction data to format and export
                     - prediction_confidence (float): Confidence score for the prediction
                     - timestamp (str): ISO-formatted timestamp for the prediction
                     - model_id (int): Unique identifier for the ML model
                     - model_name (str): Name of the ML model
-                    - model_retention (str): Model retention policy configuration
-                    - comment (str): Operational comment or error description
                     - schema (str): Database schema for data storage
                     - table_name (str): Target table for data persistence
                     - opc_output_config (dict[str, Any]): OPC server export configuration
-                    - prediction_store_policy (str, optional): Data retention policy
+                Optional keys:
+                    - transformed_data (dict[str, Any]): Transformed data to export separately
+                        Only processed when path_flag is None
+                    - transform_table_name (str): Target table for transformed data export
+                        Required if transformed_data is provided
+                    - prediction_store_policy (str): Data retention policy (e.g., 'lts:1', 'erl:2')
+                        Required when path_flag is None
+                    - comment (str): Operational comment or error description
+                        Required when path_flag is not None
 
         Returns:
-            bool: True if the workflow completes successfully, False otherwise
+            None: The workflow completes successfully when all export operations finish
+
+        Note:
+            When transformed_data is provided and path_flag is None, the workflow will:
+            1. Format the transformed data using format_transformed_data
+            2. Export it to a separate table (transform_table_name) asynchronously
+            3. Wait for both prediction and transformed data exports to complete
         """
         metadata = input_data['metadata']
         path_flag = input_data['path_flag']
         data = input_data['data']
+        transformed_data = input_data.get('transformed_data', None)
         prediction_confidence = input_data['prediction_confidence']
 
         if path_flag is None:
-            # proceed with formatting and exporting
+            # Normal prediction path: format prediction data with full metadata
             prediction = await workflow.execute_local_activity_method(
                 Activities.format_prediction,
                 {
@@ -87,8 +103,40 @@ class FormatAndExportPrediction:
                 start_to_close_timeout=timedelta(seconds=60),
             )
 
+            # Optionally format and export transformed data to separate table
+            if transformed_data is not None:
+                transformed = await workflow.execute_local_activity_method(
+                    Activities.format_transformed_data,
+                    {
+                        **metadata,
+                        'data': transformed_data,
+                        'model_id': input_data['model_id'],
+                    },
+                    retry_policy=retry_policy,
+                    start_to_close_timeout=timedelta(seconds=60),
+                )
+
+                write_transformed_handler = workflow.start_activity_method(
+                    Activities.export_data_to_postgres,
+                    {
+                        **metadata,
+                        'schema': input_data['schema'],
+                        'table_name': input_data['transform_table_name'],
+                        'data': transformed,
+                        'timestamp_conversion': {
+                            'column': 'timestamp',
+                            'format': DATETIME_FORMAT_WITH_TZ,
+                        },
+                    },
+                    retry_policy=retry_policy,
+                    start_to_close_timeout=timedelta(seconds=60),
+                )
+
+            else:
+                write_transformed_handler = None
+
         else:
-            # create default prediction
+            # Error path: create default prediction with error indicators
             prediction = await workflow.execute_local_activity_method(
                 Activities.format_default_prediction,
                 {
@@ -102,6 +150,8 @@ class FormatAndExportPrediction:
                 start_to_close_timeout=timedelta(seconds=60),
             )
 
+            write_transformed_handler = None
+
         # write to opc
         prediction, opc_metrics = await workflow.execute_activity_method(
             Activities.write_opc_data,
@@ -115,7 +165,7 @@ class FormatAndExportPrediction:
         )
 
         # write to postgres
-        await workflow.execute_activity_method(
+        prediction_handler = workflow.execute_activity_method(
             Activities.export_data_to_postgres,
             {
                 **metadata,
@@ -128,6 +178,11 @@ class FormatAndExportPrediction:
             start_to_close_timeout=timedelta(seconds=180),
         )
 
+        await prediction_handler
+
+        if write_transformed_handler is not None:
+            await write_transformed_handler
+
         await workflow.execute_activity_method(
             Activities.write_metrics,
             {

laborious/workflows/sub_workflows/prediction_process.py

@@ -82,6 +82,7 @@ class PredictionProcess:
         model_id = input_data['model_id']
         model_name = input_data['model_name']
         model_config = input_data.get('model_config', {})
+        save_transform = input_data.get('save_transform', True)
 
         # Get last timestamp for incremental processing
         last_timestamp = await workflow.execute_local_activity_method(
@@ -199,6 +200,7 @@ class PredictionProcess:
                 'metadata': metadata,
                 'path_flag': path_flag,
                 'data': response_data['content'],
+                'transformed_data': transformed_data if save_transform else None,
                 'prediction_confidence': confidence,
                 'timestamp': last_timestamp,
                 'model_id': model_id,
@@ -207,6 +209,7 @@ class PredictionProcess:
                 'opc_output_config': input_data['opc_output_config'],
                 'schema': input_data['schema'],
                 'table_name': input_data['table_name'],
+                'transform_table_name': input_data['transform_table_name'],
                 'comment': comment,
                 'prediction_store_policy': input_data['prediction_store_policy'],
             },
@@ -248,6 +251,7 @@ class PredictionProcess:
 
         schema = input_data['schema']
         table_name = input_data['table_name']
+        transform_table_name = input_data['transform_table_name']
         model_id = input_data['model_id']
         model_name = input_data['model_name']
         model_config = input_data.get('model_config', {})
@@ -287,6 +291,7 @@ class PredictionProcess:
                     'model_config': model_config,
                     'schema': schema,
                     'table_name': table_name,
+                    'transform_table_name': transform_table_name,
                     'comment': comment,
                     'opc_output_config': input_data['opc_output_config'],
                     'prediction_store_policy': input_data['prediction_store_policy'],

pyproject.toml

@@ -48,6 +48,7 @@ ignore = [
     "S101",  # use of assert (needed for tests)
     "S105",  # possible hardcoded password (false positives)
     "S106",  # possible hardcoded password (false positives)
+    "S608",  # potential sql injection (false positives)
     "N802",  # function name should be lowercase (temporal decorators)
     "N806",  # variable in function should be lowercase
 ]
@@ -152,4 +153,4 @@ directory = "htmlcov"
 
 [tool.bandit]
 exclude_dirs = ["tests", "venv", ".venv"]
-skips = ["B101", "B601"]  # Skip assert and shell injection in controlled environments
+skips = ["B101", "B601", "B608"]  # Skip assert, shell injection, and SQL injection (false positives)

requirements-light.txt

@@ -3,7 +3,7 @@ psycopg2-binary
 sqlalchemy
 asyncua
 redis
-git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.5.2
+git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.6.1
 prometheus-client
 botocore
 boto3

requirements.txt

@@ -3,8 +3,8 @@ psycopg2-binary
 sqlalchemy
 asyncua
 redis
-git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.5.2
+git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.6.1
-git+ssh://git@github.com/Aignosi/sientia-mlops-library.git@0.39.0
+git+ssh://git@github.com/Aignosi/sientia-mlops-library.git@0.40.6
 prometheus-client
 botocore
 boto3

tests.ipynb

@@ -487,6 +487,283 @@
     "except Exception as e:\n",
     "    print(e)\n"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "771ab4ee",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>a</th>\n",
+       "      <th>b</th>\n",
+       "      <th>target</th>\n",
+       "      <th>prediction</th>\n",
+       "      <th>timestamp</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>1</td>\n",
+       "      <td>4</td>\n",
+       "      <td>7</td>\n",
+       "      <td>1</td>\n",
+       "      <td>2025-01-01</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>2</td>\n",
+       "      <td>5</td>\n",
+       "      <td>8</td>\n",
+       "      <td>2</td>\n",
+       "      <td>2025-01-02</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>3</td>\n",
+       "      <td>6</td>\n",
+       "      <td>9</td>\n",
+       "      <td>3</td>\n",
+       "      <td>2025-01-03</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "   a  b  target  prediction   timestamp\n",
+       "0  1  4       7           1  2025-01-01\n",
+       "1  2  5       8           2  2025-01-02\n",
+       "2  3  6       9           3  2025-01-03"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "from pandas import DataFrame, merge\n",
+    "\n",
+    "retrain_dataset = DataFrame({\n",
+    "    'a': {'2025-01-01': 1, '2025-01-02': 2, '2025-01-03': 3},\n",
+    "    'b': {'2025-01-01': 4, '2025-01-02': 5, '2025-01-03': 6},\n",
+    "    'c': {'2025-01-01': 7, '2025-01-02': 8, '2025-01-03': 9},\n",
+    "})\n",
+    "\n",
+    "prediction_data = DataFrame({\n",
+    "    'prediction': {'1': 1, '2': 2, '3': 3},\n",
+    "})\n",
+    "\n",
+    "prediction_data.index = retrain_dataset.index\n",
+    "\n",
+    "prediction_data = merge(\n",
+    "    retrain_dataset, prediction_data, left_index=True, right_index=True, how='left')\n",
+    "\n",
+    "prediction_data.rename(columns={'c': 'target'}, inplace=True)\n",
+    "\n",
+    "prediction_data['timestamp'] = prediction_data.index\n",
+    "\n",
+    "prediction_data.reset_index(drop=True, inplace=True)\n",
+    "\n",
+    "display(prediction_data)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "486b95b3",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>a</th>\n",
+       "      <th>b</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>2025-01-01</th>\n",
+       "      <td>1</td>\n",
+       "      <td>4</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2025-01-02</th>\n",
+       "      <td>2</td>\n",
+       "      <td>5</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2025-01-03</th>\n",
+       "      <td>3</td>\n",
+       "      <td>6</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "            a  b\n",
+       "2025-01-01  1  4\n",
+       "2025-01-02  2  5\n",
+       "2025-01-03  3  6"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/plain": [
+       "{'index': ['2025-01-01', '2025-01-02', '2025-01-03'],\n",
+       " 'columns': ['a', 'b'],\n",
+       " 'data': [[1, 4], [2, 5], [3, 6]]}"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>0</th>\n",
+       "      <th>1</th>\n",
+       "      <th>2</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>index</th>\n",
+       "      <td>2025-01-01</td>\n",
+       "      <td>2025-01-02</td>\n",
+       "      <td>2025-01-03</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>columns</th>\n",
+       "      <td>a</td>\n",
+       "      <td>b</td>\n",
+       "      <td>None</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>data</th>\n",
+       "      <td>[1, 4]</td>\n",
+       "      <td>[2, 5]</td>\n",
+       "      <td>[3, 6]</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "                  0           1           2\n",
+       "index    2025-01-01  2025-01-02  2025-01-03\n",
+       "columns           a           b        None\n",
+       "data         [1, 4]      [2, 5]      [3, 6]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/plain": [
+       "3"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from pandas import DataFrame\n",
+    "\n",
+    "data = DataFrame()\n",
+    "\n",
+    "display(data.to_dict(orient='records'))\n",
+    "\n",
+    "data = DataFrame({\n",
+    "    \"a\": {\"2025-01-01\": 1, \"2025-01-02\": 2, \"2025-01-03\": 3},\n",
+    "    \"b\": {\"2025-01-01\": 4, \"2025-01-02\": 5, \"2025-01-03\": 6},\n",
+    "})\n",
+    "\n",
+    "display(data)\n",
+    "\n",
+    "data_list = data.to_dict('split')\n",
+    "\n",
+    "display(data_list)\n",
+    "\n",
+    "data_rec = DataFrame.from_dict(data_list, orient='index')\n",
+    "\n",
+    "display(data_rec)\n",
+    "\n",
+    "data.shape[0]"
+   ]
   }
  ],
  "metadata": {

tests/conftest.py

@@ -0,0 +1,16 @@
+"""
+Pytest configuration file with global mocks for external dependencies.
+
+This module mocks the 'sientia' module to avoid requiring its installation
+during unit tests. The mock is registered in sys.modules before any test
+imports are executed.
+"""
+
+import sys
+from unittest.mock import MagicMock
+
+# Mock sientia module
+sientia_mock = MagicMock()
+sientia_mock.ModelAnalysis = MagicMock
+sys.modules['sientia'] = sientia_mock
+sys.modules['sientia.ModelAnalysis'] = MagicMock()

tests/laborious/activities/test_gates.py

@@ -546,6 +546,80 @@ async def test_format_prediction_with_timestamp_invalid_policy(gates_activity):
         raise AssertionError('Expected ValueError')
 
 
+@mark.asyncio
+async def test_format_transformed_data_single_row(gates_activity):
+    # Arrange
+    input_data = {
+        **metadata,
+        'data': {
+            'var1': {'2023-05-26 11:12:27': 1.0},
+            'var2': {'2023-05-26 11:12:27': 2.0},
+        },
+        'model_id': 'test_model',
+    }
+
+    # Act
+    result = await gates_activity.format_transformed_data(input_data)
+
+    # Assert
+    assert result['timestamp'] == {0: '2023-05-26 11:12:27', 1: '2023-05-26 11:12:27'}
+    assert result['variable'] == {0: 'var1', 1: 'var2'}
+    assert result['value'] == {0: 1.0, 1: 2.0}
+    assert result['model_id'] == {0: 'test_model', 1: 'test_model'}
+    gates_activity.info.assert_called()
+
+
+@mark.asyncio
+async def test_format_transformed_data_multiple_rows(gates_activity):
+    # Arrange
+    input_data = {
+        **metadata,
+        'data': {
+            'var1': {
+                '2023-05-26 11:12:27': 1.0,
+                '2023-05-26 11:12:28': 2.0,
+            },
+            'var2': {
+                '2023-05-26 11:12:27': 3.0,
+                '2023-05-26 11:12:28': 4.0,
+            },
+        },
+        'model_id': 'test_model',
+    }
+
+    # Act
+    result = await gates_activity.format_transformed_data(input_data)
+
+    # Assert
+    assert len(result['timestamp']) == 4
+    assert len(result['variable']) == 4
+    assert len(result['value']) == 4
+    assert len(result['model_id']) == 4
+    assert all(v == 'test_model' for v in result['model_id'].values())
+    assert set(result['variable'].values()) == {'var1', 'var2'}
+    gates_activity.info.assert_called()
+
+
+@mark.asyncio
+async def test_format_transformed_data_empty_data(gates_activity):
+    # Arrange
+    input_data = {
+        **metadata,
+        'data': {},
+        'model_id': 'test_model',
+    }
+
+    # Act
+    result = await gates_activity.format_transformed_data(input_data)
+
+    # Assert
+    assert result['timestamp'] == {}
+    assert result['variable'] == {}
+    assert result['value'] == {}
+    assert result['model_id'] == {}
+    gates_activity.info.assert_called()
+
+
 @mark.asyncio
 async def test_format_default_prediction(gates_activity):
     # Arrange
@@ -603,6 +677,40 @@ async def test_format_retrain_report(gates_activity):
     assert result['mlflow_experiment_id'] == {0: 'test_mlflow_experiment_id'}
 
 
+@mark.asyncio
+async def test_format_retrain_report_failure(gates_activity):
+    # Arrange
+    input_data = {
+        **metadata,
+        'experiment_response': {
+            'success': False,
+            'timestamp': '2023-05-26 11:12:27',
+            'message': 'failure',
+        },
+        'update_report': {
+            'version': '1.0.0',
+            'mlflow_run_id': 'test_mlflow_run_id',
+            'mlflow_experiment_id': 'test_mlflow_experiment_id',
+        },
+        'model_id': 'test_model',
+        'model_name': 'test_model',
+    }
+
+    # Act
+    result = await gates_activity.format_retrain_report(input_data)
+
+    # Assert
+    assert result['model_id'] == {0: 'test_model'}
+    assert result['model_name'] == {0: 'test_model'}
+    assert result['timestamp'] == {0: '2023-05-26 11:12:27'}
+    assert result['status'] == {0: 'failure'}
+    assert 'version' not in result
+    assert 'mlflow_run_id' not in result
+    assert 'mlflow_experiment_id' not in result
+    gates_activity.info.assert_called()
+    gates_activity.debug.assert_called()
+
+
 @mark.asyncio
 async def test_get_last_timestamp_with_data(gates_activity):
     # Arrange
@@ -742,3 +850,41 @@ async def test_write_metrics(mock_metrics, gates_activity):
             ),
         ]
     )
+
+
+@mark.asyncio
+@patch('laborious.activities.gates.metrics')
+async def test_write_metrics_with_none_opc_response_time(mock_metrics, gates_activity):
+    """Test write_metrics method with None response_time in opc_metrics."""
+    input_data = {
+        **metadata,
+        'prediction': {
+            'prediction': [1],
+            'prediction_confidence': [0.9],
+            'response_time': [0.1],
+        },
+        'opc_metrics': {'server1': {'tag1': 0.1, 'tag2': None}},
+    }
+    await gates_activity.write_metrics(input_data)
+
+    # Verify that metrics for tag1 are emitted
+    gates_activity.emit_metric.assert_any_call(
+        metric_object=mock_metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR,
+        method='observe',
+        tags={
+            'pod_id': gates_activity.pod_id,
+            'model_name': metadata['metadata']['model_name'],
+            'workflow_name': metadata['metadata']['workflow_name'],
+            'opc_server_id': 'server1',
+            'tag': 'tag1',
+        },
+        value=0.1,
+    )
+
+    # Verify that metrics for tag2 (with None response_time) are NOT emitted
+    calls = [
+        c
+        for c in gates_activity.emit_metric.call_args_list
+        if len(c[1].get('tags', {})) > 0 and c[1]['tags'].get('tag') == 'tag2'
+    ]
+    assert len(calls) == 0, 'Metrics should not be emitted for None response_time'

tests/laborious/activities/test_mlflow.py

@@ -76,6 +76,11 @@ def mlflow(mock_minio_repository, mock_mlflow_repository):
     mlflow.send_notification = MagicMock()
     mlflow.emit_metric = AsyncMock()
     mlflow.send_notification_async = AsyncMock()
+    mlflow.error = MagicMock()
+    mlflow.debug = MagicMock()
+    mlflow.info = MagicMock()
+    mlflow.warning = MagicMock()
+    mlflow.critical = MagicMock()
 
     return mlflow
 
@@ -488,3 +493,87 @@ async def test_update_production_model_error(mlflow):
         )
     else:
         raise AssertionError('No exception raised')
+
+
+@mark.asyncio
+@patch('laborious.activities.mlflow.to_datetime')
+async def test_get_reference_data_success(mock_to_datetime, mlflow):
+    # Arrange
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+    }
+
+    # Mock reference data DataFrame
+    mock_reference_data = MagicMock()
+    mock_reference_data.__getitem__.return_value = MagicMock()
+    mock_to_datetime.return_value.dt.strftime.return_value = MagicMock()
+    mock_reference_data.to_dict.return_value = [
+        {'timestamp': '2023-05-26 11:12:27', 'value': 1.0},
+        {'timestamp': '2023-05-26 11:12:28', 'value': 2.0},
+    ]
+
+    mlflow.model_monitoring_repository.load_artifact_dataframe.return_value = mock_reference_data
+
+    # Act
+    result = await mlflow.get_reference_data(input_data)
+
+    # Assert
+    mlflow.model_monitoring_repository.load_artifact_dataframe.assert_called_once_with(
+        model_name='test_model',
+        artifact_path='evaluation_data.csv',
+        metadata=metadata['metadata'],
+    )
+    mock_to_datetime.assert_called_once_with(mock_reference_data.__getitem__.return_value)
+
+    mock_reference_data.to_dict.assert_called_once_with(orient='records')
+    assert result == mock_reference_data.to_dict.return_value
+
+
+@mark.asyncio
+async def test_get_reference_data_not_found(mlflow):
+    # Arrange
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+    }
+
+    mlflow.model_monitoring_repository.load_artifact_dataframe.return_value = None
+
+    # Act
+    result = await mlflow.get_reference_data(input_data)
+
+    # Assert
+    mlflow.model_monitoring_repository.load_artifact_dataframe.assert_called_once_with(
+        model_name='test_model',
+        artifact_path='evaluation_data.csv',
+        metadata=metadata['metadata'],
+    )
+    mlflow.warning.assert_called_once_with(
+        'Reference data not found for model test_model', metadata['metadata']
+    )
+    assert result is None
+
+
+@mark.asyncio
+async def test_get_reference_data_exception(mlflow):
+    # Arrange
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+    }
+
+    mlflow.model_monitoring_repository.load_artifact_dataframe.side_effect = Exception(
+        'Error loading artifact'
+    )
+
+    # Act & Assert
+    with raises(Exception) as e:
+        await mlflow.get_reference_data(input_data)
+
+    assert str(e.value) == 'Error loading artifact'
+    mlflow.model_monitoring_repository.load_artifact_dataframe.assert_called_once_with(
+        model_name='test_model',
+        artifact_path='evaluation_data.csv',
+        metadata=metadata['metadata'],
+    )

tests/laborious/activities/test_model_metrics.py

@@ -0,0 +1,989 @@
+from unittest.mock import ANY, AsyncMock, MagicMock, patch
+
+from pandas import DataFrame
+from pytest import fixture, mark
+from sientia_do.notifications.models import NotificationLevel
+
+from laborious.activities.model_metrics import ModelMetrics
+
+
+@fixture
+def model_metrics_activity():
+    model_metrics = ModelMetrics(
+        logger=MagicMock(),
+        notification_handler=MagicMock(),
+        metrics_controller=AsyncMock(),
+    )
+    model_metrics.error = MagicMock()
+    model_metrics.debug = MagicMock()
+    model_metrics.info = MagicMock()
+    model_metrics.warning = MagicMock()
+    model_metrics.critical = MagicMock()
+    model_metrics.send_notification = MagicMock()
+    model_metrics.send_notification_async = AsyncMock()
+    model_metrics.emit_metric = AsyncMock()
+    model_metrics.get_core_labels = MagicMock(
+        return_value={
+            'pod_id': 'test_pod',
+            'model_name': 'test_model',
+            'workflow_name': 'test_workflow',
+        }
+    )
+    model_metrics.observe_lag = AsyncMock()
+    model_metrics.pod_id = 'test_pod'
+    return model_metrics
+
+
+metadata = {
+    'metadata': {
+        'model_id': 'test_model',
+        'model_name': 'test_model',
+        'workflow_name': 'test_workflow',
+        'schema_name': 'test_schedule',
+    },
+}
+
+
+@mark.asyncio
+async def test_calculate_drift_invalid_chunk_period(model_metrics_activity):
+    # Arrange
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+        'model_id': 'test_model_id',
+        'reference_data': None,
+        'target_data': {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'variable': ['feature1'],
+            'value': [1.0],
+        },
+        'target_name': 'target',
+        'drift_metrics': ['ks_test'],
+        'chunk_period': 'invalid',
+    }
+
+    # Act & Assert
+    try:
+        await model_metrics_activity.calculate_drift(input_data)
+    except ValueError as e:
+        assert str(e) == 'Invalid chunk period: invalid, must be "min" or "s"'
+        model_metrics_activity.error.assert_called_once_with(
+            'Invalid chunk period: invalid', metadata['metadata']
+        )
+    else:
+        raise AssertionError('Expected ValueError')
+
+
+@mark.asyncio
+@patch('laborious.activities.model_metrics.DataFrame')
+@patch('laborious.activities.model_metrics.to_datetime')
+async def test_calculate_drift_with_reference_data(
+    mock_to_datetime, mock_dataframe, model_metrics_activity
+):
+    # Arrange
+    mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27'
+    mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = (
+        '2023-05-26 11:12:27+00:00'
+    )
+
+    mock_drift_df = MagicMock()
+    mock_drift_df.empty = False
+    mock_drift_df.drop.return_value = mock_drift_df
+    mock_drift_df.__getitem__.return_value.isin.return_value = [True]
+    mock_drift_df.__getitem__.return_value = mock_drift_df
+    mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = (
+        '2023-05-26 11:12:27+00:00'
+    )
+    mock_drift_df.rename.return_value = mock_drift_df
+    mock_drift_df.drop_duplicates.return_value = mock_drift_df
+    mock_drift_df.to_dict.return_value = [
+        {
+            'method': 'ks_test',
+            'value': 0.5,
+            'feature': 'feature1',
+            'timestamp': '2023-05-26 11:12:27+00:00',
+            'model_id': 'test_model_id',
+            'accurate': True,
+        }
+    ]
+
+    model_metrics_activity.get_drift_metrics = AsyncMock(return_value=mock_drift_df)
+
+    reference_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'target': [1.0],
+            'feature1': [1.0],
+        }
+    )
+
+    mock_target_df = MagicMock()
+    mock_target_df.pivot.return_value = mock_target_df
+    mock_target_df.index = ['2023-05-26 11:12:27']
+    mock_target_df.reset_index.return_value = mock_target_df
+    mock_target_df.dropna.return_value = mock_target_df
+    mock_target_df.__getitem__.return_value.apply.return_value = ['2023-05-26 11:12:27']
+    mock_target_df.drop.return_value.columns = ['feature1']
+    mock_dataframe.return_value = mock_target_df
+
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+        'model_id': 'test_model_id',
+        'reference_data': reference_data.to_dict(),
+        'target_data': {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'variable': ['feature1'],
+            'value': [1.0],
+        },
+        'target_name': 'target',
+        'drift_metrics': ['ks_test'],
+        'chunk_period': 'min',
+    }
+
+    # Act
+    result = await model_metrics_activity.calculate_drift(input_data)
+
+    # Assert
+    assert isinstance(result, list)
+    assert result == mock_drift_df.to_dict.return_value  # type: ignore[comparison-overlap]
+    model_metrics_activity.info.assert_called()
+    model_metrics_activity.get_drift_metrics.assert_called_once()
+    # Verify transformations were called
+    mock_drift_df.drop.assert_called_once_with(columns=['p_value'], inplace=True)
+    mock_drift_df.__getitem__.assert_called()
+    mock_drift_df.rename.assert_called_once_with(
+        columns={'metric': 'method', 'statistic': 'value'}, inplace=True
+    )
+    mock_drift_df.drop_duplicates.assert_called_once_with(
+        subset=['timestamp', 'method', 'feature'], keep='first', inplace=True
+    )
+    mock_drift_df.to_dict.assert_called_once_with(orient='records')
+
+
+@mark.asyncio
+@patch('laborious.activities.model_metrics.DataFrame')
+@patch('laborious.activities.model_metrics.to_datetime')
+async def test_calculate_drift_without_reference_data(
+    mock_to_datetime, mock_dataframe, model_metrics_activity
+):
+    # Arrange
+    mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27'
+    mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = (
+        '2023-05-26 11:12:27+00:00'
+    )
+
+    mock_drift_df = MagicMock()
+    mock_drift_df.empty = False
+    mock_drift_df.drop.return_value = mock_drift_df
+    mock_drift_df.__getitem__.return_value.isin.return_value = [True]
+    mock_drift_df.__getitem__.return_value = mock_drift_df
+    mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = (
+        '2023-05-26 11:12:27+00:00'
+    )
+    mock_drift_df.rename.return_value = mock_drift_df
+    mock_drift_df.drop_duplicates.return_value = mock_drift_df
+    mock_drift_df.to_dict.return_value = [
+        {
+            'method': 'ks_test',
+            'value': 0.5,
+            'feature': 'feature1',
+            'timestamp': '2023-05-26 11:12:27+00:00',
+            'model_id': 'test_model_id',
+            'accurate': False,
+        }
+    ]
+
+    model_metrics_activity.get_drift_metrics = AsyncMock(return_value=mock_drift_df)
+
+    target_data_dict = {
+        'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'],
+        'variable': ['feature1', 'feature1', 'feature1'],
+        'value': [1.0, 2.0, 3.0],
+    }
+
+    mock_target_df = MagicMock()
+    mock_target_df.pivot.return_value = mock_target_df
+    mock_target_df.index = ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29']
+    mock_target_df.reset_index.return_value = mock_target_df
+    mock_target_df.dropna.return_value = mock_target_df
+    mock_target_df.sort_values.return_value = mock_target_df
+    mock_target_df.head.return_value = DataFrame(
+        {'timestamp': ['2023-05-26 11:12:27'], 'feature1': [1.0]}
+    )
+    mock_target_df.__getitem__.return_value.apply.return_value = [
+        '2023-05-26 11:12:27',
+        '2023-05-26 11:12:28',
+        '2023-05-26 11:12:29',
+    ]
+    mock_target_df.drop.return_value.columns = ['feature1']
+    mock_dataframe.return_value = mock_target_df
+    mock_dataframe.side_effect = lambda x=None: mock_target_df if x is not None else mock_target_df
+
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+        'model_id': 'test_model_id',
+        'reference_data': None,
+        'target_data': target_data_dict,
+        'target_name': 'target',
+        'drift_metrics': ['ks_test'],
+        'chunk_period': 's',
+    }
+
+    # Act
+    result = await model_metrics_activity.calculate_drift(input_data)
+
+    # Assert
+    assert isinstance(result, list)
+    assert result == mock_drift_df.to_dict.return_value
+    model_metrics_activity.warning.assert_called()
+    model_metrics_activity.send_notification_async.assert_called_once_with(
+        metadata=metadata['metadata'],
+        notification_id='MODEL_METRICS_REFERENCE_DATA_WARNING',
+        message='Using 30% first rows of target data as reference data',
+        block='model_metrics',
+        level=NotificationLevel.WARNING,
+        attachment_content=ANY,
+    )
+    # Verify transformations were called
+    mock_drift_df.drop.assert_called_once_with(columns=['p_value'], inplace=True)
+    mock_drift_df.__getitem__.assert_called()
+    mock_drift_df.rename.assert_called_once_with(
+        columns={'metric': 'method', 'statistic': 'value'}, inplace=True
+    )
+    mock_drift_df.drop_duplicates.assert_called_once_with(
+        subset=['timestamp', 'method', 'feature'], keep='first', inplace=True
+    )
+    mock_drift_df.to_dict.assert_called_once_with(orient='records')
+
+
+@mark.asyncio
+@patch('laborious.activities.model_metrics.DataFrame')
+@patch('laborious.activities.model_metrics.to_datetime')
+async def test_calculate_drift_empty_drift_df(
+    mock_to_datetime, mock_dataframe, model_metrics_activity
+):
+    # Arrange
+    mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27'
+
+    model_metrics_activity.get_drift_metrics = AsyncMock(return_value=DataFrame())
+
+    reference_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'target': [1.0],
+            'feature1': [1.0],
+        }
+    )
+
+    mock_target_df = MagicMock()
+    mock_target_df.pivot.return_value = mock_target_df
+    mock_target_df.index = ['2023-05-26 11:12:27']
+    mock_target_df.reset_index.return_value = mock_target_df
+    mock_target_df.dropna.return_value = mock_target_df
+    mock_target_df.__getitem__.return_value.apply.return_value = ['2023-05-26 11:12:27']
+    mock_target_df.drop.return_value.columns = ['feature1']
+    mock_dataframe.return_value = mock_target_df
+
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+        'model_id': 'test_model_id',
+        'reference_data': reference_data.to_dict(),
+        'target_data': {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'variable': ['feature1'],
+            'value': [1.0],
+        },
+        'target_name': 'target',
+        'drift_metrics': ['ks_test'],
+        'chunk_period': 'min',
+    }
+
+    # Act
+    result = await model_metrics_activity.calculate_drift(input_data)
+
+    # Assert
+    assert result == []
+    model_metrics_activity.warning.assert_called_with(
+        'No drift metrics found', metadata['metadata']
+    )
+
+
+@mark.asyncio
+@patch('laborious.activities.model_metrics.DataFrame')
+@patch('laborious.activities.model_metrics.to_datetime')
+async def test_calculate_drift_empty_after_timestamp_filter(
+    mock_to_datetime, mock_dataframe, model_metrics_activity
+):
+    # Arrange
+    mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27'
+
+    mock_drift_df = MagicMock()
+    mock_drift_df.empty = False
+    mock_drift_df.drop.return_value = mock_drift_df
+
+    # Set up __getitem__ to handle filtering - timestamp access returns series with isin=False
+    # and filtering returns empty DataFrame
+    mock_timestamp_series = MagicMock()
+    mock_timestamp_series.isin.return_value = [False]
+    mock_empty_df = MagicMock()
+    mock_empty_df.empty = True
+
+    def getitem_side_effect(key):
+        if key == 'timestamp':
+            return mock_timestamp_series
+        else:
+            # This is the filtering operation - return empty DataFrame
+            return mock_empty_df
+
+    mock_drift_df.__getitem__.side_effect = getitem_side_effect
+
+    model_metrics_activity.get_drift_metrics = AsyncMock(return_value=mock_drift_df)
+
+    reference_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'target': [1.0],
+            'feature1': [1.0],
+        }
+    )
+
+    mock_target_df = MagicMock()
+    mock_target_df.pivot.return_value = mock_target_df
+    mock_target_df.index = ['2023-05-26 11:12:27']
+    mock_target_df.reset_index.return_value = mock_target_df
+    mock_target_df.dropna.return_value = mock_target_df
+    mock_target_df.__getitem__.return_value.apply.return_value = ['2023-05-26 11:12:27']
+    mock_target_df.drop.return_value.columns = ['feature1']
+    mock_dataframe.return_value = mock_target_df
+
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+        'model_id': 'test_model_id',
+        'reference_data': reference_data.to_dict(),
+        'target_data': {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'variable': ['feature1'],
+            'value': [1.0],
+        },
+        'target_name': 'target',
+        'drift_metrics': ['ks_test'],
+        'chunk_period': 'min',
+    }
+
+    # Act
+    result = await model_metrics_activity.calculate_drift(input_data)
+
+    # Assert
+    assert result == []
+    model_metrics_activity.warning.assert_called_with(
+        'No drift metrics found after dropping rows where timestamp is not in target data',
+        metadata['metadata'],
+    )
+    # Verify transformations were called
+    mock_drift_df.drop.assert_called_once_with(columns=['p_value'], inplace=True)
+    mock_drift_df.__getitem__.assert_called()
+
+
+@mark.asyncio
+@patch('laborious.activities.model_metrics.DataFrame')
+@patch('laborious.activities.model_metrics.to_datetime')
+async def test_calculate_drift_success_min(
+    mock_to_datetime, mock_dataframe, model_metrics_activity
+):
+    # Arrange
+    mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27'
+    mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = (
+        '2023-05-26 11:12:27+00:00'
+    )
+
+    mock_drift_df = MagicMock()
+    mock_drift_df.empty = False
+    mock_drift_df.drop.return_value = mock_drift_df
+    mock_drift_df.__getitem__.return_value.isin.return_value = [True]
+    mock_drift_df.__getitem__.return_value = mock_drift_df
+    mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = (
+        '2023-05-26 11:12:27+00:00'
+    )
+    mock_drift_df.rename.return_value = mock_drift_df
+    mock_drift_df.drop_duplicates.return_value = mock_drift_df
+    mock_drift_df.to_dict.return_value = [
+        {
+            'method': 'ks_test',
+            'value': 0.5,
+            'feature': 'feature1',
+            'timestamp': '2023-05-26 11:12:27+00:00',
+            'model_id': 'test_model_id',
+            'accurate': True,
+        }
+    ]
+
+    model_metrics_activity.get_drift_metrics = AsyncMock(return_value=mock_drift_df)
+
+    reference_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'target': [1.0],
+            'feature1': [1.0],
+        }
+    )
+
+    mock_target_df = MagicMock()
+    mock_target_df.pivot.return_value = mock_target_df
+    mock_target_df.index = ['2023-05-26 11:12:27']
+    mock_target_df.reset_index.return_value = mock_target_df
+    mock_target_df.dropna.return_value = mock_target_df
+    mock_target_df.__getitem__.return_value.isin.return_value = [True]
+    mock_target_df.drop.return_value.columns = ['feature1']
+    mock_dataframe.return_value = mock_target_df
+
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+        'model_id': 'test_model_id',
+        'reference_data': reference_data.to_dict(),
+        'target_data': {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'variable': ['feature1'],
+            'value': [1.0],
+        },
+        'target_name': 'target',
+        'drift_metrics': ['ks_test'],
+        'chunk_period': 'min',
+    }
+
+    # Act
+    result = await model_metrics_activity.calculate_drift(input_data)
+
+    # Assert
+    assert isinstance(result, list)
+    assert result == mock_drift_df.to_dict.return_value  # type: ignore[comparison-overlap]
+    model_metrics_activity.info.assert_called()
+    model_metrics_activity.get_drift_metrics.assert_called_once()
+    # Verify transformations were called
+    mock_drift_df.drop.assert_called_once_with(columns=['p_value'], inplace=True)
+    mock_drift_df.__getitem__.assert_called()
+    mock_drift_df.rename.assert_called_once_with(
+        columns={'metric': 'method', 'statistic': 'value'}, inplace=True
+    )
+    mock_drift_df.drop_duplicates.assert_called_once_with(
+        subset=['timestamp', 'method', 'feature'], keep='first', inplace=True
+    )
+    mock_drift_df.to_dict.assert_called_once_with(orient='records')
+
+
+@mark.asyncio
+@patch('laborious.activities.model_metrics.DataFrame')
+@patch('laborious.activities.model_metrics.to_datetime')
+async def test_calculate_drift_success_s(mock_to_datetime, mock_dataframe, model_metrics_activity):
+    # Arrange
+    mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27'
+    mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = (
+        '2023-05-26 11:12:27+00:00'
+    )
+
+    mock_drift_df = MagicMock()
+    mock_drift_df.empty = False
+    mock_drift_df.drop.return_value = mock_drift_df
+    mock_drift_df.__getitem__.return_value.isin.return_value = [True]
+    mock_drift_df.__getitem__.return_value = mock_drift_df
+    mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = (
+        '2023-05-26 11:12:27+00:00'
+    )
+    mock_drift_df.rename.return_value = mock_drift_df
+    mock_drift_df.drop_duplicates.return_value = mock_drift_df
+    mock_drift_df.to_dict.return_value = [
+        {
+            'method': 'ks_test',
+            'value': 0.5,
+            'feature': 'feature1',
+            'timestamp': '2023-05-26 11:12:27+00:00',
+            'model_id': 'test_model_id',
+            'accurate': True,
+        }
+    ]
+
+    model_metrics_activity.get_drift_metrics = AsyncMock(return_value=mock_drift_df)
+
+    reference_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'target': [1.0],
+            'feature1': [1.0],
+        }
+    )
+
+    mock_target_df = MagicMock()
+    mock_target_df.pivot.return_value = mock_target_df
+    mock_target_df.index = ['2023-05-26 11:12:27']
+    mock_target_df.reset_index.return_value = mock_target_df
+    mock_target_df.dropna.return_value = mock_target_df
+    mock_target_df.__getitem__.return_value.isin.return_value = [True]
+    mock_target_df.drop.return_value.columns = ['feature1']
+    mock_dataframe.return_value = mock_target_df
+
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+        'model_id': 'test_model_id',
+        'reference_data': reference_data.to_dict(),
+        'target_data': {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'variable': ['feature1'],
+            'value': [1.0],
+        },
+        'target_name': 'target',
+        'drift_metrics': ['ks_test'],
+        'chunk_period': 's',
+    }
+
+    # Act
+    result = await model_metrics_activity.calculate_drift(input_data)
+
+    # Assert
+    assert isinstance(result, list)
+    assert result == mock_drift_df.to_dict.return_value  # type: ignore[comparison-overlap]
+    model_metrics_activity.info.assert_called()
+    model_metrics_activity.get_drift_metrics.assert_called_once()
+    # Verify transformations were called
+    mock_drift_df.drop.assert_called_once_with(columns=['p_value'], inplace=True)
+    mock_drift_df.__getitem__.assert_called()
+    mock_drift_df.rename.assert_called_once_with(
+        columns={'metric': 'method', 'statistic': 'value'}, inplace=True
+    )
+    mock_drift_df.drop_duplicates.assert_called_once_with(
+        subset=['timestamp', 'method', 'feature'], keep='first', inplace=True
+    )
+    mock_drift_df.to_dict.assert_called_once_with(orient='records')
+
+
+@mark.asyncio
+@patch('laborious.activities.model_metrics.DataFrame')
+@patch('laborious.activities.model_metrics.to_datetime')
+async def test_calculate_drift_get_drift_metrics_error(
+    mock_to_datetime, mock_dataframe, model_metrics_activity
+):
+    # Arrange
+    mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27'
+
+    model_metrics_activity.get_drift_metrics = AsyncMock(
+        side_effect=Exception('Get drift metrics error')
+    )
+
+    reference_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'target': [1.0],
+            'feature1': [1.0],
+        }
+    )
+
+    mock_target_df = MagicMock()
+    mock_target_df.pivot.return_value = mock_target_df
+    mock_target_df.index = ['2023-05-26 11:12:27']
+    mock_target_df.reset_index.return_value = mock_target_df
+    mock_target_df.dropna.return_value = mock_target_df
+    mock_target_df.__getitem__.return_value.apply.return_value = ['2023-05-26 11:12:27']
+    mock_target_df.drop.return_value.columns = ['feature1']
+    mock_dataframe.return_value = mock_target_df
+
+    input_data = {
+        **metadata,
+        'model_name': 'test_model',
+        'model_id': 'test_model_id',
+        'reference_data': reference_data.to_dict(),
+        'target_data': {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'variable': ['feature1'],
+            'value': [1.0],
+        },
+        'target_name': 'target',
+        'drift_metrics': ['ks_test'],
+        'chunk_period': 'min',
+    }
+
+    # Act
+    result = await model_metrics_activity.calculate_drift(input_data)
+
+    # Assert
+    assert result == []
+    model_metrics_activity.error.assert_called_once_with(
+        'Error getting drift metrics: Get drift metrics error', metadata['metadata']
+    )
+    model_metrics_activity.send_notification_async.assert_called_once_with(
+        metadata=metadata['metadata'],
+        notification_id='MODEL_METRICS_GET_DRIFT_METRICS_ERROR',
+        message='Error getting drift metrics: Get drift metrics error',
+        block='model_metrics',
+        level=NotificationLevel.ERROR,
+        attachment_content=ANY,
+    )
+
+
+@mark.asyncio
+@patch('laborious.activities.model_metrics.to_datetime')
+@patch('laborious.activities.model_metrics.time.time')
+@patch('laborious.activities.model_metrics.ModelAnalysis')
+@patch('laborious.activities.model_metrics.metrics')
+async def test_get_drift_metrics_success(
+    mock_metrics, mock_model_analysis, mock_time, mock_to_datetime, model_metrics_activity
+):
+    # Arrange
+    mock_time.return_value = 1000.0
+
+    mock_drift_df = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'metric': ['ks_test'],
+            'statistic': [0.5],
+            'feature': ['feature1'],
+        }
+    )
+
+    mock_model_analysis.return_value.detect_univariate_drift.return_value = MagicMock()
+    mock_model_analysis.return_value.detect_multivariate_drift.return_value = MagicMock()
+    mock_model_analysis.return_value.get_drift_metrics_dataframe.return_value = mock_drift_df
+
+    reference_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'target': [1.0],
+            'feature1': [1.0],
+        }
+    )
+
+    target_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'target': [1.0],
+            'feature1': [1.0],
+        }
+    )
+
+    reference_columns = reference_data.drop(
+        columns=['target', 'timestamp'], errors='ignore'
+    ).columns
+
+    # Act
+    result = await model_metrics_activity.get_drift_metrics(
+        reference_data=reference_data,
+        target_data=target_data,
+        target_name='target',
+        reference_columns=reference_columns,
+        drift_metrics=['ks_test'],
+        chunk_period='min',
+        metadata=metadata['metadata'],
+    )
+
+    # Assert
+    assert isinstance(result, DataFrame)
+    model_metrics_activity.debug.assert_called()
+    model_metrics_activity.observe_lag.assert_called()
+    model_metrics_activity.emit_metric.assert_called()
+
+
+@mark.asyncio
+@patch('laborious.activities.model_metrics.to_datetime')
+@patch('laborious.activities.model_metrics.time.time')
+@patch('laborious.activities.model_metrics.ModelAnalysis')
+@patch('laborious.activities.model_metrics.metrics')
+async def test_get_drift_metrics_univariate_error(
+    mock_metrics, mock_model_analysis, mock_time, mock_to_datetime, model_metrics_activity
+):
+    # Arrange
+    mock_time.return_value = 1000.0
+
+    mock_model_analysis.return_value.detect_univariate_drift.side_effect = Exception(
+        'Univariate drift error'
+    )
+
+    reference_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'target': [1.0],
+            'feature1': [1.0],
+        }
+    )
+
+    target_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'target': [1.0],
+            'feature1': [1.0],
+        }
+    )
+
+    reference_columns = reference_data.drop(
+        columns=['target', 'timestamp'], errors='ignore'
+    ).columns
+
+    # Act & Assert
+    try:
+        await model_metrics_activity.get_drift_metrics(
+            reference_data=reference_data,
+            target_data=target_data,
+            target_name='target',
+            reference_columns=reference_columns,
+            drift_metrics=['ks_test'],
+            chunk_period='min',
+            metadata=metadata['metadata'],
+        )
+    except Exception as e:
+        assert str(e) == 'Univariate drift error'
+        model_metrics_activity.error.assert_called_once_with(
+            'Error detecting univariate drift: Univariate drift error', metadata['metadata']
+        )
+        model_metrics_activity.emit_metric.assert_called_with(
+            metric_object=mock_metrics.MODEL_ANALYZE_ERROR_COUNT, tags=ANY
+        )
+    else:
+        raise AssertionError('Expected Exception')
+
+
+@mark.asyncio
+async def test_calculate_simple_metrics_success_all_metrics(model_metrics_activity):
+    # Arrange
+    target_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'],
+            'target': [1.0, 2.0, 3.0],
+            'prediction': [1.1, 2.1, 2.9],
+        }
+    )
+
+    input_data = {
+        **metadata,
+        'model_id': 'test_model_id',
+        'target_data': target_data.to_dict(),
+        'metrics': ['rmse', 'mse', 'mae', 'r2'],
+        'interval_minutes': 5,
+    }
+
+    # Act
+    result = DataFrame(await model_metrics_activity.calculate_simple_metrics(input_data))
+
+    # Assert
+    assert len(result['metric']) == 4
+    assert 'rmse' in result['metric'].values
+    assert 'mse' in result['metric'].values
+    assert 'mae' in result['metric'].values
+    assert 'r2' in result['metric'].values
+    assert all(model_id == 'test_model_id' for model_id in result['model_id'].values)
+    assert all(timestamp == '2023-05-26 11:12:29' for timestamp in result['timestamp'].values)
+    assert all(data_size == 3 for data_size in result['data_size'].values)
+    assert all(interval_minutes == 5 for interval_minutes in result['interval_minutes'].values)
+    model_metrics_activity.info.assert_called_once_with(
+        "Calculating simple metrics for model test_model_id: ['rmse', 'mse', 'mae', 'r2']",
+        metadata['metadata'],
+    )
+    model_metrics_activity.debug.assert_called_once()
+
+
+@mark.asyncio
+async def test_calculate_simple_metrics_success_rmse_only(model_metrics_activity):
+    # Arrange
+    target_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
+            'target': [1.0, 2.0],
+            'prediction': [1.1, 2.1],
+        }
+    )
+
+    input_data = {
+        **metadata,
+        'model_id': 'test_model_id',
+        'target_data': target_data.to_dict(),
+        'metrics': ['rmse'],
+        'interval_minutes': 5,
+    }
+
+    # Act
+    result = DataFrame(await model_metrics_activity.calculate_simple_metrics(input_data))
+
+    # Assert
+    assert len(result['metric']) == 1
+    assert result['metric'].values[0] == 'rmse'
+    assert result['model_id'].values[0] == 'test_model_id'
+    assert result['timestamp'].values[0] == '2023-05-26 11:12:28'
+    assert result['data_size'].values[0] == 2
+    assert result['interval_minutes'].values[0] == 5
+    model_metrics_activity.info.assert_called_once_with(
+        "Calculating simple metrics for model test_model_id: ['rmse']", metadata['metadata']
+    )
+
+
+@mark.asyncio
+async def test_calculate_simple_metrics_success_mse_only(model_metrics_activity):
+    # Arrange
+    target_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
+            'target': [1.0, 2.0],
+            'prediction': [1.1, 2.1],
+        }
+    )
+
+    input_data = {
+        **metadata,
+        'model_id': 'test_model_id',
+        'target_data': target_data.to_dict(),
+        'metrics': ['mse'],
+        'interval_minutes': 5,
+    }
+
+    # Act
+    result = DataFrame(await model_metrics_activity.calculate_simple_metrics(input_data))
+
+    # Assert
+    assert len(result['metric']) == 1
+    assert result['metric'].values[0] == 'mse'
+    assert result['model_id'].values[0] == 'test_model_id'
+    assert result['timestamp'].values[0] == '2023-05-26 11:12:28'
+    assert result['data_size'].values[0] == 2
+    assert result['interval_minutes'].values[0] == 5
+    model_metrics_activity.info.assert_called_once_with(
+        "Calculating simple metrics for model test_model_id: ['mse']", metadata['metadata']
+    )
+
+
+@mark.asyncio
+async def test_calculate_simple_metrics_success_mae_only(model_metrics_activity):
+    # Arrange
+    target_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
+            'target': [1.0, 2.0],
+            'prediction': [1.1, 2.1],
+        }
+    )
+
+    input_data = {
+        **metadata,
+        'model_id': 'test_model_id',
+        'target_data': target_data.to_dict(),
+        'metrics': ['mae'],
+        'interval_minutes': 5,
+    }
+
+    # Act
+    result = DataFrame(await model_metrics_activity.calculate_simple_metrics(input_data))
+
+    # Assert
+    assert len(result['metric']) == 1
+    assert result['metric'].values[0] == 'mae'
+    assert result['model_id'].values[0] == 'test_model_id'
+    assert result['timestamp'].values[0] == '2023-05-26 11:12:28'
+    assert result['data_size'].values[0] == 2
+    assert result['interval_minutes'].values[0] == 5
+    model_metrics_activity.info.assert_called_once_with(
+        "Calculating simple metrics for model test_model_id: ['mae']", metadata['metadata']
+    )
+
+
+@mark.asyncio
+async def test_calculate_simple_metrics_success_r2_only(model_metrics_activity):
+    # Arrange
+    target_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
+            'target': [1.0, 2.0],
+            'prediction': [1.1, 2.1],
+        }
+    )
+
+    input_data = {
+        **metadata,
+        'model_id': 'test_model_id',
+        'target_data': target_data.to_dict(),
+        'metrics': ['r2'],
+        'interval_minutes': 5,
+    }
+
+    # Act
+    result = DataFrame(await model_metrics_activity.calculate_simple_metrics(input_data))
+
+    # Assert
+    assert len(result['metric']) == 1
+    assert result['metric'].values[0] == 'r2'
+    assert result['model_id'].values[0] == 'test_model_id'
+    assert result['timestamp'].values[0] == '2023-05-26 11:12:28'
+    assert result['data_size'].values[0] == 2
+    assert result['interval_minutes'].values[0] == 5
+    model_metrics_activity.info.assert_called_once_with(
+        "Calculating simple metrics for model test_model_id: ['r2']", metadata['metadata']
+    )
+
+
+@mark.asyncio
+async def test_calculate_simple_metrics_r2_zero_ss_tot(model_metrics_activity):
+    # Arrange
+    # All target values are the same, so ss_tot will be 0
+    target_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
+            'target': [1.0, 1.0],
+            'prediction': [1.1, 1.1],
+        }
+    )
+
+    input_data = {
+        **metadata,
+        'model_id': 'test_model_id',
+        'target_data': target_data.to_dict(),
+        'metrics': ['r2'],
+        'interval_minutes': 5,
+    }
+
+    # Act
+    result = DataFrame(await model_metrics_activity.calculate_simple_metrics(input_data))
+
+    # Assert
+    assert len(result['metric']) == 1
+    assert result['metric'].values[0] == 'r2'
+    assert result['value'].values[0] == 0.0  # Should return 0.0 when ss_tot == 0
+    assert result['model_id'].values[0] == 'test_model_id'
+    assert result['timestamp'].values[0] == '2023-05-26 11:12:28'
+    assert result['data_size'].values[0] == 2
+    assert result['interval_minutes'].values[0] == 5
+    model_metrics_activity.info.assert_called_once_with(
+        "Calculating simple metrics for model test_model_id: ['r2']", metadata['metadata']
+    )
+
+
+@mark.asyncio
+async def test_calculate_simple_metrics_success_multiple_metrics_subset(model_metrics_activity):
+    # Arrange
+    target_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'],
+            'target': [1.0, 2.0, 3.0],
+            'prediction': [1.1, 2.1, 2.9],
+        }
+    )
+
+    input_data = {
+        **metadata,
+        'model_id': 'test_model_id',
+        'target_data': target_data.to_dict(),
+        'metrics': ['rmse', 'mae'],
+        'interval_minutes': 5,
+    }
+
+    # Act
+    result = DataFrame(await model_metrics_activity.calculate_simple_metrics(input_data))
+
+    # Assert
+    assert len(result['metric']) == 2
+    assert 'rmse' in result['metric'].values
+    assert 'mae' in result['metric'].values
+    assert all(model_id == 'test_model_id' for model_id in result['model_id'].values)
+    assert all(timestamp == '2023-05-26 11:12:29' for timestamp in result['timestamp'].values)
+    assert all(data_size == 3 for data_size in result['data_size'].values)
+    assert all(interval_minutes == 5 for interval_minutes in result['interval_minutes'].values)
+    model_metrics_activity.info.assert_called_once_with(
+        "Calculating simple metrics for model test_model_id: ['rmse', 'mae']", metadata['metadata']
+    )

tests/laborious/utils/repository/test_model_repository.py

@@ -190,6 +190,30 @@ def test_get_model_params(mlflow, mlflow_repository):
     assert output == mlflow.get_run.return_value.data.params
 
 
+def test_check_artifact_exists_true(mlflow_repository):
+    artifact = MagicMock(path='test_artifact')
+    mlflow_repository.client.list_artifacts.return_value = [artifact]
+
+    result = mlflow_repository.check_artifact_exists(
+        'run_id', 'test_artifact', metadata['metadata']
+    )
+
+    assert result is True
+    mlflow_repository.client.list_artifacts.assert_called_once_with('run_id')
+
+
+def test_check_artifact_exists_false(mlflow_repository):
+    artifact = MagicMock(path='other_artifact')
+    mlflow_repository.client.list_artifacts.return_value = [artifact]
+
+    result = mlflow_repository.check_artifact_exists(
+        'run_id', 'test_artifact', metadata['metadata']
+    )
+
+    assert result is False
+    mlflow_repository.client.list_artifacts.assert_called_once_with('run_id')
+
+
 @pytest.mark.asyncio
 @patch('laborious.utils.repository.model_repository.path')
 @patch('laborious.utils.repository.model_repository.rmtree')
@@ -275,6 +299,70 @@ async def test_download_artifacts_error(makedirs, rmtree, path, mlflow_repositor
         mlflow_repository.observe_lag.assert_not_called()
 
 
+@pytest.mark.asyncio
+@patch('laborious.utils.repository.model_repository.mlflow')
+@patch('laborious.utils.repository.model_repository.pd')
+@patch('laborious.utils.repository.model_repository.StringIO')
+async def test_load_artifact_dataframe_success(_stringio, pd, mlflow, mlflow_repository):
+    mlflow_repository.get_model_run_id = MagicMock(return_value='run_id')
+    mlflow_repository.check_artifact_exists = MagicMock(return_value=True)
+
+    mlflow.artifacts.load_text.return_value = 'col1,col2\n1,2\n3,4'
+
+    result = await mlflow_repository.load_artifact_dataframe(
+        'model_name', 'artifact_path', metadata['metadata']
+    )
+
+    mlflow_repository.get_model_run_id.assert_called_once_with(
+        model_name='model_name', stage='Production'
+    )
+    mlflow_repository.check_artifact_exists.assert_called_once_with(
+        'run_id', 'artifact_path', metadata['metadata']
+    )
+    mlflow.artifacts.load_text.assert_called_once_with('runs:/run_id/artifact_path')
+    assert result == pd.read_csv.return_value
+    mlflow_repository.observe_lag.assert_called_once_with(ANY, metrics.MODEL_READ_LAG, ANY)
+    mlflow_repository.emit_metric.assert_called_once_with(
+        metric_object=metrics.MODEL_READ_COUNT, tags=ANY
+    )
+
+
+@pytest.mark.asyncio
+async def test_load_artifact_dataframe_not_exists(mlflow_repository):
+    mlflow_repository.get_model_run_id = MagicMock(return_value='run_id')
+    mlflow_repository.check_artifact_exists = MagicMock(return_value=False)
+
+    result = await mlflow_repository.load_artifact_dataframe(
+        'model_name', 'artifact_path', metadata['metadata']
+    )
+
+    assert result is None
+    mlflow_repository.get_model_run_id.assert_called_once_with(
+        model_name='model_name', stage='Production'
+    )
+    mlflow_repository.check_artifact_exists.assert_called_once_with(
+        'run_id', 'artifact_path', metadata['metadata']
+    )
+
+
+@pytest.mark.asyncio
+@patch('laborious.utils.repository.model_repository.mlflow')
+async def test_load_artifact_dataframe_error(mlflow, mlflow_repository):
+    mlflow_repository.get_model_run_id = MagicMock(return_value='run_id')
+    mlflow_repository.check_artifact_exists = MagicMock(return_value=True)
+    mlflow.artifacts.load_text.side_effect = ValueError('error')
+
+    with pytest.raises(ValueError):
+        await mlflow_repository.load_artifact_dataframe(
+            'model_name', 'artifact_path', metadata['metadata']
+        )
+
+    mlflow_repository.emit_metric.assert_called_once_with(
+        metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=ANY
+    )
+    mlflow_repository.observe_lag.assert_not_called()
+
+
 def test_get_experiment_error(mlflow, mlflow_repository):
     mlflow.get_experiment_by_name.return_value = None
 
@@ -719,6 +807,7 @@ async def test_fit_models_not_df_target_name_none_and_not_in_model(
     mlflow_repository.detect_and_parse_datetime_index = MagicMock(
         return_value=MagicMock(drop_duplicates=MagicMock(return_value=MagicMock(columns=[])))
     )
+    mlflow_repository.get_prediction_data = MagicMock(return_value=DataFrame())
 
     data = MagicMock()
 
@@ -775,9 +864,17 @@ async def test_fit_models_not_df_target_name_none_and_not_in_model(
 
     prediction_model.fit.assert_called_once_with(pd_merge.return_value)
 
+    mlflow_repository.get_prediction_data.assert_called_once_with(
+        prediction_model,
+        pd_merge.return_value,
+        data_model.fit.return_value.target_variable,
+        'pyfunc',
+    )
+
     assert output == {
         'prediction_model': {'model': prediction_model, 'artifact_path': 'artifact_path'},
         'data_model': {'model': data_model.fit.return_value, 'artifact_path': 'artifact_path'},
+        'prediction_data': mlflow_repository.get_prediction_data.return_value,
     }
 
 
@@ -801,6 +898,7 @@ async def test_fit_models_df_target_name_not_none_and_in_model(
             drop_duplicates=MagicMock(return_value=MagicMock(columns=['feat_1']))
         )
     )
+    mlflow_repository.get_prediction_data = MagicMock(return_value=DataFrame())
 
     data = MagicMock()
 
@@ -855,9 +953,14 @@ async def test_fit_models_df_target_name_not_none_and_in_model(
 
     prediction_model.fit.assert_called_once_with(transformed_data)
 
+    mlflow_repository.get_prediction_data.assert_called_once_with(
+        prediction_model, transformed_data, 'feat_1', 'pyfunc'
+    )
+
     assert output == {
         'prediction_model': {'model': prediction_model, 'artifact_path': 'artifact_path'},
         'data_model': {'model': data_model, 'artifact_path': 'artifact_path'},
+        'prediction_data': mlflow_repository.get_prediction_data.return_value,
     }
 
 
@@ -878,7 +981,8 @@ async def test_log_model_sklearn(mlflow, mlflow_repository):
 @patch('laborious.utils.repository.model_repository.path')
 @pytest.mark.asyncio
 async def test_log_model_pyfunc(path, mlflow, mlflow_repository):
-    model_data = {'model': MagicMock(), 'artifact_path': 'artifact_path'}
+    model_mock = MagicMock()
+    model_data = {'model': model_mock, 'artifact_path': 'artifact_path'}
     await mlflow_repository.log_model(
         model_data, 'pyfunc', 'prediction_model', metadata['metadata']
     )
@@ -887,7 +991,7 @@ async def test_log_model_pyfunc(path, mlflow, mlflow_repository):
 
     path.join.assert_called_once_with('artifact_path', 'code', 'utils')
 
-    model_data['model'].store_model.assert_called_once_with(
+    model_mock.store_model.assert_called_once_with(
         artifact_path='prediction_model', code_path=[path.join.return_value], to_disk=False
     )
 
@@ -934,9 +1038,11 @@ async def test_create_new_experiment(
 ):
     model_name = 'model_name'
     data = MagicMock()
+    prediction_data = MagicMock(spec=DataFrame)
     retrain_data = {
         'prediction_model': {'model': MagicMock(), 'artifact_path': 'artifact_path'},
         'data_model': {'model': MagicMock(), 'artifact_path': 'artifact_path'},
+        'prediction_data': prediction_data,
     }
 
     mlflow_repository.get_model_params = MagicMock(
@@ -971,7 +1077,13 @@ async def test_create_new_experiment(
         mlflow_repository.get_experiment.return_value.name
     )
 
-    data.to_csv.assert_called_once_with('./tmp/artifacts/model_name/retrain_data.csv', index=True)
+    data.to_csv.assert_called_once_with('./tmp/artifacts/model_name/retrain_data.csv', index=False)
+    # Verify prediction_data.to_csv was called with correct arguments
+    prediction_data.to_csv.assert_called_once()
+    assert (
+        prediction_data.to_csv.call_args[0][0] == './tmp/artifacts/model_name/evaluation_data.csv'
+    )
+    assert prediction_data.to_csv.call_args[1]['index'] is False
 
     mlflow.start_run.assert_called_once_with(
         experiment_id=mlflow_repository.get_experiment.return_value.experiment_id,
@@ -1000,7 +1112,12 @@ async def test_create_new_experiment(
         }
     )
 
-    mlflow.log_artifact.assert_called_once_with('./tmp/artifacts/model_name/retrain_data.csv')
+    mlflow.log_artifact.assert_has_calls(
+        [
+            call('./tmp/artifacts/model_name/retrain_data.csv'),
+            call('./tmp/artifacts/model_name/evaluation_data.csv'),
+        ]
+    )
 
     force_memory_release.assert_called_once_with(mlflow_repository.logger)
 
@@ -1026,9 +1143,11 @@ async def test_create_new_experiment_error(
     mlflow.start_run.side_effect = ValueError('error')
     model_name = 'model_name'
     data = MagicMock()
+    prediction_data = MagicMock(spec=DataFrame)
     retrain_data = {
         'prediction_model': {'model': MagicMock(), 'artifact_path': 'artifact_path'},
         'data_model': {'model': MagicMock(), 'artifact_path': 'artifact_path'},
+        'prediction_data': prediction_data,
     }
 
     mlflow_repository.get_model_params = MagicMock(
@@ -1371,3 +1490,57 @@ async def test_update_production_model(mlflow_repository):
         'mlflow_run_id': '0',
         'mlflow_experiment_id': '0',
     }
+
+
+def test_get_prediction_data_dataframe(mlflow_repository):
+    prediction_model = MagicMock()
+    retrain_dataset = DataFrame({'feat_1': [1, 2], 'target': [3, 4]}, index=['idx1', 'idx2'])
+    prediction_model.predict.return_value = DataFrame({'pred': [5, 6]}, index=['idx1', 'idx2'])
+    target_name = 'target'
+    predict_flavor = 'sklearn'
+
+    result = mlflow_repository.get_prediction_data(
+        prediction_model, retrain_dataset, target_name, predict_flavor
+    )
+
+    prediction_model.predict.assert_called_once_with(retrain_dataset)
+    assert 'prediction' in result.columns
+    assert 'target' in result.columns
+    assert 'timestamp' in result.columns
+    assert result.index.tolist() == [0, 1]
+
+
+def test_get_prediction_data_array(mlflow_repository):
+    prediction_model = MagicMock()
+    retrain_dataset = DataFrame({'feat_1': [1, 2], 'target': [3, 4]}, index=['idx1', 'idx2'])
+    prediction_model.predict.return_value = [5, 6]
+    target_name = 'target'
+    predict_flavor = 'sklearn'
+
+    result = mlflow_repository.get_prediction_data(
+        prediction_model, retrain_dataset, target_name, predict_flavor
+    )
+
+    prediction_model.predict.assert_called_once_with(retrain_dataset)
+    assert 'prediction' in result.columns
+    assert 'target' in result.columns
+    assert 'timestamp' in result.columns
+    assert result.index.tolist() == [0, 1]
+
+
+def test_get_prediction_data_pyfunc(mlflow_repository):
+    prediction_model = MagicMock()
+    retrain_dataset = DataFrame({'feat_1': [1, 2], 'target': [3, 4]}, index=['idx1', 'idx2'])
+    prediction_model.predict.return_value = DataFrame({'pred': [5, 6]}, index=['idx1', 'idx2'])
+    target_name = 'target'
+    predict_flavor = 'pyfunc'
+
+    result = mlflow_repository.get_prediction_data(
+        prediction_model, retrain_dataset, target_name, predict_flavor
+    )
+
+    prediction_model.predict.assert_called_once_with({}, retrain_dataset)
+    assert 'prediction' in result.columns
+    assert 'target' in result.columns
+    assert 'timestamp' in result.columns
+    assert result.index.tolist() == [0, 1]

tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py

@@ -125,6 +125,156 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
     assert workflow_mock.execute_local_activity_method.call_count == 1
 
 
+@mark.asyncio
+@patch(
+    'laborious.workflows.sub_workflows.format_and_export_prediction.workflow',
+    new_callable=AsyncMock,
+)
+async def test_run_none_path_flag_with_transformed_data(
+    workflow_mock, format_and_export_prediction
+):
+    # Arrange
+    input_data = {
+        'metadata': metadata,
+        'path_flag': None,
+        'data': {'test': 'data'},
+        'transformed_data': {'transformed': 'data'},
+        'timestamp': '2021-01-01',
+        'model_id': 1,
+        'prediction_confidence': 0.9,
+        'schema': 'test_schema',
+        'table_name': 'test_table',
+        'transform_table_name': 'test_transform_table',
+        'opc_servers': ['test_server'],
+        'opc_output_config': {'test': 'config'},
+        'prediction_store_policy': 'lts:1',
+    }
+
+    prediction_data = MagicMock()
+    opc_metrics = MagicMock()
+    transformed_data = MagicMock()
+
+    workflow_mock.execute_local_activity_method.side_effect = [
+        prediction_data,  # format_prediction
+        transformed_data,  # format_transformed_data
+    ]
+
+    write_transformed_handler = AsyncMock()
+    workflow_mock.start_activity_method.return_value = write_transformed_handler
+    workflow_mock.execute_activity_method.side_effect = [
+        (prediction_data, opc_metrics),  # write_opc_data
+        MagicMock(),  # export_data_to_postgres (prediction)
+        MagicMock(),  # write_metrics
+    ]
+
+    # Act
+    await format_and_export_prediction.run(input_data)
+
+    # Assert - format_prediction call
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.format_prediction,
+                {
+                    'data': input_data['data'],
+                    'timestamp': input_data['timestamp'],
+                    'model_id': input_data['model_id'],
+                    'prediction_confidence': input_data['prediction_confidence'],
+                    'prediction_store_policy': input_data['prediction_store_policy'],
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            ),
+            call(
+                Activities.format_transformed_data,
+                {
+                    'data': input_data['transformed_data'],
+                    'model_id': input_data['model_id'],
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            ),
+        ]
+    )
+
+    # Assert - start_activity_method for transformed data export
+    workflow_mock.start_activity_method.assert_called_once_with(
+        Activities.export_data_to_postgres,
+        {
+            'schema': input_data['schema'],
+            'table_name': input_data['transform_table_name'],
+            'data': transformed_data,
+            'timestamp_conversion': {
+                'column': 'timestamp',
+                'format': DATETIME_FORMAT_WITH_TZ,
+            },
+            **metadata,
+        },
+        retry_policy=ANY,
+        start_to_close_timeout=ANY,
+    )
+
+    # Assert - write_opc_data call
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.write_opc_data,
+                {
+                    'opc_output_config': input_data['opc_output_config'],
+                    'data': prediction_data,
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    # Assert - export_data_to_postgres for prediction call
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.export_data_to_postgres,
+                {
+                    'schema': input_data['schema'],
+                    'table_name': input_data['table_name'],
+                    'data': prediction_data,
+                    'timestamp_conversion': {
+                        'column': 'timestamp',
+                        'format': DATETIME_FORMAT_WITH_TZ,
+                    },
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    # Assert - write_metrics call
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.write_metrics,
+                {
+                    **metadata,
+                    'prediction': prediction_data,
+                    'opc_metrics': opc_metrics,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    # Assert - verify counts
+    assert workflow_mock.execute_activity_method.call_count == 3
+    assert workflow_mock.execute_local_activity_method.call_count == 2
+    assert workflow_mock.start_activity_method.call_count == 1
+
+
 @mark.asyncio
 @patch(
     'laborious.workflows.sub_workflows.format_and_export_prediction.workflow',

tests/laborious/workflows/subworkflows/test_prediction_process.py

@@ -31,6 +31,7 @@ async def test_run(workflow_mock, prediction_process):
         'data': {'test': 'data'},
         'schema': 'test_schema',
         'table_name': 'test_table',
+        'transform_table_name': 'test_transform_table',
         'model_id': 1,
         'input_filters': {'test': 'filter'},
         'mlflow_transform_filters': {'test': 'filter'},
@@ -175,6 +176,7 @@ async def test_run(workflow_mock, prediction_process):
             'metadata': metadata,
             'path_flag': 'continue',
             'data': 'predicted_data',
+            'transformed_data': 'transformed_data',
             'prediction_confidence': 0.95,
             'timestamp': '2024-01-01',
             'model_id': 1,
@@ -183,6 +185,7 @@ async def test_run(workflow_mock, prediction_process):
             'opc_output_config': input_data['opc_output_config'],
             'schema': input_data['schema'],
             'table_name': input_data['table_name'],
+            'transform_table_name': input_data['transform_table_name'],
             'comment': 'Error',
             'prediction_store_policy': input_data['prediction_store_policy'],
         },
@@ -199,6 +202,7 @@ async def test_run_stop_at_input_gate(workflow_mock, prediction_process):
         'data': {'test': 'data'},
         'schema': 'test_schema',
         'table_name': 'test_table',
+        'transform_table_name': 'test_transform_table',
         'model_id': 1,
         'input_filters': {'test': 'filter'},
         'mlflow_transform_filters': {'test': 'filter'},
@@ -257,6 +261,7 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_
         'data': {'test': 'data'},
         'schema': 'test_schema',
         'table_name': 'test_table',
+        'transform_table_name': 'test_transform_table',
         'model_id': 1,
         'input_filters': {'test': 'filter'},
         'mlflow_transform_filters': {'test': 'filter'},
@@ -352,6 +357,7 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process
         'data': {'test': 'data'},
         'schema': 'test_schema',
         'table_name': 'test_table',
+        'transform_table_name': 'test_transform_table',
         'model_id': 1,
         'input_filters': {'test': 'filter'},
         'mlflow_transform_filters': {'test': 'filter'},
@@ -467,6 +473,7 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p
         'data': {'test': 'data'},
         'schema': 'test_schema',
         'table_name': 'test_table',
+        'transform_table_name': 'test_transform_table',
         'model_id': 1,
         'input_filters': {'test': 'filter'},
         'mlflow_transform_filters': {'test': 'filter'},
@@ -626,6 +633,7 @@ async def test_path_flag_handler_stop(workflow_mock, prediction_process):
             'metadata': metadata,
             'schema': schema,
             'table_name': table_name,
+            'transform_table_name': 'test_transform_table',
             'model_id': model,
             'last_timestamp': last_timestamp,
             'model_name': model_name,
@@ -664,6 +672,7 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
             'metadata': metadata,
             'schema': schema,
             'table_name': table_name,
+            'transform_table_name': 'test_transform_table',
             'model_id': model,
             'last_timestamp': last_timestamp,
             'model_name': model_name,
@@ -714,6 +723,7 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
             'metadata': metadata,
             'schema': schema,
             'table_name': table_name,
+            'transform_table_name': 'test_transform_table',
             'model_id': model,
             'last_timestamp': last_timestamp,
             'model_name': model_name,
@@ -742,6 +752,7 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
             'model_config': model_config,
             'schema': schema,
             'table_name': table_name,
+            'transform_table_name': 'test_transform_table',
             'comment': 'Prediction Process',
             'opc_output_config': {'test': 'config'},
             'prediction_store_policy': prediction_store_policy,
@@ -771,6 +782,7 @@ async def test_path_flag_handler_unknown(workflow_mock, prediction_process):
             **metadata,
             'schema': schema,
             'table_name': table_name,
+            'transform_table_name': 'test_transform_table',
             'model_id': model,
             'last_timestamp': last_timestamp,
             'model_name': model_name,

tests/laborious/workflows/test_drift.py

@@ -0,0 +1,252 @@
+from unittest.mock import ANY, AsyncMock, call, patch
+
+from pytest import fixture, mark
+from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
+
+from laborious.activities.activities import Activities
+from laborious.workflows.drift import Drift
+
+
+@fixture
+def drift() -> Drift:
+    return Drift()
+
+
+metadata = {
+    'metadata': {
+        'model_id': 'test_model_id',
+        'model_name': 'test_model',
+        'workflow_name': 'drift',
+        'schedule_name': 'test_schedule',
+    },
+}
+
+
+@mark.asyncio
+@patch('laborious.workflows.drift.workflow', new_callable=AsyncMock)
+async def test_run(workflow_mock: AsyncMock, drift: Drift):
+    # Arrange
+    input_data = {
+        'schedule_name': 'test_schedule',
+        'model_name': 'test_model',
+        'model_id': 'test_model_id',
+        'schema': 'test_schema',
+        'source_table_name': 'test_source_table',
+        'target_table_name': 'test_target_table',
+        'interval': 60,
+        'model_config': {'target': 'test_target'},
+        'drift_metrics': ['psi', 'ks'],
+        'chunk_period': 'hour',
+    }
+
+    target_name = input_data['model_config']['target']
+
+    target_data = {'data': 'test_target_data'}
+    reference_data = {'data': 'test_reference_data'}
+    drift_data = {'drift': 'test_drift_data'}
+
+    workflow_mock.start_local_activity_method.side_effect = [target_data, reference_data]
+
+    workflow_mock.execute_local_activity_method.return_value = drift_data
+    workflow_mock.execute_activity_method = AsyncMock()
+
+    # Act
+    await drift.run(input_data)
+
+    # Assert - Check start_local_activity_method calls
+    # Query format matches psycopg2.sql output (identifiers with double quotes, literals with single quotes)
+    expected_gathering_query = f"""
+            SELECT *
+            FROM "{input_data['schema']}"."{input_data['source_table_name']}"
+            WHERE
+                model_id = '{input_data['model_id']}' AND
+                timestamp > NOW() - INTERVAL '{input_data['interval']} minutes'
+            ORDER BY timestamp ASC
+        """
+
+    workflow_mock.start_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.load_custom_query,
+                {
+                    **metadata,
+                    'query': expected_gathering_query,
+                    'datetime_columns': ['timestamp', 'created_at'],
+                    'orient': 'records',
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            ),
+            call(
+                Activities.get_reference_data,
+                {
+                    **metadata,
+                    'model_name': input_data['model_name'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            ),
+        ]
+    )
+
+    # Assert - Check calculate_drift call
+    workflow_mock.execute_local_activity_method.assert_called_once_with(
+        Activities.calculate_drift,
+        {
+            **metadata,
+            'target_data': target_data,
+            'reference_data': reference_data,
+            'model_name': input_data['model_name'],
+            'model_id': input_data['model_id'],
+            'target_name': target_name,
+            'drift_metrics': input_data['drift_metrics'],
+            'chunk_period': input_data['chunk_period'],
+        },
+        retry_policy=ANY,
+        start_to_close_timeout=ANY,
+    )
+
+    # Assert - Check export_data_to_postgres call
+    workflow_mock.execute_activity_method.assert_called_once_with(
+        Activities.export_data_to_postgres,
+        {
+            **metadata,
+            'data': drift_data,
+            'schema': input_data['schema'],
+            'table_name': input_data['target_table_name'],
+            'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ},
+        },
+        retry_policy=ANY,
+        start_to_close_timeout=ANY,
+    )
+
+
+@mark.asyncio
+@patch('laborious.workflows.drift.workflow', new_callable=AsyncMock)
+async def test_run_empty_target_data(workflow_mock: AsyncMock, drift: Drift):
+    # Arrange
+    input_data = {
+        'schedule_name': 'test_schedule',
+        'model_name': 'test_model',
+        'model_id': 'test_model_id',
+        'schema': 'test_schema',
+        'source_table_name': 'test_source_table',
+        'target_table_name': 'test_target_table',
+        'interval': 60,
+        'model_config': {'target': 'test_target'},
+        'drift_metrics': ['psi', 'ks'],
+    }
+
+    target_data = None
+    reference_data = {'data': 'test_reference_data'}
+
+    workflow_mock.start_local_activity_method.side_effect = [target_data, reference_data]
+
+    workflow_mock.execute_local_activity_method = AsyncMock()
+    workflow_mock.execute_activity_method = AsyncMock()
+
+    # Act
+    await drift.run(input_data)
+
+    # Assert - Should not call calculate_drift or export
+    workflow_mock.execute_local_activity_method.assert_not_called()
+    workflow_mock.execute_activity_method.assert_not_called()
+
+
+@mark.asyncio
+@patch('laborious.workflows.drift.workflow', new_callable=AsyncMock)
+async def test_run_empty_drift_data(workflow_mock: AsyncMock, drift: Drift):
+    # Arrange
+    input_data = {
+        'schedule_name': 'test_schedule',
+        'model_name': 'test_model',
+        'model_id': 'test_model_id',
+        'schema': 'test_schema',
+        'source_table_name': 'test_source_table',
+        'target_table_name': 'test_target_table',
+        'interval': 60,
+        'model_config': {'target': 'test_target'},
+        'drift_metrics': ['psi', 'ks'],
+    }
+
+    target_name = input_data['model_config']['target']
+
+    target_data = {'data': 'test_target_data'}
+    reference_data = {'data': 'test_reference_data'}
+    drift_data = None
+
+    workflow_mock.start_local_activity_method.side_effect = [target_data, reference_data]
+
+    workflow_mock.execute_local_activity_method.return_value = drift_data
+    workflow_mock.execute_activity_method = AsyncMock()
+
+    # Act
+    await drift.run(input_data)
+
+    # Assert - Should call calculate_drift but not export
+    workflow_mock.execute_local_activity_method.assert_called_once_with(
+        Activities.calculate_drift,
+        {
+            **metadata,
+            'target_data': target_data,
+            'reference_data': reference_data,
+            'model_name': input_data['model_name'],
+            'model_id': input_data['model_id'],
+            'target_name': target_name,
+            'drift_metrics': input_data['drift_metrics'],
+            'chunk_period': input_data.get('chunk_period', 'min'),
+        },
+        retry_policy=ANY,
+        start_to_close_timeout=ANY,
+    )
+
+    workflow_mock.execute_activity_method.assert_not_called()
+
+
+@mark.asyncio
+@patch('laborious.workflows.drift.workflow', new_callable=AsyncMock)
+async def test_run_default_chunk_period(workflow_mock: AsyncMock, drift: Drift):
+    # Arrange
+    input_data = {
+        'schedule_name': 'test_schedule',
+        'model_name': 'test_model',
+        'model_id': 'test_model_id',
+        'schema': 'test_schema',
+        'source_table_name': 'test_source_table',
+        'target_table_name': 'test_target_table',
+        'interval': 60,
+        'model_config': {'target': 'test_target'},
+        'drift_metrics': ['psi', 'ks'],
+        # chunk_period not provided, should default to 'min'
+    }
+
+    target_name = input_data['model_config']['target']
+
+    target_data = {'data': 'test_target_data'}
+    reference_data = {'data': 'test_reference_data'}
+    drift_data = {'drift': 'test_drift_data'}
+
+    workflow_mock.start_local_activity_method.side_effect = [target_data, reference_data]
+
+    workflow_mock.execute_local_activity_method.return_value = drift_data
+    workflow_mock.execute_activity_method = AsyncMock()
+
+    # Act
+    await drift.run(input_data)
+
+    # Assert - Check calculate_drift call with default chunk_period
+    workflow_mock.execute_local_activity_method.assert_called_once_with(
+        Activities.calculate_drift,
+        {
+            **metadata,
+            'target_data': target_data,
+            'reference_data': reference_data,
+            'model_name': input_data['model_name'],
+            'model_id': input_data['model_id'],
+            'target_name': target_name,
+            'drift_metrics': input_data['drift_metrics'],
+            'chunk_period': 'min',  # Default value
+        },
+        retry_policy=ANY,
+        start_to_close_timeout=ANY,
+    )

tests/laborious/workflows/test_predictions_batch.py

@@ -32,6 +32,7 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch
         'query': 'SELECT * FROM test',
         'schema': 'test_schema',
         'table_name': 'test_table',
+        'transform_table_name': 'test_transform_table',
         'opc_output_config': 'test_opc_output_config',
         'datetime_columns': ['timestamp', 'created_at'],
         'prediction_store_policy': 'erl:1',
@@ -59,6 +60,7 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch
         'data': {'data': 'test_data'},
         'schema': input_data['schema'],
         'table_name': input_data['table_name'],
+        'transform_table_name': input_data['transform_table_name'],
         'model_id': input_data['model_id'],
         'model_name': input_data['model_name'],
         'input_filters': input_data.get('input_filters', {'EMPTY_DATA': {'POLICY': 'STOP'}}),
@@ -71,7 +73,7 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch
         'model_config': input_data.get('model_config', {}),
         'path_priority': input_data.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT']),
         'opc_output_config': input_data.get('opc_output_config', {}),
-        'prediction_store_policy': input_data.get('prediction_store_policy', 'erl:1'),
+        'prediction_store_policy': input_data.get('prediction_store_policy', 'lts:1'),
     }
 
     workflow_mock.execute_child_workflow.assert_has_calls(

tests/laborious/workflows/test_simple_metrics.py

@@ -0,0 +1,223 @@
+from unittest.mock import ANY, AsyncMock, call, patch
+
+from pytest import fixture, mark
+from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
+
+from laborious.activities.activities import Activities
+from laborious.workflows.simple_metrics import SimpleMetrics
+
+
+@fixture
+def simple_metrics() -> SimpleMetrics:
+    return SimpleMetrics()
+
+
+metadata = {
+    'metadata': {
+        'model_id': 'test_model_id',
+        'model_name': 'test_model',
+        'workflow_name': 'simple_metrics',
+        'schedule_name': 'test_schedule',
+    },
+}
+
+
+@mark.asyncio
+@patch('laborious.workflows.simple_metrics.workflow', new_callable=AsyncMock)
+async def test_run(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics):
+    # Arrange
+    input_data = {
+        'schedule_name': 'test_schedule',
+        'model_name': 'test_model',
+        'model_id': 'test_model_id',
+        'interval_minutes': 60,
+        'model_config': {'target': 'test_target'},
+        'schema': 'test_schema',
+        'predictions_table_name': 'test_predictions_table',
+        'data_table_name': 'test_data_table',
+        'target_table_name': 'test_target_table',
+        'metrics': ['rmse', 'mse', 'mae', 'r2'],
+    }
+
+    target_data = {'data': 'test_target_data'}
+    simple_metrics_data = {'metrics': 'test_simple_metrics_data'}
+
+    workflow_mock.execute_local_activity_method.side_effect = [target_data, simple_metrics_data]
+
+    workflow_mock.execute_activity_method = AsyncMock()
+
+    # Act
+    await simple_metrics.run(input_data)
+
+    # Assert - Check load_custom_query call
+    # Query format matches psycopg2.sql output (identifiers with double quotes, literals with single quotes)
+    expected_query = f"""
+            select p."timestamp", p.prediction, ld.value as "target"
+            from "{input_data['schema']}"."{input_data['predictions_table_name']}" p
+            inner join "{input_data['schema']}"."{input_data['data_table_name']}" ld
+                on p."timestamp" = ld."timestamp"
+            where
+                p.model_id = '{input_data['model_id']}' and
+                p.prediction is not null and
+                ld.variable = '{input_data['model_config']['target']}' and
+                ld.value is not null and
+                p."timestamp" >= NOW() - INTERVAL '{input_data['interval_minutes']} minutes'
+            order by
+                p."timestamp" desc;
+        """
+
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.load_custom_query,
+                {
+                    **metadata,
+                    'query': expected_query,
+                    'datetime_columns': ['timestamp'],
+                    'orient': 'records',
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            ),
+            call(
+                Activities.calculate_simple_metrics,
+                {
+                    **metadata,
+                    'model_id': input_data['model_id'],
+                    'target_data': target_data,
+                    'metrics': input_data['metrics'],
+                    'interval_minutes': input_data['interval_minutes'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            ),
+        ]
+    )
+
+    # Assert - Check export_data_to_postgres call
+    workflow_mock.execute_activity_method.assert_called_once_with(
+        Activities.export_data_to_postgres,
+        {
+            **metadata,
+            'data': simple_metrics_data,
+            'schema': input_data['schema'],
+            'table_name': input_data['target_table_name'],
+            'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ},
+        },
+        retry_policy=ANY,
+        start_to_close_timeout=ANY,
+    )
+
+
+@mark.asyncio
+@patch('laborious.workflows.simple_metrics.workflow', new_callable=AsyncMock)
+async def test_run_empty_target_data(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics):
+    # Arrange
+    input_data = {
+        'schedule_name': 'test_schedule',
+        'model_name': 'test_model',
+        'model_id': 'test_model_id',
+        'interval_minutes': 60,
+        'model_config': {'target': 'test_target'},
+        'schema': 'test_schema',
+        'predictions_table_name': 'test_predictions_table',
+        'data_table_name': 'test_data_table',
+        'target_table_name': 'test_target_table',
+        'metrics': ['rmse', 'mse'],
+    }
+
+    target_data = None
+
+    workflow_mock.execute_local_activity_method.return_value = target_data
+    workflow_mock.execute_activity_method = AsyncMock()
+
+    # Act
+    await simple_metrics.run(input_data)
+
+    # Assert - Should not call calculate_simple_metrics or export
+    assert workflow_mock.execute_local_activity_method.call_count == 1
+    workflow_mock.execute_activity_method.assert_not_called()
+
+
+@mark.asyncio
+@patch('laborious.workflows.simple_metrics.workflow', new_callable=AsyncMock)
+async def test_run_empty_simple_metrics(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics):
+    # Arrange
+    input_data = {
+        'schedule_name': 'test_schedule',
+        'model_name': 'test_model',
+        'model_id': 'test_model_id',
+        'interval_minutes': 60,
+        'model_config': {'target': 'test_target'},
+        'schema': 'test_schema',
+        'predictions_table_name': 'test_predictions_table',
+        'data_table_name': 'test_data_table',
+        'target_table_name': 'test_target_table',
+        'metrics': ['rmse', 'mse'],
+    }
+
+    target_data = {'data': 'test_target_data'}
+    simple_metrics_data = None
+
+    workflow_mock.execute_local_activity_method.side_effect = [target_data, simple_metrics_data]
+
+    workflow_mock.execute_activity_method = AsyncMock()
+
+    # Act
+    await simple_metrics.run(input_data)
+
+    # Assert - Should call calculate_simple_metrics but not export
+    assert workflow_mock.execute_local_activity_method.call_count == 2
+    workflow_mock.execute_activity_method.assert_not_called()
+
+
+@mark.asyncio
+@patch('laborious.workflows.simple_metrics.workflow', new_callable=AsyncMock)
+async def test_run_default_metrics(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics):
+    # Arrange
+    input_data = {
+        'schedule_name': 'test_schedule',
+        'model_name': 'test_model',
+        'model_id': 'test_model_id',
+        'interval_minutes': 60,
+        'model_config': {'target': 'test_target'},
+        'schema': 'test_schema',
+        'predictions_table_name': 'test_predictions_table',
+        'data_table_name': 'test_data_table',
+        'target_table_name': 'test_target_table',
+        # metrics not provided, should default to ['rmse', 'mse', 'mae', 'r2']
+    }
+
+    target_data = {'data': 'test_target_data'}
+    simple_metrics_data = {'metrics': 'test_simple_metrics_data'}
+
+    workflow_mock.execute_local_activity_method.side_effect = [target_data, simple_metrics_data]
+
+    workflow_mock.execute_activity_method = AsyncMock()
+
+    # Act
+    await simple_metrics.run(input_data)
+
+    # Assert - Check calculate_simple_metrics call with default metrics
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.load_custom_query,
+                ANY,
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            ),
+            call(
+                Activities.calculate_simple_metrics,
+                {
+                    **metadata,
+                    'model_id': input_data['model_id'],
+                    'target_data': target_data,
+                    'metrics': ['rmse', 'mse', 'mae', 'r2'],  # Default value
+                    'interval_minutes': input_data['interval_minutes'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            ),
+        ]
+    )

validate.sh

@@ -87,7 +87,7 @@ if ! run_step "3. Type Checking (mypy)" "mypy laborious/"; then
 fi
 
 # Step 4: Security Analysis (Bandit)
-if ! run_step "4. Security Analysis (Bandit)" "bandit -r laborious/ -ll -q"; then
+if ! run_step "4. Security Analysis (Bandit)" "bandit -c pyproject.toml -r laborious/ -ll -q"; then
     FAILED_STEPS+=("Security Analysis")
 fi
 

values.yaml

@@ -11,7 +11,7 @@ image:
   # This sets the pull policy for images.
   pullPolicy: Always
   # Overrides the image tag whose default is the chart appVersion.
-  tag: "1.1.0"
+  tag: "1.1.2"
 
 0# This is for the secrets for pulling an image from a private repository more information can be found here: https://kubernetes.io/docs/tasks/configure-pod-container/pull-image-private-registry/
 imagePullSecrets:
@@ -52,17 +52,15 @@ securityContext: {}
   # runAsUser: 1000
 
 
-resources: {}
+resources:
-  # We usually recommend not to specify default resources and to leave this as a conscious
+  # Resource limits and requests are important for ResourceBasedTuner to work correctly.
-  # choice for the user. This also increases chances charts run on environments with little
+  # The tuner monitors system CPU and memory usage, so proper resource limits must be set.
-  # resources, such as Minikube. If you do want to specify resources, uncomment the following
+  limits:
-  # lines, adjust them as necessary, and remove the curly braces after 'resources:'.
+    cpu: 2000m  # 2 CPU cores
-  # limits:
+    memory: 20Gi  # 20 GB memory
-  #   cpu: 100m
+  requests:
-  #   memory: 128Mi
+    cpu: 1000m  # 1 CPU core
-  # requests:
+    memory: 2Gi  # 2 GB memory
-  #   cpu: 100m
-  #   memory: 128Mi
 
 # This is to setup the liveness and readiness probes more information can be found here: https://kubernetes.io/docs/tasks/configure-pod-container/configure-liveness-readiness-startup-probes/
 livenessProbe:
@@ -151,7 +149,7 @@ env:
   - name: GITHUB_REPO_URL
     value: "git@github.com:Aignosi/sientia-dataops-laborious_temporal.git"
   - name: GITHUB_BRANCH
-    value: "feature/SIENTIAPDE-1325-adicionar-metricas-especificas-de-operacoes-externas"
+    value: "feature/SIENTIAPDE-1273"
   - name: PYTHON_APP
     value: "laborious.worker.worker"
 
@@ -167,9 +165,11 @@ env:
   - name: POSTGRES_DBNAME
     value: "sientia"
   - name: POSTGRES_MIN_CONNECTIONS
-    value: "10"
+    value: "20"
+  # max_connections = number_of_workers * max_concurrent_activities * safety_factor
+  # Example: 4 workers * 50 activities * 0.5 = 100 connections
   - name: POSTGRES_MAX_CONNECTIONS
-    value: "30"
+    value: "100"
 
   - name: MLFLOW_HOST
     value: "http://sientia-tracker-mlflow-tracking.sientia-tracker.svc.cluster.local"
@@ -234,7 +234,7 @@ ssh:
 
 # kubectl create secret docker-registry docker-hub-secret --namespace sientia --docker-server=http://aignosi.azurecr.io --docker-username=aignosi --docker-password=5I5zpQ6sRaHqX1hD3dr+2mo647yO3FRc359/wu6gsP+ACRDRz5mp
 
-# helm upgrade --install sientia-laborious-worker sientia/sientia-module -n sientia --create-namespace -f ./values.yaml --version 0.5.0
+# helm upgrade --install sientia-laborious-worker sientia/sientia-module -n sientia --create-namespace -f ./values.yaml --version 0.6.0
 
 # kubectl create secret generic git-ssh-key-sientia-laborious-worker \
 #   --namespace sientia \