Update app.py

app.py

@@ -11,7 +11,8 @@ from datetime import datetime
 
 # --- BACKEND IMPORTS ---
 from langchain_cohere import ChatCohere
-from langchain_community.utilities.python import PythonREPL # Re-introducing the standard, robust executor
+# --- THE FIXED IMPORT IS HERE ---
+from langchain_experimental.utilities.python import PythonREPL 
 
 # --- LOCAL MODULE IMPORTS ---
 from settings import (
@@ -38,11 +39,11 @@ def _sanitize_text(s: str) -> str:
 
 def _create_python_script(user_scenario: str, schema_context: str) -> str:
     """Uses an LLM to act as an "AI Coder", writing a complete Python script."""
-    # --- THE FINAL PROMPT FIX IS HERE ---
     prompt_for_coder = f"""
 You are an expert Python data scientist. Your sole job is to write a single, complete, and executable Python script to answer the user's request.
-You have access to a list of pandas dataframes loaded into a variable named `dfs`.
+You have access to a list of pandas dataframes loaded into a variable named `dfs`. The first dataframe is `dfs[0]`, the second is `dfs[1]`, and so on.
 
+CRITICAL CONTEXT: Before writing any code, you MUST first understand the data you have been given. Here is the schema for each dataframe:
 --- DATA SCHEMA ---
 {schema_context}
 --- END SCHEMA ---
@@ -52,8 +53,8 @@ CRITICAL RULE: You MUST use the exact column names provided in the DATA SCHEMA.
 Based on the user's scenario below, write a single Python script that performs the entire analysis.
 
 RULES FOR YOUR SCRIPT:
-1.  **Use the DataFrames:** Your script MUST use the `dfs` list and the exact column names from the schema.
-2.  **Print Your Findings:** Use the `print()` function at each step to output the results as a formatted report.
+1.  **Use the DataFrames:** Your script MUST use the `dfs` list to access the data.
+2.  **Print Your Findings:** Use the `print()` function at each step of your analysis to output the results as a formatted report.
 3.  **No Placeholders:** Do not use placeholder data.
 4.  **Self-Contained:** The script must be entirely self-contained.
 
@@ -80,7 +81,7 @@ def ping_cohere() -> str:
         cli = _co_client()
         if not cli: return "Cohere client not initialized."
         vecs = cohere_embed(["hello", "world"])
-        return f"Cohere OK ✅ (model={COHERE_MODEL_PRIMARY})" if vecs else "Cohere reachable."
+        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}"
 
@@ -111,12 +112,17 @@ def handle(user_msg: str, files: list) -> str:
             schema_context = "\n".join(schema_parts)
             analysis_script = _create_python_script(safe_in, schema_context)
 
+            # Initialize the Python Executor
             python_repl = PythonREPL()
+            # Pass the dataframes into the execution environment
             local_vars = {"dfs": dataframes}
+            
             try:
+                # Execute the AI-generated script
                 res = python_repl.run(command=analysis_script, locals=local_vars)
                 return _sanitize_text(res)
             except Exception as e:
+                # If execution fails, return the error and the script for debugging
                 return f"An error occurred executing the script: {e}\n\nGenerated Script:\n```python\n{analysis_script}\n```"
         else:
             prompt = f"{GENERAL_CONVERSATION_PROMPT}\n\nUser: {safe_in}\nAssistant:"
@@ -176,7 +182,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🧠 Generating analysis script... This may take a moment.\n```")
+        thinking_message = _append_msg(chat_with_user_msg, "assistant", "```\n🧠 Generating and executing analysis script... This may take a moment.\n```")
         yield thinking_message, history_state_list, gr.update()
         
         ai_response_text = handle(prompt, files)
@@ -191,7 +197,7 @@ with gr.Blocks(theme="soft", css="style.css") as demo:
 
     def view_history(selection, history_state_list):
         if not selection or not history_state_list: return ""
-        selected_id = selection.split(" - ")
+        selected_id = selection.split(" - ")[0]
         selected_assessment = next((item for item in history_state_list if item["id"] == selected_id), None)
         if selected_assessment:
             file_list_md = "\n- ".join(selected_assessment['files'])