Update app.py

app.py

@@ -41,11 +41,14 @@ st.write(f"Valid columns: {valid_columns}")
 
 # Step 3: Set up the LLM Chain to generate SQL queries
 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.
+You are an expert data scientist assistant. Given a natural language question, the name of the table, and a list of valid columns, generate a valid SQL query or Python code that answers the question.
 
-Ensure that:
-- You only use the columns provided.
-- String comparisons in the WHERE clause are case-insensitive by using 'COLLATE NOCASE' or the LOWER() function.
+Instructions:
+- If the question involves data retrieval or simple aggregations, generate a SQL query.
+- Ensure that you only use the columns provided.
+- For case-insensitive string comparisons, use either 'LOWER(column) = LOWER(value)' or 'column = value COLLATE NOCASE', but do not use both together.
+- Do not include any import statements in the code.
+- Provide the code between <CODE> and </CODE> tags.
 
 Question: {question}
 
@@ -53,8 +56,9 @@ Table name: {table_name}
 
 Valid columns: {columns}
 
-SQL Query:
+Response:
 """
+
 prompt = PromptTemplate(template=template, input_variables=['question', 'table_name', 'columns'])
 sql_generation_chain = LLMChain(llm=OpenAI(temperature=0), prompt=prompt)
 
@@ -86,6 +90,7 @@ def process_input():
                         # It's a SQL query
                         st.write(f"Generated SQL Query:\n{code}")
                         try:
+                            # Execute the SQL query
                             result = pd.read_sql_query(code, conn)
                             assistant_response = f"Generated SQL Query:\n{code}"
                             st.session_state.history.append({"role": "assistant", "content": assistant_response})
@@ -134,6 +139,7 @@ def process_input():
         # Reset the user_input in session state
         st.session_state['user_input'] = ''
 
+
 def extract_code(response):
     """Extracts code enclosed between <CODE> and </CODE> tags."""
     import re
@@ -144,7 +150,6 @@ def extract_code(response):
     else:
         return None
 
-# Display the conversation history
 # Display the conversation history
 for message in st.session_state.history:
     if message['role'] == 'user':
@@ -160,5 +165,6 @@ for message in st.session_state.history:
             st.markdown(f"**Assistant:** {content}")
 
 
+
 # Place the input field at the bottom with the callback
 st.text_input("Enter your message:", key='user_input', on_change=process_input)