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 Data 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"])
@@ -55,6 +55,8 @@ Ensure that:
 - 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}
@@ -67,7 +69,7 @@ sql_prompt = PromptTemplate(template=sql_template, input_variables=['question',
 llm = OpenAI(temperature=0, openai_api_key=openai_api_key)
 sql_generation_chain = LLMChain(llm=llm, prompt=sql_prompt)
 
-# AnswerScript for generating insights based on query results
+# Insights Generation Chain
 insights_template = """
 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.
 
@@ -81,6 +83,18 @@ Analysis and Recommendations:
 insights_prompt = PromptTemplate(template=insights_template, input_variables=['question', 'result'])
 insights_chain = LLMChain(llm=llm, prompt=insights_prompt)
 
+# 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.
+
+Dataset Summary:
+{dataset_summary}
+
+Analysis and Recommendations:
+"""
+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):
     """Removes incorrect usage of COLLATE NOCASE from the SQL query."""
@@ -103,6 +117,42 @@ def clean_sql_query(query):
         statements.append(''.join([str(t) for t in tokens]))
     return ' '.join(statements)
 
+# Function to classify user 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 or recommendation queries.
+
+    Determine the appropriate category for the following user question.
+
+    Question: "{question}"
+
+    Category (SQL/INSIGHTS):
+    """
+    classification_prompt = PromptTemplate(template=classification_template, input_variables=['question'])
+    classification_chain = LLMChain(llm=llm, prompt=classification_prompt)
+    category = classification_chain.run({'question': question}).strip().upper()
+    if category.startswith('SQL'):
+        return 'SQL'
+    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']
@@ -112,46 +162,77 @@ def process_input():
             # Append user message to history
             st.session_state.history.append({"role": "user", "content": user_prompt})
 
-            if "columns" in user_prompt.lower():
+            # 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})
-            else:
+            elif category == 'SQL':
                 columns = ', '.join(valid_columns)
                 generated_sql = sql_generation_chain.run({
                     'question': user_prompt,
                     'table_name': table_name,
                     'columns': columns
-                })
+                }).strip()
 
-                # Clean the SQL query
-                generated_sql = clean_sql_query(generated_sql)
+                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 = 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)
+                    logging.info(f"Generated SQL Query: {cleaned_sql}")
 
-                # Attempt to execute SQL query and handle exceptions
-                try:
-                    result = pd.read_sql_query(generated_sql, conn)
+                    # 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."
+                        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 = insights_chain.run({
+                                'question': user_prompt,
+                                'result': 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:
-                        # 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
-                        insights = insights_chain.run({
-                            'question': user_prompt,
-                            'result': 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 = 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})
+        
         except Exception as e:
             logging.error(f"An error occurred: {e}")
             assistant_response = f"Error: {e}"