mcp/umls.py

import os, httpx, asyncio
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"

# ── internal helpers ────────────────────────────────────────────────
async def _get_ticket() -> str | None:
    """Return a single-use service ticket or None on auth failure."""
    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
    except Exception:
        return None

# ── public API (cached) ─────────────────────────────────────────────
@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).
    """
    ticket = await _get_ticket()
    if not ticket:
        return {"term": term, "cui": None, "name": None, "definition": None}

    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"),
            }
    except Exception:
        return {"term": term, "cui": None, "name": None, "definition": None}