Update app.py

app.py

@@ -2,6 +2,7 @@
 import joblib
 import pandas as pd
 from flask import Flask, request, jsonify
+from utils.validation import validate_and_prepare_input, InputValidationError
 
 # Initialize Flask app with a name
 pred_mainteanance_api = Flask ("Engine Maintenance Predictor")
@@ -20,32 +21,34 @@ def predict_need_maintenance ():
     # Get JSON data from the request
     engine_sensor_inputs = request.get_json ()
 
-    import datetime
 
-    current_year = datetime.datetime.now ().year   # dynamic current year
+    try:
+        input_json = request.get_json()
+        input_df = pd.DataFrame([input_json])
 
-    # Extract relevant features from the input data
-    data_info = {
-        'Engine_rpm'                : engine_sensor_inputs ['Engine_rpm'],
-        'Lub_oil_pressure'          : engine_sensor_inputs ['Lub_oil_pressure'],
-        'Fuel_pressure'             : engine_sensor_inputs ['Fuel_pressure'],
-        'Coolant_pressure'          : engine_sensor_inputs ['Coolant_pressure'],
-        'lub_oil_temp'              : engine_sensor_inputs ['lub_oil_temp'],
-        'Coolant_temp'              : engine_sensor_inputs ['Coolant_temp']
-    }
+        validated_df = validate_and_prepare_input(input_df, model)
 
-    # Convert the extracted data into a DataFrame
-    input_data = pd.DataFrame ([data_info])
+        prediction = model.predict(validated_df)[0]
 
-    # Enforce types - convert all to float
-    input_data = input_data.astype (float)
+        return jsonify({
+            "status": "success",
+            "prediction": int(prediction)
+        })
 
-    # Make prediction using the trained model
-    predicted_sales = model.predict (input_data).tolist ()[0]
-
-    # Return the prediction as a JSON response
-    return jsonify ({'NeedsMaintenance': predicted_sales})
+    except InputValidationError as e:
+        return jsonify({
+            "status": "error",
+            "error_type": "validation_error",
+            "message": str(e)
+        }), 400
 
+    except Exception as e:
+        return jsonify({
+            "status": "error",
+            "error_type": "internal_error",
+            "message": "Unexpected server error"
+        }), 500
+        
 
 # Run the Flask app
 if __name__ == "__main__":