mcp/umls.py

# mcp/umls.py

import httpx
import os

UMLS_API_KEY = os.environ.get("UMLS_KEY")
UMLS_AUTH = "https://utslogin.nlm.nih.gov/cas/v1/api-key"
UMLS_SEARCH = "https://uts-ws.nlm.nih.gov/rest/search/current"

async def _get_umls_ticket():
    async with httpx.AsyncClient() as client:
        resp = await client.post(UMLS_AUTH, data={"apikey": UMLS_API_KEY})
        if resp.status_code != 201:
            raise Exception("UMLS Authentication Failed")
        tgt = resp.text.split('action="')[1].split('"')[0]
        ticket_resp = await client.post(tgt, data={"service": "http://umlsks.nlm.nih.gov"})
        return ticket_resp.text

async def lookup_umls(term: str):
    ticket = await _get_umls_ticket()
    params = {
        "string": term,
        "ticket": ticket,
        "pageSize": 1,
    }
    async with httpx.AsyncClient() as client:
        resp = await client.get(UMLS_SEARCH, params=params)
        items = resp.json().get("result", {}).get("results", [])
        if items:
            cui = items[0].get("ui")
            name = items[0].get("name")
            definition = items[0].get("rootSource")
            return {"term": term, "cui": cui, "name": name, "definition": definition}
    return {"term": term, "cui": None, "name": None, "definition": None}