src/models/api.py

import os
import warnings
# Suppress deprecation and future warnings from third-party libraries (like MLflow)
warnings.filterwarnings("ignore", category=FutureWarning)

os.environ["MLFLOW_ALLOW_FILE_STORE"] = "true"

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import mlflow
import pandas as pd
import logging
from contextlib import asynccontextmanager

logging.basicConfig(level=logging.INFO)

# Configuration
# Force API to use local models bundled via Git LFS instead of downloading from DagsHub
MLFLOW_TRACKING_URI = "file:///app/mlruns" if os.path.exists("/app/mlruns") else "file:./mlruns"
EXPERIMENT_NAME = "SL_Commodity_Forecasting"

# Global model variable
model = None

COMMODITY_MAP = {
    "Samba": 0,
    "Kekulu": 1,
    "Big Onion": 2,
    "Potato": 3,
    "Dried Chilli": 4,
    "Coconut": 5
}

class PredictionRequest(BaseModel):
    commodity: str
    lag_prices: list[float]
    
    model_config = {
        "json_schema_extra": {
            "example": {
                "commodity": "Samba",
                "lag_prices": [220.5, 219.0, 218.5, 221.0, 222.5, 220.0, 219.5]
            }
        }
    }

def load_best_model():
    """Loads the best XGBoost model from the local MLflow registry on startup."""
    global model
    mlflow.set_tracking_uri(MLFLOW_TRACKING_URI)
    
    # Configure DagsHub credentials automatically if URI is provided
    dagshub_token = os.environ.get("DAGSHUB_USER_TOKEN")
    if dagshub_token and "dagshub.com" in MLFLOW_TRACKING_URI:
        os.environ["MLFLOW_TRACKING_USERNAME"] = MLFLOW_TRACKING_URI.split('/')[-2]
        os.environ["MLFLOW_TRACKING_PASSWORD"] = dagshub_token
        logging.info("Configured DagsHub MLflow remote for API")
    
    try:
        experiment = mlflow.get_experiment_by_name(EXPERIMENT_NAME)
        if not experiment:
            logging.error(f"Experiment '{EXPERIMENT_NAME}' not found. Make sure you trained the model.")
            return

        # Fetch the best XGBoost run by RMSE
        runs = mlflow.search_runs(
            experiment_ids=[experiment.experiment_id],
            filter_string="tags.mlflow.runName = 'XGBoost_Candidate'",
            order_by=["metrics.rmse ASC"],
            max_results=1
        )
        
        is_empty = False
        if runs is None:
            is_empty = True
        elif hasattr(runs, "empty"):
            is_empty = runs.empty
        else:
            is_empty = not runs

        if is_empty:
            logging.error("No XGBoost runs found in MLflow.")
            return
            
        if hasattr(runs, "iloc"):
            best_run_id = runs.iloc[0].run_id
        else:
            # Handle list/PagedList of Run objects robustly
            best_run = runs[0]
            best_run_id = None
            
            # 1. Try direct attribute/key 'run_id'
            if hasattr(best_run, "run_id"):
                best_run_id = best_run.run_id
            elif isinstance(best_run, dict) and "run_id" in best_run:
                best_run_id = best_run["run_id"]
                
            # 2. Try info attribute or callable info
            if not best_run_id and hasattr(best_run, "info"):
                info_attr = best_run.info
                if callable(info_attr):
                    try:
                        info_obj = info_attr()
                        best_run_id = getattr(info_obj, "run_id", None)
                    except Exception:
                        pass
                if not best_run_id:
                    best_run_id = getattr(info_attr, "run_id", None)
                    
            # 3. Try info key/dict
            if not best_run_id and isinstance(best_run, dict) and "info" in best_run:
                info_val = best_run["info"]
                if isinstance(info_val, dict):
                    best_run_id = info_val.get("run_id")
                else:
                    best_run_id = getattr(info_val, "run_id", None)
                    
            # 4. Last resort fallback
            if not best_run_id:
                best_run_id = getattr(best_run, "run_id", None)
        
        # Try finding the model in the git-tracked models folder
        model_found = False
        model_uri = None
        models_dir = os.path.join("mlruns", experiment.experiment_id, "models")
        if os.path.exists(models_dir):
            for model_id in os.listdir(models_dir):
                meta_path = os.path.join(models_dir, model_id, "meta.yaml")
                if os.path.exists(meta_path):
                    with open(meta_path, "r") as f:
                        if f"source_run_id: {best_run_id}" in f.read():
                            model_uri = os.path.join(models_dir, model_id, "artifacts")
                            model_found = True
                            break
                            
        if not model_found or model_uri is None:
            model_uri = f"runs:/{best_run_id}/xgboost_model"
            
        logging.info(f"Loading latest XGBoost model (Run ID: {best_run_id}) from {model_uri}")
        model = mlflow.xgboost.load_model(model_uri)
        logging.info("✅ Model loaded successfully and ready for inference.")
        
    except Exception as e:
        logging.error(f"❌ Failed to load model: {e}")

@asynccontextmanager
async def lifespan(app: FastAPI):
    load_best_model()
    yield

app = FastAPI(
    title="SL Commodity Forecasting API", 
    description="MLOps API serving the XGBoost model for Sri Lankan commodity prices",
    version="1.0",
    lifespan=lifespan
)

@app.get("/")
def health_check():
    return {
        "service": "MLOps Forecasting API",
        "status": "healthy", 
        "model_loaded": model is not None
    }

@app.post("/predict")
def predict_price(request: PredictionRequest):
    if not model:
        raise HTTPException(status_code=503, detail="Model is currently unavailable or failed to load.")
    
    if len(request.lag_prices) != 7:
        raise HTTPException(status_code=400, detail="Exactly 7 lag prices must be provided (chronological order: oldest to newest).")
        
    if request.commodity not in COMMODITY_MAP:
        raise HTTPException(status_code=400, detail=f"Unknown commodity. Must be one of {list(COMMODITY_MAP.keys())}")
        
    try:
        current_lags = request.lag_prices.copy()
        forecast = []
        
        # 7-day recursive forecasting
        for day in range(7):
            # Construct the feature DataFrame based on what the model expects
            expected_features = model.get_booster().feature_names
            features: dict[str, list] = {}
            if "Commodity_ID" in expected_features:
                features["Commodity_ID"] = [COMMODITY_MAP[request.commodity]]
            for i in range(1, 8):
                if f"Lag_{i}" in expected_features:
                    features[f"Lag_{i}"] = [current_lags[-i]]
                
            # Ensure columns are in the exact order the model expects
            df_features = pd.DataFrame(features)[expected_features]
            
            # Predict
            prediction = float(model.predict(df_features)[0])
            forecast.append(round(prediction, 2))
            
            # Append new prediction to the end of the lags (it becomes the new t-1 for the next iteration)
            current_lags.append(prediction)

        return {
            "predicted_price_lkr": forecast[0],
            "7_day_forecast": forecast,
            "currency": "LKR"
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Prediction error: {str(e)}")