From fdeafe97685af33abc672d76db778b3670f69f4a Mon Sep 17 00:00:00 2001
From: vitor-aignosi <vitor.santos@aignosi.com.br>
Date: Wed, 7 Oct 2026 16:47:31 -0300
Subject: [PATCH] SAM0123-320: fix: call pyfunc data model predict with wrapper
 signature

When a data model's `fit` returns the refit object instead of a DataFrame,
`fit_models` called `predict(data)` sklearn-style. A pyfunc wrapper's
`predict` takes `(context, model_input)`, so Courier Stratified retrain
failed with "missing 1 required positional argument: 'model_input'".

Mirror `get_prediction_data`: pass `{}` as context when the transform
flavor is pyfunc.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../utils/repository/model_repository.py      |  5 ++-
 .../utils/repository/test_model_repository.py | 32 +++++++++++++++++++
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py
index bddfc60..76854a4 100644
--- a/laborious/utils/repository/model_repository.py
+++ b/laborious/utils/repository/model_repository.py
@@ -1341,7 +1341,10 @@ class MLFlowRepository(SientiaMonitoring):
 
         if not isinstance(treated_data_candidate, pd.DataFrame):
             data_model = treated_data_candidate
-            treated_data = data_model.predict(data)
+            if transform_flavor == 'pyfunc':
+                treated_data = data_model.predict({}, data)
+            else:
+                treated_data = data_model.predict(data)
         else:
             treated_data = treated_data_candidate
 
diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py
index 296ba6c..cf5e0be 100644
--- a/tests/laborious/utils/repository/test_model_repository.py
+++ b/tests/laborious/utils/repository/test_model_repository.py
@@ -1043,6 +1043,38 @@ async def test_fit_models_not_df_target_name_none_and_not_in_model(
     }
 
 
+@patch('laborious.utils.repository.model_repository.pd.merge')
+@patch('laborious.utils.repository.model_repository.isinstance')
+@pytest.mark.asyncio
+async def test_fit_models_pyfunc_transform_predicts_with_wrapper_signature(
+    isinstance_mock, pd_merge, mlflow_repository
+):
+    """A pyfunc data model refit returns the wrapper, whose `predict` takes `(context, model_input)`."""
+    isinstance_mock.return_value = False
+
+    data_model = MagicMock()
+    mlflow_repository.download_model = AsyncMock(
+        side_effect=[(data_model, 'artifact_path'), (MagicMock(), 'artifact_path')],
+    )
+    mlflow_repository.detect_and_parse_datetime_index = MagicMock()
+    mlflow_repository.get_prediction_data = MagicMock(return_value=DataFrame())
+
+    data = MagicMock()
+
+    await mlflow_repository.fit_models(
+        'model_name',
+        data,
+        'latest_production_id',
+        metadata['metadata'],
+        'pyfunc',
+        False,
+        'pyfunc',
+        'target',
+    )
+
+    data_model.fit.return_value.predict.assert_called_once_with({}, data)
+
+
 @patch('laborious.utils.repository.model_repository.pd.merge')
 @patch('laborious.utils.repository.model_repository.isinstance')
 @pytest.mark.asyncio