Update app.py

app.py

@@ -1,12 +1,13 @@
-from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
+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")
 
-# Allow browser requests
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -15,61 +16,63 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Expected number of features
-EXPECTED_FEATURES = 35
+BASE_DIR = Path(__file__).resolve().parent
+MODEL_PATH = BASE_DIR / "bot_detection_model.pkl"
+
 
-# Request schema
 class InputData(BaseModel):
     features: list[float]
 
-# Load model
+
 try:
-    model = joblib.load("bot_detection_model.pkl")
-except Exception as e:
-    raise RuntimeError(f"Model failed to load: {e}")
+    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
+
 
-# Root endpoint
 @app.get("/")
-def home():
+def home() -> dict[str, str | int]:
     return {
         "message": "Bot detection model running",
-        "expected_features": EXPECTED_FEATURES
+        "expected_features": EXPECTED_FEATURES,
     }
 
-# Health check
+
 @app.get("/health")
-def health():
+def health() -> dict[str, str]:
     return {"status": "ok"}
 
-# Prediction endpoint
-@app.post("/predict")
-def predict(data: InputData):
 
-    # Validate feature count
+@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"
+            detail=f"Expected {EXPECTED_FEATURES} features",
         )
 
     try:
-        input_data = np.array(data.features).reshape(1, -1)
+        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")
 
-        prediction = model.predict(input_data)[0]
+        pred = int(model.predict(x)[0])
+        result = "bot_detected" if pred == 1 else "normal_traffic"
 
-        # Convert prediction to readable output
-        if prediction == 1:
-            result = "bot_detected"
-        else:
-            result = "normal_traffic"
-
-        return {
+        response: dict[str, int | float | str] = {
             "prediction": result,
-            "raw_prediction": int(prediction)
+            "raw_prediction": pred,
         }
 
-    except Exception as e:
-        raise HTTPException(
-            status_code=500,
-            detail=f"Prediction failed: {str(e)}"
-        )
+        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
+