mcp/pubchem_ext.py

"""
pubchem_ext.py  ·  PubChem PUG-REST helper
Docs: https://pubchem.ncbi.nlm.nih.gov/docs/pug-rest-tutorial  📄

• No API-key is required for the public endpoint.
• If you set PUBCHEM_KEY (token), we include it as X-Pubchem-ApiKey
  which lifts burst limits.  Otherwise header is omitted.
• One retry on 5xx / ReadTimeout, then returns {}.
"""

from __future__ import annotations
import os, asyncio, httpx
from urllib.parse import quote_plus
from functools import lru_cache
from typing import Dict

_TOKEN = os.getenv("PUBCHEM_KEY")           # optional
_BASE  = "https://pubchem.ncbi.nlm.nih.gov/rest/pug"
_HDR   = {"Accept": "application/json"}
if _TOKEN:
    _HDR["X-Pubchem-ApiKey"] = _TOKEN

@lru_cache(maxsize=256)
async def fetch_compound(name: str) -> Dict:
    """
    Return compound data as dict or {} on failure.
    Uses /compound/name/<drug>/JSON endpoint.
    """
    url = f"{_BASE}/compound/name/{quote_plus(name)}/JSON"

    async def _hit() -> Dict:
        async with httpx.AsyncClient(timeout=10, headers=_HDR) as cli:
            r = await cli.get(url)
            if r.status_code == 404:
                return {}                 # no such compound
            r.raise_for_status()
            return r.json()

    try:
        return await _hit()
    except (httpx.HTTPStatusError, httpx.ReadTimeout):
        await asyncio.sleep(0.5)          # quick back-off
        try:
            return await _hit()
        except Exception:
            return {}                     # graceful fallback