Update app.py

app.py

@@ -47,60 +47,48 @@ def main():
             st.info(f"Dataset contains {df.shape[0]} rows and {df.shape[1]} columns.")
 
         # --- 3. AGENT CONFIGURATION ---
-        query = st.text_area("What analysis would you like to perform?", placeholder="e.g., 'Plot the distribution of prices.'")
+        query = st.text_area("What analysis would you like to perform?", 
+                             placeholder="e.g., 'Plot the distribution of prices.'")
 
         if st.button("Run Agent") and query:
-            # Initialize LLM with max_retries at the model level
+            # Initialize LLM
             llm = ChatGoogleGenerativeAI(
-                model="gemini-2.5-flash-preview-09-2025", 
+                model="gemini-2.5-flash-preview-09-2025",
                 google_api_key=GEMINI_API_KEY,
                 temperature=0,
-                max_retries=6,
+                max_retries=6
             )
 
-            # Updated Prefix: We manually include the DF context here since we are using a custom suffix
-            # This satisfies the requirement to have DF info without triggering the ValueError
+            # Add dataframe context manually in prefix
             df_context = f"The dataframe 'df' has the following columns: {', '.join(df.columns)}"
             
             custom_prefix = f"""
             You are working with a pandas dataframe in Python. The name of the dataframe is `df`.
             {df_context}
-            
-            You should use the tools below to answer the question posed of you.
-            
-            STRICT FORMATTING RULES:
+
+            You should use the tools below to answer the question posed.
+
+            STRICT 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. Your 'Action Input' must be valid Python code and ONLY Python code. No markdown backticks.
-            4. Before creating a visualization for filtered data, always check if the filtered dataframe is empty.
-            5. For plots, use `plt.figure()` and end with `st.pyplot(plt.gcf())`.
-            6. The visual and code must be clear and beautiful.
-            7. Avoide over writing .
+            3. 'Action Input' must be valid Python code ONLY — no markdown.
+            4. Before plotting filtered data, ALWAYS check `df.empty`.
+            5. Use `plt.figure()` before every plot, and finish with `st.pyplot(plt.gcf())`.
+            6. Code must be clean and readable.
+            7. Avoid overwriting variables unless necessary.
             """
 
-            # The suffix reinforces the format right before the model generates text
-            custom_suffix = """
-            Begin!
-            Question: {input}
-            Thought: {agent_scratchpad}
-            """
-
-            # Create the Pandas Agent
-            # FIX: Ensure include_df_in_prompt is False when using a custom suffix
-            # This prevents the ValueError: "If suffix is specified, include_df_in_prompt should not be."
+            # IMPORTANT: No suffix → avoids ValueError
             agent = create_pandas_dataframe_agent(
-                 llm,
-                 df,
-                 verbose=True,
-                 agent_type="zero-shot-react",   # <-- FIXED
-                 allow_dangerous_code=True,
-                 prefix=custom_prefix,
-                 suffix=custom_suffix,
-                 include_df_in_prompt=False,     # <-- now allowed
-                 handle_parsing_errors=True,
-                 agent_executor_kwargs={
-                     "handle_parsing_errors": True
-                }
+                llm,
+                df,
+                verbose=True,
+                agent_type="zero-shot-react-description",  # required agent type
+                allow_dangerous_code=True,
+                prefix=custom_prefix,
+                include_df_in_prompt=False,
+                handle_parsing_errors=True,
+                agent_executor_kwargs={"handle_parsing_errors": True}
             )
 
             # --- 4. EXECUTION ---
@@ -126,14 +114,15 @@ def main():
                                 time.sleep(wait_time)
                                 continue
                             else:
-                                st.error("Rate limit exceeded consistently. Please try again in a minute.")
+                                st.error("Rate limit exceeded consistently. Try again later.")
                         else:
-                            st.error(f"Agent encountered a parsing or execution error.")
+                            st.error("Agent encountered a parsing or execution error.")
                             with st.expander("Show Technical Error"):
                                 st.code(str(e))
                             break
+
     else:
         st.info("👆 Upload a CSV to begin.")
 
 if __name__ == "__main__":
-    main()
+    main()