app.py

import gradio as gr
import pandas as pd
import io
import pickle
import matplotlib.pyplot as plt
import seaborn as sns
from lazypredict.Supervised import LazyClassifier, LazyRegressor
from sklearn.model_selection import train_test_split
from ydata_profiling import ProfileReport
import tempfile
import requests
import json
from typing import Optional, Tuple, Any, Union
from openai import OpenAI  # Added for Nebius AI Studio LLM integration

# Constants
NO_TASK_DETECTED = "No task detected"
NO_COLUMNS_LOADED = "No columns loaded."


def load_data(file_input: Any) -> Tuple[Optional[pd.DataFrame], Optional[str]]:
    """
    Loads CSV data from either a local file upload or a public URL.
    
    Args:
        file_input: A file object from Gradio upload or a URL string.
        
    Returns:
        Tuple containing the DataFrame and comma-separated column names,
        or (None, None) if loading fails.
    """
    if file_input is None:
        return None, None

    try:
        if hasattr(file_input, 'name'):
            file_path = file_input.name
            with open(file_path, 'rb') as f:
                file_bytes = f.read()
            df = pd.read_csv(io.BytesIO(file_bytes))
        elif isinstance(file_input, str) and file_input.startswith('http'):
            response = requests.get(file_input, timeout=30)
            response.raise_for_status()
            df = pd.read_csv(io.StringIO(response.text))
        else:
            return None, None

        # Extract column names here
        column_names = ", ".join(df.columns.tolist())
        return df, column_names
    except Exception as e:
        gr.Warning(f"Failed to load or parse data: {e}")
        return None, None


def generate_dataset_summary(df: pd.DataFrame, target_column: str) -> str:
    """
    Generates a concise summary of the dataset for LLM context.
    
    Args:
        df: The pandas DataFrame to summarize.
        target_column: The name of the target column.
        
    Returns:
        A formatted string summary of the dataset.
    """
    summary_parts = [
        f"Dataset Shape: {df.shape[0]} rows, {df.shape[1]} columns",
        f"Target Column: {target_column}",
        f"Target Unique Values: {df[target_column].nunique()}",
        f"Features: {', '.join([col for col in df.columns if col != target_column])}",
        f"Missing Values: {df.isnull().sum().sum()} total",
        f"Numeric Columns: {len(df.select_dtypes(include=['number']).columns)}",
        f"Categorical Columns: {len(df.select_dtypes(include=['object', 'category']).columns)}"
    ]
    return "\n".join(summary_parts)


def update_detected_columns_display(file_data: Any, url_data: Optional[str]) -> str:
    """
    Detects and displays column names from the uploaded file or URL
    as soon as the input changes, before the main analysis button is pressed.
    
    Args:
        file_data: File object from Gradio file upload component.
        url_data: URL string from Gradio textbox component.
        
    Returns:
        Comma-separated string of column names or error message.
    """
    data_source = file_data if file_data is not None else url_data
    if data_source is None:
        return ""

    _, column_names = load_data(data_source)
    if column_names:
        return column_names
    else:
        return "No columns detected or error loading file. Please check the file format."


