app.py

import gradio as gr
import os
import sys
import pandas as pd
import sqlite3
from pathlib import Path
import matplotlib.pyplot as plt
import re

# For Hugging Face Spaces, set project root to current directory
PROJECT_ROOT = Path(__file__).parent.resolve()
sys.path.append(str(PROJECT_ROOT))

# Import model loading and utility functions
from code.train_sqlgen_t5_local import load_model as load_sql_model, generate_sql, get_schema_from_csv
from code.train_intent_classifier_local import load_model as load_intent_model, classify_intent

# Load models
sql_model, sql_tokenizer, device = load_sql_model()
intent_model, intent_tokenizer, device, label_mapping = load_intent_model()

# Path to the built-in data file in the data folder
DATA_FILE = str(PROJECT_ROOT / "data" / "testing_sql_data.csv")

# Verify data file exists
if not os.path.exists(DATA_FILE):
    raise FileNotFoundError(f"Data file not found at {DATA_FILE}. Please ensure testing_sql_data.csv exists in the data folder.")

def process_query(question, chart_type="auto"):
    try:
        # Generate schema from CSV
        schema = get_schema_from_csv(DATA_FILE)
        # Generate SQL
        sql_query = generate_sql(question, schema, sql_model, sql_tokenizer, device)
        # --- Fix: Table and column name replacements ---
        sql_query = re.sub(r'(FROM|JOIN)\s+\w+', r'\1 data', sql_query, flags=re.IGNORECASE)
        sql_query = re.sub(r'(FROM|JOIN)\s+"[^"]+"', r'\1 data', sql_query, flags=re.IGNORECASE)
        sql_query = re.sub(r'(FROM|JOIN)\s+\'[^"]+\'', r'\1 data', sql_query, flags=re.IGNORECASE)
        sql_query = sql_query.replace('product_price', 'total_price')
        sql_query = sql_query.replace('store_name', 'store_id')
        sql_query = sql_query.replace('sales_method', 'date')
        sql_query = re.sub(r'\bsales\b', 'total_price', sql_query)
        # --- End fix ---
        # Classify intent
        intent = classify_intent(question, intent_model, intent_tokenizer, device, label_mapping)
        # Execute SQL on the CSV data
        df = pd.read_csv(DATA_FILE)
        conn = sqlite3.connect(":memory:")
        df.to_sql("data", conn, index=False, if_exists="replace")
        result_df = pd.read_sql_query(sql_query, conn)
        conn.close()
        # Defensive check for result_df columns
        if result_df.empty or len(result_df.columns) < 2:
            chart_path = None
            insights = "No results or not enough columns to display chart/insights."
            return result_df, intent, sql_query, chart_path, insights
        # Generate chart
        chart_path = os.path.join(PROJECT_ROOT, "chart.png")
        if not result_df.empty:
            plt.figure(figsize=(10, 6))
            if chart_type == "auto":
                if intent == "trend":
                    chart_type = "line"
                elif intent == "comparison":
                    chart_type = "bar"
                else:
                    chart_type = "bar"
            if chart_type == "bar":
                result_df.plot(kind="bar", x=result_df.columns[0], y=result_df.columns[1])
            elif chart_type == "line":
                result_df.plot(kind="line", x=result_df.columns[0], y=result_df.columns[1], marker='o')
            elif chart_type == "pie":
                result_df.plot(kind="pie", y=result_df.columns[1], labels=result_df[result_df.columns[0]])
            plt.title(question)
            plt.tight_layout()
            plt.savefig(chart_path)
            plt.close()
        else:
            chart_path = None
        # Generate insights
        insights = generate_insights(result_df, intent, question)
        return result_df, intent, sql_query, chart_path, insights
    except Exception as e:
        return None, "Error", str(e), None, f"Error: {str(e)}"

def generate_insights(result_df, intent, question):
    if result_df is None or result_df.empty or len(result_df.columns) < 2:
        return "No data available for insights."
    insights = []
    if intent == "summary":
        try:
            total = result_df[result_df.columns[1]].sum()
            insights.append(f"Total {result_df.columns[1]}: {total:,.2f}")
        except Exception:
            pass
    elif intent == "comparison":
        if len(result_df) >= 2:
            try:
                highest = result_df.iloc[0]
                lowest = result_df.iloc[-1]
                diff = ((highest.iloc[1] / lowest.iloc[1] - 1) * 100)
                insights.append(f"{highest.iloc[0]} is {diff:.1f}% higher than {lowest.iloc[0]}")
            except Exception:
                pass
    elif intent == "trend":
        if len(result_df) >= 2:
            try:
                first = result_df.iloc[0][result_df.columns[1]]
                last = result_df.iloc[-1][result_df.columns[1]]
                change = ((last / first - 1) * 100)
                insights.append(f"Overall change: {change:+.1f}%")
            except Exception:
                pass
    insights.append(f"Analysis covers {len(result_df)} records")
    if "category" in result_df.columns:
        insights.append(f"Number of categories: {result_df['category'].nunique()}")
    return "\n".join(f"• {insight}" for insight in insights)

# Clickable FAQs (6 only)
faqs = [
    "What are the top 5 products by quantity sold?",
    "What is the total sales amount for each category?",
    "Which store had the highest total sales?",
    "What are the most popular payment methods?",
    "What is the sales trend over time?",
    "What is the average transaction value?"
]

def fill_question(faq):
    return gr.update(value=faq)

with gr.Blocks(title="RetailGenie - Natural Language to SQL") as demo:
    gr.Markdown("""
    # RetailGenie - Natural Language to SQL
    Ask questions in natural language to generate SQL queries and visualizations. Using retail dataset with product sales information.
    """)
    with gr.Row():
        with gr.Column(scale=1):
            question = gr.Textbox(
                label="Enter your question",
                placeholder="What is the total sales amount for each product category?"
            )
            faq_radio = gr.Radio(faqs, label="FAQs (click to autofill)", interactive=True)
            faq_radio.change(fn=fill_question, inputs=faq_radio, outputs=question)
            chart_type = gr.Radio(
                ["auto", "bar", "line", "pie"],
                label="Chart Type",
                value="auto"
            )
            submit_btn = gr.Button("Generate", variant="primary")
        with gr.Column(scale=2):
            with gr.Accordion("SQL and Intent Details", open=False):
                intent_output = gr.Textbox(label="Predicted Intent")
                sql_output = gr.Textbox(label="Generated SQL", lines=3)
            results_df = gr.DataFrame(label="Query Results")
            chart_output = gr.Image(label="Chart")
            insights_output = gr.Textbox(label="Insights", lines=5)
    submit_btn.click(
        fn=process_query,
        inputs=[question, chart_type],
        outputs=[results_df, intent_output, sql_output, chart_output, insights_output]
    )

if __name__ == "__main__":
    demo.launch()