Update app.py

app.py

@@ -37,7 +37,6 @@ def _sanitize_text(s: str) -> str:
     if not isinstance(s, str): return s
     return re2.sub(r'[\p{C}--[\n\t]]+', '', s)
 
-# THIS FUNCTION IS NOW UPGRADED
 def _create_enhanced_prompt(user_scenario: str, file_context: str) -> str:
     """Uses an LLM to pre-process the user's prompt and adds critical data context."""
     prompt_for_planner = f"""
@@ -100,31 +99,34 @@ def handle(user_msg: str, files: list) -> str:
             
             if not dataframes: return "Please upload at least one CSV file."
 
-            # Create the crucial file context string
             file_context_string = "The user has provided the following data files for your analysis: " + ", ".join(file_names)
-
             llm = ChatCohere(model=COHERE_MODEL_PRIMARY, temperature=0)
             enhanced_prompt = _create_enhanced_prompt(safe_in, file_context_string)
             
+            # --- THE FINAL, STRICTEST AGENT PREFIX ---
             AGENT_PREFIX = """
-You are a data analysis agent. You have access to one or more pandas dataframes. Your task is to use the provided dataframes to answer the user's questions.
-You MUST respond in one of two formats.
+Your job is to act as a data analyst. You have access to pandas dataframes (df1, df2, etc.).
+You MUST follow these rules. This is not a suggestion.
 
-FORMAT 1: To perform a task.
-Thought: Your step-by-step reasoning for using the data.
-Action: python_repl_ast
-Action Input: The Python code to run on the dataframes (df1, df2, etc.).
+1.  Your response MUST be in one of two formats. NEVER mix them.
+2.  To run code, use this exact format:
+    Thought: Your reasoning for the code you are about to run.
+    Action: python_repl_ast
+    Action Input: The single line of python code to run.
 
-FORMAT 2: To give the final answer.
-Thought: I have now completed all the tasks and can provide the final report based on the real data.
-Final Answer: The complete answer, structured as the user requested.
+3.  To give the final answer, use this exact format:
+    Thought: I have finished all the work and have the final answer.
+    Final Answer: The complete, final answer to the user's question.
 
