Sync from GitHub: 7341e56760b533299dbaa7cc93ac040d699ffe2c

main.py

@@ -7,22 +7,26 @@ import os
 import sys
 import time
 import uuid
+import shutil
+import json
 from typing import Optional, Dict, List
 from datetime import datetime
 
 import pandas as pd
 import numpy as np
-from fastapi import FastAPI, HTTPException, Header
+from fastapi import FastAPI, HTTPException, Header, BackgroundTasks, UploadFile, File, Form
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse, StreamingResponse
 from pydantic import BaseModel, Field, validator
 import uvicorn
+from supabase import create_client, Client
 
 # Add current directory to path for imports
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
 
 from inference import FraudInference, load_inference_engine
 from simulation import simulation_manager, SimulationConfig
+from training_service import train_model_async
 import warnings
 
 # Suppress warnings
@@ -111,6 +115,32 @@ class BatchPredictResponse(BaseModel):
     processing_time_ms: int
     total_transactions: int
 
+class TrainingConfig(BaseModel):
+    """Configuration for model training"""
+    test_split: float = Field(default=0.05, ge=0.01, le=0.5)
+    n_estimators: int = Field(default=300, ge=10, le=2000)
+    max_depth: int = Field(default=6, ge=1, le=20)
+    learning_rate: float = Field(default=0.05, ge=0.001, le=1.0)
+    pagerank_limit: int = Field(default=10000, ge=0)
+    name: str = Field(default="Custom Model")
+    version: str = Field(default="1.0")
+
+class TrainingResponse(BaseModel):
+    """Response for training initiation"""
+    job_id: str
+    status: str
+    message: str
+
+class ModelItem(BaseModel):
+    """Model registry item"""
+    id: str
+    name: str
+    version: Optional[str]
+    status: str
+    metrics: Optional[Dict]
+    is_active: bool
+    created_at: str
+
 # ============================================================================
 # FASTAPI APP
 # ============================================================================
@@ -142,6 +172,11 @@ model_loading_lock = False
 MODEL_VERSION = "1.0.0"
 MODEL_THRESHOLD = float(os.getenv("MODEL_THRESHOLD", "0.0793"))
 
+# Supabase Client
+supabase_url: str = os.environ.get("SUPABASE_URL")
+supabase_key: str = os.environ.get("SUPABASE_SERVICE_ROLE_KEY") # Use Service Role Key for backend updates
+supabase: Client = create_client(supabase_url, supabase_key) if supabase_url and supabase_key else None
+
 # Risk thresholds (matching config.py)
 RISK_THRESHOLDS = {
     "pass": 0.30,
@@ -803,6 +838,159 @@ async def stream_simulation():
         media_type="text/event-stream"
     )
 
