Update genesis/api_clients/umls_api.py

genesis/api_clients/umls_api.py

@@ -3,40 +3,71 @@ import os
 import requests
 
 UMLS_API_KEY = os.getenv("UMLS_API_KEY")
-UMLS_AUTH_URL = "https://utslogin.nlm.nih.gov/cas/v1/api-key"
-UMLS_SEARCH_URL = "https://uts-ws.nlm.nih.gov/rest/search/current"
+UMLS_AUTH_URL = "https://utslogin.nlm.nih.gov"
+UMLS_API_URL = "https://uts-ws.nlm.nih.gov/rest"
 
-def get_umls_ticket():
+if not UMLS_API_KEY:
+    raise ValueError("Missing UMLS_API_KEY in environment variables")
+
+def get_tgt():
     """
-    Get a single-use ticket for UMLS API queries.
+    Get a Ticket Granting Ticket (TGT) for UMLS authentication.
     """
-    if not UMLS_API_KEY:
-        raise ValueError("UMLS_API_KEY is missing in environment variables")
-
+    url = f"{UMLS_AUTH_URL}/cas/v1/api-key"
     params = {"apikey": UMLS_API_KEY}
-    res = requests.post(UMLS_AUTH_URL, data=params)
+    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.strip()
+    return res.text
 
-def search_umls(term: str, max_results: int = 10):
+def search_umls(query: str, page: int = 1, page_size: int = 10):
     """
-    Search UMLS Metathesaurus for concepts related to a term.
+    Search for a term in UMLS.
     """
-    ticket = get_umls_ticket()
+    tgt = get_tgt()
+    ticket = get_service_ticket(tgt)
+    url = f"{UMLS_API_URL}/search/current"
     params = {
-        "string": term,
+        "string": query,
         "ticket": ticket,
-        "pageSize": max_results
+        "pageNumber": page,
+        "pageSize": page_size
     }
-    res = requests.get(UMLS_SEARCH_URL, params=params)
+    res = requests.get(url, params=params)
     res.raise_for_status()
-    data = res.json()
+    return res.json().get("result", {}).get("results", [])
 
-    results = []
-    for result in data.get("result", {}).get("results", []):
-        results.append({
-            "name": result.get("name"),
-            "ui": result.get("ui"),
-            "rootSource": result.get("rootSource")
-        })
-    return 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"
+    ]