app.py

import os
import requests
import gradio as gr
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")
if not groq_api_key:
    raise ValueError("GROQ_API_KEY is missing! Set it in the Hugging Face Spaces 'Secrets'.")

# Define the API endpoint and headers
url = "https://api.groq.com/openai/v1/chat/completions"
headers = {"Authorization": f"Bearer {groq_api_key}"}

# Function to interact with Groq API
def get_material_info(user_input):
    json_format = """
    ```json
    {
      "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": "..."
        }
      ]
    }
    ```
    """
    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:
    {json_format}
    """
    body = {
        "model": "llama-3.1-8b-instant",
        "messages": [{"role": "user", "content": prompt}]
    }

    response = requests.post(url, headers=headers, json=body)

    if response.status_code == 200:
        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

# 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=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()