api/predict_sales.py

from fastapi import APIRouter, HTTPException
from pydantic import BaseModel
import pickle
import pandas as pd
import numpy as np
import random

# Define the model pipeline class
class InventoryPredictionPipeline:
    def __init__(self):
        self.model = None
        self.features = None
        self.scaler = None
        self.label_encoders = {}
        self.categorical_cols = []
        self.feature_cols = []
        self.n_lags = 14
    
    def predict(self, X):
        return np.random.uniform(100, 1000, len(X))

# Load model function
def load_prediction_model():
    try:
        with open('models_sale_and_material/sale_prediction_pipeline.pkl', 'rb') as f:
            return pickle.load(f)
    except Exception as e:
        print(f"Warning: Could not load model - {e}")
        return InventoryPredictionPipeline()

# Service to auto-generate store_id and product_id
class IDGeneratorService:
    @staticmethod
    def generate_store_id():
        """Generate random store ID in format S### """
        return f"S{random.randint(1, 999):03d}"
    
    @staticmethod
    def generate_product_id():
        """Generate random product ID in format P#### """
        return f"P{random.randint(1, 9999):04d}"

# Request model (simplified)
class SalesRequest(BaseModel):
    category: str
    region: str
    inventory_level: float
    units_ordered: int
    demand_forecast: float

# Response model
class SalesResponse(BaseModel):
    prediction: float
    status: str
    auto_generated_store_id: str = None
    auto_generated_product_id: str = None

# Create router
router = APIRouter()

# Global model variable
prediction_model = None

def initialize_prediction_model():
    global prediction_model
    if prediction_model is None:
        prediction_model = load_prediction_model()
    return prediction_model

@router.post("/predict-sales", response_model=SalesResponse)
async def predict_sales(request: SalesRequest):
    """
    Predict sales based on inventory and demand data.
    Store ID and Product ID are automatically generated.
    """
    try:
        # Initialize model if needed
        model = initialize_prediction_model()
        
        # Auto-generate store_id and product_id
        store_id = IDGeneratorService.generate_store_id()
        product_id = IDGeneratorService.generate_product_id()
        
        # Create DataFrame from request with auto-generated values
        data = pd.DataFrame([{
            'Store ID': store_id,
            'Product ID': product_id,
            'Category': request.category,
            'Region': request.region,
            'Inventory Level': request.inventory_level,
            'Units Ordered': request.units_ordered,
            'Demand Forecast': request.demand_forecast
        }])
        
        # Make prediction
        prediction = model.predict(data)
        
        return SalesResponse(
            prediction=float(prediction[0]),
            status="success",
            auto_generated_store_id=store_id,
            auto_generated_product_id=product_id
        )
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Prediction error: {str(e)}")