Update app.py

app.py

@@ -1,9 +1,10 @@
 import os
 import requests
 import gradio as gr
-from ase.io import read
-import nglview as nv
-import tempfile
+import json  # For handling JSON responses
+from pymatgen.core import Structure, Lattice
+from pymatgen.io.cif import CifWriter
+from pymatgen.io.xyz import XYZ
 
 # Retrieve the API key from the environment variable
 groq_api_key = os.getenv("GROQ_API_KEY")
@@ -15,61 +16,150 @@ url = "https://api.groq.com/openai/v1/chat/completions"
 headers = {"Authorization": f"Bearer {groq_api_key}"}
 
 # Function to interact with Groq API
-def chat_with_groq(user_input):
+def get_material_info(user_input):
+    prompt = f"""You are a materials science expert. A user is asking for applications of materials.
+    Based on the user's request: "{user_input}", identify the 3 best materials for this application.
+    For each material, provide:
+    - Material Name:
+    - Key Properties relevant to the application: (e.g., strength, conductivity, melting point) with values if possible.
+    - A brief explanation of why this material is suitable for the application.
+    - A simplified representation of its atomic structure (if readily available and can be described textually, e.g., "FCC lattice", "HCP lattice", or a simple chemical formula with a basic structural description).
+
+    Format your response as a JSON object with the following structure:
+    {
+      "materials": [
+        {
+          "Material Name": "...",
+          "Key Properties": {
+            "Property 1": "value unit",
+            "Property 2": "value unit",
+            "...": "..."
+          },
+          "Suitability Explanation": "...",
+          "Atomic Structure": "..."
+        },
+        {
+          "Material Name": "...",
+          "Key Properties": {
+            "Property 1": "value unit",
+            "Property 2": "value unit",
+            "...": "..."
+          },
+          "Suitability Explanation": "...",
+          "Atomic Structure": "..."
+        },
+        {
+          "Material Name": "...",
+          "Key Properties": {
+            "Property 1": "value unit",
+            "Property 2": "value unit",
+            "...": "..."
+          },
+          "Suitability Explanation": "...",
+          "Atomic Structure": "..."
+        }
+      ]
+    }
+    """
     body = {
         "model": "llama-3.1-8b-instant",
-        "messages": [{"role": "user", "content": user_input}]
+        "messages": [{"role": "user", "content": prompt}]
     }
+
     response = requests.post(url, headers=headers, json=body)
+
     if response.status_code == 200:
-        return response.json()['choices'][0]['message']['content']
+        try:
+            return json.loads(response.json()['choices'][0]['message']['content'])
+        except json.JSONDecodeError:
+            return f"Error decoding JSON: {response.text}"
     else:
         return f"Error: {response.json()}"
 
