Add lines filter to remove unwanted code from Gemini's output

app.py

@@ -25,46 +25,57 @@ def answer_question(history, message):
     history: chat history (list of [user, assistant] pairs)
     message: latest user message (string)
     """
-    # Build prompt for Gemini
     prompt = f"""
-    You are a data analysis assistant. 
-    You can ONLY answer questions using the two Excel reports provided (df1 and df2). 
-    Do not hallucinate or use external knowledge.
-    If the question is irrelevant, respond with:
-    "I can only answer questions about the provided Excel reports."
-
-    The reports have the following schema:
-    {schema_info}
-
-    The user asked:
-    {message}
-
-    Rules:
-    - Use only pandas, df1, df2, and Python built-ins.
-    - Do NOT import anything (numpy, pandas, etc. are already available).
-    - Always put the result in a variable named `result`.
-
-    Generate Python code that computes the answer.
-    Return ONLY code, nothing else.
+You are a data analysis assistant. 
+You can ONLY answer questions using the two Excel reports provided (df1 and df2). 
+Do not hallucinate or use external knowledge.
+If the question is irrelevant, respond with:
+"I can only answer questions about the provided Excel reports."
+
+The reports have the following schema:
+{schema_info}
+
+The user asked:
+{message}
+
+Rules:
+- Use only pandas, df1, df2, and Python built-ins.
+- Do NOT write import statements (pandas is already imported as pd).
+- Always put the answer in a variable named `result`.
+- Return ONLY Python code, nothing else.
     """
 
     try:
-        # Call Gemini
+        # Ask Gemini
         response = model.generate_content(prompt)
-        code = response.text.strip("```python").strip("```")
-
-        # 4. Execute code safely
+        code = response.text.strip()
+
+        # Strip markdown fences if any
+        if code.startswith("```"):
+            code = code.strip("```").replace("python", "").strip()
+
+        # Remove dangerous or irrelevant lines
+        safe_lines = []
+        for line in code.splitlines():
+            if line.strip().startswith("import"):
+                continue
+            if line.strip().startswith("!"):  # e.g. shell commands
+                continue
+            safe_lines.append(line)
+        safe_code = "\n".join(safe_lines)
+
+        # Prepare sandbox
         local_vars = {"df1": df1, "df2": df2, "pd": pd}
-        try:
-            result = eval(code, {"__builtins__": {}}, local_vars)
-        except Exception as e:
-            exec(code, {"__builtins__": {}}, local_vars)
-            result = local_vars.get("result", "No result variable found")
 
+        # Execute code
+        exec(safe_code, {"__builtins__": {}}, local_vars)
+
+        # Fetch result
+        result = local_vars.get("result", "No result produced")
         return str(result)
 
     except Exception as e:
-        return f"Error: {str(e)}"
+        return f"Execution error: {e}"
 
 
 # 5. Gradio UI