def analyze_and_model(
    df: pd.DataFrame, 
    target_column: str
) -> Tuple[ProfileReport, str, str, pd.DataFrame, str, str, str]:
    """
    Internal function to perform EDA, model training, and visualization.
    
    Args:
        df: The pandas DataFrame containing the dataset.
        target_column: The name of the target column for prediction.
        
    Returns:
        Tuple containing: profile report, profile path, task type, 
        models dataframe, plot path, pickle path, and best model name.
    """
    profile = ProfileReport(df, title="EDA Report", minimal=True)
    with tempfile.NamedTemporaryFile(delete=False, suffix=".html") as temp_html:
        profile.to_file(temp_html.name)
        profile_path = temp_html.name

    X = df.drop(columns=[target_column])
    y = df[target_column]
    task = "classification" if y.nunique() <= 10 else "regression"
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    lazy_model = LazyClassifier(ignore_warnings=True, verbose=0) if task == "classification" else LazyRegressor(ignore_warnings=True, verbose=0)
    models, _ = lazy_model.fit(X_train, X_test, y_train, y_test)

    sort_metric = "Accuracy" if task == "classification" else "R-Squared"
    sorted_models = models.sort_values(by=sort_metric, ascending=False)
    best_model_name = sorted_models.index[0]
    
    # Safely access the best model with error handling
    try:
        best_model = lazy_model.models[best_model_name]
    except KeyError:
        # Fallback: try to find the model with stripped whitespace
        model_keys = list(lazy_model.models.keys())
        matching_key = next((k for k in model_keys if k.strip() == best_model_name.strip()), None)
        if matching_key:
            best_model = lazy_model.models[matching_key]
        else:
            # Use the first available model as fallback
            best_model = list(lazy_model.models.values())[0]
            gr.Warning(f"Could not find exact model '{best_model_name}', using first available model.")

    with tempfile.NamedTemporaryFile(delete=False, suffix=".pkl") as temp_pkl:
        pickle.dump(best_model, temp_pkl)
        pickle_path = temp_pkl.name

    plt.figure(figsize=(10, 6))
    plot_column = "Accuracy" if task == "classification" else "R-Squared"
    top_models = models.head(10)
    sns.barplot(x=top_models[plot_column].values, y=top_models.index.tolist())
    plt.title(f"Top 10 Models by {plot_column}")
    plt.xlabel(plot_column)
    plt.ylabel("Model")
    plt.tight_layout()
    with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as temp_png:
        plt.savefig(temp_png.name)
        plot_path = temp_png.name
    plt.close()

    models_reset = models.reset_index().rename(columns={'index': 'Model'})
    return profile, profile_path, task, models_reset, plot_path, pickle_path, best_model_name

def run_pipeline(
    data_source: Union[Any, str], 
    target_column: str, 
    nebius_api_key: Optional[str] = None
) -> Tuple[Optional[str], str, Optional[pd.DataFrame], Optional[str], Optional[str], str, str]:
    """
    Run the complete AutoML pipeline including data loading, EDA, model training, and AI explanation.
    
    This is the primary MCP tool function that orchestrates the entire AutoML workflow.
    
    Args:
        data_source: Either a file path/object from local upload or a URL string pointing to a CSV file.
        target_column: The name of the column to predict (target variable).
        nebius_api_key: Optional API key for Nebius AI Studio to enable AI-powered explanations.
        
    Returns:
        Tuple containing:
        - eda_report_path: Path to the generated HTML EDA report file.
        - task_type: Either "classification" or "regression" based on target variable.
        - models_dataframe: DataFrame with performance metrics of all trained models.
        - visualization_path: Path to the model comparison chart image.
        - model_pickle_path: Path to the serialized best model (.pkl file).
        - llm_explanation: AI-generated explanation of results (or fallback message).
        - column_names: Comma-separated list of detected column names.
    """
    # --- 1. Input Validation ---
    if not data_source or not target_column:
        error_msg = "Please provide both a data source and target column name."
        gr.Warning("Error: Data source and target column must be provided.")
        return None, NO_TASK_DETECTED, None, None, None, error_msg, NO_COLUMNS_LOADED

    gr.Info("Starting analysis...")

    # --- 2. Data Loading ---
    df, column_names = load_data(data_source)
    if df is None:
        error_msg = "Could not load data. Please check the file format or URL."
        return None, NO_TASK_DETECTED, None, None, None, error_msg, NO_COLUMNS_LOADED

    if target_column not in df.columns:
        error_msg = f"Target column '{target_column}' not found. Available columns: {column_names}"
        gr.Warning(error_msg)
        return None, NO_TASK_DETECTED, None, None, None, error_msg, column_names

    # --- 3. Analysis and Modeling ---
    _, profile_path, task, models_df, plot_path, pickle_path, best_model_name = analyze_and_model(df, target_column)

    # --- 4. Generate Dataset Summary for LLM Context ---
    dataset_summary = generate_dataset_summary(df, target_column)
    
    # Get top 5 model performance summary
    top_models_summary = models_df.head(5).to_string(index=False)

    # --- 5. Explanation with Nebius AI Studio LLM ---
    llm_explanation = "AI explanation is unavailable. Please provide a Nebius AI Studio API key to enable this feature."

    if nebius_api_key and nebius_api_key.strip():
        try:
            client = OpenAI(
                base_url="https://api.studio.nebius.com/v1/",
                api_key=nebius_api_key.strip()
            )

            # Craft an improved prompt with actual data context
            prompt_text = f"""Analyze this AutoML result and provide a concise, professional explanation:

**Dataset Overview:**
{dataset_summary}

**Task Type:** {task}

**Top 5 Performing Models:**
{top_models_summary}

**Best Model:** {best_model_name}

Please explain:
1. Why '{best_model_name}' performed best for this {task} task
2. Key insights about the dataset characteristics
3. Recommendations for model deployment or further improvement

Keep the explanation concise (3-4 paragraphs) and accessible to both technical and non-technical stakeholders."""

            response = client.chat.completions.create(
                model="meta-llama/Llama-3.3-70B-Instruct",
                messages=[
                    {"role": "system", "content": "You are an expert data scientist assistant that explains machine learning results clearly and professionally."},
                    {"role": "user", "content": prompt_text}
                ],
                temperature=0.6,
                max_tokens=512,
                top_p=0.9,
                extra_body={"top_k": 50}
            )
            # Simplified response access (no need for json.loads)
            llm_explanation = response.choices[0].message.content

        except Exception as e:
            gr.Warning(f"Failed to get AI explanation: {e}")
            llm_explanation = f"AI explanation unavailable due to an error. The best performing model is **{best_model_name}** for your {task} task."

    gr.Info("Analysis complete!")
    gr.Info(f'Profile report saved to: {profile_path}')
    return profile_path, task, models_df, plot_path, pickle_path, llm_explanation, column_names 

