app.py

import os
import streamlit as st
import pandas as pd
import sqlite3
from transformers import pipeline
import sqlparse
import logging

# Initialize conversation history
if 'history' not in st.session_state:
    st.session_state.history = []

# Load a pre-trained GPT-2 model from Hugging Face
llm = pipeline('text-generation', model='gpt2')

# Step 1: Upload CSV data file (or use default)
st.title("Natural Language to SQL Query App with Enhanced Insights")
st.write("Upload a CSV file to get started, or use the default dataset.")

csv_file = st.file_uploader("Upload your CSV file", type=["csv"])
if csv_file is None:
    data = pd.read_csv("default_data.csv")  # Ensure this file exists in your working directory
    st.write("Using default_data.csv file.")
    table_name = "default_table"
else:
    data = pd.read_csv(csv_file)
    table_name = csv_file.name.split('.')[0]
    st.write(f"Data Preview ({csv_file.name}):")
    st.dataframe(data.head())

# Step 2: Load CSV data into a persistent SQLite database
db_file = 'my_database.db'
conn = sqlite3.connect(db_file)
data.to_sql(table_name, conn, index=False, if_exists='replace')

# SQL table metadata (for validation and schema)
valid_columns = list(data.columns)
st.write(f"Valid columns: {valid_columns}")

# Function to generate SQL query using Hugging Face model
def generate_sql_query(question, table_name, columns):
    prompt = f"""
    You are an expert data scientist. Given a natural language question, the name of the table, and a list of valid columns, generate a valid SQL query that answers the question.
    Ensure that:
    - You only use the columns provided.
    - When performing string comparisons in the WHERE clause, make them case-insensitive by using 'COLLATE NOCASE' or the LOWER() function.
    - Do not use 'COLLATE NOCASE' in ORDER BY clauses unless sorting a string column.
    - Do not apply 'COLLATE NOCASE' to numeric columns.
    If the question is vague or open-ended and does not pertain to specific data retrieval, respond with "NO_SQL" to indicate that a SQL query should not be generated.
    Question: {question}
    Table name: {table_name}
    Valid columns: {columns}
    SQL Query:
    """
    response = llm(prompt, max_length=180)
    return response[0]['generated_text'].strip()

# Function to generate insights using Hugging Face model
def generate_insights(question, result):
    prompt = f"""
    You are an expert data scientist. Based on the user's question and the SQL query result provided below, generate a concise analysis that includes key data insights and actionable recommendations. Limit the response to a maximum of 150 words.
    User's Question: {question}
    SQL Query Result:
    {result}
    Concise Analysis (max 200 words):
    """
    response = llm(prompt, max_length=150)
    return response[0]['generated_text'].strip()

# Function to classify user query as SQL or Insights
def classify_query(question):
    prompt = f"""
    You are an AI assistant that classifies user queries into two categories: 'SQL' for specific data retrieval queries and 'INSIGHTS' for general analytical or recommendation queries.
    Determine the appropriate category for the following user question.
    Question: "{question}"
    Category (SQL/INSIGHTS):
    """
    response = llm(prompt, max_length=10)
    category = response[0]['generated_text'].strip().upper()
    return 'SQL' if 'SQL' in category else 'INSIGHTS'

# Function to generate dataset summary
def generate_dataset_summary(data):
    summary_template = f"""
    You are an expert data scientist. Based on the dataset provided below, generate a concise summary that includes the number of records, number of columns, data types, and any notable features.
    Dataset:
    {data.head().to_string(index=False)}
    Dataset Summary:
    """
    response = llm(summary_template, max_length=150)
    return response[0]['generated_text'].strip()

