Update app.py

app.py

@@ -37,18 +37,25 @@ def _sanitize_text(s: str) -> str:
     if not isinstance(s, str): return s
     return re2.sub(r'[\p{C}--[\n\t]]+', '', s)
 
-def _create_enhanced_prompt(user_scenario: str) -> str:
-    """Uses an LLM to pre-process the user's prompt into a structured brief."""
+# 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"""
-You are an expert data analysis project manager. Your task is to read the user's unstructured scenario below and create a clear, structured brief for a data analysis AI.
-From the user's text, extract:
-1. Primary Objective: A one-sentence summary of the user's main goal.
-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.
-CRITICAL INSTRUCTION: Tell the data analyst that it MUST answer ALL of the key tasks before providing its final answer.
+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.
+
+--- 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.
+
 --- 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.
 """
     structured_brief = cohere_chat(prompt_for_planner)
     return structured_brief if structured_brief else user_scenario
@@ -78,35 +85,41 @@ def handle(user_msg: str, files: list) -> str:
 
         if file_paths:
             dataframes = []
+            file_names = []
             for p in file_paths:
                 if p.endswith('.csv'):
                     try:
                         df = pd.read_csv(p)
                         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.")
                         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."
 
+            # 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)
+            enhanced_prompt = _create_enhanced_prompt(safe_in, file_context_string)
             
             AGENT_PREFIX = """
-You are a data analysis agent. You have access to one or more pandas dataframes.
+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.
 
-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.
+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.
+Action Input: The Python code to run on the dataframes (df1, df2, etc.).
 
-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.
+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.
 
-CRITICAL RULE: NEVER combine `Action` and `Final Answer` in the same response. Choose one format.
+CRITICAL RULE: NEVER use hypothetical data. ALWAYS use the provided dataframes to generate your results.
 """
 
             agent = create_pandas_dataframe_agent(
@@ -131,7 +144,8 @@ 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) ...
+    # ... (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():
@@ -162,7 +176,11 @@ with gr.Blocks(theme="soft", css="style.css") as demo:
         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=[])
@@ -175,37 +193,28 @@ with gr.Blocks(theme="soft", css="style.css") as demo:
         terms_link = gr.Button("Terms of Service", variant="link")
 
     # --- UI LOGIC ---
-    
-    # THIS IS THE NEW, RESPONSIVE "RUN" FUNCTION
     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.")
-            # We must yield the original state to prevent an error on empty run
             yield chat_history_list, history_state_list, gr.update()
-            return # This stops the generator
+            return
 
-        # 1. Immediately show the user's message and a "Thinking..." status
         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() # This provides immediate feedback
+        yield thinking_message, history_state_list, gr.update()
 
-        # 2. Call the powerful (and slow) backend engine
         ai_response_text = handle(prompt, files)
 
-        # 3. Replace "Thinking..." with the final AI response
         final_chat = _append_msg(chat_with_user_msg, "assistant", ai_response_text)
         
-        # 4. Save the completed assessment to our history state
         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]
         
-        # 5. Create user-friendly labels for the history dropdown
         history_labels = [f"{item['id']} - {item['prompt'][:40]}..." for item in updated_history]
         
-        # 6. Yield the final, complete state to the UI
         yield final_chat, updated_history, gr.update(choices=history_labels)
 
     def view_history(selection, history_state_list):
@@ -237,6 +246,7 @@ with gr.Blocks(theme="soft", css="style.css") as demo:
     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.")