Update app.py

app.py

@@ -37,24 +37,30 @@ def _sanitize_text(s: str) -> str:
     if not isinstance(s, str): return s
     return re2.sub(r'[\p{C}--[\n\t]]+', '', s)
 
+# --- THE FINAL FIX (PART 1): The "Senior Analyst" AI ---
 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."""
+    """
+    Uses an LLM to act as a "Senior Analyst", breaking the complex user
+    scenario into a clear, step-by-step plan for the agent.
+    """
     prompt_for_planner = f"""
-You are an expert data analysis project manager. Your task is to create a clear, structured brief for a data analysis AI based on the user's scenario and the provided data context.
+You are a Senior Data Analyst. Your job is to create a clear, step-by-step execution plan for a Junior AI Data Analyst.
+The user has provided a complex scenario and a list of data files. The Junior Analyst gets confused by long prompts and can get stuck in loops.
+
+Your plan must be simple, clear, and sequential.
 
 --- DATA CONTEXT ---
 {file_context}
-The dataframes are available in a list, indexed as df1, df2, and so on, in the order they are listed above. Your primary task is to use these dataframes to answer the user's questions. Do not use hypothetical data.
+The Junior Analyst has access to these files in a list of pandas dataframes (df1, df2, etc.), in the order listed above.
 
 --- USER'S SCENARIO ---
 {user_scenario}
 
 --- YOUR TASK ---
-Based on BOTH the user's scenario and the data context, extract the following:
-1. Primary Objective: A one-sentence summary of the user's main goal.
-2. Key Tasks: A numbered list of ALL specific questions the user wants answered using the provided data.
-3. Required Output Format: A description of how the user wants the final answer structured.
-CRITICAL INSTRUCTION: Tell the data analyst that it MUST answer ALL of the key tasks before providing its final answer.
+Create a "Step-by-Step Execution Plan" for the Junior Analyst. Tell it exactly what to do, one task at a time, referencing the correct dataframe (df1, df2, etc.).
+Instruct it to perform all data preparation first, then the analysis, then the recommendations.
+Tell it that it MUST complete ALL steps in the plan before providing the final report.
+This plan will be given to the Junior Analyst. Make it easy to follow.
 """
     structured_brief = cohere_chat(prompt_for_planner)
     return structured_brief if structured_brief else user_scenario
@@ -68,7 +74,7 @@ def ping_cohere() -> str:
         cli = _co_client()
         if not cli: return "Cohere client not initialized. Is COHERE_API_KEY set?"
         vecs = cohere_embed(["hello", "world"])
-        return f"Cohere OK ✅ (model={COHERE_MODEL_PRIMARY}, timeout={COHERE_TIMEOUT_S}s)" if vecs else "Cohere reachable, but embeddings returned no vectors."
+        return f"Cohere OK ✅ (model={COHERE_MODEL_PRIMARY}, timeout={COHERE_TIMEOUT_S}s)" if vecs else "Cohere reachable."
     except Exception as e:
         return f"Cohere ping failed: {e}"
 
@@ -92,41 +98,41 @@ def handle(user_msg: str, files: list) -> str:
                         dataframes.append(df)
                         file_names.append(os.path.basename(p))
                     except UnicodeDecodeError:
-                        print(f"Warning: Failed to read {os.path.basename(p)} with UTF-8. Falling back to latin1 encoding.")
+                        print(f"Warning: Reading {os.path.basename(p)} with fallback latin1 encoding.")
                         df = pd.read_csv(p, encoding='latin1')
                         dataframes.append(df)
                         file_names.append(os.path.basename(p))
             
             if not dataframes: return "Please upload at least one CSV file."
 
-            file_context_string = "The user has provided the following data files for your analysis: " + ", ".join(file_names)
+            file_context_string = "The user has provided the following data files: " + ", ".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 ---
+            # --- THE FINAL FIX (PART 2): Stricter Agent with Error Handling Rule ---
             AGENT_PREFIX = """
-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.
-
-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.
-
-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.
-
-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".
+You are a Junior AI Data Analyst. Your job is to execute the step-by-step plan provided by your Senior Analyst using Python and pandas.
+You have access to dataframes named df1, df2, etc.
+
+You MUST follow these rules:
+
+1.  **EXECUTE THE PLAN:** Follow the execution plan exactly, one step at a time.
+2.  **FORMATTING:** Your response MUST be in one of two formats. NEVER mix them.
+    *   **To run code:**
+        Thought: Your reasoning for the code you are about to run to complete the current step.
+        Action: python_repl_ast
+        Action Input: The single line of python code to run.
+    *   **To give the final answer:**
+        Thought: I have finished all steps in the plan and can now provide the final report.
+        Final Answer: The complete, final answer, formatted as a concise report.
+
+3.  **ERROR HANDLING:** If your code produces an error, DO NOT try the same code again. Analyze the error message and try a DIFFERENT approach to solve the step. If you are stuck, say so.
 """
 
             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,
-                # handle_parsing_errors is now less critical but a good safety net
-                handle_parsing_errors=True 
+                handle_parsing_errors=True
             )
             result = agent.invoke({"input": enhanced_prompt})
             return _sanitize_text(result.get("output", "No output generated."))
@@ -146,7 +152,7 @@ TERMS_OF_SERVICE_TEXT = load_markdown_text("terms_of_service.md")
 # ---------------- THE PROFESSIONAL UI WITH INTEGRATED LEGAL DOCS ----------------
 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) ...
+    # ... (The rest of the UI code is identical to the last working version) ...
     with gr.Group(visible=False) as privacy_modal:
         with gr.Blocks():
             gr.Markdown(PRIVACY_POLICY_TEXT)
@@ -188,7 +194,7 @@ with gr.Blocks(theme="soft", css="style.css") as demo:
             return
 
         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```")
+        thinking_message = _append_msg(chat_with_user_msg, "assistant", "```\n🧠 Formulating execution plan... Please wait.\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)