-CRITICAL RULE: NEVER use hypothetical data. ALWAYS use the provided dataframes to generate your results.
+NEVER, EVER, provide a "Final Answer" and an "Action" in the same response. This is a fatal error.
+Begin now. Analyze the user's request and provide your first "Thought" and "Action".
 """
 
             agent = create_pandas_dataframe_agent(
                 llm, dataframes, agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
-                verbose=True, allow_dangerous_code=True, prefix=AGENT_PREFIX, max_iterations=50
+                verbose=True, allow_dangerous_code=True, prefix=AGENT_PREFIX, max_iterations=50,
+                # handle_parsing_errors is now less critical but a good safety net
+                handle_parsing_errors=True 
             )
             result = agent.invoke({"input": enhanced_prompt})
             return _sanitize_text(result.get("output", "No output generated."))
@@ -145,8 +147,6 @@ TERMS_OF_SERVICE_TEXT = load_markdown_text("terms_of_service.md")
 with gr.Blocks(theme="soft", css="style.css") as demo:
     assessment_history = gr.State([])
     # ... (The rest of the UI code is identical to the last version) ...
-    # ... (For brevity, I will omit it, but you should use the full UI code from the previous step)
-    # --- MODALS (POPUPS) DEFINED FIRST, INITIALLY HIDDEN ---
     with gr.Group(visible=False) as privacy_modal:
         with gr.Blocks():
             gr.Markdown(PRIVACY_POLICY_TEXT)
@@ -157,11 +157,8 @@ with gr.Blocks(theme="soft", css="style.css") as demo:
             gr.Markdown(TERMS_OF_SERVICE_TEXT)
             close_terms_btn = gr.Button("Close")
 
-    # --- MAIN UI LAYOUT ---
     gr.Markdown("# Universal AI Data Analyst")
-
     with gr.Row(variant="panel"):
-        # --- LEFT COLUMN: CONTROLS ---
         with gr.Column(scale=1):
             gr.Markdown("## New Assessment")
             files_input = gr.Files(label="Upload Data Files (.csv)", file_count="multiple", type="filepath", file_types=[".csv"])
@@ -171,28 +168,19 @@ with gr.Blocks(theme="soft", css="style.css") as demo:
                 clear_btn = gr.Button("🗑️ Clear")
             ping_btn = gr.Button("Ping Cohere")
             ping_out = gr.Markdown()
-
-        # --- RIGHT COLUMN: RESULTS & HISTORY ---
         with gr.Column(scale=2):
             with gr.Tabs():
                 with gr.TabItem("Current Assessment", id=0):
-                    chat_history_output = gr.Chatbot(
-                        label="Analysis Output",
-                        type="messages",
-                        height=600
-                    )
+                    chat_history_output = gr.Chatbot(label="Analysis Output", type="messages", height=600)
                 with gr.TabItem("Assessment History", id=1):
                     gr.Markdown("## Review Past Assessments")
                     history_dropdown = gr.Dropdown(label="Select an assessment to review", choices=[])
                     history_display = gr.Markdown(label="Selected Assessment Details")
-
-    # --- FOOTER FOR LEGAL LINKS ---
     with gr.Row(): gr.Markdown("---")
     with gr.Row():
         privacy_link = gr.Button("Privacy Policy", variant="link")
         terms_link = gr.Button("Terms of Service", variant="link")
 
-    # --- UI LOGIC ---
     def run_analysis_wrapper(prompt, files, chat_history_list, history_state_list):
         if not prompt or not files:
             gr.Warning("Please provide both a prompt and at least one data file.")
@@ -202,19 +190,13 @@ with gr.Blocks(theme="soft", css="style.css") as demo:
         chat_with_user_msg = _append_msg(chat_history_list, "user", prompt)
         thinking_message = _append_msg(chat_with_user_msg, "assistant", "```\n🧠 Analyzing... Please wait. This may take a minute.\n```")
         yield thinking_message, history_state_list, gr.update()
-
         ai_response_text = handle(prompt, files)
-
         final_chat = _append_msg(chat_with_user_msg, "assistant", ai_response_text)
-        
         timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
         file_names = [os.path.basename(f.name if hasattr(f, 'name') else f) for f in files]
-        
         new_assessment = {"id": timestamp, "prompt": prompt, "files": file_names, "response": ai_response_text}
         updated_history = history_state_list + [new_assessment]
-        
         history_labels = [f"{item['id']} - {item['prompt'][:40]}..." for item in updated_history]
-        
         yield final_chat, updated_history, gr.update(choices=history_labels)
 
     def view_history(selection, history_state_list):
@@ -226,7 +208,6 @@ with gr.Blocks(theme="soft", css="style.css") as demo:
             return f"""### Assessment from: {selected_assessment['id']}\n**Files Used:**\n- {file_list_md}\n---\n**Original Prompt:**\n> {selected_assessment['prompt']}\n---\n**AI Generated Response:**\n{selected_assessment['response']}"""
         return "Could not find the selected assessment."
 
-    # Wire up the components
     send_btn.click(
         run_analysis_wrapper,
         inputs=[prompt_input, files_input, chat_history_output, assessment_history],
@@ -239,14 +220,11 @@ with gr.Blocks(theme="soft", css="style.css") as demo:
     )
     clear_btn.click(lambda: (None, None, [], []), outputs=[prompt_input, files_input, chat_history_output, assessment_history])
     ping_btn.click(ping_cohere, outputs=[ping_out])
-
-    # Wire up the modal popups
     privacy_link.click(lambda: gr.update(visible=True), outputs=[privacy_modal])
     close_privacy_btn.click(lambda: gr.update(visible=False), outputs=[privacy_modal])
     terms_link.click(lambda: gr.update(visible=True), outputs=[terms_modal])
     close_terms_btn.click(lambda: gr.update(visible=False), outputs=[terms_modal])
 
-
 if __name__ == "__main__":
     if not os.getenv("COHERE_API_KEY"):
         print("🔴 COHERE_API_KEY environment variable not set. Application may not function correctly.")