diff --git a/laborious/workflows/sub_workflows/prediction_process.py b/laborious/workflows/sub_workflows/prediction_process.py
index 3560d5f..2c20ce8 100644
--- a/laborious/workflows/sub_workflows/prediction_process.py
+++ b/laborious/workflows/sub_workflows/prediction_process.py
@@ -110,13 +110,13 @@ class PredictionProcess:
         self,
         input_data: dict[str, Any],
         metadata: dict[str, Any],
-        data: Any,
-        model_id: 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
+        last_timestamp = data['last_timestamp']
 
         # Apply input data quality gates
         gate_input = {