Update app.py

app.py

@@ -20,7 +20,7 @@ if not openai_api_key:
     st.stop()
 
 # Step 1: Upload CSV data file (or use default)
-st.title("Natural Language to SQL Query App with Dynamic Insights")
+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"])
@@ -71,27 +71,29 @@ sql_generation_chain = LLMChain(llm=llm, prompt=sql_prompt)
 
 # Insights Generation Chain
 insights_template = """
-You are an expert data scientist. Based on the SQL query result provided below, generate a concise and informative analysis that includes specific data-driven insights.
+You are an expert data scientist. Based on the user's question and the SQL query result provided below, generate a concise and informative analysis that includes data insights and actionable recommendations.
+
+User's Question: {question}
 
 SQL Query Result:
 {result}
 
-Analysis:
+Analysis and Recommendations:
 """
-insights_prompt = PromptTemplate(template=insights_template, input_variables=['result'])
+insights_prompt = PromptTemplate(template=insights_template, input_variables=['question', 'result'])
 insights_chain = LLMChain(llm=llm, prompt=insights_prompt)
 
-# Recommendations Generation Chain
-recommendations_template = """
-You are an expert data scientist. Based on the SQL query result provided below, generate actionable recommendations for improving performance.
+# General Insights and Recommendations Chain
+general_insights_template = """
+You are an expert data scientist. Based on the entire dataset provided below, generate a comprehensive analysis that includes key insights and actionable recommendations.
 
-SQL Query Result:
-{result}
+Dataset Summary:
+{dataset_summary}
 
-Recommendations:
+Analysis and Recommendations:
 """
-recommendations_prompt = PromptTemplate(template=recommendations_template, input_variables=['result'])
-recommendations_chain = LLMChain(llm=llm, prompt=recommendations_prompt)
+general_insights_prompt = PromptTemplate(template=general_insights_template, input_variables=['dataset_summary'])
+general_insights_chain = LLMChain(llm=llm, prompt=general_insights_prompt)
 
 # Optional: Clean up function to remove incorrect COLLATE NOCASE usage
 def clean_sql_query(query):
@@ -119,7 +121,7 @@ def clean_sql_query(query):
 def classify_query(question):
     """Classify the user query as either 'SQL' or 'INSIGHTS'."""
     classification_template = """
-    You are an AI assistant that classifies user queries into two categories: 'SQL' for specific data retrieval queries and 'INSIGHTS' for general analytical queries.
+    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.
 
@@ -135,6 +137,22 @@ def classify_query(question):
     else:
         return 'INSIGHTS'
 
+# Function to generate dataset summary
+def generate_dataset_summary(data):
+    """Generate a summary of the dataset for general insights."""
+    summary_template = """
+    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}
+
+    Dataset Summary:
+    """
+    summary_prompt = PromptTemplate(template=summary_template, input_variables=['data'])
+    summary_chain = LLMChain(llm=llm, prompt=summary_prompt)
+    summary = summary_chain.run({'data': data.head().to_string(index=False)})
+    return summary
+
 # Define the callback function
 def process_input():
     user_prompt = st.session_state['user_input']
@@ -160,8 +178,19 @@ def process_input():
                 }).strip()
 
                 if generated_sql.upper() == "NO_SQL":
-                    assistant_response = "This query is too vague for generating SQL. Please ask a more specific question."
-                    st.session_state.history.append({"role": "assistant", "content": assistant_response})
+                    # 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 = general_insights_chain.run({
+                        'dataset_summary': 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)
@@ -178,14 +207,14 @@ def process_input():
                             # 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 based on the query result
+                            # Generate insights and recommendations based on the query result
                             insights = insights_chain.run({
+                                'question': user_prompt,
                                 'result': result_str
                             })
 
-                            # Display insights in a scrollable text area
-                            st.text_area("Insights", value=insights, height=300)
-
+                            # 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:
@@ -193,25 +222,16 @@ def process_input():
                         assistant_response = f"Error executing SQL query: {e}"
                         st.session_state.history.append({"role": "assistant", "content": assistant_response})
             else:  # INSIGHTS category
-                if "recommendations" in user_prompt.lower():
-                    # Generate recommendations based on the query result
-                    dataset_summary = data.describe().to_string()  # Summary for recommendations
-                    recommendations = recommendations_chain.run({
-                        'result': dataset_summary
-                    })
+                # Generate dataset summary
+                dataset_summary = generate_dataset_summary(data)
 
-                    # Display recommendations in a scrollable text area
-                    st.text_area("Recommendations", value=recommendations, height=300)
-                    
-                else:
-                    # Generate insights based on general insights (without recommendations)
-                    dataset_summary = data.describe().to_string()  # Summary for insights
-                    insights = insights_chain.run({
-                        'result': dataset_summary
-                    })
+                # Generate general insights and recommendations
+                general_insights = general_insights_chain.run({
+                    'dataset_summary': dataset_summary
+                })
 
-                    # Display insights in a scrollable text area
-                    st.text_area("Insights", value=insights, height=300)
+                # 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}")