Update backend/api/main.py

backend/api/main.py

@@ -1,21 +1,29 @@
 from fastapi import FastAPI, UploadFile, File, HTTPException
 import pandas as pd
+import os
+
 from backend.core.orchestrator import Orchestrator
+from backend.api.schemas import PredictRequest
 
 app = FastAPI()
 orchestrator = Orchestrator()
 
+MODEL_PATH = "exports/models/trained_model.pkl"
+
+
 @app.post("/analyze")
-async def analyze_dataset(file: UploadFile = File(...), target_column: str = "target"):
+async def analyze_dataset(
+    file: UploadFile = File(...),
+    target_column: str = "target"
+):
     try:
         df = pd.read_csv(file.file)
         result = orchestrator.run(df, target_column)
 
-        # Format response for frontend
         dataset_info = result.get("dataset_info", {})
         strategy = result.get("strategy", {})
 
-        response = {
+        return {
             "columns": list(df.columns),
             "dataTypes": dataset_info.get("data_types", {}),
             "risks": dataset_info.get("risks", []),
@@ -23,62 +31,84 @@ async def analyze_dataset(file: UploadFile = File(...), target_column: str = "ta
             "confidence": strategy.get("confidence", 0),
             "strategy": strategy
         }
-        return response
+
     except Exception as e:
         raise HTTPException(status_code=400, detail=str(e))
 
+
 @app.post("/train")
-async def train_model(file: UploadFile = File(...), target_column: str = "target"):
+async def train_model(
+    file: UploadFile = File(...),
+    target_column: str = "target"
+):
     try:
         df = pd.read_csv(file.file)
         result = orchestrator.run(df, target_column, train=True)
 
-        # Ensure strategy is included in the response
-        strategy = result.get("strategy", {})
-        response = {
-            "strategy": strategy,
-            "metrics": result.get("metrics", {}),
-            "model_path": result.get("model_path", "/path/to/model.pkl"),
-            "training_time": result.get("training_time", 0),
-            "model_id": result.get("model_id", "trained_model_123")
+        return {
+            "strategy": result.get("strategy"),
+            "metrics": result.get("metrics"),
+            "model_path": MODEL_PATH,
+            "model_id": "trained_model",
         }
-        return response
+
     except Exception as e:
         raise HTTPException(status_code=400, detail=str(e))
 
+
 @app.post("/explain")
-async def explain_model(file: UploadFile = File(...), target_column: str = "target"):
+async def explain_model(
+    file: UploadFile = File(...),
+    target_column: str = "target"
+):
     try:
         df = pd.read_csv(file.file)
         result = orchestrator.run(df, target_column, train=True)
+
         return {
             "strategy_explanation": result.get("strategy_explanation"),
-            "metrics": result.get("metrics", {}),
-            "feature_importance": result.get("feature_importance", [])
+            "metrics": result.get("metrics"),
+            "feature_importance": result.get("feature_importance"),
         }
+
     except Exception as e:
         raise HTTPException(status_code=400, detail=str(e))
 
+
 @app.post("/predict")
-async def predict(data: dict):
+async def predict(request: PredictRequest):
     try:
-        # Load the trained model
-        model = orchestrator.model_io.load("exports/models/trained_model.pkl")
-        # Prepare data for prediction
-        df = pd.DataFrame([data])
-        preds = model.predict(df)
-        return {"prediction": preds.tolist()}
+        if not os.path.exists(MODEL_PATH):
+            raise HTTPException(
+                status_code=400,
+                detail="No trained model found. Train a model first."
+            )
+
+        model = orchestrator.model_io.load(MODEL_PATH)
+
+        df = pd.DataFrame(request.instances)
+
+        if df.empty:
+            raise HTTPException(
+                status_code=400,
+                detail="Prediction data is empty."
+            )
+
+        predictions = model.predict(df)
+
+        return {
+            "predictions": predictions.tolist()
+        }
+
     except Exception as e:
         raise HTTPException(status_code=400, detail=str(e))
 
+
 @app.get("/")
 def root():
     return {"message": "ModelSmith AI API", "status": "running"}
 
+
 @app.get("/health")
 def health():
     return {"status": "ok"}
-
-
-
-