Spaces:
Running
Running
File size: 2,963 Bytes
5bb077e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | """PubChem PUG REST + autocomplete helpers for compound lookup."""
from __future__ import annotations
import asyncio
import logging
from urllib.parse import quote
import httpx
logger = logging.getLogger(__name__)
PUG_BASE = "https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound"
AUTOCOMPLETE_BASE = "https://pubchem.ncbi.nlm.nih.gov/rest/autocomplete/compound"
_REQUEST_TIMEOUT = 12.0
class PubChemError(Exception):
pass
async def _get_json(url: str) -> dict:
try:
async with httpx.AsyncClient(timeout=_REQUEST_TIMEOUT, follow_redirects=True) as client:
resp = await client.get(url)
resp.raise_for_status()
return resp.json()
except httpx.HTTPStatusError as e:
if e.response.status_code == 404:
raise PubChemError("Compound not found in PubChem")
raise PubChemError(f"PubChem HTTP {e.response.status_code}")
except (httpx.RequestError, ValueError) as e:
raise PubChemError(f"PubChem request failed: {e}")
async def name_to_cid(name: str) -> int | None:
"""Resolve a chemical name to a PubChem CID (first hit)."""
url = f"{PUG_BASE}/name/{quote(name)}/cids/JSON"
try:
data = await _get_json(url)
except PubChemError:
return None
ids = data.get("IdentifierList", {}).get("CID", [])
return int(ids[0]) if ids else None
async def cid_to_record(cid: int) -> dict:
"""Fetch canonical SMILES + name metadata for a CID."""
url = (
f"{PUG_BASE}/cid/{int(cid)}/property/"
f"IsomericSMILES,CanonicalSMILES,IUPACName,MolecularFormula,InChIKey/JSON"
)
data = await _get_json(url)
props = data.get("PropertyTable", {}).get("Properties", [{}])[0]
return {
"cid": int(props.get("CID", cid)),
"smiles": props.get("SMILES") or props.get("IsomericSMILES") or props.get("CanonicalSMILES"),
"name": props.get("IUPACName"),
"formula": props.get("MolecularFormula"),
"inchikey": props.get("InChIKey"),
}
async def search_suggestions(query: str, limit: int = 10) -> list[dict]:
"""Return PubChem autocomplete suggestions for a compound search box.
The autocomplete endpoint returns names only, so CIDs are resolved in
parallel. Each entry: {cid, name}.
"""
if not query or not query.strip():
return []
url = f"{AUTOCOMPLETE_BASE}/{quote(query.strip())}/JSON?limit={int(limit)}"
try:
data = await _get_json(url)
except PubChemError:
return []
names = data.get("dictionary_terms", {}).get("compound", [])[:limit]
if not names:
return []
sem = asyncio.Semaphore(6)
async def _resolve(name: str) -> dict | None:
async with sem:
cid = await name_to_cid(name)
if not cid:
return None
return {"cid": int(cid), "name": name}
resolved = await asyncio.gather(*(_resolve(n) for n in names))
return [r for r in resolved if r]
|