Update mcp/disgenet.py

mcp/disgenet.py

@@ -1,23 +1,68 @@
-# mcp/disgenet.py
-"""
-DisGeNET disease–gene associations.
+#!/usr/bin/env python3
+"""MedGenesis – DisGeNET async helper (disease → gene associations).
+
+Features
+~~~~~~~~
+* Accepts optional Bearer token via env **`DISGENET_KEY`** (rate‑limit free).
+* Endpoint: `https://www.disgenet.org/api/gda/disease/<disease_name>`
+* Back‑off retry (2×, 4×) for 429/5xx.
+* LRU cache (24 h, 512 queries) to minimise API calls.
+* Returns top *N* rows (default = 10) as `list[dict]`.
 """
+from __future__ import annotations
 
-import os, httpx
-from typing import List, Dict
+import os, asyncio, httpx
+from functools import lru_cache
+from typing import List, Dict, Any
 
-DISGENET_KEY = os.getenv("DISGENET_KEY")
-HEADERS = {"Authorization": f"Bearer {DISGENET_KEY}"} if DISGENET_KEY else {}
+_BASE = "https://www.disgenet.org/api/gda/disease"
+_TOKEN = os.getenv("DISGENET_KEY")
+_HEADERS = {"Authorization": f"Bearer {_TOKEN}"} if (_TOKEN := os.getenv("DISGENET_KEY")) else {}
+_TIMEOUT = 15
 
-BASE = "https://www.disgenet.org/api"
+# ---------------------------------------------------------------------
+# Internal fetch with retry
+# ---------------------------------------------------------------------
+async def _fetch(url: str, *, retries: int = 3) -> List[Dict[str, Any]]:
+    delay = 2
+    for _ in range(retries):
+        async with httpx.AsyncClient(timeout=_TIMEOUT, headers=_HEADERS) as cli:
+            resp = await cli.get(url, params={"source": "ALL", "format": "json"})
+            if resp.status_code == 200:
+                return resp.json()
+            if resp.status_code in {429, 500, 503}:
+                await asyncio.sleep(delay)
+                delay *= 2
+                continue
+            resp.raise_for_status()
+    return []  # final fallback
 
-async def disease_to_genes(disease_name: str, limit: int = 10) -> List[Dict]:
-    """
-    Return top gene associations for a disease.
+# ---------------------------------------------------------------------
+# Public API – cached 24 h
+# ---------------------------------------------------------------------
+@lru_cache(maxsize=512)
+async def disease_to_genes(disease_name: str, *, limit: int = 10) -> List[Dict]:
+    """Return up to *limit* gene‑association dicts for *disease_name*.
+
+    Parameters
+    ----------
+    disease_name : str
+        Free‑text disease label (e.g. "glioblastoma"). Internally converted
+        to lowercase and URL‑encoded.
+    limit : int, optional
+        Maximum number of rows to return (default = 10).
     """
-    url = f"{BASE}/gda/disease/{disease_name.lower()}"
-    async with httpx.AsyncClient(timeout=20, headers=HEADERS) as client:
-        r = await client.get(url, params={"source": "ALL", "format": "json"})
-        r.raise_for_status()
-        data = r.json()
-        return data[:limit]
+    url = f"{_BASE}/{disease_name.lower()}"
+    data = await _fetch(url)
+    return data[:limit]
+
+
+# ---------------------------------------------------------------------
+# CLI demo
+# ---------------------------------------------------------------------
+if __name__ == "__main__":
+    import json
+    async def _demo():
+        out = await disease_to_genes("glioblastoma", limit=5)
+        print(json.dumps(out[:2], indent=2))
+    asyncio.run(_demo())