Redone the entire code.

app.py

@@ -1,102 +1,276 @@
 import gradio as gr
-from google import genai
-from ase import Atoms
+import os
+from dotenv import load_dotenv
+import google.generativeai as genai
+import ase
+from ase.build import molecule # Often used, good to have for exec scope if needed
+from ase.calculators.emt import EMT # Common default
+from ase.optimize import BFGS # Common optimizer
+from ase.visualize import plot_atoms
+import matplotlib
+matplotlib.use('Agg') # Use Agg backend for non-interactive plotting
+import matplotlib.pyplot as plt
+from PIL import Image
+import io
+import base64
+import traceback
+import contextlib
+
+# Load environment variables from .env file (for GEMINI_API_KEY)
+load_dotenv()
+
+# --- Gemini API Configuration ---
+API_KEY = "AIzaSyCXggYtx5BtkGuXZMaeuhZ9-7xbuf0s7hs"
+MODEL = None
+
+if not API_KEY:
+    print("WARNING: GEMINI_API_KEY environment variable not set. App will not function.")
+else:
+    try:
+        genai.configure(api_key=API_KEY)
+        MODEL = genai.GenerativeModel('gemini-2.0-flash') # Or specific model like 'gemini-1.5-flash'
+        print("Gemini model initialized successfully.")
+    except Exception as e:
+        print(f"ERROR: Could not initialize Gemini model: {e}")
+        MODEL = None
+
+# --- Prompts for Gemini ---
+PROMPT_FOR_ASE_CODE_GENERATION = """
+You are an expert in molecular simulations using the Atomic Simulation Environment (ASE) and Python.
+A user wants to perform the following task: "{user_query}"
+
+Generate a complete and runnable Python script using ASE to perform this task.
+The script MUST:
+1.  Include all necessary import statements (e.g., `import ase`, `from ase.build import molecule`, `from ase.calculators.emt import EMT`, `from ase.optimize import BFGS`, etc.). Do not assume any modules are pre-imported.
+2.  Define the atomic structure.
+3.  Set up an appropriate ASE calculator. If the user doesn't specify, you can use `ase.calculators.emt.EMT()` for simplicity, but acknowledge if a more advanced calculator would be better for real research.
+4.  Perform the simulation as requested (e.g., geometry optimization, molecular dynamics, single point energy calculation).
+5.  Ensure the final ASE `Atoms` object, representing the result of the simulation (e.g., optimized structure), is assigned to a variable named `atoms_object`. This is crucial for visualization.
+6.  Optionally, print any relevant information or results from the simulation to standard output (e.g., final energy, bond lengths, status messages).
+
+Example of optimizing a water molecule:
+```python
+import ase
+from ase.build import molecule
 from ase.calculators.emt import EMT
 from ase.optimize import BFGS
-import matplotlib.pyplot as plt
-import tempfile
-import os
 
-# Gemini API initialization
-client = genai.Client(api_key=("AIzaSyCXggYtx5BtkGuXZMaeuhZ9-7xbuf0s7hs"))
+# Define the atomic structure
+atoms_object = molecule('H2O')
 
-def gemini_prompt(prompt, system="You are a helpful assistant for scientific simulations."):
-    # model = genai.GenerativeModel(model_name="gemini-pro", system_instruction=system)
-    # response = model.generate_content(prompt)
+# Set up a calculator
+atoms_object.calc = EMT()
 
-    response = client.models.generate_content(
-        model="gemini-2.0-flash", contents=prompt
-    )
+# Perform the simulation
+print("Starting geometry optimization for H2O...")
+dyn = BFGS(atoms_object)
+dyn.run(fmax=0.05) # Converge when max force is below 0.05 eV/Angstrom
 
-    return response.text
-
-def parse_params_from_text(text):
-    # Use Gemini to extract ASE simulation parameters from user text
-    prompt = f"""Given the following user request, extract the simulation parameters for an ASE simulation as a Python code snippet. The code should define an 'atoms' object using ase.Atoms. Only return the Python code, do not include any explanations or extra text. Make sure all necessary imports are included in the code snippet. If the user asks for a specific calculator, use it; otherwise, use EMT().
-Request: {text}"""
-    code = gemini_prompt(prompt)
-    return code, prompt
-
-def run_simulation_from_code(code_str):
-    # Dangerous: eval! In production, use a safer execution method.
-    local_vars = {}
-    exec(code_str, globals(), local_vars)
-    atoms = local_vars.get('atoms')
-    if atoms is None:
-        raise ValueError("No 'atoms' object found in the code.")
-    dyn = BFGS(atoms, logfile=None)
-    dyn.run(fmax=0.02)
-    return atoms
-
-def analyze_results_with_llm(atoms, user_query):
-    prompt = f"""Given the following ASE Atoms object representing the simulation results, and the original user request, provide a concise analysis of the simulation. Focus on key properties like bond lengths, angles, and overall structure changes if a geometry optimization was performed. If the user asked for specific information, make sure to include it. 
-Atoms: {atoms}
-User request: {user_query}
-Respond with a clear and concise summary."""
-    summary = gemini_prompt(prompt)
-    return summary
-
-def plot_atoms(atoms):
-    # Save image to temp file
-    from ase.visualize.plot import plot_atoms
-    fig, ax = plt.subplots()
-    plot_atoms(atoms, ax=ax)
-    tmpfile = tempfile.NamedTemporaryFile(suffix='.png', delete=False)
-    plt.savefig(tmpfile.name)
-    plt.close(fig)
-    return tmpfile.name
-
-def mcp_tool(user_text):
+# Print results
+print("Optimization finished.")
+print(f"Final energy: {atoms_object.get_potential_energy()} eV")
+# The 'atoms_object' variable now holds the optimized structure.
+```
+
+Provide ONLY the Python code block. Do not include any explanations, comments, or markdown formatting before or after the code block itself.
+The script will be executed in an environment where ASE and its dependencies are installed.
+"""
+
+PROMPT_FOR_ANALYSIS = """
+You are an AI assistant specialized in interpreting results from molecular simulations performed with the Atomic Simulation Environment (ASE).
+
+The user's original simulation request was:
+"{user_query}"
+
+The following Python ASE code was generated and executed to address this request:
+```python
+{generated_code}
+```
+
+The execution of this code produced the following standard output (stdout/stderr):
+```
+{simulation_output}
+```
+
+{structure_info}
+
+Based on all the above information, please provide a concise analysis:
+1.  Briefly explain what the simulation aimed to do and what steps were performed by the code.
+2.  Summarize the key results obtained from the simulation output (e.g., final energy, convergence status, important structural parameters if mentioned in output).
+3.  If a molecular structure was visualized, comment on it in relation to the simulation (e.g., "The visualized structure shows the molecule after optimization").
+4.  Conclude whether the simulation likely achieved the user's original request.
+5.  If there were any errors or issues, please mention them.
+
+Present the analysis in a clear, easy-to-understand manner. Avoid overly technical jargon where possible.
+"""
+
+# --- Helper Functions ---
+def plot_atoms_to_pil_image(atoms_obj):
+    """Converts an ASE Atoms object to a PIL Image for display."""
+    if not isinstance(atoms_obj, ase.Atoms):
+        return None
     try:
