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")
+st.title("Natural Language to SQL Query App with Data 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"])
@@ -43,7 +43,8 @@ data.to_sql(table_name, conn, index=False, if_exists='replace')
 valid_columns = list(data.columns)
 st.write(f"Valid columns: {valid_columns}")
 
-# Step 3: Set up the LLM Chain to generate SQL queries
+# Step 3: Set up the LLM Chains
+# SQL Generation Chain
 sql_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.
 
@@ -66,6 +67,20 @@ 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_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.
+
+User's Question: {question}
+
+SQL Query Result:
+{result}
+
+Analysis and Recommendations:
+"""
+insights_prompt = PromptTemplate(template=insights_template, input_variables=['question', 'result'])
+insights_chain = LLMChain(llm=llm, prompt=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."""
@@ -115,28 +130,23 @@ def process_input():
                 try:
                     result = pd.read_sql_query(generated_sql, conn)
 
-                    # Limit the result to first 5 rows for brevity
-                    result_limited = result.head(5)
-                    result_str = result_limited.to_string(index=False)
-
-                    # Generate natural language answer
-                    answer_template = """
-                    Given the user's question and the SQL query result, provide a concise and informative answer to the question using the data from the query result.
-
-                    User's question: {question}
-                    Query result:
-                    {result}
-
-                    Answer:
-                    """
-                    answer_prompt = PromptTemplate(template=answer_template, input_variables=['question', 'result'])
-                    answer_chain = LLMChain(llm=llm, prompt=answer_prompt)
-                    assistant_answer = answer_chain.run({'question': user_prompt, 'result': result_str})
-
-                    # Append the assistant's answer to the history
-                    st.session_state.history.append({"role": "assistant", "content": assistant_answer})
-                    # Append the result DataFrame to the history
-                    st.session_state.history.append({"role": "assistant", "content": result})
+                    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
+                        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}"