# --- Gradio UI ---
with gr.Blocks(title="AutoML Trainer", theme=gr.themes.Soft()) as demo:
    gr.Markdown("## 🤖 AutoML Trainer")

    with gr.Row():
        with gr.Column(scale=1):
            file_input = gr.File(label="Upload Local CSV File")
            url_input = gr.Textbox(label="Or Enter Public CSV URL", placeholder="e.g., https://.../data.csv")
            gr.Textbox(label="Sample CSV", value="https://raw.githubusercontent.com/daniel-was-taken/MCP_Project/refs/heads/master/collegePlace.csv")
            target_column_input = gr.Textbox(label="Enter Target Column Name", placeholder="e.g., approved")
            nebius_api_key_input = gr.Textbox(label="Nebius AI Studio API Key (Optional)", type="password", placeholder="Enter your API key for AI explanations")
            run_button = gr.Button("Run Analysis & AutoML", variant="primary")

        with gr.Column(scale=2):
            column_names_output = gr.Textbox(label="Detected Columns", interactive=False, lines=2) # New Textbox for column names
            task_output = gr.Textbox(label="Detected Task", interactive=False)
            llm_output = gr.Markdown(label="AI Explanation")
            metrics_output = gr.Dataframe(label="Model Performance Metrics")

    with gr.Row():
        vis_output = gr.Image(label="Top Models Comparison")
        with gr.Column():
            eda_output = gr.File(label="Download Full EDA Report")
            model_output = gr.File(label="Download Best Model (.pkl)")

    def process_inputs(
        file_data: Any, 
        url_data: Optional[str], 
        target: str, 
        api_key: Optional[str]
    ) -> Tuple[Optional[str], str, Optional[pd.DataFrame], Optional[str], Optional[str], str, str]:
        """
        Process inputs and run the AutoML pipeline.
        
        This wrapper function handles input selection between file upload and URL,
        then delegates to the main run_pipeline function.
        """
        data_source = file_data if file_data is not None else url_data
        return run_pipeline(data_source, target, api_key)

    file_input.change(
        fn=update_detected_columns_display,
        inputs=[file_input, url_input],
        outputs=column_names_output
    )
    url_input.change(
        fn=update_detected_columns_display,
        inputs=[file_input, url_input],
        outputs=column_names_output
    )

    run_button.click(
        fn=process_inputs,
        inputs=[file_input, url_input, target_column_input, nebius_api_key_input],
        outputs=[eda_output, task_output, metrics_output, vis_output, model_output, llm_output, column_names_output],
        api_name="run_automl_pipeline"  # Explicit API name for MCP
    )

demo.launch(
    server_name="0.0.0.0",
    server_port=7860,
    share=False,
    inbrowser=True,
    mcp_server=True
)