Update app.py

app.py

@@ -13,7 +13,7 @@ openai_api_key = os.getenv("OPENAI_API_KEY")
 
 # Initialize conversation history
 if 'conversation' not in st.session_state:
-    st.session_state.conversation = []
+    st.session_state.conversation = []  # Store previous conversation messages
 
 # Step 1: Upload CSV data file (or use default)
 csv_file = st.file_uploader("Upload your CSV file", type=["csv"])
@@ -33,30 +33,11 @@ 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 extract column names from the question
-def extract_column_name(question, valid_columns):
-    for column in valid_columns:
-        if column.lower() in question.lower():
-            return column
-    return None
-
-# Function to generate statistical insights
-def generate_statistical_insights(question, data):
-    if "mean" in question.lower():
-        column = extract_column_name(question, valid_columns)
-        if column:
-            mean_value = data[column].mean()
-            st.session_state.conversation.append(f"Mean of {column}: {mean_value}")
-        else:
-            st.session_state.conversation.append(f"Could not find a valid column in the question.")
-    elif "median" in question.lower():
-        column = extract_column_name(question, valid_columns)
-        if column:
-            median_value = data[column].median()
-            st.session_state.conversation.append(f"Median of {column}: {median_value}")
-    # Add more statistical insights (mode, std, etc.)
+
+# Display the conversation thread
+st.write("### Conversation Thread:")
+for message in st.session_state.conversation:
+    st.write(message)
 
 # Step 3: Define SQL validation helpers
 def validate_sql(query, valid_columns):
@@ -65,9 +46,8 @@ def validate_sql(query, valid_columns):
     columns_in_query = parser.columns
     for column in columns_in_query:
         if column not in valid_columns:
-            st.session_state.conversation.append(f"Invalid column detected: {column}")
-            return False
-    return True
+            return False, f"Invalid column detected: {column}"
+    return True, None
 
 def validate_sql_with_sqlparse(query):
     """Validates SQL syntax using sqlparse."""
@@ -89,36 +69,51 @@ SQL Query:
 prompt = PromptTemplate(template=template, input_variables=['question', 'table_name', 'columns'])
 sql_generation_chain = LLMChain(llm=OpenAI(temperature=0), prompt=prompt)
 
-# Display conversation history like a text thread
-st.write("### Conversation Thread")
-for chat in st.session_state.conversation:
-    st.write(f"User: {chat}")
+# Step 5: Generate SQL query based on user input
+user_prompt = st.text_input("Enter your message:")
 
-# Step 5: Generate SQL query or statistical insights based on user input
-user_prompt = st.text_input("Enter your question or prompt here:")
 if user_prompt:
+    # Add user prompt to conversation history
+    st.session_state.conversation.append(f"User: {user_prompt}")
+    
     try:
-        # Step 6: Handle statistical insights or generate SQL
-        if any(stat_term in user_prompt.lower() for stat_term in ["mean", "median", "mode", "std"]):
-            generate_statistical_insights(user_prompt, data)
+        # Step 6: Adjust the logic to handle "what are the columns" query
+        if "columns" in user_prompt.lower():
+            # Custom logic to return columns
+            columns_response = f"The columns are: {', '.join(valid_columns)}"
+            st.session_state.conversation.append(f"Bot: {columns_response}")
+            st.write(f"The columns are: {', '.join(valid_columns)}")
         else:
+            # Generate SQL query based on user input
             columns = ', '.join(valid_columns)
-            generated_sql = sql_generation_chain.run({'question': user_prompt, 'table_name': table_name, 'columns': columns})
+            generated_sql = sql_generation_chain.run({
+                'question': user_prompt,
+                'table_name': table_name,
+                'columns': columns
+            })
 
-            # Display generated SQL query in the conversation thread
-            st.session_state.conversation.append(f"Generated SQL Query: {generated_sql}")
+            # Debug: Display generated SQL query for inspection
+            st.write(f"Generated SQL Query:\n{generated_sql}")
 
             # Step 7: Validate SQL query
             if not validate_sql_with_sqlparse(generated_sql):
-                st.session_state.conversation.append("Generated SQL is not valid.")
-            elif not validate_sql(generated_sql, valid_columns):
-                st.session_state.conversation.append("Generated SQL references invalid columns.")
+                error_message = "Generated SQL is not valid."
+                st.session_state.conversation.append(f"Bot: {error_message}")
+                st.write(error_message)
+            elif not validate_sql(generated_sql, valid_columns)[0]:
+                invalid_column_message = "Generated SQL references invalid columns."
+                st.session_state.conversation.append(f"Bot: {invalid_column_message}")
+                st.write(invalid_column_message)
             else:
                 # Step 8: Execute SQL query
                 result = pd.read_sql_query(generated_sql, conn)
-                st.session_state.conversation.append("Query Results:")
-                st.session_state.conversation.append(result.to_string())
-
+                st.write("Query Results:")
+                st.dataframe(result)
+                st.session_state.conversation.append("Bot: Here are the results of your query.")
     except Exception as e:
         logging.error(f"An error occurred: {e}")
-        st.session_state.conversation.append(f"Error: {e}")
+        error_message = f"Error: {e}"
+        st.session_state.conversation.append(f"Bot: {error_message}")
+        st.write(f"Error: {e}")
+
+# Persist the conversation after each message