Update app.py

app.py

@@ -1,8 +1,9 @@
 # app.py
 from __future__ import annotations
 import os
+import io
 import traceback
-import regex as re2
+from contextlib import redirect_stdout
 from typing import List, Dict, Any
 
 import gradio as gr
@@ -11,8 +12,6 @@ from datetime import datetime
 
 # --- BACKEND IMPORTS ---
 from langchain_cohere import ChatCohere
-# --- THE FIXED IMPORT IS HERE ---
-from langchain_experimental.utilities.python import PythonREPL 
 
 # --- LOCAL MODULE IMPORTS ---
 from settings import (
@@ -41,9 +40,8 @@ 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."""
     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`. The first dataframe is `dfs[0]`, the second is `dfs[1]`, and so on.
+You have access to a list of pandas dataframes loaded into a variable named `dfs`.
 
-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 ---
@@ -53,16 +51,16 @@ 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 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.
+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.
 3.  **No Placeholders:** Do not use placeholder data.
-4.  **Self-Contained:** The script must be entirely self-contained.
+4.  **Self-Contained:** The script must be entirely self-contained, starting with `import pandas as pd`.
 
 --- USER'S SCENARIO ---
 {user_scenario}
 
 --- PYTHON SCRIPT ---
-Now, write the complete Python script to be executed. The script should start with `import pandas as pd` and contain all the logic.
+Now, write the complete Python script to be executed.
 ```python
 """
     generated_text = cohere_chat(prompt_for_coder)
@@ -81,7 +79,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}, timeout={COHERE_TIMEOUT_S}s)" if vecs else "Cohere reachable."
+        return f"Cohere OK ✅ (model={COHERE_MODEL_PRIMARY})" if vecs else "Cohere reachable."
     except Exception as e:
         return f"Cohere ping failed: {e}"
 
@@ -112,15 +110,24 @@ 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}
+            # --- THE FINAL, ROBUST FIX IS HERE: Using standard Python 'exec' ---
+            # Create a safe environment for the script to run in
+            execution_namespace = {
+                "dfs": dataframes,
+                "pd": pd
+            }
+            # Create a string buffer to capture the script's `print` statements
+            output_buffer = io.StringIO()
             
             try:
-                # Execute the AI-generated script
-                res = python_repl.run(command=analysis_script, locals=local_vars)
-                return _sanitize_text(res)
+                # Use a context manager to redirect stdout to our buffer
+                with redirect_stdout(output_buffer):
+                    # Execute the AI-generated script in the safe namespace
+                    exec(analysis_script, execution_namespace)
+                
+                # Get the captured output
+                result = output_buffer.getvalue()
+                return _sanitize_text(result or "(The analysis script ran but produced no output.)")
             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```"
@@ -197,7 +204,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(" - ")[0]
+        selected_id = selection.split(" - ")
         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'])