-        prompt = f"""Given the following user request, extract the simulation parameters for an ASE simulation as a Python code snippet. The code should define an 'atoms' object using ase.Atoms. Only return the Python code, do not include any explanations or extra text. Make sure all necessary imports are included in the code snippet. If the user asks for a specific calculator, use it; otherwise, use EMT().
-Request: {user_text}"""
-        code_str = gemini_prompt(prompt)
-        atoms = run_simulation_from_code(code_str)
-        img_path = plot_atoms(atoms)
-        summary = analyze_results_with_llm(atoms, user_text)
-        return img_path, summary, prompt, code_str
+        fig, ax = plt.subplots()
+        # Use ase.visualize.plot_atoms for plotting
+        plot_atoms(atoms_obj, ax, radii=0.3, rotation=('0x,0y,0z'))
+        ax.set_xlabel("X (Å)")
+        ax.set_ylabel("Y (Å)")
+        ax.axis('equal')
+        
+        buf = io.BytesIO()
+        fig.savefig(buf, format='png', bbox_inches='tight')
+        plt.close(fig)  # Close the figure to free up memory
+        buf.seek(0)
+        return Image.open(buf)
     except Exception as e:
-        return None, f"Error: {e}", "", ""
+        print(f"Error plotting atoms: {traceback.format_exc()}")
+        return None
+
+# --- Core Simulation Logic ---
+def simulate_and_analyze(user_query_str):
+    if not MODEL:
+        error_msg = "Gemini model not initialized. Please check API key and configuration."
+        return "# Error: Model not initialized.", error_msg, None, error_msg
+
+    generated_code_str = ""
+    simulation_stdout_str = ""
+    pil_image = None
+    analysis_text_str = ""
+
+    try:
+        # 1. Generate ASE code from user query
+        prompt_for_code = PROMPT_FOR_ASE_CODE_GENERATION.format(user_query=user_query_str)
+        response_code = MODEL.generate_content(prompt_for_code)
+        generated_code_str = response_code.text.strip()
+        # Remove potential markdown backticks if Gemini adds them
+        if generated_code_str.startswith("```python"):
+            generated_code_str = generated_code_str[len("```python"):].strip()
+        if generated_code_str.endswith("```"):
+            generated_code_str = generated_code_str[:-len("```")].strip()
+
+    except Exception as e:
+        error_msg = f"Error generating ASE code with Gemini: {traceback.format_exc()}"
+        return f"# Error: Could not generate ASE code.\n{error_msg}", error_msg, None, error_msg
+
+    # 2. Execute generated ASE code
+    atoms_object_from_exec = None
+    execution_namespace = {}
+    captured_output = io.StringIO()
+    exec_error_msg = None
+
+    try:
+        with contextlib.redirect_stdout(captured_output):
+            # Provide a minimal, controlled global scope for exec
+            # ASE and its modules should be imported by the generated code itself
+            exec(generated_code_str, {'__builtins__': __builtins__, 'ase': ase}, execution_namespace)
+        
+        simulation_stdout_str = captured_output.getvalue()
+        atoms_object_from_exec = execution_namespace.get('atoms_object')
+
+        if atoms_object_from_exec is not None and not isinstance(atoms_object_from_exec, ase.Atoms):
+            warning_msg = f"Warning: 'atoms_object' was found but is not a valid ASE Atoms object (type: {type(atoms_object_from_exec)}). Visualization might fail."
+            simulation_stdout_str += "\n" + warning_msg
+            atoms_object_from_exec = None # Invalidate if not correct type
+        elif 'atoms_object' not in execution_namespace:
+            simulation_stdout_str += "\nWarning: 'atoms_object' variable not found after code execution. Visualization will not be available."
+
+    except Exception as e:
+        exec_error_msg = f"Error during ASE code execution:\n{traceback.format_exc()}"
+        simulation_stdout_str = captured_output.getvalue() # Get any output before error
+        if simulation_stdout_str:
+            simulation_stdout_str += f"\n\n{exec_error_msg}"
+        else:
+            simulation_stdout_str = exec_error_msg
+        atoms_object_from_exec = None # Ensure it's None on error
+
+    # 3. Plot atoms object if available
+    if atoms_object_from_exec and isinstance(atoms_object_from_exec, ase.Atoms):
+        pil_image = plot_atoms_to_pil_image(atoms_object_from_exec)
+        structure_info_for_analysis = "A molecular structure was generated and visualized."
+    else:
+        structure_info_for_analysis = "No molecular structure was available or successfully generated for visualization."
+        if exec_error_msg and not atoms_object_from_exec : # If exec failed before atoms_object could be assigned
+             structure_info_for_analysis += f" This might be due to the execution error: {exec_error_msg.splitlines()[0]}"
+
+    # 4. Generate analysis using Gemini
+    try:
+        prompt_for_analysis = PROMPT_FOR_ANALYSIS.format(
+            user_query=user_query_str,
+            generated_code=generated_code_str,
+            simulation_output=simulation_stdout_str,
+            structure_info=structure_info_for_analysis
+        )
+        response_analysis = MODEL.generate_content(prompt_for_analysis)
+        analysis_text_str = response_analysis.text
+    except Exception as e:
+        analysis_text_str = f"Error generating analysis with Gemini: {traceback.format_exc()}"
+        if exec_error_msg: # If code exec failed, make analysis reflect that primarily
+            analysis_text_str = f"Could not perform analysis as the ASE code execution failed.\nDetails: {exec_error_msg}\n\nGemini analysis generation also failed: {e}"
+        elif not generated_code_str:
+            analysis_text_str = f"Could not perform analysis as ASE code generation failed.\nGemini analysis generation also failed: {e}"
+
+    return generated_code_str, simulation_stdout_str, pil_image, analysis_text_str
+
+# --- Gradio Interface ---
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="teal", secondary_hue="orange")) as demo:
+    gr.Markdown("# ⚛️ GeminiSim: ASE Simulations via Natural Language 💬")
+    gr.Markdown(
+        "Describe your atomic/molecular simulation task in plain English. "
+        "Gemini AI will generate ASE Python code, run the simulation, and provide an analysis of the results.\n\n"
+        "**⚠️ SECURITY WARNING:** This tool executes AI-generated Python code using `exec()`. "
+        "This is inherently unsafe if the AI generates malicious code. Use with caution and only with trusted inputs, "
+        "especially in non-sandboxed environments."
+    )
 
