Update app.py

app.py

@@ -93,17 +93,18 @@ few_shot_prompt = FewShotPromptTemplate(
     examples=few_shot_examples,
     prefix=system_prefix,
     suffix="Question: {input}\nSQL Query:",
-    input_variables=["input"]
+    input_variables=["input", "agent_scratchpad"]
 )
 
 # Step 5: Initialize the LLM and toolkit
 llm = ChatOpenAI(temperature=0, openai_api_key=openai_api_key)
 toolkit = SQLDatabaseToolkit(db=engine, llm=llm)
 
-# Step 6: Create the agent
+# Step 6: Create the agent using 'zero-shot-react-description' agent type
 agent_prompt = ChatPromptTemplate.from_messages([
     SystemMessagePromptTemplate(prompt=few_shot_prompt),
-    HumanMessagePromptTemplate.from_template("{input}")
+    HumanMessagePromptTemplate.from_template("{input}"),
+    MessagesPlaceholder(variable_name="agent_scratchpad")
 ])
 
 sql_agent = create_sql_agent(
@@ -111,7 +112,7 @@ sql_agent = create_sql_agent(
     toolkit=toolkit,
     prompt=agent_prompt,
     verbose=True,
-    agent_type="openai-functions",
+    agent_type="zero-shot-react-description",
     max_iterations=5
 )
 
@@ -128,65 +129,46 @@ def process_input():
             with st.spinner("Generating SQL query..."):
                 response = sql_agent.run(user_prompt)
             
-            # Check if the response contains SQL code
-            if "SELECT" in response.upper():
-                sql_query = response.strip()
-                logging.info(f"Generated SQL Query: {sql_query}")
-
-                # Attempt to execute SQL query and handle exceptions
-                try:
-                    result = pd.read_sql_query(sql_query, conn)
-
-                    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:
-                        # Limit the result to first 10 rows for display
-                        result_display = result.head(10)
-                        st.session_state.history.append({"role": "assistant", "content": "Here are the results:"})
-                        st.session_state.history.append({"role": "assistant", "content": result_display})
-
-                        # Generate insights based on the query result
-                        insights_template = """
-                        You are an expert data analyst. Based on the user's question and the SQL query result provided below, generate a concise analysis that includes key data insights and actionable recommendations. Limit the response to a maximum of 150 words.
-
-                        User's Question: {question}
-
-                        SQL Query Result:
-                        {result}
-
-                        Concise Analysis:
-                        """
-                        insights_prompt = PromptTemplate(template=insights_template, input_variables=['question', 'result'])
-                        insights_chain = LLMChain(llm=llm, prompt=insights_prompt)
-
-                        result_str = result_display.to_string(index=False)
-                        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})
-                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:
-                # Handle responses that do not contain SQL queries
-                assistant_response = response
-                st.session_state.history.append({"role": "assistant", "content": assistant_response})
+            # Extract the SQL query from the agent's response
+            sql_query = response.strip()
+            logging.info(f"Generated SQL Query: {sql_query}")
 
-            # Evaluate the response for harmful content
+            # Attempt to execute SQL query and handle exceptions
             try:
-                evaluator = load_evaluator("harmful_content", llm=llm)
-                eval_result = evaluator.evaluate_strings(
-                    input=user_prompt,
-                    prediction=response
-                )
-                if eval_result['flagged']:
-                    st.warning("The assistant's response may not be appropriate.")
+                result = pd.read_sql_query(sql_query, conn)
+
+                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:
-                    logging.info("Response evaluated as appropriate.")
+                    # Limit the result to first 10 rows for display
+                    result_display = result.head(10)
+                    st.session_state.history.append({"role": "assistant", "content": "Here are the results:"})
+                    st.session_state.history.append({"role": "assistant", "content": result_display})
+
+                    # Generate insights based on the query result
+                    insights_template = """
+                    You are an expert data analyst. Based on the user's question and the SQL query result provided below, generate a concise analysis that includes key data insights and actionable recommendations. Limit the response to a maximum of 150 words.
+
+                    User's Question: {question}
+
+                    SQL Query Result:
+                    {result}
+
+                    Concise Analysis:
+                    """
+                    insights_prompt = PromptTemplate(template=insights_template, input_variables=['question', 'result'])
+                    insights_chain = LLMChain(llm=llm, prompt=insights_prompt)
+
+                    result_str = result_display.to_string(index=False)
+                    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})
             except Exception as e:
-                logging.error(f"An error occurred during evaluation: {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})
         except Exception as e:
             logging.error(f"An error occurred: {e}")
             assistant_response = f"Error: {e}"