+# ============================================================================
+# TRAINING ENDPOINTS
+# ============================================================================
+
+@app.post("/train", response_model=TrainingResponse)
+async def train_model_endpoint(
+    background_tasks: BackgroundTasks,
+    config: str = Form(...),
+    file: UploadFile = File(...)
+):
+    """
+    Start a background training job.
+    Accepts a CSV file and a JSON string for configuration.
+    """
+    try:
+        config_dict = json.loads(config)
+        # Validate config using Pydantic
+        training_config = TrainingConfig(**config_dict)
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=f"Invalid configuration JSON: {str(e)}")
+
+    if not file.filename.endswith('.csv'):
+        raise HTTPException(status_code=400, detail="Only CSV files are supported")
+
+    job_id = str(uuid.uuid4())
+    
+    # Save uploaded file to temp location
+    temp_dir = "temp_uploads"
+    os.makedirs(temp_dir, exist_ok=True)
+    file_path = os.path.join(temp_dir, f"{job_id}_{file.filename}")
+    
+    try:
+        with open(file_path, "wb") as buffer:
+            shutil.copyfileobj(file.file, buffer)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Failed to save file: {str(e)}")
+
+    # Create initial record in Supabase
+    if supabase:
+        try:
+            supabase.table("model_registry").insert({
+                "id": job_id,
+                "name": training_config.name,
+                "version": training_config.version,
+                "status": "pending",
+                "is_active": False
+            }).execute()
+        except Exception as e:
+            # Clean up file if DB insert fails
+            if os.path.exists(file_path):
+                os.remove(file_path)
+            raise HTTPException(status_code=500, detail=f"Failed to create job record: {str(e)}")
+    else:
+        print("⚠️ Supabase not connected. Training will proceed but status won't be persisted.")
+
+    # Trigger background task
+    params = training_config.dict()
+    background_tasks.add_task(train_model_async, job_id, file_path, params)
+
+    return TrainingResponse(
+        job_id=job_id,
+        status="pending",
+        message="Training job started in background"
+    )
+
+@app.get("/models", response_model=List[ModelItem])
+async def list_models():
+    """List all trained models from registry"""
+    if not supabase:
+        # Return mock data if no DB (or raise error)
+        # For development ease, we can return checking a local dir, but let's stick to DB
+        return []
+        
+    try:
+        response = supabase.table("model_registry").select("*").order("created_at", desc=True).execute()
+        
+        # Convert created_at to string if needed or Pydantic handles it
+        models = []
+        for item in response.data:
+            models.append(ModelItem(
+                id=item['id'],
+                name=item['name'],
+                version=item.get('version'),
+                status=item['status'],
+                metrics=item.get('metrics'),
+                is_active=item.get('is_active', False),
+                created_at=item['created_at']
+            ))
+        return models
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Failed to fetch models: {str(e)}")
+
+@app.post("/models/{model_id}/activate")
+async def activate_model(model_id: str):
+    """Activate a specific model version"""
+    if not supabase:
+        raise HTTPException(status_code=503, detail="Database not available")
+
+    # 1. Get model details
+    try:
+        response = supabase.table("model_registry").select("*").eq("id", model_id).single().execute()
+        model_data = response.data
+    except Exception:
+        raise HTTPException(status_code=404, detail="Model not found")
+
+    if model_data['status'] != 'ready':
+        raise HTTPException(status_code=400, detail="Model is not in 'ready' state")
+
+    model_path = model_data.get('file_path')
+    if not model_path:
+        raise HTTPException(status_code=400, detail="Model file path missing")
+
+    # 2. Update DB (set active=True for this, False for others)
+    # The trigger 'trigger_ensure_single_active_model' handles the "False for others" part!
+    try:
+        supabase.table("model_registry").update({"is_active": True}).eq("id", model_id).execute()
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Failed to update database: {str(e)}")
+
+    # 3. Hot-swap in memory
+    # We need to set the environment variable or just reload logic
+    # Best way: Update global variable directly
+    full_path = os.path.abspath(model_path) if os.path.isabs(model_path) else os.path.join(os.path.dirname(__file__), model_path)
+    
+    if not os.path.exists(full_path):
+        # Try relative to Models/ if not found
+        full_path = os.path.join(os.path.dirname(__file__), "Models", os.path.basename(model_path))
+
+    if not os.path.exists(full_path):
+         raise HTTPException(status_code=500, detail=f"Model file not found on disk: {full_path}")
+
+    # Set env var for future reloads
+    os.environ["MODEL_PATH"] = full_path
+    
+    # Trigger reload
+    global inference_engine
+    try:
+        # Load new model
+        new_engine = load_inference_engine(
+            model_path=full_path,
+            test_dataset_path=os.getenv("TEST_DATASET_PATH"),
+            threshold=float(os.getenv("MODEL_THRESHOLD", "0.0793")),
+            groq_api_key=os.getenv("GROQ_API_KEY"),
+            pagerank_limit=int(os.getenv("PAGERANK_LIMIT", "10000"))
+        )
+        inference_engine = new_engine
+        print(f"✅ Hot-swapped to model {model_id} ({full_path})")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Failed to load model into memory: {str(e)}")
+
+    return {"status": "success", "message": f"Model {model_id} activated"}
+
+
 # ============================================================================
 # ERROR HANDLERS
 # ============================================================================

requirements.txt

@@ -32,4 +32,5 @@ scipy>=1.10.0
 
 # CORS support (included in FastAPI)
 # python-multipart for file uploads (if needed)
+supabase>=2.0.0
 

training_service.py

