api/model/model_controller.py

from datetime import datetime
from typing import Any, Dict

from fastapi.responses import JSONResponse

from api.model.model_service import ModelService
from packages import percentage


class ModelController:
    def __init__(self, model_service: ModelService):
        self.model_service = model_service

    async def predict(self, input_data: Dict[str, Any]) -> JSONResponse:
        if not self.model_service.is_model_loaded():
            raise RuntimeError("Model not loaded on server.")

        df = await self.model_service.preprocess_data(input_data)

        pred, proba_yes, proba_no, priority = await self.model_service.predict(df)

        return JSONResponse(
            status_code=200,
            content={
                "status_code": 200,
                "success": True,
                "timestamp": datetime.utcnow().isoformat(),
                "data": {
                    "predicted_class": "YES" if pred == 1 else "NO",
                    "probability_yes": round(proba_yes, 4)
                    if proba_yes is not None
                    else None,
                    "probability_no": round(proba_no, 4)
                    if proba_no is not None
                    else None,
                    "probability_yes_percentage": percentage(proba_yes)
                    if proba_yes is not None
                    else None,
                    "probability_no_percentage": percentage(proba_no)
                    if proba_no is not None
                    else None,
                    "priority": priority,
                },
                "meta": {
                    "input_shape": list(df.shape),
                },
            },
        )

    async def get_complete_model_info(self) -> JSONResponse:
        model_metadata = self.model_service.get_complete_model_info()
        return JSONResponse(
            status_code=200,
            content={
                "status_code": 200,
                "success": True,
                "timestamp": datetime.utcnow().isoformat(),
                "data": model_metadata,
            },
        )