Update app.py

app.py

@@ -8,9 +8,32 @@ from typing import List
 
 app = FastAPI(title="Headache Predictor API")
 
+# Global variables for model and threshold
 clf = None
 threshold = 0.5
 
+# --- Pydantic Models ---
+
+class SinglePredictionRequest(BaseModel):
+    features: List[float]
+
+class SinglePredictionResponse(BaseModel):
+    prediction: int
+    probability: float
+
+class BatchPredictionRequest(BaseModel):
+    instances: List[List[float]]
+
+class DayPrediction(BaseModel):
+    day: int
+    prediction: int
+    probability: float
+
+class BatchPredictionResponse(BaseModel):
+    predictions: List[DayPrediction]
+
+# --- Startup Event ---
+
 @app.on_event("startup")
 async def load_model():
     global clf, threshold
@@ -32,6 +55,7 @@ async def load_model():
 
         if isinstance(model_data, dict):
             clf = model_data["model"]
+            # Load threshold if available, otherwise default to 0.5
             threshold = float(model_data.get("optimal_threshold", 0.5))
             print(f"✅ Model loaded (optimal_threshold={threshold})")
         else:
@@ -44,19 +68,67 @@ async def load_model():
         import traceback
         traceback.print_exc()
 
-class BatchPredictionRequest(BaseModel):
-    instances: List[List[float]]
+# --- Endpoints ---
+
+@app.get("/")
+def read_root():
+    return {
+        "message": "Headache Predictor API",
+        "status": "running",
+        "endpoints": {
+            "predict": "/predict - Single day prediction",
+            "predict_batch": "/predict/batch - 7-day forecast",
+            "health": "/health"
+        },
+        "examples": {
+            "single": {
+                "url": "/predict",
+                # Example shortened for brevity in display
+                "body": {"features": [1, 0, 0, 0, 1, 0, 1005.0, -9.5, 85.0, 15.5, 64.0, 5.5, 41.0, 0.0, 1, 1, 1, 0, 0, 1, 0, 1, 1, 1, 1, 0, 0, 0, 10, 40, 4, 7.0, 50.0, 60.0, 3.5, 1.5, 6.8]}
+            },
+            "batch": {
+                "url": "/predict/batch",
+                "body": {"instances": [["array of 37 features for day 1"], ["array for day 2"]]}
+            }
+        }
+    }
+
+@app.get("/health")
+def health_check():
+    return {
+        "status": "healthy",
+        "model_loaded": clf is not None
+    }
+
+@app.post("/predict", response_model=SinglePredictionResponse)
+def predict(request: SinglePredictionRequest):
+    """Predict headache risk for a single day"""
+    if clf is None:
+        raise HTTPException(status_code=503, detail="Model not loaded")
 
-class DayPrediction(BaseModel):
-    day: int
-    prediction: int
-    probability: float
+    try:
+        # Convert input to numpy array
+        features = np.array(request.features).reshape(1, -1)
 
-class BatchPredictionResponse(BaseModel):
-    predictions: List[DayPrediction]
+        # Get probability array for both classes [prob_0, prob_1]
+        prob_array = clf.predict_proba(features)[0]
+        
+        # Always return probability of headache (class 1)
+        headache_probability = float(prob_array[1])
+
+        # Make prediction using the loaded threshold
+        prediction = 1 if headache_probability >= threshold else 0
+
+        return SinglePredictionResponse(
+            prediction=int(prediction),
+            probability=headache_probability
+        )
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=f"Prediction error: {str(e)}")
 
 @app.post("/predict/batch", response_model=BatchPredictionResponse)
 def predict_batch(request: BatchPredictionRequest):
+    """Predict headache risk for multiple days (batch)"""
     if clf is None:
         raise HTTPException(status_code=503, detail="Model not loaded")
 
@@ -66,6 +138,8 @@ def predict_batch(request: BatchPredictionRequest):
             raise ValueError(f"Expected 2D array, got shape {X.shape}")
 
         probas = clf.predict_proba(X)[:, 1]   # class-1 prob
+        
+        # Use the global threshold
         preds = (probas >= threshold).astype(int)
 
         results = [
@@ -75,4 +149,4 @@ def predict_batch(request: BatchPredictionRequest):
         return BatchPredictionResponse(predictions=results)
 
     except Exception as e:
-        raise HTTPException(status_code=400, detail=f"Batch prediction error: {str(e)}")
+        raise HTTPException(status_code=400, detail=f"Batch prediction error: {str(e)}")