Update mcp/wikidata.py

mcp/wikidata.py

@@ -1,23 +1,68 @@
-# mcp/wikidata.py
-"""
-Minimal Wikidata entity lookup for biomedical concepts.
+#!/usr/bin/env python3
+"""MedGenesis – minimal **Wikidata** lookup helper (async).
+
+Features
+~~~~~~~~
+* `simple_search(term)` – return first matching entity dict `{id, label, description}`.
+* `fetch_entity(qid)`  – return full entity data (`claims`, `labels`, etc.).
+* Uses public Wikidata REST API (no key). 15‑second timeout with `httpx`.
+* Least‑recently‑used cache (128) to avoid repeated hits when the same
+  concept appears across multiple papers.
 """
+from __future__ import annotations
 
-import httpx
+import httpx, asyncio
+from functools import lru_cache
 from typing import Dict, Optional
 
-API = "https://www.wikidata.org/w/api.php"
+_API = "https://www.wikidata.org/w/api.php"
+_TIMEOUT = 15
+_HEADERS = {"User-Agent": "MedGenesis/1.0 (https://huggingface.co/spaces)"}
+
+# ---------------------------------------------------------------------
+# Public helpers
+# ---------------------------------------------------------------------
 
+@lru_cache(maxsize=128)
 async def simple_search(term: str) -> Optional[Dict]:
+    """Return top search hit for *term* or None."""
     params = {
         "action": "wbsearchentities",
         "search": term,
         "language": "en",
         "format": "json",
-        "limit": 1
+        "limit": 1,
+    }
+    async with httpx.AsyncClient(timeout=_TIMEOUT, headers=_HEADERS) as client:
+        resp = await client.get(_API, params=params)
+        resp.raise_for_status()
+        hits = resp.json().get("search", [])
+        return hits[0] if hits else None
+
+
+@lru_cache(maxsize=128)
+async def fetch_entity(qid: str) -> Dict:
+    """Fetch full entity JSON for a Wikidata Q‑ID (e.g. `Q12136`)."""
+    params = {
+        "action": "wbgetentities",
+        "ids": qid,
+        "format": "json",
+        "languages": "en",
     }
-    async with httpx.AsyncClient(timeout=15) as client:
-        r = await client.get(API, params=params)
-        r.raise_for_status()
-        data = r.json()["search"]
-        return data[0] if data else None
+    async with httpx.AsyncClient(timeout=_TIMEOUT, headers=_HEADERS) as client:
+        resp = await client.get(_API, params=params)
+        resp.raise_for_status()
+        return resp.json().get("entities", {}).get(qid, {})
+
+
+# ---------------------------------------------------------------------
+# Demo / manual test
+# ---------------------------------------------------------------------
+if __name__ == "__main__":
+    async def _demo():
+        hit = await simple_search("glioblastoma")
+        print("Top hit:", hit)
+        if hit:
+            full = await fetch_entity(hit["id"])
+            print("Labels:", full.get("labels", {}).get("en", {}))
+    asyncio.run(_demo())