SIENTIAPDE-1273

Enhance prediction data retrieval in MLFlowRepository by adding predict_flavor parameter

- Updated get_prediction_data method to accept a predict_flavor argument, allowing for different prediction model handling.
- Adjusted calls to get_prediction_data throughout the codebase to include the new parameter.
- Added new test cases to validate behavior for different predict_flavor values, ensuring robust functionality.

tests/laborious/utils/repository/test_model_repository.py

@@ -865,7 +865,7 @@ async def test_fit_models_not_df_target_name_none_and_not_in_model(
     prediction_model.fit.assert_called_once_with(pd_merge.return_value)
 
     mlflow_repository.get_prediction_data.assert_called_once_with(
-        prediction_model, pd_merge.return_value, data_model.fit.return_value.target_variable
+        prediction_model, pd_merge.return_value, data_model.fit.return_value.target_variable, 'pyfunc'
     )
 
     assert output == {
@@ -951,7 +951,7 @@ async def test_fit_models_df_target_name_not_none_and_in_model(
     prediction_model.fit.assert_called_once_with(transformed_data)
 
     mlflow_repository.get_prediction_data.assert_called_once_with(
-        prediction_model, transformed_data, 'feat_1'
+        prediction_model, transformed_data, 'feat_1', 'pyfunc'
     )
 
     assert output == {
@@ -1494,8 +1494,11 @@ def test_get_prediction_data_dataframe(mlflow_repository):
     retrain_dataset = DataFrame({'feat_1': [1, 2], 'target': [3, 4]}, index=['idx1', 'idx2'])
     prediction_model.predict.return_value = DataFrame({'pred': [5, 6]}, index=['idx1', 'idx2'])
     target_name = 'target'
+    predict_flavor = 'sklearn'
 
-    result = mlflow_repository.get_prediction_data(prediction_model, retrain_dataset, target_name)
+    result = mlflow_repository.get_prediction_data(
+        prediction_model, retrain_dataset, target_name, predict_flavor
+    )
 
     prediction_model.predict.assert_called_once_with(retrain_dataset)
     assert 'prediction' in result.columns
@@ -1509,11 +1512,32 @@ def test_get_prediction_data_array(mlflow_repository):
     retrain_dataset = DataFrame({'feat_1': [1, 2], 'target': [3, 4]}, index=['idx1', 'idx2'])
     prediction_model.predict.return_value = [5, 6]
     target_name = 'target'
+    predict_flavor = 'sklearn'
 
-    result = mlflow_repository.get_prediction_data(prediction_model, retrain_dataset, target_name)
+    result = mlflow_repository.get_prediction_data(
+        prediction_model, retrain_dataset, target_name, predict_flavor
+    )
 
     prediction_model.predict.assert_called_once_with(retrain_dataset)
     assert 'prediction' in result.columns
     assert 'target' in result.columns
     assert 'timestamp' in result.columns
     assert result.index.tolist() == [0, 1]
+
+
+def test_get_prediction_data_pyfunc(mlflow_repository):
+    prediction_model = MagicMock()
+    retrain_dataset = DataFrame({'feat_1': [1, 2], 'target': [3, 4]}, index=['idx1', 'idx2'])
+    prediction_model.predict.return_value = DataFrame({'pred': [5, 6]}, index=['idx1', 'idx2'])
+    target_name = 'target'
+    predict_flavor = 'pyfunc'
+
+    result = mlflow_repository.get_prediction_data(
+        prediction_model, retrain_dataset, target_name, predict_flavor
+    )
+
+    prediction_model.predict.assert_called_once_with({}, retrain_dataset)
+    assert 'prediction' in result.columns
+    assert 'target' in result.columns
+    assert 'timestamp' in result.columns
+    assert result.index.tolist() == [0, 1]