Update app.py

app.py

@@ -6,11 +6,15 @@ from langchain import OpenAI, LLMChain, PromptTemplate
 from langchain_community.utilities import SQLDatabase
 import sqlparse
 import logging
-from sql_metadata import Parser  # Added import
+from sql_metadata import Parser
 
 # OpenAI API key (ensure it is securely stored)
 openai_api_key = os.getenv("OPENAI_API_KEY")
 
+# Initialize conversation history
+if 'conversation' not in st.session_state:
+    st.session_state.conversation = []
+
 # 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:
@@ -31,6 +35,29 @@ data.to_sql(table_name, conn, index=False, if_exists='replace')
 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.)
+
 # Step 3: Define SQL validation helpers
 def validate_sql(query, valid_columns):
     """Validates the SQL query by ensuring it references only valid columns."""
@@ -38,7 +65,7 @@ def validate_sql(query, valid_columns):
     columns_in_query = parser.columns
     for column in columns_in_query:
         if column not in valid_columns:
-            st.write(f"Invalid column detected: {column}")
+            st.session_state.conversation.append(f"Invalid column detected: {column}")
             return False
     return True
 
@@ -62,32 +89,40 @@ 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 natural language prompt:")
 if user_prompt:
     try:
-        # Step 6: Adjust the logic to handle "what are the columns" query
-        if "columns" in user_prompt.lower():
-            # Custom logic to return columns
-            st.write(f"The columns are: {', '.join(valid_columns)}")
+        # 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)
         else:
             columns = ', '.join(valid_columns)
             generated_sql = sql_generation_chain.run({'question': user_prompt, 'table_name': table_name, 'columns': columns})
 
-            # Debug: Display generated SQL query for inspection
-            st.write(f"Generated SQL Query:\n{generated_sql}")
+            # Display generated SQL query in the conversation thread
+            st.session_state.conversation.append(f"Generated SQL Query: {generated_sql}")
 
             # Step 7: Validate SQL query
             if not validate_sql_with_sqlparse(generated_sql):
-                st.write("Generated SQL is not valid.")
+                st.session_state.conversation.append("Generated SQL is not valid.")
             elif not validate_sql(generated_sql, valid_columns):
-                st.write("Generated SQL references invalid columns.")
+                st.session_state.conversation.append("Generated SQL references invalid columns.")
             else:
                 # Step 8: Execute SQL query
                 result = pd.read_sql_query(generated_sql, conn)
-                st.write("Query Results:")
-                st.dataframe(result)
+                st.session_state.conversation.append("Query Results:")
+                st.session_state.conversation.append(result.to_string())
 
     except Exception as e:
         logging.error(f"An error occurred: {e}")
-        st.write(f"Error: {e}")
+        st.session_state.conversation.append(f"Error: {e}")
+
+# Display the text input box below the conversation thread
+user_input = st.text_input("Enter a question to ask the data:", key="user_input")
+