Update app.py

app.py

@@ -10,8 +10,8 @@ import pandas as pd
 
 # New additions for data analysis agent
 from langchain.agents.agent_types import AgentType
+from langchain_cohere import ChatCohere
 from langchain_experimental.agents.agent_toolkits import create_pandas_dataframe_agent
-from langchain_cohere import ChatCohere # <-- NEW, CORRECT IMPORT
 
 # ---- Local modules
 from settings import (
@@ -47,11 +47,11 @@ You are an expert data analysis project manager. Your task is to read the user's
 
 From the user's text, extract the following:
 1.  **Primary Objective:** A one-sentence summary of the user's main goal.
-2.  **Key Tasks:** A numbered list of the specific questions the user wants answered.
-3.  **Expert Guidelines & Assumptions:** A bulleted list of EVERY specific number, metric, calculation method, or assumption mentioned in the text. This is critical for high-quality analysis.
-4.  **Required Output Format:** A description of how the user wants the final answer to be structured.
+2.  **Key Tasks:** A numbered list of ALL the specific questions the user wants answered.
+3.  **Expert Guidelines & Assumptions:** A bulleted list of any specific numbers, metrics, or calculation methods mentioned.
+4.  **Required Output Format:** A description of how the user wants the final answer structured.
 
-Present this as a clean brief. Then, include the user's original text at the end.
+CRITICAL INSTRUCTION: Tell the data analyst that it MUST answer ALL of the key tasks before providing its final answer.
 
 --- USER'S SCENARIO ---
 {user_scenario}
@@ -66,7 +66,7 @@ def is_healthcare_scenario(text: str, has_files: bool) -> bool:
     Dynamic detection: require uploaded files AND either structured scenario sections
     or healthcare keywords (configured in settings).
     """
-    t = (text or "").lower() # <-- INDENTATION IS NOW FIXED
+    t = (text or "").lower()
     kws = HEALTHCARE_SETTINGS["healthcare_keywords"]
     structured = any(s in t for s in ["background", "situation", "tasks", "deliverables"])
     return has_files and (structured or any(k in t for k in kws))
@@ -113,7 +113,22 @@ def handle(user_msg: str, history_messages: List[Dict[str, str]], files: list) -
 
                 llm = ChatCohere(model=COHERE_MODEL_PRIMARY, temperature=0)
                 enhanced_prompt = _create_enhanced_prompt(safe_in)
-                AGENT_PREFIX = """...""" # Prefix content remains the same
+                
+                AGENT_PREFIX = """
+You are a data analysis agent. You have access to one or more pandas dataframes.
+You MUST respond in one of two formats.
+
+FORMAT 1: To perform a task. Your response must be a single block of text with ONLY these three sections:
+Thought: Your step-by-step reasoning.
+Action: python_repl_ast
+Action Input: The Python code to run.
+
+FORMAT 2: To give the final answer. Your response must be a single block of text with ONLY these two sections:
+Thought: I have now answered all the user's questions and can provide the final report.
+Final Answer: The complete answer, structured as the user requested.
+
+CRITICAL RULE: NEVER combine `Action` and `Final Answer` in the same response. Choose one format.
+"""
 
                 agent = create_pandas_dataframe_agent(
                     llm,
@@ -121,10 +136,10 @@ def handle(user_msg: str, history_messages: List[Dict[str, str]], files: list) -
                     agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
                     verbose=True,
                     allow_dangerous_code=True,
-                    prefix=AGENT_PREFIX
+                    prefix=AGENT_PREFIX,
+                    max_iterations=50  # <-- THE FINAL FIX IS HERE
                 )
 
-                # Use the new .invoke() method
                 result = agent.invoke({"input": enhanced_prompt})
                 reply = _sanitize_text(result.get("output", "No output generated."))
 
@@ -188,4 +203,4 @@ if __name__ == "__main__":
         "open_fallbacks": USE_OPEN_FALLBACKS,
         "timeout_s": COHERE_TIMEOUT_S
     })
-    demo.launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", "7860")))
+    demo.launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", "7860")))