Update app.py with 3D molecule viewer

app.py

@@ -1,7 +1,11 @@
 import gradio as gr
 from rdkit import Chem
-from rdkit.Chem import Descriptors, Draw
+from rdkit.Chem import Descriptors, Draw, AllChem
 import cirpy
+import py3Dmol
+
+
+# RDKit API with multiple endpoints
 
 
 def _mol_from_smiles(smiles: str):
@@ -47,6 +51,46 @@ def name_to_smiles(name: str) -> str:
         raise gr.Error(f"Error converting name to SMILES: {str(e)}")
 
 
+def name_to_3d_molecule(name: str) -> str:
+    """Convert chemical name to 3D molecule visualization"""
+    try:
+        # Convert name to SMILES
+        smiles = cirpy.resolve(name, 'smiles')
+        if smiles is None:
+            raise gr.Error(f"Could not find SMILES for chemical name: {name}")
+
+        # Create molecule from SMILES
+        mol = Chem.MolFromSmiles(smiles)
+        if mol is None:
+            raise gr.Error(f"Could not create molecule from SMILES: {smiles}")
+
+        # Add hydrogens for better 3D structure
+        mol = Chem.AddHs(mol)
+
+        # Generate 3D coordinates
+        success = AllChem.EmbedMolecule(mol, AllChem.ETKDG())
+        if success == -1:
+            raise gr.Error(f"Could not generate 3D coordinates for: {name}")
+
+        # Optimize geometry
+        AllChem.MMFFOptimizeMolecule(mol)
+
+        # Convert to PDB format for py3Dmol
+        pdb_block = Chem.MolToPDBBlock(mol)
+
+        # Create 3D viewer
+        viewer = py3Dmol.view(width=400, height=400)
+        viewer.addModel(pdb_block, 'pdb')
+        viewer.setStyle({'stick': {}})
+        viewer.zoomTo()
+
+        # Return HTML representation
+        return viewer._repr_html_()
+
+    except Exception as e:
+        raise gr.Error(f"Error creating 3D molecule: {str(e)}")
+
+
 smiles_interface = gr.Interface(
     fn=smiles_to_canonical,
     inputs=gr.Textbox(label="SMILES"),
@@ -88,11 +132,21 @@ tpsa_interface = gr.Interface(
     description="Calculate the topological polar surface area (TPSA).",
 )
 
+molecule_3d_interface = gr.Interface(
+    fn=name_to_3d_molecule,
+    inputs=gr.Textbox(label="Chemical Name", placeholder="e.g., benzene, aspirin, caffeine"),
+    outputs=gr.HTML(label="3D Molecule Viewer"),
+    api_name="name_to_3d_molecule",
+    description="Convert a chemical name to an interactive 3D molecule visualization.",
+    examples=[["benzene"], ["aspirin"], ["caffeine"], ["ethanol"]],
+)
+
 
 demo = gr.TabbedInterface(
-    [name_interface, smiles_interface, mw_interface, logp_interface, tpsa_interface],
+    [name_interface, molecule_3d_interface, smiles_interface, mw_interface, logp_interface, tpsa_interface],
     [
         "Name to SMILES",
+        "3D Molecule Viewer",
         "SMILES to Canonical",
         "Molecular Weight",
         "LogP",