Update mcp/umls.py

mcp/umls.py

@@ -1,48 +1,66 @@
-import os, httpx, asyncio
+# mcp/umls.py
+"""
+Async UMLS lookup via NLM REST API.
+Graceful fallback: always returns a dict with all keys even on failure.
+"""
+
+import os
+import httpx
 from functools import lru_cache
 
-UMLS_KEY   = os.getenv("UMLS_KEY")          # MUST be set in HF Secrets
-_AUTH_URL  = "https://utslogin.nlm.nih.gov/cas/v1/api-key"
-_SEARCH_URL= "https://uts-ws.nlm.nih.gov/rest/search/current"
+UMLS_KEY    = os.getenv("UMLS_KEY")          # Should be set as a secret (Hugging Face/colab/etc)
+_AUTH_URL   = "https://utslogin.nlm.nih.gov/cas/v1/api-key"
+_SEARCH_URL = "https://uts-ws.nlm.nih.gov/rest/search/current"
 
-# ── internal helpers ────────────────────────────────────────────────
+# ── Internal helper: get a single-use service ticket ──────────────
 async def _get_ticket() -> str | None:
-    """Return a single-use service ticket or None on auth failure."""
+    """
+    Return a single-use UMLS service ticket, or None if auth fails.
+    """
     if not UMLS_KEY:
         return None
     try:
-        async with httpx.AsyncClient(timeout=10) as c:
-            tgt = await c.post(_AUTH_URL, data={"apikey": UMLS_KEY})
-            tgt.raise_for_status()
-            action = tgt.text.split('action="')[1].split('"')[0]
-            st   = await c.post(action, data={"service": "http://umlsks.nlm.nih.gov"})
-            return st.text
+        async with httpx.AsyncClient(timeout=10) as client:
+            # 1. Get ticket-granting ticket (TGT)
+            tgt_resp = await client.post(_AUTH_URL, data={"apikey": UMLS_KEY})
+            tgt_resp.raise_for_status()
+            tgt_url = tgt_resp.text.split('action="')[1].split('"')[0]
+            # 2. Exchange for single-use service ticket
+            st_resp = await client.post(tgt_url, data={"service": "http://umlsks.nlm.nih.gov"})
+            return st_resp.text.strip()
     except Exception:
         return None
 
-# ── public API (cached) ─────────────────────────────────────────────
+# ── Public API: async, returns safe dict ──────────────────────────
 @lru_cache(maxsize=512)
 async def lookup_umls(term: str) -> dict:
     """
-    Return {term,cui,name,definition}.   Never raises.
-    If auth/quota fails → returns everything as None (safe for UI).
+    Lookup a UMLS concept by text term.
+    Returns dict with {term, cui, name, definition}, always.
+    If quota/auth fails: all values None except term.
     """
     ticket = await _get_ticket()
     if not ticket:
         return {"term": term, "cui": None, "name": None, "definition": None}
-
-    params = {"string": term, "ticket": ticket, "pageSize": 1}
+    params = {
+        "string": term,
+        "ticket": ticket,
+        "pageSize": 1,
+    }
     try:
-        async with httpx.AsyncClient(timeout=8) as c:
-            r = await c.get(_SEARCH_URL, params=params)
-            r.raise_for_status()
-            items = r.json()["result"]["results"]
-            hit   = items[0] if items else {}
-            return {
-                "term": term,
-                "cui":  hit.get("ui"),
-                "name": hit.get("name"),
-                "definition": hit.get("rootSource"),
-            }
+        async with httpx.AsyncClient(timeout=8) as client:
+            resp = await client.get(_SEARCH_URL, params=params)
+            resp.raise_for_status()
+            items = resp.json().get("result", {}).get("results", [])
+            if items:
+                hit = items[0]
+                return {
+                    "term": term,
+                    "cui": hit.get("ui"),
+                    "name": hit.get("name"),
+                    "definition": hit.get("rootSource"),
+                }
+            else:
+                return {"term": term, "cui": None, "name": None, "definition": None}
     except Exception:
         return {"term": term, "cui": None, "name": None, "definition": None}