Update app.py

app.py

@@ -58,17 +58,22 @@ def main():
                 max_retries=6, # LangChain's built-in retry logic
             )
 
+            # Updated Prefix to strictly separate Thought, Action, and Final Answer
             custom_prefix = """
             You are working with a pandas dataframe in Python. The name of the dataframe is `df`.
             You should use the tools below to answer the question posed of you.
             
-            IMPORTANT RULES:
-            1. Your 'Action Input' must be valid Python code and ONLY Python code. 
-            2. Do not wrap the code in markdown code blocks within the Action Input.
-            3. Before creating a visualization for filtered data, always check if the filtered dataframe is empty.
-            4. For plots, use `plt.figure()` and end with `st.pyplot(plt.gcf())` if you are using matplotlib/seaborn.
+            STRICT FORMATTING RULES:
+            1. NEVER provide a 'Final Answer' and an 'Action' in the same response.
+            2. If you need to run code, provide ONLY the 'Action' and 'Action Input'.
+            3. Once you have the result, provide ONLY the 'Final Answer'.
+            4. Your 'Action Input' must be valid Python code and ONLY Python code. No markdown backticks.
+            5. Before creating a visualization for filtered data, always check if the filtered dataframe is empty.
+            6. For plots, use `plt.figure()` and end with `st.pyplot(plt.gcf())`.
             """
 
+            # Create the Pandas Agent
+            # Using agent_executor_kwargs to pass handle_parsing_errors to the executor directly
             agent = create_pandas_dataframe_agent(
                 llm,
                 df,
@@ -77,6 +82,9 @@ def main():
                 allow_dangerous_code=True,
                 prefix=custom_prefix,
                 handle_parsing_errors=True,
+                agent_executor_kwargs={
+                    "handle_parsing_errors": True
+                }
             )
 
             # --- 4. EXECUTION ---
@@ -105,7 +113,8 @@ def main():
                             else:
                                 st.error("Rate limit exceeded consistently. Please try again in a minute.")
                         else:
-                            st.error(f"Agent failed to complete the task.")
+                            # Enhanced error display for parsing issues
+                            st.error(f"Agent encountered a parsing or execution error.")
                             with st.expander("Show Technical Error"):
                                 st.code(str(e))
                             break