app.py

import gradio as gr
import os
import subprocess
import logging
import asyncio 
from src.agents.nl_sql_agent import NLSQLAgent

logging.basicConfig(level=logging.INFO)
logging.getLogger("httpx").setLevel(logging.WARNING)
logging.getLogger("chromadb").setLevel(logging.WARNING)

print("--- Starting Hugging Face Space setup ---")

SETUP_DATABASE_SCRIPT = os.path.join('src', 'setup_database.py')
RAG_INDEX_SCRIPT = os.path.join('src', 'rag_index.py')

if not os.environ.get("NEBIUS_API_KEY"):
    print("FATAL ERROR: NEBIUS_API_KEY environment variable not set. This is required for Nebius LLMs and Embeddings.")

try:
    print(f"Running database setup script: {SETUP_DATABASE_SCRIPT}...")
    subprocess.run(["python", SETUP_DATABASE_SCRIPT], check=True, capture_output=True)
    print("Database setup complete.")

    print(f"Running RAG indexing script: {RAG_INDEX_SCRIPT}...")
    subprocess.run(["python", RAG_INDEX_SCRIPT], check=True, capture_output=True)
    print("RAG indexing complete.")

except subprocess.CalledProcessError as e:
    print(f"Error during initial setup. Script failed with exit code {e.returncode}.")
    print(f"Stdout:\n{e.stdout.decode()}")
    print(f"Stderr:\n{e.stderr.decode()}")
    print("Exiting application due to critical setup failure.")
    exit(1) 

print("--- Hugging Face Space setup complete. Initializing Agent ---")

# --- Initialize the NL-to-SQL Agent ---
nl_sql_agent_instance = NLSQLAgent() 
print("NLSQLAgent initialized.")

# --- Define Gradio Interface Functions ---
async def query_agent_gradio(user_query: str) -> str:
    if not user_query.strip():
        return "Please enter a question to get started!"

    try:
        response = await nl_sql_agent_instance.process_query(user_query)
        return response
    except Exception as e:
        logging.error(f"Error processing query in Gradio app: {e}", exc_info=True) 
        return f"An internal error occurred: {type(e).__name__}: {str(e)}. Please check the Space logs for more details."

# --- Create Gradio Interface ---
# --- Define the list of examples ---
example_list = [
    ["What is the total number of sales?"],
    ["What are the names of customers in the North region?"],
    ["How much revenue did we generate from Electronics products?"],
    ["Which customer has the most sales in the past one month?"],
    ["Show me the total revenue for each month over the last six months, ordered by month."],
    ["Which are our top 5 products by total quantity sold across all time?"],
    ["What is the average amount per sale for each product category?"],
    ["How many unique customers have made a purchase in each region over the last year?"],
    ["What is the total revenue generated by customers from each region?"]
]

# --- Create Gradio Interface ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # AI-Powered Data Analyst Assistant
        Ask natural language questions to get insights from your sales database. 
        This demo exclusively uses the NL-to-SQL Agent, leveraging fine-tuned LLMs, RAG for schema context, and SQLite.
        """
    )
    
    with gr.Row():
        user_query = gr.Textbox(
            lines=2, 
            placeholder="e.g., 'What is the total number of sales?'",
            label="Your Question"
        )

    with gr.Row():
        clear_btn = gr.ClearButton(value="Clear")
        submit_btn = gr.Button("Submit", variant="primary")

    with gr.Box(label="Answer"):
        output_markdown = gr.Markdown()

    gr.Examples(
        examples=example_list,
        inputs=user_query,
        label="Example Questions (Click to try one)"
    )

    # Define what happens when you click the submit button
    submit_btn.click(
        fn=query_agent_gradio,
        inputs=user_query,
        outputs=output_markdown,
        show_progress="minimal"
    )

    # Define what happens when you click the clear button
    clear_btn.add(components=[user_query, output_markdown])


if __name__ == "__main__":
    print("Launching Gradio app...")
    demo.launch()