Update genesis/molecule_viewer.py

genesis/molecule_viewer.py

@@ -1,45 +1,48 @@
 # genesis/molecule_viewer.py
+"""
+Molecule Viewer for GENESIS-AI
+Fetches 3D molecular structure data and metadata from public chemical databases.
+"""
+
 import requests
-from typing import Optional
 
-def fetch_pdb(pdb_id: str) -> Optional[str]:
-    """Fetch PDB file from RCSB."""
-    url = f"https://files.rcsb.org/download/{pdb_id.upper()}.pdb"
-    r = requests.get(url)
-    if r.status_code == 200:
-        return r.text
-    return None
+PDBE_API = "https://www.ebi.ac.uk/pdbe/api/pdb/entry/summary/"
+PUBCHEM_API = "https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/name"
+CHEMBL_API = "https://www.ebi.ac.uk/chembl/api/data/molecule"
+
+def get_pdbe_structure(pdb_id):
+    """Fetch structure summary from PDBe."""
+    try:
+        r = requests.get(f"{PDBE_API}{pdb_id}")
+        r.raise_for_status()
+        return r.json()
+    except Exception as e:
+        return {"error": str(e)}
 
-def fetch_alphafold(uniprot_id: str) -> Optional[str]:
-    """Fetch PDB model from AlphaFold."""
-    url = f"https://alphafold.ebi.ac.uk/files/AF-{uniprot_id}-F1-model_v4.pdb"
-    r = requests.get(url)
-    if r.status_code == 200:
-        return r.text
-    return None
+def get_pubchem_structure(compound_name):
+    """Fetch compound data from PubChem."""
+    try:
+        url = f"{PUBCHEM_API}/{compound_name}/JSON"
+        r = requests.get(url)
+        r.raise_for_status()
+        return r.json()
+    except Exception as e:
+        return {"error": str(e)}
 
-def render_3dmol(pdb_data: str) -> str:
-    """Generate HTML for 3Dmol.js viewer with biotech neon styling."""
-    return f"""
-    <div id="viewer" style="width: 100%; height: 600px; background-color: black; border: 2px solid #00ff99; box-shadow: 0 0 20px #ff6600;"></div>
-    <script src="https://3dmol.csb.pitt.edu/build/3Dmol-min.js"></script>
-    <script>
-      let viewer = $3Dmol.createViewer('viewer', {{ backgroundColor: 'black' }});
-      let pdbData = `{pdb_data}`;
-      viewer.addModel(pdbData, "pdb");
-      viewer.setStyle({{}}, {{cartoon: {{color: 'spectrum'}}}});
-      viewer.zoomTo();
-      viewer.render();
-    </script>
-    """
+def get_chembl_structure(query):
+    """Search ChEMBL for molecule info."""
+    try:
+        r = requests.get(f"{CHEMBL_API}?molecule_synonyms__icontains={query}")
+        r.raise_for_status()
+        return r.json()
+    except Exception as e:
+        return {"error": str(e)}
 
-def get_molecule_view(pdb_or_gene: str) -> str:
-    """Fetch molecule structure and return 3D viewer HTML."""
-    # Try PDB ID first
-    pdb_data = fetch_pdb(pdb_or_gene)
-    if not pdb_data:
-        # Try AlphaFold with UniProt ID or gene name
-        pdb_data = fetch_alphafold(pdb_or_gene)
-    if pdb_data:
-        return render_3dmol(pdb_data)
-    return f"<p style='color:red;'>No structure found for {pdb_or_gene}</p>"
+def fetch_structure(query):
+    """Try PDBe (if PDB ID), else PubChem, else ChEMBL."""
+    if len(query) == 4 and query.isalnum():
+        return get_pdbe_structure(query)
+    pubchem_result = get_pubchem_structure(query)
+    if "error" not in pubchem_result:
+        return pubchem_result
+    return get_chembl_structure(query)