genesis/api_clients/umls_api.py

# genesis/api_clients/umls_api.py
import os
import requests

UMLS_API_KEY = os.getenv("UMLS_API_KEY")
UMLS_AUTH_URL = "https://utslogin.nlm.nih.gov"
UMLS_API_URL = "https://uts-ws.nlm.nih.gov/rest"

if not UMLS_API_KEY:
    raise ValueError("Missing UMLS_API_KEY in environment variables")

def get_tgt():
    """
    Get a Ticket Granting Ticket (TGT) for UMLS authentication.
    """
    url = f"{UMLS_AUTH_URL}/cas/v1/api-key"
    params = {"apikey": UMLS_API_KEY}
    res = requests.post(url, data=params)
    res.raise_for_status()
    # Extract TGT URL from HTML form action
    return res.text.split('action="')[1].split('"')[0]

def get_service_ticket(tgt):
    """
    Get a Service Ticket (ST) using TGT.
    """
    params = {"service": "http://umlsks.nlm.nih.gov"}
    res = requests.post(tgt, data=params)
    res.raise_for_status()
    return res.text

def search_umls(query: str, page: int = 1, page_size: int = 10):
    """
    Search for a term in UMLS.
    """
    tgt = get_tgt()
    ticket = get_service_ticket(tgt)
    url = f"{UMLS_API_URL}/search/current"
    params = {
        "string": query,
        "ticket": ticket,
        "pageNumber": page,
        "pageSize": page_size
    }
    res = requests.get(url, params=params)
    res.raise_for_status()
    return res.json().get("result", {}).get("results", [])

def get_cui_details(cui: str):
    """
    Get detailed info for a CUI (Concept Unique Identifier).
    """
    tgt = get_tgt()
    ticket = get_service_ticket(tgt)
    url = f"{UMLS_API_URL}/content/current/CUI/{cui}"
    params = {"ticket": ticket}
    res = requests.get(url, params=params)
    res.raise_for_status()
    return res.json()

def map_to_umls(term: str):
    """
    Given a term, return the matching UMLS CUI(s) and definitions.
    """
    matches = search_umls(term)
    return [
        {
            "name": m.get("name"),
            "cui": m.get("ui"),
            "rootSource": m.get("rootSource")
        }
        for m in matches if m.get("ui") != "NONE"
    ]