diff --git a/laborious/workflows/simple_metrics.py b/laborious/workflows/simple_metrics.py index e8cd9fb..3253517 100644 --- a/laborious/workflows/simple_metrics.py +++ b/laborious/workflows/simple_metrics.py @@ -32,6 +32,7 @@ class SimpleMetrics: model_config = input_data['model_config'] target_name = model_config['target'] model_type = model_config.get('model_type') + thresholds = model_config.get('simple_metrics_thresholds') query = f""" select p."timestamp", p.prediction, ld.value as "target" @@ -72,6 +73,7 @@ class SimpleMetrics: 'metrics': input_data.get('metrics', ['rmse', 'mse', 'mae', 'r2']), 'interval_minutes': interval_minutes, 'model_type': model_type, + 'thresholds': thresholds, }, retry_policy=retry_policy, start_to_close_timeout=timedelta(seconds=300), diff --git a/tests/laborious/workflows/test_simple_metrics.py b/tests/laborious/workflows/test_simple_metrics.py index 2423cc5..5d94f60 100644 --- a/tests/laborious/workflows/test_simple_metrics.py +++ b/tests/laborious/workflows/test_simple_metrics.py @@ -104,6 +104,7 @@ async def test_run(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics): 'metrics': input_data['metrics'], 'interval_minutes': input_data['interval_minutes'], 'model_type': None, + 'thresholds': None, }, retry_policy=ANY, start_to_close_timeout=ANY, @@ -211,6 +212,7 @@ async def test_run_default_metrics(workflow_mock: AsyncMock, simple_metrics: Sim 'metrics': ['rmse', 'mse', 'mae', 'r2'], # Default value 'interval_minutes': input_data['interval_minutes'], 'model_type': None, + 'thresholds': None, }, retry_policy=ANY, start_to_close_timeout=ANY,