SIENTIAPDE-2072: bump sientia_do pin to 1.12.2

Picks up the notification timestamp -> native datetime fix.

.env.example

@@ -26,4 +26,10 @@ MONGODB_USERNAME="mongo_user"
 MONGODB_PASSWORD="mongo_db_password"
 MONGODB_URL="my-release-mongodb.mongodb.svc.cluster.local:27017"
 MONGODB_DATABASE="sientia"
-MONGODB_TTL_INDEX_HOURS="1"
+MONGODB_TTL_INDEX_HOURS="1"
+
+MINIO_ENDPOINT_URL="http://localhost:9000"
+MINIO_ACCESS_KEY="sientia"
+MINIO_SECRET_KEY="sientia"
+MINIO_REGION_NAME="sa-east-1"
+MINIO_DEFAULT_BUCKET="sientia"

.github/workflows/quality-gate.yml

@@ -4,71 +4,28 @@ on:
   push:
     branches:
       - main
+      - 'release/**'
+      - 'feature/**'
   pull_request:
     branches:
       - main
+      - 'release/**'
+      - 'feature/**'
     types: [ opened, synchronize, reopened ]
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
 jobs:
-  sonar:
-    name: SonarQube Analysis
-    runs-on: ubuntu-latest
-    permissions: write-all
-    steps:
-      - name: ⬇️ Checkout Code
-        uses: actions/checkout@v4
-        with:
-          fetch-depth: 0
-          persist-credentials: false
-
-      - name: Generate App Token
-        id: generate-app-token
-        uses: actions/create-github-app-token@v1
-        with:
-          app-id: ${{ secrets.APP_ID }}
-          private-key: ${{ secrets.APP_PRIVATE_KEY }}
-          owner: 'Aignosi'
-          repositories: 'sientia-dataops-library,sientia-mlops-library'
-
-      - name: Prepare requirements.txt
-        id: prepare-requirements
-        run: |
-          sed -e "s|git+ssh://git@github.com/|git+https://github.com/|g" \
-              -e "s|git@github.com:|git+https://github.com/|g" \
-              requirements.txt > requirements_prepared.txt
-          echo "PROCESSED_REQUIREMENTS_FILE=requirements_prepared.txt" >> $GITHUB_OUTPUT
-
-      - name: Configure Git to use App Token
-        env:
-          GH_APP_TOKEN: ${{ steps.generate-app-token.outputs.token }}
-        run: |
-          git config --global url."https://oauth2:${GH_APP_TOKEN}@github.com/".insteadOf "https://github.com/"
-
-      - name: 🔧 Setup Python
-        uses: actions/setup-python@v4
-        with:
-          python-version: "3.11"
-
-      - name: 🗄️ Cache Python dependencies
-        uses: actions/cache@v3
-        with:
-          path: ~/.cache/pip
-          key: ${{ runner.os }}-pip-${{ hashFiles(steps.prepare-requirements.outputs.PROCESSED_REQUIREMENTS_FILE) }}
-          restore-keys: |
-            ${{ runner.os }}-pip-
-
-      - name: 📦 Install Dependencies
-        run: |
-          python -m pip install --upgrade pip
-          pip install -r ${{ steps.prepare-requirements.outputs.PROCESSED_REQUIREMENTS_FILE }}
-          pip install pytest pytest-cov pytest-asyncio
-
-      - name: 🧪 Run Tests with Pytest
-        run: |
-          pytest tests --junitxml=pytest.xml --cov=laborious --cov-report=xml --cov-report=term
-
-      - name: Run SonarQube Analysis
-        uses: SonarSource/sonarqube-scan-action@v5
-        env:
-          SONAR_TOKEN: ${{ secrets.SONAR_TOKEN }}
-          SONAR_HOST_URL: ${{ secrets.SONAR_HOST_URL }}
+  quality-gate:
+    uses: Aignosi/github_workflow_templates/.github/workflows/python-quality-gate.yml@main
+    permissions:
+      contents: read
+      pull-requests: write
+      issues: write
+    with:
+      project_name: 'laborious'
+      repositories: 'sientia-dataops-library, sientia-mlops-library'
+      requirements_file: 'requirements-light.txt'
+    secrets: inherit

.github/workflows/release.yml

@@ -0,0 +1,25 @@
+name: Create Release on Merge to Main
+
+on:
+  pull_request:
+    types: [closed]
+    branches:
+      - main
+  workflow_dispatch:
+    inputs:
+      version:
+        description: 'Version to release'
+        required: false
+        type: string
+
+jobs:
+  release:
+    if: |
+      (github.event_name == 'pull_request' && github.event.pull_request.merged == true) ||
+      github.event_name == 'workflow_dispatch'
+    uses: Aignosi/github_workflow_templates/.github/workflows/dataops-module-release.yml@main
+    permissions: write-all
+    with:
+      project_name: 'laborious'
+      release_version: ${{ github.event.inputs.version || '' }}
+    secrets: inherit

.gitignore

@@ -37,10 +37,22 @@ __pycache__/
 # Ignorar coverage
 htmlcov/
 .coverage
+coverage.xml
 
 # git keys
 git_key*
 
 git_log
 
-.env
+.env
+
+tmp/
+catboost_info/
+
+.ruff_cache/
+.mypy_cache/
+mlruns/
+
+relatorio*
+openspec/
+.cursor/

docs/opc-communication.md

@@ -0,0 +1,170 @@
+# OPC UA communication (Laborious)
+
+Laborious exports predictions to OPC UA servers through `OpcRepository` ([`laborious/utils/repository/opc_repository.py`](../laborious/utils/repository/opc_repository.py)) and the Temporal activity layer in [`laborious/activities/opc.py`](../laborious/activities/opc.py).
+
+Implementation plan for session/channel recovery on Tier-1 `Bad*` errors: [`.cursor/plans/opc_bad_reconnect_ac4c6045.plan.md`](../.cursor/plans/opc_bad_reconnect_ac4c6045.plan.md).
+
+## Architecture
+
+```text
+Worker (long-lived)
+  └── OpcRepository per OPC server id (from OPC_CONFIG / env)
+        ├── connect / disconnect / validate_connection (read-only)
+        ├── _connect_locked / _reconnect_locked (under _connection_lock)
+        ├── write_data (single attempt per call)
+        └── background reconnect on Tier-1 Bad*, protocol closed, or stale session
+
+Temporal activity write_opc_data
+  └── OPC.manage_output_tags → write_data per tag (sequential per activity)
+```
+
+One worker process holds one `OpcRepository` instance per configured server. Multiple Temporal activities can call `write_data` concurrently on the same repository.
+
+## Connection lifecycle
+
+| Phase | Behavior |
+|-------|----------|
+| Startup | `init_opc()` creates repositories and calls `connect()` → `_connect_locked()` |
+| Steady state | `validate_connection()` is read-only (`protocol.state` only); `_session_ready` is checked in `write_data` |
+| Tier-1 Bad* / protocol closed | `_start_reconnect` → `_run_reconnect` → `_reconnect_locked()` (respects `reconnection_interval`) |
+| Write | `write_data()` checks reconnect task, `_session_ready`, validates protocol, then one `get_node` + `write_value` |
+| Shutdown | `close()` disconnects all repositories |
+
+### Session and channel timeouts
+
+Requested session and secure-channel lifetime: **10 minutes** (`OPC_UA_SESSION_AND_CHANNEL_TIMEOUT_MS` in `opc_repository.py`). The server may revise these values; negotiated values are logged after connect and exposed as `opc_session_revised_timeout_milliseconds`.
+
+### Reconnection interval
+
+`OPC_RECONNECTION_INTERVAL` is in **seconds** (default `120`). It gates **background** reconnect after Tier-1 `Bad*`, closed protocol, or stale session (`last_reconnection_time` is updated only in `_reconnect_locked()`). It limits load on the OPC server when many workflows fail at once.
+
+## Concurrency: connection lock and session readiness
+
+To allow **multiple concurrent writes** when the session is healthy, but **block all writes** while the connection is being torn down or re-established:
+
+| Primitive | Role |
+|-----------|------|
+| `_connection_lock` (`asyncio.Lock`) | Held for the entire `disconnect` → `connect` path. Only one connection-maintenance task at a time. |
+| `_session_ready` (`asyncio.Event`) | Set when a session is ready for writes; cleared before reconnect starts and set again after a successful connect. |
+
+**Connection methods (caller holds `_connection_lock` for `_*_locked` helpers):**
+
+| Method | Role |
+|--------|------|
+| `_create_client()` | Create asyncua `Client` + optional `set_security`; raises if `client` already exists |
+| `_open_session()` | `client.connect()` + metrics; raises if session already open or client missing |
+| `_connect_locked()` | `_create_client()` (when needed) + `_open_session()`; raises if already connected |
+| `_disconnect_locked()` | Teardown session and clear `client` |
+| `_reconnect_locked()` | `_disconnect_locked()` + `_connect_locked()`; sets `last_reconnection_time` |
+
+Public `connect()` / `disconnect()` acquire the lock and call `_connect_locked()` / `_disconnect_locked()`.
+
+**Write path (`write_data`):**
+
+1. If a reconnect task is **in flight** → **fail immediately** (`opc_error_kind=reconnect_in_progress`).
+2. If `_session_ready` is cleared and no task is running → schedule reconnect (`SessionNotReady`); fail with `connection_lost` or `reconnect_in_progress` if a task started.
+3. `validate_connection()` checks `protocol.state` only (read-only). If closed → schedule reconnect (`ProtocolClosed`) and fail with `opc_error_kind=connection_lost`.
+4. Single `get_node` + `write_value` (no retry). Tier-1 `Bad*` on write also schedules reconnect.
+
+**Reconnect path (`_run_reconnect`):**
+
+1. `_start_reconnect` clears `_session_ready` and schedules the task when the interval allows and `_allow_reconnect` is true.
+2. `async with _connection_lock:` → `_reconnect_locked()`.
+3. `_session_ready` is set on successful `_open_session()`.
+4. `disconnect()` sets `_allow_reconnect=False` so shutdown does not respawn sessions.
+
+A second `_connect_locked()` while a session is already open raises `OpcSessionAlreadyConnectedError` (disconnect first).
+
+**asyncua note:** Concurrent `write_value` on the same session is only safe if the stack tolerates it. If production shows issues, serialize writes with an optional `asyncio.Semaphore(1)` while keeping the connection lock semantics above.
+
+**Future threads:** replace `asyncio.Lock` / `Event` with `threading` primitives or route all OPC I/O through one dedicated loop.
+
+## Reconnect triggers
+
+Background reconnect is scheduled when:
+
+- `validate_connection()` sees a closed or missing protocol (`ProtocolClosed`).
+- `_session_ready` is clear after a failed reconnect (`SessionNotReady`).
+- A write raises a Tier-1 `UaStatusCodeError` in `RECONNECTABLE_OPC_BAD_NAMES`.
+
+For Tier-1 `Bad*` when the server invalidates the session (e.g. `BadSessionIdInvalid`) but the client still sees transport as open, `write_data` fails once, records the OPC status in metrics, and **schedules** reconnect if:
+
+- The exception is a `UaStatusCodeError` whose name is in `RECONNECTABLE_OPC_BAD_NAMES` (see plan), and
+- `reconnection_interval` has elapsed since `last_reconnection_time`, and
+- No reconnect task is already running.
+
+There is **no write retry**: the failed export is not sent again in the same activity.
+
+## Prediction confidence and PostgreSQL comments
+
+| `prediction_confidence` | Meaning |
+|-------------------------|---------|
+| (unchanged) | Successful OPC export |
+| **12** | Generic OPC write failure (`OPC_WRITTING_ERROR_CONFIDENCE`) |
+| **14** | Tier-1 session/channel `Bad*` on export (`OPC_SESSION_BAD_CONFIDENCE`) |
+| **14** | Write while reconnect in progress (`OPC_SESSION_BAD_CONFIDENCE`, comment `OPC UA reconnect in progress`) |
+| **13** | PI Web API write failure (separate path) |
+
+Session/channel errors use a stable comment for counting:
+
+```text
+OPC UA session/channel error: BadSessionIdInvalid
+```
+
+Reconnect-in-progress exports use:
+
+```text
+OPC UA reconnect in progress
+```
+
+Example SQL:
+
+```sql
+SELECT count(*) FROM predictions WHERE prediction_confidence = 14;
+SELECT count(*) FROM predictions WHERE comments LIKE 'OPC UA session/channel error:%';
+```
+
+## Prometheus metrics (`opc_*`)
+
+Defined in [`laborious/metrics.py`](../laborious/metrics.py). Do not rename in production without a dashboard migration.
+
+| Metric | Purpose |
+|--------|---------|
+| `opc_connections_initiated_total` | Connection attempts |
+| `opc_connections_failed_total` | Failed connects |
+| `opc_connection_status` | Gauge 1=connected, 0=disconnected |
+| `opc_session_created_total` | Session established after connect |
+| `opc_session_closed_total` | Disconnect initiated |
+| `opc_session_revised_timeout_milliseconds` | Negotiated session timeout (ms) |
+| `opc_write_attempts_total` | Per write; label `result` = `OK` or exception name |
+| `opc_write_inter_arrival_over_session_timeout_total` | Successful writes spaced longer than revised session timeout |
+
+Legacy activity metrics: `laborious_prediction_opc_writing_count`, `laborious_prediction_opc_writing_response_time_monitor`.
+
+## Environment variables
+
+| Variable | Default | Description |
+|----------|---------|-------------|
+| `OPC_CONFIG` | — | JSON map of server configs (overrides single-server env) |
+| `OPC_ID` | `1` | Server id |
+| `OPC_URL` | `opc.tcp://localhost:4840` | Endpoint |
+| `OPC_SERVER_NAME` | `default_server` | Label for metrics/logs |
+| `OPC_SERVER_URI` | same as URL | Application URI / cert SAN |
+| `OPC_CERT_PATH` | — | Client certificate (secure mode) |
+| `OPC_PRIVATE_KEY_PATH` | — | Client private key |
+| `OPC_SERVER_CERT_PATH` | — | Server certificate |
+| `OPC_RECONNECTION_INTERVAL` | `120` | Minimum seconds between reconnects |
+
+## Operations checklist
+
+- Correlate `BadSessionIdInvalid` in `opc_write_attempts_total` with `opc_session_closed_total` / `opc_session_created_total` (reconnect may finish after the row is stored with confidence 14).
+- Use confidence **14** and comment prefix for session invalidation rates; use **12** for other OPC failures.
+- Respect `OPC_RECONNECTION_INTERVAL` under parallel load; bursts of confidence 14 are expected until the next successful cycle.
+
+## Related tests
+
+- Unit: [`tests/laborious/utils/repository/test_opc_repository.py`](../tests/laborious/utils/repository/test_opc_repository.py)
+- Unit: [`tests/laborious/activities/test_opc.py`](../tests/laborious/activities/test_opc.py)
+- E2E (mock OPC): [`e2e/test_predictions_batch_format_export.py`](../e2e/test_predictions_batch_format_export.py)
+- E2E (in-process asyncua server + real `OpcRepository`): [`e2e/test_opc_real_server.py`](../e2e/test_opc_real_server.py) — scenarios 3.1.2, 3.2.2, 3.2.4, 3.2.5
+- Scenarios: [`e2e/scenarios.md`](../e2e/scenarios.md)

e2e/__init__.py

@@ -0,0 +1,3 @@
+"""
+End-to-end tests for laborious temporal workflows.
+"""

e2e/conftest.py

@@ -0,0 +1,681 @@
+"""
+Pytest configuration and fixtures for E2E tests.
+"""
+
+import sys
+from unittest.mock import AsyncMock, MagicMock, patch
+
+# E2E workflows under test do not run ModelAnalysis; stub before Activities import.
+_model_analysis_module = MagicMock()
+_model_analysis_module.ModelAnalysis = MagicMock
+sys.modules.setdefault('sientia', MagicMock())
+sys.modules.setdefault('sientia.ModelAnalysis', _model_analysis_module)
+from io import BytesIO
+
+import pandas as pd
+import pytest
+import pytest_asyncio
+from sqlalchemy import create_engine, text
+from testcontainers.minio import MinioContainer
+from testcontainers.postgres import PostgresContainer
+from temporalio.testing import WorkflowEnvironment
+from temporalio.worker import Worker
+
+from e2e.opc_test_server import OpcE2ETestServer
+from laborious.activities.activities import Activities
+from laborious.workflows.predictions_batch import PredictionsBatch
+from laborious.workflows.sub_workflows.prediction_process import PredictionProcess
+from laborious.workflows.sub_workflows.format_and_export_prediction import FormatAndExportPrediction
+from sientia_do.notifications.handlers import CoreNotificationHandler
+from sientia_do.observability.logger import Logger
+from sientia_do.observability.metrics_controller import MetricsController
+
+# Test constants
+TEST_MONGODB_CONNECTION_STRING = 'mongodb://localhost:27017'
+TEST_DATABASE_NAME = 'test_db'
+
+
+@pytest_asyncio.fixture(scope='session')
+def minio_container():
+    """
+    MinIO S3-compatible storage for E2E tests that exercise real offload uploads.
+    """
+    minio = MinioContainer()
+    minio.start()
+    yield minio
+    minio.stop()
+
+
+@pytest_asyncio.fixture(scope='session')
+def postgres_container():
+    """
+    Create a PostgreSQL container using testcontainers.
+    
+    This fixture creates a real PostgreSQL database in a Docker container
+    that will be used for all tests in the session.
+    """
+    postgres = PostgresContainer('postgres:15')
+    postgres.start()
+    yield postgres
+    postgres.stop()
+
+
+@pytest_asyncio.fixture
+def postgres_engine(postgres_container):
+    """
+    Create SQLAlchemy engine for PostgreSQL test database.
+    
+    This fixture creates a connection to the PostgreSQL container
+    created by the postgres_container fixture.
+    """
+    engine = create_engine(postgres_container.get_connection_url())
+    
+    yield engine
+    
+    engine.dispose()
+
+
+def _create_schema_and_tables(engine):
+    """
+    Helper function to create schema and tables in the given engine.
+    
+    Creates predictions_schema with:
+    - laborious_data: Input data table for queries
+    - predictions: Output predictions table
+    - transformed_data: Output transformed data table
+    """
+    # Use begin() to ensure transaction is properly committed
+    with engine.begin() as conn:
+        # Create predictions_schema
+        conn.execute(text("CREATE SCHEMA IF NOT EXISTS predictions_schema"))
+        
+        # Create laborious_data table (input data from sensors)
+        create_laborious_data_sql = """
+            CREATE TABLE IF NOT EXISTS predictions_schema.laborious_data (
+                id SERIAL NOT NULL,
+                model_id int4 NOT NULL,
+                variable text NOT NULL,
+                value numeric NULL,
+                "timestamp" timestamptz NOT NULL,
+                created_at timestamptz NOT NULL,
+                PRIMARY KEY (id)
+            );
+        """
+        conn.execute(text(create_laborious_data_sql))
+        
+        # Create predictions table
+        create_predictions_sql = """
+            CREATE TABLE if not exists predictions_schema.predictions (
+                id SERIAL NOT NULL ,
+                model_id int4 NOT NULL,
+                prediction numeric NULL,
+                prediction_confidence numeric NOT NULL,
+                response_time numeric NOT NULL,
+                prediction_status text NOT NULL,
+                "timestamp" timestamptz NOT NULL,
+                created_at timestamptz DEFAULT CURRENT_TIMESTAMP NOT NULL,
+                "comments" text NULL,
+                PRIMARY KEY (id, created_at)
+            );
+        """
+        conn.execute(text(create_predictions_sql))
+        
+        # Create transformed_data table
+        create_transformed_sql = """
+            CREATE TABLE IF NOT EXISTS predictions_schema.transformed_data (
+                id SERIAL NOT NULL,
+                model_id int4 NOT NULL,
+                variable text NOT NULL,
+                value numeric NULL,
+                "timestamp" timestamptz NOT NULL,
+                created_at timestamptz DEFAULT CURRENT_TIMESTAMP NOT NULL,
+                PRIMARY KEY (id)
+            );
+        """
+        conn.execute(text(create_transformed_sql))
+
+
+@pytest_asyncio.fixture(autouse=True)
+def setup_postgres_schema_and_tables(postgres_engine):
+    """
+    Automatically create necessary schema and tables before each test.
+    
+    This fixture runs automatically (autouse=True) and ensures
+    that the predictions_schema and tables exist with the correct structure.
+    """
+    _create_schema_and_tables(postgres_engine)
+    yield
+
+
+@pytest_asyncio.fixture
+def mock_logger():
+    """Mock logger for testing."""
+    def message(message):
+        print(f"[LOG] {message}")
+    def custom_message(message, _metadata={}):
+        print(f"[LOG] {message}")
+    logger = MagicMock()
+    logger.info = MagicMock(
+        side_effect=message
+    )
+    logger.debug = MagicMock(
+        side_effect=message
+    )
+    logger.error = MagicMock(
+        side_effect=message
+    )
+    logger.warning = MagicMock(
+        side_effect=message
+    )
+    logger.custom_info = MagicMock(
+        side_effect=custom_message
+    )
+    logger.custom_debug = MagicMock(
+        side_effect=custom_message
+    )
+    logger.custom_error = MagicMock(
+        side_effect=custom_message
+    )
+    logger.custom_warning = MagicMock(
+        side_effect=custom_message
+    )
+    return logger
+
+
+@pytest_asyncio.fixture
+def mock_mongo_client():
+    """
+    Mock MongoDB client to avoid real connections.
+    
+    This fixture mocks the pymongo.MongoClient used by CoreNotificationHandler,
+    allowing us to use a real NotificationHandler instance without connecting to MongoDB.
+    """
+    mock_client = MagicMock()
+    mock_db = MagicMock()
+    mock_collection = MagicMock()
+    
+    # Configure the mock chain: client[database] -> db[collection] -> collection
+    mock_client.__getitem__.return_value = mock_db
+    mock_db.__getitem__.return_value = mock_collection
+    
+    # Mock server_info() to avoid connection attempts
+    mock_client.server_info = MagicMock()
+    
+    # Mock insert_one for notifications
+    mock_collection.insert_one = MagicMock()
+    
+    return mock_client
+
+
+@pytest.fixture
+def notification_inserts(mock_mongo_client):
+    """
+    Mongo insert_one mock used by CoreNotificationHandler for notification persistence.
+
+    Yields:
+        MagicMock for insert_one, reset before each test.
+    """
+    mock_db = mock_mongo_client.__getitem__.return_value
+    mock_collection = mock_db.__getitem__.return_value
+    mock_collection.insert_one.reset_mock()
+    yield mock_collection.insert_one
+
+
+@pytest_asyncio.fixture
+def notification_handler(mock_logger, mock_mongo_client):
+    """
+    Create a real NotificationHandler instance with mocked MongoDB client.
+    
+    This fixture creates a real CoreNotificationHandler instance but mocks
+    the underlying MongoDB connection to avoid real database connections.
+    """
+    # Patch MongoClient where it's imported in the handlers module
+    with patch('sientia_do.notifications.handlers.MongoClient', return_value=mock_mongo_client):
+        handler = CoreNotificationHandler(
+            connection_string=TEST_MONGODB_CONNECTION_STRING,
+            database=TEST_DATABASE_NAME,
+            logger=mock_logger,
+            project_name='laborious',
+        )
+        yield handler
+        handler.shutdown()
+
+
+@pytest_asyncio.fixture
+def metrics_controller(mock_logger):
+    """Create a real MetricsController instance."""
+    return MetricsController(logger=mock_logger)
+
+
+@pytest_asyncio.fixture
+def mock_minio_repository():
+    """Mock MinIO repository for object storage operations."""
+    mock_repo = MagicMock()
+
+    # Provide at least valid parquet bytes so that MinioDataFramePayload.retrieve()
+    # can decode the payload if offloading is exercised in an integration scenario.
+    parquet_df = pd.DataFrame({'a': [1]})
+    parquet_buffer = BytesIO()
+    parquet_df.to_parquet(parquet_buffer, engine='pyarrow', index=True)
+    parquet_bytes = parquet_buffer.getvalue()
+    
+    # sientia_do MinioRepository API
+    mock_repo.bucket = 'test-bucket'
+    mock_repo.upload_file = AsyncMock(
+        side_effect=lambda file_bytes, relative_key, content_type='application/octet-stream', bucket=None, metadata=None: {
+            'minio_object_name': f'sientia/streamlit-connectors/{relative_key}',
+            'original_filename': relative_key.rsplit('/', 1)[-1],
+            'uploaded_at': '2024-01-01T00:00:00Z',
+            'sha256_hash': 'deadbeef',
+        }
+    )
+    mock_repo.download_file = AsyncMock(return_value=parquet_bytes)
+    mock_repo.list_objects = AsyncMock(return_value=[])
+    mock_repo.delete_file = AsyncMock()
+    mock_repo.close = MagicMock()
+    
+    return mock_repo
+
+
+@pytest_asyncio.fixture
+def mock_pi_web_api_repository():
+    """Mock PI Web API repository for PI Web API operations."""
+    mock_repo = MagicMock()
+
+    async def _write_value(web_ids, value, metadata=None, **kwargs):
+        """
+        Mirror successful PI writes: one response item per requested web_id.
+
+        write_pi_web_api_data passes the list into process_pi_web_api_response (not a
+        wrapped {'Items': ...} envelope).
+        """
+        return [{'WebId': wid, 'Errors': []} for wid in web_ids]
+
+    mock_repo.write_value = AsyncMock(side_effect=_write_value)
+    mock_repo.close = MagicMock()
+    return mock_repo
+
+@pytest_asyncio.fixture
+def mock_opc_repository():
+    """Mock OPC repository for OPC operations."""
+    mock_repo = MagicMock()
+    mock_repo.write_data = AsyncMock(
+        return_value=(True, {'response_time': 0.1})
+    )
+    mock_repo.disconnect = AsyncMock()
+    return mock_repo
+
+
+@pytest_asyncio.fixture
+async def opc_e2e_server():
+    """
+    In-process asyncua OPC UA server for E2E tests against OpcRepository.
+    """
+    server = OpcE2ETestServer()
+    await server.start()
+    try:
+        yield server
+    finally:
+        await server.stop()
+
+@pytest_asyncio.fixture
+def patch_create_engine(postgres_engine):
+    """Patch create_engine to return test postgres_engine."""
+    with patch('sientia_do.temporal.activities.postgres.create_engine', return_value=postgres_engine):
+        yield
+
+
+@pytest_asyncio.fixture
+def patch_minio_repository(mock_minio_repository):
+    """Patch MinioRepository to return mock."""
+    # Patch where Activities resolves the symbol (import binds the original class).
+    with patch('laborious.activities.activities.MinioRepository', return_value=mock_minio_repository):
+        yield
+
+@pytest_asyncio.fixture
+def patch_pi_web_api_repository(mock_pi_web_api_repository):
+    """Patch MLflowRepository to return mock."""
+    with patch('laborious.activities.api.PIWebAPIClient', return_value=mock_pi_web_api_repository):
+        yield
+
+@pytest_asyncio.fixture
+def mock_mlflow_models():
+    """Create mock models for MLflow load_model methods."""
+    # Mock transform model - returns DataFrame with same index as input
+    mock_transform_model = MagicMock()
+    def mock_transform_predict(data):
+        num_rows = max(len(data), 1) if hasattr(data, '__len__') else 1
+        print(data.to_csv())
+        print(data.index)
+        result = pd.DataFrame({
+            'feature_1': [0.234] * num_rows,
+            'feature_2': [0.783] * num_rows,
+        })
+        result.index = data.index
+        return result
+    mock_transform_model.predict = MagicMock(side_effect=mock_transform_predict)
+    
+    # Mock predict model - returns array/list of predictions
+    mock_predict_model = MagicMock()
+    def mock_predict_predict(data):
+        num_rows = max(len(data), 1) if hasattr(data, '__len__') else 1
+        return [0.5] * num_rows
+    mock_predict_model.predict = MagicMock(side_effect=mock_predict_predict)
+    
+    # Mock PyFuncModel for compressed models
+    mock_pyfunc_model = MagicMock()
+    mock_pyfunc_model._model_impl = MagicMock()
+    mock_pyfunc_model._model_impl.python_model = mock_transform_model
+    
+    return {
+        'transform_model': mock_transform_model,
+        'predict_model': mock_predict_model,
+        'pyfunc_model': mock_pyfunc_model,
+    }
+
+
+@pytest_asyncio.fixture
+def patch_mlflow(mock_mlflow_models):
+    """Patch mlflow module in repository with load_model mocks."""
+    mock_mlflow = MagicMock()
+    
+    # Mock sklearn.load_model
+    def mock_sklearn_load_model(model_uri):
+        if 'data_model' in model_uri or 'transform' in model_uri.lower():
+            return mock_mlflow_models['transform_model']
+        return mock_mlflow_models['predict_model']
+    mock_mlflow.sklearn = MagicMock()
+    mock_mlflow.sklearn.load_model = MagicMock(side_effect=mock_sklearn_load_model)
+    
+    # Mock pyfunc.load_model
+    def mock_pyfunc_load_model(model_uri):
+        if 'artifacts' in model_uri or 'tmp' in model_uri:
+            return mock_mlflow_models['pyfunc_model']
+        if 'data_model' in model_uri or 'transform' in model_uri.lower():
+            return mock_mlflow_models['transform_model']
+        return mock_mlflow_models['predict_model']
+    mock_mlflow.pyfunc = MagicMock()
+    mock_mlflow.pyfunc.load_model = MagicMock(side_effect=mock_pyfunc_load_model)
+    
+    # Mock pytorch.load_model
+    mock_mlflow.pytorch = MagicMock()
+    mock_mlflow.pytorch.load_model = MagicMock(return_value=mock_mlflow_models['predict_model'])
+    
+    # Mock other mlflow methods that might be called
+    mock_mlflow.set_tracking_uri = MagicMock()
+    mock_mlflow.get_run = MagicMock(return_value=MagicMock(info=MagicMock(artifact_uri='mlflow-artifacts:/test_run_id')))
+    mock_mlflow.tracking = MagicMock()
+    mock_mlflow.tracking.MlflowClient = MagicMock(return_value=MagicMock(
+        search_registered_models=MagicMock(return_value=[MagicMock(name='test_model')]),
+        search_model_versions=MagicMock(return_value=[MagicMock(
+            current_stage='Production',
+            version='1',
+            source='runs:/artifacts/test_run_id'
+        )])
+    ))
+    
+    with patch('laborious.utils.repository.model_repository.mlflow', new=mock_mlflow):
+        yield mock_mlflow
+
+
+@pytest_asyncio.fixture(scope='function')
+async def test_activities(
+    postgres_engine,
+    postgres_container,
+    mock_logger,
+    notification_handler,
+    metrics_controller,
+    mock_minio_repository,
+    patch_create_engine,
+    patch_minio_repository,
+    patch_mlflow,
+    patch_pi_web_api_repository,
+    mock_opc_repository
+):
+    """
+    Create Activities instance with test dependencies.
+
+    This fixture creates a real Activities instance with:
+    - PostgreSQL database (via testcontainers)
+    - Mocked MinIO client
+    - Real NotificationHandler and MetricsController (with mocked underlying services)
+    """
+    activities = Activities(
+        postgres_config={
+            'host': 'localhost',
+            'port': postgres_container.get_exposed_port(5432),
+            'user': 'test',
+            'password': 'test',
+            'dbname': 'test',
+            'min_connections': 1,
+            'max_connections': 5,
+        },
+        mlflow_config={
+            'host': 'http://localhost',
+            'port': '5000',
+            'username': 'test',
+            'password': 'test',
+        },
+        minio_config={
+            # Host:port only; Minio() prepends http(s):// from the secure flag.
+            'endpoint_url': 'localhost:9000',
+            'access_key': 'test',
+            'secret_key': 'test',
+            'default_bucket': 'test-bucket',
+            'retention_hours': 24,
+            'secure': False,
+        },
+        opc_config={},
+        pi_web_api_config={
+            'base_url': 'http://localhost:8080',
+            'auth_type': 'bearer',
+            'auth_token': 'test_token',
+        },
+        logger=mock_logger,
+        notification_handler=notification_handler,
+    )
+
+    activities.opc_repository = {
+        '1': mock_opc_repository,
+    }
+
+    try:
+        yield activities
+    finally:
+        # Cleanup - ALWAYS runs, even if test fails
+        await activities.shutdown()
+
+
+@pytest_asyncio.fixture(scope='function')
+async def test_activities_real_minio(
+    postgres_engine,
+    postgres_container,
+    minio_container,
+    mock_logger,
+    notification_handler,
+    metrics_controller,
+    patch_create_engine,
+    patch_mlflow,
+    patch_pi_web_api_repository,
+    mock_opc_repository,
+):
+    """
+    Activities with a real MinIO testcontainer (no MinioRepository patch) for offload tests.
+    """
+    minio_client = minio_container.get_client()
+    if not minio_client.bucket_exists('test-bucket'):
+        minio_client.make_bucket('test-bucket')
+    minio_port = minio_container.get_exposed_port(9000)
+    activities = Activities(
+        postgres_config={
+            'host': 'localhost',
+            'port': postgres_container.get_exposed_port(5432),
+            'user': 'test',
+            'password': 'test',
+            'dbname': 'test',
+            'min_connections': 1,
+            'max_connections': 5,
+        },
+        mlflow_config={
+            'host': 'http://localhost',
+            'port': '5000',
+            'username': 'test',
+            'password': 'test',
+        },
+        minio_config={
+            'endpoint_url': f'localhost:{minio_port}',
+            'access_key': 'minioadmin',
+            'secret_key': 'minioadmin',
+            'default_bucket': 'test-bucket',
+            'retention_hours': 24,
+            'secure': False,
+        },
+        opc_config={},
+        pi_web_api_config={
+            'base_url': 'http://localhost:8080',
+            'auth_type': 'bearer',
+            'auth_token': 'test_token',
+        },
+        logger=mock_logger,
+        notification_handler=notification_handler,
+    )
+    activities.opc_repository = {'1': mock_opc_repository}
+    try:
+        yield activities
+    finally:
+        await activities.shutdown()
+
+
+def _worker_activity_list(test_activities: Activities):
+    return [
+        test_activities.load_custom_query,
+        test_activities.load_query_with_minio_offload,
+        test_activities.cleanup_minio_objects_expired,
+        test_activities.input_gate,
+        test_activities.request_transform,
+        test_activities.mlflow_response_gate,
+        test_activities.mlflow_content_gate,
+        test_activities.request_predict,
+        test_activities.repeat_last_prediction,
+        test_activities.format_prediction,
+        test_activities.format_transformed_data,
+        test_activities.format_default_prediction,
+        test_activities.write_pi_web_api_data,
+        test_activities.write_opc_data,
+        test_activities.export_data_to_postgres,
+        test_activities.export_payload_to_postgres,
+        test_activities.write_metrics,
+    ]
+
+
+@pytest_asyncio.fixture(scope='function')
+async def temporal_test_env():
+    """Create Temporal test environment."""
+    env = await WorkflowEnvironment.start_time_skipping()
+    async with env:
+        yield env
+
+
+@pytest_asyncio.fixture(scope='function')
+async def temporal_worker(temporal_test_env, test_activities):
+    """Create Temporal worker with test activities."""
+    async with Worker(
+        temporal_test_env.client,
+        task_queue='test-queue',
+        workflows=[PredictionsBatch, PredictionProcess, FormatAndExportPrediction],
+        activities=_worker_activity_list(test_activities),
+    ) as worker:
+        yield worker
+
+
+@pytest_asyncio.fixture(scope='function')
+async def temporal_worker_real_minio(temporal_test_env, test_activities_real_minio):
+    """Temporal worker backed by Activities using real MinIO testcontainer."""
+    async with Worker(
+        temporal_test_env.client,
+        task_queue='test-queue',
+        workflows=[PredictionsBatch, PredictionProcess, FormatAndExportPrediction],
+        activities=_worker_activity_list(test_activities_real_minio),
+    ) as worker:
+        yield worker
+
+
+@pytest_asyncio.fixture(scope='function')
+async def test_activities_real_opc(
+    postgres_engine,
+    postgres_container,
+    opc_e2e_server: OpcE2ETestServer,
+    mock_logger,
+    notification_handler,
+    metrics_controller,
+    patch_create_engine,
+    patch_minio_repository,
+    patch_mlflow,
+    patch_pi_web_api_repository,
+):
+    """
+    Activities with a real OpcRepository connected to the in-process OPC UA server.
+    """
+    activities = Activities(
+        postgres_config={
+            'host': 'localhost',
+            'port': postgres_container.get_exposed_port(5432),
+            'user': 'test',
+            'password': 'test',
+            'dbname': 'test',
+            'min_connections': 1,
+            'max_connections': 5,
+        },
+        mlflow_config={
+            'host': 'http://localhost',
+            'port': '5000',
+            'username': 'test',
+            'password': 'test',
+        },
+        minio_config={
+            'endpoint_url': 'localhost:9000',
+            'access_key': 'test',
+            'secret_key': 'test',
+            'default_bucket': 'test-bucket',
+            'retention_hours': 24,
+            'secure': False,
+        },
+        opc_config={
+            '1': {
+                'id': '1',
+                'server_name': 'e2e-opc',
+                'url': opc_e2e_server.url,
+                'server_uri': opc_e2e_server.url,
+                'cert_path': None,
+                'private_key_path': None,
+                'server_cert_path': None,
+                'reconnection_interval': 0,
+            }
+        },
+        pi_web_api_config={
+            'base_url': 'http://localhost:8080',
+            'auth_type': 'bearer',
+            'auth_token': 'test_token',
+        },
+        logger=mock_logger,
+        notification_handler=notification_handler,
+    )
+    await activities.init_opc()
+    repo = activities.opc_repository['1']
+    assert repo._session_ready.is_set(), 'OPC E2E server connection failed during init_opc'
+    try:
+        yield activities
+    finally:
+        await activities.shutdown()
+
+
+@pytest_asyncio.fixture(scope='function')
+async def temporal_worker_real_opc(temporal_test_env, test_activities_real_opc):
+    """Temporal worker backed by Activities using the in-process OPC UA server."""
+    async with Worker(
+        temporal_test_env.client,
+        task_queue='test-queue',
+        workflows=[PredictionsBatch, PredictionProcess, FormatAndExportPrediction],
+        activities=_worker_activity_list(test_activities_real_opc),
+    ) as worker:
+        yield worker

e2e/helpers.py

@@ -0,0 +1,181 @@
+"""
+Shared helpers for E2E tests (Temporal workflows + PostgreSQL).
+"""
+
+import asyncio
+from datetime import datetime
+from decimal import Decimal
+from typing import Any
+
+from sqlalchemy import text
+from sqlalchemy.engine import Engine
+
+
+async def start_and_await_workflow(client, workflow_run, input_data: dict, workflow_id: str, timeout: float = 60.0):
+    """
+    Start a workflow and wait for its result.
+
+    Args:
+        client: Temporal client from WorkflowEnvironment.
+        workflow_run: Workflow run method (e.g. PredictionsBatch.run).
+        input_data: Workflow input payload.
+        workflow_id: Unique workflow id.
+        timeout: Max seconds to wait for completion.
+
+    Return:
+        Workflow result value.
+    """
+    handle = await client.start_workflow(
+        workflow_run,
+        input_data,
+        id=workflow_id,
+        task_queue='test-queue',
+    )
+    return await asyncio.wait_for(handle.result(), timeout=timeout)
+
+
+def insert_sample_data(postgres_engine: Engine, model_id: int, values: list[Any]) -> None:
+    """
+    Replace laborious_data rows for a model_id with one row per value (sensor_1..n).
+
+    Args:
+        postgres_engine: SQLAlchemy engine.
+        model_id: Model id column value.
+        values: Per-sensor values; use string 'NULL' for SQL NULL.
+    """
+    with postgres_engine.begin() as conn:
+        conn.execute(text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}'))
+        values_sql = []
+        for i, value in enumerate(values):
+            values_sql.append(f"""
+                ({model_id}, 'sensor_{i + 1}', {value}, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
+            """)
+        insert_sql = f"""
+            INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
+            VALUES
+                {', '.join(values_sql)}
+        """
+        conn.execute(text(insert_sql))
+
+
+def assert_prediction(
+    postgres_engine: Engine,
+    model_id: int,
+    prediction: float = 0.5,
+    prediction_confidence: int | Decimal = 0,
+    prediction_status: str = 'Good',
+    comments: str | None = None,
+    comments_contains: str | None = None,
+) -> None:
+    """
+    Assert exactly one prediction row exists for model_id with expected columns.
+
+    Args:
+        postgres_engine: SQLAlchemy engine.
+        model_id: Expected model_id.
+        prediction: Expected prediction value.
+        prediction_confidence: Expected confidence (int or Decimal for numeric column).
+        prediction_status: Expected status string.
+        comments: Expected exact comments string (optional).
+        comments_contains: Substring expected in comments when queued (optional).
+    """
+    import pytest
+
+    with postgres_engine.connect() as conn:
+        result_query = conn.execute(
+            text(
+                f'SELECT model_id, prediction, prediction_confidence, prediction_status, comments '
+                f'FROM predictions_schema.predictions WHERE model_id = {model_id} '
+                f'ORDER BY created_at ASC'
+            )
+        )
+        prediction_rows = result_query.fetchall()
+        assert len(prediction_rows) == 1, f'Expected one prediction record, got {len(prediction_rows)}'
+        row = prediction_rows[0]
+        assert row[0] == model_id, f'Expected model_id={model_id}, got {row[0]}'
+        assert row[1] == prediction or Decimal(str(row[1])) == Decimal(str(prediction)), (
+            f'Expected prediction={prediction}, got {row[1]}'
+        )
+        assert row[2] == prediction_confidence or Decimal(str(row[2])) == Decimal(
+            str(prediction_confidence)
+        ), f'Expected prediction_confidence={prediction_confidence}, got {row[2]}'
+        assert row[3] == prediction_status, f"Expected prediction_status='{prediction_status}', got {row[3]}"
+        if comments is not None:
+            assert row[4] == comments, f"Expected comments='{comments}', got {row[4]}"
+        if comments_contains is not None:
+            assert comments_contains in row[4], (
+                f"Expected comments to contain '{comments_contains}', got {row[4]}"
+            )
+
+
+def assert_continue(
+    postgres_engine: Engine,
+    model_id: int,
+    prediction_confidence: Decimal = Decimal(2),
+    comments: str = 'Input data with bad quality',
+) -> None:
+    """Assert one default-style prediction row after CONTINUE gate path."""
+    with postgres_engine.connect() as conn:
+        result_query = conn.execute(
+            text(
+                f'SELECT model_id, prediction, prediction_confidence, prediction_status, comments '
+                f'FROM predictions_schema.predictions WHERE model_id = {model_id}'
+            )
+        )
+        prediction_rows = result_query.fetchall()
+        assert len(prediction_rows) == 1, 'Expected one prediction record despite warnings'
+        row = prediction_rows[0]
+        assert row[1] == 0, f'Expected prediction=0, got {row[1]}'
+        assert row[2] == prediction_confidence, (
+            f'Expected prediction_confidence={prediction_confidence}, got {row[2]}'
+        )
+        assert row[3] == 'Bad', f"Expected prediction_status='Bad', got {row[3]}"
+        assert row[4] == comments, f"Expected comments='{comments}', got {row[4]}"
+
+
+def assert_stop(postgres_engine: Engine, model_id: int) -> None:
+    """Assert no prediction rows for model_id."""
+    import pytest
+
+    with postgres_engine.connect() as conn:
+        result_query = conn.execute(
+            text(f'SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = {model_id}')
+        )
+        count = result_query.scalar()
+        assert count == 0, f'Expected no predictions, but found {count} records'
+
+
+def assert_repeat(postgres_engine: Engine, model_id: int, last_prediction: tuple) -> None:
+    """
+    Assert two prediction rows for model_id both match last_prediction.
+
+    Rows are compared in created_at order for stability.
+
+    Args:
+        postgres_engine: SQLAlchemy engine.
+        model_id: Model id.
+        last_prediction: Tuple (model_id, prediction, confidence, status) to match both rows.
+    """
+    import pytest
+
+    with postgres_engine.connect() as conn:
+        result_query = conn.execute(
+            text(
+                f'SELECT model_id, prediction, prediction_confidence, prediction_status '
+                f'FROM predictions_schema.predictions WHERE model_id = {model_id} '
+                f'ORDER BY created_at ASC'
+            )
+        )
+        prediction_rows = result_query.fetchall()
+        assert len(prediction_rows) == 2, 'Expected two prediction records'
+        assert prediction_rows[0] == last_prediction, (
+            f'Expected first row {last_prediction}, got {prediction_rows[0]}'
+        )
+        assert prediction_rows[1] == last_prediction, (
+            f'Expected second row {last_prediction}, got {prediction_rows[1]}'
+        )
+
+
+def make_workflow_id(prefix: str) -> str:
+    """Build a unique workflow id using a prefix and current timestamp."""
+    return f'{prefix}-{datetime.now().timestamp()}'

e2e/opc_test_server.py

@@ -0,0 +1,189 @@
+"""
+In-process OPC UA server for E2E tests (asyncua).
+
+Provides writable prediction/confidence nodes and optional write faults
+(Tier-1 BadSessionIdInvalid via PreWrite callback).
+"""
+
+from __future__ import annotations
+
+import socket
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from asyncua import Server, ua
+from asyncua.common.callback import CallbackType
+from asyncua.common.utils import ServiceError
+
+if TYPE_CHECKING:
+    from asyncua.common.node import Node
+
+
+UNKNOWN_NODE_ID = 'ns=99;i=9999'
+
+
+@dataclass(frozen=True)
+class OpcE2ENodeIds:
+    """NodeId strings used in opc_output_config for E2E workflows."""
+
+    prediction: str
+    confidence: str
+    unknown: str = UNKNOWN_NODE_ID
+
+
+class OpcE2ETestServer:
+    """
+    Ephemeral asyncua server with Laborious E2E variables and controllable faults.
+
+    Args:
+        host: Bind address (default 127.0.0.1).
+    """
+
+    def __init__(self, host: str = '127.0.0.1') -> None:
+        self._host = host
+        self._server: Server | None = None
+        self._prediction_node: Node | None = None
+        self._confidence_node: Node | None = None
+        self._session_bad_on_write = False
+        self._url: str | None = None
+        self._node_ids: OpcE2ENodeIds | None = None
+
+    @property
+    def url(self) -> str:
+        if self._url is None:
+            raise RuntimeError('OPC E2E server is not started')
+        return self._url
+
+    @property
+    def node_ids(self) -> OpcE2ENodeIds:
+        if self._node_ids is None:
+            raise RuntimeError('OPC E2E server is not started')
+        return self._node_ids
+
+    def set_session_bad_on_write(self, enabled: bool) -> None:
+        """
+        When enabled, every client Write is rejected with BadSessionIdInvalid.
+
+        Args:
+            enabled (bool): Turn Tier-1 session fault injection on or off.
+        """
+        self._session_bad_on_write = enabled
+
+    async def start(self) -> OpcE2ENodeIds:
+        """
+        Start the OPC UA server on a free TCP port.
+
+        Return:
+            OpcE2ENodeIds: NodeId strings for prediction and confidence tags.
+        """
+        port = _free_port(self._host)
+        self._url = f'opc.tcp://{self._host}:{port}/freeopcua/server/'
+
+        server = Server()
+        server.set_endpoint(self._url)
+        await server.init()
+        server.iserver.callback_service.addListener(
+            CallbackType.PreWrite,
+            self._pre_write_callback,
+        )
+
+        idx = await server.register_namespace('http://sientia.test/laborious-e2e')
+        e2e_object = await server.nodes.objects.add_object(idx, 'LaboriousE2E')
+        prediction = await e2e_object.add_variable(
+            idx,
+            'Prediction',
+            ua.Variant(0.0, ua.VariantType.Float),
+        )
+        confidence = await e2e_object.add_variable(
+            idx,
+            'Confidence',
+            ua.Variant(0.0, ua.VariantType.Float),
+        )
+        await prediction.set_writable()
+        await confidence.set_writable()
+
+        await server.start()
+        self._server = server
+        self._prediction_node = prediction
+        self._confidence_node = confidence
+        self._node_ids = OpcE2ENodeIds(
+            prediction=prediction.nodeid.to_string(),
+            confidence=confidence.nodeid.to_string(),
+        )
+        return self._node_ids
+
+    async def stop(self) -> None:
+        """Stop the OPC UA server and release the listening port."""
+        if self._server is not None:
+            await self._server.stop()
+            self._server = None
+        self._prediction_node = None
+        self._confidence_node = None
+        self._url = None
+        self._node_ids = None
+        self._session_bad_on_write = False
+
+    async def read_prediction(self) -> float:
+        """
+        Read the current prediction variable value from the address space.
+
+        Return:
+            float: Stored prediction value.
+        """
+        if self._prediction_node is None:
+            raise RuntimeError('OPC E2E server is not started')
+        value = await self._prediction_node.read_value()
+        return float(value)
+
+    async def read_confidence(self) -> float:
+        """
+        Read the current confidence variable value from the address space.
+
+        Return:
+            float: Stored confidence value.
+        """
+        if self._confidence_node is None:
+            raise RuntimeError('OPC E2E server is not started')
+        value = await self._confidence_node.read_value()
+        return float(value)
+
+    async def _pre_write_callback(self, _event, _service) -> None:
+        if self._session_bad_on_write:
+            raise ServiceError(ua.StatusCodes.BadSessionIdInvalid)
+
+
+def _free_port(host: str) -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+        sock.bind((host, 0))
+        return int(sock.getsockname()[1])
+
+
+def build_opc_output_config(
+    node_ids: OpcE2ENodeIds,
+    *,
+    prediction_tag: str | None = None,
+    confidence_tag: str | None = None,
+    prediction_only: bool = False,
+    server_key: str = '1',
+) -> dict[str, dict]:
+    """
+    Build opc_output_config for PredictionsBatch using real server NodeIds.
+
+    Args:
+        node_ids (OpcE2ENodeIds): Node ids from OpcE2ETestServer.
+        prediction_tag (str | None): Override prediction NodeId (default: node_ids.prediction).
+        confidence_tag (str | None): Override confidence NodeId (default: node_ids.confidence).
+        prediction_only (bool): When True, omit confidence_tags (single write per activity).
+        server_key (str): OPC server id key in opc_output_config.
+
+    Return:
+        dict: opc_output_config payload for workflow input.
+    """
+    pred = prediction_tag if prediction_tag is not None else node_ids.prediction
+    conf = confidence_tag if confidence_tag is not None else node_ids.confidence
+    server_config: dict = {
+        'prediction_tags': {pred: {'data_type': 'float'}},
+    }
+    if not prediction_only:
+        server_config['confidence_tags'] = {conf: {'data_type': 'float'}}
+    return {server_key: server_config}

e2e/scenarios.md

@@ -0,0 +1,564 @@
+# Test Scenarios for Predictions Batch Workflow
+
+This document describes all possible test scenarios for the `predictions_batch` workflow and its child workflows `prediction_process` and `format_and_export_prediction`.
+
+## Running automated E2E tests (`e2e/`)
+
+- **Runtime**: Docker (or a Docker-compatible daemon) must be available so [testcontainers](https://testcontainers.com/) can start **PostgreSQL** and **MinIO** containers.
+- **Dependencies**: install dev requirements (includes `testcontainers[postgres,minio]`).
+- **Invocation**: run only integration-marked tests, for example: `pytest e2e/ -m integration`.
+- **MinIO tests**: `e2e/test_minio_offload.py` exercises real S3 uploads; other E2E modules continue to mock MinIO on the worker used by most scenarios.
+- **OPC tests (real server)**: `e2e/test_opc_real_server.py` uses an in-process **asyncua** server and real `OpcRepository` (`test_activities_real_opc`). Scenarios 3.1.2, 3.2.2, 3.2.4, and 3.2.5 are covered there. Other E2E modules keep the OPC mock.
+- Run only OPC real-server tests: `pytest e2e/test_opc_real_server.py -m "integration and opc"`.
+
+## Workflow Overview
+
+The `predictions_batch` workflow:
+1. Loads data using a custom SQL query
+2. Prepares prediction configuration
+3. Delegates to `prediction_process` child workflow which:
+   - Retrieves last timestamp for incremental processing
+   - Applies input data quality gates
+   - Executes MLFlow transform operation
+   - Validates transform response
+   - Executes MLFlow predict operation
+   - Validates predict response
+   - Delegates to `format_and_export_prediction` child workflow
+4. The `format_and_export_prediction` workflow:
+   - Formats prediction data (normal or default)
+   - Exports to PI Web API (optional)
+   - Exports to OPC server (optional)
+   - Exports to PostgreSQL
+   - Writes metrics
+
+---
+
+## 1. Predictions Batch - Main Workflow Scenarios
+
+### 1.1 Success Scenarios
+
+#### Scenario 1.1.1: Happy Path - Complete Success
+**Description**: Workflow completes successfully with valid SQL query and all activities succeed
+
+**Input**:
+- Valid `schedule_name`, `model_name`, `model_id`
+- Valid `query` returning non-empty DataFrame
+- Valid `schema`, `table_name`, `transform_table_name`
+- Optional `datetime_columns` for timestamp parsing
+- Optional `input_filters`, `mlflow_transform_filters`, `mlflow_predict_filters`
+- Optional `path_priority`, `opc_output_config`, `pi_web_api_output_config`
+
+**Expected Behavior**:
+- `load_custom_query` returns DataFrame with data
+- Workflow prepares prediction input with all configurations
+- `prediction_process` child workflow executes successfully
+- All gates pass with no issues
+- Transform and predict operations succeed
+- Data exported to PostgreSQL
+- Metrics written
+
+**Assertions**:
+- SQL query executed once
+- `prediction_process` workflow called with correct parameters
+- Data exists in PostgreSQL (predictions table)
+- Metrics recorded
+- No errors raised
+
+---
+
+### 1.2 Error Scenarios
+
+#### Scenario 1.2.1: SQL Query Execution Error
+**Description**: SQL query fails due to syntax error or connection issue
+
+**Input**:
+- Invalid SQL query (syntax error)
+- Or database connection unavailable
+
+**Expected Behavior**:
+- `load_custom_query` raises exception (caught by Temporal retry policy)
+- Notification sent with SQL error details
+- After retries, activity may return empty data or workflow may fail
+- If empty data returned, workflow completes with early exit via input gate
+
+**Assertions**:
+- Error notification sent
+- Workflow completes (either fails or exits early)
+- No data in predictions table
+
+---
+
+#### Scenario 1.2.2: Missing Required Parameters
+**Description**: Essential parameters missing from input
+
+**Input**:
+- Missing `query` or `model_id` or `schema` or `table_name`
+
+**Expected Behavior**:
+- Workflow or activity raises KeyError or validation error
+- Workflow fails immediately
+
+**Assertions**:
+- Workflow fails with parameter error
+- Error notification sent
+- No child workflow called
+
+---
+
+#### Scenario 1.2.3: Invalid Datetime Column Specification
+**Description**: Datetime column specified doesn't exist in query results
+
+**Input**:
+- `datetime_columns: ['nonexistent_column']`
+- Query results don't have this column
+
+**Expected Behavior**:
+- `load_custom_query` may raise KeyError or warning
+- Depending on implementation, workflow may fail or continue
+- Error notification sent
+
+**Assertions**:
+- Error raised or warning logged
+- Workflow behavior depends on error handling policy
+
+---
+
+## 2. Prediction Process - Child Workflow Scenarios
+
+### 2.1 Input gate Early Exit Scenarios
+
+#### Scenario 2.1.1: Input Gate Triggers CONTINUE
+**Description**: Input gate determines data should use previous prediction
+
+**Input**:
+- Data that should continue with input data as prediction
+- `input_filters` configured with `POLICY: 'CONTINUE'`
+- `path_priority` includes CONTINUE
+
+**Expected Behavior**:
+- `input_gate` returns `path_flag='CONTINUE'`
+- `path_flag_handler` calls export workflow with input data directly
+- MLFlow transform and predict skipped
+- Data exported as-is
+
+**Assertions**:
+- `input_gate` called
+- MLFlow operations NOT called
+- Export workflow called with original data
+- Workflow completes
+
+
+#### Scenario 2.1.2: Input Gate Triggers STOP
+**Description**: Input data quality gate fails with STOP policy
+
+**Input**:
+- Data with EMPTY_DATA or other critical issues
+- `input_filters` configured with `POLICY: 'STOP'`
+
+**Expected Behavior**:
+- `input_gate` returns `path_flag='STOP'`
+- `path_flag_handler` detects STOP
+- Workflow returns early without calling MLFlow
+- No prediction exported
+
+**Assertions**:
+- `input_gate` called
+- `path_flag_handler` returns True (early exit)
+- MLFlow transform NOT called
+- Export workflow NOT called
+- Workflow completes without error
+
+
+#### Scenario 2.1.3: Input Gate Triggers REPEAT
+**Description**: Input gate determines data should repeat last prediction
+
+**Input**:
+- Data with quality issues that require using previous prediction
+- `input_filters` configured with `POLICY: 'REPEAT'`
+- `path_priority` includes REPEAT
+
+**Expected Behavior**:
+- `input_gate` returns `path_flag='REPEAT'`
+- `path_flag_handler` calls `repeat_last_prediction` activity
+- MLFlow transform and predict skipped
+- Last prediction repeated and exported
+
+**Assertions**:
+- `input_gate` called
+- MLFlow operations NOT called
+- `repeat_last_prediction` activity called
+- Workflow completes
+
+---
+
+### 2.2 Transform gate Early Exit Scenarios
+
+#### Scenario 2.2.1: Transform Gate Triggers CONTINUE
+**Description**: Transform response gate determines data should continue despite issues
+
+**Input**:
+- Valid input data
+- Transform response has quality issues but policy is CONTINUE
+- `mlflow_transform_filters` configured with `POLICY: 'CONTINUE'`
+- `path_priority` includes CONTINUE
+
+**Expected Behavior**:
+- `request_transform` succeeds
+- `mlflow_response_gate` for transform returns `path_flag='CONTINUE'`
+- `path_flag_handler` calls export workflow with transform data
+- MLFlow predict skipped
+- Transform data exported as-is
+
+**Assertions**:
+- Transform completed
+- `mlflow_response_gate` called for transform
+- MLFlow predict NOT called
+- Export workflow called with transform data
+- Workflow completes
+
+---
+
+#### Scenario 2.2.2: Transform Gate Triggers STOP
+**Description**: Transform response validation fails with STOP policy
+
+**Input**:
+- Valid input data
+- Transform response has critical errors
+- `mlflow_transform_filters` configured with `POLICY: 'STOP'`
+
+**Expected Behavior**:
+- `request_transform` succeeds but response invalid
+- `mlflow_response_gate` for transform returns `path_flag='STOP'`
+- Workflow exits without calling predict or export
+
+**Assertions**:
+- Transform completed but validation failed
+- `mlflow_response_gate` called for transform
+- MLFlow predict NOT called
+- Export workflow NOT called
+- Workflow completes without error
+
+---
+
+#### Scenario 2.2.3: Transform Gate Triggers REPEAT
+**Description**: Transform response gate determines data should repeat last prediction
+
+**Input**:
+- Valid input data
+- Transform response has quality issues that require using previous prediction
+- `mlflow_transform_filters` configured with `POLICY: 'REPEAT'`
+- `path_priority` includes REPEAT
+
+**Expected Behavior**:
+- `request_transform` succeeds but response has issues
+- `mlflow_response_gate` for transform returns `path_flag='REPEAT'`
+- `path_flag_handler` calls `repeat_last_prediction` activity
+- MLFlow predict skipped
+- Last prediction repeated and exported
+
+**Assertions**:
+- Transform completed but validation triggered REPEAT
+- `mlflow_response_gate` called for transform
+- MLFlow predict NOT called
+- `repeat_last_prediction` activity called
+- Workflow completes
+
+---
+
+### 2.3 Predict gate Early Exit Scenarios
+
+#### Scenario 2.3.1: Predict Gate Triggers CONTINUE
+**Description**: Predict response gate determines data should continue despite issues
+
+**Input**:
+- Valid input and transform data
+- Predict response has quality issues but policy is CONTINUE
+- `mlflow_predict_filters` configured with `POLICY: 'CONTINUE'`
+- `path_priority` includes CONTINUE
+
+**Expected Behavior**:
+- `request_predict` succeeds
+- `mlflow_response_gate` for predict returns `path_flag='CONTINUE'`
+- `path_flag_handler` calls export workflow with predict data
+- Prediction exported despite quality issues
+
+**Assertions**:
+- Transform and predict completed
+- `mlflow_response_gate` called for predict
+- Export workflow called with predict data
+- Workflow completes
+
+---
+
+#### Scenario 2.3.2: Predict Gate Triggers STOP
+**Description**: Prediction validation fails with STOP policy
+
+**Input**:
+- Valid input and transform
+- Predict response has critical errors
+- `mlflow_predict_filters` configured with `POLICY: 'STOP'`
+
+**Expected Behavior**:
+- `request_predict` succeeds but response invalid
+- `mlflow_response_gate` for predict returns `path_flag='STOP'`
+- Workflow exits without export
+
+**Assertions**:
+- Transform completed
+- Predict completed but validation failed
+- Export workflow NOT called
+- Workflow completes without error
+
+---
+
+#### Scenario 2.3.3: Predict Gate Triggers REPEAT
+**Description**: Predict response gate determines data should repeat last prediction
+
+**Input**:
+- Valid input and transform data
+- Predict response has quality issues that require using previous prediction
+- `mlflow_predict_filters` configured with `POLICY: 'REPEAT'`
+- `path_priority` includes REPEAT
+
+**Expected Behavior**:
+- `request_predict` succeeds but response has issues
+- `mlflow_response_gate` for predict returns `path_flag='REPEAT'`
+- `path_flag_handler` calls `repeat_last_prediction` activity
+- Last prediction repeated and exported
+
+**Assertions**:
+- Transform and predict completed but validation triggered REPEAT
+- `mlflow_response_gate` called for predict
+- `repeat_last_prediction` activity called
+- Export workflow NOT called with current prediction
+- Workflow completes
+
+---
+
+## 3. Format and Export Prediction - Child Workflow Scenarios
+
+### 3.1 Success Scenarios
+
+#### Scenario 3.1.1: Default Prediction Export
+**Description**: Error prediction path creates default prediction
+
+**Input**:
+- `path_flag: 'ERROR'` or other non-None value (not STOP/CONTINUE/REPEAT)
+- `comment` provided with error details
+
+**Expected Behavior**:
+- `format_default_prediction` called instead of `format_prediction`
+- Default prediction created with error metadata
+- Exported to PostgreSQL only
+- Transformed data NOT processed
+- Metrics written
+
+**Assertions**:
+- `format_default_prediction` called
+- `format_prediction` NOT called
+- `format_transformed_data` NOT called
+- One PostgreSQL export only
+- Default values in prediction data
+- Comment included
+
+---
+
+#### Scenario 3.1.2: Export with OPC only
+**Description**: Export to PostgreSQL and OPC server only (no PI Web API)
+
+**Input**:
+- `path_flag: None`
+- `opc_output_config` configured with valid OPC settings
+- `pi_web_api_output_config: None` or `{}`
+
+**Expected Behavior**:
+- Normal formatting
+- PostgreSQL export executed
+- OPC export executed
+- PI Web API activity skipped
+- Metrics written with OPC metrics
+
+**Assertions**:
+- PI Web API activity NOT called
+- OPC activity called
+- PostgreSQL export called
+- Metrics written with `opc_metrics` populated
+
+---
+
+#### Scenario 3.1.3: Export with PI Web API only
+**Description**: Export to PostgreSQL and PI Web API only (no OPC)
+
+**Input**:
+- `path_flag: None`
+- `pi_web_api_output_config` configured with valid PI Web API settings
+- `opc_output_config: None` or `{}`
+
+**Expected Behavior**:
+- Normal formatting
+- PostgreSQL export executed
+- PI Web API export executed
+- OPC activity skipped
+- Metrics written without OPC metrics
+
+**Assertions**:
+- OPC activity NOT called
+- PI Web API activity called
+- PostgreSQL export called
+- Metrics written with empty `opc_metrics`
+
+---
+
+#### Scenario 3.1.4: Export Without Optional Outputs
+**Description**: Export only to PostgreSQL (no OPC or PI Web API)
+
+**Input**:
+- `path_flag: None`
+- `opc_output_config: None` or `{}`
+- `pi_web_api_output_config: None` or `{}`
+
+**Expected Behavior**:
+- Normal formatting
+- Only PostgreSQL export executed
+- OPC and PI Web API activities skipped
+- Metrics written without OPC metrics
+
+**Assertions**:
+- PI Web API activity NOT called
+- OPC activity NOT called
+- PostgreSQL export called
+- Metrics written with empty `opc_metrics`
+
+---
+
+#### Scenario 3.1.5: Export Without Transformed Data
+**Description**: Only prediction exported, no transform table
+
+**Input**:
+- `path_flag: None`
+- `transformed_data: None` or `save_transform: False`
+- `opc_output_config: None` or `{}`
+- `pi_web_api_output_config: None` or `{}`
+
+**Expected Behavior**:
+- Only prediction formatted and exported
+- Transform export skipped
+- Single PostgreSQL write
+
+**Assertions**:
+- `format_transformed_data` NOT called
+- One PostgreSQL export
+- Transform table remains empty
+
+---
+
+### 3.2 Error Scenarios
+
+These paths do **not** rely on Temporal activity retries for export failures: the write activities run once, errors are handled inside the activity, and the **workflow completes successfully** with degraded metadata on the persisted prediction (`prediction_confidence` and `comments`).
+
+#### Scenario 3.2.1: PI Web API Write Error
+**Description**: PI Web API export fails
+
+**Input**:
+- Valid prediction
+- PI Web API service unavailable or invalid config
+
+**Expected Behavior**:
+- `write_pi_web_api_data` surfaces the failure (exception handled in the activity layer)
+- Notification may be sent
+- Workflow **completes** (does not fail)
+- Prediction row is still written to PostgreSQL with error confidence **13** and a comment describing the PI error
+- Subsequent steps (e.g. OPC, Postgres) still run per workflow order with the updated prediction payload
+
+**Assertions**:
+- PI Web API error notification sent (when applicable)
+- Workflow completes
+- PostgreSQL contains the prediction with `prediction_confidence` 13 and expected `comments`
+
+---
+
+#### Scenario 3.2.2: OPC Write Error
+**Description**: OPC server write fails
+
+**Input**:
+- Valid prediction
+- OPC server unavailable or invalid configuration
+
+**Expected Behavior**:
+- `write_opc_data` reports failure without aborting the workflow
+- Notification may be sent
+- Workflow **completes** (does not fail)
+- Prediction row is written to PostgreSQL with OPC error confidence **12** and a comment indicating OPC write issues
+
+**Assertions**:
+- OPC error notification sent (when applicable)
+- Workflow completes
+- PostgreSQL contains the prediction with `prediction_confidence` 12 and expected `comments`
+
+---
+
+#### Scenario 3.2.4: OPC Session / Channel Bad* (Tier-1)
+**Description**: OPC write fails with a Tier-1 session or channel status (e.g. `BadSessionIdInvalid`) while transport may still appear open on the client
+
+**Input**:
+- Valid prediction and OPC output config
+- Mock or server returning Tier-1 `UaStatusCodeError` on write (no write retry in the same activity)
+
+**Expected Behavior**:
+- `write_opc_data` fails forward for affected tags; background reconnect may be scheduled if `OPC_RECONNECTION_INTERVAL` allows
+- Workflow **completes**
+- PostgreSQL row uses **`prediction_confidence` 14** and comment prefix `OPC UA session/channel error:` (including OPC status name)
+- `opc_write_attempts_total` records `result=BadSessionIdInvalid` (or matching status); no second write attempt in the same activity
+
+**Assertions**:
+- Workflow completes
+- `prediction_confidence = 14`
+- `comments` matches `OPC UA session/channel error:%`
+- Generic OPC error confidence **12** is not used for this case
+
+**Reference**: [docs/opc-communication.md](../docs/opc-communication.md), plan `.cursor/plans/opc_bad_reconnect_ac4c6045.plan.md`
+
+---
+
+#### Scenario 3.2.5: OPC Write Blocked During Reconnect
+**Description**: A write is attempted while the repository is reconnecting (session not ready)
+
+**Input**:
+- Valid prediction
+- Simulated slow reconnect (e.g. delayed `connect`) or concurrent writes where the first triggers reconnect
+
+**Expected Behavior**:
+- Second write (or parallel write) is rejected **immediately** when reconnect is in progress or `_session_ready` is cleared — **without** calling `write_value`
+- No wait/sleep on the write path; no duplicate `connect` from parallel writers (connection lock)
+- `prediction_confidence = 14`, `comments = OPC UA reconnect in progress` (distinguish from Tier-1 `Bad*` via comment prefix in SQL)
+
+**Assertions**:
+- At most one reconnect sequence (`disconnect` + `connect`) for the overlapping window
+- No write retry after failure
+- Tests in `test_opc_repository` (unit) and optional e2e in `test_predictions_batch_format_export.py`
+
+---
+
+#### Scenario 3.2.3: PI Web API Partial Write Error
+**Description**: Two prediction tags attempt to be written to PI Web API, but only one succeeds
+
+**Input**:
+- Valid prediction
+- Two prediction tags configured
+- PI Web API returns partial success (one tag succeeds, one fails)
+
+**Expected Behavior**:
+- `write_pi_web_api_data` processes response
+- `process_pi_web_api_response` detects partial failure
+- Error confidence set (13)
+- Notification sent for failed tag
+- Workflow completes with error confidence (single activity attempt; no retry loop)
+
+**Assertions**:
+- One tag written successfully
+- One tag failed
+- Error confidence set in prediction
+- Error notification sent
+- Workflow completes
+
+---

e2e/test_child_workflows_e2e.py

@@ -0,0 +1,74 @@
+"""
+Direct E2E execution of child workflows (smaller surface than PredictionsBatch).
+"""
+
+from decimal import Decimal
+
+import pytest
+from sqlalchemy import text
+from temporalio.testing import WorkflowEnvironment
+from temporalio.worker import Worker
+
+from e2e.helpers import make_workflow_id, start_and_await_workflow
+from laborious.workflows.sub_workflows.format_and_export_prediction import FormatAndExportPrediction
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_format_and_export_prediction_default_path_e2e(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker: Worker,
+    postgres_engine,
+):
+    """
+    Run FormatAndExportPrediction with path_flag set (format_default_prediction path).
+    """
+    client = temporal_test_env.client
+    model_id = 401
+    with postgres_engine.begin() as conn:
+        conn.execute(text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}'))
+
+    metadata = {
+        'metadata': {
+            'model_id': model_id,
+            'model_name': 'test_model',
+            'schedule_name': 'test-schedule',
+            'workflow_name': 'subworkflow.format_and_export_prediction',
+        }
+    }
+    input_data = {
+        'metadata': metadata,
+        'path_flag': 'CONTINUE',
+        'data': {'last_timestamp': '2024-01-01 12:00:00+00:00'},
+        'prediction_confidence': 2,
+        'timestamp': '2024-01-01 12:00:00+00:00',
+        'model_id': model_id,
+        'model_name': 'test_model',
+        'schema': 'predictions_schema',
+        'table_name': 'predictions',
+        'transform_table_name': 'transformed_data',
+        'comment': 'e2e child workflow default path',
+        'opc_output_config': {},
+        'pi_web_api_output_config': {},
+        'prediction_store_policy': 'lts:1',
+    }
+
+    await start_and_await_workflow(
+        client,
+        FormatAndExportPrediction.run,
+        input_data,
+        make_workflow_id('e2e-format-export-child'),
+    )
+
+    with postgres_engine.connect() as conn:
+        row = conn.execute(
+            text(
+                f'SELECT prediction, prediction_confidence, prediction_status, comments '
+                f'FROM predictions_schema.predictions WHERE model_id = {model_id}'
+            )
+        ).fetchone()
+    assert row is not None
+    assert row[0] == 0
+    assert row[1] == Decimal(2)
+    assert row[2] == 'Bad'
+    assert row[3] == 'e2e child workflow default path'

e2e/test_minio_offload.py

@@ -0,0 +1,124 @@
+"""
+E2E-style tests for MinIO offload using a real MinIO testcontainer.
+"""
+
+from unittest.mock import patch
+
+import pytest
+from sqlalchemy import text
+from temporalio.testing import WorkflowEnvironment
+from temporalio.worker import Worker
+
+from e2e.helpers import insert_sample_data, make_workflow_id, start_and_await_workflow
+from laborious.activities.activities import Activities
+from laborious.utils.models import minio_dataframe_payload as mdp
+from laborious.workflows.predictions_batch import PredictionsBatch
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_load_query_with_minio_offload_writes_object_to_bucket(
+    postgres_engine,
+    minio_container,
+    test_activities_real_minio: Activities,
+):
+    """
+    With a tiny offload threshold, query results are uploaded as Parquet to MinIO.
+
+    Uses real MinioRepository against testcontainers MinIO (no MinIO mock).
+    """
+    model_id = 501
+    with postgres_engine.begin() as conn:
+        conn.execute(text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}'))
+    insert_sample_data(postgres_engine, model_id, [1.0, 2.0])
+
+    metadata = {
+        'metadata': {
+            'schedule_name': 'test-schedule',
+            'model_name': 'test_model',
+            'model_id': model_id,
+            'workflow_name': 'predictions_batch',
+        }
+    }
+    with patch.object(mdp, 'OFFLOAD_THRESHOLD_BYTES', 1):
+        payload = await test_activities_real_minio.load_query_with_minio_offload(
+            {
+                **metadata,
+                'query': (
+                    'SELECT timestamp, variable, value, created_at '
+                    f'FROM predictions_schema.laborious_data WHERE model_id = {model_id}'
+                ),
+                'model_name': 'test_model',
+                'datetime_columns': ['timestamp', 'created_at'],
+            }
+        )
+    assert payload.object_key, 'offloaded payload must reference a MinIO object'
+    assert payload.data is None or payload.data == {}, 'large payloads should not inline tabular dict'
+
+    df = await payload.retrieve(test_activities_real_minio.minio_repository, metadata['metadata'])
+    assert len(df) >= 1
+
+    client = minio_container.get_client()
+    listed = list(client.list_objects('test-bucket', recursive=True))
+    names = [getattr(o, 'object_name', None) or getattr(o, '_object_name', '') for o in listed]
+    assert any(n and 'prediction_datasets' in n for n in names), f'unexpected object listing: {names!r}'
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_predictions_batch_with_minio_offload_path(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker_real_minio: Worker,
+    postgres_engine,
+    test_activities_real_minio: Activities,
+):
+    """
+    Full PredictionsBatch run with offload: load step stores Parquet in MinIO; pipeline completes.
+    """
+    model_id = 502
+    with postgres_engine.begin() as conn:
+        conn.execute(text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}'))
+        conn.execute(text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}'))
+        conn.execute(text(f'DELETE FROM predictions_schema.transformed_data WHERE model_id = {model_id}'))
+    insert_sample_data(postgres_engine, model_id, [10.0, 20.0, 30.0])
+
+    input_data = {
+        'schedule_name': 'test-schedule',
+        'model_name': 'test_model',
+        'model_id': model_id,
+        'query': (
+            'SELECT timestamp, variable, value, created_at '
+            f'FROM predictions_schema.laborious_data WHERE model_id = {model_id}'
+        ),
+        'schema': 'predictions_schema',
+        'table_name': 'predictions',
+        'transform_table_name': 'transformed_data',
+        'input_filters': {'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}}},
+        'mlflow_transform_filters': {'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}},
+        'mlflow_predict_filters': {'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}},
+        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
+        'opc_output_config': {},
+        'pi_web_api_output_config': {},
+        'save_transform': True,
+        'prediction_store_policy': 'lts:1',
+        'model_config': {
+            'retention_minutes': 0,
+            'transform_flavor': 'sklearn',
+            'predict_flavor': 'sklearn',
+        },
+        'datetime_columns': ['timestamp', 'created_at'],
+    }
+
+    with patch.object(mdp, 'OFFLOAD_THRESHOLD_BYTES', 1):
+        await start_and_await_workflow(
+            temporal_test_env.client,
+            PredictionsBatch.run,
+            input_data,
+            make_workflow_id('test-batch-minio-offload'),
+        )
+
+    with postgres_engine.connect() as conn:
+        count = conn.execute(
+            text(f'SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = {model_id}')
+        ).scalar()
+    assert count == 1

e2e/test_opc_real_server.py

@@ -0,0 +1,194 @@
+"""
+E2E tests for OPC export using an in-process asyncua server and real OpcRepository.
+
+Covers scenarios 3.1.2, 3.2.2, 3.2.4, and 3.2.5 from e2e/scenarios.md.
+Mock-based OPC tests remain in test_predictions_batch_format_export.py.
+"""
+
+import asyncio
+
+import pytest
+from temporalio.testing import WorkflowEnvironment
+from temporalio.worker import Worker
+
+from e2e.helpers import assert_prediction, insert_sample_data, make_workflow_id, start_and_await_workflow
+from e2e.opc_test_server import UNKNOWN_NODE_ID, OpcE2ETestServer, build_opc_output_config
+from e2e.test_predictions_batch_format_export import get_base_input_data
+from laborious.activities.activities import Activities
+from laborious.activities.opc import OPC_RECONNECT_IN_PROGRESS_COMMENT
+from laborious.utils.repository.opc_repository import OpcRepository
+from laborious.workflows.predictions_batch import PredictionsBatch
+
+
+async def _slow_reconnect_under_lock(repo: OpcRepository, hold_seconds: float = 0.75) -> None:
+    """
+    Hold the connection lock briefly so concurrent writes see reconnect_in_progress.
+
+    Args:
+        repo (OpcRepository): Connected repository.
+        hold_seconds (float): Time to keep the lock before reconnecting.
+    """
+    async with repo._connection_lock:
+        await asyncio.sleep(hold_seconds)
+        await repo._reconnect_locked()
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+@pytest.mark.opc
+async def test_scenario_3_1_2_export_with_opc_only_real_server(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker_real_opc: Worker,
+    test_activities_real_opc: Activities,
+    opc_e2e_server: OpcE2ETestServer,
+    postgres_engine,
+):
+    """
+    Scenario 3.1.2 (real OPC): connect, write prediction and confidence, verify server values.
+    """
+    client = temporal_test_env.client
+    model_id = 412
+    insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
+
+    input_data = get_base_input_data(model_id)
+    input_data['opc_output_config'] = build_opc_output_config(opc_e2e_server.node_ids)
+    input_data['pi_web_api_output_config'] = None
+
+    await start_and_await_workflow(
+        client,
+        PredictionsBatch.run,
+        input_data,
+        make_workflow_id('test-opc-real-happy'),
+    )
+
+    test_activities_real_opc.pi_web_api_client.write_value.assert_not_called()
+    assert await opc_e2e_server.read_prediction() == pytest.approx(0.5)
+    assert await opc_e2e_server.read_confidence() == pytest.approx(0.0)
+    assert_prediction(postgres_engine, model_id)
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+@pytest.mark.opc
+async def test_scenario_3_2_2_opc_write_error_real_server(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker_real_opc: Worker,
+    test_activities_real_opc: Activities,
+    opc_e2e_server: OpcE2ETestServer,
+    postgres_engine,
+):
+    """
+    Scenario 3.2.2 (real OPC): unknown NodeId yields generic write failure (confidence 12).
+    """
+    client = temporal_test_env.client
+    model_id = 422
+    insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
+
+    node_ids = opc_e2e_server.node_ids
+    input_data = get_base_input_data(model_id)
+    input_data['opc_output_config'] = build_opc_output_config(
+        node_ids,
+        prediction_tag=UNKNOWN_NODE_ID,
+        confidence_tag=UNKNOWN_NODE_ID,
+    )
+    input_data['pi_web_api_output_config'] = None
+
+    await start_and_await_workflow(
+        client,
+        PredictionsBatch.run,
+        input_data,
+        make_workflow_id('test-opc-real-bad-node'),
+    )
+
+    assert_prediction(
+        postgres_engine,
+        model_id,
+        prediction_confidence=12,
+        comments='Some data could not be written to OPC servers',
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+@pytest.mark.opc
+async def test_scenario_3_2_4_opc_session_bad_real_server(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker_real_opc: Worker,
+    test_activities_real_opc: Activities,
+    opc_e2e_server: OpcE2ETestServer,
+    postgres_engine,
+):
+    """
+    Scenario 3.2.4 (real OPC): server PreWrite fault injects BadSessionIdInvalid (confidence 14).
+    """
+    client = temporal_test_env.client
+    model_id = 424
+    insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
+
+    opc_e2e_server.set_session_bad_on_write(True)
+    try:
+        input_data = get_base_input_data(model_id)
+        input_data['opc_output_config'] = build_opc_output_config(
+            opc_e2e_server.node_ids,
+            prediction_only=True,
+        )
+        input_data['pi_web_api_output_config'] = None
+
+        await start_and_await_workflow(
+            client,
+            PredictionsBatch.run,
+            input_data,
+            make_workflow_id('test-opc-real-session-bad'),
+        )
+    finally:
+        opc_e2e_server.set_session_bad_on_write(False)
+
+    assert_prediction(
+        postgres_engine,
+        model_id,
+        prediction_confidence=14,
+        comments_contains='OPC UA session/channel error: BadSessionIdInvalid',
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+@pytest.mark.opc
+async def test_scenario_3_2_5_opc_write_blocked_during_reconnect_real_server(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker_real_opc: Worker,
+    test_activities_real_opc: Activities,
+    opc_e2e_server: OpcE2ETestServer,
+    postgres_engine,
+):
+    """
+    Scenario 3.2.5 (real OPC): writes rejected while reconnect holds the connection lock.
+    """
+    client = temporal_test_env.client
+    model_id = 425
+    insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
+
+    repo = test_activities_real_opc.opc_repository['1']
+    repo._session_ready.clear()
+    reconnect_task = asyncio.create_task(_slow_reconnect_under_lock(repo))
+
+    input_data = get_base_input_data(model_id)
+    input_data['opc_output_config'] = build_opc_output_config(opc_e2e_server.node_ids)
+    input_data['pi_web_api_output_config'] = None
+
+    try:
+        await start_and_await_workflow(
+            client,
+            PredictionsBatch.run,
+            input_data,
+            make_workflow_id('test-opc-real-reconnect-block'),
+        )
+    finally:
+        await reconnect_task
+
+    assert_prediction(
+        postgres_engine,
+        model_id,
+        prediction_confidence=14,
+        comments_contains=OPC_RECONNECT_IN_PROGRESS_COMMENT,
+    )

e2e/test_predictions_batch_format_export.py

@@ -0,0 +1,786 @@
+"""
+End-to-end tests for PredictionsBatch workflow - Format and Export scenarios.
+"""
+
+from decimal import Decimal
+from typing import Any, cast
+from unittest.mock import ANY, AsyncMock, call
+
+import pytest
+from sientia_do.notifications.models import NotificationLevel
+from sqlalchemy import text
+from temporalio.testing import WorkflowEnvironment
+from temporalio.worker import Worker
+
+from e2e.helpers import assert_prediction, insert_sample_data, make_workflow_id, start_and_await_workflow
+from laborious.activities.activities import Activities
+from laborious.workflows.predictions_batch import PredictionsBatch
+
+base_input_data = {
+    'schedule_name': 'test-schedule',
+    'model_name': 'test_model',
+    'model_id': 301,
+    'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 301',
+    'schema': 'predictions_schema',
+    'table_name': 'predictions',
+    'transform_table_name': 'transformed_data',
+    'input_filters': {
+        'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}},
+    },
+    'mlflow_transform_filters': {
+        'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
+    },
+    'mlflow_predict_filters': {
+        'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
+    },
+    'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
+    'opc_output_config': {},
+    'pi_web_api_output_config': {},
+    'save_transform': True,
+    'prediction_store_policy': 'lts:1',
+    'model_config': {
+        'retention_minutes': 0,
+        'transform_flavor': 'sklearn',
+        'predict_flavor': 'sklearn',
+    },
+    'datetime_columns': ['timestamp', 'created_at'],
+}
+
+base_query = "SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = {model_id}"
+
+def get_base_input_data(model_id):
+    return {
+        **base_input_data,
+        'model_id': model_id,
+        'query': base_query.format(model_id=model_id),
+    }
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_scenario_3_1_1_default_prediction_export(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker: Worker,
+    test_activities: Activities,
+    postgres_engine,
+):
+    """
+    Scenario 3.1.1: Default prediction export (non-None path_flag).
+
+    Triggers input_gate CONTINUE via SPECIFIC_VARIABLES_NULL_VALUES so
+    PredictionProcess calls FormatAndExportPrediction with path_flag set.
+    That workflow uses format_default_prediction (not format_prediction) and
+    skips format_transformed_data / transform Postgres export.
+
+    Optional PI Web API and OPC outputs still run when configured.
+    """
+    client = temporal_test_env.client
+
+    model_id = 311
+
+    with postgres_engine.begin() as conn:
+        conn.execute(text(f"DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}"))
+        conn.execute(text(f"DELETE FROM predictions_schema.transformed_data WHERE model_id = {model_id}"))
+    insert_sample_data(postgres_engine, model_id, ['NULL', 78.2])
+
+    input_data = get_base_input_data(model_id)
+    input_data['input_filters'] = {
+        'SPECIFIC_VARIABLES_NULL_VALUES': {
+            'POLICY': 'CONTINUE',
+            'CONFIG': {'variables': ['sensor_1']},
+        },
+    }
+    input_data['pi_web_api_output_config'] = {
+        'endpoint': 'test_endpoint',
+        'prediction_tags': {'tag_1': 'web_id_1'},
+        'confidence_tags': {'tag_2': 'web_id_2'},
+    }
+    input_data['opc_output_config'] = {
+        '1': {
+            'prediction_tags': {
+                'addr_1': {
+                    'data_type': 'float',
+                }
+            },
+            'confidence_tags': {
+                'addr_2': {
+                    'data_type': 'float',
+                }
+            },
+        }
+    }
+
+    wid = make_workflow_id('test-default-prediction')
+
+    await start_and_await_workflow(client, PredictionsBatch.run, input_data, wid)
+
+    test_activities.pi_web_api_client.write_value.assert_has_calls(
+        [
+            call(
+                web_ids=['web_id_1'],
+                value={
+                    'Timestamp': '2024-01-01 12:00:00+0000',
+                    'Value': 0,
+                },
+                metadata={
+                    'model_id': 311,
+                    'model_name': 'test_model',
+                    'schedule_name': 'test-schedule',
+                    'workflow_name': 'predictions_batch',
+                },
+            ),
+            call(
+                web_ids=['web_id_2'],
+                value={
+                    'Timestamp': '2024-01-01 12:00:00+0000',
+                    'Value': 2,
+                },
+                metadata={
+                    'model_id': 311,
+                    'model_name': 'test_model',
+                    'schedule_name': 'test-schedule',
+                    'workflow_name': 'predictions_batch',
+                },
+            ),
+        ],
+        any_order=True,
+    )
+
+    opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data)
+    opc_write_data.assert_has_calls(
+        [
+            call(
+                'addr_1',
+                0,
+                'float',
+                {
+                    'model_id': 311,
+                    'model_name': 'test_model',
+                    'schedule_name': 'test-schedule',
+                    'workflow_name': 'predictions_batch',
+                },
+            ),
+            call(
+                'addr_2',
+                2,
+                'float',
+                {
+                    'model_id': 311,
+                    'model_name': 'test_model',
+                    'schedule_name': 'test-schedule',
+                    'workflow_name': 'predictions_batch',
+                },
+            ),
+        ]
+    )
+
+    with postgres_engine.connect() as conn:
+        tf_count = conn.execute(
+            text(f"SELECT COUNT(*) FROM predictions_schema.transformed_data WHERE model_id = {model_id}")
+        ).scalar()
+        assert tf_count == 0, 'transform export must be skipped when path_flag is set'
+
+    assert_prediction(
+        postgres_engine,
+        model_id,
+        prediction=0,
+        prediction_confidence=Decimal(2),
+        prediction_status='Bad',
+        comments='Input data with bad quality',
+    )
+
+
+
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_scenario_3_1_2_export_with_opc_only(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker: Worker,
+    test_activities: Activities,
+    postgres_engine,
+):
+    """
+    Scenario 3.1.2: Export with OPC only
+    
+    Description:
+    Export to PostgreSQL and OPC server only (no PI Web API).
+    
+    Expected Behavior:
+    - Normal formatting
+    - PostgreSQL export executed
+    - OPC export executed
+    - PI Web API activity skipped
+    - Metrics written with OPC metrics
+    
+    Assertions:
+    - PI Web API activity NOT called
+    - OPC activity called
+    - PostgreSQL export called
+    - Metrics written with opc_metrics populated
+    """
+    client = temporal_test_env.client
+
+    model_id = 312
+
+    insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
+
+    input_data = get_base_input_data(model_id)
+    input_data['opc_output_config'] = {
+        '1': {
+            'prediction_tags': {
+                'addr_1': {
+                    'data_type': 'float',
+                }
+            },
+            'confidence_tags': {
+                'addr_2': {
+                    'data_type': 'float',
+                }
+            },
+        }
+    }
+    input_data['pi_web_api_output_config'] = None  # No PI Web API config
+
+    await start_and_await_workflow(
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-opc-only')
+    )
+
+    opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data)
+    opc_write_data.assert_has_calls(
+        [
+            call(
+                'addr_1',
+                0.5,
+                'float',
+                {
+                    'model_id': 312,
+                    'model_name': 'test_model',
+                    'schedule_name': 'test-schedule',
+                    'workflow_name': 'predictions_batch',
+                },
+            ),
+            call(
+                'addr_2',
+                0,
+                'float',
+                {
+                    'model_id': 312,
+                    'model_name': 'test_model',
+                    'schedule_name': 'test-schedule',
+                    'workflow_name': 'predictions_batch',
+                },
+            ),
+        ]
+    )
+
+    test_activities.pi_web_api_client.write_value.assert_not_called()
+
+    assert_prediction(postgres_engine, model_id)
+
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_scenario_3_1_3_export_with_pi_web_api_only(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker: Worker,
+    test_activities: Activities,
+    postgres_engine,
+):
+    """
+    Scenario 3.1.3: Export with PI Web API only
+    
+    Description:
+    Export to PostgreSQL and PI Web API only (no OPC).
+    
+    Expected Behavior:
+    - Normal formatting
+    - PostgreSQL export executed
+    - PI Web API export executed
+    - OPC activity skipped
+    - Metrics written without OPC metrics
+    
+    Assertions:
+    - OPC activity NOT called
+    - PI Web API activity called
+    - PostgreSQL export called
+    - Metrics written with empty opc_metrics
+    """
+    client = temporal_test_env.client
+
+    model_id = 313
+
+    insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
+
+    input_data = get_base_input_data(model_id)
+    input_data['pi_web_api_output_config'] = {
+        'endpoint': 'test_endpoint',
+        'prediction_tags': {'tag_1': 'web_id_1'},
+        'confidence_tags': {'tag_2': 'web_id_2'},
+    }
+    input_data['opc_output_config'] = None  # No OPC config
+
+    await start_and_await_workflow(
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-pi-api-only')
+    )
+
+    test_activities.pi_web_api_client.write_value.assert_has_calls(
+        [
+            call(
+                web_ids=['web_id_1'],
+                value={
+                    'Timestamp': '2024-01-01 12:00:00+0000',
+                    'Value': 0.5,
+                },
+                metadata={
+                    'model_id': 313,
+                    'model_name': 'test_model',
+                    'schedule_name': 'test-schedule',
+                    'workflow_name': 'predictions_batch',
+                },
+            ),
+            call(
+                web_ids=['web_id_2'],
+                value={
+                    'Timestamp': '2024-01-01 12:00:00+0000',
+                    'Value': 0,
+                },
+                metadata={
+                    'model_id': 313,
+                    'model_name': 'test_model',
+                    'schedule_name': 'test-schedule',
+                    'workflow_name': 'predictions_batch',
+                },
+            ),
+        ],
+        any_order=True,
+    )
+
+    opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data)
+    opc_write_data.assert_not_called()
+
+    assert_prediction(postgres_engine, model_id)
+
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_scenario_3_1_4_export_without_optional_outputs(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker: Worker,
+    test_activities: Activities,
+    postgres_engine,
+):
+    """
+    Scenario 3.1.4: Export Without Optional Outputs
+    
+    Description:
+    Export only to PostgreSQL (no OPC or PI Web API).
+    
+    Expected Behavior:
+    - Normal formatting
+    - Only PostgreSQL export executed
+    - OPC and PI Web API activities skipped
+    - Metrics written without OPC metrics
+    
+    Assertions:
+    - PI Web API activity NOT called
+    - OPC activity NOT called
+    - PostgreSQL export called
+    - Metrics written with empty opc_metrics
+    """
+    client = temporal_test_env.client
+
+    model_id = 314
+
+    insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
+
+    input_data = get_base_input_data(model_id)
+    input_data['opc_output_config'] = None  # No OPC config
+    input_data['pi_web_api_output_config'] = None  # No PI Web API config
+
+    await start_and_await_workflow(
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-no-optional-outputs')
+    )
+
+    test_activities.pi_web_api_client.write_value.assert_not_called()
+    opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data)
+    opc_write_data.assert_not_called()
+
+    assert_prediction(postgres_engine, model_id)
+
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_scenario_3_1_5_export_without_transformed_data(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker: Worker,
+    test_activities: Activities,
+    postgres_engine,
+):
+    """
+    Scenario 3.1.5: Export Without Transformed Data
+    
+    Description:
+    Only prediction exported, no transform table.
+    
+    Expected Behavior:
+    - Only prediction formatted and exported
+    - Transform export skipped
+    - Single PostgreSQL write
+    
+    Assertions:
+    - format_transformed_data NOT called
+    - One PostgreSQL export
+    - Transform table remains empty
+    """
+    client = temporal_test_env.client
+
+    model_id = 315
+
+    insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
+    with postgres_engine.begin() as conn:
+        conn.execute(text(f"DELETE FROM predictions_schema.transformed_data WHERE model_id = {model_id}"))
+
+    input_data = get_base_input_data(model_id)
+    input_data['save_transform'] = False  # Don't save transformed data
+    input_data['pi_web_api_output_config'] = {
+        'endpoint': 'test_endpoint',
+        'prediction_tags': {'tag_1': 'web_id_1'},
+        'confidence_tags': {'tag_2': 'web_id_2'},
+    }
+    input_data['opc_output_config'] = {
+        '1': {
+            'prediction_tags': {
+                'addr_1': {
+                    'data_type': 'float',
+                }
+            },
+            'confidence_tags': {
+                'addr_2': {
+                    'data_type': 'float',
+                }
+            },
+        }
+    }
+
+    await start_and_await_workflow(
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-no-transform-export')
+    )
+
+    test_activities.pi_web_api_client.write_value.assert_has_calls(
+        [
+            call(
+                web_ids=['web_id_1'],
+                value={
+                    'Timestamp': '2024-01-01 12:00:00+0000',
+                    'Value': 0.5,
+                },
+                metadata={
+                    'model_id': 315,
+                    'model_name': 'test_model',
+                    'schedule_name': 'test-schedule',
+                    'workflow_name': 'predictions_batch',
+                },
+            ),
+            call(
+                web_ids=['web_id_2'],
+                value={
+                    'Timestamp': '2024-01-01 12:00:00+0000',
+                    'Value': 0,
+                },
+                metadata={
+                    'model_id': 315,
+                    'model_name': 'test_model',
+                    'schedule_name': 'test-schedule',
+                    'workflow_name': 'predictions_batch',
+                },
+            ),
+        ],
+        any_order=True,
+    )
+
+    opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data)
+    opc_write_data.assert_has_calls(
+        [
+            call(
+                'addr_1',
+                0.5,
+                'float',
+                {
+                    'model_id': 315,
+                    'model_name': 'test_model',
+                    'schedule_name': 'test-schedule',
+                    'workflow_name': 'predictions_batch',
+                },
+            ),
+            call(
+                'addr_2',
+                0,
+                'float',
+                {
+                    'model_id': 315,
+                    'model_name': 'test_model',
+                    'schedule_name': 'test-schedule',
+                    'workflow_name': 'predictions_batch',
+                },
+            ),
+        ]
+    )
+
+    with postgres_engine.connect() as conn:
+        result_query = conn.execute(
+            text(f"SELECT COUNT(*) FROM predictions_schema.transformed_data WHERE model_id = {model_id}")
+        )
+        count = result_query.scalar()
+        assert count == 0, f"Expected transform table to be empty, but found {count} records"
+
+    assert_prediction(postgres_engine, model_id)
+
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_scenario_3_2_1_pi_web_api_write_error(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker: Worker,
+    test_activities: Activities,
+    postgres_engine,
+    notification_inserts,
+):
+    """
+    Scenario 3.2.1: PI Web API Write Error
+
+    Export failure is handled inside the activity; there is no retry loop. The
+    workflow completes and PostgreSQL stores prediction_confidence 13 and the
+    error message in comments.
+    """
+    client = temporal_test_env.client
+
+    model_id = 321
+
+    insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
+
+    test_activities.pi_web_api_client.write_value.side_effect = Exception(
+        "PI Web API service unavailable")
+
+    input_data = get_base_input_data(model_id)
+    input_data['pi_web_api_output_config'] = {
+        'endpoint': 'test_endpoint',
+        'prediction_tags': {'tag_1': 'web_id_1'},
+        'confidence_tags': {'tag_2': 'web_id_2'},
+    }
+    input_data['opc_output_config'] = {
+        '1': {
+            'prediction_tags': {
+                'addr_1': {
+                    'data_type': 'float',
+                }
+            },
+            'confidence_tags': {
+                'addr_2': {
+                    'data_type': 'float',
+                }
+            },
+        }
+    }
+
+    await start_and_await_workflow(
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-pi-api-error')
+    )
+
+    assert_prediction(
+        postgres_engine, model_id,
+        prediction_confidence=13,
+        comments='PI Web API service unavailable',
+    )
+    assert notification_inserts.call_count >= 1
+
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_scenario_3_2_2_opc_write_error(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker: Worker,
+    test_activities: Activities,
+    postgres_engine,
+):
+    """
+    Scenario 3.2.2: OPC Write Error
+
+    OPC failure is reported without failing the workflow; there is no retry
+    loop. PostgreSQL stores prediction_confidence 12 and OPC error comments.
+    """
+    client = temporal_test_env.client
+
+    model_id = 322
+
+    insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
+
+    opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data)
+    opc_write_data.return_value = (False, {
+        'notification_id': 'OPC_WRITE_DATA_ERROR_1',
+        'message': 'OPC server unavailable',
+        'block': 'opc_repository',
+        'level': NotificationLevel.ERROR,
+        'attachment_content': 'OPC server unavailable',
+    })
+    
+    input_data = get_base_input_data(model_id)
+    input_data['opc_output_config'] = {
+        '1': {
+            'prediction_tags': {
+                'addr_1': {
+                    'data_type': 'float',
+                }
+            },
+            'confidence_tags': {
+                'addr_2': {
+                    'data_type': 'float',
+                }
+            },
+        }
+    }
+    input_data['pi_web_api_output_config'] = {
+        'endpoint': 'test_endpoint',
+        'prediction_tags': {'tag_1': 'web_id_1'},
+        'confidence_tags': {'tag_2': 'web_id_2'},
+    }
+
+    await start_and_await_workflow(
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-opc-error')
+    )
+
+    assert_prediction(
+        postgres_engine, model_id,
+        prediction_confidence=12,
+        comments='Some data could not be written to OPC servers',
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_scenario_3_2_4_opc_session_bad_error(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker: Worker,
+    test_activities: Activities,
+    postgres_engine,
+):
+    """
+    Scenario 3.2.4: OPC session/channel Tier-1 Bad* (e.g. BadSessionIdInvalid).
+
+    PostgreSQL stores prediction_confidence 14 and a stable session error comment.
+    """
+    client = temporal_test_env.client
+
+    model_id = 324
+
+    insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
+
+    opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data)
+    opc_write_data.return_value = (
+        False,
+        {
+            'notification_id': 'OPC_WRITE_DATA_ERROR_1',
+            'message': 'BadSessionIdInvalid',
+            'block': 'opc_repository',
+            'level': NotificationLevel.ERROR,
+            'attachment_content': 'BadSessionIdInvalid',
+            'opc_error_kind': 'session_bad',
+            'opc_status': 'BadSessionIdInvalid',
+        },
+    )
+
+    input_data = get_base_input_data(model_id)
+    input_data['opc_output_config'] = {
+        '1': {
+            'prediction_tags': {
+                'addr_1': {
+                    'data_type': 'float',
+                }
+            },
+            'confidence_tags': {
+                'addr_2': {
+                    'data_type': 'float',
+                }
+            },
+        }
+    }
+    input_data['pi_web_api_output_config'] = None
+
+    await start_and_await_workflow(
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-opc-session-bad')
+    )
+
+    assert_prediction(
+        postgres_engine,
+        model_id,
+        prediction_confidence=14,
+        comments='OPC UA session/channel error: BadSessionIdInvalid',
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_scenario_3_2_3_pi_web_api_partial_write_error(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker: Worker,
+    test_activities: Activities,
+    postgres_engine,
+):
+    """
+    Scenario 3.2.3: PI Web API Partial Write Error
+
+    Partial PI write: confidence 13, descriptive comments, workflow completes
+    without an activity retry loop.
+    """
+    client = temporal_test_env.client
+
+    model_id = 323
+
+    insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
+
+    test_activities.pi_web_api_client.write_value = AsyncMock(
+        side_effect=[
+            # Prediction batch: two web_ids requested, only one acknowledged.
+            [{'WebId': 'web_id_1', 'Errors': []}],
+            # Confidence write succeeds.
+            [{'WebId': 'web_id_2', 'Errors': []}],
+        ]
+    )
+
+    input_data = get_base_input_data(model_id)
+    input_data['pi_web_api_output_config'] = {
+        'endpoint': 'test_endpoint',
+        'prediction_tags': {'tag_1': 'web_id_1', 'tag_3': 'web_id_3'},
+        'confidence_tags': {'tag_2': 'web_id_2'},
+    }
+    input_data['opc_output_config'] = {
+        '1': {
+            'prediction_tags': {
+                'addr_1': {
+                    'data_type': 'float',
+                }
+            },
+            'confidence_tags': {
+                'addr_2': {
+                    'data_type': 'float',
+                }
+            },
+        }
+    }
+
+    await start_and_await_workflow(
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-pi-api-partial-error')
+    )
+
+    assert_prediction(
+        postgres_engine, model_id,
+        prediction_confidence=13,
+        comments="The number of written tags does not match the number of tag names: Expected ['tag_1', 'tag_3'] tags, but ['tag_1'] tags were written.",
+    )
+    

e2e/test_predictions_batch_main_workflow.py

@@ -0,0 +1,303 @@
+"""
+End-to-end tests for PredictionsBatch workflow - Main workflow scenarios.
+"""
+
+import asyncio
+
+from sqlalchemy import text
+from temporalio.testing import WorkflowEnvironment
+from temporalio.worker import Worker
+
+import pytest
+
+from e2e.helpers import make_workflow_id, start_and_await_workflow
+from laborious.activities.activities import Activities
+from laborious.workflows.predictions_batch import PredictionsBatch
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_scenario_1_1_1_happy_path_complete_success(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker: Worker,
+    test_activities: Activities,
+    postgres_engine,
+):
+    """Scenario 1.1.1: Happy path with SQL load, MLflow mocks, Postgres predictions and transforms."""
+    client = temporal_test_env.client
+
+    with postgres_engine.begin() as conn:
+        conn.execute(text('DELETE FROM predictions_schema.laborious_data WHERE model_id = 123'))
+        insert_sql = """
+            INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
+            VALUES
+                (123, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'),
+                (123, 'sensor_2', 78.2, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'),
+                (123, 'sensor_3', 120.8, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
+        """
+        conn.execute(text(insert_sql))
+
+    input_data = {
+        'metadata': {
+            'metadata': {
+                'model_id': 123,
+                'model_name': 'test_model',
+                'schedule_name': 'test-schedule',
+                'workflow_name': 'predictions_batch',
+            }
+        },
+        'schedule_name': 'test-schedule',
+        'model_name': 'test_model',
+        'model_id': 123,
+        'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 123',
+        'schema': 'predictions_schema',
+        'table_name': 'predictions',
+        'transform_table_name': 'transformed_data',
+        'input_filters': {
+            'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}},
+        },
+        'mlflow_transform_filters': {
+            'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
+        },
+        'mlflow_predict_filters': {
+            'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
+        },
+        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
+        'opc_output_config': {},
+        'pi_web_api_output_config': {},
+        'save_transform': True,
+        'prediction_store_policy': 'lts:1',
+        'model_config': {
+            'retention_minutes': 0,
+            'transform_flavor': 'sklearn',
+            'predict_flavor': 'sklearn',
+        },
+        'datetime_columns': ['timestamp', 'created_at'],
+    }
+
+    await start_and_await_workflow(
+        client,
+        PredictionsBatch.run,
+        input_data,
+        make_workflow_id('test-predictions-batch'),
+    )
+
+    schema_name = 'predictions_schema'
+    with postgres_engine.connect() as conn:
+        result_query = conn.execute(
+            text(
+                f'SELECT model_id, prediction, prediction_confidence, response_time, prediction_status, comments '
+                f'FROM {schema_name}.predictions WHERE model_id = 123'
+            )
+        )
+        prediction_rows = result_query.fetchall()
+        assert len(prediction_rows) == 1
+        row = prediction_rows[0]
+        assert row[0] == 123
+        assert row[1] == 0.5
+        assert row[2] == 0, f'Expected prediction_confidence=0, got {row[2]}'
+        assert row[3] is not None
+        assert row[4] == 'Good'
+        assert row[5] == ''
+
+        result_query = conn.execute(
+            text(
+                f'SELECT model_id, variable, value FROM {schema_name}.transformed_data WHERE model_id = 123'
+            )
+        )
+        transformed_rows = result_query.fetchall()
+        assert len(transformed_rows) == 2
+        assert transformed_rows[0][0] == 123
+        assert transformed_rows[0][1] == 'feature_1'
+        assert float(transformed_rows[0][2]) == 0.234
+        assert transformed_rows[1][0] == 123
+        assert transformed_rows[1][1] == 'feature_2'
+        assert float(transformed_rows[1][2]) == 0.783
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_scenario_1_2_1_sql_query_execution_error(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker: Worker,
+    test_activities: Activities,
+    postgres_engine,
+):
+    """Invalid SQL: workflow may complete with early exit; no prediction rows."""
+    client = temporal_test_env.client
+
+    input_data = {
+        'metadata': {
+            'metadata': {
+                'model_id': 128,
+                'model_name': 'test_model',
+                'schedule_name': 'test-schedule',
+                'workflow_name': 'predictions_batch',
+            }
+        },
+        'schedule_name': 'test-schedule',
+        'model_name': 'test_model',
+        'model_id': 128,
+        'query': 'SELECT * FROM nonexistent_table WHERE invalid_syntax =',
+        'schema': 'predictions_schema',
+        'table_name': 'predictions',
+        'transform_table_name': 'transformed_data',
+        'input_filters': {
+            'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}},
+        },
+        'mlflow_transform_filters': {
+            'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
+        },
+        'mlflow_predict_filters': {
+            'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
+        },
+        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
+        'opc_output_config': {},
+        'pi_web_api_output_config': {},
+        'save_transform': True,
+        'prediction_store_policy': 'lts:1',
+        'model_config': {
+            'retention_minutes': 0,
+            'transform_flavor': 'sklearn',
+            'predict_flavor': 'sklearn',
+        },
+    }
+
+    await start_and_await_workflow(
+        client,
+        PredictionsBatch.run,
+        input_data,
+        make_workflow_id('test-sql-error'),
+    )
+
+    with postgres_engine.connect() as conn:
+        count = conn.execute(
+            text('SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 128')
+        ).scalar()
+    assert count == 0
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_scenario_1_2_2_missing_required_parameters(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker: Worker,
+    test_activities: Activities,
+    postgres_engine,
+):
+    """Missing query: workflow does not produce predictions and is terminated explicitly."""
+    client = temporal_test_env.client
+
+    input_data = {
+        'metadata': {
+            'metadata': {
+                'model_id': 129,
+                'model_name': 'test_model',
+                'schedule_name': 'test-schedule',
+                'workflow_name': 'predictions_batch',
+            }
+        },
+        'schedule_name': 'test-schedule',
+        'model_name': 'test_model',
+        'model_id': 129,
+        'schema': 'predictions_schema',
+        'table_name': 'predictions',
+        'transform_table_name': 'transformed_data',
+    }
+
+    handle = await client.start_workflow(
+        PredictionsBatch.run,
+        input_data,
+        id=make_workflow_id('test-missing-param'),
+        task_queue='test-queue',
+    )
+
+    # Let Temporal process a few workflow tasks; for this case, result() can hang.
+    await asyncio.sleep(2.0)
+
+    with postgres_engine.connect() as conn:
+        count = conn.execute(
+            text('SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 129')
+        ).scalar()
+    assert count == 0
+
+    await handle.terminate('expected failure path in e2e test (missing required parameters)')
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_scenario_1_2_3_invalid_datetime_column_specification(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker: Worker,
+    test_activities: Activities,
+    postgres_engine,
+):
+    """Invalid datetime column: no predictions persisted; workflow terminated after validation."""
+    client = temporal_test_env.client
+
+    with postgres_engine.begin() as conn:
+        conn.execute(text('DELETE FROM predictions_schema.laborious_data WHERE model_id = 130'))
+        conn.execute(
+            text(
+                """
+            INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
+            VALUES (130, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
+        """
+            )
+        )
+
+    input_data = {
+        'metadata': {
+            'metadata': {
+                'model_id': 130,
+                'model_name': 'test_model',
+                'schedule_name': 'test-schedule',
+                'workflow_name': 'predictions_batch',
+            }
+        },
+        'schedule_name': 'test-schedule',
+        'model_name': 'test_model',
+        'model_id': 130,
+        'query': 'SELECT timestamp, variable, value FROM predictions_schema.laborious_data WHERE model_id = 130',
+        'schema': 'predictions_schema',
+        'table_name': 'predictions',
+        'transform_table_name': 'transformed_data',
+        'input_filters': {
+            'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}},
+        },
+        'mlflow_transform_filters': {
+            'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
+        },
+        'mlflow_predict_filters': {
+            'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
+        },
+        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
+        'opc_output_config': {},
+        'pi_web_api_output_config': {},
+        'save_transform': True,
+        'prediction_store_policy': 'lts:1',
+        'model_config': {
+            'retention_minutes': 0,
+            'transform_flavor': 'sklearn',
+            'predict_flavor': 'sklearn',
+        },
+        'datetime_columns': ['nonexistent_column'],
+    }
+
+    handle = await client.start_workflow(
+        PredictionsBatch.run,
+        input_data,
+        id=make_workflow_id('test-invalid-datetime-col'),
+        task_queue='test-queue',
+    )
+
+    # Let Temporal process and surface the failure path internally.
+    await asyncio.sleep(2.0)
+
+    with postgres_engine.connect() as conn:
+        count = conn.execute(
+            text('SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 130')
+        ).scalar()
+    assert count == 0
+
+    await handle.terminate('expected failure path in e2e test (invalid datetime column)')

e2e/test_predictions_batch_prediction_process.py

@@ -0,0 +1,379 @@
+"""
+End-to-end tests for PredictionsBatch workflow - Prediction Process scenarios.
+"""
+
+from decimal import Decimal
+from unittest.mock import MagicMock, patch
+
+import numpy as np
+import pandas as pd
+import pytest
+from sqlalchemy import text
+from temporalio.testing import WorkflowEnvironment
+from temporalio.worker import Worker
+
+from e2e.helpers import (
+    assert_continue,
+    assert_repeat,
+    assert_stop,
+    insert_sample_data,
+    make_workflow_id,
+    start_and_await_workflow,
+)
+from laborious.activities.activities import Activities
+from laborious.workflows.predictions_batch import PredictionsBatch
+
+base_input_data = {
+    'schedule_name': 'test-schedule',
+    'model_name': 'test_model',
+    'model_id': 201,
+    'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 201',
+    'schema': 'predictions_schema',
+    'table_name': 'predictions',
+    'transform_table_name': 'transformed_data',
+    'input_filters': {
+        'SPECIFIC_VARIABLES_NULL_VALUES': {
+            'POLICY': 'CONTINUE',
+            'CONFIG': {'variables': ['sensor_1']},
+        },
+    },
+    'mlflow_transform_filters': {
+        'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
+    },
+    'mlflow_predict_filters': {
+        'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
+    },
+    'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
+    'opc_output_config': {},
+    'pi_web_api_output_config': {},
+    'save_transform': True,
+    'prediction_store_policy': 'lts:1',
+    'model_config': {
+        'retention_minutes': 0,
+        'transform_flavor': 'sklearn',
+        'predict_flavor': 'sklearn',
+    },
+    'datetime_columns': ['timestamp', 'created_at'],
+}
+
+base_query = 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = {model_id}'
+
+
+def get_base_input_data(model_id):
+    return {
+        **base_input_data,
+        'model_id': model_id,
+        'query': base_query.format(model_id=model_id),
+    }
+
+
+def insert_sample_prediction(postgres_engine, model_id):
+    with postgres_engine.begin() as conn:
+        conn.execute(text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}'))
+        insert_sql = f"""
+            INSERT INTO predictions_schema.predictions (model_id, timestamp, prediction, prediction_confidence, prediction_status, comments, response_time)
+            VALUES
+               ({model_id}, '2024-01-01 12:00:00+00:00', 10, 0, 'Good', '', 0.1)
+        """
+        conn.execute(text(insert_sql))
+    return (model_id, Decimal(10), Decimal(0), 'Good')
+
+
+@pytest.fixture
+def bad_data_model(patch_mlflow):
+    model = MagicMock(predict=MagicMock(side_effect=Exception('Bad data model')))
+    patch_mlflow.sklearn.load_model = MagicMock(return_value=model)
+    return model
+
+
+@pytest.fixture
+def bad_predict_model(patch_mlflow, mock_mlflow_models):
+    model = MagicMock(predict=MagicMock(side_effect=Exception('Bad predict model')))
+
+    def mock_sklearn_load_model(model_uri):
+        if 'data_model' in model_uri or 'transform' in model_uri.lower():
+            return mock_mlflow_models['transform_model']
+        return model
+
+    patch_mlflow.sklearn = MagicMock()
+    patch_mlflow.sklearn.load_model = MagicMock(side_effect=mock_sklearn_load_model)
+    return model
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_scenario_2_1_1_input_gate_triggers_continue(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker: Worker,
+    test_activities: Activities,
+    postgres_engine,
+    mock_mlflow_models,
+):
+    """Input gate CONTINUE: export default prediction; MLflow transform/predict not used."""
+    client = temporal_test_env.client
+    model_id = 211
+    insert_sample_data(postgres_engine, model_id, ['NULL', 78.2])
+    input_data = get_base_input_data(model_id)
+    await start_and_await_workflow(
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-continue-policy')
+    )
+    assert_continue(postgres_engine, model_id)
+    mock_mlflow_models['transform_model'].predict.assert_not_called()
+    mock_mlflow_models['predict_model'].predict.assert_not_called()
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_scenario_2_1_2_input_gate_triggers_stop(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker: Worker,
+    test_activities: Activities,
+    postgres_engine,
+    mock_mlflow_models,
+):
+    """Input gate STOP: no export, no MLflow."""
+    client = temporal_test_env.client
+    model_id = 212
+    insert_sample_data(postgres_engine, model_id, ['NULL', 78.2])
+    input_data = get_base_input_data(model_id)
+    input_data['input_filters']['SPECIFIC_VARIABLES_NULL_VALUES']['POLICY'] = 'STOP'
+    await start_and_await_workflow(
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-input-stop')
+    )
+    assert_stop(postgres_engine, model_id)
+    mock_mlflow_models['transform_model'].predict.assert_not_called()
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_scenario_2_1_3_input_gate_triggers_repeat(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker: Worker,
+    test_activities: Activities,
+    postgres_engine,
+    mock_mlflow_models,
+):
+    """Input gate REPEAT with existing history."""
+    client = temporal_test_env.client
+    model_id = 213
+    insert_sample_data(postgres_engine, model_id, ['NULL', 78.2])
+    data = insert_sample_prediction(postgres_engine, model_id)
+    input_data = get_base_input_data(model_id)
+    input_data['input_filters']['SPECIFIC_VARIABLES_NULL_VALUES']['POLICY'] = 'REPEAT'
+    await start_and_await_workflow(
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-input-repeat')
+    )
+    assert_repeat(postgres_engine, model_id, data)
+    mock_mlflow_models['transform_model'].predict.assert_not_called()
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_scenario_2_1_4_input_gate_repeat_without_prior_prediction(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker: Worker,
+    test_activities: Activities,
+    postgres_engine,
+):
+    """REPEAT when no prior row in predictions: repeat_last_prediction runs; still no new duplicate export path."""
+    client = temporal_test_env.client
+    model_id = 214
+    insert_sample_data(postgres_engine, model_id, ['NULL', 78.2])
+    with postgres_engine.begin() as conn:
+        conn.execute(text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}'))
+    input_data = get_base_input_data(model_id)
+    input_data['input_filters']['SPECIFIC_VARIABLES_NULL_VALUES']['POLICY'] = 'REPEAT'
+    await start_and_await_workflow(
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-input-repeat-no-history')
+    )
+    assert_stop(postgres_engine, model_id)
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_scenario_2_2_1_transform_gate_triggers_continue(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker: Worker,
+    test_activities: Activities,
+    postgres_engine,
+    bad_data_model,
+):
+    client = temporal_test_env.client
+    model_id = 221
+    insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
+    input_data = get_base_input_data(model_id)
+    input_data['mlflow_transform_filters']['API_ERROR']['POLICY'] = 'CONTINUE'
+    await start_and_await_workflow(
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-transform-continue')
+    )
+    assert_continue(
+        postgres_engine=postgres_engine,
+        model_id=model_id,
+        prediction_confidence=Decimal(10),
+        comments='Unknown MLFlow API error',
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_scenario_2_2_2_transform_gate_triggers_stop(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker: Worker,
+    test_activities: Activities,
+    postgres_engine,
+    bad_data_model,
+    mock_mlflow_models,
+):
+    client = temporal_test_env.client
+    model_id = 222
+    insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
+    input_data = get_base_input_data(model_id)
+    input_data['mlflow_transform_filters']['API_ERROR']['POLICY'] = 'STOP'
+    await start_and_await_workflow(
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-transform-stop')
+    )
+    assert_stop(postgres_engine, model_id)
+    mock_mlflow_models['predict_model'].predict.assert_not_called()
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_scenario_2_2_3_transform_gate_triggers_repeat(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker: Worker,
+    test_activities: Activities,
+    postgres_engine,
+    bad_data_model,
+):
+    client = temporal_test_env.client
+    model_id = 223
+    insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
+    data = insert_sample_prediction(postgres_engine, model_id)
+    input_data = get_base_input_data(model_id)
+    input_data['mlflow_transform_filters']['API_ERROR']['POLICY'] = 'REPEAT'
+    await start_and_await_workflow(
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-transform-repeat')
+    )
+    assert_repeat(postgres_engine, model_id, data)
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_scenario_2_2_4_transform_content_gate_nan_values_stop(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker: Worker,
+    test_activities: Activities,
+    postgres_engine,
+    mock_mlflow_models,
+):
+    """mlflow_content_gate triggers STOP when transform output is all NaN (NAN_VALUES filter)."""
+    client = temporal_test_env.client
+    model_id = 224
+
+    def all_nan_transform(data):
+        num_rows = max(len(data), 1) if hasattr(data, '__len__') else 1
+        result = pd.DataFrame({'feature_1': [np.nan] * num_rows, 'feature_2': [np.nan] * num_rows})
+        result.index = data.index
+        return result
+
+    mock_mlflow_models['transform_model'].predict = MagicMock(side_effect=all_nan_transform)
+
+    insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
+    input_data = get_base_input_data(model_id)
+    input_data['mlflow_transform_filters'] = {
+        'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
+        'NAN_VALUES': {'POLICY': 'STOP', 'CONFIG': {}},
+    }
+    await start_and_await_workflow(
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-transform-content-stop')
+    )
+    assert_stop(postgres_engine, model_id)
+    mock_mlflow_models['predict_model'].predict.assert_not_called()
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_scenario_2_3_1_predict_gate_triggers_continue(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker: Worker,
+    test_activities: Activities,
+    postgres_engine,
+    bad_predict_model,
+):
+    client = temporal_test_env.client
+    model_id = 231
+    insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
+    input_data = get_base_input_data(model_id)
+    input_data['mlflow_predict_filters']['API_ERROR']['POLICY'] = 'CONTINUE'
+    await start_and_await_workflow(
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-predict-continue')
+    )
+    assert_continue(
+        postgres_engine=postgres_engine,
+        model_id=model_id,
+        prediction_confidence=Decimal(10),
+        comments='Unknown MLFlow API error',
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_scenario_2_3_2_predict_gate_triggers_stop(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker: Worker,
+    test_activities: Activities,
+    postgres_engine,
+    bad_predict_model,
+):
+    client = temporal_test_env.client
+    model_id = 232
+    insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
+    input_data = get_base_input_data(model_id)
+    input_data['mlflow_predict_filters']['API_ERROR']['POLICY'] = 'STOP'
+    await start_and_await_workflow(
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-predict-stop')
+    )
+    assert_stop(postgres_engine, model_id)
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_scenario_2_3_3_predict_gate_triggers_repeat(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker: Worker,
+    test_activities: Activities,
+    postgres_engine,
+    bad_predict_model,
+):
+    client = temporal_test_env.client
+    model_id = 233
+    insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
+    data = insert_sample_prediction(postgres_engine, model_id)
+    input_data = get_base_input_data(model_id)
+    input_data['mlflow_predict_filters']['API_ERROR']['POLICY'] = 'REPEAT'
+    input_data['path_priority'] = ['REPEAT', 'STOP', 'CONTINUE']
+    await start_and_await_workflow(
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-predict-repeat')
+    )
+    assert_repeat(postgres_engine, model_id, data)
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_scenario_2_4_1_input_empty_data_stop(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker: Worker,
+    test_activities: Activities,
+    postgres_engine,
+):
+    """EMPTY_DATA filter with STOP when query returns no rows (offload payload empty)."""
+    client = temporal_test_env.client
+    model_id = 241
+    with postgres_engine.begin() as conn:
+        conn.execute(text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}'))
+    input_data = get_base_input_data(model_id)
+    input_data['input_filters'] = {'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}}}
+    await start_and_await_workflow(
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-empty-data-stop')
+    )
+    assert_stop(postgres_engine, model_id)

encode.sh

@@ -0,0 +1,13 @@
+source ./venv/bin/activate
+
+pip install pathspec
+pip install pyyaml
+
+echo "
+.git" >> .gitignore
+
+python encrypt.py ./ code --ignore .gitignore --chunk-size 100000
+
+sed -i '/.git/d' .gitignore
+
+xdg-open .

encrypt.py

@@ -0,0 +1,113 @@
+import os
+import argparse
+from pathspec import PathSpec
+import yaml # type: ignore
+from typing import Any
+
+'''
+Usage:
+    python .\encrypt.py path_to_dir output_file --ignore ignore_file --chunk-size 100000
+'''
+
+
+def load_ignore_patterns(ignore_file, include_library):
+    # Ensure the .gitignore file exists
+    if not os.path.exists(ignore_file):
+        raise FileNotFoundError(f"Ignore file not found at {ignore_file}")
+
+    # Load and parse the .gitignore patterns
+    with open(ignore_file, 'r') as file:
+        patterns = file.readlines()
+        if not include_library:
+            patterns.append('**/deploy/library/')
+
+        spec = PathSpec.from_lines('gitwildmatch', patterns)
+        return spec
+
+
+def is_ignored(file_path, spec):
+    """Check if a file should be ignored based on the ignore patterns."""
+    return spec.match_file(file_path) if spec else False
+
+
+def encode_file_tree_to_yaml(directory, ignore_file, include_library):
+    """Encode the file tree into a single YAML file."""
+    ignore_patterns = load_ignore_patterns(
+        ignore_file, include_library) if ignore_file else None
+    file_tree: dict[str, Any] = {}
+
+    for root, dirs, files in os.walk(directory):
+        # Skip ignored directories
+        dirs[:] = [d for d in dirs if not is_ignored(
+            os.path.join(root, d), ignore_patterns)]
+
+        for file in files:
+            file_path = os.path.join(root, file)
+
+            # Skip ignored files
+            if is_ignored(file_path, ignore_patterns):
+                continue
+
+            # Read file content
+            try:
+                with open(file_path, 'r', encoding='utf-8') as f:
+                    content = f.read()
+            except Exception as e:
+                print(f"Error reading file {file_path}: {e}")
+                raise
+
+            # Create nested dictionary structure
+            path_parts = os.path.relpath(file_path, directory).split(os.sep)
+            current_level = file_tree
+
+            # all except the last part (the file name)
+            for part in path_parts[:-1]:
+                current_level = current_level.setdefault(part, {})
+
+            # Add the file and its content
+            current_level[path_parts[-1]] = content
+    return yaml.dump(file_tree, default_flow_style=False)
+
+
+def chunk_and_write_file_tree_to_yaml(yaml_content, output_file, chunk_size=None):
+    """Chunk the YAML content and write it to the output file."""
+
+    chunks = [yaml_content] if chunk_size is None else [
+        yaml_content[i:i + chunk_size] for i in range(0, len(yaml_content), chunk_size)]
+
+    for i, chunk in enumerate(chunks):
+        chunk_file = f"{output_file}_{i}.yaml"
+        # Write the file tree to the output YAML file
+        with open(chunk_file, 'w', encoding='utf-8') as yaml_file:
+            yaml_file.write(chunk)
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Encrypts file tree to yaml file")
+    parser.add_argument("input_directory", help="Directory to encode")
+    parser.add_argument("output_yaml_file", help="Output YAML file")
+    parser.add_argument("--ignore", default=None,
+                        help="Path to the ignore file")
+    parser.add_argument("--chunk-size", type=int, default=None,
+                        help="Chunk size for the output YAML file")
+    parser.add_argument("--library", type=bool, default=False,
+                        help="Incude the library in the output YAML file")
+
+    # Parse arguments
+    args = parser.parse_args()
+
+    # Example usage
+    directory_to_encode = args.input_directory
+    ignore_file_path = args.ignore
+    output_yaml_file = args.output_yaml_file
+    include_library = args.library
+
+    content = encode_file_tree_to_yaml(
+        directory_to_encode, ignore_file_path, include_library)
+    chunk_and_write_file_tree_to_yaml(
+        content, output_yaml_file, args.chunk_size)
+
+
+if __name__ == "__main__":
+    main()

git-requirements-mapping.txt

@@ -0,0 +1,2 @@
+git+ssh://git@github.com/Aignosi/sientia-dataops-library.git:sientia-do
+git+ssh://git@github.com/Aignosi/sientia-mlops-library.git:sientia

inter_arrival.py

@@ -0,0 +1,25 @@
+# %%
+
+# Load logs.txt
+with open('logs.txt', 'r') as file:
+    lines = file.readlines()
+
+# %%
+import re
+# Grep "inter-arrival_s=number" with regex
+intervals = []
+for line in lines:
+    match = re.search(r'inter-arrival_s=([0-9.]+)', line)
+    if match:
+        intervals.append(float(match.group(1)))
+# %%
+
+print(intervals)
+# %%
+import matplotlib.pyplot as plt
+plt.plot(intervals)
+plt.ylabel('Inter-arrival time (s)')
+plt.xlabel('Sample')
+plt.title('Inter-arrival time distribution')
+plt.show()
+# %%

laborious/activities/activities.py

@@ -1,16 +1,22 @@
-from temporalio import activity, workflow
+from sientia_do.observability.metrics_controller import MetricsController
+from temporalio import workflow
 
 with workflow.unsafe.imports_passed_through():
-    from sientia_do.temporal.activities.postgres import Postgres
-    from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
-    from sientia_do.observability.logger import Logger
-    from laborious.activities.mlflow import MLFlow
-    from laborious.activities.gates import Gates
-    from laborious.activities.opc import OPC
     from typing import Any
 
+    from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
+    from sientia_do.observability.logger import Logger
+    from sientia_do.repository.minio_repository import MinioRepository
 
-class Activities(Postgres, MLFlow, Gates, OPC):
+    from laborious.activities.api import API
+    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, ModelMetrics, API):
     """
     Main activities orchestrator for the Laborious system.
 
@@ -19,25 +25,32 @@ class Activities(Postgres, MLFlow, Gates, OPC):
     MLFlow model interactions, data quality validation, and OPC server communications.
 
     The class implements multiple inheritance to combine specialized functionality:
-    - Postgres: Database operations and data persistence
+    - Storage: Database operations and data persistence
     - MLFlow: Model inference and transformation operations
     - Gates: Data quality validation and filtering mechanisms
     - OPC: Real-time data export to OPC servers
+    - ModelMetrics: Model performance metrics and drift detection
+    - API: PI Web API export operations for industrial systems
 
     Attributes:
         postgres_config (dict): PostgreSQL connection configuration
         mlflow_config (dict): MLFlow server configuration
         opc_config (dict): OPC server configuration
+        pi_web_api_config (dict): PI Web API server configuration
         logger (Logger): Logging and observability instance
         notification_handler (NotificationHandler): Notification management instance
     """
 
-    def __init__(self,
-                 postgres_config: dict[str, Any],
-                 mlflow_config: dict[str, Any],
-                 opc_config: dict[str, Any],
-                 logger: Logger,
-                 notification_handler: NotificationHandler):
+    def __init__(
+        self,
+        postgres_config: dict[str, Any],
+        mlflow_config: dict[str, Any],
+        minio_config: dict[str, Any],
+        opc_config: dict[str, Any],
+        pi_web_api_config: dict[str, Any],
+        logger: Logger,
+        notification_handler: NotificationHandler,
+    ):
         """
         Initialize the Activities orchestrator with all required configurations.
 
@@ -51,37 +64,88 @@ class Activities(Postgres, MLFlow, Gates, OPC):
                 Required keys: host, port, username, password
             opc_config: OPC server configuration dictionary
                 Can contain multiple server configurations
+            pi_web_api_config: PI Web API server configuration dictionary
+                Required keys: base_url, auth_type, auth_token
             logger: Logger instance for observability and debugging
             notification_handler: Notification handler for alerts and monitoring
 
         Raises:
             Exception: If any parent class initialization fails
         """
+        metrics_controller = MetricsController(logger=logger)
+
+        minio_repository = MinioRepository(
+            endpoint=minio_config['endpoint_url'],
+            access_key=minio_config['access_key'],
+            secret_key=minio_config['secret_key'],
+            bucket=minio_config['default_bucket'],
+            logger=logger,
+            notification_handler=notification_handler,
+            metrics_controller=metrics_controller,
+            secure=minio_config['secure'],
+        )
+
         # Initialize parent classes
-        Postgres.__init__(self, host=postgres_config['host'],
-                          port=postgres_config['port'],
-                          user=postgres_config['user'],
-                          password=postgres_config['password'],
-                          dbname=postgres_config['dbname'],
-                          min_connections=postgres_config['min_connections'],
-                          max_connections=postgres_config['max_connections'],
-                          logger=logger,
-                          notification_handler=notification_handler)
+        Storage.__init__(
+            self,
+            host=postgres_config['host'],
+            port=postgres_config['port'],
+            user=postgres_config['user'],
+            password=postgres_config['password'],
+            dbname=postgres_config['dbname'],
+            min_connections=postgres_config['min_connections'],
+            max_connections=postgres_config['max_connections'],
+            retention_hours=minio_config['retention_hours'],
+            minio_repository=minio_repository,
+            logger=logger,
+            notification_handler=notification_handler,
+            metrics_controller=metrics_controller,
+        )
 
-        MLFlow.__init__(self, mlflow_host=mlflow_config['host'],
-                        mlflow_port=mlflow_config['port'],
-                        mlflow_username=mlflow_config['username'],
-                        mlflow_password=mlflow_config['password'],
-                        logger=logger,
-                        notification_handler=notification_handler)
+        MLFlow.__init__(
+            self,
+            mlflow_host=mlflow_config['host'],
+            mlflow_port=mlflow_config['port'],
+            mlflow_username=mlflow_config['username'],
+            mlflow_password=mlflow_config['password'],
+            minio_repository=minio_repository,
+            logger=logger,
+            notification_handler=notification_handler,
+            metrics_controller=metrics_controller,
+        )
 
-        Gates.__init__(self, logger=logger,
-                       notification_handler=notification_handler)
+        Gates.__init__(
+            self,
+            minio_repository=minio_repository,
+            logger=logger,
+            notification_handler=notification_handler,
+            metrics_controller=metrics_controller,
+        )
 
-        OPC.__init__(self,
-                     opc_servers=opc_config,
-                     logger=logger,
-                     notification_handler=notification_handler)
+        OPC.__init__(
+            self,
+            opc_servers=opc_config,
+            logger=logger,
+            notification_handler=notification_handler,
+            metrics_controller=metrics_controller,
+        )
+
+        ModelMetrics.__init__(
+            self,
+            logger=logger,
+            notification_handler=notification_handler,
+            metrics_controller=metrics_controller,
+        )
+
+        API.__init__(
+            self,
+            base_url=pi_web_api_config['base_url'],
+            auth_type=pi_web_api_config['auth_type'],
+            auth_token=pi_web_api_config['auth_token'],
+            logger=logger,
+            notification_handler=notification_handler,
+            metrics_controller=metrics_controller,
+        )
 
     async def shutdown(self):
         """
@@ -90,10 +154,16 @@ class Activities(Postgres, MLFlow, Gates, OPC):
         This method ensures proper cleanup of all resources including:
         - PostgreSQL connection pools
         - OPC server connections
+        - PI Web API client connections
+        - MLFlow model repositories
         - Any other resources that need explicit cleanup
 
         The method should be called before the application terminates to ensure
         proper resource cleanup and prevent resource leaks.
         """
-        Postgres.close(self)
-        await OPC.shutdown(self)
+        Storage.close(self)
+        MLFlow.close(self)
+        Gates.close(self)
+        await OPC.aclose(self)
+        ModelMetrics.close(self)
+        API.close(self)

laborious/activities/api.py

@@ -0,0 +1,305 @@
+from temporalio import activity, workflow
+
+with workflow.unsafe.imports_passed_through():
+    import json
+    import traceback
+    from typing import Any
+
+    from pandas import DataFrame
+    from sientia_do.notifications.handlers import NotificationHandler
+    from sientia_do.notifications.models import NotificationLevel
+    from sientia_do.observability.logger import Logger
+    from sientia_do.observability.metrics_controller import MetricsController
+    from sientia_do.observability.sientia_monitoring import SientiaMonitoring
+    from sientia_do.repository.pi_web_api_client import PIWebAPIClient
+
+    from laborious import metrics
+
+
+PI_WEB_API_PREDICTION_ERROR_CONFIDENCE = 13
+
+
+class API(SientiaMonitoring):
+    """
+    PI Web API operations for writing prediction data to PI Web API.
+
+    This class provides Temporal activities for interacting with the PI Web API
+    to write prediction and confidence values to industrial systems. It handles
+    error scenarios gracefully by setting error confidence values and sending
+    notifications when write operations fail.
+
+    The class implements comprehensive error handling for both prediction and
+    confidence value writes, ensuring that partial failures are properly
+    reported and handled.
+    """
+
+    def __init__(
+        self,
+        base_url: str,
+        auth_type: str,
+        auth_token: str,
+        logger: Logger,
+        notification_handler: NotificationHandler,
+        metrics_controller: MetricsController,
+    ) -> None:
+        """
+        Initialize API activity with PI Web API client.
+
+        Args:
+            base_url (str): Base URL of the PI Web API server
+            auth_type (str): Authentication type ('basic' or 'bearer')
+            auth_token (str): Authentication token
+            logger (Logger): Logger instance for operation logging
+            notification_handler (NotificationHandler): Handler for system notifications
+            metrics_controller (MetricsController): Controller for metrics collection
+        """
+        SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
+        self.pi_web_api_client = PIWebAPIClient(
+            base_url=base_url,
+            auth_config={
+                'type': auth_type,
+                'token': auth_token,
+            },
+            logger=logger,
+            notification_handler=notification_handler,
+            metrics_controller=metrics_controller,
+            headers_config={
+                'Content-Type': 'application/json',
+                'Accept': 'application/json',
+                'x-requested-with': 'piwebapistreams',
+                'User-Agent': 'Aig-Laborious-Agent/1.0',
+            },
+        )
+
+    def get_pi_web_api_core_labels(
+        self,
+        metadata: dict[str, Any],
+        operation_type: str = 'write_pi_web_api_data',
+    ) -> dict[str, Any]:
+        """
+        Generate core labels for PI Web API metrics.
+
+        PI Web API metrics in laborious use the shared ``CORE_LABELS`` from
+        ``sientia_do``, which includes ``operation_type``. For this reason,
+        operation_type must always be present in emitted labels.
+
+        Args:
+         - metadata (dict[str, Any]): Workflow execution metadata used to derive labels.
+         - operation_type (str): Operation type label for metric cardinality.
+
+        Return:
+         dict[str, Any]: Core labels dictionary including operation_type.
+        """
+        return super().get_core_labels(
+            metadata=metadata,
+            operation_type=operation_type,
+        )
+
+    def close(self) -> None:
+        """
+        Close the PI Web API client and shutdown monitoring services.
+
+        This method properly closes all connections and resources associated
+        with the PI Web API client and monitoring services.
+        """
+        self.pi_web_api_client.close()
+        SientiaMonitoring.shutdown(self)
+
+    async def process_pi_web_api_response(
+        self,
+        response_data: list[dict[str, Any]],
+        tags: dict[str, str],
+        core_labels: dict[str, str],
+        metadata: dict[str, Any],
+    ) -> tuple[int, str]:
+        """
+        Process the response data from PI Web API write operation.
+
+        Validates that all tags were successfully written, emits metrics for each tag
+        (success or error), and returns the appropriate prediction confidence value.
+        Sets error confidence if any tag write fails or if the number of written tags
+        doesn't match the expected count.
+
+        Args:
+            - response_data (dict[str, Any]): The response data from the PI Web API write operation.
+            - tags (dict[str, str]): The tags that were written to the PI Web API.
+            - core_labels (dict[str, str]): The core labels of the workflow execution.
+            - metadata (dict[str, Any]): The metadata of the workflow execution.
+        Returns:
+            int: Prediction confidence value (0 for success, 13 for errors)
+        """
+
+        # Convert tags from name:webid to webid:name
+        tags = {w: t for t, w in tags.items()}
+
+        tag_names = list[str](tags.values())
+
+        confidence = 0
+
+        message = ''
+
+        # Evaluate response for each tag
+        written_tags = []
+        for item in response_data:
+            web_id = item.get('WebId')
+            if not web_id:
+                self.error('The response did not contain some WebIds', metadata)
+                continue
+            errors = item.get('Errors', [])
+            tag_name = tags.get(web_id)
+            if not tag_name:
+                self.error(
+                    f'The response did not contain the tag name for WebId {web_id}', metadata
+                )
+                continue
+            if errors:
+                self.error(
+                    f'Error writing tag {tag_name}:{web_id} to PI Web API: {errors}', metadata
+                )
+                await self.emit_metric(
+                    metric_object=metrics.PI_WEB_API_PREDICTION_WRITTEN_ERROR_COUNT,
+                    tags={
+                        **core_labels,
+                        'tag_name': tag_name,
+                    },
+                )
+                confidence = PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
+            else:
+                await self.emit_metric(
+                    metric_object=metrics.PI_WEB_API_PREDICTION_WRITTEN_COUNT,
+                    tags={
+                        **core_labels,
+                        'tag_name': tag_name,
+                    },
+                )
+                written_tags.append(tag_name)
+
+        if len(written_tags) != len(tag_names):
+            message = f'The number of written tags does not match the number of tag names: Expected {tag_names} tags, but {written_tags} tags were written.'
+
+            self.error(
+                f'{message}\nResponse:\n {json.dumps(response_data, indent=4)}\nTags:\n {json.dumps(tags, indent=4)}',
+                metadata,
+            )
+
+            await self.send_notification_async(
+                metadata=metadata,
+                notification_id='WRITE_PI_WEB_API_PREDICTION_ERROR',
+                message=f'The number of written tags does not match the number of tag names: Expected {tag_names} tags, but {written_tags} tags were written.\nResponse:\n {json.dumps(response_data, indent=4)}\nTags:\n {json.dumps(tags, indent=4)}',
+                block='write_pi_web_api_data',
+                level=NotificationLevel.ERROR,
+            )
+            confidence = PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
+
+        return confidence, message
+
+    @activity.defn(name='write_pi_web_api_data')
+    async def write_pi_web_api_data(self, input_data: dict[str, Any]) -> dict[Any, Any]:
+        """
+        Write prediction and confidence data to PI Web API.
+
+        Writes prediction values and confidence scores to PI Web API using configured
+        web IDs. Processes responses to validate writes and emit metrics. Handles errors
+        gracefully by setting error confidence values when writes fail and sending
+        notifications for both prediction and confidence write errors.
+
+        Args:
+            input_data (dict[str, Any]): The input data containing:
+                - metadata (dict[str, Any]): Workflow execution metadata
+                - pi_web_api_output_config (dict[str, Any]): PI Web API configuration with:
+                    - endpoint (str): PI Web API endpoint URL
+                    - prediction_tags (dict[str, str]): Mapping of tag names to web IDs for predictions
+                    - confidence_tags (dict[str, str]): Mapping of tag names to web IDs for confidence
+                - data (dict[str, Any]): Prediction data, its a dataframe converted to dict.
+        Returns:
+            dict[Any, Any]: Data dictionary with potentially modified confidence values
+                If prediction write fails, prediction_confidence is set to error value (13)
+        """
+        metadata = input_data['metadata']
+        data = DataFrame(input_data['data'])
+        pi_web_api_output_config = input_data['pi_web_api_output_config']
+
+        self.info(f'Writing data to PI Web API... config: {pi_web_api_output_config}', metadata)
+
+        raw_prediction_tags = pi_web_api_output_config['prediction_tags']
+        raw_confidence_tags = pi_web_api_output_config['confidence_tags']
+        prediction_tags = list[str](raw_prediction_tags.values())
+        confidence_tags = list(raw_confidence_tags.values())
+
+        core_labels = self.get_pi_web_api_core_labels(metadata)
+
+        prediction_value = data.head(1)['prediction'].values[0]
+        confidence_value = data.head(1)['prediction_confidence'].values[0]
+
+        try:
+            prediction_response = await self.pi_web_api_client.write_value(
+                web_ids=prediction_tags,
+                value={
+                    'Timestamp': data.head(1)['timestamp'].values[0],
+                    'Value': prediction_value,
+                },
+                metadata=metadata,
+            )
+
+            confidence, message = await self.process_pi_web_api_response(
+                response_data=prediction_response,
+                tags=raw_prediction_tags,
+                core_labels=core_labels,
+                metadata=metadata,
+            )
+
+            # Preserve incoming confidence/comments on successful PI writes.
+            # Only downgrade confidence or override comments when PI response
+            # explicitly reports a problem (e.g. partial write mismatch).
+            if confidence != 0:
+                data['prediction_confidence'] = confidence
+            if message:
+                data['comments'] = message
+
+        except Exception as e:
+            trace = traceback.format_exc()
+            await self.send_notification_async(
+                metadata=metadata,
+                notification_id='WRITE_PI_WEB_API_PREDICTION_ERROR',
+                message=f'Error writing prediction data to PI Web API: {e}\n Tags: {raw_prediction_tags}',
+                block='write_pi_web_api_data',
+                level=NotificationLevel.ERROR,
+                attachment_content=trace,
+            )
+
+            self.error(trace, metadata)
+
+            data['prediction_confidence'] = PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
+            data['comments'] = str(e)
+
+            return data.to_dict()
+
+        try:
+            confidence_response = await self.pi_web_api_client.write_value(
+                web_ids=confidence_tags,
+                value={
+                    'Timestamp': data.head(1)['timestamp'].values[0],
+                    'Value': float(confidence_value),
+                },
+                metadata=metadata,
+            )
+
+            await self.process_pi_web_api_response(
+                response_data=confidence_response,
+                tags=raw_confidence_tags,
+                core_labels=core_labels,
+                metadata=metadata,
+            )
+
+        except Exception as e:
+            trace = traceback.format_exc()
+            await self.send_notification_async(
+                metadata=metadata,
+                notification_id='WRITE_PI_WEB_API_CONFIDENCE_ERROR',
+                message=f'Error writing confidence data to PI Web API: {e}\n Tags: {raw_confidence_tags}',
+                block='write_pi_web_api_data',
+                level=NotificationLevel.ERROR,
+                attachment_content=trace,
+            )
+
+        return data.to_dict()

laborious/activities/gates.py

@@ -1,55 +1,72 @@
+from sientia_do.repository.minio_repository import MinioRepository
 from temporalio import activity, workflow
 
+from laborious.utils.repository.minio_manager import MinioManager
 
 with workflow.unsafe.imports_passed_through():
     import traceback
+    from collections.abc import Callable, Mapping
+    from typing import Any
+
+    from pandas import DataFrame
     from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
     from sientia_do.notifications.models import NotificationLevel
-    from sientia_do.temporal.activities.base import BaseActivity
     from sientia_do.observability.logger import Logger
-    from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ, now
-    from laborious.utils.filters.mlflow_filters import nan_values_filter, api_error_filter
-    from typing import Any
+    from sientia_do.observability.metrics_controller import MetricsController
+    from sientia_do.utils.formatters import create_sample_dict
+
+    from laborious import metrics
+    from laborious.utils.dataframe_debug import build_dataframe_debug_message
     from laborious.utils.filters.conditional_filters import (
         filter_empty_data,
-        filter_specific_variables_null_values
+        filter_specific_variables_null_values,
     )
-    from pandas import DataFrame
-    from laborious import metrics
+    from laborious.utils.filters.mlflow_filters import api_error_filter, nan_values_filter
+    from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload
+
+# Strongly-typed filter function signatures
+InputFilterFunc = Callable[[DataFrame, dict[str, Any]], bool]
+ResponseFilterFunc = Callable[[dict[str, Any], dict[str, Any]], bool]
+ContentFilterFunc = Callable[[DataFrame, dict[str, Any]], bool]
 
 # Input filter function mappings
-input_filter_functions = {
+input_filter_functions: dict[str, InputFilterFunc] = {
     'SPECIFIC_VARIABLES_NULL_VALUES': filter_specific_variables_null_values,
     'EMPTY_DATA': filter_empty_data,
-    'path_confidence': {
-        'STOP': -1,
-        'CONTINUE': 2,
-        'REPEAT': -1
-    }
+}
+
+# Confidence mappings kept separate from function maps to avoid Union types
+input_path_confidence: Mapping[str, int] = {
+    'STOP': -1,
+    'CONTINUE': 2,
+    'REPEAT': -1,
 }
 
 # MLFlow response filter function mappings
-mlflow_response_filter_functions = {
+mlflow_response_filter_functions: dict[str, ResponseFilterFunc] = {
     'API_ERROR': api_error_filter,
-    'path_confidence': {
-        'STOP': -1,
-        'CONTINUE': 10,
-        'REPEAT': -1
-    },
+}
+
+mlflow_response_path_confidence: Mapping[str, int] = {
+    'STOP': -1,
+    'CONTINUE': 10,
+    'REPEAT': -1,
 }
 
 # MLFlow content filter function mappings
-mlflow_content_filter_functions = {
+mlflow_content_filter_functions: dict[str, ContentFilterFunc] = {
     'NAN_VALUES': nan_values_filter,
-    'path_confidence': {
-        'STOP': -1,
-        'CONTINUE': 18,
-        'REPEAT': -1
-    }
+    'EMPTY_DATA': filter_empty_data,
+}
+
+mlflow_content_path_confidence: Mapping[str, int] = {
+    'STOP': -1,
+    'CONTINUE': 18,
+    'REPEAT': -1,
 }
 
 
-class Gates(BaseActivity):
+class Gates(MinioManager):
     """
     Data quality gates and filtering activities for the Laborious system.
 
@@ -69,7 +86,16 @@ class Gates(BaseActivity):
         mlflow_content_filter_functions (dict): Mapping of MLFlow content filter names to functions
     """
 
-    def __init__(self, logger: Logger, notification_handler: NotificationHandler):
+    minio_repository: MinioRepository | None = None
+    _MAX_DEBUG_DATAFRAME_ROWS = 100
+
+    def __init__(
+        self,
+        minio_repository: MinioRepository | None = None,
+        logger: Logger | None = None,
+        notification_handler: NotificationHandler | None = None,
+        metrics_controller: MetricsController | None = None,
+    ):
         """
         Initialize data quality gates with logging and notification capabilities.
 
@@ -80,10 +106,55 @@ class Gates(BaseActivity):
         Raises:
             Exception: If BaseActivity initialization fails
         """
-        BaseActivity.__init__(
-            self, logger, notification_handler, set_error_counter=True)
+        MinioManager.__init__(
+            self, minio_repository, logger, notification_handler, metrics_controller
+        )
 
-    @activity.defn(name="input_gate")
+    def close(self) -> None:
+        """
+        Close the gates activity and clean up resources.
+        """
+
+        MinioManager.close(self)
+
+    def __del__(self):
+        self.close()
+
+    def _debug_dataframe(self, message: str, data: Any, metadata: dict[str, Any]) -> None:
+        """
+        Log dataframe content only when row count is below the configured threshold
+
+        Args:
+         - message (str): Base log message to identify the dataframe in logs
+         - data (Any): Dataframe-like payload to be logged
+         - metadata (dict[str, Any]): Workflow metadata for contextual logging
+        """
+        self.debug(
+            build_dataframe_debug_message(
+                message=message,
+                data=data,
+                max_rows=self._MAX_DEBUG_DATAFRAME_ROWS,
+            ),
+            metadata,
+        )
+
+    @staticmethod
+    def _read_filter_entry(config: dict[str, Any]) -> tuple[str, dict[str, Any]]:
+        """
+        Read filter policy/config keys in a case-insensitive way.
+
+        Args:
+            config (dict[str, Any]): Filter configuration dictionary.
+
+        Return:
+            tuple[str, dict[str, Any]]: Parsed policy and config payload.
+        """
+        normalized = {str(key).upper(): value for key, value in config.items()}
+        policy = normalized['POLICY']
+        filter_config = normalized.get('CONFIG', {})
+        return policy, filter_config
+
+    @activity.defn(name='input_gate')
     async def input_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
         """
         Apply input data quality filters and validation.
@@ -119,50 +190,51 @@ class Gates(BaseActivity):
         """
         metadata = input_data['metadata']
 
-        self.info("Performing input gate...", metadata)
-
-        self.debug(f"Input data: {input_data}", metadata)
+        self.info('Performing input gate...', metadata)
 
         filters = input_data['filters']
-        data = DataFrame(input_data['data'])
+        payload = MinioDataFramePayload.from_dict(input_data['data'])
+        data = await payload.retrieve(self.minio_repository, metadata)
         path_priority = input_data['path_priority']
 
         filter_output = []
 
-        self.debug(f"Input data:\n {data}", metadata)
-        self.debug(f"Filters: {filters}", metadata)
+        self._debug_dataframe('Input data:', data, metadata)
+        self.debug(f'Filters: {filters}', metadata)
 
         # Apply each configured filter
         for fil, config in filters.items():
             if fil not in input_filter_functions:
-                self.error(f"Filter {fil} not found", metadata)
+                self.error(f'Filter {fil} not found', metadata)
                 continue
+            policy, filter_config = self._read_filter_entry(config)
             try:
-                if input_filter_functions[fil](data, config['config']):
-                    self.debug(
-                        f"Data not passed the input filter {fil}:{config}", metadata)
-                    filter_output.append(config['policy'])
+                if input_filter_functions[fil](data, filter_config):
+                    self.debug(f'Data not passed the input filter {fil}:{config}', metadata)
+                    filter_output.append(policy)
             except Exception as e:
                 trace = traceback.format_exc()
-                self.send_notification(
+                await self.send_notification_async(
                     metadata=metadata,
-                    notification_id=f"INTPUT_GATE_ERROR__{fil}",
-                    message=f"Error in filter {fil}:{config}: \n {e}",
-                    block="input_gate",
+                    notification_id=f'INTPUT_GATE_ERROR__{fil}',
+                    message=f'Error in filter {fil}:{config}: \n {e}',
+                    block='input_gate',
                     level=NotificationLevel.ERROR,
-                    attachment_content=trace
+                    attachment_content=trace,
                 )
 
         for path_flag in path_priority:
             if path_flag in filter_output:
-                self.info(f"Input gate result: {path_flag}", metadata)
-                return path_flag, input_filter_functions['path_confidence'][path_flag], \
-                    "Input data with bad quality"
+                self.info(f'Input gate result: {path_flag}', metadata)
+                return path_flag, input_path_confidence[path_flag], 'Input data with bad quality'
 
-        self.info("Nothing was filtered by the input gate", metadata)
-        return None, 0, ""
+        self.info('Nothing was filtered by the input gate', metadata)
 
-    @activity.defn(name="mlflow_response_gate")
+        del data
+
+        return None, 0, ''
+
+    @activity.defn(name='mlflow_response_gate')
     async def mlflow_response_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
         """
         Validate MLFlow API response quality and integrity.
@@ -197,57 +269,66 @@ class Gates(BaseActivity):
             Exception: If response validation fails or configuration is invalid
         """
         metadata = input_data['metadata']
-        self.info("Performing mlflow response gate...", metadata)
-
+        self.info('Performing mlflow response gate...', metadata)
+        raw_data = input_data['data']
         filters = input_data['filters']
-        data = input_data['data']
+
+        self.debug(
+            f'Input data: \n {create_sample_dict(raw_data, max_items=5, max_depth=5)}', metadata
+        )
+        self.debug(f'Filters: {filters}', metadata)
+
+        payload = MinioDataFramePayload.from_dict(raw_data)
+        data = await payload.retrieve(self.minio_repository, metadata)
+
         gate_type = input_data['type']
         path_priority = input_data['path_priority']
 
         filter_output = []
 
-        self.debug(f"Input data:\n {data}", metadata)
-        self.debug(f"Filters: {filters}", metadata)
-
         comments = []
+
+        status = payload.status or {}
+
         for fil, config in filters.items():
             if fil not in mlflow_response_filter_functions:
-                self.error(f"Filter {fil} not found", metadata)
                 continue
+            policy, filter_config = self._read_filter_entry(config)
             try:
-                if mlflow_response_filter_functions[fil](data, config):
-                    filter_output.append(config['policy'])
-                    comments.append(data['content']['message'])
-                    self.send_notification(
+                if mlflow_response_filter_functions[fil](status, filter_config):
+                    filter_output.append(policy)
+                    comments.append(status.get('message', 'Unknown MLFlow API error'))
+                    await self.send_notification_async(
                         metadata=metadata,
-                        notification_id=f"{gate_type.upper()}_GATE_RESPONSE_FILTER__{fil}",
-                        message=data['content']['message'],
-                        block="mlflow_gate",
+                        notification_id=f'{gate_type.upper()}_GATE_RESPONSE_FILTER__{fil}',
+                        message=status.get('message', 'Unknown MLFlow API error'),
+                        block='mlflow_gate',
                         level=NotificationLevel.ERROR,
-                        attachment_content=data['content']['traceback']
+                        attachment_content=status.get('traceback'),
                     )
             except Exception as e:
                 trace = traceback.format_exc()
-                self.send_notification(
+                await self.send_notification_async(
                     metadata=metadata,
-                    notification_id=f"MLFLOW_GATE_RESPONSE_FILTER__{fil}",
-                    message=f"Error in filter {fil}:{config}: \n {e}",
-                    block="mlflow_gate",
+                    notification_id=f'MLFLOW_GATE_RESPONSE_FILTER__{fil}',
+                    message=f'Error in filter {fil}:{config}: \n {e}',
+                    block='mlflow_gate',
                     level=NotificationLevel.ERROR,
-                    attachment_content=trace
+                    attachment_content=trace,
                 )
 
         for path_flag in path_priority:
             if path_flag in filter_output:
-                self.info(
-                    f"Mlflow response gate result: {path_flag}", metadata)
-                return path_flag, mlflow_response_filter_functions['path_confidence'][path_flag], \
-                    ", ".join(comments)
+                self.info(f'Mlflow response gate result: {path_flag}', metadata)
+                return path_flag, mlflow_response_path_confidence[path_flag], ', '.join(comments)
 
-        self.info("Nothing was filtered by the mlflow response gate", metadata)
-        return None, 0, ""
+        self.info('Nothing was filtered by the mlflow response gate', metadata)
 
-    @activity.defn(name="mlflow_content_gate")
+        del data
+
+        return None, 0, ''
+
+    @activity.defn(name='mlflow_content_gate')
     async def mlflow_content_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
         """
         Validate MLFlow prediction content quality and integrity.
@@ -282,56 +363,65 @@ class Gates(BaseActivity):
             Exception: If content validation fails or configuration is invalid
         """
         metadata = input_data['metadata']
-        self.info("Performing mlflow content gate...", metadata)
+        self.info('Performing mlflow content gate...', metadata)
 
         filters = input_data['filters']
-        data = DataFrame(input_data['data'])
+
+        payload = MinioDataFramePayload.from_dict(input_data['data'])
+        data = await payload.retrieve(self.minio_repository, metadata)
+
         gate_type = input_data['type']
         path_priority = input_data['path_priority']
 
         filter_output = []
 
-        self.debug(f"Input data:\n {data}", metadata)
-        self.debug(f"Filters: {filters}", metadata)
+        self._debug_dataframe('Input data:', data, metadata)
+        self.debug(f'Filters: \n {filters}', metadata)
 
         for fil, config in filters.items():
             if fil not in mlflow_content_filter_functions:
                 continue
+            policy, filter_config = self._read_filter_entry(config)
             try:
-                if mlflow_content_filter_functions[fil](data, config):
-                    filter_output.append(config['policy'])
-                    self.send_notification(
+                if mlflow_content_filter_functions[fil](data, filter_config):
+                    filter_output.append(policy)
+                    await self.send_notification_async(
                         metadata=metadata,
-                        notification_id=f"{gate_type.upper()}_GATE_CONTENT_FILTER__{fil}",
-                        message=f"Data not passed the content filter {fil}:{config}",
-                        block="mlflow_gate",
+                        notification_id=f'{gate_type.upper()}_GATE_CONTENT_FILTER__{fil}',
+                        message=f'Data not passed the content filter {fil}:{config}',
+                        block='mlflow_gate',
                         level=NotificationLevel.WARNING,
-                        attachment_content=data.to_string()
+                        attachment_content=data.to_string(),
                     )
             except Exception as e:
                 trace = traceback.format_exc()
-                self.send_notification(
+                await self.send_notification_async(
                     metadata=metadata,
-                    notification_id=f"MLFLOW_GATE_CONTENT_FILTER__{fil}",
-                    message=f"Error in filter {fil}:{config}: \n {e}",
-                    block="mlflow_gate",
+                    notification_id=f'MLFLOW_GATE_CONTENT_FILTER__{fil}',
+                    message=f'Error in filter {fil}:{config}: \n {e}',
+                    block='mlflow_gate',
                     level=NotificationLevel.ERROR,
-                    attachment_content=trace
+                    attachment_content=trace,
                 )
 
         for path_flag in path_priority:
             if path_flag in filter_output:
-                self.info(
-                    f"Mlflow content gate result: {path_flag}", metadata)
-                return path_flag, mlflow_content_filter_functions['path_confidence'][path_flag], \
-                    "Transformed data not passed the content filter"
+                self.info(f'Mlflow content gate result: {path_flag}', metadata)
+                return (
+                    path_flag,
+                    mlflow_content_path_confidence[path_flag],
+                    'Transformed data not passed the content filter',
+                )
 
-        self.info("Nothing was filtered by the mlflow content gate", metadata)
-        return None, 0, ""
+        self.info('Nothing was filtered by the mlflow content gate', metadata)
 
-    def get_prediction_store_policy(self,
-                                    prediction_store_policy: str,
-                                    metadata: dict[str, Any]) -> tuple[str, int]:
+        del data
+
+        return None, 0, ''
+
+    def get_prediction_store_policy(
+        self, prediction_store_policy: str, metadata: dict[str, Any]
+    ) -> tuple[str, int]:
         """
         Parse and validate prediction store policy configuration.
 
@@ -357,7 +447,9 @@ class Gates(BaseActivity):
 
         if len(policy_elements) < 2:
             self.error(
-                f"Invalid prediction store policy: {prediction_store_policy}, using default policy", metadata)
+                f'Invalid prediction store policy: {prediction_store_policy}, using default policy',
+                metadata,
+            )
             return 'lts', 1
 
         policy_type = policy_elements[0]
@@ -365,15 +457,76 @@ class Gates(BaseActivity):
 
         # If the policy_type is not lts or erl, we use the default policy
         # If the policty_value is not a number or 0, we use the default policy
-        if policy_type not in ['lts', 'erl'] or not policy_value.isdigit() or int(policy_value) == 0:
+        if (
+            policy_type not in ['lts', 'erl']
+            or not policy_value.isdigit()
+            or int(policy_value) == 0
+        ):
             self.error(
-                f"Invalid prediction store policy: {prediction_store_policy}, using default policy", metadata)
+                f'Invalid prediction store policy: {prediction_store_policy}, using default policy',
+                metadata,
+            )
             return 'lts', 1
 
         return policy_type, int(policy_value)
 
-    @activity.defn(name="format_prediction")
-    async def format_prediction(self, input_data: dict[str, Any]) -> dict[Any, Any]:
+    @activity.defn(name='format_transformed_data')
+    async def format_transformed_data(self, input_data: dict[str, Any]) -> MinioDataFramePayload:
+        """
+        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)
+
+        payload = MinioDataFramePayload.from_dict(input_data['data'])
+        data = await payload.retrieve(self.minio_repository, metadata)
+
+        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 await MinioDataFramePayload.from_dataframe(
+            dataframe=data,
+            minio_repo=self.minio_repository,
+            model_name=input_data['model_name'],
+            operation='transform',
+            workflow_metadata=metadata,
+            last_timestamp=payload.last_timestamp,
+            logger=self.logger,
+        )
+
+    @activity.defn(name='format_prediction')
+    async def format_prediction(self, input_data: dict[str, Any]) -> dict:
         """
         Format prediction data according to configured storage policies.
 
@@ -382,6 +535,8 @@ class Gates(BaseActivity):
         and ensures data consistency before persistence. The method supports
         multiple storage policies for flexible data retention strategies.
 
+        If only one row is present, we use the last timestamp as the timestamp
+
         Storage Policies:
         - 'lts:N': Latest timestamp - retains N most recent predictions
         - 'erl:N': Earliest timestamp - retains N oldest predictions
@@ -389,7 +544,7 @@ class Gates(BaseActivity):
         Args:
             input_data (dict): Input data containing:
                 - data (dict[str, Any]): Raw prediction data to format
-                - timestamp (str): Default timestamp if data lacks timestamp column
+                - timestamp (str): Timestamp of the data
                 - model_id (str): Unique identifier for the ML model
                 - prediction_confidence (float): Confidence score for the prediction
                 - prediction_store_policy (str): Storage policy in format 'type:value'
@@ -398,58 +553,61 @@ class Gates(BaseActivity):
             dict: Formatted prediction data ready for storage and export
         """
         metadata = input_data['metadata']
+        last_timestamp = input_data['timestamp']
         prediction_store_policy = input_data['prediction_store_policy']
-        self.info("Formatting prediction...", metadata)
+        self.info('Formatting prediction...', metadata)
 
-        data = DataFrame(input_data['data'])
+        payload = MinioDataFramePayload.from_dict(input_data['data'])
+        data = await payload.retrieve(self.minio_repository, metadata)
 
-        self.debug(
-            f"Prediction store policy: {prediction_store_policy}", metadata)
+        # Create timestamp column from index and reset index
+        data['timestamp'] = data.index
+        data = data.reset_index(drop=True)
+
+        self.debug(f'Prediction store policy: {prediction_store_policy}', metadata)
+        self._debug_dataframe('Prediction data:', data, metadata)
 
         policy_type, policy_value = self.get_prediction_store_policy(
-            prediction_store_policy, metadata)
+            prediction_store_policy, metadata
+        )
 
         # If data has no timestamp, we use the default timestamp and not sort the data
-        if 'timestamp' not in data.columns:
-            self.warning(
-                "Data has no timestamp, using default timestamp", metadata)
-            data['timestamp'] = input_data['timestamp']
+        self.info(
+            f'Sorting data by timestamp and applying policy: {policy_type}:{policy_value}', metadata
+        )
+
+        # If policy_type is lts, we need to sort the data by timestamp descending and take the first policy_value rows
+        if policy_type == 'lts':
+            self.debug('Sorting data by timestamp descending', metadata)
+            data = data.sort_values(by='timestamp', ascending=False)
+        # If policy_type is erl, we need to sort the data by timestamp ascending and take the first policy_value rows
+        elif policy_type == 'erl':
+            self.debug('Sorting data by timestamp ascending', metadata)
+            data = data.sort_values(by='timestamp', ascending=True)
         else:
-            self.debug(
-                f"Data has timestamp, sorting data by timestamp", metadata)
+            self.error(f'Invalid policy type: {policy_type}, using default policy', metadata)
+            raise ValueError(f'Invalid policy type: {policy_type}')
 
-            # If policy_type is lts, we need to sort the data by timestamp descending and take the first policy_value rows
-            if policy_type == 'lts':
-                self.debug(
-                    f"Sorting data by timestamp descending", metadata)
-                data = data.sort_values(by='timestamp', ascending=False)
-            # If policy_type is erl, we need to sort the data by timestamp ascending and take the first policy_value rows
-            elif policy_type == 'erl':
-                self.debug(
-                    f"Sorting data by timestamp ascending", metadata)
-                data = data.sort_values(by='timestamp', ascending=True)
-            else:
-                self.error(
-                    f"Invalid policy type: {policy_type}, using default policy", metadata)
-                raise ValueError(
-                    f"Invalid policy type: {policy_type}")
+        int_policy_value = int(policy_value)
 
-        data = data.head(int(policy_value))
+        data = data.head(int_policy_value)
+        if int_policy_value == 1:
+            data['timestamp'] = last_timestamp
 
         data['model_id'] = input_data['model_id']
         data['prediction_confidence'] = input_data['prediction_confidence']
         data['prediction_status'] = 'Good'
-        data['comments'] = ""
+        data['comments'] = ''
         data = data.sort_values(by='timestamp', ascending=False)
         data = data.reset_index(drop=True)
 
-        self.info(f"Prediction formatted: {data.size} rows", metadata)
-        self.debug(f"Prediction data: {data.to_string()}", metadata)
+        self.info(f'Prediction formatted: {len(data)} rows', metadata)
+        self._debug_dataframe('Prediction data:', data, metadata)
 
         return data.to_dict()
 
-    @activity.defn(name="format_default_prediction")
-    async def format_default_prediction(self, input_data: dict[str, Any]) -> dict[Any, Any]:
+    @activity.defn(name='format_default_prediction')
+    async def format_default_prediction(self, input_data: dict[str, Any]) -> dict:
         """
         Create and format default prediction data for error conditions.
 
@@ -476,64 +634,94 @@ class Gates(BaseActivity):
         """
 
         metadata = input_data['metadata']
-        self.debug("Formatting default prediction...", metadata)
+        self.debug('Formatting default prediction...', metadata)
 
-        data = DataFrame({
-            'prediction': [0],
-            'response_time': [0],
-            'timestamp': [input_data['timestamp']],
-            'model_id': [input_data['model_id']],
-            'prediction_confidence': [input_data['prediction_confidence']],
-            'prediction_status': ['Bad'],
-            'comments': [input_data['comment']]
-        })
+        data = DataFrame(
+            {
+                'prediction': [0],
+                'response_time': [0],
+                'timestamp': [input_data['timestamp']],
+                'model_id': [input_data['model_id']],
+                'prediction_confidence': [input_data['prediction_confidence']],
+                'prediction_status': ['Bad'],
+                'comments': [input_data['comment']],
+            }
+        )
 
-        self.info(f"Default prediction formatted: {data.size} rows", metadata)
+        self.info(f'Default prediction formatted: {data.size} rows', metadata)
         return data.to_dict()
 
-    @activity.defn(name="get_last_timestamp")
-    async def get_last_timestamp(self, input_data: dict[str, Any]) -> str:
+    @activity.defn(name='format_retrain_report')
+    async def format_retrain_report(self, input_data: dict[str, Any]) -> dict:
         """
-        Extract the most recent timestamp from prediction data.
+        Format retrain report data for storage and audit trail maintenance.
 
-        This method analyzes prediction data to find the latest timestamp,
-        enabling incremental processing and data continuity tracking.
-        It handles empty datasets gracefully by returning the current time
-        as a fallback timestamp.
+        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 method is essential for:
-        1. Incremental data processing workflows
-        2. Data continuity validation
-        3. Timestamp-based data loading optimization
-        4. Workflow execution tracking
+        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:
-                - data (dict[str, Any]): Prediction data to analyze
+                - 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:
-            str: Formatted timestamp string in UTC with timezone
+            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)
 
-        self.info("Getting last timestamp...", metadata)
+        experiment_response = input_data['experiment_response']
+        update_report = input_data['update_report']
+        model_id = input_data['model_id']
+        model_name = input_data['model_name']
 
-        data = DataFrame(input_data['data'])
+        report = DataFrame(
+            {
+                'model_id': [model_id],
+                'model_name': [model_name],
+                'timestamp': [experiment_response['timestamp']],
+                'status': [experiment_response['message']],
+            }
+        )
 
-        self.debug(f"Input data: {data.to_string()}", metadata)
+        if experiment_response['success']:
+            # Retrain was successfull
+            report['version'] = update_report['version']
+            report['mlflow_run_id'] = update_report['mlflow_run_id']
+            report['mlflow_experiment_id'] = update_report['mlflow_experiment_id']
 
-        if data.empty:
-            return now().strftime(DATETIME_FORMAT_WITH_TZ)
+        self._debug_dataframe('Retrain report:', report, metadata)
 
-        max_timestamp = max(
-            data['timestamp'].values.tolist())
+        return report.to_dict()
 
-        self.info(
-            f"Last timestamp: {max_timestamp}", metadata)
-
-        return max_timestamp
-
-    @activity.defn(name="write_metrics")
+    @activity.defn(name='write_metrics')
     async def write_metrics(self, input_data: dict[str, Any]):
         """
         Write prediction performance metrics to Prometheus monitoring system.
@@ -560,27 +748,57 @@ class Gates(BaseActivity):
         prediction = DataFrame(input_data['prediction'])
         prediction_confidence = prediction['prediction_confidence'].values[0]
         response_time = prediction['response_time'].values[0]
+        opc_metrics = input_data['opc_metrics']
 
-        self.info(
-            f"Writing metrics for model {metadata['model_name']}", metadata)
+        self.info(f'Writing metrics for model {metadata["model_name"]}', metadata)
 
-        metrics.PREDICTIONS_WRITTEN_COUNT.labels(
-            pod_id=self.pod_id,
-            model_name=metadata['model_name'],
-            pipeline_name=metadata['workflow_name']
-        ).inc()
+        core_tags = {
+            'pod_id': self.pod_id,
+            'runtime': self.runtime,
+            'operation_type': 'predict',
+            'model_name': metadata['model_name'],
+            'workflow_name': metadata['workflow_name'],
+        }
+        await self.emit_metric(
+            metric_object=metrics.PREDICTIONS_WRITTEN_COUNT,
+            tags=core_tags,
+        )
 
-        metrics.PREDICTION_CONFIDENCE_MONITOR.labels(
-            pod_id=self.pod_id,
-            model_name=metadata['model_name'],
-            pipeline_name=metadata['workflow_name']
-        ).set(prediction_confidence)
+        await self.emit_metric(
+            metric_object=metrics.PREDICTION_CONFIDENCE_MONITOR,
+            method='set',
+            tags=core_tags,
+            value=prediction_confidence,
+        )
 
-        metrics.PREDICTION_RESPONSE_TIME_MONITOR.labels(
-            pod_id=self.pod_id,
-            model_name=metadata['model_name'],
-            pipeline_name=metadata['workflow_name']
-        ).observe(response_time)
+        await self.emit_metric(
+            metric_object=metrics.PREDICTION_RESPONSE_TIME_MONITOR,
+            method='observe',
+            tags=core_tags,
+            value=response_time,
+        )
 
-        self.info(
-            f"Metrics written for model {metadata['model_name']}", metadata)
+        for server_id, tags in opc_metrics.items():
+            for tag, response_time in tags.items():
+                if response_time is not None:
+                    await self.emit_metric(
+                        metric_object=metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR,
+                        method='observe',
+                        tags={
+                            **core_tags,
+                            'opc_server_id': server_id,
+                            'tag': tag,
+                        },
+                        value=response_time,
+                    )
+
+                    await self.emit_metric(
+                        metric_object=metrics.PREDICTION_OPC_WRITING_COUNT,
+                        tags={
+                            **core_tags,
+                            'opc_server_id': server_id,
+                            'tag': tag,
+                        },
+                    )
+
+        self.info(f'Metrics written for model {metadata["model_name"]}', metadata)

laborious/activities/mlflow.py

@@ -1,20 +1,31 @@
-import json
 from temporalio import activity, workflow
 
-
 with workflow.unsafe.imports_passed_through():
-    from sientia_do.temporal.activities.base import BaseActivity
+    import traceback
+    from typing import Any
+
+    import numpy as np
+    from pandas import to_datetime
     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.repository.minio_repository import MinioRepository
+    from sientia_do.temporal.constants import (
+        DATETIME_FORMAT,
+        DATETIME_FORMAT_MS_WITH_TZ,
+        DATETIME_FORMAT_WITH_TZ,
+        now,
+    )
+    from sientia_do.utils.formatters import create_sample_dict
+
+    from laborious.utils.dataframe_debug import build_dataframe_debug_message
+    from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload
+    from laborious.utils.repository.minio_manager import MinioManager
     from laborious.utils.repository.model_repository import MLFlowRepository
-    from typing import Any
-    import numpy as np
-    from pandas import DataFrame
-    import traceback
 
 
-class MLFlow(BaseActivity):
+class MLFlow(MinioManager):
     """
     MLFlow integration activities for model inference operations.
 
@@ -33,8 +44,19 @@ class MLFlow(BaseActivity):
         model_monitoring_repository (MLFlowRepository): Repository for MLFlow operations
     """
 
-    def __init__(self, mlflow_host: str, mlflow_port: int, mlflow_username: str,
-                 mlflow_password: str, logger: Logger, notification_handler: NotificationHandler):
+    _MAX_DEBUG_DATAFRAME_ROWS = 100
+
+    def __init__(
+        self,
+        mlflow_host: str,
+        mlflow_port: int,
+        mlflow_username: str,
+        mlflow_password: str,
+        minio_repository: MinioRepository | None = None,
+        logger: Logger | None = None,
+        notification_handler: NotificationHandler | None = None,
+        metrics_controller: MetricsController | None = None,
+    ):
         """
         Initialize MLFlow activities with server configuration.
 
@@ -49,19 +71,52 @@ class MLFlow(BaseActivity):
         Raises:
             Exception: If MLFlowRepository initialization fails
         """
-        BaseActivity.__init__(
-            self, logger, notification_handler, set_error_counter=True)
+        MinioManager.__init__(
+            self, minio_repository, logger, notification_handler, metrics_controller
+        )
         self.mlflow_host = mlflow_host
         self.mlflow_port = mlflow_port
         self.mlflow_username = mlflow_username
         self.mlflow_password = mlflow_password
 
         self.model_monitoring_repository = MLFlowRepository(
-            f"{mlflow_host}:{mlflow_port}", mlflow_username, mlflow_password, logger.base_logger
+            f'{mlflow_host}:{mlflow_port}',
+            mlflow_username,
+            mlflow_password,
+            logger,
+            notification_handler,
+            metrics_controller,
         )
 
-    @activity.defn(name="request_transform")
-    async def request_transform(self, input_data: dict[str, Any]) -> dict[str, Any]:
+    def close(self) -> None:
+        """
+        Close the MLFlow activity and clean up resources.
+        """
+        MinioManager.close(self)
+
+    def __del__(self):
+        self.close()
+
+    def _debug_dataframe(self, message: str, data: Any, metadata: dict[str, Any]) -> None:
+        """
+        Log dataframe content only when row count is below the configured threshold
+
+        Args:
+         - message (str): Base log message to identify the dataframe in logs
+         - data (Any): Dataframe-like object expected to expose shape and to_csv
+         - metadata (dict[str, Any]): Workflow metadata for contextual logging
+        """
+        self.debug(
+            build_dataframe_debug_message(
+                message=message,
+                data=data,
+                max_rows=self._MAX_DEBUG_DATAFRAME_ROWS,
+            ),
+            metadata,
+        )
+
+    @activity.defn(name='request_transform')
+    async def request_transform(self, input_data: dict[str, Any]) -> MinioDataFramePayload:
         """
         Transform input data using MLFlow models.
 
@@ -92,12 +147,14 @@ class MLFlow(BaseActivity):
         """
         metadata = input_data['metadata']
         self.info('Transforming data...', metadata)
-        data = DataFrame(input_data['data'])
-        model_name = input_data['model_name']
-        model_retention = input_data['model_retention']
 
-        self.debug("Raw input data:", metadata)
-        self.debug(data, metadata)
+        payload = MinioDataFramePayload.from_dict(input_data['data'])
+        data = await payload.retrieve(self.minio_repository, metadata)
+
+        model_name = input_data['model_name']
+        model_config = input_data.get('model_config', {})
+
+        self._debug_dataframe('Raw input data:', data, metadata)
 
         # Sort by created_at in descending order and keep first occurrence of each variable/timestamp pair
         data = data.sort_values('created_at', ascending=False).drop_duplicates(
@@ -105,29 +162,60 @@ class MLFlow(BaseActivity):
         )
 
         # Pivot data for model input format
-        data = data.pivot(
-            index='timestamp', columns='variable',
-            values='value')
+        data = data.pivot(index='timestamp', columns='variable', values='value')
         data.fillna(np.nan, inplace=True)
-        # data.reset_index(inplace=True)
-        data.columns.name = None
 
-        self.debug("Processed input data:", metadata)
-        self.debug(data, metadata)
+        data.columns.name = None
+        data.index.name = None
+
+        data['timestamp'] = data.index
+
+        self._debug_dataframe('Processed input data:', data, metadata)
 
         # Request transformation from MLFlow model
-        response_data = self.model_monitoring_repository.transform(
-            model_name, data, model_retention)
+        response_data = await self.model_monitoring_repository.transform(
+            model_name, data, model_config, metadata
+        )
 
-        self.debug("Transform response data:", metadata)
-        self.debug(json.dumps(response_data, indent=4), metadata)
+        self.debug(
+            f'Transform raw response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
+            metadata,
+        )
 
-        self.info("Data transformed successfully", metadata)
+        self.debug(
+            f'Transform response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
+            metadata,
+        )
 
-        return response_data
+        self.info('Data transformed successfully', metadata)
 
-    @activity.defn(name="request_predict")
-    async def request_predict(self, input_data: dict[str, Any]) -> dict[str, Any]:
+        if not response_data.get('success', False):
+            return await MinioDataFramePayload.from_dataframe(
+                dataframe=None,
+                minio_repo=self.minio_repository,
+                model_name=model_name,
+                operation='transform',
+                status=response_data,
+                workflow_metadata=metadata,
+                last_timestamp=payload.last_timestamp,
+                logger=self.logger,
+            )
+
+        return await MinioDataFramePayload.from_dataframe(
+            dataframe=response_data['content'],
+            minio_repo=self.minio_repository,
+            model_name=model_name,
+            operation='transform',
+            workflow_metadata=metadata,
+            status={
+                'success': True,
+            },
+            last_timestamp=payload.last_timestamp,
+            logger=self.logger,
+        )
+
+    @activity.defn(name='request_predict')
+    async def request_predict(self, input_data: dict[str, Any]) -> MinioDataFramePayload:
         """
         Execute predictions using MLFlow models.
 
@@ -158,27 +246,61 @@ class MLFlow(BaseActivity):
         """
         metadata = input_data['metadata']
         self.info('Predicting data...', metadata)
-        data = DataFrame(input_data['data'])
-        model_name = input_data['model_name']
-        model_retention = input_data['model_retention']
 
-        self.debug(data, metadata)
+        payload = MinioDataFramePayload.from_dict(input_data['data'])
+        data = await payload.retrieve(self.minio_repository, metadata)
+
+        model_name = input_data['model_name']
+        model_config = input_data.get('model_config', {})
+
+        self._debug_dataframe('Input data for prediction:', data, metadata)
 
         # Convert numpy.nan to None for model compatibility
         data.replace(np.nan, None, inplace=True)
 
+        data['timestamp'] = data.index
+        data['timestamp'] = to_datetime(
+            data['timestamp'], format=DATETIME_FORMAT_WITH_TZ
+        ).dt.strftime(DATETIME_FORMAT)
+
         # Request prediction from MLFlow model
-        response_data = self.model_monitoring_repository.predict(
-            model_name, data, model_retention)
+        response_data = await self.model_monitoring_repository.predict(
+            model_name, data, model_config, metadata
+        )
 
-        self.debug("Prediction response data:", metadata)
-        self.debug(json.dumps(response_data, indent=4), metadata)
+        self.debug(
+            f'Prediction response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
+            metadata,
+        )
 
-        self.info("Data predicted successfully", metadata)
+        self.info('Data predicted successfully', metadata)
 
-        return response_data
+        if not response_data.get('success', False):
+            return await MinioDataFramePayload.from_dataframe(
+                dataframe=None,
+                minio_repo=self.minio_repository,
+                model_name=model_name,
+                operation='predict',
+                status=response_data,
+                workflow_metadata=metadata,
+                last_timestamp=payload.last_timestamp,
+                logger=self.logger,
+            )
 
-    @activity.defn(name="retrain_model")
+        return await MinioDataFramePayload.from_dataframe(
+            dataframe=response_data['content'],
+            minio_repo=self.minio_repository,
+            model_name=model_name,
+            operation='predict',
+            workflow_metadata=metadata,
+            status={
+                'success': True,
+            },
+            last_timestamp=payload.last_timestamp,
+            logger=self.logger,
+        )
+
+    @activity.defn(name='retrain_model')
     async def retrain_model(self, input_data: dict[str, Any]) -> dict[str, Any]:
         """
         Retrain MLFlow models with updated training data.
@@ -210,51 +332,89 @@ class MLFlow(BaseActivity):
         Raises:
             Exception: If retraining fails or encounters critical errors
         """
+
+        if self.minio_repository is None:
+            raise ValueError('Minio repository not initialized')
+
         metadata = input_data['metadata']
-        data = DataFrame(input_data['data'])
+
+        try:
+            # Payload-based retrain input (inline dict or MinIO offloaded).
+            payload = MinioDataFramePayload.from_dict(input_data['data'])
+            data = await payload.retrieve(self.minio_repository, metadata)
+
+        except Exception as e:
+            trace = traceback.format_exc()
+            await self.send_notification_async(
+                metadata=metadata,
+                notification_id='ERROR_LOADING_RETRAIN_DATA',
+                message=f'Error loading retrain data: {e}',
+                block='retrain_model',
+                level=NotificationLevel.ERROR,
+                attachment_content=trace,
+            )
+            self.error(trace, metadata)
+            return {
+                'success': False,
+                'message': f'Error loading retrain data: {e}',
+                'traceback': trace,
+                'timestamp': now().strftime(DATETIME_FORMAT_MS_WITH_TZ),
+            }
+
+        self.debug(f'Retrain data loaded successfully: shape {data.shape}', metadata)
+
         model_name = input_data['model_name']
+        model_config = input_data.get('model_config', {})
 
         self.info(f'Retraining model {model_name}...', metadata)
 
         timestamp = data['timestamp'].max()
         self.debug(f'Timestamp: {timestamp}', metadata)
 
+        # Sort by created_at in descending order and keep first occurrence of each variable/timestamp pair
+        if 'created_at' in data.columns:
+            data = data.sort_values('created_at', ascending=False).drop_duplicates(
+                subset=['variable', 'timestamp'], keep='first'
+            )
+        else:
+            data = data.drop_duplicates(subset=['variable', 'timestamp'], keep='first')
+
         data.drop(columns=['model_id'], inplace=True, errors='ignore')
         data.drop(columns=['created_at'], inplace=True, errors='ignore')
 
-        data = data.pivot(index='timestamp', columns='variable',
-                          values='value')
-        data.sort_index(inplace=True)
-        data.reset_index(inplace=True)
-
-        data = data.dropna()
+        # Pivot data for model input format
+        data = data.pivot(index='timestamp', columns='variable', values='value')
+        data.fillna(np.nan, inplace=True)
+        # data.reset_index(inplace=True)
         data.columns.name = None
 
-        try:
-            retrain_output, experiment = self.model_monitoring_repository.retrain_model(
-                data=data,
-                model_name=model_name
-            )
+        data['timestamp'] = data.index
+        data['timestamp'] = to_datetime(
+            data['timestamp'], format=DATETIME_FORMAT_WITH_TZ
+        ).dt.strftime(DATETIME_FORMAT)
+        data['timestamp'] = to_datetime(data['timestamp'], format=DATETIME_FORMAT)
 
-            return {
-                'status': retrain_output,
-                'timestamp': timestamp,
-                'experiment': experiment
-            }
-        except Exception as e:
-            trace = traceback.format_exc()
-            self.send_notification(
+        data.columns.name = None
+
+        retrain_output = await self.model_monitoring_repository.retrain_model(
+            data=data, model_name=model_name, model_config=model_config, metadata=metadata
+        )
+
+        if not retrain_output['success']:
+            trace = retrain_output['traceback']
+            await self.send_notification_async(
                 metadata=metadata,
                 notification_id='RETRAIN_MODEL_ERROR',
-                message=f'Error retraining model {model_name}: {e}',
+                message=f'Error retraining model {model_name}: {retrain_output["message"]}',
                 block='retrain_model',
                 level=NotificationLevel.ERROR,
-                attachment_content=trace
+                attachment_content=trace,
             )
             self.error(trace, metadata=metadata)
-            raise e
 
-    @activity.defn(name="update_production_model")
+        return {**retrain_output, 'timestamp': timestamp}
+
+    @activity.defn(name='update_production_model')
     async def update_production_model(self, input_data: dict[str, Any]) -> dict[Any, Any]:
         """
         Update production model with newly trained model version.
@@ -293,39 +453,76 @@ class MLFlow(BaseActivity):
         """
         metadata = input_data['metadata']
         model_name = input_data['model_name']
-        model_id = input_data['model_id']
         experiment = input_data['experiment']
-        timestamp = input_data['timestamp']
-        status = input_data['status']
-
         self.info(
-            f'Updating production model {model_name} from experiment {experiment}...', metadata)
+            f'Updating production model {model_name} from experiment {experiment}...', metadata
+        )
 
         try:
-            response = self.model_monitoring_repository.update_production_model(
-                experiment=experiment,
-                model_name=model_name
+            response = await self.model_monitoring_repository.update_production_model(
+                experiment=experiment, model_name=model_name, metadata=metadata
             )
 
-            report = DataFrame([response])
-            report['model_id'] = model_id
-            report['model_name'] = model_name
-            report['timestamp'] = timestamp
-            report['status'] = status
-
-            self.info(
-                f'Production model {model_name} updated successfully', metadata)
-            return report.to_dict()
+            self.info(f'Production model {model_name} updated successfully', metadata)
+            return response
 
         except Exception as e:
             trace = traceback.format_exc()
-            self.send_notification(
+            await self.send_notification_async(
                 metadata=metadata,
                 notification_id='UPDATE_PRODUCTION_MODEL_ERROR',
                 message=f'Error updating production model {model_name}: {e}',
                 block='update_production_model',
                 level=NotificationLevel.ERROR,
-                attachment_content=trace
+                attachment_content=trace,
             )
             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,364 @@
+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
+    from laborious.utils.dataframe_debug import build_dataframe_debug_message
+
+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.
+    """
+
+    _MAX_DEBUG_DATAFRAME_ROWS = 100
+
+    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()
+
+    def _debug_dataframe(self, message: str, data: Any, metadata: dict[str, Any]) -> None:
+        """
+        Log dataframe content only when row count is below the configured threshold
+
+        Args:
+         - message (str): Base log message to identify the dataframe in logs
+         - data (Any): Dataframe-like payload to be logged
+         - metadata (dict[str, Any]): Workflow metadata for contextual logging
+        """
+        self.debug(
+            build_dataframe_debug_message(
+                message=message,
+                data=data,
+                max_rows=self._MAX_DEBUG_DATAFRAME_ROWS,
+            ),
+            metadata,
+        )
+
+    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_dataframe(
+            f'Reference data: Size {reference_data.shape}', reference_data, metadata
+        )
+
+        self._debug_dataframe(f'Target data: Size {target_data.shape}', target_data, 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_dataframe(f'Drift dataframe: Size {drift_df.shape}', drift_df, 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_dataframe(f'Drift dataframe: Size {drift_df.shape}', drift_df, 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_dataframe(f'Simple metrics dataframe: Size {data.shape}', data, metadata)
+
+        return data.to_dict(orient='records')

laborious/activities/opc.py

@@ -1,20 +1,62 @@
 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
     from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
     from sientia_do.notifications.models import NotificationLevel
-    from sientia_do.temporal.activities.base import BaseActivity
     from sientia_do.observability.logger import Logger
+    from sientia_do.observability.metrics_controller import MetricsController
+    from sientia_do.observability.sientia_monitoring import SientiaMonitoring
+
     from laborious.utils.repository.opc_repository import OpcRepository
-    from typing import Any
-    import traceback
-    from pandas import DataFrame
 
 OPC_WRITTING_ERROR_CONFIDENCE = 12
+OPC_SESSION_BAD_CONFIDENCE = 14
+OPC_SESSION_BAD_COMMENT_PREFIX = 'OPC UA session/channel error:'
+OPC_WRITTING_ERROR_MESSAGE = 'Some data could not be written to OPC servers'
+OPC_RECONNECT_IN_PROGRESS_COMMENT = 'OPC UA reconnect in progress'
+OPC_COMMENT_SEPARATOR = ' | '
 
 
-class OPC(BaseActivity):
+def _opc_session_bad_comment(opc_status: str | None) -> str:
+    status = opc_status or 'Unknown'
+    return f'{OPC_SESSION_BAD_COMMENT_PREFIX} {status}'
+
+
+def _apply_opc_write_error(
+    error_info: dict[str, Any] | None,
+    session_bad_seen: bool,
+    session_bad_status: str | None,
+    reconnect_in_progress_seen: bool,
+) -> tuple[bool, str | None, bool]:
+    """
+    Update session/reconnect flags from an OPC write error payload.
+
+    Args:
+        error_info: Repository error details, or None when the write succeeded.
+        session_bad_seen: Whether a session_bad error was seen so far.
+        session_bad_status: Last known OPC status for session errors.
+        reconnect_in_progress_seen: Whether reconnect_in_progress was seen so far.
+
+    Return:
+        Updated (session_bad_seen, session_bad_status, reconnect_in_progress_seen).
+    """
+    if not error_info:
+        return session_bad_seen, session_bad_status, reconnect_in_progress_seen
+
+    kind = error_info.get('opc_error_kind')
+    if kind == 'session_bad':
+        return True, error_info.get('opc_status', session_bad_status), reconnect_in_progress_seen
+    if kind == 'reconnect_in_progress':
+        return session_bad_seen, session_bad_status, True
+    return session_bad_seen, session_bad_status, reconnect_in_progress_seen
+
+
+class OPC(SientiaMonitoring):
     """
     OPC server integration activities for real-time data export.
 
@@ -33,18 +75,20 @@ class OPC(BaseActivity):
         notification_handler (NotificationHandler): Notification management instance
     """
 
-    def __init__(self, opc_servers: dict[str, dict[str, Any]],
-                 logger: Logger, notification_handler: NotificationHandler):
-
+    def __init__(
+        self,
+        opc_servers: dict[str, dict[str, Any]],
+        logger: Logger,
+        notification_handler: NotificationHandler,
+        metrics_controller: MetricsController,
+    ):
         self.logger = logger
         self.notification_handler = notification_handler
         self.opc_servers = opc_servers
 
-        BaseActivity.__init__(
-            self, logger, notification_handler, set_error_counter=True)
+        SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
 
         self.opc_repository: dict[str, OpcRepository] = {}
-        self.opc_servers = opc_servers
 
     async def init_opc(self):
         """
@@ -70,10 +114,11 @@ class OPC(BaseActivity):
             the initialization of other OPC servers. Each server is handled
             independently to ensure maximum availability.
         """
-        self.logger.info("Initializing OPC servers...")
-        for id, server in self.opc_servers.items():
-            self.opc_repository[id] = OpcRepository(
-                id=server['id'],
+        self.logger.info('Initializing OPC servers...')
+        for opc_id, server in self.opc_servers.items():
+            self.opc_repository[opc_id] = OpcRepository(
+                opc_id=server['id'],
+                server_name=server['server_name'],
                 url=server['url'],
                 logger=self.logger,
                 server_uri=server['server_uri'],
@@ -82,77 +127,73 @@ class OPC(BaseActivity):
                 server_cert_path=server['server_cert_path'],
                 notification_handler=self.notification_handler,
                 reconnection_interval=server['reconnection_interval'],
-                pod_id=self.pod_id
+                metrics_controller=self.metrics_controller,
             )
-            is_connected, error_data = await self.opc_repository[id].connect()
+            is_connected, error_data = await self.opc_repository[opc_id].connect()
             if not is_connected:
-                self.send_notification(
+                await self.send_notification_async(
                     metadata={
                         'model_id': '-',
                         'model_name': '-',
                         'workflow_name': '-',
-                        'schedule_name': 'INITIALIZATION'
+                        'schedule_name': 'INITIALIZATION',
                     },
                     notification_id=error_data['notification_id'],
                     message=error_data['message'],
                     block=error_data['block'],
                     level=error_data.get('level', NotificationLevel.ERROR),
-                    attachment_content=error_data.get(
-                        'attachment_content', None)
+                    attachment_content=error_data.get('attachment_content', None),
                 )
             else:
                 self.logger.info(
-                    f"OPC server {id} connected successfully.")
+                    f'OPC server {opc_id}:{server["server_name"]} connected successfully.'
+                )
 
-    async def write_data(self, server_id: str, tag: str, data: Any,
-                         data_type: str, tag_type: str, metadata: dict[str, Any]) -> bool:
+    async def write_data(
+        self,
+        server_id: str,
+        tag: str,
+        data: Any,
+        data_type: str,
+        tag_type: str,
+        metadata: dict[str, Any],
+    ) -> tuple[float | None, dict[str, Any] | None]:
         """
         Write data to a specific OPC server tag with comprehensive error handling.
 
-        This method provides a secure and reliable way to write data to OPC servers
-        with automatic error handling, notification integration, and detailed logging.
-        It validates server availability before attempting write operations and
-        provides comprehensive error reporting for operational monitoring.
-
-        Args:
-            - server_id (str): The id of the OPC server.
-            - tag (str): The tag to write to.
-            - data (Any): The data to write.
-            - data_type (str): The data type.
-            - tag_type (str): The tag type.
-
-        Returns:
-            - bool: True if the data was written successfully, False otherwise.
+        Return:
+            tuple[float | None, dict[str, Any] | None]: Response time on success, or
+                (None, error info_data) on repository failure.
         """
 
         try:
-            is_success, error_data = await self.opc_repository[server_id].write_data(
-                tag, data, data_type, self.logger, metadata)
+            is_success, info_data = await self.opc_repository[server_id].write_data(
+                tag, data, data_type, metadata
+            )
             if not is_success:
-                self.send_notification(
+                await self.send_notification_async(
                     metadata=metadata,
-                    notification_id=error_data['notification_id'],
-                    message=error_data['message'],
-                    block=error_data['block'],
-                    level=error_data.get('level', NotificationLevel.ERROR),
-                    attachment_content=error_data.get(
-                        'attachment_content', None)
+                    notification_id=info_data['notification_id'],
+                    message=info_data['message'],
+                    block=info_data['block'],
+                    level=info_data.get('level', NotificationLevel.ERROR),
+                    attachment_content=info_data.get('attachment_content', None),
                 )
-                return False
-            return True
+                return None, info_data
+            return info_data['response_time'], None
         except Exception as e:
             trace = traceback.format_exc()
-            self.send_notification(
+            await self.send_notification_async(
                 metadata=metadata,
-                notification_id=f"WRITE_OPC_{tag_type.upper()}_ERROR",
-                message=f"Error writing data to OPC server: {e}",
-                block="write_opc_data",
+                notification_id=f'WRITE_OPC_{tag_type.upper()}_ERROR',
+                message=f'Error writing data to OPC server: {e}',
+                block='write_opc_data',
                 level=NotificationLevel.ERROR,
-                attachment_content=trace
+                attachment_content=trace,
             )
             raise e
 
-    def validate_server(self, server_id: str, metadata: dict[str, Any]) -> bool:
+    async def validate_server(self, server_id: str, metadata: dict[str, Any]) -> bool:
         """
         Validate that an OPC server is available and configured for write operations.
 
@@ -174,21 +215,81 @@ class OPC(BaseActivity):
             This helps operators quickly identify configuration issues.
         """
         if self.opc_repository.get(server_id) is None:
-            message = f"OPC server {server_id} not found to perform write operation."
-            self.send_notification(
+            message = f'OPC server {server_id} not found to perform write operation.'
+            await self.send_notification_async(
                 metadata=metadata,
-                notification_id="OPC_SERVER_NOT_FOUND",
+                notification_id='OPC_SERVER_NOT_FOUND',
                 message=message,
-                block="write_opc_data",
+                block='write_opc_data',
                 level=NotificationLevel.ERROR,
-                attachment_content=f"OPC servers: {list(self.opc_repository.keys())}"
+                attachment_content=f'OPC servers: {list(self.opc_repository.keys())}',
             )
             return False
         return True
 
+    async def _write_tags_from_config(
+        self,
+        server_id: str,
+        tags_config: dict[str, dict[str, Any]],
+        data: DataFrame,
+        data_column: str,
+        tag_type: str,
+        log_label: str,
+        metadata: dict[str, Any],
+    ) -> tuple[dict[str, float | None], bool, str | None, bool]:
+        """
+        Write a group of OPC tags and collect response times and error flags.
+
+        Args:
+            server_id: Target OPC server identifier.
+            tags_config: Tag name to configuration mapping.
+            data: DataFrame with prediction/confidence columns.
+            data_column: Column name whose first row value is written.
+            tag_type: Tag category passed to write_data ('prediction' or 'confidence').
+            log_label: Human-readable label for success logs.
+            metadata: Context metadata for logging and notifications.
+
+        Return:
+            (response_times, session_bad_seen, session_bad_status, reconnect_in_progress_seen)
+        """
+        response_times: dict[str, float | None] = {}
+        session_bad_seen = False
+        session_bad_status: str | None = None
+        reconnect_in_progress_seen = False
+
+        for tag, tag_config in tags_config.items():
+            response_time, error_info = await self.write_data(
+                server_id=server_id,
+                tag=tag,
+                data=data.head(1)[data_column].values[0],
+                data_type=tag_config['data_type'],
+                tag_type=tag_type,
+                metadata=metadata,
+            )
+            session_bad_seen, session_bad_status, reconnect_in_progress_seen = (
+                _apply_opc_write_error(
+                    error_info,
+                    session_bad_seen,
+                    session_bad_status,
+                    reconnect_in_progress_seen,
+                )
+            )
+            if response_time is not None:
+                self.info(
+                    f'{log_label} written to OPC server {server_id} for tag {tag}.',
+                    metadata,
+                )
+            response_times[tag] = response_time
+
+        return response_times, session_bad_seen, session_bad_status, reconnect_in_progress_seen
+
     async def manage_output_tags(
-            self, server_id: str, config: dict[str, Any], data: DataFrame,
-            metadata: dict[str, Any], success: bool) -> tuple[bool, int]:
+        self,
+        server_id: str,
+        config: dict[str, Any],
+        data: DataFrame,
+        metadata: dict[str, Any],
+    ) -> tuple[bool, dict[str, float | None], bool, str | None, bool]:
         """
         Manage the writing of prediction and confidence data to OPC server tags.
 
@@ -215,44 +316,52 @@ class OPC(BaseActivity):
                 - overall_success: True if all configured tags were written successfully
                 - total_tags_written: Count of successfully written tags
         """
+        response_times: dict[str, float | None] = {}
+        session_bad_seen = False
+        session_bad_status: str | None = None
+        reconnect_in_progress_seen = False
 
-        count = 0
-        if 'prediction_tags' in config:
-            for tag, tag_config in config['prediction_tags'].items():
-                local_success = await self.write_data(
-                    server_id=server_id,
-                    tag=tag,
-                    data=data.head(1)['prediction'].values[0],
-                    data_type=tag_config['data_type'],
-                    tag_type='prediction',
-                    metadata=metadata
-                )
-                if local_success:
-                    self.info(
-                        f"Prediction data written to OPC server {server_id} for tag {tag}.", metadata)
-                    count += 1
-                success = success and local_success
+        tag_groups = (
+            ('prediction_tags', 'prediction', 'prediction', 'Prediction data'),
+            ('confidence_tags', 'prediction_confidence', 'confidence', 'Confidence data'),
+        )
+        for config_key, data_column, tag_type, log_label in tag_groups:
+            if config_key not in config:
+                continue
+            (
+                group_times,
+                group_session_bad,
+                group_status,
+                group_reconnect,
+            ) = await self._write_tags_from_config(
+                server_id=server_id,
+                tags_config=config[config_key],
+                data=data,
+                data_column=data_column,
+                tag_type=tag_type,
+                log_label=log_label,
+                metadata=metadata,
+            )
+            response_times.update(group_times)
+            if group_session_bad:
+                session_bad_seen = True
+                session_bad_status = group_status or session_bad_status
+            if group_reconnect:
+                reconnect_in_progress_seen = True
 
-        if 'confidence_tags' in config:
-            for tag, tag_config in config['confidence_tags'].items():
-                local_success = await self.write_data(
-                    server_id=server_id,
-                    tag=tag,
-                    data=data.head(1)['prediction_confidence'].values[0],
-                    data_type=tag_config['data_type'],
-                    tag_type='confidence',
-                    metadata=metadata
-                )
-                if local_success:
-                    self.info(
-                        f"Confidence data written to OPC server {server_id} for tag {tag}.", metadata)
-                    count += 1
-                success = success and local_success
-
-        return success, count
+        success = None not in response_times.values()
+        return (
+            success,
+            response_times,
+            session_bad_seen,
+            session_bad_status,
+            reconnect_in_progress_seen,
+        )
 
     @activity.defn(name='write_opc_data')
-    async def write_opc_data(self, input_data: dict[str, Any]) -> dict[Any, Any]:
+    async def write_opc_data(
+        self, input_data: dict[str, Any]
+    ) -> 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.
@@ -271,29 +380,66 @@ class OPC(BaseActivity):
 
         """
         metadata = input_data['metadata']
-        self.info("Writing data to OPC servers...", metadata)
+        self.info('Writing data to OPC servers...', metadata)
         data = DataFrame(input_data['data'])
         opc_output_config = input_data['opc_output_config']
-        self.info(f"Data to write: {data.size} rows", metadata)
+        self.info(f'Data to write: {data.size} rows', metadata)
 
         success = True
+        session_bad_seen = False
+        session_bad_status: str | None = None
+        reconnect_in_progress_seen = False
+
+        metrics: dict[str, dict[str, float | None]] = {}
 
         for server_id, config in opc_output_config.items():
-
-            if not self.validate_server(server_id, metadata):
+            if not await self.validate_server(server_id, metadata):
                 success = False
                 continue
 
-            local_success, local_count = await self.manage_output_tags(
-                server_id, config, data, metadata, success)
+            (
+                local_success,
+                local_response_times,
+                local_session_bad,
+                local_status,
+                local_reconnect_in_progress,
+            ) = await self.manage_output_tags(server_id, config, data, metadata)
+            metrics[server_id] = local_response_times
+            local_count = len(local_response_times)
             success = success and local_success
+            if local_session_bad:
+                session_bad_seen = True
+                session_bad_status = local_status or session_bad_status
+            if local_reconnect_in_progress:
+                reconnect_in_progress_seen = True
 
             self.info(
-                f"Process completed for OPC server {server_id}: {local_count} of {len(config['prediction_tags'])} prediction tags and {len(config['confidence_tags'])} confidence tags", metadata)
+                f'Process completed for OPC server {server_id}: {local_count} of {len(config.get("prediction_tags", []))} prediction tags and {len(config.get("confidence_tags", []))} confidence tags',
+                metadata,
+            )
 
-        return self.process_confidence(data, success, metadata)
+        return (
+            self.process_confidence(
+                data,
+                success,
+                metadata,
+                session_bad=session_bad_seen,
+                opc_status=session_bad_status,
+                reconnect_in_progress=reconnect_in_progress_seen,
+            ),
+            metrics,
+        )
 
-    def process_confidence(self, data: DataFrame, success: bool, metadata: dict[str, Any]) -> dict[Any, Any]:
+    def process_confidence(
+        self,
+        data: DataFrame,
+        success: bool,
+        metadata: dict[str, Any],
+        *,
+        session_bad: bool = False,
+        opc_status: str | None = None,
+        reconnect_in_progress: bool = False,
+    ) -> dict[Hashable, Any]:
         """
         Process prediction confidence based on OPC write operation success.
 
@@ -321,18 +467,31 @@ class OPC(BaseActivity):
         """
 
         if not success:
-            data['prediction_confidence'] = OPC_WRITTING_ERROR_CONFIDENCE
-            self.debug(
-                f"Some data could not be written to OPC servers, setting confidence to  {OPC_WRITTING_ERROR_CONFIDENCE}.",
-                metadata
-            )
+            comment_parts: list[str] = []
+            confidence = OPC_WRITTING_ERROR_CONFIDENCE
 
+            if session_bad:
+                comment_parts.append(_opc_session_bad_comment(opc_status))
+                confidence = OPC_SESSION_BAD_CONFIDENCE
+            if reconnect_in_progress:
+                comment_parts.append(OPC_RECONNECT_IN_PROGRESS_COMMENT)
+                confidence = OPC_SESSION_BAD_CONFIDENCE
+            if not comment_parts:
+                comment_parts.append(OPC_WRITTING_ERROR_MESSAGE)
+
+            comments = OPC_COMMENT_SEPARATOR.join(comment_parts)
+            data['prediction_confidence'] = confidence
+            data['comments'] = comments
+            self.debug(
+                f'OPC write issues, confidence={confidence}, comments={comments}',
+                metadata,
+            )
         else:
-            self.debug("Data written to OPC servers successfully.", metadata)
+            self.debug('Data written to OPC servers successfully.', metadata)
 
         return data.to_dict()
 
-    async def shutdown(self):
+    async def aclose(self):
         """
         Gracefully shutdown all OPC server connections and cleanup resources.
 

laborious/activities/storage.py

@@ -0,0 +1,210 @@
+from temporalio import activity, workflow
+
+from laborious.utils.repository.minio_manager import MinioManager
+
+with workflow.unsafe.imports_passed_through():
+    # Extend the Temporal Postgres activities for convenient query -> MinIO export
+    import traceback
+    from datetime import timedelta
+    from typing import Any
+
+    import pandas as pd
+    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.repository.minio_repository import MinioRepository
+    from sientia_do.temporal.activities.postgres import Postgres
+    from sientia_do.temporal.constants import now
+
+    from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload
+
+_LOAD_QUERY_OFFLOAD_SKIP_KEYS = frozenset({'model_name', 'key_prefix', 'size_threshold_bytes'})
+
+
+class Storage(Postgres, MinioManager):
+    """
+    Extensions for Postgres activities with a helper to export query results
+    directly to MinIO as Parquet and return the object name.
+    """
+
+    minio_repository: MinioRepository | None = None
+
+    def __init__(
+        self,
+        host: str,
+        port: int,
+        user: str,
+        password: str,
+        dbname: str,
+        min_connections: int,
+        max_connections: int,
+        retention_hours: int = 24,
+        minio_repository: MinioRepository | None = None,
+        logger: Logger | None = None,
+        notification_handler: NotificationHandler | None = None,
+        metrics_controller: MetricsController | None = None,
+    ):
+        self.retention_hours = retention_hours
+        Postgres.__init__(
+            self,
+            host=host,
+            port=port,
+            user=user,
+            password=password,
+            dbname=dbname,
+            min_connections=min_connections,
+            max_connections=max_connections,
+            logger=logger,
+            notification_handler=notification_handler,
+            metrics_controller=metrics_controller,
+        )
+
+        MinioManager.__init__(
+            self, minio_repository, logger, notification_handler, metrics_controller
+        )
+
+    @activity.defn(name='load_query_with_minio_offload')
+    async def load_query_with_minio_offload(
+        self, input_data: dict[str, Any]
+    ) -> MinioDataFramePayload:
+        """
+        Run the custom SQL load, then return a MinIO-aware dataframe wire dict.
+
+        Args (input_data):
+            metadata (dict): Workflow metadata (same as load_custom_query).
+            query (str): SQL query.
+            datetime_columns (list[str], optional): Datetime column names.
+            model_name (str): Model name for object key basename.
+            key_prefix (str, optional): Directory prefix inside the bucket.
+            size_threshold_bytes (int, optional): Override env offload threshold.
+
+        Returns:
+            dict[str, Any]: Flat ``MinioDataFramePayload`` dict or ``success: False`` on failure.
+        """
+        if self.minio_repository is None:
+            raise ValueError('Minio repository not initialized')
+
+        metadata: dict = input_data.get('metadata', {})
+        model_name = input_data['model_name']
+
+        rows = await self.load_custom_query(
+            input_data,
+        )
+        if not rows:
+            self.error(
+                'load_query_with_minio_offload failed: No data returned from query', metadata
+            )
+            dataframe = None
+        else:
+            dataframe = pd.DataFrame(rows)
+
+        return await MinioDataFramePayload.from_dataframe(
+            dataframe,
+            minio_repo=self.minio_repository,
+            workflow_metadata=metadata,
+            model_name=model_name,
+            operation='initial',
+            logger=self.logger,
+        )
+
+    @activity.defn(name='export_payload_to_postgres')
+    async def export_payload_to_postgres(self, input_data: dict[str, Any]) -> dict:
+        """
+        Export a payload to PostgreSQL.
+        """
+        metadata = input_data.get('metadata')
+        payload = MinioDataFramePayload.from_dict(input_data['data'])
+        data = await payload.retrieve(self.minio_repository, metadata)
+
+        return await self.export_data_to_postgres(
+            {
+                **input_data,
+                'data': data,
+            }
+        )
+
+    @activity.defn(name='cleanup_minio_objects_expired')
+    async def cleanup_minio_objects_expired(self, input_data: dict[str, Any]) -> dict[str, Any]:
+        """
+        Delete objects under the given prefixes that are older than the retention window.
+
+        Args (input_data):
+            metadata (dict): Workflow metadata for logging and metrics.
+            prefixes (list[str]): Key prefixes to scan (one level or subtree per prefix).
+
+        Returns:
+            dict[str, Any]: ``success``, ``deleted_count``, and optional ``message``.
+        """
+        if self.minio_repository is None:
+            raise ValueError('Minio repository not initialized')
+
+        metadata = input_data.get('metadata', {})
+        payload = MinioDataFramePayload.from_dict(input_data['data'])
+        prefix = payload.cleanup_prefix()
+        base = now()
+        cutoff = (base.replace(tzinfo=None) if base.tzinfo else base) - timedelta(
+            hours=self.retention_hours
+        )
+
+        report: dict[str, Any] = {
+            'failed': {},
+            'deleted': {},
+            'failed_count': 0,
+            'deleted_count': 0,
+        }
+        try:
+            keys = await self.minio_repository.list_objects(
+                prefix=prefix,
+                recursive=True,
+                metadata=metadata,
+            )
+            for key in keys:
+                try:
+                    ts = MinioDataFramePayload.parse_object_timestamp(key)
+                    if ts is None:
+                        continue
+                    if ts >= cutoff:
+                        continue
+                    await self.minio_repository.delete_file(
+                        object_name=key,
+                        metadata=metadata,
+                    )
+                except Exception as e:
+                    report['failed'][key] = {
+                        'success': False,
+                        'message': str(e),
+                    }
+                    report['failed_count'] += 1
+                    continue
+                report['deleted'][key] = {
+                    'success': True,
+                    'message': 'Deleted',
+                }
+                report['deleted_count'] += 1
+        except Exception as e:
+            trace = traceback.format_exc()
+            await self.send_notification_async(
+                metadata=metadata,
+                notification_id='ERROR_CLEANUP_MINIO_OBJECTS_EXPIRED',
+                message=f'Error cleaning up MinIO objects: {e}',
+                block='cleanup_minio_objects_expired',
+                level=NotificationLevel.ERROR,
+                attachment_content=trace,
+            )
+            self.error(trace, metadata)
+        else:
+            # Cleanup success is expected in normal flow; avoid noisy INFO notifications
+            # that do not impact behavior and can flood observability in test runs.
+            self.info('MinIO objects cleaned up successfully', metadata)
+
+        return report
+
+    def close(self) -> None:
+        """Close Storage resources (MinIO client and Postgres engine)."""
+        if hasattr(self, 'engine'):
+            Postgres.close(self)
+        MinioManager.close(self)
+
+    def __del__(self):
+        self.close()

laborious/metrics.py

@@ -18,55 +18,183 @@ Key Metric Categories:
 
 Metric Labels:
 - pod_id: Kubernetes pod identifier for multi-instance deployments
+- runtime: Runtime / environment identifier (matches ``RUNTIME`` env, see ``SientiaMonitoring``)
 - model_name: Name of the ML model being used
-- pipeline_name: Name of the prediction pipeline
+- workflow_name: Name of the prediction pipeline
 - opc_server_id: Identifier for OPC server operations
 """
 
-from prometheus_client import Gauge, Counter, Histogram
+from prometheus_client import Counter, Gauge, Histogram
+from sientia_do.observability.metrics import CORE_LABELS
 
 # Application health metric
 APP_UP = Gauge(
-    "app_up",
-    "Indicates if the application is running (1) or shutting down (0)",
-    ["pod_id"],
+    'app_up',
+    'Indicates if the application is running (1) or shutting down (0)',
+    ['pod_id'],
 )
 
-# Core labels used across multiple metrics
-CORE_LABELS = ["pod_id", "model_name", "pipeline_name"]
-
 # Prediction operation metrics
 PREDICTIONS_WRITTEN_COUNT = Counter(
-    "laborious_predictions_written_count",
-    "Number of predictions written to the database table predictions",
+    'laborious_predictions_written_count',
+    'Number of predictions written to the database table predictions',
     CORE_LABELS,
 )
 
 # Prediction quality metrics
 PREDICTION_CONFIDENCE_MONITOR = Gauge(
-    "laborious_prediction_confidence_monitor",
-    "Current confidence of each prediction",
+    'laborious_prediction_confidence_monitor',
+    'Current confidence of each prediction',
     CORE_LABELS,
 )
 
-# Performance monitoring metrics
+# Prediction total response time
 PREDICTION_RESPONSE_TIME_MONITOR = Histogram(
-    "laborious_prediction_response_time_monitor",
-    "Current response time of each prediction",
+    'laborious_prediction_response_time_monitor',
+    'Current response time of each prediction',
     CORE_LABELS,
-    buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0]
+    buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0],
 )
 
-# OPC export metrics
+
+# ================== OPC metrics ==================
+
+
 PREDICTION_OPC_WRITING_COUNT = Counter(
-    "laborious_prediction_opc_writing_count",
-    "Number of predictions written to the OPC server",
-    [*CORE_LABELS, "opc_server_id"],
+    'laborious_prediction_opc_writing_count',
+    'Number of predictions written to the OPC server',
+    [*CORE_LABELS, 'opc_server_id', 'tag'],
 )
 
 PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR = Histogram(
-    "laborious_prediction_opc_writing_response_time_monitor",
-    "Current response time of each prediction written to the OPC server",
-    [*CORE_LABELS, "opc_server_id"],
-    buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0]
+    'laborious_prediction_opc_writing_response_time_monitor',
+    'Current response time of each prediction written to the OPC server',
+    [*CORE_LABELS, 'opc_server_id', 'tag'],
+    buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0],
+)
+
+OPC_CONNECTIONS_TOTAL = Counter(
+    'opc_connections_initiated_total',
+    'Total connection attempts to OPC servers',
+    ['pod_id', 'server_name'],
+)
+OPC_CONNECTIONS_FAILED = Counter(
+    'opc_connections_failed_total',
+    'Total failed connection attempts to OPC servers',
+    ['pod_id', 'server_name'],
+)
+OPC_CONNECTION_STATUS = Gauge(
+    'opc_connection_status',
+    'Connection status with the OPC server (1=connected, 0=disconnected)',
+    ['pod_id', 'server_name', 'server_url'],
+)
+
+_OPC_SESSION_DEBUG_LABELS = ['pod_id', 'server_name', 'runtime', 'opc_server_id', 'session_id']
+
+OPC_SESSION_CREATED_TOTAL = Counter(
+    'opc_session_created_total',
+    'OPC UA sessions established (after successful connect)',
+    _OPC_SESSION_DEBUG_LABELS,
+)
+
+OPC_SESSION_CLOSED_TOTAL = Counter(
+    'opc_session_closed_total',
+    'OPC UA client disconnects completed (session tear-down initiated)',
+    _OPC_SESSION_DEBUG_LABELS,
+)
+
+OPC_SESSION_REVISED_TIMEOUT_MS = Gauge(
+    'opc_session_revised_timeout_milliseconds',
+    'Server-revised OPC UA session timeout (RevisedSessionTimeout) in ms after connect',
+    _OPC_SESSION_DEBUG_LABELS,
+)
+
+OPC_WRITE_ATTEMPT_LABELS = [*_OPC_SESSION_DEBUG_LABELS, 'model_id', 'model_name', 'result']
+
+OPC_WRITE_ATTEMPTS_TOTAL = Counter(
+    'opc_write_attempts_total',
+    'OPC UA write attempts with session and outcome (result=OK or exception class name)',
+    OPC_WRITE_ATTEMPT_LABELS,
+)
+
+OPC_WRITE_INTER_ARRIVAL_OVER_SESSION_TIMEOUT_TOTAL = Counter(
+    'opc_write_inter_arrival_over_session_timeout_total',
+    'Successful writes where seconds since the previous successful write exceeded RevisedSessionTimeout (ms)',
+    _OPC_SESSION_DEBUG_LABELS,
+)
+
+# ================== Model metrics ==================
+
+MODEL_READ_LAG = Histogram(
+    'laborious_model_read_lag',
+    'Lag between the start and read of read operations',
+    CORE_LABELS,
+    buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0],
+)
+
+MODEL_WRITE_LAG = Histogram(
+    'laborious_model_write_lag',
+    'Lag between the start and end of write operations',
+    CORE_LABELS,
+    buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0],
+)
+
+MODEL_READ_COUNT = Counter(
+    'laborious_model_read_count',
+    'Number of reads from the model',
+    CORE_LABELS,
+)
+
+MODEL_WRITE_COUNT = Counter(
+    'laborious_model_write_count',
+    'Number of writes to the model',
+    CORE_LABELS,
+)
+
+MODEL_READ_ERROR_COUNT = Counter(
+    'laborious_model_read_error_count',
+    'Number of errors reading from the model',
+    CORE_LABELS,
+)
+
+MODEL_WRITE_ERROR_COUNT = Counter(
+    'laborious_model_write_error_count',
+    'Number of errors writing to the model',
+    CORE_LABELS,
+)
+
+MODEL_ANALYZE_LAG = Histogram(
+    'laborious_model_analyze_lag',
+    'Lag between the start and end of analyze operations',
+    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',
+    CORE_LABELS,
+)
+
+MODEL_ANALYZE_ERROR_COUNT = Counter(
+    'laborious_model_analyze_error_count',
+    'Number of errors during analyze operations',
+    CORE_LABELS,
+)
+
+
+# ================== PI Web API metrics ==================
+
+PI_WEB_API_LABELS = [*CORE_LABELS, 'tag_name']
+
+PI_WEB_API_PREDICTION_WRITTEN_COUNT = Counter(
+    'laborious_pi_web_api_prediction_written_count',
+    'Number of predictions written to the PI Web API',
+    PI_WEB_API_LABELS,
+)
+
+PI_WEB_API_PREDICTION_WRITTEN_ERROR_COUNT = Counter(
+    'laborious_pi_web_api_prediction_written_error_count',
+    'Number of errors writing predictions to the PI Web API',
+    PI_WEB_API_LABELS,
 )

laborious/utils/connectors_config.py

@@ -1,238 +1,92 @@
-"""
-Connectors Configuration Module
-
-This module provides configuration management for all external service connectors
-used by the Sientia DataOps Laborious system. It centralizes configuration
-for databases, MLFlow servers, OPC servers, and other external dependencies.
-
-The module implements configuration builders for:
-1. PostgreSQL database connections
-2. MLFlow model serving endpoints
-3. OPC server configurations
-4. MongoDB notification systems
-
-Key Features:
-- Environment variable-based configuration
-- Default value management for development
-- Connection pool configuration
-- Security credential management
-- Configuration validation and error handling
-- Support for multiple service instances
-
-Configuration Sources:
-- Environment variables for production deployment
-- Default values for local development
-- Kubernetes secrets integration
-- Configurable connection parameters
-
-Environment Variables:
-- POSTGRES_*: PostgreSQL connection parameters
-- MLFLOW_*: MLFlow server parameters
-- OPC_*: OPC server configuration
-- MONGODB_*: MongoDB connection parameters
-
-Dependencies:
-- os: Environment variable access
-- typing: Type hints and annotations
-"""
-
-import os
-from typing import Dict, Any
+import json
+from os import getenv
+from typing import Any
 
 
-def build_postgres_config() -> Dict[str, Any]:
-    """
-    Build PostgreSQL database configuration from environment variables.
-    
-    This function constructs a PostgreSQL configuration dictionary from
-    environment variables with sensible defaults for local development.
-    It handles connection pool configuration and security parameters.
-    
-    Environment Variables:
-        POSTGRES_HOST: Database hostname (default: localhost)
-        POSTGRES_PORT: Database port (default: 5432)
-        POSTGRES_USER: Database username (default: sientia)
-        POSTGRES_PASSWORD: Database password (default: sientia)
-        POSTGRES_DBNAME: Database name (default: sientia)
-        POSTGRES_MIN_CONNECTIONS: Minimum connection pool size (default: 1)
-        POSTGRES_MAX_CONNECTIONS: Maximum connection pool size (default: 10)
-    
-    Returns:
-        dict: PostgreSQL configuration dictionary with all required parameters
-        
-    Example:
-        >>> config = build_postgres_config()
-        >>> print(config)
-        {
-            'host': 'localhost',
-            'port': 5432,
-            'user': 'sientia',
-            'password': 'sientia',
-            'dbname': 'sientia',
-            'min_connections': 1,
-            'max_connections': 10
-        }
-    
-    Note:
-        In production, ensure all required environment variables are set
-        with appropriate values for your database environment.
-    """
-    return {
-        'host': os.getenv('POSTGRES_HOST', 'localhost'),
-        'port': int(os.getenv('POSTGRES_PORT', '5432')),
-        'user': os.getenv('POSTGRES_USER', 'sientia'),
-        'password': os.getenv('POSTGRES_PASSWORD', 'sientia'),
-        'dbname': os.getenv('POSTGRES_DBNAME', 'sientia'),
-        'min_connections': int(os.getenv('POSTGRES_MIN_CONNECTIONS', '1')),
-        'max_connections': int(os.getenv('POSTGRES_MAX_CONNECTIONS', '10'))
-    }
-
-
-def build_mlflow_config() -> Dict[str, Any]:
+def build_mlflow_config() -> dict[str, Any]:
     """
     Build MLFlow server configuration from environment variables.
-    
+
     This function constructs an MLFlow configuration dictionary from
     environment variables with sensible defaults for local development.
     It handles server connection and authentication parameters.
-    
+
     Environment Variables:
-        MLFLOW_HOST: MLFlow server hostname (default: localhost)
-        MLFLOW_PORT: MLFlow server port (default: 5000)
-        MLFLOW_USERNAME: MLFlow username (default: admin)
-        MLFLOW_PASSWORD: MLFlow password (default: admin)
-    
+        MLFLOW_HOST: MLFlow server hostname (default: http://localhost)
+        MLFLOW_PORT: MLFlow server port (default: 5080)
+        MLFLOW_USERNAME: MLFlow username (default: aignosi)
+        MLFLOW_PASSWORD: MLFlow password (default: aignosi)
+
     Returns:
         dict: MLFlow configuration dictionary with all required parameters
-        
-    Example:
-        >>> config = build_mlflow_config()
-        >>> print(config)
-        {
-            'host': 'localhost',
-            'port': 5000,
-            'username': 'admin',
-            'password': 'admin'
-        }
-    
-    Note:
-        In production, ensure all required environment variables are set
-        with appropriate values for your MLFlow server environment.
-        Consider using secure authentication methods for production deployments.
     """
     return {
-        'host': os.getenv('MLFLOW_HOST', 'localhost'),
-        'port': int(os.getenv('MLFLOW_PORT', '5000')),
-        'username': os.getenv('MLFLOW_USERNAME', 'admin'),
-        'password': os.getenv('MLFLOW_PASSWORD', 'admin')
+        'host': getenv('MLFLOW_HOST', 'http://localhost'),
+        'port': int(getenv('MLFLOW_PORT', '5080')),
+        'username': getenv('MLFLOW_USERNAME', 'aignosi'),
+        'password': getenv('MLFLOW_PASSWORD', 'aignosi'),
     }
 
 
-def build_opc_config() -> Dict[str, Any]:
+def build_opc_config() -> dict[str, Any]:
     """
     Build OPC server configuration from environment variables.
-    
+
     This function constructs an OPC server configuration dictionary from
     environment variables. It supports both single server and multi-server
     configurations with flexible parameter handling.
-    
+
     Environment Variables:
         OPC_CONFIG: JSON string containing multiple OPC server configurations
-        OPC_URL: Single OPC server URL (fallback)
-        OPC_NAME: Single OPC server name (fallback)
-        OPC_SERVER_URI: Single OPC server URI (fallback)
-        OPC_CERT_PATH: Client certificate path (fallback)
-        OPC_PRIVATE_KEY_PATH: Client private key path (fallback)
-        OPC_SERVER_CERT_PATH: Server certificate path (fallback)
-        OPC_RECONNECTION_INTERVAL: Reconnection interval in milliseconds (fallback)
-    
+        OPC_ID: OPC server ID (fallback, default: 1)
+        OPC_URL: Single OPC server URL (fallback, default: opc.tcp://localhost:4840)
+        OPC_SERVER_URI: Single OPC server URI (fallback, default: opc.tcp://localhost:4840)
+        OPC_CERT_PATH: Client certificate path (fallback, default: None)
+        OPC_PRIVATE_KEY_PATH: Client private key path (fallback, default: None)
+        OPC_SERVER_CERT_PATH: Server certificate path (fallback, default: None)
+        OPC_RECONNECTION_INTERVAL: Reconnection interval in seconds (fallback, default: 120)
+
     Returns:
         dict: OPC server configuration dictionary
-        
-    Configuration Modes:
-        1. Multi-server: Use OPC_CONFIG environment variable with JSON string
-        2. Single server: Use individual OPC_* environment variables
-        
-    Example Multi-server Configuration:
-        >>> # Set OPC_CONFIG environment variable
-        >>> os.environ['OPC_CONFIG'] = '''
-        ... {
-        ...     "opc_server_1": {
-        ...         "url": "opc.tcp://server1:4840",
-        ...         "name": "Server1",
-        ...         "server_uri": "urn:server1:opcua",
-        ...         "cert_path": "/path/to/cert.pem",
-        ...         "private_key_path": "/path/to/key.pem",
-        ...         "server_cert_path": "/path/to/server_cert.pem",
-        ...         "reconnection_interval": 5000
-        ...     }
-        ... }
-        ... '''
-        >>> config = build_opc_config()
-        
-    Example Single Server Configuration:
-        >>> # Set individual environment variables
-        >>> os.environ['OPC_URL'] = 'opc.tcp://localhost:4840'
-        >>> os.environ['OPC_NAME'] = 'LocalServer'
-        >>> config = build_opc_config()
-    
-    Note:
-        For production deployments, prefer the OPC_CONFIG approach for
-        multiple servers and ensure all certificate paths are properly configured.
     """
-    # Check for multi-server configuration
-    opc_config = os.getenv('OPC_CONFIG')
-    if opc_config:
-        try:
-            import json
-            return json.loads(opc_config)
-        except (json.JSONDecodeError, ImportError) as e:
-            # Fall back to single server configuration if JSON parsing fails
-            pass
-    
-    # Single server configuration fallback
+    opc_raw = getenv('OPC_CONFIG', None)
+
+    if opc_raw:
+        return json.loads(opc_raw)
+
     return {
-        'default': {
-            'url': os.getenv('OPC_URL', 'opc.tcp://localhost:4840'),
-            'name': os.getenv('OPC_NAME', 'DefaultServer'),
-            'server_uri': os.getenv('OPC_SERVER_URI', 'urn:default:opcua'),
-            'cert_path': os.getenv('OPC_CERT_PATH', ''),
-            'private_key_path': os.getenv('OPC_PRIVATE_KEY_PATH', ''),
-            'server_cert_path': os.getenv('OPC_SERVER_CERT_PATH', ''),
-            'reconnection_interval': int(os.getenv('OPC_RECONNECTION_INTERVAL', '5000'))
+        getenv('OPC_ID', '1'): {
+            'id': getenv('OPC_ID', '1'),
+            'server_name': getenv('OPC_SERVER_NAME', 'default_server'),
+            'url': getenv('OPC_URL', 'opc.tcp://localhost:4840'),
+            'server_uri': getenv('OPC_SERVER_URI', 'opc.tcp://localhost:4840'),
+            'cert_path': getenv('OPC_CERT_PATH', None),
+            'private_key_path': getenv('OPC_PRIVATE_KEY_PATH', None),
+            'server_cert_path': getenv('OPC_SERVER_CERT_PATH', None),
+            'reconnection_interval': int(getenv('OPC_RECONNECTION_INTERVAL', '120')),
         }
     }
 
 
-def build_mongodb_config() -> Dict[str, Any]:
+def build_minio_config() -> dict[str, Any]:
     """
-    Build MongoDB configuration from environment variables.
-    
-    This function constructs a MongoDB configuration dictionary from
-    environment variables with sensible defaults for local development.
-    It handles connection string and database name configuration.
-    
+    Build MinIO (S3-compatible) configuration from environment variables.
+
     Environment Variables:
-        MONGODB_URL: MongoDB connection URI (default: localhost:27017)
-        MONGODB_DATABASE: MongoDB database name (default: sientia)
-    
+        MINIO_ENDPOINT: MinIO endpoint including scheme (default: http://localhost:9000)
+        MINIO_ACCESS_KEY: Access key (default: minioadmin)
+        MINIO_SECRET_KEY: Secret key (default: minioadmin)
+        MINIO_REGION: Region name for S3 client (default: us-east-1)
+        MINIO_BUCKET_DEFAULT: Default bucket for uploads (default: laborious)
+        MINIO_SECURE: Whether to use HTTPS (default: false)
     Returns:
-        dict: MongoDB configuration dictionary with connection parameters
-        
-    Example:
-        >>> config = build_mongodb_config()
-        >>> print(config)
-        {
-            'connection_string': 'localhost:27017',
-            'database_name': 'sientia'
-        }
-    
-    Note:
-        In production, ensure the MONGODB_URL environment variable is set
-        with a proper MongoDB connection string including authentication
-        if required by your MongoDB deployment.
+        dict: MinIO configuration dictionary
     """
     return {
-        'connection_string': os.getenv('MONGODB_URL', 'localhost:27017'),
-        'database_name': os.getenv('MONGODB_DATABASE', 'sientia')
+        'endpoint_url': getenv('MINIO_ENDPOINT_URL', 'http://localhost:9000'),
+        'access_key': getenv('MINIO_ACCESS_KEY', 'minioadmin'),
+        'secret_key': getenv('MINIO_SECRET_KEY', 'minioadmin'),
+        'default_bucket': getenv('MINIO_DEFAULT_BUCKET', 'laborious'),
+        'retention_hours': int(getenv('MINIO_RETENTION_HOURS', '24')),
+        'secure': getenv('MINIO_SECURE', 'false') == 'true',
     }

laborious/utils/dataframe_debug.py

@@ -0,0 +1,34 @@
+from typing import Any
+
+from pandas import DataFrame
+
+DEFAULT_MAX_DEBUG_DATAFRAME_ROWS = 100
+
+
+def build_dataframe_debug_message(
+    message: str,
+    data: Any,
+    max_rows: int = DEFAULT_MAX_DEBUG_DATAFRAME_ROWS,
+) -> str:
+    """
+    Build a safe debug message for dataframe payloads
+
+    Args:
+     - message (str): Base message to identify the logged payload
+     - data (Any): Payload to evaluate for dataframe-aware logging
+     - max_rows (int): Maximum dataframe row count allowed for full payload logging
+
+    Return:
+     Formatted debug message with full dataframe content or compact summary
+    """
+    if not isinstance(data, DataFrame):
+        return f'{message} {data}'
+
+    rows = data.shape[0]
+    if rows <= max_rows:
+        return f'{message}\n{data.to_csv()}'
+
+    return (
+        f'{message} skipped because dataframe has {rows} rows '
+        f'(max: {max_rows}). Shape: {data.shape}'
+    )

laborious/utils/filters/conditional_filters.py

@@ -1,197 +1,48 @@
-"""
-Conditional Data Filters Module
-
-This module provides conditional data filtering functions for the Sientia DataOps Laborious system.
-It implements data quality validation filters that can be applied to input data before
-ML operations to ensure data integrity and quality.
-
-The module implements filters for:
-1. Empty data detection and validation
-2. Specific variable null value checking
-3. Configurable data quality rules
-4. Flexible filter configuration
-
-Key Features:
-- Configurable filter policies and thresholds
-- Multiple data quality validation rules
-- Flexible configuration options
-- Comprehensive error handling
-- Performance-optimized filtering
-
-Filter Types:
-- EMPTY_DATA: Detects empty or insufficient data sets
-- SPECIFIC_VARIABLES_NULL_VALUES: Validates specific variable null values
-- Custom filters can be added for specific validation needs
-
-Dependencies:
-- pandas.DataFrame: Data manipulation and processing
-- typing: Type hints and annotations
-"""
-
-from typing import Any, Dict, List
 from pandas import DataFrame
 
 
-def filter_empty_data(data: DataFrame, config: Dict[str, Any]) -> bool:
+def filter_specific_variables_null_values(data: DataFrame, config: dict) -> bool:
     """
-    Filter data based on empty data conditions.
-    
-    This function checks if the input data meets minimum requirements for
-    processing. It can validate data size, completeness, and other quality
-    metrics to ensure sufficient data is available for ML operations.
-    
-    The filter implements multiple validation criteria:
-    1. Data frame size validation
-    2. Row count validation
-    3. Column completeness validation
-    4. Configurable threshold checking
-    
+    Filter to check if specific variables contain null values.
+
+    This function examines a DataFrame to determine if any of the specified variables
+    contain null (NaN) values. It returns True if null values are found for any of
+    the specified variables, False otherwise.
+
     Args:
-        data: Input data as pandas DataFrame
-        config: Filter configuration dictionary
-            Required keys:
-                - min_rows (int, optional): Minimum number of rows required
-                - min_columns (int, optional): Minimum number of columns required
-                - min_data_points (int, optional): Minimum total data points required
-    
+        data (DataFrame): The pandas DataFrame to be examined. Must contain columns
+            named 'variable' and 'value'.
+        config (dict): Configuration dictionary containing the following key:
+            - variables (list): List of variable names to check for null values
+
     Returns:
-        bool: True if data should be filtered (fails quality check), False otherwise
-        
-    Filter Logic:
-        - Returns True (filter) if data is empty or below thresholds
-        - Returns False (pass) if data meets quality requirements
-        - Handles missing configuration gracefully with defaults
-        
-    Example:
-        >>> import pandas as pd
-        >>> df = pd.DataFrame({'A': [1, 2, 3], 'B': [4, 5, 6]})
-        >>> config = {'min_rows': 2, 'min_columns': 2}
-        >>> result = filter_empty_data(df, config)
-        >>> print(result)
-        False  # Data passes filter
-        
-        >>> empty_df = pd.DataFrame()
-        >>> result = filter_empty_data(empty_df, config)
-        >>> print(result)
-        True   # Data fails filter
-        
-    Default Thresholds:
-        - min_rows: 1 (at least one row required)
-        - min_columns: 1 (at least one column required)
-        - min_data_points: 1 (at least one data point required)
+        bool: True if any of the specified variables contain null values,
+              False if none of the specified variables contain null values.
+
     """
-    # Check if data is completely empty
+
     if data.empty:
-        return True
-    
-    # Get configuration with defaults
-    min_rows = config.get('min_rows', 1)
-    min_columns = config.get('min_columns', 1)
-    min_data_points = config.get('min_data_points', 1)
-    
-    # Check row count
-    if len(data) < min_rows:
-        return True
-    
-    # Check column count
-    if len(data.columns) < min_columns:
-        return True
-    
-    # Check total data points
-    if data.size < min_data_points:
-        return True
-    
-    # Data passes all quality checks
-    return False
+        return False
+
+    return not data[data['variable'].isin(config['variables']) & data['value'].isna()].empty
 
 
-def filter_specific_variables_null_values(data: DataFrame, config: Dict[str, Any]) -> bool:
+def filter_empty_data(data: DataFrame, _config: dict) -> bool:
     """
-    Filter data based on null values in specific variables.
-    
-    This function checks for null values in specified variables and determines
-    if the data quality is sufficient for processing. It can validate
-    individual columns or groups of columns for data completeness.
-    
-    The filter implements variable-specific validation:
-    1. Individual variable null value checking
-    2. Configurable null value thresholds
-    3. Multiple variable validation
-    4. Flexible threshold configuration
-    
+    Filter to check if the DataFrame is empty.
+
+    This function determines whether the provided DataFrame contains any data.
+    It's a simple utility function that can be used in conditional logic to
+    handle cases where no data is available.
+
     Args:
-        data: Input data as pandas DataFrame
-        config: Filter configuration dictionary
-            Required keys:
-                - variables (list): List of variable names to check
-                - max_null_ratio (float, optional): Maximum allowed null value ratio (0.0 to 1.0)
-                - max_null_count (int, optional): Maximum allowed null value count
-    
+        data (DataFrame): The pandas DataFrame to be checked for emptiness.
+        _config (dict): Configuration dictionary (unused in this function).
+            The underscore prefix indicates this parameter is required for
+            interface consistency but not used in the implementation.
+
     Returns:
-        bool: True if data should be filtered (fails quality check), False otherwise
-        
-    Filter Logic:
-        - Returns True (filter) if null value thresholds are exceeded
-        - Returns False (pass) if null values are within acceptable limits
-        - Handles missing variables gracefully
-        - Supports both ratio and count-based thresholds
-        
-    Example:
-        >>> import pandas as pd
-        >>> df = pd.DataFrame({
-        ...     'temperature': [25.5, None, 27.0, 26.5],
-        ...     'humidity': [60.0, 65.0, None, 62.0]
-        ... })
-        >>> config = {
-        ...     'variables': ['temperature', 'humidity'],
-        ...     'max_null_ratio': 0.25
-        ... }
-        >>> result = filter_specific_variables_null_values(df, config)
-        >>> print(result)
-        False  # Data passes filter (null ratio = 0.25, which equals max)
-        
-        >>> config = {
-        ...     'variables': ['temperature', 'humidity'],
-        ...     'max_null_ratio': 0.20
-        ... }
-        >>> result = filter_specific_variables_null_values(df, config)
-        >>> print(result)
-        True   # Data fails filter (null ratio = 0.25, exceeds max of 0.20)
-        
-    Default Thresholds:
-        - max_null_ratio: 0.5 (50% null values allowed)
-        - max_null_count: None (no count-based limit by default)
-        
-    Note:
-        If both max_null_ratio and max_null_count are specified, the filter
-        will trigger if either threshold is exceeded.
+        bool: True if the DataFrame is empty (has no rows), False if it contains data.
+
     """
-    # Get configuration
-    variables = config.get('variables', [])
-    max_null_ratio = config.get('max_null_ratio', 0.5)
-    max_null_count = config.get('max_null_count', None)
-    
-    # Check if variables exist in data
-    if not variables:
-        return False  # No variables specified, pass filter
-    
-    # Validate each specified variable
-    for variable in variables:
-        if variable not in data.columns:
-            continue  # Skip variables that don't exist in data
-        
-        # Calculate null value statistics
-        null_count = data[variable].isnull().sum()
-        total_count = len(data[variable])
-        null_ratio = null_count / total_count if total_count > 0 else 0.0
-        
-        # Check ratio threshold
-        if null_ratio > max_null_ratio:
-            return True
-        
-        # Check count threshold (if specified)
-        if max_null_count is not None and null_count > max_null_count:
-            return True
-    
-    # All variables pass null value checks
-    return False
+    return data.empty

laborious/utils/filters/mlflow_filters.py

@@ -52,8 +52,11 @@ def nan_values_filter(predictions: DataFrame, _config: dict) -> bool:
         bool: True if data should be filtered (too many NaN values), False otherwise
 
     """
-    data = predictions.replace({None: np.nan}).drop(
-        columns=['timestamp'], errors='ignore').infer_objects()
+    data = (
+        predictions.replace({None: np.nan})
+        .drop(columns=['timestamp'], errors='ignore')
+        .infer_objects()
+    )
 
     if data.isna().all().all():
         return True

laborious/utils/models/minio_dataframe_payload.py

@@ -0,0 +1,348 @@
+"""
+MinIO-backed DataFrame payload for Temporal workflows.
+
+Data is never stored as a pandas ``DataFrame`` field on the dataclass.
+Instead, the DataFrame is only provided as an input to:
+`from_dataframe` / `from_dataframe_to_dict`.
+
+At build time, the DataFrame is evaluated for its serialized size; if it exceeds
+the configured threshold, it is serialized to parquet bytes and uploaded to MinIO.
+Otherwise, it is inlined as a Temporal-friendly ``dict``.
+"""
+
+import pickle
+import re
+from collections.abc import Hashable
+from dataclasses import dataclass
+from datetime import datetime
+from io import BytesIO
+from os import getenv
+from typing import Any, Literal
+
+from pandas import DataFrame, read_parquet
+from sientia_do.observability.logger import Logger
+from sientia_do.repository.minio_repository import MinioRepository
+from sientia_do.temporal.constants import DATETIME_FORMAT_FILENAME, DATETIME_FORMAT_WITH_TZ, now
+
+# Keys that are part of the serialized wire format (not arbitrary metadata).
+_SERIALIZED_FIELD_KEYS = frozenset({'data', 'bucket', 'object_key', 'object_prefix', 'uri'})
+
+_OBJECT_TIMESTAMP_PATTERN = re.compile(
+    r'-(?:initial|transform)-(\d{4}-\d{2}-\d{2}_\d{2}-\d{2}-\d{2})\.parquet$'
+)
+
+OFFLOAD_THRESHOLD_BYTES = int(
+    float(getenv('SIENTIA_MINIO_OFFLOAD_THRESHOLD_MEGABYTES', '1.5')) * 1024 * 1024
+)
+
+# Relative prefix used for storing offloaded prediction datasets in MinIO.
+# It is also the root directory for retention cleanup listing.
+PREDICTION_DATASETS_PREFIX = 'prediction_datasets'
+
+OperationKind = Literal['initial', 'transform', 'predict']
+
+
+def _build_object_key(
+    model_name: str, operation: OperationKind, timestamp: str
+) -> tuple[str, str | None]:
+    """
+    Build the MinIO object key and the directory prefix used for retention listing.
+
+    Args:
+        model_name: Registered model name used in the pipeline.
+        operation: Either initial (pre-transform load) or transform (post-MLFlow transform).
+        timestamp: Filename timestamp segment from DATETIME_FORMAT_FILENAME.
+
+    Return:
+        tuple[str, str | None]: Full object key and normalized prefix (or None if at bucket root).
+    """
+    # Naming convention:
+    #   - Directory is always `prediction_datasets/<model_name>`
+    #   - Filename follows the retention-parsing pattern
+    basename = f'{model_name}-{operation}-{timestamp}.parquet'
+    model_dir = model_name.strip().strip('/')
+    prefix = f'{PREDICTION_DATASETS_PREFIX}/{model_dir}'
+    return f'{prefix}/{basename}', prefix
+
+
+@dataclass
+class MinioDataFramePayload:
+    """
+    Serializable payload after a DataFrame was evaluated: inline tabular dict and/or MinIO keys.
+
+    Build from a live DataFrame only via `from_dataframe` / `from_dataframe_to_dict`.
+    Rehydrate from Temporal via `from_dict`. The DataFrame is not a field on this class.
+    """
+
+    last_timestamp: str
+    status: dict[str, Any] | None = None
+    data: dict[Hashable, Any] | None = None
+    bucket: str | None = None
+    object_key: str | None = None
+    object_prefix: str | None = None
+    uri: str | None = None
+
+    @staticmethod
+    def _debug(
+        logger: Logger | None,
+        message: str,
+        metadata: dict[str, Any] | None = None,
+    ) -> None:
+        """
+        Emit debug logs only when logger is provided
+
+        Args:
+         - logger (Logger | None): Logger instance used for debug messages
+         - message (str): Message to be logged
+         - metadata (dict[str, Any] | None): Optional workflow metadata context
+        """
+        if logger is None:
+            return
+        logger.custom_debug(message, metadata)
+
+    @classmethod
+    def from_dict(cls, raw: 'dict[str, Any] | MinioDataFramePayload') -> 'MinioDataFramePayload':
+        """
+        Reconstruct a MinioDataFramePayload from a plain dict produced by Temporal serialization.
+
+        Temporal converts dataclass return values into plain dicts when crossing
+        workflow/activity boundaries. This method rebuilds the typed instance so
+        that methods like ``retrieve``, ``cleanup_prefix`` and ``has_data`` are
+        available on the receiving side.
+
+        If the argument is already a MinioDataFramePayload, it is returned as-is.
+
+        Args:
+            raw: Dict with keys matching the dataclass fields
+                (last_timestamp, status, data, bucket, object_key, object_prefix, uri),
+                or an existing MinioDataFramePayload instance.
+
+        Return:
+            MinioDataFramePayload: Reconstructed (or original) instance.
+        """
+        if isinstance(raw, MinioDataFramePayload):
+            return raw
+        return cls(
+            last_timestamp=raw['last_timestamp'],
+            status=raw.get('status'),
+            data=raw.get('data'),
+            bucket=raw.get('bucket'),
+            object_key=raw.get('object_key'),
+            object_prefix=raw.get('object_prefix'),
+            uri=raw.get('uri'),
+        )
+
+    @staticmethod
+    def estimate_size_bytes(
+        df: DataFrame,
+        metadata: dict[str, Any] | None = None,
+        logger: Logger | None = None,
+    ) -> int:
+        """
+        Approximate serialized size of the DataFrame as the default-orient dict.
+
+        Args:
+            df: DataFrame whose tabular content size is estimated.
+
+        Return:
+            int: Estimated size in bytes (pickle of dict representation).
+        """
+        try:
+            size = len(pickle.dumps(df.to_dict()))
+        except Exception:
+            size = len(pickle.dumps(df))
+
+        MinioDataFramePayload._debug(
+            logger,
+            f'DataFrame size: {size} bytes',
+            metadata,
+        )
+        return size
+
+    @staticmethod
+    def parse_object_timestamp(object_key: str) -> datetime | None:
+        """
+        Parse the timestamp embedded in the object key basename (before .parquet).
+
+        Args:
+            object_key: S3/MinIO object key whose basename follows
+                ``{model}-{initial|transform}-{DATETIME_FORMAT_FILENAME}.parquet``.
+
+        Return:
+            datetime | None: Parsed UTC-naive datetime from the key, or None if not matched.
+        """
+        basename = object_key.rsplit('/', 1)[-1]
+        match = _OBJECT_TIMESTAMP_PATTERN.search(basename)
+        if not match:
+            return None
+        try:
+            return datetime.strptime(match.group(1), DATETIME_FORMAT_FILENAME)
+        except ValueError:
+            return None
+
+    def cleanup_prefix(self) -> str | None:
+        """
+        Return True if cleanup is enabled for this payload.
+        """
+        if self.object_key is not None and self.data is None:
+            return self.object_prefix
+        return None
+
+    def has_data(self) -> bool:
+        """
+        Return True if the payload has some data internally or in MinIO.
+        """
+        return (self.data is not None and self.data != {}) or self.object_key is not None
+
+    @classmethod
+    async def from_dataframe(
+        cls,
+        dataframe: DataFrame | None,
+        minio_repo: MinioRepository,
+        model_name: str,
+        operation: OperationKind,
+        status: dict[str, Any] | None = None,
+        workflow_metadata: dict | None = None,
+        last_timestamp: str | None = None,
+        logger: Logger | None = None,
+    ) -> 'MinioDataFramePayload':
+        """
+        Evaluate the DataFrame size, then either inline dict or upload parquet to MinIO.
+
+        The DataFrame is not stored on the returned instance.
+
+        Args:
+            dataframe: Tabular data to evaluate and persist (inline or MinIO).
+            metadata: Small metadata dict merged into the payload (e.g. success, message).
+            minio_repo: sientia_do MinioRepository (or compatible) with `upload_file()`.
+            workflow_metadata: Metadata passed to MinIO store for logging/metrics.
+            model_name: Registered model name used in the object basename.
+            operation: Either ``initial`` (query load) or ``transform`` (post-transform).
+            key_prefix: Backward-compatible parameter (currently ignored for object naming).
+            size_threshold_bytes: Byte limit before offload. When None, the module-level
+                environment-derived default is used.
+
+        Return:
+            MinioDataFramePayload: Instance with data and/or MinIO fields set.
+        """
+
+        if dataframe is None or dataframe.empty:
+            cls._debug(
+                logger,
+                'MinioDataFramePayload.from_dataframe received empty dataframe, returning empty payload',
+                workflow_metadata,
+            )
+            return cls(
+                data=None, last_timestamp=now().strftime(DATETIME_FORMAT_WITH_TZ), status=status
+            )
+
+        if last_timestamp is None:
+            last_timestamp = max(dataframe['timestamp'].values.tolist())
+
+        dataframe_size = cls.estimate_size_bytes(dataframe, workflow_metadata, logger)
+        cls._debug(
+            logger,
+            (
+                f'MinioDataFramePayload.from_dataframe estimated size: {dataframe_size} bytes '
+                f'(threshold: {OFFLOAD_THRESHOLD_BYTES} bytes)'
+            ),
+            workflow_metadata,
+        )
+
+        if dataframe_size <= OFFLOAD_THRESHOLD_BYTES:
+            cls._debug(
+                logger,
+                'MinioDataFramePayload.from_dataframe using inline payload',
+                workflow_metadata,
+            )
+            return cls(data=dataframe.to_dict(), last_timestamp=last_timestamp, status=status)
+
+        timestamp = now().strftime(DATETIME_FORMAT_FILENAME)
+        object_key, object_prefix = _build_object_key(model_name, operation, timestamp)
+        cls._debug(
+            logger,
+            (
+                'MinioDataFramePayload.from_dataframe offloading payload to MinIO '
+                f'with key {object_key}'
+            ),
+            workflow_metadata,
+        )
+
+        # Upload using the relative object key. The upstream repository will
+        # prefix it internally under its MinIO namespace.
+        parquet_buffer = BytesIO()
+        dataframe.to_parquet(parquet_buffer, engine='pyarrow', index=True)
+        file_bytes = parquet_buffer.getvalue()
+
+        upload_result = await minio_repo.upload_file(
+            file_bytes=file_bytes,
+            relative_key=object_key,
+            metadata=workflow_metadata,
+        )
+
+        bucket = minio_repo.bucket
+        object_key_full = upload_result.get('minio_object_name', object_key)
+        uri = f's3://{bucket}/{object_key_full}' if bucket else None
+        cls._debug(
+            logger,
+            f'MinioDataFramePayload.from_dataframe upload completed: {uri}',
+            workflow_metadata,
+        )
+
+        return cls(
+            data=None,
+            bucket=bucket,
+            object_key=object_key_full,
+            object_prefix=object_prefix,
+            uri=uri,
+            last_timestamp=last_timestamp,
+            status=status,
+        )
+
+    async def retrieve(
+        self,
+        minio_repo: MinioRepository,
+        workflow_metadata: dict[str, Any] | None = None,
+        logger: Logger | None = None,
+    ) -> DataFrame:
+        """
+        Load parquet from MinIO when object_key is set and populate inline data.
+
+        Args:
+            minio_repo: sientia_do MinioRepository (or compatible) with download_file().
+            workflow_metadata: Metadata passed to MinIO read for logging/metrics.
+
+        Return:
+            dict[str, Any]: Flat dict with data filled (same keys as to_dict after load).
+        """
+        if self.data is not None:
+            self._debug(
+                logger,
+                'MinioDataFramePayload.retrieve using inline payload data',
+                workflow_metadata,
+            )
+            return DataFrame(self.data)
+
+        if not self.has_data():
+            self._debug(
+                logger,
+                'MinioDataFramePayload.retrieve found no payload data, returning empty dataframe',
+                workflow_metadata,
+            )
+            return DataFrame()
+
+        self._debug(
+            logger,
+            f'MinioDataFramePayload.retrieve downloading object from MinIO: {self.object_key}',
+            workflow_metadata,
+        )
+        file_bytes = await minio_repo.download_file(
+            object_name=self.object_key, metadata=workflow_metadata
+        )
+        df = read_parquet(BytesIO(file_bytes))
+        self._debug(
+            logger,
+            f'MinioDataFramePayload.retrieve loaded dataframe from MinIO with shape {df.shape}',
+            workflow_metadata,
+        )
+        return df

laborious/utils/repository/minio_manager.py

@@ -0,0 +1,32 @@
+from sientia_do.notifications.handlers import NotificationHandler
+from sientia_do.observability.logger import Logger
+from sientia_do.observability.metrics_controller import MetricsController
+from sientia_do.observability.sientia_monitoring import SientiaMonitoring
+from sientia_do.repository.minio_repository import MinioRepository
+
+
+class MinioManager(SientiaMonitoring):
+    minio_repository: MinioRepository | None = None
+
+    def __init__(
+        self,
+        minio_repository: MinioRepository | None = None,
+        logger: Logger | None = None,
+        notification_handler: NotificationHandler | None = None,
+        metrics_controller: MetricsController | None = None,
+    ):
+        if self.minio_repository is None:
+            self.minio_repository = minio_repository
+        SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
+
+    def close(self) -> None:
+        """
+        Close the MinioManager and clean up resources.
+        """
+        if self.minio_repository is not None:
+            try:
+                self.minio_repository.close()
+            finally:
+                self.minio_repository = None
+
+        SientiaMonitoring.shutdown(self)

laborious/worker/worker.py

@@ -5,18 +5,25 @@ 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:
-- predictions_batch-queue: Handles batch prediction workflows
-- minimal_retrain-queue: Handles model retraining workflows
+The worker supports multiple runtime-scoped task queues (via ``sientia_do.temporal.worker.prepare_worker``):
+- predictions_batch-{runtime}-queue: Batch prediction workflows (heavy workload)
+- minimal_retrain-{runtime}-queue: Model retraining workflows
+- drift-{runtime}-queue: Drift detection workflows
+- simple_metrics-{runtime}-queue: Simple metrics workflows
+
+``RUNTIME`` must be set; it is passed to every ``prepare_worker`` call. Schedulers must use the
+same queue names (breaking change vs legacy ``drift-queue`` / ``simple_metrics-queue``).
 
 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
 - Multiple worker instances for different workflow types
 
 Environment Variables:
+- RUNTIME: Required non-empty string; suffix for all task queue names
 - TEMPORAL_HOST: Temporal server address (default: localhost:7233)
 - TEMPORAL_NAMESPACE: Temporal namespace (default: laborious)
 - POD_ID: Kubernetes pod identifier for metrics
@@ -25,33 +32,42 @@ Environment Variables:
 - PROJECT_NAME: Project name for notifications (default: laborious)
 """
 
-from temporalio import workflow, client
-from temporalio.worker import Worker, PollerBehaviorAutoscaling
-from temporalio.runtime import Runtime, TelemetryConfig, PrometheusConfig
+from temporalio import client, workflow
+from temporalio.runtime import PrometheusConfig, Runtime, TelemetryConfig
 
 with workflow.unsafe.imports_passed_through():
+    import asyncio
     import os
     import sys
-    import asyncio
-    from laborious.workflows.minimal_retrain import MinimalRetrain
-    from laborious.workflows.predictions_batch import PredictionsBatch
-    from laborious.workflows.sub_workflows.prediction_process import PredictionProcess
-    from laborious.workflows.sub_workflows.format_and_export_prediction import \
-        FormatAndExportPrediction
-    from laborious.activities.activities import Activities
-    from laborious.utils.connectors_config import (
-        build_postgres_config,
-        build_mlflow_config,
-        build_opc_config,
-        build_mongodb_config
-    )
+
+    from prometheus_client import start_http_server
     from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
     from sientia_do.observability.logger import get_logger
-    from laborious import metrics
-    from prometheus_client import start_http_server
+    from sientia_do.temporal.worker.prepare_worker import prepare_worker
+    from sientia_do.utils.connectors_config import (
+        build_api_config,
+        build_mongodb_config,
+        build_postgres_config,
+    )
 
-POD_ID = os.getenv('POD_ID')
-SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', "9091"))
+    from laborious import metrics
+    from laborious.activities.activities import Activities
+    from laborious.utils.connectors_config import (
+        build_minio_config,
+        build_mlflow_config,
+        build_opc_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
+
+POD_ID = os.getenv('HOSTNAME')
+SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', '9091'))
 
 
 async def main():
@@ -87,75 +103,110 @@ async def main():
 
     logger.custom_info(f'Starting Worker with POD_ID: {POD_ID}', metadata)
 
-    logger.custom_info("Starting prometheus client...", metadata)
+    runtime = os.getenv('RUNTIME', '').strip()
+    if not runtime:
+        logger.custom_critical(
+            'RUNTIME environment variable is required and must be non-empty',
+            metadata,
+        )
+        metrics.APP_UP.labels(pod_id=POD_ID).set(0)
+        sys.exit(1)
+
+    metadata_runtime = {**metadata, 'runtime': runtime}
+
+    logger.custom_info('Starting prometheus client...', metadata_runtime)
     start_prometheus_server()
 
-    logger.custom_info('Starting Notification Handler...', metadata)
+    logger.custom_info('Starting Notification Handler...', metadata_runtime)
 
     mongo_config = build_mongodb_config()
     notification_handler = NotificationHandler(
         connection_string=mongo_config['connection_string'],
         database=mongo_config['database_name'],
         logger=logger,
-        project_name=os.getenv('PROJECT_NAME', 'laborious')
+        project_name=os.getenv('PROJECT_NAME', 'laborious'),
     )
 
-    logger.custom_info('Starting Activities...', metadata)
+    logger.custom_info('Starting Activities...', metadata_runtime)
 
     activities = Activities(
         postgres_config=build_postgres_config(),
         mlflow_config=build_mlflow_config(),
+        minio_config=build_minio_config(),
         opc_config=build_opc_config(),
+        pi_web_api_config=build_api_config(),
         logger=logger,
-        notification_handler=notification_handler
+        notification_handler=notification_handler,
     )
 
-    logger.custom_info('Initializing OPC...', metadata)
+    logger.custom_info('Initializing OPC...', metadata_runtime)
     await activities.init_opc()
 
     logger.custom_info(
-        f'Starting SDK Metrics Server on port {SDK_METRICS_PORT}...', metadata)
+        f'Starting SDK Metrics Server on port {SDK_METRICS_PORT}...',
+        metadata_runtime,
+    )
 
     new_runtime = Runtime(
         telemetry=TelemetryConfig(
-            metrics=PrometheusConfig(
-                bind_address=f"0.0.0.0:{SDK_METRICS_PORT}")
+            metrics=PrometheusConfig(bind_address=f'0.0.0.0:{SDK_METRICS_PORT}')
         )
     )
 
-    logger.custom_info(f'Starting Temporal Client at {host}...', metadata)
+    logger.custom_info(f'Starting Temporal Client at {host}...', metadata_runtime)
 
     temporal_client = await client.Client.connect(
         target_host=host,
         namespace=os.getenv('TEMPORAL_NAMESPACE', 'laborious'),
-        runtime=new_runtime
+        runtime=new_runtime,
     )
 
-    logger.custom_info('Starting Workers...', metadata)
+    logger.custom_info(f'Starting Workers (runtime={runtime})...', metadata_runtime)
 
     workers = [
-        Worker(
-            temporal_client,
-            task_queue='minimal_retrain-queue',
-            workflows=[MinimalRetrain],
+        prepare_worker(
+            temporal_client=temporal_client,
+            main_workflow=MinimalRetrain,
+            other_workflows=[],
             activities=[
-                activities.load_custom_query,
+                activities.load_query_with_minio_offload,
                 activities.retrain_model,
                 activities.update_production_model,
-                activities.export_data_to_postgres
+                activities.format_retrain_report,
+                activities.export_data_to_postgres,
             ],
-            max_concurrent_workflow_tasks=50,
-            max_concurrent_activities=50,
-            max_concurrent_local_activities=50,
-            max_cached_workflows=200,
-            workflow_task_poller_behavior=PollerBehaviorAutoscaling(),
-            activity_task_poller_behavior=PollerBehaviorAutoscaling()
+            logger=logger,
+            runtime=runtime,
         ),
-        Worker(
-            temporal_client,
-            task_queue='predictions_batch-queue',
-            workflows=[PredictionsBatch, PredictionProcess,
-                       FormatAndExportPrediction],
+        prepare_worker(
+            temporal_client=temporal_client,
+            main_workflow=SimpleMetrics,
+            other_workflows=[],
+            activities=[
+                activities.load_custom_query,
+                activities.calculate_simple_metrics,
+                activities.export_data_to_postgres,
+            ],
+            logger=logger,
+            runtime=runtime,
+        ),
+        prepare_worker(
+            temporal_client=temporal_client,
+            main_workflow=Drift,
+            other_workflows=[],
+            activities=[
+                activities.load_custom_query,
+                activities.get_reference_data,
+                activities.calculate_drift,
+                activities.export_data_to_postgres,
+            ],
+            logger=logger,
+            runtime=runtime,
+        ),
+        prepare_worker(
+            temporal_client=temporal_client,
+            main_workflow=PredictionsBatch,
+            other_workflows=[PredictionProcess, FormatAndExportPrediction],
             activities=[
                 # MLFlow
                 activities.request_predict,
@@ -164,46 +215,47 @@ 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,
                 # OPC
                 activities.write_opc_data,
-                # Postgres
-                activities.load_custom_query,
+                # Postgres / MinIO offload
+                activities.load_query_with_minio_offload,
+                activities.cleanup_minio_objects_expired,
                 activities.repeat_last_prediction,
                 activities.export_data_to_postgres,
-                activities.write_metrics
+                activities.export_payload_to_postgres,
+                activities.write_metrics,
+                # Pi Web API
+                activities.write_pi_web_api_data,
             ],
-            max_concurrent_workflow_tasks=50,
-            max_concurrent_activities=50,
-            max_concurrent_local_activities=50,
-            max_cached_workflows=200,
-            workflow_task_poller_behavior=PollerBehaviorAutoscaling(),
-            activity_task_poller_behavior=PollerBehaviorAutoscaling()
-        )
+            logger=logger,
+            runtime=runtime,
+        ),
     ]
 
     handlers = []
     for w in workers:
         handlers.append(w.run())
 
-    logger.custom_info('Workers started successfully', metadata)
+    logger.custom_info('Workers started successfully', metadata_runtime)
 
+    exit_code = 0
     try:
         # This will run the workers and wait for them to complete.
         # If an exception occurs in any of the worker handlers, it will be propagated here.
         await asyncio.gather(*handlers)
     except BaseException as e:  # NOSONAR
-        logger.custom_error(f"An unhandled exception occurred: {e}", metadata)
+        logger.custom_error(f'An unhandled exception occurred: {e}', metadata)
+        exit_code = 1
     finally:
         if notification_handler:
             notification_handler.shutdown()
         if activities:
             await activities.shutdown()
-        # Exit with a non-zero status code to indicate failure to Kubernetes
         metrics.APP_UP.labels(pod_id=POD_ID).set(0)  # Mark app as DOWN
-        sys.exit(1)
+        sys.exit(exit_code)
 
 
 def start_prometheus_server():
@@ -224,12 +276,12 @@ def start_prometheus_server():
         SystemExit: If the metrics server fails to start
     """
     try:
-        port = int(os.getenv("HTTP_METRICS_PORT", 9090))
+        port = int(os.getenv('HTTP_METRICS_PORT', 9090))
         start_http_server(port)
-        print(f"Prometheus server started on port {port}.")
+        print(f'Prometheus server started on port {port}.')
         metrics.APP_UP.labels(pod_id=POD_ID).set(1)  # Mark app as UP
     except Exception as e:
-        print(f"Failed to start Prometheus server: {e}")
+        print(f'Failed to start Prometheus server: {e}')
         os._exit(1)
 
 

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
+        """  # nosec B608 - values come from internal Temporal workflow config, not user input
+
+        target_data_handler = workflow.start_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_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/minimal_retrain.py

@@ -1,14 +1,17 @@
 from temporalio import workflow
 
 with workflow.unsafe.imports_passed_through():
-    from laborious.activities.activities import Activities
-    from typing import Any
-    from sientia_do.temporal.policies import retry_policy
     from datetime import timedelta
+    from typing import Any
+
+    from sientia_do.temporal.policies import retry_policy
+
+    from laborious.activities.activities import Activities
+    from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload
 
 
-@workflow.defn(name="minimal_retrain")
-class MinimalRetrain():
+@workflow.defn(name='minimal_retrain')
+class MinimalRetrain:
     """
     Automated model retraining workflow for the Laborious system.
 
@@ -63,44 +66,62 @@ class MinimalRetrain():
                 'schedule_name': input_data['schedule_name'],
                 'model_name': input_data['model_name'],
                 'model_id': input_data['model_id'],
-                'workflow_name': 'minimal_retrain'
+                'workflow_name': 'minimal_retrain',
             }
         }
 
         model_name = input_data['model_name']
+        model_config = input_data.get('model_config', {})
 
-        data = await workflow.execute_local_activity_method(
-            Activities.load_custom_query,
+        storage_result = await workflow.execute_activity_method(
+            Activities.load_query_with_minio_offload,
             {
                 **metadata,
                 'query': input_data['query'],
-                'datetime_columns': input_data.get('datetime_columns', [])
+                'datetime_columns': input_data.get('datetime_columns', []),
+                'model_name': model_name,
             },
             retry_policy=retry_policy,
-            start_to_close_timeout=timedelta(seconds=60)
+            start_to_close_timeout=timedelta(seconds=600),
         )
 
+        storage_payload = MinioDataFramePayload.from_dict(storage_result)
+        if not storage_payload.has_data():
+            raise ValueError('No data returned from query')
+
         experiment_response = await workflow.execute_activity_method(
             Activities.retrain_model,
             {
                 **metadata,
-                'data': data,
-                'model_name': model_name
+                'data': storage_result,
+                'model_name': model_name,
+                'model_config': model_config,
             },
             retry_policy=retry_policy,
-            start_to_close_timeout=timedelta(seconds=60)
+            start_to_close_timeout=timedelta(hours=1),
         )
 
-        report = await workflow.execute_activity_method(
-            Activities.update_production_model,
+        if experiment_response['success']:
+            update_report = await workflow.execute_activity_method(
+                Activities.update_production_model,
+                {**metadata, 'model_name': model_name, **experiment_response},
+                retry_policy=retry_policy,
+                start_to_close_timeout=timedelta(seconds=60),
+            )
+        else:
+            update_report = {}
+
+        report = await workflow.execute_local_activity_method(
+            Activities.format_retrain_report,
             {
                 **metadata,
+                'experiment_response': experiment_response,
                 'model_name': model_name,
                 'model_id': input_data['model_id'],
-                **experiment_response
+                'update_report': update_report,
             },
             retry_policy=retry_policy,
-            start_to_close_timeout=timedelta(seconds=60)
+            start_to_close_timeout=timedelta(seconds=60),
         )
 
         await workflow.execute_activity_method(
@@ -109,8 +130,8 @@ class MinimalRetrain():
                 **metadata,
                 'data': report,
                 'schema': input_data['schema'],
-                'table_name': input_data['table_name']
+                'table_name': input_data['table_name'],
             },
             retry_policy=retry_policy,
-            start_to_close_timeout=timedelta(seconds=60)
+            start_to_close_timeout=timedelta(seconds=600),
         )

laborious/workflows/predictions_batch.py

@@ -1,14 +1,16 @@
 from temporalio import workflow
 
 with workflow.unsafe.imports_passed_through():
-    from laborious.activities.activities import Activities
-    from typing import Any
-    from sientia_do.temporal.policies import retry_policy
     from datetime import timedelta
+    from typing import Any
+
+    from sientia_do.temporal.policies import retry_policy
+
+    from laborious.activities.activities import Activities
 
 
-@workflow.defn(name="predictions_batch")
-class PredictionsBatch():
+@workflow.defn(name='predictions_batch')
+class PredictionsBatch:
     """
     Main batch prediction workflow for the Laborious system.
 
@@ -29,22 +31,6 @@ class PredictionsBatch():
         2. Configuration Preparation: Sets up prediction parameters and filters
         3. Workflow Delegation: Spawns PredictionProcess child workflow
         4. Error Handling: Implements comprehensive error handling and retry policies
-
-    Example:
-        >>> # Start the workflow
-        >>> await client.start_workflow(
-        ...     PredictionsBatch.run,
-        ...     id="batch_pred_001",
-        ...     task_queue="predictions_batch-queue",
-        ...     input_data={
-        ...         "schedule_name": "hourly_predictions",
-        ...         "model_name": "temperature_model",
-        ...         "model_id": "temp_001",
-        ...         "query": "SELECT * FROM sensor_data WHERE timestamp > NOW() - INTERVAL '1 hour'",
-        ...         "schema": {"timestamp": "datetime", "temperature": "float"},
-        ...         "table_name": "predictions"
-        ...     }
-        ... )
     """
 
     @workflow.run
@@ -75,7 +61,10 @@ class PredictionsBatch():
                     - model_retention (int, optional): Model retention period in minutes
                     - path_priority (list[str]): Decision path priority configuration
                     - opc_output_config (dict, optional): OPC server export configuration
+                    - pi_web_api_output_config (dict, optional): PI Web API export configuration
                     - datetime_columns (list[str], optional): Columns to treat as datetime
+                    - save_transform (bool, optional): Whether to save transformed data (default: True)
+                    - prediction_store_policy (str, optional): Data retention policy (default: 'lts:1')
 
         Returns:
             None: The workflow completes successfully when the child workflow finishes
@@ -90,20 +79,21 @@ class PredictionsBatch():
                 'schedule_name': input_data['schedule_name'],
                 'model_name': input_data['model_name'],
                 'model_id': input_data['model_id'],
-                'workflow_name': 'predictions_batch'
+                'workflow_name': 'predictions_batch',
             }
         }
 
-        # Load data using custom query
-        data = await workflow.execute_local_activity_method(
-            Activities.load_custom_query,
+        # Load data using custom query with optional MinIO offload for large frames
+        data = await workflow.execute_activity_method(
+            Activities.load_query_with_minio_offload,
             {
                 **metadata,
                 'query': input_data['query'],
-                'datetime_columns': input_data.get('datetime_columns', [])
+                'datetime_columns': input_data.get('datetime_columns', []),
+                'model_name': input_data['model_name'],
             },
             retry_policy=retry_policy,
-            start_to_close_timeout=timedelta(seconds=60)
+            start_to_close_timeout=timedelta(seconds=300),
         )
 
         # Prepare input for prediction_process workflow
@@ -112,28 +102,26 @@ 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'
-                }
-            }),
-            'mlflow_transform_filters': input_data.get('mlflow_transform_filters', {
-                'API_ERROR': {
-                    'POLICY': 'STOP'
-                }
-            }),
-            'mlflow_predict_filters': input_data.get('mlflow_predict_filters', {
-                'API_ERROR': {
-                    'POLICY': 'STOP'
-                }
-            }),
-            'model_retention': input_data.get('model_retention', 60),
+            'input_filters': input_data.get(
+                'input_filters', {'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}}}
+            ),
+            'mlflow_transform_filters': input_data.get(
+                'mlflow_transform_filters', {'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}}
+            ),
+            'mlflow_predict_filters': input_data.get(
+                'mlflow_predict_filters', {'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}}
+            ),
+            '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', {})
+            'opc_output_config': input_data.get('opc_output_config', {}),
+            'on_conflict': input_data.get('on_conflict', 'error'),
+            'pi_web_api_output_config': input_data.get('pi_web_api_output_config', {}),
+            'prediction_store_policy': input_data.get('prediction_store_policy', 'lts:1'),
+            'save_transform': input_data.get('save_transform', True),
         }
 
         # Execute prediction process workflow
-        await workflow.execute_child_workflow(
-            'prediction_process', prediction_input)
+        await workflow.execute_child_workflow('subworkflow.prediction_process', prediction_input)

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;
+        """  # nosec B608 - values come from internal Temporal workflow config, not user input
+
+        target_data = await workflow.execute_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

@@ -1,15 +1,17 @@
 from temporalio import workflow
 
 with workflow.unsafe.imports_passed_through():
-    from laborious.activities.activities import Activities
-    from typing import Any
     from datetime import timedelta
-    from sientia_do.temporal.policies import retry_policy
+    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="format_and_export_prediction")
-class FormatAndExportPrediction():
+@workflow.defn(name='subworkflow.format_and_export_prediction')
+class FormatAndExportPrediction:
     """
     Data formatting and export workflow for prediction results.
 
@@ -24,6 +26,7 @@ class FormatAndExportPrediction():
 
     Export Destinations:
     - PostgreSQL Database: Persistent storage with timestamp conversion
+    - PI Web API: Real-time industrial system integration for prediction and confidence values
     - OPC Servers: Real-time industrial system integration
     - Prometheus Metrics: Performance monitoring and operational visibility
     """
@@ -36,9 +39,10 @@ class FormatAndExportPrediction():
         This method orchestrates the complete data export process by:
         1. Determining the appropriate formatting strategy based on path_flag
         2. Formatting prediction data according to quality and requirements
-        3. Exporting data to OPC servers for real-time industrial access
-        4. Persisting data to PostgreSQL database with comprehensive metadata
-        5. Recording performance metrics for operational monitoring
+        3. Exporting data to PI Web API for real-time industrial access (if configured)
+        4. Exporting data to OPC servers for real-time industrial access (if configured)
+        5. Persisting data to PostgreSQL database with comprehensive metadata
+        6. Recording performance metrics for operational monitoring
 
         The method implements flexible formatting strategies:
         - Normal predictions: Full data formatting with confidence scores
@@ -48,29 +52,50 @@ 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
+                Optional keys:
                     - opc_output_config (dict[str, Any]): OPC server export configuration
-                    - prediction_store_policy (str, optional): Data retention policy
+                    - pi_web_api_output_config (dict[str, Any]): PI Web API export configuration
+                        Contains endpoint, prediction_tags, and confidence_tags mappings
+                    - 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']
 
+        opc_output_config = input_data.get('opc_output_config', None)
+        pi_web_api_output_config = input_data.get('pi_web_api_output_config', None)
+
         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,
                 {
@@ -79,15 +104,48 @@ class FormatAndExportPrediction():
                     'timestamp': input_data['timestamp'],
                     'model_id': input_data['model_id'],
                     'prediction_confidence': prediction_confidence,
-                    'prediction_store_policy': input_data.get(
-                        'prediction_store_policy', 'lts:1')
+                    'prediction_store_policy': input_data['prediction_store_policy'],
+                    'model_name': input_data['model_name'],
                 },
                 retry_policy=retry_policy,
-                start_to_close_timeout=timedelta(seconds=60)
+                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'],
+                        'model_name': input_data['model_name'],
+                    },
+                    retry_policy=retry_policy,
+                    start_to_close_timeout=timedelta(seconds=60),
+                )
+
+                write_transformed_handler = workflow.start_activity_method(
+                    Activities.export_payload_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,
                 {
@@ -95,23 +153,41 @@ class FormatAndExportPrediction():
                     'timestamp': input_data['timestamp'],
                     'model_id': input_data['model_id'],
                     'prediction_confidence': prediction_confidence,
-                    'comment': input_data['comment']
+                    'comment': input_data['comment'],
                 },
                 retry_policy=retry_policy,
-                start_to_close_timeout=timedelta(seconds=60)
+                start_to_close_timeout=timedelta(seconds=60),
+            )
+
+            write_transformed_handler = None
+
+        opc_metrics = {}
+
+        # write to pi web api
+        if pi_web_api_output_config:
+            prediction = await workflow.execute_activity_method(
+                Activities.write_pi_web_api_data,
+                {
+                    'pi_web_api_output_config': pi_web_api_output_config,
+                    'data': prediction,
+                    **metadata,
+                },
+                retry_policy=retry_policy,
+                start_to_close_timeout=timedelta(seconds=60),
             )
 
         # write to opc
-        prediction = await workflow.execute_activity_method(
-            Activities.write_opc_data,
-            {
-                **metadata,
-                'opc_output_config': input_data['opc_output_config'],
-                'data': prediction
-            },
-            retry_policy=retry_policy,
-            start_to_close_timeout=timedelta(seconds=60)
-        )
+        if opc_output_config:
+            prediction, opc_metrics = await workflow.execute_activity_method(
+                Activities.write_opc_data,
+                {
+                    'opc_output_config': opc_output_config,
+                    'data': prediction,
+                    **metadata,
+                },
+                retry_policy=retry_policy,
+                start_to_close_timeout=timedelta(seconds=60),
+            )
 
         # write to postgres
         await workflow.execute_activity_method(
@@ -121,21 +197,24 @@ class FormatAndExportPrediction():
                 'schema': input_data['schema'],
                 'table_name': input_data['table_name'],
                 'data': prediction,
-                'timestamp_conversion': {
-                    'column': 'timestamp',
-                    'format': DATETIME_FORMAT_WITH_TZ
-                }
+                'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ},
+                'on_conflict': input_data.get('on_conflict', 'error'),
+                'unique_columns': ['model_id', 'timestamp'],
             },
             retry_policy=retry_policy,
-            start_to_close_timeout=timedelta(seconds=60)
+            start_to_close_timeout=timedelta(seconds=180),
         )
 
+        if write_transformed_handler is not None:
+            await write_transformed_handler
+
         await workflow.execute_activity_method(
             Activities.write_metrics,
             {
                 **metadata,
-                'prediction': prediction
+                'prediction': prediction,
+                'opc_metrics': opc_metrics,
             },
             retry_policy=retry_policy,
-            start_to_close_timeout=timedelta(seconds=60)
+            start_to_close_timeout=timedelta(seconds=60),
         )

laborious/workflows/sub_workflows/prediction_process.py

@@ -1,14 +1,16 @@
 from temporalio import workflow
 
 with workflow.unsafe.imports_passed_through():
-    from laborious.activities.activities import Activities
-    from typing import Any
-    from sientia_do.temporal.policies import retry_policy
     from datetime import timedelta
+    from typing import Any
+
+    from sientia_do.temporal.policies import retry_policy
+
+    from laborious.activities.activities import Activities
 
 
-@workflow.defn(name="prediction_process")
-class PredictionProcess():
+@workflow.defn(name='subworkflow.prediction_process')
+class PredictionProcess:
     """
     Core prediction processing workflow for the Laborious system.
 
@@ -64,7 +66,10 @@ class PredictionProcess():
                     - mlflow_predict_filters (dict): MLFlow prediction filters
                     - model_retention (int): Model retention period in minutes
                     - path_priority (list[str]): Decision path priority configuration
-                    - opc_output_config (dict): OPC server export configuration
+                    - opc_output_config (dict, optional): OPC server export configuration
+                    - pi_web_api_output_config (dict, optional): PI Web API export configuration
+                    - save_transform (bool, optional): Whether to save transformed data (default: True)
+                    - prediction_store_policy (str, optional): Data retention policy (default: 'lts:1')
 
         Returns:
             None: The workflow completes successfully when export workflow finishes
@@ -79,25 +84,52 @@ class PredictionProcess():
         data = input_data['data']
         model_id = input_data['model_id']
         model_name = input_data['model_name']
-        model_retention = input_data['model_retention']
+        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(
-            Activities.get_last_timestamp,
-            {
-                **metadata,
-                'data': data
-            },
-            retry_policy=retry_policy,
-            start_to_close_timeout=timedelta(minutes=1),
-        )
+        try:
+            await self._run_prediction_pipeline(
+                input_data,
+                metadata,
+                data,
+                model_id,
+                model_name,
+                model_config,
+                save_transform,
+            )
+            await workflow.execute_activity_method(
+                Activities.cleanup_minio_objects_expired,
+                {**metadata, 'data': data},
+                retry_policy=retry_policy,
+                start_to_close_timeout=timedelta(minutes=5),
+            )
+        except Exception as e:
+            await workflow.execute_activity_method(
+                Activities.cleanup_minio_objects_expired,
+                {**metadata, 'data': data},
+                retry_policy=retry_policy,
+                start_to_close_timeout=timedelta(minutes=5),
+            )
+            raise e
+
+    async def _run_prediction_pipeline(
+        self,
+        input_data: dict[str, Any],
+        metadata: dict[str, Any],
+        data: dict[str, Any],
+        model_id: str,
+        model_name: str,
+        model_config: dict[str, Any],
+        save_transform: bool,
+    ) -> None:
+        last_timestamp = data['last_timestamp']
 
         # Apply input data quality gates
         gate_input = {
             **metadata,
             'filters': input_data['input_filters'],
             'data': data,
-            'path_priority': input_data['path_priority']
+            'path_priority': input_data['path_priority'],
         }
 
         path_flag, confidence, comment = await workflow.execute_local_activity_method(
@@ -114,16 +146,11 @@ class PredictionProcess():
             return
 
         # Request MLFlow model transformation
-        response_data = await workflow.execute_local_activity_method(
+        transformed_data = await workflow.execute_activity_method(
             Activities.request_transform,
-            {
-                **metadata,
-                'data': data,
-                'model_name': model_name,
-                'model_retention': model_retention
-            },
+            {**metadata, 'data': data, 'model_name': model_name, 'model_config': model_config},
             retry_policy=retry_policy,
-            start_to_close_timeout=timedelta(minutes=1),
+            start_to_close_timeout=timedelta(minutes=5),
         )
 
         # Validate MLFlow transform response
@@ -132,9 +159,9 @@ class PredictionProcess():
             {
                 **metadata,
                 'filters': input_data['mlflow_transform_filters'],
-                'data': response_data,
+                'data': transformed_data,
                 'type': 'transform',
-                'path_priority': input_data['path_priority']
+                'path_priority': input_data['path_priority'],
             },
             retry_policy=retry_policy,
             start_to_close_timeout=timedelta(minutes=1),
@@ -146,8 +173,6 @@ class PredictionProcess():
         ):
             return
 
-        transformed_data = response_data['content']
-
         path_flag, confidence, comment = await workflow.execute_local_activity_method(
             Activities.mlflow_content_gate,
             {
@@ -155,7 +180,7 @@ class PredictionProcess():
                 'filters': input_data['mlflow_transform_filters'],
                 'data': transformed_data,
                 'type': 'transform',
-                'path_priority': input_data['path_priority']
+                'path_priority': input_data['path_priority'],
             },
             retry_policy=retry_policy,
             start_to_close_timeout=timedelta(minutes=1),
@@ -166,16 +191,16 @@ class PredictionProcess():
         ):
             return
 
-        response_data = await workflow.execute_local_activity_method(
+        predicted_data = await workflow.execute_activity_method(
             Activities.request_predict,
             {
                 **metadata,
                 'data': transformed_data,
                 'model_name': model_name,
-                'model_retention': model_retention
+                'model_config': model_config,
             },
             retry_policy=retry_policy,
-            start_to_close_timeout=timedelta(minutes=1),
+            start_to_close_timeout=timedelta(minutes=5),
         )
 
         # Validate MLFlow prediction response
@@ -184,9 +209,9 @@ class PredictionProcess():
             {
                 **metadata,
                 'filters': input_data['mlflow_predict_filters'],
-                'data': response_data,
+                'data': predicted_data,
                 'type': 'predict',
-                'path_priority': input_data['path_priority']
+                'path_priority': input_data['path_priority'],
             },
             retry_policy=retry_policy,
             start_to_close_timeout=timedelta(minutes=1),
@@ -200,25 +225,37 @@ class PredictionProcess():
 
         # Delegate to export workflow for data persistence
         await workflow.execute_child_workflow(
-            'format_and_export_prediction',
+            'subworkflow.format_and_export_prediction',
             {
                 'metadata': metadata,
+                'on_conflict': input_data.get('on_conflict', 'error'),
                 'path_flag': path_flag,
-                'data': response_data['content'],
+                'data': predicted_data,
+                'transformed_data': transformed_data if save_transform else None,
                 'prediction_confidence': confidence,
                 'timestamp': last_timestamp,
                 'model_id': model_id,
                 'model_name': model_name,
-                'model_retention': model_retention,
+                'model_config': model_config,
                 'opc_output_config': input_data['opc_output_config'],
+                'pi_web_api_output_config': input_data['pi_web_api_output_config'],
                 'schema': input_data['schema'],
                 'table_name': input_data['table_name'],
-                'comment': comment
-            }
+                'transform_table_name': input_data['transform_table_name'],
+                'comment': comment,
+                'prediction_store_policy': input_data['prediction_store_policy'],
+            },
         )
 
-    async def path_flag_handler(self, data: dict, path_flag: str, input_data: dict,
-                                confidence: int, last_timestamp: str, comment: str) -> bool:
+    async def path_flag_handler(
+        self,
+        data: dict[str, Any],
+        path_flag: str,
+        input_data: dict,
+        confidence: int,
+        last_timestamp: str,
+        comment: str,
+    ) -> bool:
         """
         Handle path decisions based on filter results and confidence levels.
 
@@ -229,7 +266,17 @@ class PredictionProcess():
         Args:
             data: Input data for processing
             path_flag: Path decision from filter (STOP, CONTINUE, REPEAT)
-            input_data: Complete workflow input configuration
+            input_data: Complete workflow input configuration including:
+                - metadata (dict): Workflow execution metadata
+                - schema (str): Database schema
+                - table_name (str): Target table for predictions
+                - transform_table_name (str): Target table for transformed data
+                - model_id (str): ML model identifier
+                - model_name (str): ML model name
+                - model_config (dict, optional): Model configuration
+                - opc_output_config (dict, optional): OPC server export configuration
+                - pi_web_api_output_config (dict, optional): PI Web API export configuration
+                - prediction_store_policy (str, optional): Data retention policy
             confidence: Confidence level from filter validation
             last_timestamp: Last processed timestamp
             comment: Additional information about the filter result
@@ -239,16 +286,17 @@ class PredictionProcess():
 
         Path Handling:
             - STOP: Terminates workflow execution
-            - CONTINUE: Proceeds with normal processing
+            - CONTINUE: Delegates to FormatAndExportPrediction workflow with current data
             - REPEAT: Repeats last prediction if available
         """
         metadata = input_data['metadata']
 
         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_retention = input_data['model_retention']
+        model_config = input_data.get('model_config', {})
 
         path_flag = path_flag.upper() if path_flag else ''
 
@@ -264,7 +312,7 @@ class PredictionProcess():
                     'schema': schema,
                     'table_name': table_name,
                     'model': model_id,
-                    'last_timestamp': last_timestamp
+                    'last_timestamp': last_timestamp,
                 },
                 retry_policy=retry_policy,
                 start_to_close_timeout=timedelta(minutes=1),
@@ -273,7 +321,7 @@ class PredictionProcess():
         elif path_flag == 'CONTINUE':
             # call write workflow
             await workflow.execute_child_workflow(
-                'format_and_export_prediction',
+                'subworkflow.format_and_export_prediction',
                 {
                     'metadata': metadata,
                     'path_flag': path_flag,
@@ -282,12 +330,16 @@ class PredictionProcess():
                     'timestamp': last_timestamp,
                     'model_id': model_id,
                     'model_name': model_name,
-                    'model_retention': model_retention,
+                    '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']
-                }
+                    'opc_output_config': input_data['opc_output_config'],
+                    'pi_web_api_output_config': input_data['pi_web_api_output_config'],
+                    'prediction_store_policy': input_data['prediction_store_policy'],
+                    'on_conflict': input_data.get('on_conflict', 'error'),
+                },
             )
             return True
 

model_convert.ipynb

@@ -0,0 +1,106 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 23,
+   "id": "e838ff21",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import csv\n",
+    "\n",
+    "def csv_to_tag_lists(csv_path: str) -> dict:\n",
+    "    read_tags = []\n",
+    "    write_tags = []\n",
+    "\n",
+    "    def to_float(val):\n",
+    "        try:\n",
+    "            return float(str(val).strip())\n",
+    "        except Exception:\n",
+    "            return None\n",
+    "\n",
+    "    with open(csv_path, newline=\"\", encoding=\"utf-8\") as f:\n",
+    "        reader = csv.DictReader(f)\n",
+    "        for row in reader:\n",
+    "            # Basic normalization\n",
+    "            op = (row.get(\"operation\") or \"\").strip()\n",
+    "\n",
+    "            if op == \"READ\":\n",
+    "                # Build common tag payload with required mappings\n",
+    "                tag = {\n",
+    "                    \"server_id\": \"1\",\n",
+    "                    \"tag_address\": row.get(\"opc_tag\"),\n",
+    "                    \"tag_name\": row.get(\"name\"),\n",
+    "                    \"data_range\": [to_float(row.get(\"min_value\")), to_float(row.get(\"max_value\"))],\n",
+    "                    \"aggr_func\": row.get(\"aggregation_func\").lower(),\n",
+    "                    # keep other fields with their original names\n",
+    "                    \"frequency\": row.get(\"frequency\"),\n",
+    "                    \"local\": row.get(\"local\"),\n",
+    "                    \"area\": row.get(\"area\"),\n",
+    "                    \"description\": row.get(\"description\"),\n",
+    "                }\n",
+    "\n",
+    "                read_tags.append(tag)\n",
+    "\n",
+    "            else:\n",
+    "                tag = {\n",
+    "                    \"server_id\": \"1\",\n",
+    "                    \"addr\": row.get(\"opc_tag\"),\n",
+    "                    \"tag_name\": row.get(\"name\"),\n",
+    "                    \"local\": row.get(\"local\"),\n",
+    "                    \"area\": row.get(\"area\"),\n",
+    "                    \"description\": row.get(\"description\"),\n",
+    "                }\n",
+    "                \n",
+    "            if op == \"WRITE_PREDICTION\":\n",
+    "                tag[\"type\"] = \"prediction\"\n",
+    "                write_tags.append(tag)\n",
+    "            elif op == \"WRITE_CONFIDENCE\":\n",
+    "                tag[\"type\"] = \"confidence\"\n",
+    "                write_tags.append(tag)\n",
+    "            # ignore any other operation values silently\n",
+    "\n",
+    "    return {\"read_tags\": read_tags, \"write_tags\": write_tags}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4621cd43",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import json\n",
+    "\n",
+    "file_names = [\"Courier - Página1.csv\"]\n",
+    "\n",
+    "for file_name in file_names:\n",
+    "    write_file = file_name.replace(\".csv\", \".json\")\n",
+    "\n",
+    "    with open(write_file, \"w\", encoding=\"utf-8\") as f:\n",
+    "        json.dump(csv_to_tag_lists(file_name), f, indent=2, ensure_ascii=False)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

pyproject.toml

@@ -0,0 +1,164 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "laborious"
+version = "0.0.0"
+description = "Sientia DataOps Laborious - ML Model Orchestration System"
+readme = "README.md"
+requires-python = ">=3.11"
+authors = [
+    {name = "Aignosi", email = "dev@aignosi.com"}
+]
+
+[tool.ruff]
+line-length = 100
+target-version = "py311"
+exclude = [
+    ".git",
+    ".venv",
+    "venv",
+    "__pycache__",
+    "*.pyc",
+    ".pytest_cache",
+    "htmlcov",
+    "tests/laborious/workflows/subworkflows/test_prediction_process.py",
+]
+
+[tool.ruff.lint]
+select = [
+    "E",   # pycodestyle errors
+    "W",   # pycodestyle warnings
+    "F",   # pyflakes
+    "I",   # isort
+    "B",   # flake8-bugbear
+    "C4",  # flake8-comprehensions
+    "UP",  # pyupgrade
+    "N",   # pep8-naming
+    "YTT", # flake8-2020
+    "S",   # flake8-bandit
+    "BLE", # flake8-blind-except
+    "A",   # flake8-builtins
+    "C90", # mccabe complexity
+]
+
+ignore = [
+    "BLE001",  # ignore blind except, we need to send notifications with any error
+    "E501",  # line too long (handled by formatter)
+    "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
+]
+
+[tool.ruff.lint.per-file-ignores]
+"tests/**/*.py" = [
+    "S101",  # assert allowed in tests
+    "S105",  # hardcoded passwords ok in tests
+    "S106",  # hardcoded passwords ok in tests
+]
+
+[tool.ruff.lint.mccabe]
+max-complexity = 15
+
+[tool.ruff.format]
+quote-style = "single"
+indent-style = "space"
+line-ending = "auto"
+
+[tool.mypy]
+python_version = "3.11"
+warn_return_any = false
+warn_unused_configs = true
+disallow_untyped_defs = false
+disallow_incomplete_defs = false
+check_untyped_defs = true
+no_implicit_optional = true
+warn_redundant_casts = true
+warn_unused_ignores = false
+warn_no_return = true
+strict_equality = true
+ignore_missing_imports = true
+
+# Ignore missing imports for external packages
+[[tool.mypy.overrides]]
+module = "temporalio.*"
+ignore_missing_imports = true
+
+[[tool.mypy.overrides]]
+module = "sientia_do.*"
+ignore_missing_imports = true
+
+[[tool.mypy.overrides]]
+module = "mlflow.*"
+ignore_missing_imports = true
+
+[[tool.mypy.overrides]]
+module = "prometheus_client.*"
+ignore_missing_imports = true
+
+[[tool.mypy.overrides]]
+module = "sientia.*"
+ignore_missing_imports = true
+
+[[tool.mypy.overrides]]
+module = "pandas.*"
+ignore_missing_imports = true
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+python_files = ["test_*.py"]
+python_classes = ["Test*"]
+python_functions = ["test_*"]
+addopts = [
+    "-v",
+    "--strict-markers",
+    # pytest>=9.1 has a known bug where its unraisableexception plugin crashes
+    # (tracemalloc partially-initialized AttributeError) when 2+ unraisable
+    # exceptions land close together — e.g. "coroutine was never awaited" from
+    # AsyncMock-mocked sync methods (metrics_controller, minio_repository) being
+    # GC'd. Harmless mock artifacts turned into a hard ERROR by the plugin itself.
+    "-p", "no:unraisableexception",
+]
+markers = [
+    "asyncio: marks tests as async",
+    "integration: marks tests as integration tests",
+    "unit: marks tests as unit tests",
+    "opc: marks tests that use the in-process OPC UA server (OpcRepository E2E)",
+]
+
+[tool.coverage.run]
+source = ["laborious"]
+omit = [
+    "*/tests/*",
+    "*/venv/*",
+    "*/__pycache__/*",
+    "*/site-packages/*",
+]
+branch = true
+
+[tool.coverage.report]
+precision = 2
+show_missing = true
+skip_covered = false
+exclude_lines = [
+    "pragma: no cover",
+    "def __repr__",
+    "def __str__",
+    "raise AssertionError",
+    "raise NotImplementedError",
+    "if __name__ == .__main__.:",
+    "if TYPE_CHECKING:",
+    "class .*\\bProtocol\\):",
+    "@(abc\\.)?abstractmethod",
+]
+
+[tool.coverage.html]
+directory = "htmlcov"
+
+[tool.bandit]
+exclude_dirs = ["tests", "venv", ".venv"]
+skips = ["B101", "B601", "B608"]  # Skip assert, shell injection, and SQL injection (false positives)

requirements-dev.txt

@@ -0,0 +1,21 @@
+# Development and Testing Dependencies
+# These packages are only needed for development, testing, and code quality checks
+# Install with: pip install -r requirements-dev.txt
+
+# Code Quality & Linting
+ruff>=0.1.0              # Fast Python linter and formatter (replaces flake8, black, isort)
+mypy>=1.7.0              # Static type checker
+bandit>=1.7.5            # Security vulnerability scanner
+pandas-stubs>=2.0.0      # Type stubs for pandas
+types-requests>=2.31.0   # Type stubs for requests
+
+# Testing
+pytest>=7.4.0            # Testing framework
+pytest-cov>=4.1.0        # Coverage plugin for pytest
+pytest-asyncio>=0.21.0   # Async test support (already in main requirements)
+testcontainers[postgres,minio] # PostgreSQL and MinIO containers for E2E tests
+
+# Development Tools
+ipython>=8.12.0          # Enhanced Python shell
+ipdb>=0.13.13            # IPython debugger
+ipykernel==6.30.1        # IPython kernel for Jupyter notebooks

requirements-light.txt

@@ -0,0 +1,18 @@
+temporalio
+psycopg2-binary
+sqlalchemy
+asyncua==1.0.6
+redis
+sientia_do>=1.12.2
+mlflow
+prometheus-client
+botocore
+boto3
+s3fs
+pyarrow
+kaleido
+hyperopt
+shap
+pycurl
+scipy<1.14.0
+scikit-learn==1.5.2

requirements-local.txt

@@ -0,0 +1,18 @@
+temporalio
+psycopg2-binary
+sqlalchemy
+asyncua==1.0.6
+redis
+git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.12.2
+git+ssh://git@github.com/Aignosi/sientia-model-library.git@0.10.0
+prometheus-client   
+botocore
+boto3
+s3fs
+pyarrow
+kaleido
+hyperopt
+shap
+pycurl
+scipy<1.14.0
+scikit-learn==1.5.2

requirements.txt

@@ -1,8 +1,18 @@
 temporalio
 psycopg2-binary
 sqlalchemy
-asyncua
+asyncua==1.0.6
 redis
-git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.4.4
-git+ssh://git@github.com/Aignosi/sientia-mlops-library.git@0.38.12
+sientia_do>=1.12.2
+sientia>0.40.0
 prometheus-client
+botocore
+boto3
+s3fs
+pyarrow
+kaleido
+hyperopt
+shap
+pycurl
+scipy<1.14.0
+scikit-learn==1.5.2

requirements_prepared.txt

@@ -1,8 +0,0 @@
-temporalio
-psycopg2-binary
-sqlalchemy
-asyncua
-redis
-git+https://github.com/Aignosi/sientia-dataops-library.git@1.4.4
-git+https://github.com/Aignosi/sientia-mlops-library.git@0.38.12
-prometheus-client

simulator/Dockerfile

@@ -1,30 +0,0 @@
-# syntax=docker/dockerfile:1.4
-
-FROM python:3.11-slim
-
-# Enable use of SSH agent/socket
-# This line enables SSH during build
-# (don't forget the syntax header above)
-RUN apt-get update && apt-get install -y git openssh-client && rm -rf /var/lib/apt/lists/*
-
-# Use build-time SSH mount for Git clone
-# The SSH key will NOT remain in the image
-# IMPORTANT: this block requires BuildKit
-# and the --ssh flag during docker build
-
-# SSH config to skip host key check (safe in CI/local dev)
-RUN mkdir -p /root/.ssh && echo "StrictHostKeyChecking no" > /root/.ssh/config
-
-WORKDIR /app
-
-# Clone using SSH
-ARG GIT_REPO
-ARG GIT_BRANCH=main
-
-# Mount SSH key just for this RUN
-RUN --mount=type=ssh git clone --branch ${GIT_BRANCH} ${GIT_REPO} .
-
-# Install requirements if exists
-RUN if [ -f requirements.txt ]; then pip install --no-cache-dir -r requirements.txt; fi
-
-CMD ["python", "server.py"]

sonar-project.properties

@@ -1,9 +1,9 @@
-sonar.projectKey=Aignosi_sientia-dataops-laborious_temporal_beaec423-6c42-4f26-8134-b676287b499d
+sonar.projectKey=Aignosi_sientia-dataops-laborious_temporal_ca1a7039-6db9-49e5-be78-54d29bc93e4f
 sonar.projectName=sientia-dataops-laborious_temporal
 sonar.sources=laborious
 sonar.tests=tests
 sonar.projectVersion=1.0.0
-sonar.coverage.exclusions=laborious/worker/worker.py
+sonar.coverage.exclusions=laborious/worker/*
 sonar.qualitygate.wait=true
 sonar.qualitygate.timeout=300
 sonar.python.coverage.reportPaths=coverage.xml

tests/conftest.py

@@ -0,0 +1,59 @@
+import os
+import sys
+from unittest.mock import MagicMock
+
+# The production code converts SIENTIA_MINIO_OFFLOAD_THRESHOLD_MEGABYTES to int at import-time.
+# Tests must set it to a valid integer string to avoid import errors.
+os.environ.setdefault('SIENTIA_MINIO_OFFLOAD_THRESHOLD_MEGABYTES', '1')
+
+
+class DummyMinioDataFramePayload:
+    """
+    Minimal payload double used by unit tests.
+
+    The production workflow/gates expect a MinioDataFramePayload-like object with:
+    - async retrieve(minio_repo, workflow_metadata) -> DataFrame | dict
+    - has_data() -> bool
+    - cleanup_prefix() -> str | None
+    - last_timestamp: attribute
+    - status: attribute
+    """
+
+    def __init__(
+        self,
+        *,
+        retrieve_return=None,
+        has_data: bool = True,
+        cleanup_prefix: str | None = None,
+        last_timestamp: str = '2024-01-01',
+        status: dict | None = None,
+    ):
+        self._retrieve_return = retrieve_return
+        self._has_data = has_data
+        self._cleanup_prefix = cleanup_prefix
+        self.last_timestamp = last_timestamp
+        self.status = status
+
+    async def retrieve(self, _minio_repo, _workflow_metadata=None):
+        return self._retrieve_return
+
+    def has_data(self) -> bool:
+        return self._has_data
+
+    def cleanup_prefix(self) -> str | None:
+        return self._cleanup_prefix
+
+
+"""
+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.
+"""
+
+# Mock sientia module
+sientia_mock = MagicMock()
+sientia_mock.ModelAnalysis = MagicMock
+sys.modules['sientia'] = sientia_mock
+sys.modules['sientia.ModelAnalysis'] = MagicMock()

tests/laborious/activities/test_activities.py

@@ -1,18 +1,34 @@
+from unittest.mock import ANY, AsyncMock, MagicMock, patch
+
 from pytest import mark
-from unittest.mock import patch, MagicMock, ANY
-from sientia_do.temporal.activities.postgres import Postgres
+
 from laborious.activities.activities import Activities
-from laborious.activities.mlflow import MLFlow
+from laborious.activities.api import API
 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
 
 
-@patch('laborious.activities.activities.Postgres.__init__')
+@patch('laborious.activities.activities.Storage.__init__')
 @patch('laborious.activities.activities.MLFlow.__init__')
 @patch('laborious.activities.activities.OPC.__init__')
 @patch('laborious.activities.activities.Gates.__init__')
-def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgres_init):
-
+@patch('laborious.activities.activities.ModelMetrics.__init__')
+@patch('laborious.activities.activities.API.__init__')
+@patch('laborious.activities.activities.MinioRepository')
+@patch('laborious.activities.activities.MetricsController')
+def test___init__(
+    mock_metrics_controller,
+    mock_minio_repository,
+    mock_api_init,
+    mock_model_metrics_init,
+    mock_gates_init,
+    mock_opc_init,
+    mock_mlflow_init,
+    mock_storage_init,
+):
     postgres_config = {
         'host': 'localhost',
         'port': 5432,
@@ -20,20 +36,30 @@ def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgre
         'password': 'postgres',
         'dbname': 'postgres',
         'min_connections': 1,
-        'max_connections': 10
+        'max_connections': 10,
     }
 
-    mlflow_config = {
-        'host': 'localhost',
-        'port': 5000,
-        'username': 'mlflow',
-        'password': 'mlflow'
+    minio_config = {
+        'endpoint_url': 'localhost:9000',
+        'access_key': 'minio',
+        'secret_key': 'minio123',
+        'default_bucket': 'test',
+        'retention_hours': 24,
+        'secure': False,
     }
 
+    mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
+
     opc_config = {
         'bootstrap_servers': 'localhost:9092',
         'polling_time': 1000,
-        'group_id': 'test-group'
+        'group_id': 'test-group',
+    }
+
+    pi_web_api_config = {
+        'base_url': 'https://test-pi-server.com',
+        'auth_type': 'bearer',
+        'auth_token': 'test_token',
     }
 
     logger = MagicMock()
@@ -42,18 +68,22 @@ def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgre
     activities = Activities(
         postgres_config=postgres_config,
         mlflow_config=mlflow_config,
+        minio_config=minio_config,
         opc_config=opc_config,
+        pi_web_api_config=pi_web_api_config,
         logger=logger,
-        notification_handler=notification_handler
+        notification_handler=notification_handler,
     )
 
     assert isinstance(activities, Activities)
-    assert isinstance(activities, Postgres)
+    assert isinstance(activities, Storage)
     assert isinstance(activities, MLFlow)
     assert isinstance(activities, OPC)
     assert isinstance(activities, Gates)
+    assert isinstance(activities, ModelMetrics)
+    assert isinstance(activities, API)
 
-    mock_postgres_init.assert_called_once_with(
+    mock_storage_init.assert_called_once_with(
         ANY,
         host=postgres_config['host'],
         port=postgres_config['port'],
@@ -62,8 +92,11 @@ def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgre
         dbname=postgres_config['dbname'],
         min_connections=postgres_config['min_connections'],
         max_connections=postgres_config['max_connections'],
+        retention_hours=minio_config['retention_hours'],
+        minio_repository=mock_minio_repository.return_value,
         logger=logger,
-        notification_handler=notification_handler
+        notification_handler=notification_handler,
+        metrics_controller=mock_metrics_controller.return_value,
     )
 
     mock_mlflow_init.assert_called_once_with(
@@ -72,30 +105,75 @@ def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgre
         mlflow_port=mlflow_config['port'],
         mlflow_username=mlflow_config['username'],
         mlflow_password=mlflow_config['password'],
+        minio_repository=mock_minio_repository.return_value,
         logger=logger,
-        notification_handler=notification_handler
+        notification_handler=notification_handler,
+        metrics_controller=mock_metrics_controller.return_value,
     )
 
     mock_opc_init.assert_called_once_with(
         ANY,
         opc_servers=opc_config,
         logger=logger,
-        notification_handler=notification_handler
+        notification_handler=notification_handler,
+        metrics_controller=mock_metrics_controller.return_value,
     )
 
     mock_gates_init.assert_called_once_with(
         ANY,
+        minio_repository=mock_minio_repository.return_value,
         logger=logger,
-        notification_handler=notification_handler
+        notification_handler=notification_handler,
+        metrics_controller=mock_metrics_controller.return_value,
+    )
+
+    mock_model_metrics_init.assert_called_once_with(
+        ANY,
+        logger=logger,
+        notification_handler=notification_handler,
+        metrics_controller=mock_metrics_controller.return_value,
+    )
+
+    mock_api_init.assert_called_once_with(
+        ANY,
+        base_url=pi_web_api_config['base_url'],
+        auth_type=pi_web_api_config['auth_type'],
+        auth_token=pi_web_api_config['auth_token'],
+        logger=logger,
+        notification_handler=notification_handler,
+        metrics_controller=mock_metrics_controller.return_value,
+    )
+
+    mock_minio_repository.assert_called_once_with(
+        endpoint=minio_config['endpoint_url'],
+        access_key=minio_config['access_key'],
+        secret_key=minio_config['secret_key'],
+        bucket=minio_config['default_bucket'],
+        logger=logger,
+        notification_handler=notification_handler,
+        metrics_controller=mock_metrics_controller.return_value,
+        secure=minio_config['secure'],
     )
 
 
 @mark.asyncio
-@patch('laborious.activities.activities.Postgres', return_value=MagicMock())
-@patch('laborious.activities.activities.MLFlow', return_value=MagicMock())
-@patch('laborious.activities.activities.OPC', return_value=MagicMock())
-async def test_shutdown(mock_opc_init,
-                        _mock_mlflow_init, mock_postgres_init):
+@patch('laborious.activities.activities.Storage')
+@patch('laborious.activities.activities.MLFlow')
+@patch('laborious.activities.activities.OPC')
+@patch('laborious.activities.activities.Gates')
+@patch('laborious.activities.activities.ModelMetrics')
+@patch('laborious.activities.activities.API')
+@patch('laborious.activities.activities.MinioRepository')
+async def test_shutdown(
+    _mock_minio_repository,
+    mock_api_init,
+    mock_model_metrics_init,
+    mock_gates_init,
+    mock_opc_init,
+    mock_mlflow_init,
+    mock_storage_init,
+):
+    mock_opc_init.aclose = AsyncMock()
     postgres_config = {
         'host': 'localhost',
         'port': 5432,
@@ -103,20 +181,30 @@ async def test_shutdown(mock_opc_init,
         'password': 'postgres',
         'dbname': 'postgres',
         'min_connections': 1,
-        'max_connections': 10
+        'max_connections': 10,
     }
 
-    mlflow_config = {
-        'host': 'localhost',
-        'port': 5000,
-        'username': 'mlflow',
-        'password': 'mlflow'
+    minio_config = {
+        'endpoint_url': 'localhost:9000',
+        'access_key': 'minio',
+        'secret_key': 'minio123',
+        'default_bucket': 'test',
+        'retention_hours': 24,
+        'secure': False,
     }
 
+    mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
+
     opc_config = {
         'bootstrap_servers': 'localhost:9092',
         'polling_time': 1000,
-        'group_id': 'test-group'
+        'group_id': 'test-group',
+    }
+
+    pi_web_api_config = {
+        'base_url': 'https://test-pi-server.com',
+        'auth_type': 'bearer',
+        'auth_token': 'test_token',
     }
 
     logger = MagicMock()
@@ -125,11 +213,17 @@ async def test_shutdown(mock_opc_init,
     activities = Activities(
         postgres_config=postgres_config,
         mlflow_config=mlflow_config,
+        minio_config=minio_config,
         opc_config=opc_config,
+        pi_web_api_config=pi_web_api_config,
         logger=logger,
-        notification_handler=notification_handler
+        notification_handler=notification_handler,
     )
 
     await activities.shutdown()
-    mock_opc_init.shutdown.assert_called_once()
-    mock_postgres_init.close.assert_called_once()
+    mock_opc_init.aclose.assert_called_once()
+    mock_storage_init.close.assert_called_once()
+    mock_mlflow_init.close.assert_called_once()
+    mock_gates_init.close.assert_called_once()
+    mock_model_metrics_init.close.assert_called_once()
+    mock_api_init.close.assert_called_once()

tests/laborious/activities/test_api.py

@@ -0,0 +1,492 @@
+from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
+
+import pytest_asyncio
+from pytest import fixture, mark
+from sientia_do.notifications.models import NotificationLevel
+
+from laborious.activities.api import API, PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
+
+metadata = {
+    'metadata': {
+        'model_id': 'test_model',
+        'model_name': 'test_model',
+        'workflow_name': 'test_workflow',
+        'schema_name': 'test_schedule',
+    },
+}
+
+
+def _create_mock_dataframe(to_dict_return=None):
+    """Helper function to create a mocked DataFrame for testing."""
+    mock_df = MagicMock()
+    mock_head = MagicMock()
+
+    def get_column_values(key):
+        if key == 'prediction':
+            return MagicMock(values=[0.75])
+        elif key == 'prediction_confidence':
+            return MagicMock(values=[0.95])
+        else:
+            return MagicMock(values=['2024-01-01T00:00:00+00:00'])
+
+    mock_head.__getitem__.side_effect = get_column_values
+    mock_df.head.return_value = mock_head
+
+    if to_dict_return is None:
+        to_dict_return = {
+            'prediction': [0.75],
+            'prediction_confidence': [0.95],
+            'timestamp': ['2024-01-01T00:00:00+00:00'],
+        }
+    mock_df.to_dict.return_value = to_dict_return
+
+    return mock_df
+
+
+@fixture
+def base_input_data():
+    """Base input data for PI Web API tests."""
+    return {
+        **metadata,
+        'data': {
+            'prediction': [0.75],
+            'prediction_confidence': [0.95],
+            'timestamp': ['2024-01-01T00:00:00+00:00'],
+        },
+        'pi_web_api_output_config': {
+            'endpoint': 'https://test-pi-server.com/piwebapi',
+            'prediction_tags': {'tag1': 'web_id_1'},
+            'confidence_tags': {'tag2': 'web_id_2'},
+        },
+    }
+
+
+@patch('laborious.activities.api.PIWebAPIClient')
+def test_get_pi_web_api_core_labels_without_operation_type(mock_pi_web_api_client):
+    from sientia_do.observability.sientia_monitoring import SientiaMonitoring
+
+    api_instance = API(
+        base_url='https://test-pi-server.com',
+        auth_type='bearer',
+        auth_token='test_token',
+        logger=MagicMock(),
+        notification_handler=MagicMock(),
+        metrics_controller=AsyncMock(),
+    )
+    with patch.object(
+        SientiaMonitoring,
+        'get_core_labels',
+        return_value={
+            'pod_id': 'test_pod',
+            'runtime': 'local',
+            'model_name': 'test_model',
+            'workflow_name': 'test_workflow',
+            'operation_type': 'write_pi_web_api_data',
+        },
+    ):
+        labels = api_instance.get_pi_web_api_core_labels(metadata=metadata['metadata'])
+    assert labels['operation_type'] == 'write_pi_web_api_data'
+    assert labels == {
+        'pod_id': 'test_pod',
+        'runtime': 'local',
+        'model_name': 'test_model',
+        'workflow_name': 'test_workflow',
+        'operation_type': 'write_pi_web_api_data',
+    }
+
+
+@patch('laborious.activities.api.PIWebAPIClient')
+def test_get_pi_web_api_core_labels_with_operation_type(mock_pi_web_api_client):
+    from sientia_do.observability.sientia_monitoring import SientiaMonitoring
+
+    api_instance = API(
+        base_url='https://test-pi-server.com',
+        auth_type='bearer',
+        auth_token='test_token',
+        logger=MagicMock(),
+        notification_handler=MagicMock(),
+        metrics_controller=AsyncMock(),
+    )
+    with patch.object(
+        SientiaMonitoring,
+        'get_core_labels',
+        return_value={
+            'pod_id': 'test_pod',
+            'runtime': 'k8s',
+            'model_name': 'test_model',
+            'workflow_name': 'test_workflow',
+            'operation_type': 'write',
+        },
+    ):
+        labels = api_instance.get_pi_web_api_core_labels(
+            metadata=metadata['metadata'], operation_type='write'
+        )
+    assert labels['operation_type'] == 'write'
+    assert labels['runtime'] == 'k8s'
+
+
+def test__init__():
+    api = API(
+        base_url='https://test-pi-server.com',
+        auth_type='bearer',
+        auth_token='test_token',
+        logger=MagicMock(),
+        notification_handler=MagicMock(),
+        metrics_controller=AsyncMock(),
+    )
+
+    assert api.pi_web_api_client is not None
+
+
+@pytest_asyncio.fixture
+@patch('laborious.activities.api.PIWebAPIClient')
+def api(mock_pi_web_api_client):
+    mock_client = MagicMock()
+    mock_client.write_value = AsyncMock()
+    mock_client.close = MagicMock()
+    mock_client.base_url = 'https://test-pi-server.com'
+    mock_pi_web_api_client.return_value = mock_client
+
+    api_instance = API(
+        base_url='https://test-pi-server.com',
+        auth_type='bearer',
+        auth_token='test_token',
+        logger=MagicMock(),
+        notification_handler=MagicMock(),
+        metrics_controller=AsyncMock(),
+    )
+    api_instance.send_notification_async = AsyncMock()
+    api_instance.info = MagicMock()
+    api_instance.error = MagicMock()
+    api_instance.emit_metric = AsyncMock()
+    api_instance.get_core_labels = MagicMock(
+        return_value={
+            'pod_id': 'test_pod',
+            'runtime': 'local',
+            'model_name': 'test_model',
+            'workflow_name': 'test_workflow',
+        }
+    )
+    return api_instance
+
+
+@mark.asyncio
+@patch('laborious.activities.api.DataFrame')
+async def test_write_pi_web_api_data_success(mock_dataframe, api, base_input_data):
+    input_data = {
+        **base_input_data,
+        'pi_web_api_output_config': {
+            'endpoint': 'https://test-pi-server.com/piwebapi',
+            'prediction_tags': {'tag1': 'web_id_1', 'tag2': 'web_id_2'},
+            'confidence_tags': {'tag3': 'web_id_3', 'tag4': 'web_id_4'},
+        },
+    }
+
+    mock_dataframe.return_value = _create_mock_dataframe()
+
+    # Mock successful responses
+    api.pi_web_api_client.write_value.side_effect = [
+        [{'WebId': 'web_id_1', 'Errors': []}, {'WebId': 'web_id_2', 'Errors': []}],
+        [{'WebId': 'web_id_3', 'Errors': []}, {'WebId': 'web_id_4', 'Errors': []}],
+    ]
+
+    result = await api.write_pi_web_api_data(input_data)
+
+    api.pi_web_api_client.write_value.assert_has_calls(
+        [
+            call(
+                web_ids=['web_id_1', 'web_id_2'],
+                value={
+                    'Timestamp': '2024-01-01T00:00:00+00:00',
+                    'Value': 0.75,
+                },
+                metadata=metadata['metadata'],
+            ),
+            call(
+                web_ids=['web_id_3', 'web_id_4'],
+                value={
+                    'Timestamp': '2024-01-01T00:00:00+00:00',
+                    'Value': 0.95,
+                },
+                metadata=metadata['metadata'],
+            ),
+        ]
+    )
+
+    assert result == {
+        'prediction': [0.75],
+        'prediction_confidence': [0.95],
+        'timestamp': ['2024-01-01T00:00:00+00:00'],
+    }
+
+
+@mark.asyncio
+@patch('laborious.activities.api.DataFrame')
+async def test_write_pi_web_api_data_prediction_error(mock_dataframe, api, base_input_data):
+    mock_dataframe.return_value = _create_mock_dataframe(
+        {
+            'prediction': [0.75],
+            'prediction_confidence': [PI_WEB_API_PREDICTION_ERROR_CONFIDENCE],
+            'timestamp': ['2024-01-01T00:00:00+00:00'],
+        }
+    )
+
+    api.pi_web_api_client.write_value.side_effect = Exception('Prediction write failed')
+
+    result = await api.write_pi_web_api_data(base_input_data)
+
+    api.send_notification_async.assert_called_once_with(
+        metadata=metadata['metadata'],
+        notification_id='WRITE_PI_WEB_API_PREDICTION_ERROR',
+        message="Error writing prediction data to PI Web API: Prediction write failed\n Tags: {'tag1': 'web_id_1'}",
+        block='write_pi_web_api_data',
+        level=NotificationLevel.ERROR,
+        attachment_content=ANY,
+    )
+
+    assert result['prediction_confidence'][0] == PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
+    assert api.pi_web_api_client.write_value.call_count == 1
+
+
+@mark.asyncio
+@patch('laborious.activities.api.DataFrame')
+async def test_write_pi_web_api_data_confidence_error(mock_dataframe, api, base_input_data):
+    mock_dataframe.return_value = _create_mock_dataframe()
+
+    # First call succeeds, second fails
+    api.pi_web_api_client.write_value.side_effect = [
+        [{'WebId': 'web_id_1', 'Errors': []}],
+        Exception('Confidence write failed'),
+    ]
+
+    result = await api.write_pi_web_api_data(base_input_data)
+
+    api.send_notification_async.assert_called_once_with(
+        metadata=metadata['metadata'],
+        notification_id='WRITE_PI_WEB_API_CONFIDENCE_ERROR',
+        message="Error writing confidence data to PI Web API: Confidence write failed\n Tags: {'tag2': 'web_id_2'}",
+        block='write_pi_web_api_data',
+        level=NotificationLevel.ERROR,
+        attachment_content=ANY,
+    )
+
+    assert result == {
+        'prediction': [0.75],
+        'prediction_confidence': [0.95],
+        'timestamp': ['2024-01-01T00:00:00+00:00'],
+    }
+    assert api.pi_web_api_client.write_value.call_count == 2
+
+
+@mark.asyncio
+@patch('laborious.activities.api.DataFrame')
+async def test_write_pi_web_api_data_empty_tags(mock_dataframe, api, base_input_data):
+    input_data = {
+        **base_input_data,
+        'pi_web_api_output_config': {
+            'endpoint': 'https://test-pi-server.com/piwebapi',
+            'prediction_tags': {},
+            'confidence_tags': {},
+        },
+    }
+
+    mock_dataframe.return_value = _create_mock_dataframe()
+
+    # Mock empty responses
+    api.pi_web_api_client.write_value.side_effect = [
+        [],
+        [],
+    ]
+
+    result = await api.write_pi_web_api_data(input_data)
+
+    api.pi_web_api_client.write_value.assert_has_calls(
+        [
+            call(
+                web_ids=[],
+                value={
+                    'Timestamp': '2024-01-01T00:00:00+00:00',
+                    'Value': 0.75,
+                },
+                metadata=metadata['metadata'],
+            ),
+            call(
+                web_ids=[],
+                value={
+                    'Timestamp': '2024-01-01T00:00:00+00:00',
+                    'Value': 0.95,
+                },
+                metadata=metadata['metadata'],
+            ),
+        ]
+    )
+
+    assert result == {
+        'prediction': [0.75],
+        'prediction_confidence': [0.95],
+        'timestamp': ['2024-01-01T00:00:00+00:00'],
+    }
+
+
+@mark.asyncio
+async def test_close(api):
+    api.close()
+
+    api.pi_web_api_client.close.assert_called_once()
+
+
+@mark.asyncio
+async def test_process_pi_web_api_response_success(api):
+    """Test successful processing of PI Web API response with all tags written."""
+    response_data = [
+        {'WebId': 'web_id_1', 'Errors': []},
+        {'WebId': 'web_id_2', 'Errors': []},
+    ]
+    tags = {'tag1': 'web_id_1', 'tag2': 'web_id_2'}
+    core_labels = {
+        'pod_id': 'test_pod',
+        'runtime': 'local',
+        'model_name': 'test_model',
+        'workflow_name': 'test_workflow',
+    }
+
+    confidence, message = await api.process_pi_web_api_response(
+        response_data=response_data,
+        tags=tags,
+        core_labels=core_labels,
+        metadata=metadata['metadata'],
+    )
+
+    assert confidence == 0
+    assert message == ''
+    assert api.emit_metric.call_count == 2
+    # Verify that emit_metric was called with correct tags structure
+    call_args_list = api.emit_metric.call_args_list
+    assert len(call_args_list) == 2
+    # Check that all calls include core_labels and tag_name
+    for call_args in call_args_list:
+        assert 'tag_name' in call_args.kwargs['tags']
+        assert call_args.kwargs['tags']['tag_name'] in ['tag1', 'tag2']
+
+
+@mark.asyncio
+async def test_process_pi_web_api_response_with_errors(api):
+    """Test processing response with errors in some tags."""
+    response_data = [
+        {'WebId': 'web_id_1', 'Errors': ['Error writing tag']},
+        {'WebId': 'web_id_2', 'Errors': []},
+    ]
+    tags = {'tag1': 'web_id_1', 'tag2': 'web_id_2'}
+    core_labels = {
+        'pod_id': 'test_pod',
+        'runtime': 'local',
+        'model_name': 'test_model',
+        'workflow_name': 'test_workflow',
+    }
+
+    confidence, message = await api.process_pi_web_api_response(
+        response_data=response_data,
+        tags=tags,
+        core_labels=core_labels,
+        metadata=metadata['metadata'],
+    )
+
+    assert confidence == PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
+    assert (
+        message
+        == "The number of written tags does not match the number of tag names: Expected ['tag1', 'tag2'] tags, but ['tag2'] tags were written."
+    )
+    assert api.emit_metric.call_count == 2
+
+
+@mark.asyncio
+async def test_process_pi_web_api_response_missing_tags(api):
+    """Test processing response when number of written tags doesn't match expected."""
+    response_data = [
+        {'WebId': 'web_id_1', 'Errors': []},
+    ]
+    tags = {'tag1': 'web_id_1', 'tag2': 'web_id_2'}
+    core_labels = {
+        'pod_id': 'test_pod',
+        'runtime': 'local',
+        'model_name': 'test_model',
+        'workflow_name': 'test_workflow',
+    }
+
+    confidence, message = await api.process_pi_web_api_response(
+        response_data=response_data,
+        tags=tags,
+        core_labels=core_labels,
+        metadata=metadata['metadata'],
+    )
+
+    assert confidence == PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
+    assert (
+        message
+        == "The number of written tags does not match the number of tag names: Expected ['tag1', 'tag2'] tags, but ['tag1'] tags were written."
+    )
+    api.send_notification_async.assert_called_once()
+    call_args = api.send_notification_async.call_args
+    assert call_args.kwargs['notification_id'] == 'WRITE_PI_WEB_API_PREDICTION_ERROR'
+    assert call_args.kwargs['level'] == NotificationLevel.ERROR
+
+
+@mark.asyncio
+async def test_process_pi_web_api_response_missing_webid(api):
+    """Test processing response when WebId is missing in response item."""
+    response_data = [
+        {'Errors': []},
+        {'WebId': 'web_id_2', 'Errors': []},
+    ]
+    tags = {'tag1': 'web_id_1', 'tag2': 'web_id_2'}
+    core_labels = {
+        'pod_id': 'test_pod',
+        'runtime': 'local',
+        'model_name': 'test_model',
+        'workflow_name': 'test_workflow',
+    }
+
+    confidence, message = await api.process_pi_web_api_response(
+        response_data=response_data,
+        tags=tags,
+        core_labels=core_labels,
+        metadata=metadata['metadata'],
+    )
+
+    assert confidence == PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
+    assert (
+        message
+        == "The number of written tags does not match the number of tag names: Expected ['tag1', 'tag2'] tags, but ['tag2'] tags were written."
+    )
+    api.error.assert_any_call('The response did not contain some WebIds', metadata['metadata'])
+
+
+@mark.asyncio
+async def test_process_pi_web_api_response_missing_tag_name(api):
+    """Test processing response when tag name is not found for WebId."""
+    response_data = [
+        {'WebId': 'unknown_web_id', 'Errors': []},
+    ]
+    tags = {'tag1': 'web_id_1'}
+    core_labels = {
+        'pod_id': 'test_pod',
+        'runtime': 'local',
+        'model_name': 'test_model',
+        'workflow_name': 'test_workflow',
+    }
+
+    confidence, message = await api.process_pi_web_api_response(
+        response_data=response_data,
+        tags=tags,
+        core_labels=core_labels,
+        metadata=metadata['metadata'],
+    )
+
+    assert confidence == PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
+    assert (
+        message
+        == "The number of written tags does not match the number of tag names: Expected ['tag1'] tags, but [] tags were written."
+    )
+    api.error.assert_any_call(
+        'The response did not contain the tag name for WebId unknown_web_id', metadata['metadata']
+    )

tests/laborious/activities/test_mlflow.py

@@ -1,248 +1,517 @@
-from unittest.mock import ANY, MagicMock, patch
+from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
 
 import numpy as np
-from pandas import DataFrame
-from pytest import fixture, mark
-from laborious.activities.mlflow import MLFlow
+from pytest import fixture, mark, raises
 from sientia_do.notifications.models import NotificationLevel
+from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
+
+from laborious.activities.mlflow import MLFlow
 
 
-@patch("laborious.activities.mlflow.MLFlowRepository")
-def test___init__(mock_mlflow_repository):
+@fixture(autouse=True)
+def _passthrough_from_dict():
+    with patch(
+        'laborious.activities.mlflow.MinioDataFramePayload.from_dict', side_effect=lambda x: x
+    ):
+        yield
+
+
+@patch('laborious.activities.mlflow.MLFlowRepository')
+@patch('laborious.activities.mlflow.MinioRepository')
+def test___init__(mock_minio_repository, mock_mlflow_repository):
+    logger = MagicMock()
+    notification_handler = MagicMock()
+    metrics_controller = AsyncMock()
+    minio_repo = mock_minio_repository(
+        endpoint='localhost:9000',
+        access_key='minio',
+        secret_key='minio123',
+        logger=logger,
+        notification_handler=notification_handler,
+        metrics_controller=metrics_controller,
+        bucket='test',
+    )
     mlflow = MLFlow(
-        mlflow_host="http://localhost",
+        mlflow_host='http://localhost',
         mlflow_port=5000,
-        mlflow_username="admin",
-        mlflow_password="admin",
-        logger=MagicMock(),
-        notification_handler=MagicMock()
+        mlflow_username='admin',
+        mlflow_password='admin',
+        minio_repository=minio_repo,
+        logger=logger,
+        notification_handler=notification_handler,
+        metrics_controller=metrics_controller,
     )
 
-    assert mlflow.mlflow_host == "http://localhost"
+    assert mlflow.mlflow_host == 'http://localhost'
     assert mlflow.mlflow_port == 5000
-    assert mlflow.mlflow_username == "admin"
-    assert mlflow.mlflow_password == "admin"
+    assert mlflow.mlflow_username == 'admin'
+    assert mlflow.mlflow_password == 'admin'
 
     mock_mlflow_repository.assert_called_once_with(
-        "http://localhost:5000", "admin", "admin", ANY
+        'http://localhost:5000', 'admin', 'admin', ANY, ANY, ANY
+    )
+
+    mock_minio_repository.assert_called_once_with(
+        endpoint='localhost:9000',
+        access_key='minio',
+        secret_key='minio123',
+        logger=ANY,
+        notification_handler=ANY,
+        metrics_controller=ANY,
+        bucket='test',
     )
 
 
 @fixture
-@patch("laborious.activities.mlflow.MLFlowRepository")
-def mlflow(mock_mlflow_repository):
+@patch('laborious.activities.mlflow.MLFlowRepository')
+@patch('laborious.activities.mlflow.MinioRepository')
+def mlflow(mock_minio_repository, mock_mlflow_repository):
+    logger = MagicMock()
+    notification_handler = MagicMock()
+    metrics_controller = AsyncMock()
+    minio_repo = mock_minio_repository(
+        endpoint='localhost:9000',
+        access_key='minio',
+        secret_key='minio123',
+        logger=logger,
+        notification_handler=notification_handler,
+        metrics_controller=metrics_controller,
+        bucket='test',
+    )
     mlflow = MLFlow(
-        mlflow_host="http://localhost:5000",
+        mlflow_host='http://localhost:5000',
         mlflow_port=5000,
-        mlflow_username="admin",
-        mlflow_password="admin",
-        logger=MagicMock(),
-        notification_handler=MagicMock()
+        mlflow_username='admin',
+        mlflow_password='admin',
+        minio_repository=minio_repo,
+        logger=logger,
+        notification_handler=notification_handler,
+        metrics_controller=metrics_controller,
     )
 
+    mlflow.model_monitoring_repository = AsyncMock()
+    mlflow.minio_repository = AsyncMock()
+
     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
 
 
 metadata = {
-    "metadata": {
-        "model_id": "test_model",
-        "model_name": "test_model",
-        "workflow_name": "test_workflow",
-        "schema_name": "test_schedule",
+    'metadata': {
+        'model_id': 'test_model',
+        'model_name': 'test_model',
+        'workflow_name': 'test_workflow',
+        'schema_name': 'test_schedule',
     },
 }
 
 
 @mark.asyncio
-@patch("laborious.activities.mlflow.DataFrame")
-@patch("laborious.activities.mlflow.max")
-async def test_request_transform(mock_max, mock_dataframe, mlflow):
-    mock_max.return_value = '2024-01-02'
-    # Mock input data
+@patch(
+    'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
+    new_callable=AsyncMock,
+)
+async def test_request_transform_success(mock_from_dataframe, mlflow):
+    data_mock = MagicMock()
+    payload = AsyncMock()
+    payload.retrieve = AsyncMock(return_value=data_mock)
+
     input_data = {
         **metadata,
-        'data': [
-            {'timestamp': '2024-01-01', 'variable': 'var1',
-                'value': 1.0, 'created_at': '2024-01-01 12:00:00'},
-            {'timestamp': '2024-01-01', 'variable': 'var2',
-                'value': 2.0, 'created_at': '2024-01-01 12:00:00'},
-            {'timestamp': '2024-01-02', 'variable': 'var1',
-                'value': 3.0, 'created_at': '2024-01-02 12:00:00'},
-            {'timestamp': '2024-01-02', 'variable': 'var2',
-                'value': 4.0, 'created_at': '2024-01-02 12:00:00'},
-            {'timestamp': '2024-01-02', 'variable': 'var1',
-                'value': 1.0, 'created_at': '2024-01-01 12:00:00'},
-            {'timestamp': '2024-01-02', 'variable': 'var2',
-                'value': 1.0, 'created_at': '2024-01-01 12:00:00'}
-        ],
+        'data': payload,
         'model_name': 'test_model',
-        'model_retention': 30
+        'model_config': {},
     }
 
-    # Mock the transform response
-    expected_response = {'prediction': [0.5, 0.6], 'timestamp': [
-        '2024-01-01', '2024-01-02']}
-    mlflow.model_monitoring_repository.transform.return_value = expected_response
+    transform_response = {'success': True, 'content': MagicMock()}
+    mlflow.model_monitoring_repository.transform.return_value = transform_response
 
-    mock_dataframe.return_value.sort_values.return_value = mock_dataframe.return_value
-    mock_dataframe.return_value.drop_duplicates.return_value = mock_dataframe.return_value
+    data_mock.sort_values.return_value = data_mock
+    data_mock.drop_duplicates.return_value = data_mock
+    data_mock.pivot.return_value = data_mock
 
-    # Call the method
     response_data = await mlflow.request_transform(input_data)
 
-    # Verify the data was correctly transformed
-    mock_dataframe.assert_called_once_with(input_data['data'])
-    mock_dataframe.return_value.pivot.assert_called_once_with(
-        index='timestamp', columns='variable', values='value'
-    )
-    mock_dataframe = mock_dataframe.return_value.pivot.return_value
-    mock_dataframe.fillna.assert_called_once_with(np.nan, inplace=True)
-    # mock_dataframe.reset_index.assert_called_once()
-    mock_dataframe.columns.name = None
-
-    # Verify the response
-    assert response_data == expected_response
-
-    # Verify the repository was called with correct arguments
     mlflow.model_monitoring_repository.transform.assert_called_once_with(
-        'test_model', mock_dataframe, 30
+        'test_model', data_mock, {}, metadata['metadata']
     )
+    mock_from_dataframe.assert_called_once()
+    assert response_data == mock_from_dataframe.return_value
 
 
 @mark.asyncio
-@patch("laborious.activities.mlflow.DataFrame")
-@patch("laborious.activities.mlflow.max")
-async def test_request_predict(mock_max, mock_dataframe, mlflow):
-    mock_max.return_value = '2024-01-02'
-    # Mock input data
+@patch(
+    'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
+    new_callable=AsyncMock,
+)
+async def test_request_transform_failure(mock_from_dataframe, mlflow):
+    data_mock = MagicMock()
+    payload = AsyncMock()
+    payload.retrieve = AsyncMock(return_value=data_mock)
+
     input_data = {
         **metadata,
-        'data': [
-            {'timestamp': '2024-01-01', 'variable': 'var1', 'value': 1.0},
-            {'timestamp': '2024-01-01', 'variable': 'var2', 'value': 2.0},
-            {'timestamp': '2024-01-02', 'variable': 'var1', 'value': 3.0},
-            {'timestamp': '2024-01-02', 'variable': 'var2', 'value': 4.0}
-        ],
+        'data': payload,
         'model_name': 'test_model',
-        'model_retention': 30
+        'model_config': {},
     }
 
-    # Mock the predict response
-    expected_response = {'prediction': [0.5, 0.6]}
-    mlflow.model_monitoring_repository.predict.return_value = expected_response
+    transform_response = {'success': False, 'message': 'Transform failed'}
+    mlflow.model_monitoring_repository.transform.return_value = transform_response
+
+    data_mock.sort_values.return_value = data_mock
+    data_mock.drop_duplicates.return_value = data_mock
+    data_mock.pivot.return_value = data_mock
+
+    response_data = await mlflow.request_transform(input_data)
+
+    mock_from_dataframe.assert_called_once_with(
+        dataframe=None,
+        minio_repo=mlflow.minio_repository,
+        model_name='test_model',
+        operation='transform',
+        status=transform_response,
+        workflow_metadata=metadata['metadata'],
+        last_timestamp=payload.last_timestamp,
+        logger=mlflow.logger,
+    )
+    assert response_data == mock_from_dataframe.return_value
+
+
+@mark.asyncio
+@patch(
+    'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
+    new_callable=AsyncMock,
+)
+@patch('laborious.activities.mlflow.to_datetime')
+async def test_request_predict(mock_to_datetime, mock_from_dataframe, mlflow):
+    data_mock = MagicMock()
+    payload = AsyncMock()
+    payload.retrieve = AsyncMock(return_value=data_mock)
+
+    input_data = {
+        **metadata,
+        'data': payload,
+        'model_name': 'test_model',
+        'model_config': {},
+    }
+
+    predict_response = {'success': True, 'content': MagicMock()}
+    mlflow.model_monitoring_repository.predict.return_value = predict_response
 
-    # Call the method
     response_data = await mlflow.request_predict(input_data)
 
-    mock_dataframe.assert_called_once_with(input_data['data'])
-    mock_dataframe.return_value.replace.assert_called_once_with(
-        np.nan, None, inplace=True
+    data_mock.replace.assert_called_once_with(np.nan, None, inplace=True)
+    mock_to_datetime.assert_called_once_with(
+        data_mock.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ
     )
+    mock_to_datetime.return_value.dt.strftime.assert_called_once_with(DATETIME_FORMAT)
 
-    # Verify the response
-    assert response_data == expected_response
-
-    # Verify the repository was called with correct arguments
     mlflow.model_monitoring_repository.predict.assert_called_once_with(
-        'test_model', mock_dataframe.return_value, 30
+        'test_model', data_mock, {}, metadata['metadata']
     )
+    mock_from_dataframe.assert_called_once()
+    assert response_data == mock_from_dataframe.return_value
 
 
 @mark.asyncio
-async def test_retrain_model(mlflow):
-    data = {
-        "model_id": [4, 5, 6, 7],
-        "created_at": [1, 2, 3, 4],
-        "timestamp": [1, 1, 2, 2],
-        "variable": ["var1", "var2", "var1", "var2"],
-        "value": [1, 2, 3, 4]
+@patch(
+    'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
+    new_callable=AsyncMock,
+)
+@patch('laborious.activities.mlflow.to_datetime')
+async def test_request_predict_failure(mock_to_datetime, mock_from_dataframe, mlflow):
+    data_mock = MagicMock()
+    payload = AsyncMock()
+    payload.retrieve = AsyncMock(return_value=data_mock)
+
+    input_data = {
+        **metadata,
+        'data': payload,
+        'model_name': 'test_model',
+        'model_config': {},
     }
 
-    mlflow.model_monitoring_repository.retrain_model.return_value = (
-        'Model retrained successfully', 'test')
+    predict_response = {'success': False, 'message': 'Predict failed'}
+    mlflow.model_monitoring_repository.predict.return_value = predict_response
 
-    response = await mlflow.retrain_model({
-        **metadata,
-        'data': data,
-        'model_name': 'test_model'
-    })
+    response_data = await mlflow.request_predict(input_data)
 
-    mlflow.model_monitoring_repository.retrain_model.assert_called_once()
+    mock_from_dataframe.assert_called_once_with(
+        dataframe=None,
+        minio_repo=mlflow.minio_repository,
+        model_name='test_model',
+        operation='predict',
+        status=predict_response,
+        workflow_metadata=metadata['metadata'],
+        last_timestamp=payload.last_timestamp,
+        logger=mlflow.logger,
+    )
+    assert response_data == mock_from_dataframe.return_value
+
+
+@mark.asyncio
+@patch('laborious.activities.mlflow.to_datetime')
+async def test_retrain_model_success_data_success_retrain(mock_to_datetime, mlflow):
+    mlflow.model_monitoring_repository.retrain_model.return_value = {
+        'success': True,
+        'experiment': 'test_experiment',
+        'message': 'Model retrained successfully.',
+    }
+
+    raw_data = MagicMock(columns=['variable', 'timestamp', 'value'])
+    payload = AsyncMock()
+    payload.retrieve = AsyncMock(return_value=raw_data)
+
+    response = await mlflow.retrain_model(
+        {
+            **metadata,
+            'data': payload,
+            'model_name': 'test_model',
+            'model_config': {
+                'target': 'target',
+                'transform_flavor': 'sklearn',
+                'predict_flavor': 'pyfunc',
+            },
+        }
+    )
+
+    timestamp = raw_data.__getitem__.return_value.max.return_value
+
+    raw_data.sort_values.assert_not_called()
+    raw_data.drop_duplicates.assert_called_once_with(subset=['variable', 'timestamp'], keep='first')
+    raw_data = raw_data.drop_duplicates.return_value
+
+    raw_data.drop.assert_has_calls(
+        [
+            call(columns=['model_id'], inplace=True, errors='ignore'),
+            call(columns=['created_at'], inplace=True, errors='ignore'),
+        ]
+    )
+    raw_data.pivot.assert_called_once_with(index='timestamp', columns='variable', values='value')
+    raw_data.pivot.return_value.fillna.assert_called_once_with(np.nan, inplace=True)
+
+    raw_data = raw_data.pivot.return_value
+
+    raw_data.__setitem__.assert_has_calls(
+        [
+            call('timestamp', raw_data.index),
+            call('timestamp', mock_to_datetime.return_value.dt.strftime.return_value),
+            call('timestamp', mock_to_datetime.return_value),
+        ]
+    )
+    mock_to_datetime.assert_has_calls(
+        [call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ)]
+    )
+    mock_to_datetime.assert_has_calls(
+        [call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT)]
+    )
+
+    mlflow.model_monitoring_repository.retrain_model.assert_called_once_with(
+        data=raw_data,
+        model_name='test_model',
+        model_config={
+            'target': 'target',
+            'transform_flavor': 'sklearn',
+            'predict_flavor': 'pyfunc',
+        },
+        metadata=metadata['metadata'],
+    )
 
     assert response == {
-        "status": 'Model retrained successfully',
-        "timestamp": 2,
-        "experiment": 'test'
+        'success': True,
+        'experiment': 'test_experiment',
+        'message': 'Model retrained successfully.',
+        'timestamp': timestamp,
     }
 
 
 @mark.asyncio
-async def test_retrain_model_error(mlflow):
-    mlflow.model_monitoring_repository.retrain_model.side_effect = Exception(
-        'Error retraining model'
-    )
-
-    data = {
-        "model_id": [4, 5, 6, 7],
-        "created_at": [1, 2, 3, 4],
-        "timestamp": [1, 1, 2, 2],
-        "variable": ["var1", "var2", "var1", "var2"],
-        "value": [1, 2, 3, 4]
+@patch('laborious.activities.mlflow.to_datetime')
+async def test_retrain_model_success_with_payload_data(mock_to_datetime, mlflow):
+    mlflow.model_monitoring_repository.retrain_model.return_value = {
+        'success': True,
+        'experiment': 'test_experiment',
+        'message': 'Model retrained successfully.',
     }
 
-    try:
-        await mlflow.retrain_model({
+    raw_data = MagicMock(columns=['variable', 'timestamp', 'value', 'created_at'])
+    payload = AsyncMock()
+    payload.retrieve = AsyncMock(return_value=raw_data)
+
+    response = await mlflow.retrain_model(
+        {
             **metadata,
-            'data': data,
-            'model_name': 'test_model'
-        })
-    except Exception as e:
-        assert str(e) == 'Error retraining model'
-        mlflow.send_notification.assert_called_once_with(
-            metadata=metadata['metadata'],
-            notification_id='RETRAIN_MODEL_ERROR',
-            message='Error retraining model test_model: Error retraining model',
-            block='retrain_model',
-            level=NotificationLevel.ERROR,
-            attachment_content=ANY
+            'data': payload,
+            'model_name': 'test_model',
+            'model_config': {
+                'target': 'target',
+                'transform_flavor': 'sklearn',
+                'predict_flavor': 'pyfunc',
+            },
+        }
+    )
+
+    assert response['success'] is True
+    mlflow.minio_repository.download_file.assert_not_called()
+
+
+@mark.asyncio
+@patch('laborious.activities.mlflow.to_datetime')
+async def test_retrain_model_success_data_fail_retrain(mock_to_datetime, mlflow):
+    mlflow.model_monitoring_repository.retrain_model.return_value = {
+        'success': False,
+        'traceback': 'test_traceback',
+        'message': 'Model retrained failed.',
+    }
+
+    raw_data = MagicMock(columns=['variable', 'timestamp', 'value', 'created_at'])
+    payload = AsyncMock()
+    payload.retrieve = AsyncMock(return_value=raw_data)
+
+    response = await mlflow.retrain_model(
+        {
+            **metadata,
+            'data': payload,
+            'model_name': 'test_model',
+            'model_config': {
+                'target': 'target',
+                'transform_flavor': 'sklearn',
+                'predict_flavor': 'pyfunc',
+            },
+        }
+    )
+
+    timestamp = raw_data.__getitem__.return_value.max.return_value
+
+    raw_data.sort_values.assert_called_once_with('created_at', ascending=False)
+    raw_data.sort_values.return_value.drop_duplicates.assert_called_once_with(
+        subset=['variable', 'timestamp'], keep='first'
+    )
+    raw_data = raw_data.sort_values.return_value.drop_duplicates.return_value
+
+    raw_data.drop.assert_has_calls(
+        [
+            call(columns=['model_id'], inplace=True, errors='ignore'),
+            call(columns=['created_at'], inplace=True, errors='ignore'),
+        ]
+    )
+    raw_data.pivot.assert_called_once_with(index='timestamp', columns='variable', values='value')
+    raw_data.pivot.return_value.fillna.assert_called_once_with(np.nan, inplace=True)
+
+    raw_data = raw_data.pivot.return_value
+
+    raw_data.__setitem__.assert_has_calls(
+        [
+            call('timestamp', raw_data.index),
+            call('timestamp', mock_to_datetime.return_value.dt.strftime.return_value),
+            call('timestamp', mock_to_datetime.return_value),
+        ]
+    )
+    mock_to_datetime.assert_has_calls(
+        [call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ)]
+    )
+    mock_to_datetime.assert_has_calls(
+        [call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT)]
+    )
+
+    mlflow.model_monitoring_repository.retrain_model.assert_called_once_with(
+        data=raw_data,
+        model_name='test_model',
+        model_config={
+            'target': 'target',
+            'transform_flavor': 'sklearn',
+            'predict_flavor': 'pyfunc',
+        },
+        metadata=metadata['metadata'],
+    )
+
+    mlflow.send_notification_async.assert_called_once_with(
+        metadata=metadata['metadata'],
+        notification_id='RETRAIN_MODEL_ERROR',
+        message='Error retraining model test_model: Model retrained failed.',
+        block='retrain_model',
+        level=NotificationLevel.ERROR,
+        attachment_content=ANY,
+    )
+
+    assert response == {
+        'success': False,
+        'traceback': 'test_traceback',
+        'message': 'Model retrained failed.',
+        'timestamp': timestamp,
+    }
+
+
+@mark.asyncio
+async def test_retrain_model_data_error(mlflow):
+    response = await mlflow.retrain_model(
+        {
+            **metadata,
+            'model_name': 'test_model',
+            'model_config': {
+                'target': 'target',
+                'transform_flavor': 'sklearn',
+                'predict_flavor': 'pyfunc',
+            },
+        }
+    )
+
+    assert response == {
+        'success': False,
+        'message': "Error loading retrain data: 'data'",
+        'traceback': ANY,
+        'timestamp': ANY,
+    }
+
+
+@mark.asyncio
+async def test_retrain_model_data_error_no_minio_repository(mlflow):
+    mlflow.minio_repository = None
+
+    with raises(ValueError) as e:
+        await mlflow.retrain_model(
+            {
+                **metadata,
+                'object_key': 'test_object_key',
+                'model_name': 'test_model',
+                'model_config': {
+                    'target': 'target',
+                    'transform_flavor': 'sklearn',
+                    'predict_flavor': 'pyfunc',
+                },
+            }
         )
-    else:
-        assert False, "No exception raised"
+
+        assert str(e.value) == 'Minio repository not initialized'
 
 
 @mark.asyncio
 async def test_update_production_model(mlflow):
-    mlflow.model_monitoring_repository.update_production_model.return_value = (
-        {
-            "data1": 1,
-            "data2": 2
-        }
-    )
-
     input_data = {
         **metadata,
         'model_name': 'test_model',
         'model_id': 1,
         'experiment': 'test',
         'timestamp': 2,
-        'status': 'success'
+        'status': 'success',
     }
 
     response = await mlflow.update_production_model(input_data)
 
     mlflow.model_monitoring_repository.update_production_model.assert_called_once_with(
-        experiment='test', model_name='test_model')
+        experiment='test', model_name='test_model', metadata=metadata['metadata']
+    )
 
-    assert response == {
-        'data1': {0: 1},
-        'data2': {0: 2},
-        'model_id': {0: 1},
-        'model_name': {0: 'test_model'},
-        'timestamp': {0: 2},
-        'status': {0: 'success'}
-    }
+    assert response == mlflow.model_monitoring_repository.update_production_model.return_value
 
 
 @mark.asyncio
@@ -257,20 +526,104 @@ async def test_update_production_model_error(mlflow):
         'model_id': 1,
         'experiment': 'test',
         'timestamp': 2,
-        'status': 'success'
+        'status': 'success',
     }
 
     try:
         await mlflow.update_production_model(input_data)
     except Exception as e:
         assert str(e) == 'Error updating production model'
-        mlflow.send_notification.assert_called_once_with(
+        mlflow.send_notification_async.assert_called_once_with(
             metadata=metadata['metadata'],
             notification_id='UPDATE_PRODUCTION_MODEL_ERROR',
             message='Error updating production model test_model: Error updating production model',
             block='update_production_model',
             level=NotificationLevel.ERROR,
-            attachment_content=ANY
+            attachment_content=ANY,
         )
     else:
-        assert False, "No exception raised"
+        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_opc.py

@@ -1,29 +1,38 @@
-from unittest.mock import patch, MagicMock, ANY, call, AsyncMock
-from pandas import DataFrame
-from pytest import fixture, mark
+from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
+
 import pytest_asyncio
+from pandas import DataFrame
+from pytest import mark
 from sientia_do.notifications.models import NotificationLevel
 
-from laborious.activities.opc import OPC
+from laborious.activities.opc import (
+    OPC,
+    OPC_COMMENT_SEPARATOR,
+    OPC_RECONNECT_IN_PROGRESS_COMMENT,
+    OPC_SESSION_BAD_COMMENT_PREFIX,
+    OPC_SESSION_BAD_CONFIDENCE,
+    OPC_WRITTING_ERROR_CONFIDENCE,
+    OPC_WRITTING_ERROR_MESSAGE,
+    _apply_opc_write_error,
+)
 
 metadata = {
-    "metadata": {
-        "model_id": "test_model",
-        "model_name": "test_model",
-        "workflow_name": "test_workflow",
-        "schema_name": "test_schedule",
+    'metadata': {
+        'model_id': 'test_model',
+        'model_name': 'test_model',
+        'workflow_name': 'test_workflow',
+        'schema_name': 'test_schedule',
     },
 }
 
 
 def test__init__():
-    servers = {
-        'server1': 'config'
-    }
+    servers = {'server1': {'id': 'server1'}}
     opc = OPC(
         opc_servers=servers,
         logger=MagicMock(),
-        notification_handler=MagicMock()
+        notification_handler=MagicMock(),
+        metrics_controller=AsyncMock(),
     )
 
     assert opc.opc_servers == servers
@@ -31,32 +40,37 @@ def test__init__():
 
 
 @mark.asyncio
-@patch("laborious.activities.opc.OpcRepository")
-@patch("laborious.activities.opc.OPC.send_notification")
+@patch('laborious.activities.opc.OpcRepository')
+@patch('laborious.activities.opc.OPC.send_notification_async')
 async def test_init_opc(mock_send_notification, mock_opc_repository):
     mock_logger = MagicMock()
+    mock_metrics_controller = AsyncMock()
     server1 = MagicMock(
-        connect=AsyncMock(return_value=(True, {})),
-        write_data=AsyncMock(return_value=(True, {}))
+        connect=AsyncMock(return_value=(True, {})), write_data=AsyncMock(return_value=(True, {}))
     )
     server2 = MagicMock(
-        connect=AsyncMock(return_value=(True, {})),
-        write_data=AsyncMock(return_value=(True, {}))
+        connect=AsyncMock(return_value=(True, {})), write_data=AsyncMock(return_value=(True, {}))
     )
     server3 = MagicMock(
-        connect=AsyncMock(return_value=(False, {
-            'notification_id': 'OPC_CONNECTION_ERROR_server3',
-            'message': 'Failed to connect to OPC server: Test error',
-            'block': 'opc_repository',
-            'level': NotificationLevel.ERROR,
-            'attachment_content': 'Test error'
-        })),
-        write_data=AsyncMock(return_value=(True, {}))
+        connect=AsyncMock(
+            return_value=(
+                False,
+                {
+                    'notification_id': 'OPC_CONNECTION_ERROR_server3',
+                    'message': 'Failed to connect to OPC server: Test error',
+                    'block': 'opc_repository',
+                    'level': NotificationLevel.ERROR,
+                    'attachment_content': 'Test error',
+                },
+            )
+        ),
+        write_data=AsyncMock(return_value=(True, {})),
     )
     mock_opc_repository.side_effect = [server1, server2, server3]
     mock_notification_handler = MagicMock()
     servers = {
         'server1': {
+            'server_name': 'server1',
             'id': 'server1',
             'url': 'http://localhost:8080',
             'server_uri': 'opc.tcp://localhost:4840',
@@ -66,6 +80,7 @@ async def test_init_opc(mock_send_notification, mock_opc_repository):
             'reconnection_interval': 60,
         },
         'server2': {
+            'server_name': 'server2',
             'id': 'server2',
             'url': 'http://localhost:8080',
             'server_uri': 'opc.tcp://localhost:4840',
@@ -75,6 +90,7 @@ async def test_init_opc(mock_send_notification, mock_opc_repository):
             'reconnection_interval': 60,
         },
         'server3': {
+            'server_name': 'server3',
             'id': 'server3',
             'url': 'http://localhost:8080',
             'server_uri': 'opc.tcp://localhost:4840',
@@ -82,12 +98,13 @@ async def test_init_opc(mock_send_notification, mock_opc_repository):
             'private_key_path': '',
             'server_cert_path': '',
             'reconnection_interval': 60,
-        }
+        },
     }
     opc = OPC(
         opc_servers=servers,
         logger=mock_logger,
-        notification_handler=mock_notification_handler
+        notification_handler=mock_notification_handler,
+        metrics_controller=mock_metrics_controller,
     )
     await opc.init_opc()
 
@@ -97,61 +114,70 @@ async def test_init_opc(mock_send_notification, mock_opc_repository):
     assert opc.opc_repository['server1'] == server1
     assert opc.opc_repository['server2'] == server2
 
-    mock_opc_repository.assert_has_calls([
-        call(
-            id="server1",
-            url="http://localhost:8080",
-            logger=mock_logger,
-            server_uri="opc.tcp://localhost:4840",
-            cert_path="",
-            private_key_path="",
-            server_cert_path="",
-            notification_handler=mock_notification_handler,
-            reconnection_interval=60,
-            pod_id='localhost'
-        ),
-    ])
-    mock_opc_repository.assert_has_calls([
-        call(
-            id="server2",
-            url="http://localhost:8080",
-            logger=mock_logger,
-            server_uri="opc.tcp://localhost:4840",
-            cert_path="",
-            private_key_path="",
-            server_cert_path="",
-            notification_handler=mock_notification_handler,
-            reconnection_interval=60,
-            pod_id='localhost'
-        )
-    ])
+    mock_opc_repository.assert_has_calls(
+        [
+            call(
+                opc_id='server1',
+                server_name='server1',
+                url='http://localhost:8080',
+                logger=mock_logger,
+                server_uri='opc.tcp://localhost:4840',
+                cert_path='',
+                private_key_path='',
+                server_cert_path='',
+                notification_handler=mock_notification_handler,
+                reconnection_interval=60,
+                metrics_controller=mock_metrics_controller,
+            ),
+        ]
+    )
+    mock_opc_repository.assert_has_calls(
+        [
+            call(
+                opc_id='server2',
+                server_name='server2',
+                url='http://localhost:8080',
+                logger=mock_logger,
+                server_uri='opc.tcp://localhost:4840',
+                cert_path='',
+                private_key_path='',
+                server_cert_path='',
+                notification_handler=mock_notification_handler,
+                reconnection_interval=60,
+                metrics_controller=mock_metrics_controller,
+            )
+        ]
+    )
 
     server1.connect.assert_called_once()
     server2.connect.assert_called_once()
 
-    mock_send_notification.assert_has_calls([
-        call(
-            metadata={
-                'model_id': '-',
-                'model_name': '-',
-                'workflow_name': '-',
-                'schedule_name': 'INITIALIZATION'
-            },
-            notification_id="OPC_CONNECTION_ERROR_server3",
-            message="Failed to connect to OPC server: Test error",
-            block="opc_repository",
-            level=NotificationLevel.ERROR,
-            attachment_content=ANY
-        )
-    ])
+    mock_send_notification.assert_has_calls(
+        [
+            call(
+                metadata={
+                    'model_id': '-',
+                    'model_name': '-',
+                    'workflow_name': '-',
+                    'schedule_name': 'INITIALIZATION',
+                },
+                notification_id='OPC_CONNECTION_ERROR_server3',
+                message='Failed to connect to OPC server: Test error',
+                block='opc_repository',
+                level=NotificationLevel.ERROR,
+                attachment_content=ANY,
+            )
+        ]
+    )
 
 
 @pytest_asyncio.fixture
-@patch("laborious.activities.opc.OpcRepository")
+@patch('laborious.activities.opc.OpcRepository')
 async def opc(mock_opc_repository):
     servers = {
         'server1': {
             'id': 'server1',
+            'server_name': 'server1',
             'url': 'http://localhost:8080',
             'server_uri': 'opc.tcp://localhost:4840',
             'cert_path': '',
@@ -161,19 +187,18 @@ async def opc(mock_opc_repository):
         }
     }
 
-    mock_opc_repository.return_value.write_data = AsyncMock(
-        return_value=(True, {})
-    )
-    mock_opc_repository.return_value.connect = AsyncMock(
-        return_value=(True, {})
-    )
+    mock_opc_repository.return_value.write_data = AsyncMock(return_value=(True, {}))
+    mock_opc_repository.return_value.connect = AsyncMock(return_value=(True, {}))
     opc = OPC(
         opc_servers=servers,
         logger=MagicMock(),
-        notification_handler=MagicMock()
+        notification_handler=MagicMock(),
+        metrics_controller=AsyncMock(),
     )
     await opc.init_opc()
     opc.send_notification = MagicMock()
+    opc.send_notification_async = AsyncMock()
+    opc.emit_metric = AsyncMock()
     return opc
 
 
@@ -188,108 +213,385 @@ WRITE_DATA_CASES = [
 @mark.parametrize('tag,data_type,data', WRITE_DATA_CASES)
 @mark.asyncio
 async def test_write_data_success(opc, tag, data_type, data):
-    result = await opc.write_data(server_id='server1', tag=tag, data=data,
-                                  data_type=data_type, tag_type='prediction', metadata=metadata)
-    assert result is True
-    opc.opc_repository['server1'].write_data.assert_called_once_with(
-        tag, data, data_type, opc.logger, metadata)
+    opc.opc_repository['server1'].write_data.return_value = (True, {'response_time': 0.1})
+
+    response_time, error_info = await opc.write_data(
+        server_id='server1',
+        tag=tag,
+        data=data,
+        data_type=data_type,
+        tag_type='prediction',
+        metadata=metadata,
+    )
+    assert response_time == 0.1
+    assert error_info is None
+    opc.opc_repository['server1'].write_data.assert_called_once_with(tag, data, data_type, metadata)
 
 
 @mark.asyncio
 async def test_write_data_failed(opc):
-    opc.opc_repository['server1'].write_data.return_value = (False, {
-        'notification_id': 'OPC_WRITE_DATA_ERROR_server1',
-        'message': 'Failed to write data to OPC server: Test error',
-        'block': 'opc_repository',
-        'level': NotificationLevel.ERROR,
-        'attachment_content': 'Test error'
-    })
+    opc.opc_repository['server1'].write_data.return_value = (
+        False,
+        {
+            'notification_id': 'OPC_WRITE_DATA_ERROR_server1',
+            'message': 'Failed to write data to OPC server: Test error',
+            'block': 'opc_repository',
+            'level': NotificationLevel.ERROR,
+            'attachment_content': 'Test error',
+        },
+    )
 
-    result = await opc.write_data(server_id='server1', tag='tag1', data=50,
-                                  data_type='int', tag_type='prediction', metadata=metadata)
-    assert result is False
-
-    opc.send_notification.assert_called_once_with(
+    response_time, error_info = await opc.write_data(
+        server_id='server1',
+        tag='tag1',
+        data=50,
+        data_type='int',
+        tag_type='prediction',
         metadata=metadata,
-        notification_id="OPC_WRITE_DATA_ERROR_server1",
-        message="Failed to write data to OPC server: Test error",
-        block="opc_repository",
+    )
+    assert response_time is None
+    assert error_info is not None
+
+    opc.send_notification_async.assert_called_once_with(
+        metadata=metadata,
+        notification_id='OPC_WRITE_DATA_ERROR_server1',
+        message='Failed to write data to OPC server: Test error',
+        block='opc_repository',
         level=NotificationLevel.ERROR,
-        attachment_content=ANY
+        attachment_content=ANY,
     )
 
 
 @mark.asyncio
 async def test_write_data_exception(opc):
-    opc.opc_repository['server1'].write_data.side_effect = Exception(
-        "Test error")
+    opc.opc_repository['server1'].write_data.side_effect = Exception('Test error')
 
     try:
-        await opc.write_data(server_id='server1', tag='tag1', data=50,
-                             data_type='int', tag_type='prediction', metadata=metadata)
+        await opc.write_data(
+            server_id='server1',
+            tag='tag1',
+            data=50,
+            data_type='int',
+            tag_type='prediction',
+            metadata=metadata,
+        )
 
     except Exception:
-        opc.send_notification.assert_called_once_with(
+        opc.send_notification_async.assert_called_once_with(
             metadata=metadata,
-            notification_id="WRITE_OPC_PREDICTION_ERROR",
-            message="Error writing data to OPC server: Test error",
-            block="write_opc_data",
+            notification_id='WRITE_OPC_PREDICTION_ERROR',
+            message='Error writing data to OPC server: Test error',
+            block='write_opc_data',
             level=NotificationLevel.ERROR,
-            attachment_content=ANY
+            attachment_content=ANY,
         )
 
     else:
-        assert False, "Expected an exception to be raised"
+        raise AssertionError('Expected an exception to be raised')
+
+
+@mark.parametrize(
+    'error_info,initial_seen,initial_status,initial_reconnect,expected',
+    [
+        (None, False, None, False, (False, None, False)),
+        ({}, False, None, False, (False, None, False)),
+        (
+            {'opc_error_kind': 'session_bad', 'opc_status': 'BadSessionIdInvalid'},
+            False,
+            None,
+            False,
+            (True, 'BadSessionIdInvalid', False),
+        ),
+        (
+            {'opc_error_kind': 'session_bad', 'opc_status': 'NewStatus'},
+            True,
+            'OldStatus',
+            False,
+            (True, 'NewStatus', False),
+        ),
+        (
+            {'opc_error_kind': 'session_bad'},
+            True,
+            'KeptStatus',
+            False,
+            (True, 'KeptStatus', False),
+        ),
+        (
+            {'opc_error_kind': 'reconnect_in_progress'},
+            False,
+            None,
+            False,
+            (False, None, True),
+        ),
+        (
+            {'opc_error_kind': 'other'},
+            True,
+            'Status',
+            True,
+            (True, 'Status', True),
+        ),
+    ],
+)
+def test_apply_opc_write_error(
+    error_info, initial_seen, initial_status, initial_reconnect, expected
+):
+    result = _apply_opc_write_error(
+        error_info,
+        initial_seen,
+        initial_status,
+        initial_reconnect,
+    )
+    assert result == expected
 
 
 @mark.asyncio
-async def test_write_opc_data_success(opc):
+async def test_write_tags_from_config_prediction_success(opc):
+    opc.write_data = AsyncMock(return_value=(0.1, None))
+    data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
+    tags_config = {'tag1': {'data_type': 'float'}}
+
+    response_times, session_bad, opc_status, reconnect = await opc._write_tags_from_config(
+        server_id='server1',
+        tags_config=tags_config,
+        data=data,
+        data_column='prediction',
+        tag_type='prediction',
+        log_label='Prediction data',
+        metadata=metadata['metadata'],
+    )
+
+    assert response_times == {'tag1': 0.1}
+    assert session_bad is False
+    assert opc_status is None
+    assert reconnect is False
+    opc.write_data.assert_called_once_with(
+        server_id='server1',
+        tag='tag1',
+        data=0.75,
+        data_type='float',
+        tag_type='prediction',
+        metadata=metadata['metadata'],
+    )
+
+
+@mark.asyncio
+async def test_write_tags_from_config_confidence_success(opc):
+    opc.write_data = AsyncMock(return_value=(0.2, None))
+    data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
+    tags_config = {'tag2': {'data_type': 'float'}}
+
+    response_times, session_bad, opc_status, reconnect = await opc._write_tags_from_config(
+        server_id='server1',
+        tags_config=tags_config,
+        data=data,
+        data_column='prediction_confidence',
+        tag_type='confidence',
+        log_label='Confidence data',
+        metadata=metadata['metadata'],
+    )
+
+    assert response_times == {'tag2': 0.2}
+    assert session_bad is False
+    assert opc_status is None
+    assert reconnect is False
+    opc.write_data.assert_called_once_with(
+        server_id='server1',
+        tag='tag2',
+        data=0.95,
+        data_type='float',
+        tag_type='confidence',
+        metadata=metadata['metadata'],
+    )
+
+
+@mark.asyncio
+async def test_write_tags_from_config_write_failure(opc):
+    opc.write_data = AsyncMock(return_value=(None, {}))
+    data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
+
+    response_times, session_bad, opc_status, reconnect = await opc._write_tags_from_config(
+        server_id='server1',
+        tags_config={'tag1': {'data_type': 'float'}},
+        data=data,
+        data_column='prediction',
+        tag_type='prediction',
+        log_label='Prediction data',
+        metadata=metadata['metadata'],
+    )
+
+    assert response_times == {'tag1': None}
+    assert session_bad is False
+    assert opc_status is None
+    assert reconnect is False
+
+
+@mark.asyncio
+async def test_write_tags_from_config_session_bad(opc):
+    opc.write_data = AsyncMock(
+        return_value=(
+            None,
+            {
+                'opc_error_kind': 'session_bad',
+                'opc_status': 'BadSessionIdInvalid',
+            },
+        )
+    )
+    data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
+
+    response_times, session_bad, opc_status, reconnect = await opc._write_tags_from_config(
+        server_id='server1',
+        tags_config={'tag1': {'data_type': 'float'}},
+        data=data,
+        data_column='prediction',
+        tag_type='prediction',
+        log_label='Prediction data',
+        metadata=metadata['metadata'],
+    )
+
+    assert response_times == {'tag1': None}
+    assert session_bad is True
+    assert opc_status == 'BadSessionIdInvalid'
+    assert reconnect is False
+
+
+@mark.asyncio
+async def test_write_tags_from_config_reconnect_in_progress(opc):
+    opc.write_data = AsyncMock(
+        return_value=(
+            None,
+            {'opc_error_kind': 'reconnect_in_progress'},
+        )
+    )
+    data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
+
+    response_times, session_bad, opc_status, reconnect = await opc._write_tags_from_config(
+        server_id='server1',
+        tags_config={'tag1': {'data_type': 'float'}},
+        data=data,
+        data_column='prediction',
+        tag_type='prediction',
+        log_label='Prediction data',
+        metadata=metadata['metadata'],
+    )
+
+    assert response_times == {'tag1': None}
+    assert session_bad is False
+    assert opc_status is None
+    assert reconnect is True
+
+
+@mark.asyncio
+async def test_manage_output_tags_success(opc):
+    opc._write_tags_from_config = AsyncMock(
+        side_effect=[
+            ({'tag1': 0.1}, False, None, False),
+            ({'tag2': 0.1}, False, None, False),
+        ]
+    )
+    data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
+    config = {
+        'prediction_tags': {'tag1': {'data_type': 'float'}},
+        'confidence_tags': {'tag2': {'data_type': 'float'}},
+    }
+
+    output_data, opc_metrics, session_bad, opc_status, reconnect = await opc.manage_output_tags(
+        server_id='server1',
+        config=config,
+        data=data,
+        metadata=metadata['metadata'],
+    )
+
+    assert output_data is True
+    assert opc_metrics == {'tag1': 0.1, 'tag2': 0.1}
+    assert session_bad is False
+    assert opc_status is None
+    assert reconnect is False
+    assert opc._write_tags_from_config.await_count == 2
+
+
+@mark.asyncio
+async def test_manage_output_tags_failed(opc):
+    opc._write_tags_from_config = AsyncMock(
+        side_effect=[
+            ({'tag1': 0.1}, False, None, False),
+            ({'tag2': None}, False, None, False),
+        ]
+    )
+    data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
+    config = {
+        'prediction_tags': {'tag1': {'data_type': 'float'}},
+        'confidence_tags': {'tag2': {'data_type': 'float'}},
+    }
+
+    output_data, opc_metrics, _, _, _ = await opc.manage_output_tags(
+        server_id='server1',
+        config=config,
+        data=data,
+        metadata=metadata['metadata'],
+    )
+
+    assert output_data is False
+    assert opc_metrics == {'tag1': 0.1, 'tag2': None}
+
+
+@mark.asyncio
+async def test_manage_output_tags_do_nothing(opc):
+    opc._write_tags_from_config = AsyncMock()
+    data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
+    config = {'_invalid_key': {'tag1': {'data_type': 'float'}}}
+
+    output_data, opc_metrics, _, _, _ = await opc.manage_output_tags(
+        server_id='server1',
+        config=config,
+        data=data,
+        metadata=metadata['metadata'],
+    )
+
+    assert output_data is True
+    assert opc_metrics == {}
+    opc._write_tags_from_config.assert_not_called()
+
+
+@mark.asyncio
+@patch('laborious.activities.opc.DataFrame')
+async def test_write_opc_data_success(mock_dataframe, opc):
     # Arrange
     input_data = {
         **metadata,
-        'data': {
-            'prediction': [0.75],
-            'prediction_confidence': [0.95]
-        },
+        'data': {'prediction': [0.75], 'prediction_confidence': [0.95]},
         'opc_output_config': {
             'server1': {
-                'prediction_tags': {
-                    'tag1': {'data_type': 'float'}
-                },
-                'confidence_tags': {
-                    'tag2': {'data_type': 'float'}
-                }
+                'prediction_tags': {'tag1': {'data_type': 'float'}},
+                'confidence_tags': {'tag2': {'data_type': 'float'}},
             }
-        }
+        },
     }
 
     # Act
-    opc.write_data = AsyncMock(return_value=True)
+    opc.manage_output_tags = AsyncMock(
+        return_value=(True, {'tag1': 0.1, 'tag2': 0.2}, False, None, False)
+    )
+
     opc.process_confidence = MagicMock(return_value={'data': 'data'})
-    output = await opc.write_opc_data(input_data)
+    output_data, opc_metrics = await opc.write_opc_data(input_data)
 
     # Assert
-    assert output == {'data': 'data'}
-    opc.write_data.assert_has_calls([
-        call(
-            server_id='server1',
-            tag='tag1',
-            data=0.75,
-            data_type='float',
-            tag_type='prediction',
-            metadata=metadata['metadata']
-        )])
-    opc.write_data.assert_has_calls([
-        call(
-            server_id='server1',
-            tag='tag2',
-            data=0.95,
-            data_type='float',
-            tag_type='confidence',
-            metadata=metadata['metadata']
-        )
-    ])
-    assert opc.write_data.call_count == 2
+    assert output_data == {'data': 'data'}
+    assert opc_metrics == {'server1': {'tag1': 0.1, 'tag2': 0.2}}
+    opc.manage_output_tags.assert_called_once_with(
+        'server1',
+        input_data['opc_output_config']['server1'],
+        mock_dataframe.return_value,
+        metadata['metadata'],
+    )
+    opc.process_confidence.assert_called_once_with(
+        mock_dataframe.return_value,
+        True,
+        metadata['metadata'],
+        session_bad=False,
+        opc_status=None,
+        reconnect_in_progress=False,
+    )
 
 
 @mark.asyncio
@@ -297,17 +599,9 @@ async def test_write_opc_data_empty_config(opc):
     # Arrange
     input_data = {
         **metadata,
-        'data': {
-            'prediction': [0.75],
-            'prediction_confidence': [0.95]
-        },
+        'data': {'prediction': [0.75], 'prediction_confidence': [0.95]},
         'opc_servers': ['server1'],
-        'opc_output_config': {
-            'server1': {
-                'prediction_tags': {},
-                'confidence_tags': {}
-            }
-        }
+        'opc_output_config': {'server1': {'prediction_tags': {}, 'confidence_tags': {}}},
     }
 
     # Act
@@ -319,23 +613,16 @@ async def test_write_opc_data_empty_config(opc):
 
 @mark.asyncio
 async def test_write_opc_data_no_validate_server(opc):
-    opc.validate_server = MagicMock(return_value=False)
+    opc.validate_server = AsyncMock(return_value=False)
     input_data = {
         **metadata,
-        'data': {
-            'prediction': [0.75],
-            'prediction_confidence': [0.95]
-        },
+        'data': {'prediction': [0.75], 'prediction_confidence': [0.95]},
         'opc_output_config': {
             'server1': {
-                'prediction_tags': {
-                    'tag1': {'data_type': 'float'}
-                },
-                'confidence_tags': {
-                    'tag2': {'data_type': 'float'}
-                }
+                'prediction_tags': {'tag1': {'data_type': 'float'}},
+                'confidence_tags': {'tag2': {'data_type': 'float'}},
             }
-        }
+        },
     }
 
     # Act
@@ -345,25 +632,107 @@ async def test_write_opc_data_no_validate_server(opc):
     opc.opc_repository['server1'].write_data.assert_not_called()
 
 
-@mark.parametrize('data,success,expected', [
-    (DataFrame({'prediction_confidence': [0]}), True, 0),
-    (DataFrame({'prediction_confidence': [0]}), False, 12),
-])
+@mark.parametrize(
+    'data,success,expected',
+    [
+        (DataFrame({'prediction_confidence': [0]}), True, 0),
+        (DataFrame({'prediction_confidence': [0]}), False, 12),
+    ],
+)
 def test_process_confidence(opc, data, success, expected):
-    # Act
-    result = opc.process_confidence(data, success, metadata)
-
-    # Assert
+    result = opc.process_confidence(data, success, metadata['metadata'])
     assert result['prediction_confidence'][0] == expected
 
 
-def test_validate_server(opc):
-    assert opc.validate_server('server1', metadata) is True
-    assert opc.validate_server('server2', metadata) is False
+def test_process_confidence_session_bad(opc):
+    data = DataFrame({'prediction_confidence': [0.9]})
+    result = opc.process_confidence(
+        data,
+        False,
+        metadata['metadata'],
+        session_bad=True,
+        opc_status='BadSessionIdInvalid',
+    )
+    assert result['prediction_confidence'][0] == OPC_SESSION_BAD_CONFIDENCE
+    assert result['comments'][0].startswith(OPC_SESSION_BAD_COMMENT_PREFIX)
+    assert 'BadSessionIdInvalid' in result['comments'][0]
+
+
+def test_process_confidence_generic_failure(opc):
+    data = DataFrame({'prediction_confidence': [0.9]})
+    result = opc.process_confidence(data, False, metadata['metadata'])
+    assert result['prediction_confidence'][0] == OPC_WRITTING_ERROR_CONFIDENCE
+    assert result['comments'][0] == OPC_WRITTING_ERROR_MESSAGE
 
 
 @mark.asyncio
-async def test_shutdown(opc):
+async def test_manage_output_tags_merges_error_flags(opc):
+    opc._write_tags_from_config = AsyncMock(
+        side_effect=[
+            ({'tag1': None}, True, 'BadSessionIdInvalid', False),
+            ({'tag2': 0.2}, False, None, True),
+        ]
+    )
+    data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
+    config = {
+        'prediction_tags': {'tag1': {'data_type': 'float'}},
+        'confidence_tags': {'tag2': {'data_type': 'float'}},
+    }
+
+    (
+        success,
+        metrics,
+        session_bad_seen,
+        opc_status,
+        reconnect_in_progress,
+    ) = await opc.manage_output_tags('server1', config, data, metadata['metadata'])
+
+    assert success is False
+    assert session_bad_seen is True
+    assert reconnect_in_progress is True
+    assert opc_status == 'BadSessionIdInvalid'
+    assert metrics == {'tag1': None, 'tag2': 0.2}
+
+
+def test_process_confidence_reconnect_in_progress(opc):
+    data = DataFrame({'prediction_confidence': [0.9]})
+    result = opc.process_confidence(
+        data,
+        False,
+        metadata['metadata'],
+        reconnect_in_progress=True,
+    )
+    assert result['prediction_confidence'][0] == OPC_SESSION_BAD_CONFIDENCE
+    assert result['comments'][0] == OPC_RECONNECT_IN_PROGRESS_COMMENT
+
+
+def test_process_confidence_concatenates_multiple_comments(opc):
+    data = DataFrame({'prediction_confidence': [0.9]})
+    session_comment = f'{OPC_SESSION_BAD_COMMENT_PREFIX} BadSessionIdInvalid'
+
+    result = opc.process_confidence(
+        data,
+        False,
+        metadata['metadata'],
+        session_bad=True,
+        opc_status='BadSessionIdInvalid',
+        reconnect_in_progress=True,
+    )
+
+    assert result['prediction_confidence'][0] == OPC_SESSION_BAD_CONFIDENCE
+    assert result['comments'][0] == OPC_COMMENT_SEPARATOR.join(
+        [session_comment, OPC_RECONNECT_IN_PROGRESS_COMMENT]
+    )
+
+
+@mark.asyncio
+async def test_validate_server(opc):
+    assert await opc.validate_server('server1', metadata) is True
+    assert await opc.validate_server('server2', metadata) is False
+
+
+@mark.asyncio
+async def test_close(opc):
     opc.opc_repository['server1'].disconnect = AsyncMock(return_value=True)
-    await opc.shutdown()
+    await opc.aclose()
     opc.opc_repository['server1'].disconnect.assert_called_once()

tests/laborious/activities/test_storage.py

@@ -0,0 +1,323 @@
+import datetime
+import os
+from unittest.mock import ANY, AsyncMock, MagicMock, patch
+
+from pytest import fixture, mark, raises
+from sientia_do.notifications.models import NotificationLevel
+from sientia_do.temporal.activities.postgres import Postgres
+
+from laborious.activities.storage import Storage
+
+
+@fixture(autouse=True)
+def _passthrough_from_dict():
+    with patch(
+        'laborious.activities.storage.MinioDataFramePayload.from_dict', side_effect=lambda x: x
+    ):
+        yield
+
+
+metadata = {
+    'metadata': {
+        'model_id': 'test_model_id',
+        'model_name': 'test_model',
+        'workflow_name': 'test_workflow',
+        'schedule_name': 'test_schedule',
+    }
+}
+
+
+@fixture
+@patch('laborious.activities.storage.MinioRepository')
+def storage(mock_minio_repository):
+    return Storage(
+        host='localhost',
+        port=5432,
+        user='postgres',
+        password='postgres',
+        dbname='postgres',
+        min_connections=1,
+        max_connections=10,
+        retention_hours=24,
+        minio_repository=mock_minio_repository.return_value,
+        logger=MagicMock(),
+        notification_handler=MagicMock(),
+        metrics_controller=AsyncMock(),
+    )
+
+
+@patch('laborious.activities.storage.MinioRepository')
+def test___init___not_hasattr(mock_minio_repository):
+    logger = MagicMock()
+    notification_handler = MagicMock()
+    metrics_controller = AsyncMock()
+    minio_repo = mock_minio_repository.return_value
+    storage = Storage(
+        host='localhost',
+        port=5432,
+        user='postgres',
+        password='postgres',
+        dbname='postgres',
+        min_connections=1,
+        max_connections=10,
+        retention_hours=24,
+        minio_repository=minio_repo,
+        logger=logger,
+        notification_handler=notification_handler,
+        metrics_controller=metrics_controller,
+    )
+    assert isinstance(storage, Postgres)
+
+    assert storage.minio_repository is minio_repo
+    mock_minio_repository.assert_not_called()
+
+
+@patch('laborious.activities.storage.MinioRepository')
+def test___init___none_minio_repository(mock_minio_repository, storage):
+    storage.minio_repository = None
+    logger = MagicMock()
+    notification_handler = MagicMock()
+    metrics_controller = AsyncMock()
+    storage.__init__(
+        host='localhost',
+        port=5432,
+        user='postgres',
+        password='postgres',
+        dbname='postgres',
+        min_connections=1,
+        max_connections=10,
+        retention_hours=24,
+        minio_repository=None,
+        logger=logger,
+        notification_handler=notification_handler,
+        metrics_controller=metrics_controller,
+    )
+
+    assert storage.minio_repository is None
+    mock_minio_repository.assert_not_called()
+
+
+@patch('laborious.activities.storage.MinioRepository')
+def test___init___done_repository(mock_minio_repository, storage):
+    storage.__init__(
+        host='localhost',
+        port=5432,
+        user='postgres',
+        password='postgres',
+        dbname='postgres',
+        min_connections=1,
+        max_connections=10,
+        retention_hours=24,
+        minio_repository=mock_minio_repository.return_value,
+        logger=MagicMock(),
+        notification_handler=MagicMock(),
+        metrics_controller=AsyncMock(),
+    )
+    mock_minio_repository.assert_not_called()
+    assert storage.minio_repository is not None
+
+
+def test_close(storage):
+    storage.minio_repository = MagicMock()
+
+    storage.close()
+
+    assert storage.minio_repository is None
+
+
+def test___del__(storage):
+    storage.close = MagicMock()
+
+    storage.__del__()
+
+    storage.close.assert_called_once()
+
+
+@mark.asyncio
+async def test_load_query_with_minio_offload_no_rows(storage):
+    storage.load_custom_query = AsyncMock(return_value=None)
+    storage_result = {'success': False}
+    with patch(
+        'laborious.activities.storage.MinioDataFramePayload.from_dataframe',
+        new_callable=AsyncMock,
+        return_value=storage_result,
+    ) as mock_from_dataframe:
+        result = await storage.load_query_with_minio_offload(
+            {**metadata, 'query': 'SELECT 1', 'model_name': 'm', 'key_prefix': 'predictions/s'}
+        )
+    assert result == storage_result
+    mock_from_dataframe.assert_awaited_once()
+
+
+@mark.asyncio
+async def test_load_query_with_minio_offload_inline(storage):
+    storage.load_custom_query = AsyncMock(return_value=[{'a': 1}])
+    storage_result = {'success': True, 'data': {'a': [1]}, 'object_key': None}
+    with patch(
+        'laborious.activities.storage.MinioDataFramePayload.from_dataframe',
+        new_callable=AsyncMock,
+        return_value=storage_result,
+    ) as mock_from_dataframe:
+        result = await storage.load_query_with_minio_offload(
+            {
+                **metadata,
+                'query': 'SELECT 1',
+                'model_name': 'my-model',
+                'key_prefix': 'predictions/s',
+            }
+        )
+    assert result == storage_result
+    mock_from_dataframe.assert_awaited_once()
+
+
+@mark.asyncio
+async def test_load_query_with_minio_offload_minio(storage):
+    storage.load_custom_query = AsyncMock(return_value=[{'a': 1}])
+    storage_result = {'success': True, 'data': None, 'object_key': 'object-key'}
+    with patch(
+        'laborious.activities.storage.MinioDataFramePayload.from_dataframe',
+        new_callable=AsyncMock,
+        return_value=storage_result,
+    ) as mock_from_dataframe:
+        result = await storage.load_query_with_minio_offload(
+            {**metadata, 'query': 'SELECT 1', 'model_name': 'm', 'key_prefix': 'predictions/s'}
+        )
+
+    assert result == storage_result
+    mock_from_dataframe.assert_awaited_once()
+
+
+@mark.asyncio
+@patch.dict(os.environ, {'SIENTIA_MINIO_RETENTION_HOURS': '1'})
+@patch('laborious.activities.storage.now')
+async def test_cleanup_minio_objects_expired(mock_now, storage):
+    mock_now.return_value = datetime.datetime(2025, 1, 10, 12, 0, 0)
+    storage.minio_repository.list_objects = AsyncMock(
+        return_value=[
+            'sientia/streamlit-connectors/training_datasets/m/m-initial-2024-12-01_00-00-00.parquet',
+            'sientia/streamlit-connectors/training_datasets/m/m-initial-2025-01-10_12-00-00.parquet',
+        ]
+    )
+    storage.minio_repository.delete_file = AsyncMock()
+    storage.send_notification_async = AsyncMock()
+
+    data_mock = MagicMock()
+    data_mock.cleanup_prefix.return_value = 'training_datasets/m'
+
+    result = await storage.cleanup_minio_objects_expired({**metadata, 'data': data_mock})
+
+    assert result['deleted_count'] == 1
+    assert result['failed_count'] == 0
+    deleted_key = (
+        'sientia/streamlit-connectors/training_datasets/m/m-initial-2024-12-01_00-00-00.parquet'
+    )
+    assert deleted_key in result['deleted']
+    assert result['deleted'][deleted_key]['success'] is True
+    storage.minio_repository.list_objects.assert_called_once_with(
+        prefix='training_datasets/m',
+        recursive=True,
+        metadata=metadata['metadata'],
+    )
+    storage.minio_repository.delete_file.assert_called_once_with(
+        object_name='sientia/streamlit-connectors/training_datasets/m/m-initial-2024-12-01_00-00-00.parquet',
+        metadata=metadata['metadata'],
+    )
+
+
+@mark.asyncio
+async def test_load_query_with_minio_offload_minio_not_initialized(storage):
+    storage.minio_repository = None
+
+    with raises(ValueError, match='Minio repository not initialized'):
+        await storage.load_query_with_minio_offload(
+            {**metadata, 'query': 'SELECT 1', 'model_name': 'm'}
+        )
+
+
+@mark.asyncio
+async def test_export_payload_to_postgres(storage):
+    payload = AsyncMock()
+    payload.retrieve = AsyncMock(return_value=MagicMock())
+    storage.export_data_to_postgres = AsyncMock(return_value={'success': True})
+
+    result = await storage.export_payload_to_postgres(
+        {**metadata, 'data': payload, 'schema': 'public', 'table': 't'}
+    )
+
+    payload.retrieve.assert_awaited_once_with(storage.minio_repository, metadata['metadata'])
+    storage.export_data_to_postgres.assert_awaited_once()
+    assert result == {'success': True}
+
+
+@mark.asyncio
+async def test_cleanup_minio_objects_expired_minio_not_initialized(storage):
+    storage.minio_repository = None
+
+    data_mock = MagicMock()
+    data_mock.cleanup_prefix.return_value = 'test'
+    with raises(ValueError, match='Minio repository not initialized'):
+        await storage.cleanup_minio_objects_expired({**metadata, 'data': data_mock})
+
+
+@mark.asyncio
+@patch('laborious.activities.storage.now')
+async def test_cleanup_minio_objects_expired_unparseable_key(mock_now, storage):
+    mock_now.return_value = datetime.datetime(2025, 1, 10, 12, 0, 0)
+    storage.minio_repository.list_objects = AsyncMock(
+        return_value=['some/random/key-without-timestamp.parquet']
+    )
+    storage.minio_repository.delete_file = AsyncMock()
+    storage.send_notification_async = AsyncMock()
+
+    data_mock = MagicMock()
+    data_mock.cleanup_prefix.return_value = 'test'
+    result = await storage.cleanup_minio_objects_expired({**metadata, 'data': data_mock})
+
+    assert result['deleted_count'] == 0
+    assert result['failed_count'] == 0
+    storage.minio_repository.delete_file.assert_not_called()
+
+
+@mark.asyncio
+@patch('laborious.activities.storage.now')
+async def test_cleanup_minio_objects_expired_delete_fails(mock_now, storage):
+    mock_now.return_value = datetime.datetime(2025, 1, 10, 12, 0, 0)
+    old_key = 'training_datasets/m/m-initial-2024-12-01_00-00-00.parquet'
+    storage.minio_repository.list_objects = AsyncMock(return_value=[old_key])
+    storage.minio_repository.delete_file = AsyncMock(side_effect=Exception('delete error'))
+    storage.send_notification_async = AsyncMock()
+
+    data_mock = MagicMock()
+    data_mock.cleanup_prefix.return_value = 'training_datasets/m'
+    result = await storage.cleanup_minio_objects_expired({**metadata, 'data': data_mock})
+
+    assert result['deleted_count'] == 0
+    assert result['failed_count'] == 1
+    assert old_key in result['failed']
+    assert result['failed'][old_key]['success'] is False
+    assert result['failed'][old_key]['message'] == 'delete error'
+
+
+@mark.asyncio
+@patch('laborious.activities.storage.now')
+async def test_cleanup_minio_objects_expired_list_objects_error(mock_now, storage):
+    mock_now.return_value = datetime.datetime(2025, 1, 10, 12, 0, 0)
+    storage.minio_repository.list_objects = AsyncMock(side_effect=Exception('list error'))
+    storage.send_notification_async = AsyncMock()
+    storage.error = MagicMock()
+
+    data_mock = MagicMock()
+    data_mock.cleanup_prefix.return_value = 'training_datasets/m'
+    result = await storage.cleanup_minio_objects_expired({**metadata, 'data': data_mock})
+
+    assert result['deleted_count'] == 0
+    assert result['failed_count'] == 0
+    storage.send_notification_async.assert_called_once_with(
+        metadata=metadata['metadata'],
+        notification_id='ERROR_CLEANUP_MINIO_OBJECTS_EXPIRED',
+        message='Error cleaning up MinIO objects: list error',
+        block='cleanup_minio_objects_expired',
+        level=NotificationLevel.ERROR,
+        attachment_content=ANY,
+    )
+    storage.error.assert_called_once()

tests/laborious/utils/filters/test_conditional_filters.py

@@ -1,23 +1,36 @@
 from pandas import DataFrame
 
 from laborious.utils.filters.conditional_filters import (
+    filter_empty_data,
     filter_specific_variables_null_values,
-    filter_empty_data
 )
 
 
 def test_filter_specific_variables_null_values():
-    assert filter_specific_variables_null_values(
-        DataFrame(
-            {'variable': ['variable1', 'variable2'], 'value': [1, 2]}),
-        config={'variables': ['variable2']}) is False
+    assert (
+        filter_specific_variables_null_values(
+            DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, 2]}),
+            config={'variables': ['variable2']},
+        )
+        is False
+    )
+
+
+def test_filter_specific_variables_null_values_with_empty_data():
+    assert (
+        filter_specific_variables_null_values(DataFrame(), config={'variables': ['variable2']})
+        is False
+    )
 
 
 def test_filter_specific_variables_null_values_with_null_values():
-    assert filter_specific_variables_null_values(
-        DataFrame(
-            {'variable': ['variable1', 'variable2'], 'value': [1, None]}),
-        config={'variables': ['variable2']}) is True
+    assert (
+        filter_specific_variables_null_values(
+            DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, None]}),
+            config={'variables': ['variable2']},
+        )
+        is True
+    )
 
 
 def test_filter_empty_data():
@@ -25,6 +38,7 @@ def test_filter_empty_data():
 
 
 def test_filter_empty_data_with_data():
-    assert filter_empty_data(
-        DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, 2]}),
-        {}) is False
+    assert (
+        filter_empty_data(DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, 2]}), {})
+        is False
+    )

tests/laborious/utils/filters/test_mlflow_filters.py

@@ -1,22 +1,23 @@
 from pandas import DataFrame
+
 from laborious.utils.filters.mlflow_filters import api_error_filter, nan_values_filter
 
 
 def test_api_error_filter_invalid_response():
-    assert api_error_filter(None, {}) == True  # NOSONAR
+    assert api_error_filter(None, {}) is True  # NOSONAR
 
 
 def test_api_error_filter_valid_response_fail():
-    assert api_error_filter({'success': False}, {}) == True
+    assert api_error_filter({'success': False}, {}) is True
 
 
 def test_api_error_filter_valid_response_success():
-    assert api_error_filter({'success': True}, {}) == False
+    assert api_error_filter({'success': True}, {}) is False
 
 
 def test_nan_values_filter_all_nan_values():
-    assert nan_values_filter(DataFrame({'variable': [None, None]}), {}) == True
+    assert nan_values_filter(DataFrame({'variable': [None, None]}), {}) is True
 
 
 def test_nan_values_filter_no_nan_values():
-    assert nan_values_filter(DataFrame({'variable': [1, 2]}), {}) == False
+    assert nan_values_filter(DataFrame({'variable': [1, 2]}), {}) is False

tests/laborious/utils/models/test_minio_dataframe_payload.py

@@ -0,0 +1,266 @@
+from datetime import datetime
+from io import BytesIO
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+from pandas import DataFrame
+
+from laborious.utils.models.minio_dataframe_payload import (
+    MinioDataFramePayload,
+    _build_object_key,
+)
+
+
+def test_parse_object_timestamp_hyphenated_model():
+    key = 'predictions/sched/my-long-model-initial-2024-06-15_10-30-45.parquet'
+    ts = MinioDataFramePayload.parse_object_timestamp(key)
+    assert ts == datetime(2024, 6, 15, 10, 30, 45)
+
+
+def test_parse_object_timestamp_transform():
+    key = 'p/m-transform-2024-01-02_03-04-05.parquet'
+    ts = MinioDataFramePayload.parse_object_timestamp(key)
+    assert ts == datetime(2024, 1, 2, 3, 4, 5)
+
+
+def test_parse_object_timestamp_invalid():
+    assert MinioDataFramePayload.parse_object_timestamp('bad.parquet') is None
+
+
+def test_estimate_size_bytes_returns_positive_for_nonempty_frame():
+    df = DataFrame({'a': [1, 2]})
+    size = MinioDataFramePayload.estimate_size_bytes(df)
+    assert isinstance(size, int)
+    assert size > 0
+
+
+def test_cleanup_prefix_when_offloaded_returns_object_prefix():
+    payload = MinioDataFramePayload(
+        last_timestamp='t',
+        data=None,
+        object_key='training_datasets/m/m-initial-2024-01-01_00-00-00.parquet',
+        object_prefix='training_datasets/m',
+    )
+    assert MinioDataFramePayload.cleanup_prefix(payload) == 'training_datasets/m'
+
+
+def test_cleanup_prefix_when_inline_returns_none():
+    payload = MinioDataFramePayload(last_timestamp='t', data={'x': [1]}, object_key=None)
+    assert MinioDataFramePayload.cleanup_prefix(payload) is None
+
+
+def test_has_data_true_when_object_key_set():
+    payload = MinioDataFramePayload(last_timestamp='t', data=None, object_key='k')
+    assert payload.has_data() is True
+
+
+@pytest.mark.asyncio
+async def test_retrieve_inline_dict_as_dataframe():
+    payload = MinioDataFramePayload(last_timestamp='t', data={'a': [1, 2]})
+    minio = AsyncMock()
+    out = await payload.retrieve(minio, {'metadata': {}})
+    assert list(out.columns) == ['a']
+    minio.download_file.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_retrieve_downloads_parquet_when_offloaded():
+    source = DataFrame({'a': [1, 2]})
+    buf = BytesIO()
+    source.to_parquet(buf, engine='pyarrow', index=True)
+    file_bytes = buf.getvalue()
+
+    payload = MinioDataFramePayload(
+        last_timestamp='t',
+        data=None,
+        object_key='training_datasets/m/f.parquet',
+        object_prefix='training_datasets/m',
+    )
+    minio = AsyncMock()
+    minio.download_file = AsyncMock(return_value=file_bytes)
+
+    out = await payload.retrieve(minio, {'metadata': {}})
+
+    minio.download_file.assert_awaited_once_with(
+        object_name='training_datasets/m/f.parquet',
+        metadata={'metadata': {}},
+    )
+    assert list(out.columns) == ['a']
+
+
+def test_build_object_key():
+    key, prefix = _build_object_key('my-model', 'initial', '2024-01-01_00-00-00')
+    assert key == 'prediction_datasets/my-model/my-model-initial-2024-01-01_00-00-00.parquet'
+    assert prefix == 'prediction_datasets/my-model'
+
+
+def test_build_object_key_strips_slashes():
+    key, prefix = _build_object_key(' /my-model/ ', 'transform', '2024-06-15_10-30-45')
+    assert prefix == 'prediction_datasets/my-model'
+    assert key.startswith('prediction_datasets/my-model/')
+
+
+def test_estimate_size_bytes_fallback():
+    df = DataFrame({'a': [1, 2]})
+    with patch.object(df, 'to_dict', side_effect=RuntimeError('to_dict failed')):
+        size = MinioDataFramePayload.estimate_size_bytes(df)
+    assert isinstance(size, int)
+    assert size > 0
+
+
+def test_parse_object_timestamp_bad_datetime():
+    key = 'p/m-initial-9999-99-99_99-99-99.parquet'
+    assert MinioDataFramePayload.parse_object_timestamp(key) is None
+
+
+@pytest.mark.asyncio
+async def test_retrieve_empty_when_no_data():
+    payload = MinioDataFramePayload(last_timestamp='t', data=None, object_key=None)
+    minio = AsyncMock()
+    out = await payload.retrieve(minio, {})
+    assert out.empty
+    minio.download_file.assert_not_called()
+
+
+@pytest.mark.asyncio
+@patch('laborious.utils.models.minio_dataframe_payload.now')
+async def test_from_dataframe_none(mock_now):
+    mock_now.return_value = datetime(2024, 1, 1, 0, 0, 0)
+    minio = AsyncMock()
+    result = await MinioDataFramePayload.from_dataframe(
+        dataframe=None,
+        minio_repo=minio,
+        model_name='m',
+        operation='initial',
+        status={'success': False, 'message': 'no data'},
+    )
+    assert result.data is None
+    assert result.status == {'success': False, 'message': 'no data'}
+    assert result.object_key is None
+
+
+@pytest.mark.asyncio
+@patch('laborious.utils.models.minio_dataframe_payload.now')
+async def test_from_dataframe_empty(mock_now):
+    mock_now.return_value = datetime(2024, 1, 1, 0, 0, 0)
+    minio = AsyncMock()
+    mock_df = MagicMock()
+    mock_df.__bool__ = MagicMock(return_value=True)
+    mock_df.empty = True
+    result = await MinioDataFramePayload.from_dataframe(
+        dataframe=mock_df,
+        minio_repo=minio,
+        model_name='m',
+        operation='initial',
+    )
+    assert result.data is None
+    assert result.object_key is None
+
+
+def _mock_dataframe(data_dict, timestamp_values=None):
+    """Build a MagicMock that behaves enough like a DataFrame for from_dataframe."""
+    mock_df = MagicMock()
+    mock_df.__bool__ = MagicMock(return_value=True)
+    mock_df.empty = False
+    if timestamp_values is None:
+        timestamp_values = data_dict.get('timestamp', ['2024-01-01'])
+    ts_col = MagicMock()
+    ts_col.values.tolist.return_value = timestamp_values
+    mock_df.__getitem__ = MagicMock(return_value=ts_col)
+    mock_df.to_dict.return_value = data_dict
+    buf = BytesIO()
+    DataFrame(data_dict).to_parquet(buf, engine='pyarrow', index=True)
+    mock_df.to_parquet = MagicMock(side_effect=lambda b, **kw: b.write(buf.getvalue()))
+    return mock_df
+
+
+@pytest.mark.asyncio
+@patch('laborious.utils.models.minio_dataframe_payload.OFFLOAD_THRESHOLD_BYTES', 10**9)
+async def test_from_dataframe_inline():
+    minio = AsyncMock()
+    df = _mock_dataframe({'timestamp': ['2024-01-01'], 'value': [42]})
+    result = await MinioDataFramePayload.from_dataframe(
+        dataframe=df,
+        minio_repo=minio,
+        model_name='m',
+        operation='initial',
+    )
+    assert result.data is not None
+    assert result.object_key is None
+    assert result.last_timestamp == '2024-01-01'
+
+
+@pytest.mark.asyncio
+@patch('laborious.utils.models.minio_dataframe_payload.now')
+@patch('laborious.utils.models.minio_dataframe_payload.OFFLOAD_THRESHOLD_BYTES', 0)
+async def test_from_dataframe_offloaded(mock_now):
+    mock_now.return_value = datetime(2024, 1, 1, 0, 0, 0)
+    minio = AsyncMock()
+    minio.upload_file = AsyncMock(return_value={'minio_object_name': 'full/key.parquet'})
+    minio.bucket = 'test-bucket'
+
+    df = _mock_dataframe({'timestamp': ['2024-01-01'], 'value': [42]})
+    result = await MinioDataFramePayload.from_dataframe(
+        dataframe=df,
+        minio_repo=minio,
+        model_name='m',
+        operation='initial',
+        workflow_metadata={'wf': 'data'},
+    )
+    assert result.data is None
+    assert result.object_key == 'full/key.parquet'
+    assert result.bucket == 'test-bucket'
+    assert result.uri == 's3://test-bucket/full/key.parquet'
+    minio.upload_file.assert_awaited_once()
+
+
+def test_from_dict_inline():
+    raw = {
+        'last_timestamp': '2024-01-01T00:00:00+00:00',
+        'status': None,
+        'data': {'col1': {0: 'val1'}},
+        'bucket': None,
+        'object_key': None,
+        'object_prefix': None,
+        'uri': None,
+    }
+    payload = MinioDataFramePayload.from_dict(raw)
+    assert isinstance(payload, MinioDataFramePayload)
+    assert payload.last_timestamp == '2024-01-01T00:00:00+00:00'
+    assert payload.data == {'col1': {0: 'val1'}}
+    assert payload.object_key is None
+
+
+def test_from_dict_offloaded():
+    raw = {
+        'last_timestamp': '2024-06-15T10:30:45+00:00',
+        'status': {'success': True},
+        'data': None,
+        'bucket': 'my-bucket',
+        'object_key': 'training_datasets/model/model-initial-2024-06-15_10-30-45.parquet',
+        'object_prefix': 'training_datasets/model',
+        'uri': 's3://my-bucket/training_datasets/model/model-initial-2024-06-15_10-30-45.parquet',
+    }
+    payload = MinioDataFramePayload.from_dict(raw)
+    assert isinstance(payload, MinioDataFramePayload)
+    assert payload.data is None
+    assert payload.bucket == 'my-bucket'
+    assert payload.object_key == raw['object_key']
+    assert payload.object_prefix == 'training_datasets/model'
+    assert payload.uri == raw['uri']
+    assert payload.status == {'success': True}
+
+
+def test_from_dict_minimal_keys():
+    raw = {'last_timestamp': '2024-01-01'}
+    payload = MinioDataFramePayload.from_dict(raw)
+    assert payload.last_timestamp == '2024-01-01'
+    assert payload.data is None
+    assert payload.bucket is None
+    assert payload.object_key is None
+
+
+def test_from_dict_passthrough_existing_instance():
+    original = MinioDataFramePayload(last_timestamp='2024-01-01', data={'a': 1}, bucket='b')
+    result = MinioDataFramePayload.from_dict(original)
+    assert result is original

tests/laborious/utils/repository/test_opc_repository.py

@@ -1,9 +1,20 @@
-import pytest
-from unittest.mock import AsyncMock, Mock, patch, MagicMock, ANY, call
-from asyncua.crypto.security_policies import SecurityPolicyBasic256
-from laborious.utils.repository.opc_repository import OpcRepository
-from sientia_do.notifications.models import NotificationLevel
+import asyncio
+import json
 from datetime import datetime
+from unittest.mock import ANY, AsyncMock, MagicMock, Mock, patch
+
+import pytest
+from asyncua.crypto.security_policies import SecurityPolicyBasic256
+from asyncua.ua.uaerrors import BadNodeIdUnknown, BadSessionIdInvalid
+from sientia_do.notifications.models import NotificationLevel
+
+from laborious.utils.repository.opc_repository import (
+    OpcClientAlreadyExistsError,
+    OpcClientNotInitializedError,
+    OpcRepository,
+    OpcSessionAlreadyConnectedError,
+    is_reconnectable_opcua_bad,
+)
 
 
 @pytest.fixture
@@ -13,17 +24,29 @@ def mock_logger():
 
 @pytest.fixture
 def opc_repository(mock_logger):
-    return OpcRepository(
-        id="test_repo",
-        url="opc.tcp://localhost:4840",
+    repository = OpcRepository(
+        opc_id='test_repo',
+        server_name='test_server',
+        url='opc.tcp://localhost:4840',
         logger=mock_logger,
         notification_handler=Mock(),
         reconnection_interval=60,
-        server_uri="urn:test:server",
-        cert_path="/path/to/cert.pem",
-        private_key_path="/path/to/key.pem",
-        server_cert_path="/path/to/server_cert.pem"
+        server_uri='urn:test:server',
+        cert_path='/path/to/cert.pem',
+        private_key_path='/path/to/key.pem',
+        server_cert_path='/path/to/server_cert.pem',
+        metrics_controller=AsyncMock(),
     )
+    repository.disconnection_interval = 0.1
+    repository.send_notification = MagicMock()
+    repository.send_notification_async = AsyncMock()
+    repository.emit_metric = AsyncMock()
+    repository.info = MagicMock()
+    repository.error = MagicMock()
+    repository.warning = MagicMock()
+    repository.debug = MagicMock()
+    repository._session_ready.set()
+    return repository
 
 
 @pytest.fixture
@@ -35,26 +58,26 @@ def mock_client():
 
 
 metadata = {
-    "metadata": {
-        "model_id": "test_model",
-        "model_name": "test_model",
-        "workflow_name": "test_workflow",
-        "schema_name": "test_schedule",
+    'metadata': {
+        'model_id': 'test_model',
+        'model_name': 'test_model',
+        'workflow_name': 'test_workflow',
+        'schema_name': 'test_schedule',
     },
 }
 
 
 def test_init(opc_repository):
-    assert opc_repository.id == "test_repo"
-    assert opc_repository.url == "opc.tcp://localhost:4840"
-    assert opc_repository.server_uri == "urn:test:server"
-    assert opc_repository.cert_path == "/path/to/cert.pem"
-    assert opc_repository.private_key_path == "/path/to/key.pem"
-    assert opc_repository.server_cert_path == "/path/to/server_cert.pem"
+    assert opc_repository.id == 'test_repo'
+    assert opc_repository.server_name == 'test_server'
+    assert opc_repository.url == 'opc.tcp://localhost:4840'
+    assert opc_repository.server_uri == 'urn:test:server'
+    assert opc_repository.cert_path == '/path/to/cert.pem'
+    assert opc_repository.private_key_path == '/path/to/key.pem'
+    assert opc_repository.server_cert_path == '/path/to/server_cert.pem'
     assert opc_repository.reconnection_interval == 60
     assert opc_repository.client is None
     assert opc_repository.last_reconnection_time is None
-    assert opc_repository.error_count == 0
 
 
 @pytest.mark.asyncio
@@ -62,15 +85,15 @@ async def test_set_security(opc_repository, mock_client):
     opc_repository.client = mock_client
     await opc_repository.set_security()
 
-    mock_client.application_uri = "urn:test:server"
+    mock_client.application_uri = 'urn:test:server'
     mock_client.set_security.assert_called_once_with(
         SecurityPolicyBasic256,
-        certificate="/path/to/cert.pem",
-        private_key="/path/to/key.pem",
-        server_certificate="/path/to/server_cert.pem"
+        certificate='/path/to/cert.pem',
+        private_key='/path/to/key.pem',
+        server_certificate='/path/to/server_cert.pem',
     )
-    assert mock_client.secure_channel_timeout == 10000000
-    assert mock_client.session_timeout == 10000000
+    assert mock_client.secure_channel_timeout == 600_000
+    assert mock_client.session_timeout == 600_000
 
 
 @pytest.mark.asyncio
@@ -81,68 +104,160 @@ async def test_set_security_missing_certificates(opc_repository):
     try:
         await opc_repository.set_security()
     except ValueError as e:
-        assert str(
-            e) == "Certificate and private key paths must be provided for secure connection."
+        assert str(e) == 'Certificate and private key paths must be provided for secure connection.'
+
+
+@pytest.mark.asyncio
+async def test_set_security_missing_client(opc_repository):
+    opc_repository.client = None
+    try:
+        await opc_repository.set_security()
+    except ValueError as e:
+        assert str(e) == 'Client must be initialized before setting security'
 
 
 @pytest.mark.asyncio
 async def test_connect_with_security(opc_repository, mock_client):
-    opc_repository.try_connect = AsyncMock(return_value=(True, {}))
+    opc_repository._create_client = AsyncMock()
+    opc_repository._open_session = AsyncMock(return_value=(True, {}))
     result = await opc_repository.connect()
 
-    opc_repository.try_connect.assert_called_once()
-    assert opc_repository.client == mock_client
+    opc_repository._create_client.assert_called_once()
+    opc_repository._open_session.assert_called_once()
     assert result == (True, {})
 
 
 @pytest.mark.asyncio
 async def test_connect_without_security(opc_repository, mock_client):
     opc_repository.cert_path = None
-    opc_repository.try_connect = AsyncMock(return_value=(True, {}))
+    opc_repository._create_client = AsyncMock()
+    opc_repository._open_session = AsyncMock(return_value=(True, {}))
     opc_repository.set_security = AsyncMock()
     result = await opc_repository.connect()
 
-    opc_repository.try_connect.assert_called_once()
+    opc_repository._create_client.assert_called_once()
+    opc_repository._open_session.assert_called_once()
     opc_repository.set_security.assert_not_called()
-    assert opc_repository.client == mock_client
     assert result == (True, {})
 
 
 @pytest.mark.asyncio
-async def test_try_connect_success(opc_repository):
-    opc_repository.last_reconnection_time = None
+async def test_connect_raises_when_session_already_open(opc_repository, mock_client):
+    opc_repository.client = mock_client
+    proto = MagicMock()
+    proto.state = 'open'
+    mock_client.uaclient = MagicMock(protocol=proto)
+
+    with pytest.raises(OpcSessionAlreadyConnectedError, match='disconnect'):
+        await opc_repository.connect()
+
+
+@pytest.mark.asyncio
+async def test_create_client_raises_when_client_exists(opc_repository, mock_client):
+    opc_repository.client = mock_client
+
+    with pytest.raises(OpcClientAlreadyExistsError, match='already exists'):
+        await opc_repository._create_client()
+
+
+@pytest.mark.asyncio
+async def test_open_session_success(opc_repository):
+    closed_proto = MagicMock()
+    closed_proto.state = 'closed'
     opc_repository.client = AsyncMock()
-    result = await opc_repository.try_connect()
+    opc_repository.client.uaclient = MagicMock(protocol=closed_proto)
+    opc_repository.client.session_timeout = 600_000
+    opc_repository.client.secure_channel_timeout = 600_000
+
+    open_proto = MagicMock()
+    open_proto.state = 'open'
+    open_proto.authentication_token = 'tok'
+
+    async def connect_side_effect():
+        opc_repository.client.uaclient.protocol = open_proto
+
+    opc_repository.client.connect = AsyncMock(side_effect=connect_side_effect)
+
+    result = await opc_repository._open_session()
 
     opc_repository.client.connect.assert_called_once()
-    assert opc_repository.last_reconnection_time is not None
+    assert opc_repository.last_reconnection_time is None
     assert result == (True, {})
+    assert opc_repository._session_ready.is_set()
 
 
 @pytest.mark.asyncio
-async def test_try_connect_fail(opc_repository):
-    opc_repository.last_reconnection_time = None
+async def test_open_session_raises_when_already_connected(opc_repository, mock_client):
+    opc_repository.client = mock_client
+    proto = MagicMock()
+    proto.state = 'open'
+    mock_client.uaclient = MagicMock(protocol=proto)
+
+    with pytest.raises(OpcSessionAlreadyConnectedError, match='disconnect'):
+        await opc_repository._open_session()
+
+
+@pytest.mark.asyncio
+async def test_open_session_fail(opc_repository):
+    opc_repository._disconnect_locked = AsyncMock()
     opc_repository.client = MagicMock()
-    opc_repository.client.connect.side_effect = Exception("Test error")
+    opc_repository.client.uaclient = MagicMock(protocol=MagicMock(state='closed'))
+    opc_repository.client.connect = AsyncMock(side_effect=Exception('Test error'))
 
-    is_connected, error_data = await opc_repository.try_connect()
+    is_connected, error_data = await opc_repository._open_session()
 
+    opc_repository._disconnect_locked.assert_called_once()
     opc_repository.client.connect.assert_called_once()
     assert is_connected is False
-    assert error_data['notification_id'] == f"OPC_CONNECTION_ERROR_{opc_repository.id}"
-    assert error_data['message'] == "Failed to connect to OPC server: Test error"
-    assert error_data['block'] == "opc_repository"
+    assert error_data['notification_id'] == f'OPC_CONNECTION_ERROR_{opc_repository.id}'
+    assert error_data['message'] == 'Failed to connect to OPC server: Test error'
+    assert error_data['block'] == 'opc_repository'
     assert error_data['level'] == NotificationLevel.ERROR
     assert error_data['attachment_content'] is not None
 
 
 @pytest.mark.asyncio
-async def test_disconnect(opc_repository, mock_client):
+async def test_open_session_raises_when_no_client(opc_repository):
+    opc_repository.client = None
+
+    with pytest.raises(OpcClientNotInitializedError, match='not initialized'):
+        await opc_repository._open_session()
+
+
+@pytest.mark.asyncio
+async def test_disconnection_fallback_success(opc_repository, mock_client):
     opc_repository.client = mock_client
-    await opc_repository.disconnect()
+    mock_client.disconnect.return_value = True
+    result = await opc_repository._disconnection_fallback()
 
     mock_client.disconnect.assert_called_once()
+    assert result == []
+
+
+@pytest.mark.asyncio
+async def test_disconnection_fallback_fail(opc_repository, mock_client):
+    opc_repository.client = mock_client
+    mock_client.disconnect.side_effect = Exception('Test error')
+    result = await opc_repository._disconnection_fallback()
+    assert result == [
+        {'attempt': 1, 'error': 'Test error', 'traceback': ANY},
+        {'attempt': 2, 'error': 'Test error', 'traceback': ANY},
+        {'attempt': 3, 'error': 'Test error', 'traceback': ANY},
+        {'attempt': 4, 'error': 'Test error', 'traceback': ANY},
+        {'attempt': 5, 'error': 'Test error', 'traceback': ANY},
+    ]
+    assert mock_client.disconnect.call_count == 5
+
+
+@pytest.mark.asyncio
+async def test_disconnect(opc_repository, mock_client):
+    opc_repository.client = mock_client
+    opc_repository._disconnection_fallback = AsyncMock(return_value=[])
+    await opc_repository.disconnect()
+
+    opc_repository._disconnection_fallback.assert_called_once()
     assert opc_repository.client is None
+    assert opc_repository._allow_reconnect is False
 
 
 @pytest.mark.asyncio
@@ -154,12 +269,21 @@ async def test_disconnect_no_client(opc_repository):
 @pytest.mark.asyncio
 async def test_disconnect_error(opc_repository, mock_client):
     opc_repository.client = mock_client
-    mock_client.disconnect.side_effect = Exception("Test error")
+    opc_repository._disconnection_fallback = AsyncMock(
+        return_value=[{'attempt': 1, 'error': 'Test error', 'traceback': 'text'}]
+    )
     await opc_repository.disconnect()
 
-    opc_repository.logger.custom_error.assert_called_once_with(
-        "Failed to disconnect from OPC server: Test error",
-        ANY
+    opc_repository._disconnection_fallback.assert_called_once()
+    opc_repository.send_notification_async.assert_called_once_with(
+        metadata=opc_repository.metadata,
+        notification_id=f'OPC_DISCONNECTION_ERROR_{opc_repository.id}',
+        message='Failed to disconnect from OPC server in 5 attempts.',
+        block='opc_repository',
+        level=NotificationLevel.ERROR,
+        attachment_content=json.dumps(
+            [{'attempt': 1, 'error': 'Test error', 'traceback': 'text'}], indent=4
+        ),
     )
     assert opc_repository.client is None
 
@@ -167,93 +291,26 @@ async def test_disconnect_error(opc_repository, mock_client):
 @pytest.mark.asyncio
 async def test_validate_connection_none_client(opc_repository):
     opc_repository.client = None
-    opc_repository.connect = AsyncMock(return_value=(True, {}))
     response = await opc_repository.validate_connection()
-    assert response == (True, {})
-    opc_repository.connect.assert_called_once()
+    assert response == (False, opc_repository._not_connected_error())
 
 
 @pytest.mark.asyncio
-async def test_validate_connection_error_count_disconnect_error(opc_repository):
-    opc_repository.error_count = 6
-    opc_repository.client = AsyncMock()
-    opc_repository.disconnect = AsyncMock(
-        side_effect=Exception("Test error")
-    )
-    opc_repository.connect = AsyncMock(return_value=(True, {}))
-
-    response = await opc_repository.validate_connection()
-    assert response == opc_repository.connect.return_value
-    opc_repository.disconnect.assert_called_once()
-    opc_repository.connect.assert_called_once()
-    opc_repository.logger.custom_error.assert_has_calls(
-        [
-            call("Failed to disconnect from OPC server: Test error", ANY),
-        ]
-    )
-
-
-@pytest.mark.asyncio
-async def test_validate_connection_error_validate_connection_error(opc_repository):
-    opc_repository.client = MagicMock(
-        uaclient=Exception("Test error")
-    )
-    opc_repository.error_count = 0
-
-    response = await opc_repository.validate_connection()
-
-    assert response == (False, {
-        "notification_id": f"OPC_CONNECTION_CHECK_ERROR_{opc_repository.id}",
-        "message": "Failed to validate connection to OPC server: 'Exception' object has no attribute 'protocol'",
-        "block": "opc_repository",
-        "level": NotificationLevel.ERROR,
-        "attachment_content": ANY
-    })
-
-
-@pytest.mark.asyncio
-@patch('laborious.utils.repository.opc_repository.datetime')
-async def test_validate_connection_lost_not_time_to_reconnect(_mock_datetime, opc_repository):
-    _mock_datetime.now = MagicMock(
-        return_value=datetime(2025, 1, 1, 0, 0, 0))
-    opc_repository.error_count = 0
+async def test_validate_connection_session_not_open(opc_repository):
     opc_repository.client = MagicMock()
     opc_repository.client.uaclient.protocol = None
-    opc_repository.last_reconnection_time = datetime(2025, 1, 1, 0, 0, 0)
-    opc_repository.connect = MagicMock(return_value=(True, {}))
 
     response = await opc_repository.validate_connection()
-    opc_repository.connect.assert_not_called()
-    assert response == (False, {
-        "notification_id": f"OPC_CONNECTION_AWAITING_RECONNECTION_WINDOW_{opc_repository.id}",
-        "message": f"OPC server {opc_repository.id} is not connected, waiting for next reconnection window...",
-        "block": "opc_repository",
-        "level": NotificationLevel.WARNING
-    })
 
-
-@pytest.mark.asyncio
-@patch('laborious.utils.repository.opc_repository.datetime')
-async def test_validate_connection_lost_time_to_reconnect(mock_datetime, opc_repository):
-    mock_datetime.now = MagicMock(
-        return_value=datetime(2025, 1, 1, 1, 0, 0))
-    opc_repository.error_count = 0
-    opc_repository.client = AsyncMock()
-    opc_repository.client.uaclient.protocol = None
-    opc_repository.last_reconnection_time = datetime(2025, 1, 1, 0, 0, 0)
-    opc_repository.connect = AsyncMock(return_value=(True, {}))
-
-    response = await opc_repository.validate_connection()
-    opc_repository.connect.assert_called_once()
-    assert response == opc_repository.connect.return_value
+    assert response == (False, opc_repository._not_connected_error())
+    opc_repository.error.assert_called_once()
 
 
 @pytest.mark.asyncio
 async def test_validate_connection_success(opc_repository):
     opc_repository.client = MagicMock()
-    opc_repository.error_count = 0
     opc_repository.client.uaclient.protocol = MagicMock()
-    opc_repository.client.uaclient.protocol.state = "open"
+    opc_repository.client.uaclient.protocol.state = 'open'
 
     output = await opc_repository.validate_connection()
     assert output == (True, {})
@@ -262,51 +319,53 @@ async def test_validate_connection_success(opc_repository):
 @pytest.mark.asyncio
 async def test_write_data_validate_connection_do_nothing(opc_repository):
     opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
-    opc_repository.client = AsyncMock(
-        get_node=MagicMock()
-    )
+    opc_repository.client = AsyncMock(get_node=MagicMock())
     mock_node = AsyncMock()
     opc_repository.client.get_node.return_value = mock_node
 
-    result = await opc_repository.write_data("ns=2;s=TestNode", 42.0,
-                                             "float", opc_repository.logger, metadata['metadata'])
+    result = await opc_repository.write_data('ns=2;s=TestNode', 42.0, 'float', metadata['metadata'])
 
     opc_repository.validate_connection.assert_called_once()
-    opc_repository.client.get_node.assert_called_once_with("ns=2;s=TestNode")
-    assert result == (True, {})
+    opc_repository.client.get_node.assert_called_once_with('ns=2;s=TestNode')
+    assert result == (True, {'response_time': ANY})
 
 
 @pytest.mark.asyncio
 async def test_write_data_validate_connection_failed(opc_repository):
-    opc_repository.validate_connection = AsyncMock(return_value=(False, {}))
-    opc_repository.client = AsyncMock()
-    opc_repository.error_count = 0
+    opc_repository.client = MagicMock()
+    opc_repository.client.uaclient.protocol = MagicMock(state='closed')
+    opc_repository._start_reconnect = AsyncMock()
 
-    result = await opc_repository.write_data("ns=2;s=TestNode", 42.0,
-                                             "float", opc_repository.logger, metadata['metadata'])
+    is_success, error_data = await opc_repository.write_data(
+        'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
+    )
 
-    opc_repository.validate_connection.assert_called_once()
-    opc_repository.client.get_node.assert_not_called()
-    assert result == (False, {})
+    opc_repository._start_reconnect.assert_called_once()
+    assert opc_repository._start_reconnect.call_args.args[0] == 'ProtocolClosed'
+    assert is_success is False
+    assert error_data['opc_error_kind'] == 'connection_lost'
+    assert error_data['opc_status'] == 'ProtocolClosed'
 
 
 @pytest.mark.asyncio
 async def test_write_data_get_node_failed(opc_repository):
     opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
     opc_repository.client = AsyncMock()
-    opc_repository.error_count = 0
-    opc_repository.client.get_node = MagicMock(
-        side_effect=Exception("Test error"))
+    opc_repository.client.get_node = MagicMock(side_effect=Exception('Test error'))
 
-    is_success, error_data = await opc_repository.write_data("ns=2;s=TestNode", 42.0,
-                                                             "float", opc_repository.logger, metadata['metadata'])
+    is_success, error_data = await opc_repository.write_data(
+        'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
+    )
 
     opc_repository.validate_connection.assert_called_once()
-    opc_repository.client.get_node.assert_called_once_with("ns=2;s=TestNode")
+    opc_repository.client.get_node.assert_called_once_with('ns=2;s=TestNode')
     assert is_success is False
-    assert error_data['notification_id'] == f"OPC_WRITE_GET_NODE_ERROR_{opc_repository.id}"
-    assert error_data['message'] == "Failed to get node from OPC server: Test error | metadata: {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}"
-    assert error_data['block'] == "opc_repository"
+    assert error_data['notification_id'] == f'OPC_WRITE_GET_NODE_ERROR_{opc_repository.id}'
+    assert (
+        error_data['message']
+        == "Failed to get node from OPC server: Test error | metadata: {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}"
+    )
+    assert error_data['block'] == 'opc_repository'
     assert error_data['level'] == NotificationLevel.ERROR
     assert error_data['attachment_content'] is not None
 
@@ -318,51 +377,36 @@ async def test_write_data_invalid_data_type(opc_repository, mock_client):
     mock_node = AsyncMock()
     mock_client.get_node = MagicMock(return_value=mock_node)
 
-    is_success, error_data = await opc_repository.write_data("ns=2;s=TestNode", 42.0,
-                                                             "invalid_type", opc_repository.logger, metadata['metadata'])
+    is_success, error_data = await opc_repository.write_data(
+        'ns=2;s=TestNode', 42.0, 'invalid_type', metadata['metadata']
+    )
 
     opc_repository.validate_connection.assert_called_once()
-    mock_client.get_node.assert_called_once_with("ns=2;s=TestNode")
+    mock_client.get_node.assert_called_once_with('ns=2;s=TestNode')
 
     assert is_success is False
-    assert error_data['notification_id'] == f"OPC_WRITE_DATA_TYPE_ERROR_{opc_repository.id}"
-    assert error_data['message'] == "Unsupported data type: invalid_type | metadata: {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}"
-    assert error_data['block'] == "opc_repository"
+    assert error_data['notification_id'] == f'OPC_WRITE_DATA_TYPE_ERROR_{opc_repository.id}'
+    assert (
+        error_data['message']
+        == "Unsupported data type: invalid_type | metadata: {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}"
+    )
+    assert error_data['block'] == 'opc_repository'
     assert error_data['level'] == NotificationLevel.ERROR
     assert error_data.get('attachment_content') is None
 
 
 @pytest.mark.asyncio
-@patch('laborious.utils.repository.opc_repository.metrics')
-async def test_write_data(mock_metrics, opc_repository, mock_client):
+async def test_write_data(opc_repository, mock_client):
     opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
     opc_repository.client = mock_client
     mock_node = AsyncMock()
     mock_client.get_node = MagicMock(return_value=mock_node)
 
-    result = await opc_repository.write_data("ns=2;s=TestNode", 42.0,
-                                             "float", opc_repository.logger, metadata['metadata'])
+    result = await opc_repository.write_data('ns=2;s=TestNode', 42.0, 'float', metadata['metadata'])
 
-    mock_client.get_node.assert_called_once_with("ns=2;s=TestNode")
+    mock_client.get_node.assert_called_once_with('ns=2;s=TestNode')
     mock_node.write_value.assert_called_once()
-    assert result == (True, {})
-
-    mock_metrics.PREDICTION_OPC_WRITING_COUNT.labels.assert_called_once_with(
-        pod_id=opc_repository.pod_id,
-        model_name=metadata['metadata']['model_name'],
-        pipeline_name=metadata['metadata']['workflow_name'],
-        opc_server_id=opc_repository.id
-    )
-    mock_metrics.PREDICTION_OPC_WRITING_COUNT.labels.return_value.inc.assert_called_once_with()
-
-    mock_metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR.labels.assert_called_once_with(
-        pod_id=opc_repository.pod_id,
-        model_name=metadata['metadata']['model_name'],
-        pipeline_name=metadata['metadata']['workflow_name'],
-        opc_server_id=opc_repository.id
-    )
-    mock_metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR.labels.return_value.observe.assert_called_once_with(
-        ANY)
+    assert result == (True, {'response_time': ANY})
 
 
 @pytest.mark.asyncio
@@ -370,19 +414,197 @@ async def test_write_data_write_value_failed(opc_repository, mock_client):
     opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
     opc_repository.client = mock_client
     mock_node = AsyncMock()
-    opc_repository.error_count = 0
     mock_client.get_node = MagicMock(return_value=mock_node)
-    mock_node.write_value.side_effect = Exception("Test error")
+    mock_node.write_value.side_effect = Exception('Test error')
 
-    is_success, error_data = await opc_repository.write_data("ns=2;s=TestNode", 42.0,
-                                                             "float", opc_repository.logger, metadata['metadata'])
+    is_success, error_data = await opc_repository.write_data(
+        'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
+    )
 
     opc_repository.validate_connection.assert_called_once()
-    mock_client.get_node.assert_called_once_with("ns=2;s=TestNode")
+    mock_client.get_node.assert_called_once_with('ns=2;s=TestNode')
     mock_node.write_value.assert_called_once()
     assert is_success is False
-    assert error_data['notification_id'] == f"OPC_WRITE_DATA_ERROR_{opc_repository.id}"
-    assert error_data['message'] == "Failed to write data to OPC server: Test error | metadata: {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}"
-    assert error_data['block'] == "opc_repository"
+    assert error_data['notification_id'] == f'OPC_WRITE_DATA_ERROR_{opc_repository.id}'
+    assert (
+        error_data['message']
+        == "Failed to write data to OPC server: Test error | metadata: {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}"
+    )
+    assert error_data['block'] == 'opc_repository'
     assert error_data['level'] == NotificationLevel.ERROR
     assert error_data['attachment_content'] is not None
+
+
+def test_is_reconnectable_opcua_bad():
+    assert is_reconnectable_opcua_bad(BadSessionIdInvalid()) is True
+    assert is_reconnectable_opcua_bad(BadNodeIdUnknown()) is False
+    assert is_reconnectable_opcua_bad(Exception('other')) is False
+
+
+@pytest.mark.asyncio
+async def test_write_data_bad_session_id_invalid_schedules_reconnect(opc_repository, mock_client):
+    opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
+    opc_repository.client = mock_client
+    opc_repository._start_reconnect = AsyncMock()
+    mock_node = AsyncMock()
+    mock_client.get_node = MagicMock(return_value=mock_node)
+    mock_node.write_value.side_effect = BadSessionIdInvalid()
+
+    is_success, error_data = await opc_repository.write_data(
+        'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
+    )
+
+    mock_node.write_value.assert_called_once()
+    opc_repository._start_reconnect.assert_called_once()
+    assert is_success is False
+    assert error_data['opc_error_kind'] == 'session_bad'
+    assert error_data['opc_status'] == 'BadSessionIdInvalid'
+
+
+@pytest.mark.asyncio
+async def test_write_data_reconnect_in_progress_immediate(opc_repository):
+    opc_repository._session_ready.clear()
+    opc_repository._reconnect_task = asyncio.create_task(asyncio.sleep(60))
+    opc_repository.validate_connection = AsyncMock()
+
+    is_success, error_data = await opc_repository.write_data(
+        'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
+    )
+
+    opc_repository._reconnect_task.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await opc_repository._reconnect_task
+    opc_repository._reconnect_task = None
+
+    opc_repository.validate_connection.assert_not_called()
+    assert is_success is False
+    assert error_data['opc_error_kind'] == 'reconnect_in_progress'
+
+
+@pytest.mark.asyncio
+async def test_start_reconnect_skips_within_interval(opc_repository):
+    opc_repository.last_reconnection_time = datetime.now()
+    opc_repository.reconnection_interval = 3600
+
+    await opc_repository._start_reconnect('BadSessionIdInvalid', 'tok')
+
+    assert opc_repository._reconnect_task is None
+
+
+@pytest.mark.asyncio
+async def test_write_data_protocol_closed_schedules_reconnect(opc_repository):
+    opc_repository.client = MagicMock()
+    opc_repository.client.uaclient.protocol = MagicMock(state='closed')
+    opc_repository._start_reconnect = AsyncMock()
+
+    is_success, error_data = await opc_repository.write_data(
+        'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
+    )
+
+    opc_repository._start_reconnect.assert_called_once()
+    assert opc_repository._start_reconnect.call_args.args[0] == 'ProtocolClosed'
+    assert is_success is False
+    assert error_data['opc_error_kind'] == 'connection_lost'
+    assert error_data['opc_status'] == 'ProtocolClosed'
+
+
+@pytest.mark.asyncio
+async def test_write_data_protocol_closed_skips_reconnect_within_interval(opc_repository):
+    opc_repository.client = MagicMock()
+    opc_repository.client.uaclient.protocol = MagicMock(state='closed')
+    opc_repository.last_reconnection_time = datetime.now()
+    opc_repository.reconnection_interval = 3600
+
+    is_success, error_data = await opc_repository.write_data(
+        'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
+    )
+
+    assert opc_repository._reconnect_task is None
+    assert is_success is False
+    assert error_data['opc_error_kind'] == 'connection_lost'
+
+
+@pytest.mark.asyncio
+async def test_write_data_after_failed_reconnect_schedules_again(opc_repository):
+    opc_repository._session_ready.clear()
+    opc_repository.reconnection_interval = 0
+    opc_repository.last_reconnection_time = None
+    opc_repository._reconnect_locked = AsyncMock(
+        return_value=(False, {'message': 'connect failed'})
+    )
+
+    await opc_repository.write_data('ns=2;s=TestNode', 42.0, 'float', metadata['metadata'])
+    await asyncio.sleep(0.1)
+    assert opc_repository._reconnect_locked.call_count == 1
+    assert not opc_repository._reconnect_task_in_progress()
+
+    await opc_repository.write_data('ns=2;s=TestNode', 42.0, 'float', metadata['metadata'])
+    await asyncio.sleep(0.1)
+    assert opc_repository._reconnect_locked.call_count == 2
+
+
+@pytest.mark.asyncio
+async def test_write_data_after_disconnect_does_not_schedule_reconnect(opc_repository, mock_client):
+    opc_repository.client = mock_client
+    proto = MagicMock()
+    proto.state = 'closed'
+    mock_client.uaclient = MagicMock(protocol=proto)
+    opc_repository._disconnection_fallback = AsyncMock(return_value=[])
+    await opc_repository.disconnect()
+
+    is_success, error_data = await opc_repository.write_data(
+        'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
+    )
+
+    assert opc_repository._reconnect_task is None
+    assert is_success is False
+    assert error_data['opc_error_kind'] == 'connection_lost'
+
+
+@pytest.mark.asyncio
+async def test_parallel_bad_writes_single_reconnect_task(opc_repository, mock_client):
+    opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
+    opc_repository.client = mock_client
+    opc_repository.reconnection_interval = 0
+    opc_repository.last_reconnection_time = None
+    mock_node = AsyncMock()
+    mock_client.get_node = MagicMock(return_value=mock_node)
+    mock_node.write_value.side_effect = BadSessionIdInvalid()
+
+    connect_count = 0
+
+    async def slow_reconnect():
+        nonlocal connect_count
+        connect_count += 1
+        await asyncio.sleep(0.05)
+        opc_repository._session_ready.set()
+        return True, {}
+
+    opc_repository._reconnect_locked = slow_reconnect
+
+    results = await asyncio.gather(
+        opc_repository.write_data('ns=2;s=TestNode', 1.0, 'float', metadata['metadata']),
+        opc_repository.write_data('ns=2;s=TestNode2', 2.0, 'float', metadata['metadata']),
+    )
+    await asyncio.sleep(0.15)
+
+    assert connect_count <= 1
+    assert 1 <= mock_node.write_value.call_count <= 2
+    error_kinds = [r[1].get('opc_error_kind') for r in results]
+    assert error_kinds.count('session_bad') >= 1
+    assert all(k in ('session_bad', 'reconnect_in_progress') for k in error_kinds)
+
+
+@pytest.mark.asyncio
+@patch('laborious.utils.repository.opc_repository.datetime')
+async def test_reconnect_locked_sets_last_reconnection_time(mock_datetime, opc_repository):
+    mock_datetime.now = MagicMock(return_value=datetime(2025, 1, 1, 12, 0, 0))
+    opc_repository._disconnect_locked = AsyncMock()
+    opc_repository._connect_locked = AsyncMock(return_value=(True, {}))
+
+    result = await opc_repository._reconnect_locked()
+
+    opc_repository._disconnect_locked.assert_called_once()
+    opc_repository._connect_locked.assert_called_once()
+    assert result == (True, {})
+    assert opc_repository.last_reconnection_time == datetime(2025, 1, 1, 12, 0, 0)

tests/laborious/utils/test_connectors_config.py

@@ -1,8 +1,10 @@
 from os import environ
-from laborious.utils.connectors_config import (build_mlflow_config,
-                                               build_opc_config,
-                                               build_postgres_config,
-                                               build_mongodb_config)
+
+from laborious.utils.connectors_config import (
+    build_minio_config,
+    build_mlflow_config,
+    build_opc_config,
+)
 
 
 def test_build_mlflow_config_with_env_vars():
@@ -88,74 +90,33 @@ def test_build_opc_config_with_defaults():
     assert config['1']['reconnection_interval'] == 120
 
 
-def test_build_postgres_config_with_env_vars():
-    # Arrange
-    environ['POSTGRES_HOST'] = 'test-host'
-    environ['POSTGRES_PORT'] = '5433'
-    environ['POSTGRES_USER'] = 'test-user'
-    environ['POSTGRES_PASSWORD'] = 'test-pass'
-    environ['POSTGRES_DBNAME'] = 'test-db'
-    environ['POSTGRES_MIN_CONNECTIONS'] = '10'
-    environ['POSTGRES_MAX_CONNECTIONS'] = '30'
-
-    # Act
-    config = build_postgres_config()
-
-    # Assert
-    assert config['host'] == 'test-host'
-    assert config['port'] == 5433
-    assert config['user'] == 'test-user'
-    assert config['password'] == 'test-pass'
-    assert config['dbname'] == 'test-db'
-    assert config['min_connections'] == 10
-    assert config['max_connections'] == 30
-
-
-def test_build_postgres_config_with_defaults():
-    # Arrange
-    environ.pop('POSTGRES_HOST', None)
-    environ.pop('POSTGRES_PORT', None)
-    environ.pop('POSTGRES_USER', None)
-    environ.pop('POSTGRES_PASSWORD', None)
-    environ.pop('POSTGRES_DBNAME', None)
-    environ.pop('POSTGRES_MIN_CONNECTIONS', None)
-    environ.pop('POSTGRES_MAX_CONNECTIONS', None)
-
-    # Act
-    config = build_postgres_config()
-
-    # Assert
-    assert config['host'] == 'localhost'
-    assert config['port'] == 5432
-    assert config['user'] == 'sientia'
-    assert config['password'] == 'sientia'
-    assert config['dbname'] == 'sientia'
-    assert config['min_connections'] == 5
-    assert config['max_connections'] == 20
-
-
-def test_build_mongo_db_config_with_env_vars():
-    environ['MONGODB_USERNAME'] = 'sientia1'
-    environ['MONGODB_PASSWORD'] = 'sientia1'
-    environ['MONGODB_URL'] = 'localhost:27018'
-    environ['MONGODB_DATABASE_NAME'] = 'test_db'
-    environ['MONGODB_TTL_INDEX_HOURS'] = '1'
-
-    assert build_mongodb_config() == {
-        'connection_string': 'mongodb://sientia1:sientia1@localhost:27018',
-        'database_name': 'test_db',
-        'ttl_index_seconds': 3600
+def test_build_minio_config_with_env_vars():
+    environ['MINIO_ENDPOINT_URL'] = 'http://test-host'
+    environ['MINIO_ACCESS_KEY'] = 'test-key'
+    environ['MINIO_SECRET_KEY'] = 'test-secret'
+    environ['MINIO_REGION_NAME'] = 'test-region'
+    environ['MINIO_DEFAULT_BUCKET'] = 'test-bucket'
+    assert build_minio_config() == {
+        'endpoint_url': 'http://test-host',
+        'access_key': 'test-key',
+        'secret_key': 'test-secret',
+        'default_bucket': 'test-bucket',
+        'retention_hours': 24,
+        'secure': False,
     }
 
 
-def test_build_mongo_db_config_with_defaults():
-    environ.pop('MONGODB_USERNAME', None)
-    environ.pop('MONGODB_PASSWORD', None)
-    environ.pop('MONGODB_DATABASE_NAME', None)
-    environ.pop('MONGODB_URL', None)
-    environ.pop('MONGODB_TTL_INDEX_HOURS', None)
-    assert build_mongodb_config() == {
-        'connection_string': 'mongodb://root:wKZDbMNU1c@localhost:27018',
-        'database_name': 'sientia',
-        'ttl_index_seconds': 3600
+def test_build_minio_config_with_defaults():
+    environ.pop('MINIO_ENDPOINT_URL', None)
+    environ.pop('MINIO_ACCESS_KEY', None)
+    environ.pop('MINIO_SECRET_KEY', None)
+    environ.pop('MINIO_REGION_NAME', None)
+    environ.pop('MINIO_DEFAULT_BUCKET', None)
+    assert build_minio_config() == {
+        'endpoint_url': 'http://localhost:9000',
+        'access_key': 'minioadmin',
+        'secret_key': 'minioadmin',
+        'default_bucket': 'laborious',
+        'retention_hours': 24,
+        'secure': False,
     }

tests/laborious/worker/test_runtime_task_queues.py

@@ -0,0 +1,14 @@
+from sientia_do.temporal.worker.prepare_worker import build_queue_name
+
+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
+
+
+def test_runtime_scoped_queue_names():
+    runtime = 'prod-a'
+    assert build_queue_name(PredictionsBatch.__name__, runtime) == 'predictions_batch-prod-a-queue'
+    assert build_queue_name(MinimalRetrain.__name__, runtime) == 'minimal_retrain-prod-a-queue'
+    assert build_queue_name(Drift.__name__, runtime) == 'drift-prod-a-queue'
+    assert build_queue_name(SimpleMetrics.__name__, runtime) == 'simple_metrics-prod-a-queue'

tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py

@@ -1,9 +1,10 @@
-from unittest.mock import call, patch, AsyncMock, ANY
-from pytest import mark, fixture
+from unittest.mock import ANY, AsyncMock, MagicMock, 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.sub_workflows.format_and_export_prediction import FormatAndExportPrediction
-from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
 
 
 @fixture
@@ -12,149 +13,667 @@ def format_and_export_prediction():
 
 
 metadata = {
-    "metadata": {
-        "model_id": "test_model",
-        "model_name": "test_model",
-        "workflow_name": "test_workflow",
-        "schema_name": "test_schedule",
+    'metadata': {
+        'model_id': 'test_model',
+        'model_name': 'test_model',
+        'workflow_name': 'test_workflow',
+        'schema_name': 'test_schedule',
     },
 }
 
 
 @mark.asyncio
-@patch("laborious.workflows.sub_workflows.format_and_export_prediction.workflow", new_callable=AsyncMock)
+@patch(
+    'laborious.workflows.sub_workflows.format_and_export_prediction.workflow',
+    new_callable=AsyncMock,
+)
 async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
-
     input_data = {
         'metadata': metadata,
-        "path_flag": None,
-        "data": {"test": "data"},
-        "timestamp": "2021-01-01",
-        "model_id": 1,
-        "prediction_confidence": 0,
-        "schema": "test_schema",
-        "table_name": "test_table",
-        "opc_servers": ["test_server"],
-        "opc_output_config": {"test": "config"},
-        "prediction_store_policy": "erl:1"
+        'path_flag': None,
+        'data': {'test': 'data'},
+        'timestamp': '2021-01-01',
+        'model_id': 1,
+        'model_name': metadata['metadata']['model_name'],
+        'prediction_confidence': 0,
+        'schema': 'test_schema',
+        'table_name': 'test_table',
+        'opc_servers': ['test_server'],
+        'opc_output_config': {'test': 'config'},
+        'prediction_store_policy': 'erl:1',
     }
 
+    prediction_data = MagicMock()
+    opc_metrics = MagicMock()
+
+    workflow_mock.execute_activity_method.side_effect = [
+        (prediction_data, opc_metrics),
+        MagicMock(),
+        MagicMock(),
+    ]
+
     await format_and_export_prediction.run(input_data)
 
-    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
-        )])
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.format_prediction,
+                {
+                    **metadata,
+                    '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'],
+                    'model_name': input_data['model_name'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
 
-    workflow_mock.execute_activity_method.assert_has_calls([
-        call(
-            Activities.write_opc_data,
-            {
-                'opc_output_config': input_data['opc_output_config'],
-                'data': workflow_mock.execute_local_activity_method.return_value,
-                **metadata
-            },
-            retry_policy=ANY,
-            start_to_close_timeout=ANY
-        )
-    ])
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.write_opc_data,
+                {
+                    'opc_output_config': input_data['opc_output_config'],
+                    'data': workflow_mock.execute_local_activity_method.return_value,
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
 
-    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': workflow_mock.execute_activity_method.return_value,
-                **metadata,
-                'timestamp_conversion': {
-                    'column': 'timestamp',
-                    'format': DATETIME_FORMAT_WITH_TZ
-                }
-            },
-            retry_policy=ANY,
-            start_to_close_timeout=ANY
-        )])
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.export_data_to_postgres,
+                {
+                    **metadata,
+                    'schema': input_data['schema'],
+                    'table_name': input_data['table_name'],
+                    'data': prediction_data,
+                    'timestamp_conversion': {
+                        'column': 'timestamp',
+                        'format': DATETIME_FORMAT_WITH_TZ,
+                    },
+                    'on_conflict': 'error',
+                    'unique_columns': ['model_id', 'timestamp'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    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 workflow_mock.execute_activity_method.call_count == 3
     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_default_path_flag(workflow_mock, format_and_export_prediction):
-
+@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": "default",
-        "data": {"test": "data"},
-        "timestamp": "2021-01-01",
-        "model_id": 1,
-        "prediction_confidence": 0,
-        "schema": "test_schema",
-        "table_name": "test_table",
-        "opc_servers": ["test_server"],
-        "opc_output_config": {"test": "config"},
-        "comment": "test_comment"
+        'path_flag': None,
+        'data': {'test': 'data'},
+        'transformed_data': {'transformed': 'data'},
+        'timestamp': '2021-01-01',
+        'model_id': 1,
+        'model_name': metadata['metadata']['model_name'],
+        '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,
+                {
+                    **metadata,
+                    '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'],
+                    'model_name': input_data['model_name'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            ),
+            call(
+                Activities.format_transformed_data,
+                {
+                    **metadata,
+                    'data': input_data['transformed_data'],
+                    'model_id': input_data['model_id'],
+                    'model_name': input_data['model_name'],
+                },
+                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_payload_to_postgres,
+        {
+            **metadata,
+            'schema': input_data['schema'],
+            'table_name': input_data['transform_table_name'],
+            'data': transformed_data,
+            'timestamp_conversion': {
+                'column': 'timestamp',
+                'format': DATETIME_FORMAT_WITH_TZ,
+            },
+        },
+        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,
+                {
+                    **metadata,
+                    'schema': input_data['schema'],
+                    'table_name': input_data['table_name'],
+                    'data': prediction_data,
+                    'timestamp_conversion': {
+                        'column': 'timestamp',
+                        'format': DATETIME_FORMAT_WITH_TZ,
+                    },
+                    'on_conflict': 'error',
+                    'unique_columns': ['model_id', 'timestamp'],
+                },
+                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',
+    new_callable=AsyncMock,
+)
+async def test_run_default_path_flag(workflow_mock, format_and_export_prediction):
+    input_data = {
+        'metadata': metadata,
+        'path_flag': 'default',
+        'data': {'test': 'data'},
+        'timestamp': '2021-01-01',
+        'model_id': 1,
+        'model_name': metadata['metadata']['model_name'],
+        'prediction_confidence': 0,
+        'schema': 'test_schema',
+        'table_name': 'test_table',
+        'opc_servers': ['test_server'],
+        'opc_output_config': {'test': 'config'},
+        'comment': 'test_comment',
+    }
+
+    prediction_data = MagicMock()
+    opc_metrics = MagicMock()
+
+    workflow_mock.execute_activity_method.side_effect = [
+        (prediction_data, opc_metrics),
+        MagicMock(),
+        MagicMock(),
+    ]
+
     await format_and_export_prediction.run(input_data)
 
-    workflow_mock.execute_local_activity_method.assert_has_calls([
-        call(
-            Activities.format_default_prediction,
-            {
-                'timestamp': input_data['timestamp'],
-                'model_id': input_data['model_id'],
-                'prediction_confidence': input_data['prediction_confidence'],
-                'comment': input_data['comment'],
-                **metadata
-            },
-            retry_policy=ANY,
-            start_to_close_timeout=ANY
-        )
-    ])
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.format_default_prediction,
+                {
+                    **metadata,
+                    'timestamp': input_data['timestamp'],
+                    'model_id': input_data['model_id'],
+                    'prediction_confidence': input_data['prediction_confidence'],
+                    'comment': input_data['comment'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
 
-    workflow_mock.execute_activity_method.assert_has_calls([
-        call(
-            Activities.write_opc_data,
-            {
-                'opc_output_config': input_data['opc_output_config'],
-                'data': workflow_mock.execute_local_activity_method.return_value,
-                **metadata
-            },
-            retry_policy=ANY,
-            start_to_close_timeout=ANY
-        )
-    ])
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.write_opc_data,
+                {
+                    'opc_output_config': input_data['opc_output_config'],
+                    'data': workflow_mock.execute_local_activity_method.return_value,
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
 
-    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': workflow_mock.execute_activity_method.return_value,
-                **metadata,
-                'timestamp_conversion': {
-                    'column': 'timestamp',
-                    'format': DATETIME_FORMAT_WITH_TZ
-                }
-            },
-            retry_policy=ANY,
-            start_to_close_timeout=ANY
-        )
-    ])
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.export_data_to_postgres,
+                {
+                    **metadata,
+                    'schema': input_data['schema'],
+                    'table_name': input_data['table_name'],
+                    'data': prediction_data,
+                    'timestamp_conversion': {
+                        'column': 'timestamp',
+                        'format': DATETIME_FORMAT_WITH_TZ,
+                    },
+                    'on_conflict': 'error',
+                    'unique_columns': ['model_id', 'timestamp'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    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 workflow_mock.execute_activity_method.call_count == 3
+    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_pi_web_api(workflow_mock, format_and_export_prediction):
+    input_data = {
+        'metadata': metadata,
+        'path_flag': None,
+        'data': {'test': 'data'},
+        'timestamp': '2021-01-01',
+        'model_id': 1,
+        'model_name': metadata['metadata']['model_name'],
+        'prediction_confidence': 0,
+        'schema': 'test_schema',
+        'table_name': 'test_table',
+        'pi_web_api_output_config': {
+            'endpoint': 'https://test-pi-server.com',
+            'prediction_tags': {},
+            'confidence_tags': {},
+        },
+        'prediction_store_policy': 'erl:1',
+    }
+
+    pi_web_api_data = MagicMock()
+
+    workflow_mock.execute_activity_method.side_effect = [
+        pi_web_api_data,  # write_pi_web_api_data
+        MagicMock(),  # export_data_to_postgres
+        MagicMock(),  # write_metrics
+    ]
+
+    await format_and_export_prediction.run(input_data)
+
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.format_prediction,
+                {
+                    **metadata,
+                    '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'],
+                    'model_name': input_data['model_name'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.write_pi_web_api_data,
+                {
+                    'pi_web_api_output_config': input_data['pi_web_api_output_config'],
+                    'data': workflow_mock.execute_local_activity_method.return_value,
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.export_data_to_postgres,
+                {
+                    **metadata,
+                    'schema': input_data['schema'],
+                    'table_name': input_data['table_name'],
+                    'data': pi_web_api_data,
+                    'timestamp_conversion': {
+                        'column': 'timestamp',
+                        'format': DATETIME_FORMAT_WITH_TZ,
+                    },
+                    'on_conflict': 'error',
+                    'unique_columns': ['model_id', 'timestamp'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.write_metrics,
+                {
+                    **metadata,
+                    'prediction': pi_web_api_data,
+                    'opc_metrics': {},
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    assert workflow_mock.execute_activity_method.call_count == 3
+    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_pi_web_api_and_opc(
+    workflow_mock, format_and_export_prediction
+):
+    input_data = {
+        'metadata': metadata,
+        'path_flag': None,
+        'data': {'test': 'data'},
+        'timestamp': '2021-01-01',
+        'model_id': 1,
+        'model_name': metadata['metadata']['model_name'],
+        'prediction_confidence': 0,
+        'schema': 'test_schema',
+        'table_name': 'test_table',
+        'opc_output_config': {'test': 'config'},
+        'pi_web_api_output_config': {
+            'endpoint': 'https://test-pi-server.com',
+            'prediction_tags': {},
+            'confidence_tags': {},
+        },
+        'prediction_store_policy': 'erl:1',
+    }
+
+    prediction_data = MagicMock()
+    pi_web_api_data = MagicMock()
+    opc_metrics = MagicMock()
+
+    workflow_mock.execute_activity_method.side_effect = [
+        pi_web_api_data,  # write_pi_web_api_data
+        (prediction_data, opc_metrics),  # write_opc_data
+        MagicMock(),  # export_data_to_postgres
+        MagicMock(),  # write_metrics
+    ]
+
+    await format_and_export_prediction.run(input_data)
+
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.write_pi_web_api_data,
+                {
+                    'pi_web_api_output_config': input_data['pi_web_api_output_config'],
+                    'data': workflow_mock.execute_local_activity_method.return_value,
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            ),
+            call(
+                Activities.write_opc_data,
+                {
+                    'opc_output_config': input_data['opc_output_config'],
+                    'data': pi_web_api_data,
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            ),
+        ]
+    )
+
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.export_data_to_postgres,
+                {
+                    **metadata,
+                    'schema': input_data['schema'],
+                    'table_name': input_data['table_name'],
+                    'data': prediction_data,
+                    'timestamp_conversion': {
+                        'column': 'timestamp',
+                        'format': DATETIME_FORMAT_WITH_TZ,
+                    },
+                    'on_conflict': 'error',
+                    'unique_columns': ['model_id', 'timestamp'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    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 workflow_mock.execute_activity_method.call_count == 4
+    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_default_path_flag_with_pi_web_api(workflow_mock, format_and_export_prediction):
+    input_data = {
+        'metadata': metadata,
+        'path_flag': 'default',
+        'data': {'test': 'data'},
+        'timestamp': '2021-01-01',
+        'model_id': 1,
+        'model_name': metadata['metadata']['model_name'],
+        'prediction_confidence': 0,
+        'schema': 'test_schema',
+        'table_name': 'test_table',
+        'pi_web_api_output_config': {
+            'endpoint': 'https://test-pi-server.com',
+            'prediction_tags': {},
+            'confidence_tags': {},
+        },
+        'comment': 'test_comment',
+    }
+
+    pi_web_api_data = MagicMock()
+
+    workflow_mock.execute_activity_method.side_effect = [
+        pi_web_api_data,  # write_pi_web_api_data
+        MagicMock(),  # export_data_to_postgres
+        MagicMock(),  # write_metrics
+    ]
+
+    await format_and_export_prediction.run(input_data)
+
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.write_pi_web_api_data,
+                {
+                    'pi_web_api_output_config': input_data['pi_web_api_output_config'],
+                    'data': workflow_mock.execute_local_activity_method.return_value,
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.export_data_to_postgres,
+                {
+                    **metadata,
+                    'schema': input_data['schema'],
+                    'table_name': input_data['table_name'],
+                    'data': pi_web_api_data,
+                    'timestamp_conversion': {
+                        'column': 'timestamp',
+                        'format': DATETIME_FORMAT_WITH_TZ,
+                    },
+                    'on_conflict': 'error',
+                    'unique_columns': ['model_id', 'timestamp'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
 
     assert workflow_mock.execute_activity_method.call_count == 3
     assert workflow_mock.execute_local_activity_method.call_count == 1

tests/laborious/workflows/test_drift.py

@@ -0,0 +1,248 @@
+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_activity_method.side_effect = [target_data, reference_data]
+    workflow_mock.execute_local_activity_method = AsyncMock(return_value=drift_data)
+    workflow_mock.execute_activity_method = AsyncMock(return_value=None)
+
+    # 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_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_activity_method.side_effect = [target_data, reference_data]
+
+    workflow_mock.execute_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_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_activity_method.side_effect = [target_data, reference_data]
+    workflow_mock.execute_local_activity_method = AsyncMock(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_activity_method.side_effect = [target_data, reference_data]
+    workflow_mock.execute_local_activity_method = AsyncMock(return_value=drift_data)
+    workflow_mock.execute_activity_method = AsyncMock(return_value=None)
+
+    # 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_minimal_retrain.py

@@ -1,5 +1,7 @@
-from unittest.mock import AsyncMock, MagicMock, call, patch, ANY
+from unittest.mock import ANY, AsyncMock, call, patch
+
 from pytest import fixture, mark
+
 from laborious.activities.activities import Activities
 from laborious.workflows.minimal_retrain import MinimalRetrain
 
@@ -10,11 +12,11 @@ def minimal_retrain() -> MinimalRetrain:
 
 
 metadata = {
-    "metadata": {
-        "model_id": "test_model_id",
-        "model_name": "test_model",
-        "workflow_name": "minimal_retrain",
-        "schedule_name": "test_schedule",
+    'metadata': {
+        'model_id': 'test_model_id',
+        'model_name': 'test_model',
+        'workflow_name': 'minimal_retrain',
+        'schedule_name': 'test_schedule',
     },
 }
 
@@ -23,76 +25,303 @@ metadata = {
 @patch('laborious.workflows.minimal_retrain.workflow', new_callable=AsyncMock)
 async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain):
     input_data = {
-        "model_id": "test_model_id",
-        "model_name": "test_model",
-        "workflow_name": "minimal_retrain",
-        "schedule_name": "test_schedule",
-        "query": "test_query",
-        "schema": "test_schema",
-        "table_name": "test_table",
+        'model_id': 'test_model_id',
+        'model_name': 'test_model',
+        'workflow_name': 'minimal_retrain',
+        'schedule_name': 'test_schedule',
+        'query': 'test_query',
+        'schema': 'test_schema',
+        'table_name': 'test_table',
+        'model_config': {
+            'target': 'test_target',
+            'transform_flavor': 'test_transform_flavor',
+            'predict_flavor': 'test_predict_flavor',
+        },
+    }
+
+    storage_result = {
+        'last_timestamp': '2024-01-01 00:00:00+0000',
+        'status': {'success': True},
+        'data': {'timestamp': {0: '2024-01-01 00:00:00+0000'}, 'value': {0: 1.0}},
+        'bucket': None,
+        'object_key': None,
+        'object_prefix': None,
+        'uri': None,
     }
 
     workflow_mock.execute_activity_method = AsyncMock(
-        return_value={
-            "data1": "1",
-            "data2": "2",
-        }
+        side_effect=[
+            storage_result,
+            {'success': True, 'experiment': 'test_experiment'},
+            {
+                'success': True,
+                'version': 'test_version',
+                'mlflow_run_id': 'test_mlflow_run_id',
+                'mlflow_experiment_id': 'test_mlflow_experiment_id',
+            },
+            {'report': 'test_report'},
+        ]
     )
 
     await minimal_retrain.run(input_data)
 
-    workflow_mock.execute_local_activity_method.assert_has_calls(
+    workflow_mock.execute_activity_method.assert_has_calls(
         [
             call(
-                Activities.load_custom_query,
+                Activities.load_query_with_minio_offload,
                 {
                     **metadata,
-                    "query": input_data["query"],
-                    'datetime_columns': input_data.get('datetime_columns', [])
+                    'query': input_data['query'],
+                    'datetime_columns': input_data.get('datetime_columns', []),
+                    'model_name': input_data['model_name'],
                 },
                 retry_policy=ANY,
-                start_to_close_timeout=ANY
+                start_to_close_timeout=ANY,
             )
         ]
     )
 
-    workflow_mock.execute_activity_method.assert_has_calls([
-        call(
-            Activities.retrain_model,
-            {
-                **metadata,
-                'data': workflow_mock.execute_local_activity_method.return_value,
-                'model_name': input_data['model_name'],
-            },
-            retry_policy=ANY,
-            start_to_close_timeout=ANY
-        )
-    ])
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.retrain_model,
+                {
+                    **metadata,
+                    'data': storage_result,
+                    'model_name': input_data['model_name'],
+                    'model_config': input_data['model_config'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
 
-    workflow_mock.execute_activity_method.assert_has_calls([
-        call(
-            Activities.update_production_model,
-            {
-                **metadata,
-                'model_name': input_data['model_name'],
-                'model_id': input_data['model_id'],
-                **workflow_mock.execute_activity_method.return_value,
-            },
-            retry_policy=ANY,
-            start_to_close_timeout=ANY
-        )
-    ])
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.update_production_model,
+                {
+                    **metadata,
+                    'model_name': input_data['model_name'],
+                    'success': True,
+                    'experiment': 'test_experiment',
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
 
-    workflow_mock.execute_activity_method.assert_has_calls([
-        call(
-            Activities.export_data_to_postgres,
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.format_retrain_report,
+                {
+                    **metadata,
+                    'experiment_response': {'success': True, 'experiment': 'test_experiment'},
+                    'model_name': input_data['model_name'],
+                    'model_id': input_data['model_id'],
+                    'update_report': {
+                        'success': True,
+                        'version': 'test_version',
+                        'mlflow_run_id': 'test_mlflow_run_id',
+                        'mlflow_experiment_id': 'test_mlflow_experiment_id',
+                    },
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.export_data_to_postgres,
+                {
+                    **metadata,
+                    'data': workflow_mock.execute_local_activity_method.return_value,
+                    'schema': input_data['schema'],
+                    'table_name': input_data['table_name'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+
+@mark.asyncio
+@patch('laborious.workflows.minimal_retrain.workflow', new_callable=AsyncMock)
+async def test_run_storage_fail(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain):
+    input_data = {
+        'model_id': 'test_model_id',
+        'model_name': 'test_model',
+        'workflow_name': 'minimal_retrain',
+        'schedule_name': 'test_schedule',
+        'query': 'test_query',
+        'schema': 'test_schema',
+        'table_name': 'test_table',
+        'model_config': {
+            'target': 'test_target',
+            'transform_flavor': 'test_transform_flavor',
+            'predict_flavor': 'test_predict_flavor',
+        },
+    }
+
+    storage_result = {
+        'last_timestamp': '2024-01-01 00:00:00+0000',
+        'status': {'success': True},
+        'data': {},
+        'bucket': None,
+        'object_key': None,
+        'object_prefix': None,
+        'uri': None,
+    }
+
+    workflow_mock.execute_activity_method = AsyncMock(
+        side_effect=[
+            storage_result,
+            {'success': True, 'experiment': 'test_experiment'},
             {
-                **metadata,
-                'data': workflow_mock.execute_activity_method.return_value,
-                'schema': input_data['schema'],
-                'table_name': input_data['table_name'],
+                'success': True,
+                'version': 'test_version',
+                'mlflow_run_id': 'test_mlflow_run_id',
+                'mlflow_experiment_id': 'test_mlflow_experiment_id',
             },
-            retry_policy=ANY,
-            start_to_close_timeout=ANY
-        )
-    ])
+            {'report': 'test_report'},
+        ]
+    )
+
+    from pytest import raises
+
+    with raises(ValueError, match='No data returned from query'):
+        await minimal_retrain.run(input_data)
+
+    workflow_mock.execute_activity_method.assert_called_once_with(
+        Activities.load_query_with_minio_offload,
+        {
+            **metadata,
+            'query': input_data['query'],
+            'datetime_columns': input_data.get('datetime_columns', []),
+            'model_name': input_data['model_name'],
+        },
+        retry_policy=ANY,
+        start_to_close_timeout=ANY,
+    )
+
+    workflow_mock.execute_local_activity_method.assert_not_called()
+
+
+@mark.asyncio
+@patch('laborious.workflows.minimal_retrain.workflow', new_callable=AsyncMock)
+async def test_run_fail_retrain(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain):
+    input_data = {
+        'model_id': 'test_model_id',
+        'model_name': 'test_model',
+        'workflow_name': 'minimal_retrain',
+        'schedule_name': 'test_schedule',
+        'query': 'test_query',
+        'schema': 'test_schema',
+        'table_name': 'test_table',
+        'model_config': {
+            'target': 'test_target',
+            'transform_flavor': 'test_transform_flavor',
+            'predict_flavor': 'test_predict_flavor',
+        },
+    }
+
+    storage_result = {
+        'last_timestamp': '2024-01-01 00:00:00+0000',
+        'status': {'success': True},
+        'data': {'timestamp': {0: '2024-01-01 00:00:00+0000'}, 'value': {0: 1.0}},
+        'bucket': None,
+        'object_key': None,
+        'object_prefix': None,
+        'uri': None,
+    }
+
+    workflow_mock.execute_activity_method = AsyncMock(
+        side_effect=[
+            storage_result,
+            {'success': False, 'experiment': 'test_experiment'},
+            {
+                'success': True,
+                'version': 'test_version',
+                'mlflow_run_id': 'test_mlflow_run_id',
+                'mlflow_experiment_id': 'test_mlflow_experiment_id',
+            },
+            {'report': 'test_report'},
+        ]
+    )
+
+    await minimal_retrain.run(input_data)
+
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.load_query_with_minio_offload,
+                {
+                    **metadata,
+                    'query': input_data['query'],
+                    'datetime_columns': input_data.get('datetime_columns', []),
+                    'model_name': input_data['model_name'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.retrain_model,
+                {
+                    **metadata,
+                    'data': storage_result,
+                    'model_name': input_data['model_name'],
+                    'model_config': input_data['model_config'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.format_retrain_report,
+                {
+                    **metadata,
+                    'experiment_response': {'success': False, 'experiment': 'test_experiment'},
+                    'model_name': input_data['model_name'],
+                    'model_id': input_data['model_id'],
+                    'update_report': {},
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.export_data_to_postgres,
+                {
+                    **metadata,
+                    'data': workflow_mock.execute_local_activity_method.return_value,
+                    'schema': input_data['schema'],
+                    'table_name': input_data['table_name'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    assert workflow_mock.execute_activity_method.call_count == 3
+    assert workflow_mock.execute_local_activity_method.call_count == 1

tests/laborious/workflows/test_predictions_batch.py

@@ -1,5 +1,7 @@
-from unittest.mock import AsyncMock, call, patch, ANY
+from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
+
 from pytest import fixture, mark
+
 from laborious.activities.activities import Activities
 from laborious.workflows.predictions_batch import PredictionsBatch
 
@@ -10,21 +12,19 @@ def predictions_batch() -> PredictionsBatch:
 
 
 metadata = {
-    "metadata": {
-        "model_id": "test_model_id",
-        "model_name": "test_model",
-        "workflow_name": "predictions_batch",
-        "schedule_name": "test_schedule",
-    },
+    'model_id': 'test_model_id',
+    'model_name': 'test_model',
+    'workflow_name': 'predictions_batch',
+    'schedule_name': 'test_schedule',
 }
 
 
 @mark.asyncio
 @patch('laborious.workflows.predictions_batch.workflow', new_callable=AsyncMock)
 async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch):
-    workflow_mock.execute_local_activity_method.return_value = {
-        'data': 'test_data'
-    }
+    activity_return = MagicMock()
+    workflow_mock.execute_activity_method.return_value = activity_return
+
     input_data = {
         'schedule_name': 'test_schedule',
         'model_name': 'test_model',
@@ -32,52 +32,75 @@ 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']
+        'pi_web_api_output_config': 'test_pi_web_api_output_config',
+        'datetime_columns': ['timestamp', 'created_at'],
+        'prediction_store_policy': 'erl:1',
+        'model_config': {'retention': '30'},
     }
 
     await predictions_batch.run(input_data)
 
-    workflow_mock.execute_local_activity_method.assert_has_calls([
-        call(
-            Activities.load_custom_query,
-            {
-                **metadata,
-                'query': input_data['query'],
-                'datetime_columns': input_data.get('datetime_columns', [])
-            },
-            retry_policy=ANY,
-            start_to_close_timeout=ANY
-        )
-    ])
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.load_query_with_minio_offload,
+                {
+                    'metadata': metadata,
+                    'query': input_data['query'],
+                    'datetime_columns': input_data.get('datetime_columns', []),
+                    'model_name': input_data['model_name'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
     prediction_input = {
-        'metadata': metadata,
-        'data': {'data': 'test_data'},
+        'metadata': {'metadata': metadata},
+        'data': activity_return,
         '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'
-            }
-        }),
-        'mlflow_transform_filters': input_data.get('mlflow_transform_filters', {
-            'API_ERROR': {
-                'POLICY': 'STOP'
-            }
-        }),
-        'mlflow_predict_filters': input_data.get('mlflow_predict_filters', {
-            'API_ERROR': {
-                'POLICY': 'STOP'
-            }
-        }),
-        'model_retention': input_data.get('model_retention', 60),
+        'input_filters': input_data.get(
+            'input_filters',
+            {
+                'EMPTY_DATA': {
+                    'POLICY': 'STOP',
+                    'CONFIG': {},
+                }
+            },
+        ),
+        'mlflow_transform_filters': input_data.get(
+            'mlflow_transform_filters',
+            {
+                'API_ERROR': {
+                    'POLICY': 'STOP',
+                    'CONFIG': {},
+                }
+            },
+        ),
+        'mlflow_predict_filters': input_data.get(
+            'mlflow_predict_filters',
+            {
+                'API_ERROR': {
+                    'POLICY': 'STOP',
+                    'CONFIG': {},
+                }
+            },
+        ),
+        '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', {})
+        'opc_output_config': input_data.get('opc_output_config', {}),
+        'on_conflict': input_data.get('on_conflict', 'error'),
+        'pi_web_api_output_config': input_data.get('pi_web_api_output_config', {}),
+        'prediction_store_policy': input_data.get('prediction_store_policy', 'lts:1'),
+        'save_transform': input_data.get('save_transform', True),
     }
 
-    workflow_mock.execute_child_workflow.assert_has_calls([
-        call(
-            'prediction_process', prediction_input)
-    ])
+    workflow_mock.execute_child_workflow.assert_has_calls(
+        [call('subworkflow.prediction_process', prediction_input)]
+    )

tests/laborious/workflows/test_simple_metrics.py

@@ -0,0 +1,215 @@
+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_activity_method = AsyncMock(side_effect=[target_data, None])
+    workflow_mock.execute_local_activity_method = AsyncMock(return_value=simple_metrics_data)
+
+    # 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_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.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,
+            ),
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_called_once_with(
+        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,
+    )
+
+
+@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_activity_method = AsyncMock(return_value=target_data)
+
+    # Act
+    await simple_metrics.run(input_data)
+
+    # Assert - Should not call calculate_simple_metrics or export
+    assert workflow_mock.execute_activity_method.call_count == 1
+
+
+@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_activity_method = AsyncMock(return_value=target_data)
+    workflow_mock.execute_local_activity_method = AsyncMock(return_value=simple_metrics_data)
+
+    # Act
+    await simple_metrics.run(input_data)
+
+    # Assert - Should call calculate_simple_metrics but not export
+    workflow_mock.execute_activity_method.assert_called_once()
+    workflow_mock.execute_local_activity_method.assert_called_once()
+
+
+@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_activity_method = AsyncMock(side_effect=[target_data, None])
+    workflow_mock.execute_local_activity_method = AsyncMock(return_value=simple_metrics_data)
+
+    # Act
+    await simple_metrics.run(input_data)
+
+    # Assert - Check calculate_simple_metrics call with default metrics
+    workflow_mock.execute_activity_method.assert_any_call(
+        Activities.load_custom_query,
+        ANY,
+        retry_policy=ANY,
+        start_to_close_timeout=ANY,
+    )
+    workflow_mock.execute_local_activity_method.assert_called_once_with(
+        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,
+    )

values.yaml

@@ -1,229 +0,0 @@
-# Default values for sientia-module.
-# This is a YAML-formatted file.
-# Declare variables to be passed into your templates.
-
-# This will set the replicaset count more information can be found here: https://kubernetes.io/docs/concepts/workloads/controllers/replicaset/
-replicaCount: 1
-
-# This sets the container image more information can be found here: https://kubernetes.io/docs/concepts/containers/images/
-image:
-  repository: aignosi.azurecr.io/sientia-module
-  # This sets the pull policy for images.
-  pullPolicy: Always
-  # Overrides the image tag whose default is the chart appVersion.
-  tag: "0.4.4"
-
-# 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:
-- name: docker-hub-secret
-# This is to override the chart name.
-nameOverride: "sientia-laborious-worker"
-fullnameOverride: "sientia-laborious-worker"
-namespace: sientia
-
-# This section builds out the service account more information can be found here: https://kubernetes.io/docs/concepts/security/service-accounts/
-serviceAccount:
-  # Specifies whether a service account should be created
-  create: true
-  # Automatically mount a ServiceAccount's API credentials?
-  automount: true
-  # Annotations to add to the service account
-  annotations: {}
-  # The name of the service account to use.
-  # If not set and create is true, a name is generated using the fullname template
-  name: "sientia-laborious-worker"
-
-# This is for setting Kubernetes Annotations to a Pod.
-# For more information checkout: https://kubernetes.io/docs/concepts/overview/working-with-objects/annotations/
-podAnnotations: {}
-# This is for setting Kubernetes Labels to a Pod.
-# For more information checkout: https://kubernetes.io/docs/concepts/overview/working-with-objects/labels/
-podLabels: {}
-
-podSecurityContext: {}
-  # fsGroup: 2000
-
-securityContext: {}
-  # capabilities:
-  #   drop:
-  #   - ALL
-  # readOnlyRootFilesystem: true
-  # runAsNonRoot: true
-  # runAsUser: 1000
-
-
-resources: {}
-  # We usually recommend not to specify default resources and to leave this as a conscious
-  # choice for the user. This also increases chances charts run on environments with little
-  # resources, such as Minikube. If you do want to specify resources, uncomment the following
-  # lines, adjust them as necessary, and remove the curly braces after 'resources:'.
-  # limits:
-  #   cpu: 100m
-  #   memory: 128Mi
-  # requests:
-  #   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:
-  exec:
-    command:
-      - sh
-      - -c
-      - pgrep -f "laborious.worker.worker"
-  initialDelaySeconds: 20
-  periodSeconds: 30
-
-readinessProbe:
-  exec:
-    command:
-      - sh
-      - -c
-      - pgrep -f "laborious.worker.worker"
-  initialDelaySeconds: 10
-  periodSeconds: 15
-
-
-# This section is for setting up autoscaling more information can be found here: https://kubernetes.io/docs/concepts/workloads/autoscaling/
-autoscaling:
-  enabled: false
-  minReplicas: 1
-  maxReplicas: 100
-  targetCPUUtilizationPercentage: 80
-  # targetMemoryUtilizationPercentage: 80
-
-# Additional volumes on the output Deployment definition.
-volumes: []
-# - name: foo
-#   secret:
-#     secretName: mysecret
-#     optional: false
-
-# Additional volumeMounts on the output Deployment definition.
-volumeMounts: []
-# - name: foo
-#   mountPath: "/etc/foo"
-#   readOnly: true
-
-nodeSelector: {}
-
-tolerations: []
-
-affinity: {}
-
-services:
-  sdk-metrics:
-    enabled: true
-    type: ClusterIP
-    port: 9091
-    targetPort: 9091
-    name: sdk-metrics
-
-  metrics:
-    enabled: true
-    type: ClusterIP
-    port: 9090
-    targetPort: 9090
-    name: metrics
-
-# Configuração do ServiceMonitor para o Prometheus Operator
-# ref: https://github.com/prometheus-operator/prometheus-operator
-serviceMonitor:
-  # Se true, um recurso ServiceMonitor será criado.
-  enabled: true
-  # O intervalo no qual as métricas devem ser coletadas (ex: 30s, 1m).
-  endpoints:
-  - port: metrics
-    path: /metrics
-    interval: 30s
-    relabelings: []
-  - port: sdk-metrics
-    path: /metrics
-    interval: 30s
-    relabelings: []
-
-  additionalLabels:
-    release: kube-prometheus-stack
-  
-
-env:
-  # Entrypoint variables 
-  - name: GITHUB_REPO_URL
-    value: "git@github.com:Aignosi/sientia-dataops-laborious_temporal.git"
-  - name: GITHUB_BRANCH
-    value: "SIENTIAPDE-1205-alterar-opc-para-assincrono"
-  - name: PYTHON_APP
-    value: "laborious.worker.worker"
-
-  # Application variables
-  - name: POSTGRES_HOST
-    value: "paradedb-rw.paradedb.svc.cluster.local"
-  - name: POSTGRES_PORT
-    value: "5432"
-  - name: POSTGRES_USER
-    value: "sientia"
-  - name: POSTGRES_PASSWORD
-    value: "sientia"
-  - name: POSTGRES_DBNAME
-    value: "sientia"
-  - name: POSTGRES_MIN_CONNECTIONS
-    value: "10"
-  - name: POSTGRES_MAX_CONNECTIONS
-    value: "30"
-
-  - name: MLFLOW_HOST
-    value: "http://sientia-tracker-mlflow-tracking.sientia-tracker.svc.cluster.local"
-  - name: MLFLOW_PORT
-    value: "80"
-  - name: MLFLOW_USERNAME
-    value: "aignosi"
-  - name: MLFLOW_PASSWORD
-    value: "aignosi"
-
-  - name: OPC_ID
-    value: "1"
-  - name: OPC_URL
-    value: "opc.tcp://sientia-opc-simulator-opc.sientia.svc.cluster.local:4840"
-
-  - name: KAFKA_BOOTSTRAP_SERVERS
-    value: "kafka.kafka.svc.cluster.local:9092"
-
-  - name: LOG_LEVEL
-    value: "DEBUG"
-  - name: HTTP_METRICS_PORT
-    value: "9090"
-  - name: HTTP_SDK_METRICS_PORT
-    value: "9091"
-  - name: PROJECT_NAME
-    value: "sientia-laborious"
-
-  - name: TEMPORAL_HOST
-    value: "temporal-frontend.temporal.svc.cluster.local:7233"
-  - name: TEMPORAL_NAMESPACE
-    value: "laborious"
-
-  - name: MONGODB_USERNAME
-    value: "root"
-  - name: MONGODB_PASSWORD
-    value: "wKZDbMNU1c"
-  - name: MONGODB_URL
-    value: "my-release-mongodb.mongodb.svc.cluster.local:27017"
-  - name: MONGODB_DATABASE
-    value: "sientia"
-  - name: MONGODB_TTL_INDEX_HOURS
-    value: "1"
-
-ssh:
-  enabled: true
-  secretName: git-ssh-key-sientia-laborious-worker
-  sshPath: /mnt/.ssh
-  knownHostsPath: /mnt/known_hosts
-
-# 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
-
-# kubectl create secret generic git-ssh-key-sientia-laborious-worker \
-#   --namespace sientia \
-#   --from-file=ssh-privatekey=git_key \
-#   --type=kubernetes.io/ssh-auth