@@ -0,0 +1,166 @@
+import os
+import time
+import json
+import joblib
+import pandas as pd
+import numpy as np
+import xgboost as xgb
+from sklearn.pipeline import Pipeline
+from sklearn.model_selection import train_test_split, RandomizedSearchCV, cross_val_predict
+from sklearn.metrics import precision_recall_curve, classification_report, accuracy_score, f1_score, precision_score, recall_score
+from supabase import create_client, Client
+from dotenv import load_dotenv
+
+# Import local modules
+# Assuming feature_engineering.py is in the same directory
+try:
+    from feature_engineering import FraudFeatureEngineer
+except ImportError:
+    import sys
+    sys.path.append(os.path.dirname(os.path.abspath(__file__)))
+    from feature_engineering import FraudFeatureEngineer
+
+load_dotenv()
+
+# Initialize Supabase client
+url: str = os.environ.get("SUPABASE_URL")
+key: str = os.environ.get("SUPABASE_SERVICE_ROLE_KEY") # Use Service Role Key for backend updates
+supabase: Client = create_client(url, key) if url and key else None
+
+def update_job_status(job_id: str, status: str, metrics: dict = None, file_path: str = None):
+    """Update job status in Supabase"""
+    if not supabase:
+        print(f"⚠️ Supabase not configured. Job {job_id} status: {status}")
+        return
+
+    data = {"status": status}
+    if metrics:
+        data["metrics"] = metrics
+    if file_path:
+        data["file_path"] = file_path
+    
+    try:
+        supabase.table("model_registry").update(data).eq("id", job_id).execute()
+    except Exception as e:
+        print(f"❌ Failed to update Supabase for job {job_id}: {str(e)}")
+
+def make_splits(df, test_frac=0.05, time_col='step', min_test_fraud=100, random_state=42):
+    """
+    Returns: X_train, X_test, y_train, y_test
+    """
+    # Temporal test split
+    if time_col in df.columns:
+        cutoff = df[time_col].quantile(1 - test_frac)
+        train_df = df[df[time_col] <= cutoff].reset_index(drop=True)
+        test_df = df[df[time_col] > cutoff].reset_index(drop=True)
+    else:
+        train_df = df.copy()
+        test_df = pd.DataFrame(columns=df.columns)
+
+    # Fallback to stratified split if test set is poor
+    if (test_df.empty) or (test_df['isFraud'].sum() < min_test_fraud):
+        X_temp = df.drop(columns=['isFraud'])
+        y_temp = df['isFraud']
+        X_train, X_test, y_train, y_test = train_test_split(
+            X_temp, y_temp, test_size=test_frac, stratify=y_temp, random_state=random_state
+        )
+        return X_train.reset_index(drop=True), X_test.reset_index(drop=True), y_train.reset_index(drop=True), y_test.reset_index(drop=True)
+    else:
+        X_train = train_df.drop(columns=['isFraud'])
+        y_train = train_df['isFraud']
+        X_test = test_df.drop(columns=['isFraud'])
+        y_test = test_df['isFraud']
+        return X_train.reset_index(drop=True), X_test.reset_index(drop=True), y_train.reset_index(drop=True), y_test.reset_index(drop=True)
+
+async def train_model_async(job_id: str, file_path: str, params: dict):
+    """
+    Background task to train model
+    """
+    print(f"🚀 Starting training job {job_id} with params: {params}")
+    update_job_status(job_id, "training")
+    
+    try:
+        # 1. Load Data
+        if not os.path.exists(file_path):
+            raise FileNotFoundError(f"Dataset not found at {file_path}")
+            
+        df = pd.read_csv(file_path)
+        
+        # Basic Preprocessing
+        if 'oldbalanceOrg' in df.columns:
+            df = df.rename(columns={'oldbalanceOrg':'oldBalanceOrig', 'newbalanceOrig':'newBalanceOrig', 
+                                    'oldbalanceDest':'oldBalanceDest', 'newbalanceDest':'newBalanceDest'})
+            
+        # Filter types if needed (Notebook filters to TRANSFER/CASH_OUT)
+        # We'll stick to what the notebook did for consistency
+        if 'type' in df.columns:
+            df = df[df['type'].isin(['TRANSFER', 'CASH_OUT'])].reset_index(drop=True)
+            
+        # 2. Split Data
+        test_frac = params.get('test_split', 0.05)
+        X_train, X_test, y_train, y_test = make_splits(df, test_frac=test_frac)
+        
+        print(f"📊 Data split: Train={len(X_train)}, Test={len(X_test)}")
+        
+        # 3. Feature Engineering
+        pagerank_limit = params.get('pagerank_limit', 10000)
+        fe = FraudFeatureEngineer(pagerank_limit=pagerank_limit)
+        
+        print("⚙️ Fitting feature engineer...")
+        fe.fit(X_train, y_train)
+        X_train_trans = fe.transform(X_train)
+        X_test_trans = fe.transform(X_test)
+        
+        # 4. Configure XGBoost
+        # Calculate scale_pos_weight
+        neg = (y_train == 0).sum()
+        pos = (y_train == 1).sum()
+        spw = int(max(1, neg / max(1, pos)))
+        
+        xgb_params = {
+            'objective': 'binary:logistic',
+            'tree_method': 'hist',
+            'random_state': 42,
+            'n_jobs': 1,
+            'n_estimators': int(params.get('n_estimators', 300)),
+            'max_depth': int(params.get('max_depth', 6)),
+            'learning_rate': float(params.get('learning_rate', 0.05)),
+            'scale_pos_weight': spw
+        }
+        
+        clf = xgb.XGBClassifier(**xgb_params)
+        
+        # 5. Train
+        print("🏋️ Training classifier...")
+        clf.fit(X_train_trans, y_train)
+        
+        # 6. Evaluate
+        print("📝 Evaluating...")
+        y_pred = clf.predict(X_test_trans)
+        
+        metrics = {
+            "accuracy": float(accuracy_score(y_test, y_pred)),
+            "precision": float(precision_score(y_test, y_pred)),
+            "recall": float(recall_score(y_test, y_pred)),
+            "f1": float(f1_score(y_test, y_pred))
+        }
+        print(f"✅ Metrics: {metrics}")
+        
+        # 7. Save Model
+        # Ensure Models directory exists
+        models_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "Models")
+        os.makedirs(models_dir, exist_ok=True)
+        
+        model_filename = f"model_{job_id}.pkl"
+        model_save_path = os.path.join(models_dir, model_filename)
+        
+        # Save ONLY the classifier (as per notebook strategy)
+        joblib.dump(clf, model_save_path)
+        print(f"💾 Model saved to {model_save_path}")
+        
+        # 8. Update Registry
+        update_job_status(job_id, "ready", metrics=metrics, file_path=f"Models/{model_filename}")
+        
+    except Exception as e:
+        print(f"❌ Training failed: {str(e)}")
+        update_job_status(job_id, "failed", metrics={"error": str(e)})