Update app.py

app.py

@@ -3,7 +3,7 @@ import streamlit as st
 import pandas as pd
 import sqlite3
 from langchain import OpenAI, LLMChain, PromptTemplate
-from langchain_community.utilities import SQLDatabase
+# Removed unused import: from langchain_community.utilities import SQLDatabase
 import sqlparse
 import logging
 from sql_metadata import Parser
@@ -12,47 +12,104 @@ from sql_metadata import Parser
 if 'history' not in st.session_state:
     st.session_state.history = []
 
-# Process user input
-user_prompt = st.text_input("Enter your message:", key='user_input')
+# OpenAI API key (ensure it is securely stored)
+openai_api_key = os.getenv("OPENAI_API_KEY")
 
-if user_prompt:
-    try:
-        if "columns" in user_prompt.lower():
-            assistant_response = f"The columns are: {', '.join(valid_columns)}"
+# Step 1: Upload CSV data file (or use default)
+csv_file = st.file_uploader("Upload your CSV file", type=["csv"])
+if csv_file is None:
+    data = pd.read_csv("default_data.csv")  # Use default CSV if no file is uploaded
+    st.write("Using default_data.csv file.")
+else:
+    data = pd.read_csv(csv_file)
+    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)
+table_name = csv_file.name.split('.')[0] if csv_file else "default_table"
+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}")
+
+# Step 3: Define SQL validation helpers
+def validate_sql(query, valid_columns):
+    """Validates the SQL query by ensuring it references only valid columns."""
+    parser = Parser(query)
+    columns_in_query = parser.columns
+    for column in columns_in_query:
+        if column not in valid_columns:
+            st.write(f"Invalid column detected: {column}")
+            return False
+    return True
+
+def validate_sql_with_sqlparse(query):
+    """Validates SQL syntax using sqlparse."""
+    parsed_query = sqlparse.parse(query)
+    return len(parsed_query) > 0
+
+# Step 4: Set up the LLM Chain to generate SQL queries
+template = """
+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.
+
+Question: {question}
+
+Table name: {table_name}
+
+Valid columns: {columns}
+
+SQL Query:
+"""
+prompt = PromptTemplate(template=template, input_variables=['question', 'table_name', 'columns'])
+sql_generation_chain = LLMChain(llm=OpenAI(temperature=0), prompt=prompt)
+
+# 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})
-            st.session_state.history.append({"role": "assistant", "content": assistant_response})
-        else:
-            columns = ', '.join(valid_columns)
-            generated_sql = sql_generation_chain.run({
-                'question': user_prompt,
-                'table_name': table_name,
-                'columns': columns
-            })
-
-            # Validate SQL query
-            if not validate_sql_with_sqlparse(generated_sql):
-                assistant_response = "Generated SQL is not valid."
-            elif not validate_sql(generated_sql, valid_columns):
-                assistant_response = "Generated SQL references invalid columns."
+
+            if "columns" in user_prompt.lower():
+                assistant_response = f"The columns are: {', '.join(valid_columns)}"
+                st.session_state.history.append({"role": "assistant", "content": assistant_response})
             else:
-                # Execute SQL query
-                result = pd.read_sql_query(generated_sql, conn)
-                assistant_response = f"Generated SQL Query:\n{generated_sql}\n\nQuery Results:"
-                st.session_state.history.append({"role": "assistant", "content": result})
+                columns = ', '.join(valid_columns)
+                generated_sql = sql_generation_chain.run({
+                    'question': user_prompt,
+                    'table_name': table_name,
+                    'columns': columns
+                })
 
-            # Append user and assistant messages to history
-            st.session_state.history.append({"role": "user", "content": user_prompt})
-            st.session_state.history.append({"role": "assistant", "content": assistant_response})
+                # Debug: Display generated SQL query for inspection
+                # st.write(f"Generated SQL Query:\n{generated_sql}")
+
+                # Validate SQL query
+                if not validate_sql_with_sqlparse(generated_sql):
+                    assistant_response = "Generated SQL is not valid."
+                    st.session_state.history.append({"role": "assistant", "content": assistant_response})
+                elif not validate_sql(generated_sql, valid_columns):
+                    assistant_response = "Generated SQL references invalid columns."
+                    st.session_state.history.append({"role": "assistant", "content": assistant_response})
+                else:
+                    # Execute SQL query
+                    result = pd.read_sql_query(generated_sql, conn)
+                    assistant_response = f"Generated SQL Query:\n{generated_sql}"
+                    st.session_state.history.append({"role": "assistant", "content": assistant_response})
+                    st.session_state.history.append({"role": "assistant", "content": result})
 
-    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})
+        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})
 
-    # Clear the input field
-    st.session_state.user_input = ''
-    # Rerun the script to update the conversation display
-    st.experimental_rerun()
+        # Reset the user_input in session state
+        st.session_state['user_input'] = ''
 
 # Display the conversation history
 for message in st.session_state.history:
@@ -64,3 +121,6 @@ for message in st.session_state.history:
             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)