Update genesis/api_clients/umls_api.py

genesis/api_clients/umls_api.py

@@ -1,89 +1,105 @@
 # genesis/api_clients/umls_api.py
 import os
 import requests
-from typing import List, Dict, Optional
+from typing import Dict, List, Optional
 
-UMLS_API_KEY = os.getenv("UMLS_API_KEY")  # Set in Hugging Face secrets or .env
-UMLS_AUTH_ENDPOINT = "https://utslogin.nlm.nih.gov/cas/v1/api-key"
-UMLS_BASE = "https://uts-ws.nlm.nih.gov/rest"
+UMLS_API_KEY = os.getenv("UMLS_API_KEY")
+UMLS_BASE = "https://uts-ws.nlm.nih.gov"
+
+if not UMLS_API_KEY:
+    raise EnvironmentError("Missing UMLS_API_KEY in environment variables.")
 
 # -------------------------
 # AUTHENTICATION
 # -------------------------
 def get_tgt() -> str:
     """
-    Get Ticket Granting Ticket (TGT) from UMLS API key.
+    Get a Ticket Granting Ticket (TGT) from UMLS.
     """
-    params = {"apikey": UMLS_API_KEY}
-    r = requests.post(UMLS_AUTH_ENDPOINT, data=params)
-    r.raise_for_status()
-    # Extract TGT from HTML form action
-    tgt_url = r.text.split('action="')[1].split('"')[0]
-    return tgt_url
+    r = requests.post(
+        f"{UMLS_BASE}/restful/isValidService",
+        data={"apikey": UMLS_API_KEY}
+    )
+    if r.status_code != 201:
+        raise Exception(f"Failed to get UMLS TGT: {r.text}")
+    return r.headers.get("location")
 
-def get_service_ticket(tgt_url: str) -> str:
+def get_service_ticket(tgt: str) -> str:
     """
-    Get a one-time Service Ticket (ST) for API calls.
+    Get a single-use Service Ticket from the TGT.
     """
-    params = {"service": "http://umlsks.nlm.nih.gov"}
-    r = requests.post(tgt_url, data=params)
+    r = requests.post(tgt, data={"service": "http://umlsks.nlm.nih.gov"})
     r.raise_for_status()
     return r.text
 
 # -------------------------
-# SEARCH FUNCTIONS
+# CORE FUNCTIONS
 # -------------------------
-def search_umls(term: str, max_results: int = 10) -> List[Dict]:
+def search_cui(term: str, page_size: int = 10) -> List[Dict]:
     """
-    Search UMLS for a medical concept by term and return CUIs, preferred names, and semantic types.
+    Search for UMLS CUIs for a given term.
     """
     tgt = get_tgt()
-    st = get_service_ticket(tgt)
+    ticket = get_service_ticket(tgt)
+
     params = {
         "string": term,
-        "ticket": st,
-        "pageSize": max_results
+        "ticket": ticket,
+        "pageSize": page_size
     }
-    r = requests.get(f"{UMLS_BASE}/search/current", params=params)
+    r = requests.get(f"{UMLS_BASE}/rest/search/current", params=params)
     r.raise_for_status()
     data = r.json()
 
     results = []
-    for item in data.get("result", {}).get("results", []):
+    for result in data.get("result", {}).get("results", []):
         results.append({
-            "ui": item.get("ui"),  # CUI
-            "name": item.get("name"),
-            "rootSource": item.get("rootSource")
+            "ui": result.get("ui"),
+            "name": result.get("name"),
+            "rootSource": result.get("rootSource")
         })
     return results
 
 def get_cui_details(cui: str) -> Dict:
     """
-    Retrieve details for a specific CUI.
+    Retrieve detailed information about a specific CUI.
     """
     tgt = get_tgt()
-    st = get_service_ticket(tgt)
-    params = {"ticket": st}
-    r = requests.get(f"{UMLS_BASE}/content/current/CUI/{cui}", params=params)
+    ticket = get_service_ticket(tgt)
+
+    params = {"ticket": ticket}
+    r = requests.get(f"{UMLS_BASE}/rest/content/current/CUI/{cui}", params=params)
     r.raise_for_status()
     return r.json()
 
-def get_cui_relations(cui: str) -> List[Dict]:
+def get_related_concepts(cui: str, relation_type: str = "RO") -> List[Dict]:
     """
-    Get relationships for a CUI (e.g., associated diseases, drugs, symptoms).
+    Get related concepts for a given CUI.
+    relation_type: 'RO' (related other), 'RQ' (related and qualified), etc.
     """
     tgt = get_tgt()
-    st = get_service_ticket(tgt)
-    params = {"ticket": st}
-    r = requests.get(f"{UMLS_BASE}/content/current/CUI/{cui}/relations", params=params)
+    ticket = get_service_ticket(tgt)
+
+    params = {"ticket": ticket}
+    r = requests.get(
+        f"{UMLS_BASE}/rest/content/current/CUI/{cui}/relations",
+        params=params
+    )
     r.raise_for_status()
     data = r.json()
 
-    relations = []
+    related = []
     for rel in data.get("result", []):
-        relations.append({
-            "relationLabel": rel.get("relationLabel"),
-            "relatedId": rel.get("relatedId"),
-            "relatedName": rel.get("relatedIdName")
-        })
-    return relations
+        if rel.get("relationLabel") == relation_type:
+            related.append({
+                "cui": rel.get("relatedIdName"),
+                "name": rel.get("relatedIdName"),
+                "relation": rel.get("relationLabel")
+            })
+    return related
+
+def link_to_bioportal(cui: str) -> str:
+    """
+    Construct a BioPortal link for the given CUI.
+    """
+    return f"https://bioportal.bioontology.org/search?query={cui}"