Update genesis/structures.py

genesis/structures.py

@@ -1,33 +1,40 @@
 # genesis/structures.py
+"""
+Structure Fetching for GENESIS-AI
+Integrates PDBe (PDB) and ChEMBL molecule lookups.
+"""
+
 import requests
-from typing import List, Dict
 
-def fetch_structures_for_terms(terms: List[str]) -> List[Dict[str, str]]:
-    """Fetch PDB structures from PDBe or RCSB for given search terms."""
-    structures = []
+PDBe_API = "https://www.ebi.ac.uk/pdbe/api/pdb/entry/summary/"
+CHEMBL_API = "https://www.ebi.ac.uk/chembl/api/data/molecule/"
 
-    for term in terms:
-        try:
-            # PDBe search
-            pdbe_url = f"https://www.ebi.ac.uk/pdbe/search/pdb/select"
-            params = {"q": term, "rows": 2, "wt": "json"}
-            r = requests.get(pdbe_url, params=params, timeout=10)
-            r.raise_for_status()
-            data = r.json()
+def fetch_pdb_structure(pdb_id):
+    """Fetch PDB structure summary from PDBe."""
+    try:
+        r = requests.get(f"{PDBe_API}{pdb_id.lower()}")
+        r.raise_for_status()
+        return r.json()
+    except Exception as e:
+        return {"error": str(e)}
 
-            docs = data.get("response", {}).get("docs", [])
-            for doc in docs:
-                pdb_id = doc.get("pdb_id")
-                title = doc.get("title", term)
-                if pdb_id:
-                    structures.append({
-                        "term": term,
-                        "pdb_id": pdb_id,
-                        "title": title,
-                        "pdbe_url": f"https://www.ebi.ac.uk/pdbe/entry/pdb/{pdb_id}",
-                        "rcsb_url": f"https://www.rcsb.org/structure/{pdb_id}"
-                    })
-        except Exception as e:
-            print(f"[Structures] Failed to fetch for {term}: {e}")
+def fetch_chembl_molecule(chembl_id):
+    """Fetch ChEMBL molecule data."""
+    try:
+        r = requests.get(f"{CHEMBL_API}{chembl_id}.json")
+        r.raise_for_status()
+        return r.json()
+    except Exception as e:
+        return {"error": str(e)}
 
-    return structures
+def fetch_structures_for_terms(terms):
+    """
+    Searches for possible PDB or ChEMBL IDs in terms and fetches details.
+    """
+    results = []
+    for term in terms:
+        if term.upper().startswith("CHEMBL"):
+            results.append({"term": term, "data": fetch_chembl_molecule(term)})
+        elif len(term) == 4 and term.isalnum():
+            results.append({"term": term, "data": fetch_pdb_structure(term)})
+    return results