src/agents/l3_classifier.py

import joblib
from fastapi import HTTPException

class Classifier:
    def __init__(self, model_path: str = "src/models/l2_XGBOOST_Supervisor_V5.joblib"):
        try:
            saved = joblib.load(model_path)
            self.model = saved["model"]
            self.scaler = saved.get("scaler", None)  # certains modèles peuvent ne pas avoir de scaler
            self.features = [
                "Header_Length",
                "Time_To_Live",
                "Rate",
                "Tot sum",
                "Tot size",
                "Min",
                "Max",
                "AVG",
                "Std",
                "Variance",
                "IAT",
                "Number",
                "syn_ratio",
                "ack_ratio",
                "fin_ratio",
                "rst_ratio",
                "mean_pkt_size",
                "pkt_size_range",
                "pkt_size_ratio",
                "mean_iat",
                "pkt_rate",
                "throughput",
                "bytes_per_sec",
                "coef_var",
                "tcp_udp_ratio",
                "flag_entropy",
            ]
        except FileNotFoundError:
            raise HTTPException(status_code=500, detail=f"Modèle '{model_path}' introuvable")
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"Erreur lors du chargement du modèle: {e}")

    def predict(self, data):
        try:
            # Préparer les features
            X = data[self.features]

            # Appliquer le scaler si existant
            if self.scaler is not None:
                X = self.scaler.transform(X)

            # Prédictions
            preds = self.model.predict(X)
            return preds.tolist()
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"Erreur lors de la prédiction: {e}")