app.py

from pathlib import Path

import joblib
import numpy as np
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

app = FastAPI(title="Bot Detection API")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

BASE_DIR = Path(__file__).resolve().parent
MODEL_PATH = BASE_DIR / "bot_detection_model.pkl"


class InputData(BaseModel):
    features: list[float]


try:
    model = joblib.load(MODEL_PATH)
    EXPECTED_FEATURES = int(getattr(model, "n_features_in_", 14))
except Exception as ex:
    raise RuntimeError(f"Model failed to load: {ex}") from ex


@app.get("/")
def home() -> dict[str, str | int]:
    return {
        "message": "Bot detection model running",
        "expected_features": EXPECTED_FEATURES,
    }


@app.get("/health")
def health() -> dict[str, str]:
    return {"status": "ok"}


@app.post("/predict")
def predict(data: InputData) -> dict[str, int | float | str]:
    if len(data.features) != EXPECTED_FEATURES:
        raise HTTPException(
            status_code=400,
            detail=f"Expected {EXPECTED_FEATURES} features",
        )

    try:
        x = np.asarray(data.features, dtype=np.float64).reshape(1, -1)
        if not np.isfinite(x).all():
            raise HTTPException(status_code=400, detail="Features contain NaN or Inf")

        pred = int(model.predict(x)[0])
        result = "bot_detected" if pred == 1 else "normal_traffic"

        response: dict[str, int | float | str] = {
            "prediction": result,
            "raw_prediction": pred,
        }

        if hasattr(model, "predict_proba"):
            proba = model.predict_proba(x)[0]
            bot_idx = list(model.classes_).index(1)
            response["bot_probability"] = float(proba[bot_idx])

        return response
    except HTTPException:
        raise
    except Exception as ex:
        raise HTTPException(status_code=500, detail=f"Prediction failed: {ex}") from ex