Update backend/api/main.py

backend/api/main.py

@@ -1,80 +1,84 @@
-from fastapi import FastAPI, UploadFile, File, HTTPException
-import pandas as pd
-from backend.core.orchestrator import Orchestrator
-
-app = FastAPI()
-orchestrator = Orchestrator()
-
-@app.post("/analyze")
-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 = {
-            "columns": list(df.columns),
-            "dataTypes": dataset_info.get("data_types", {}),
-            "risks": dataset_info.get("risks", []),
-            "problemType": result.get("problem_type"),
-            "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"):
-    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 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"):
-    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", [])
-        }
-    except Exception as e:
-        raise HTTPException(status_code=400, detail=str(e))
-
-@app.post("/predict")
-async def predict(data: dict):
-    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()}
-    except Exception as e:
-        raise HTTPException(status_code=400, detail=str(e))
-
-@app.get("/health")
-def health():
-    return {"status": "ok"}
-
-
-
-
+from fastapi import FastAPI, UploadFile, File, HTTPException
+import pandas as pd
+from backend.core.orchestrator import Orchestrator
+
+app = FastAPI()
+orchestrator = Orchestrator()
+
+@app.post("/analyze")
+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 = {
+            "columns": list(df.columns),
+            "dataTypes": dataset_info.get("data_types", {}),
+            "risks": dataset_info.get("risks", []),
+            "problemType": result.get("problem_type"),
+            "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"):
+    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 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"):
+    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", [])
+        }
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=str(e))
+
+@app.post("/predict")
+async def predict(data: dict):
+    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()}
+    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"}
+
+
+
+