# Optional: Clean up function to remove incorrect COLLATE NOCASE usage
def clean_sql_query(query):
    """Removes incorrect usage of COLLATE NOCASE from the SQL query."""
    parsed = sqlparse.parse(query)
    statements = []
    for stmt in parsed:
        tokens = []
        idx = 0
        while idx < len(stmt.tokens):
            token = stmt.tokens[idx]
            if (token.ttype is sqlparse.tokens.Keyword and token.value.upper() == 'COLLATE'):
                # Check if the next token is 'NOCASE'
                next_token = stmt.tokens[idx + 2] if idx + 2 < len(stmt.tokens) else None
                if next_token and next_token.value.upper() == 'NOCASE':
                    # Skip 'COLLATE' and 'NOCASE' tokens
                    idx += 3  # Skip 'COLLATE', whitespace, 'NOCASE'
                    continue
            tokens.append(token)
            idx += 1
        statements.append(''.join([str(t) for t in tokens]))
    return ' '.join(statements)

# Define the callback function
def process_input():
    user_prompt = st.session_state['user_input']

    if user_prompt:
        try:
            # Append user message to history
            st.session_state.history.append({"role": "user", "content": user_prompt})

            # Classify the user query
            category = classify_query(user_prompt)
            logging.info(f"User query classified as: {category}")

            if "COLUMNS" in user_prompt.upper():
                assistant_response = f"The columns are: {', '.join(valid_columns)}"
                st.session_state.history.append({"role": "assistant", "content": assistant_response})
            elif category == 'SQL':
                columns = ', '.join(valid_columns)
                generated_sql = generate_sql_query(user_prompt, table_name, columns)

                if generated_sql.upper() == "NO_SQL":
                    # Handle cases where no SQL should be generated
                    assistant_response = "Sure, let's discuss some general insights and recommendations based on the data."
                    
                    # Generate dataset summary
                    dataset_summary = generate_dataset_summary(data)
                    
                    # Generate general insights and recommendations
                    general_insights = generate_insights(user_prompt, dataset_summary)
                    
                    # Append the assistant's insights to the history
                    st.session_state.history.append({"role": "assistant", "content": general_insights})
                else:
                    # Clean the SQL query
                    cleaned_sql = clean_sql_query(generated_sql)
                    logging.info(f"Generated SQL Query: {cleaned_sql}")

                    # Attempt to execute SQL query and handle exceptions
                    try:
                        result = pd.read_sql_query(cleaned_sql, conn)

                        if result.empty:
                            assistant_response = "The query returned no results. Please try a different question."
                            st.session_state.history.append({"role": "assistant", "content": assistant_response})
                        else:
                            # Convert the result to a string for the insights prompt
                            result_str = result.head(10).to_string(index=False)  # Limit to first 10 rows

                            # Generate insights and recommendations based on the query result
                            insights = generate_insights(user_prompt, result_str)

                            # Append the assistant's insights to the history
                            st.session_state.history.append({"role": "assistant", "content": insights})
                            # Append the result DataFrame to the history
                            st.session_state.history.append({"role": "assistant", "content": result})
                    except Exception as e:
                        logging.error(f"An error occurred during SQL execution: {e}")
                        assistant_response = f"Error executing SQL query: {e}"
                        st.session_state.history.append({"role": "assistant", "content": assistant_response})
            else:  # INSIGHTS category
                # Generate dataset summary
                dataset_summary = generate_dataset_summary(data)

                # Generate general insights and recommendations
                general_insights = generate_insights(user_prompt, dataset_summary)

                # Append the assistant's insights to the history
                st.session_state.history.append({"role": "assistant", "content": general_insights})
        
        except Exception as e:
            logging.error(f"An error occurred: {e}")
            assistant_response = f"Error: {e}"
            st.session_state.history.append({"role": "assistant", "content": assistant_response})

        # Reset the user_input in session state
        st.session_state['user_input'] = ''

# Display the conversation history
for message in st.session_state.history:
    if message['role'] == 'user':
        st.markdown(f"**User:** {message['content']}")
    elif message['role'] == 'assistant':
        if isinstance(message['content'], pd.DataFrame):
            st.markdown("**Assistant:** Query Results:")
            st.dataframe(message['content'])
        else:
            st.markdown(f"**Assistant:** {message['content']}")

# Place the input field at the bottom with the callback
st.text_input("Enter your message:", key='user_input', on_change=process_input)