+def create_structure_file(material_info, file_format="xyz"):
+    if not isinstance(material_info, dict) or "materials" not in material_info or not material_info["materials"]:
+        return "No material information found to create structure file.", None
+
+    # Choose the first material for structure generation (can be modified to let user choose)
+    first_material = material_info["materials"][0]
+    structure_description = first_material.get("Atomic Structure", "")
+    material_name = first_material.get("Material Name", "unknown")
+
+    if not structure_description:
+        return f"No atomic structure information available for {material_name}.", None
+
+    # Attempt to create a basic structure based on the description
+    try:
+        if "FCC lattice" in structure_description.lower():
+            lattice = Lattice.cubic(4.0)  # Example lattice parameter
+            structure = Structure(lattice, ["A", "A", "A", "A"], [[0, 0, 0], [0.5, 0.5, 0], [0.5, 0, 0.5], [0, 0.5, 0.5]])
+        elif "HCP lattice" in structure_description.lower():
+            lattice = Lattice.hexagonal(3.0, 5.0) # Example lattice parameters
+            structure = Structure(lattice, ["A", "A"], [[0, 0, 0], [1/3, 2/3, 1/2]])
+        elif structure_description and len(structure_description.split()) == 2: # Attempt simple diatomic
+            formula, struct = structure_description.split()
+            if len(formula) == 2 and len(struct.lower()) == 3 and "unit" in struct.lower():
+                lattice = Lattice.cubic(3.5) # Another example
+                structure = Structure(lattice, [formula[0], formula[1]], [[0, 0, 0], [0.5, 0.5, 0.5]])
+        else:
+            return f"Could not interpret the atomic structure description for {material_name}.", None
+
+        # Limit to 20 atoms if the generated structure has more
+        if structure.num_sites > 20:
+            structure = structure[:20]
+
+        if file_format == "xyz":
+            filepath = f"{material_name.replace(' ', '_')}_20atoms.xyz"
+            XYZ(structure).write_file(filepath)
+        elif file_format == "cif":
+            filepath = f"{material_name.replace(' ', '_')}_20atoms.cif"
+            CifWriter(structure, symprec=1e-5).write_file(filepath)
+        else:
+            return f"Unsupported file format: {file_format}", None
+
+        with open(filepath, 'r') as f:
+            file_content = f.read()
+        os.remove(filepath) # Clean up the temporary file
+        return f"Successfully created {file_format} file for {material_name} (first 20 atoms).", file_content
+
+    except Exception as e:
+        return f"Error creating structure file for {material_name}: {e}", None
+
+def chat_and_generate(user_input, file_format):
+    material_info = get_material_info(user_input)
+    structure_message, file_content = create_structure_file(material_info, file_format)
+
+    output_text = ""
+    if isinstance(material_info, dict) and "materials" in material_info:
+        for i, material in enumerate(material_info["materials"]):
+            output_text += f"**Material {i+1}: {material['Material Name']}**\n"
+            output_text += f"Key Properties: {', '.join([f'{k}: {v}' for k, v in material['Key Properties'].items()])}\n"
+            output_text += f"Suitability: {material['Suitability Explanation']}\n"
+            output_text += f"Atomic Structure: {material['Atomic Structure']}\n\n"
+    else:
+        output_text = str(material_info)
+
+    return output_text, structure_message, file_content
 
-def parse_response(response):
-    # Dummy parser - replace this with a proper LLM response parser as needed
-    lines = response.split('\n')
-    materials = []
-    for line in lines:
-        if any(keyword in line.lower() for keyword in ["1.", "2.", "3."]):
-            name = line.split(".")[1].split("-")[0].strip()
-            materials.append(name)
-    return materials
-
-
-# Function to generate and visualize CIF files
-def generate_and_visualize_cifs(application):
-    # Get response from Groq API
-    response = chat_with_groq(f"Suggest 3 best materials for {application} with their properties and provide CIF data.")
-    
-    # Parse the response to extract materials and their CIF contents
-    # This is a placeholder; actual implementation depends on response format
-    materials = parse_response(response)
-    
-    views = []
-    for material in materials:
-        cif_content = material['cif_content']
-        with tempfile.NamedTemporaryFile(mode='w+', suffix='.cif', delete=False) as tmp:
-            tmp.write(cif_content)
-            tmp_path = tmp.name
-        atoms = read(tmp_path)
-        view = nv.show_ase(atoms)
-        views.append(view)
-    return views
-
-# Gradio interface
-def interface_function(application):
-    views = generate_and_visualize_cifs(application)
-    return views
+# Create Gradio interface
+inputs = [
+    gr.Textbox(lines=2, placeholder="Ask for material applications (e.g., 'materials for high-temperature superconductors')."),
+    gr.Radio(["xyz", "cif"], label="Generate Structure File (first material)", value="xyz")
+]
+outputs = [
+    gr.Markdown(),
+    gr.Textbox(label="Structure Generation Status"),
+    gr.Code(label="Generated Structure File Content")
+]
 
 interface = gr.Interface(
-    fn=interface_function,
-    inputs=gr.Textbox(lines=2, placeholder="Enter material application..."),
-    outputs="text",
-    title="Material Selector and Visualizer",
-    description="Enter a material application to get top 3 materials and visualize their atomic structures."
+    fn=chat_and_generate,
+    inputs=inputs,
+    outputs=outputs,
+    title="Materials Science Expert AI",
+    description="Ask about applications of materials and get information on the top 3 candidates with their properties and a generated atomic structure file (first 20 atoms).",
 )
 
+# Launch Gradio app
 if __name__ == "__main__":
-    interface.launch()
+    interface.launch()