-with gr.Blocks(title="ASE+Gemini MCP Tool") as iface:
-    gr.Markdown("# ASE+Gemini Molecular Simulation Tool")
-    
     with gr.Row():
         with gr.Column(scale=2):
-            user_prompt = gr.Textbox(
-                label="Enter your simulation prompt",
-                placeholder="e.g., 'Optimize a water molecule with DFTB+ calculator'",
-                lines=3
+            user_query_input = gr.Textbox(
+                label="Your Simulation Request:",
+                placeholder="e.g., 'Optimize a water molecule (H2O) using EMT and show the final geometry and energy.'",
+                lines=4,
+                elem_id="user_query_textbox"
             )
-            submit_btn = gr.Button("Run Simulation")
-            
-    with gr.Row():
-        with gr.Column():
-            result_image = gr.Image(label="Simulation Result")
-        with gr.Column():
-            analysis = gr.Textbox(label="LLM Analysis")
-    
-    with gr.Row():
-        with gr.Column():
-            prompt_sent = gr.Textbox(label="Prompt Sent to Gemini")
-        with gr.Column():
-            generated_code = gr.Code(label="Generated ASE Code")
+        with gr.Column(scale=1):
+            gr.Markdown("### ✨ Example Prompts:")
+            gr.Examples(
+                examples=[
+                    "Optimize a CO molecule using BFGS and show its bond length.",
+                    "Create a 2x2x2 supercell of Silicon (Si) and get its potential energy using EMT.",
+                    "Build a methane molecule (CH4) and calculate its energy with EMT.",
+                    "Simulate a single gold (Au) atom."
+                ],
+                inputs=[user_query_input]
+            )
+
+    submit_button = gr.Button("🚀 Simulate & Analyze", variant="primary")
     
-    submit_btn.click(
-        fn=mcp_tool,
-        inputs=user_prompt,
-        outputs=[result_image, analysis, prompt_sent, generated_code]
+    with gr.Accordion("🔬 Results", open=True):
+        with gr.Row():
+            with gr.Column(scale=1):
+                gr.Markdown("### 🤖 Generated ASE Code")
+                generated_code_display = gr.Code(language="python", label="Generated ASE Python Code", lines=15, interactive=False)
+                gr.Markdown("### 🖥️ Simulation Standard Output")
+                simulation_stdout_display = gr.Textbox(label="Simulation stdout/stderr", lines=8, interactive=False, show_copy_button=True)
+            
+            with gr.Column(scale=1):
+                gr.Markdown("### 🖼️ Molecular Visualization")
+                molecule_plot_display = gr.Image(label="Molecular Structure", type="pil", interactive=False, height=350, show_download_button=True)
+                gr.Markdown("### 🧠 AI Analysis")
+                ai_analysis_display = gr.Textbox(label="Gemini's Analysis", lines=10, interactive=False, show_copy_button=True)
+
+    gr.Markdown("---")
+    gr.Markdown("Built with [Gradio](https://gradio.app/), [ASE](https://wiki.fysik.dtu.dk/ase/), and [Google Gemini](https://ai.google.dev/).")
+
+    submit_button.click(
+        fn=simulate_and_analyze,
+        inputs=[user_query_input],
+        outputs=[generated_code_display, simulation_stdout_display, molecule_plot_display, ai_analysis_display]
     )
 
 if __name__ == "__main__":
-    iface.launch()
+    if not API_KEY or not MODEL:
+        print("ERROR: Gemini API key not set or model failed to initialize. Gradio server will not start.")
+        # Optionally, could display an error in Gradio itself if it were to launch.
+    else:
+        print("Launching Gradio server...")
+        demo.launch() # Add share=True if you want a public link (be mindful of security warning!)

requirements.txt

@@ -1,6 +1,7 @@
 gradio
 ase
 matplotlib
-google-genai
+google-generativeai
 numpy
-scipy
+scipy
+python-dotenv