app.py

import sys
from fastapi import FastAPI
from controller import Controller
from database import Database
import logging
import warnings
from sklearn.exceptions import InconsistentVersionWarning
from contextlib import asynccontextmanager

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(sys.stdout)
    ]
)

logger = logging.getLogger(__name__)
warnings.filterwarnings('ignore', category=InconsistentVersionWarning)
logging.getLogger('absl').setLevel(logging.ERROR)

@asynccontextmanager
async def lifespan(app: FastAPI):
    global controller
    try:
        logger.info("Initializing Controller and Models on startup...")
        controller = Controller(database_sensor={})
        logger.info("Controller and Models loaded successfully!")
    except Exception as e:
        logger.critical(f"Failed to initialize Controller on startup: {str(e)}")
        raise
    yield
    logger.info("Shutting down application...")

app = FastAPI(lifespan=lifespan)
database = Database()
controller: Controller | None = None
        
@app.get("/")
def greet_json():
    return {"Hello": "World!"}

@app.get("/predict-machine")
def predict_machine():
    logger.info("Prediction request received")
    
    machine_ids = database.get_all_machine_id()
    
    if not machine_ids or not hasattr(machine_ids, 'data') or not machine_ids.data:
        logger.warning("No machines found in database")
        return {
            "success": False,
            "error": "No machines found in database"
        }
    
    all_results = []
    
    for machine in machine_ids.data:  
        machine_id = machine.get('id')
        machine_name = machine.get('name')
        sensor = database.get_sensor_readings(1, machine_id)

        # Check if sensor is a duplicate error response
        if isinstance(sensor, dict) and sensor.get("message") == "Data already predicted":
            logger.warning("Data already predicted - returning error")
            all_results.append({
                "machine_name": machine_name,
                "success": False,
                "error": "Data already predicted"
            })
            continue
    
        if sensor is None:
            logger.warning("No sensor data available")
            all_results.append({
                "machine_name": machine_name,
                "success": False,
                "error": "No sensor data available"
            })
            continue
        
        sensor_udi = sensor.get("udi")
        logger.debug(f"Processing sensor data for : {machine_name}, UDI: {sensor_udi}")
        
        controller.set_sensor_data(sensor)
        
        binary_result = controller.predict_binary()
        if not binary_result.get("success"):
            logger.error(f"Binary prediction failed for {machine_name}: {binary_result.get('error')}")
            if binary_result.get("error") != "Data already predicted":
                database.reset_last_processed_id(machine_id)
            all_results.append(binary_result)
            continue
        
        if binary_result.get("failure_predicted"):
            logger.info("Failure predicted - running classification and time series analysis")
            classification_result = controller.predict_classification()
            
            sensor_lstm = database.get_sensor_readings(30, machine_id=machine_id)
            if sensor_lstm is None or (isinstance(sensor_lstm, dict) and sensor_lstm.get("message")):
                logger.warning(f"Not enough time-series data available for {machine_name}")
                sensor_lstm = []
            
            # Set time series data for RUL prediction
            if sensor_lstm:
                controller.set_sensor_data(sensor_lstm)
            else:
                controller.set_sensor_data(sensor)
            
            time_series_result = controller.predict_time_series()
            
            if classification_result.get("success") and time_series_result.get("success"):
                logger.info(f"Prediction successful for {machine_name}- failure_type: {classification_result.get('failure_type')}")

                save_result = database.update_or_create_new_predictions(
                    machine_id=sensor.get("machine_id"),
                    timestamp=sensor.get("timestamp"),
                    risk_score=classification_result.get("risk_score"),
                    failure_predicted=True,
                    failure_type=classification_result.get("failure_type"),
                    predicted_failure_time=time_series_result.get("predictions", {}).get("predicted_failure_date"),
                    confidence=classification_result.get("confidence")
                )
                
                if save_result is None or save_result.get("error"):
                    logger.error(f"Failed to save prediction for UDI {sensor_udi} - resetting for retry")
                    database.reset_last_processed_id(machine_id)
                    all_results.append({
                        "machine_name": machine_name,
                        "success": False,
                        "error": "Failed to save prediction to database"
                    })
                    continue
                logger.info(f"Prediction successfully saved for UDI {sensor_udi}")
                database.mark_as_processed(machine_id, sensor_udi)
                
                all_results.append({
                    "machine_name": machine_name,
                    "success": True,
                    "failure_predicted": True,
                    "failure_type": classification_result.get("failure_type"),
                    "confidence": classification_result.get("confidence"),
                    "risk_score": classification_result.get("risk_score"),
                    "risk_level": classification_result.get("risk_level"),
                    "all_probabilities": classification_result.get("all_probabilities"),
                    "rul_prediction": time_series_result.get("predictions"),
                    "timestamp": sensor.get("timestamp"),
                    "save_data": save_result
                })
            else:
                logger.error(f"Classification or time series prediction failed for UDI {sensor_udi}")
                if not (isinstance(classification_result, dict) and classification_result.get("error") == "Data already predicted"):
                    database.reset_last_processed_id(machine_id)
                all_results.append(binary_result)
        else:
            logger.info(f"No failure predicted for UDI {sensor_udi}")
            save_result = database.update_or_create_new_predictions(
                machine_id=sensor.get("machine_id"),
                timestamp=sensor.get("timestamp"),
                risk_score=binary_result.get("risk_score"),
                failure_predicted=binary_result.get("failure_predicted"),
                failure_type=None,
                predicted_failure_time=None,
                confidence=binary_result.get("confidence")
            )
            
            if save_result and save_result.get("success"):
                database.mark_as_processed(machine_id, sensor_udi)
            
            all_results.append({
                "machine_name": machine_name,
                "success": binary_result.get("success"),
                "failure_predicted": binary_result.get("failure_predicted"),
                "risk_score": binary_result.get("risk_score"),
                "confidence": binary_result.get("confidence"),
                "timestamp": sensor.get("timestamp"),
                "save_data": save_result
            })
    return {
        "success": True,
        "machines_processed": len([r for r in all_results if r.get("success")]),
        "total_machines": len(all_results),
        "results": all_results
    }