Update app.py

app.py

@@ -6,6 +6,7 @@ import openai
 import matplotlib.pyplot as plt
 from dotenv import load_dotenv
 from PIL import Image
+import traceback
 
 # Load your OpenAI API key from the environment (Hugging Face Spaces secrets will populate it)
 load_dotenv()
@@ -140,17 +141,12 @@ User Query: "{nl_query}"
     code = response.choices[0].message.content.strip()
     return code
 
+
 def safe_exec_code(code, df, capture_charts=True, interactive=False, extra_globals=None):
     """
     Execute the generated code in a restricted namespace.
-    
-    - Remove any lines starting with triple backticks.
-    - First try to evaluate the code as an expression using eval().
-    - If that fails, execute it via exec() and expect a variable 'result' or 'basic_info'.
-    - If capture_charts is True, override plt.show() to capture charts as PIL Images.
-      After executing the code, if interactive is True, open the images.
-    - The extra_globals parameter can be used to inject additional global variables.
-    - Returns a tuple (output, charts) where output is the value of 'result' or 'basic_info'.
+    Returns a tuple (output, charts) where output is the value of 'result' or 'basic_info'.
+    In case of an error, returns a detailed error message.
     """
     # Remove markdown code fences (lines starting with ```).
     code_lines = code.splitlines()
@@ -195,14 +191,16 @@ def safe_exec_code(code, df, capture_charts=True, interactive=False, extra_globa
     try:
         compiled_expr = compile(clean_code, "<string>", "eval")
         output = eval(compiled_expr, safe_globals, safe_locals)
-    except Exception:
+    except Exception as ex:
         try:
             exec(clean_code, safe_globals, safe_locals)
             output = safe_locals.get("result", None)
             if output is None:
                 output = safe_locals.get("basic_info", None)
         except Exception as ex:
-            return "An error occurred during code execution. Please check your code or try again.", safe_globals["charts"]
+            # Capture the full traceback and exception details
+            error_details = traceback.format_exc()
+            return f"An error occurred during code execution:\n{error_details}", safe_globals["charts"]
     
     if capture_charts and not safe_globals["charts"]:
         fig_nums = plt.get_fignums()
@@ -223,6 +221,7 @@ def safe_exec_code(code, df, capture_charts=True, interactive=False, extra_globa
         output = "No output variable ('result' or 'basic_info') was set by the code."
     return output, safe_globals["charts"]
 
+
 def generate_interpretation(analysis_result, nl_query):
     """
     Use OpenAI to generate a detailed interpretation of the analysis result.