Update genesis/api_clients/umls_api.py

genesis/api_clients/umls_api.py

@@ -1,105 +1,126 @@
 # genesis/api_clients/umls_api.py
 import os
 import requests
-from typing import Dict, List, Optional
+from typing import List, Dict
 
-UMLS_API_KEY = os.getenv("UMLS_API_KEY")
-UMLS_BASE = "https://uts-ws.nlm.nih.gov"
+UMLS_API_KEY = os.getenv("UMLS_API_KEY")  # Your UMLS API key from NIH UTS
+UMLS_AUTH_ENDPOINT = "https://utslogin.nlm.nih.gov/cas/v1/api-key"
+UMLS_BASE = "https://uts-ws.nlm.nih.gov/rest"
 
-if not UMLS_API_KEY:
-    raise EnvironmentError("Missing UMLS_API_KEY in environment variables.")
+session = requests.Session()
 
 # -------------------------
-# AUTHENTICATION
+# UMLS Authentication
 # -------------------------
-def get_tgt() -> str:
+def get_umls_ticket_granting_ticket() -> str:
     """
-    Get a Ticket Granting Ticket (TGT) from UMLS.
+    Get the Ticket Granting Ticket (TGT) for UMLS API authentication.
     """
-    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}")
+    if not UMLS_API_KEY:
+        raise ValueError("UMLS API key not found in environment variables.")
+
+    params = {"apikey": UMLS_API_KEY}
+    r = session.post(UMLS_AUTH_ENDPOINT, data=params)
+    r.raise_for_status()
+
+    # The TGT is returned in the Location header of the POST response
     return r.headers.get("location")
 
-def get_service_ticket(tgt: str) -> str:
+def get_umls_service_ticket(tgt: str) -> str:
     """
-    Get a single-use Service Ticket from the TGT.
+    Use TGT to request a one-time Service Ticket (ST).
     """
-    r = requests.post(tgt, data={"service": "http://umlsks.nlm.nih.gov"})
+    params = {"service": "http://umlsks.nlm.nih.gov"}
+    r = session.post(tgt, data=params)
     r.raise_for_status()
     return r.text
 
 # -------------------------
-# CORE FUNCTIONS
+# Search UMLS Concepts
 # -------------------------
-def search_cui(term: str, page_size: int = 10) -> List[Dict]:
+def search_umls(term: str, page_size: int = 5) -> List[Dict]:
     """
-    Search for UMLS CUIs for a given term.
+    Search the UMLS Metathesaurus for a given term and return related CUIs.
     """
-    tgt = get_tgt()
-    ticket = get_service_ticket(tgt)
+    tgt = get_umls_ticket_granting_ticket()
+    st = get_umls_service_ticket(tgt)
 
     params = {
         "string": term,
-        "ticket": ticket,
+        "ticket": st,
         "pageSize": page_size
     }
-    r = requests.get(f"{UMLS_BASE}/rest/search/current", params=params)
+    r = session.get(f"{UMLS_BASE}/search/current", params=params)
     r.raise_for_status()
-    data = r.json()
-
-    results = []
-    for result in data.get("result", {}).get("results", []):
-        results.append({
-            "ui": result.get("ui"),
-            "name": result.get("name"),
-            "rootSource": result.get("rootSource")
+
+    results = r.json().get("result", {}).get("results", [])
+    concepts = []
+    for res in results:
+        concepts.append({
+            "name": res.get("name"),
+            "cui": res.get("ui"),
+            "rootSource": res.get("rootSource")
         })
-    return results
+    return concepts
 
-def get_cui_details(cui: str) -> Dict:
+# -------------------------
+# Get UMLS Concept Details
+# -------------------------
+def get_concept_details(cui: str) -> Dict:
     """
-    Retrieve detailed information about a specific CUI.
+    Fetch full details for a UMLS concept (definitions, atoms, semantic types).
     """
-    tgt = get_tgt()
-    ticket = get_service_ticket(tgt)
+    tgt = get_umls_ticket_granting_ticket()
+    st = get_umls_service_ticket(tgt)
 
-    params = {"ticket": ticket}
-    r = requests.get(f"{UMLS_BASE}/rest/content/current/CUI/{cui}", params=params)
+    r = session.get(f"{UMLS_BASE}/content/current/CUI/{cui}", params={"ticket": st})
     r.raise_for_status()
-    return r.json()
+    return r.json().get("result", {})
 
-def get_related_concepts(cui: str, relation_type: str = "RO") -> List[Dict]:
+# -------------------------
+# Get Related Concepts
+# -------------------------
+def get_related_concepts(cui: str) -> List[Dict]:
     """
-    Get related concepts for a given CUI.
-    relation_type: 'RO' (related other), 'RQ' (related and qualified), etc.
+    Fetch related concepts for a given CUI.
     """
-    tgt = get_tgt()
-    ticket = get_service_ticket(tgt)
-
-    params = {"ticket": ticket}
-    r = requests.get(
-        f"{UMLS_BASE}/rest/content/current/CUI/{cui}/relations",
-        params=params
-    )
+    tgt = get_umls_ticket_granting_ticket()
+    st = get_umls_service_ticket(tgt)
+
+    r = session.get(f"{UMLS_BASE}/content/current/CUI/{cui}/relations", params={"ticket": st})
     r.raise_for_status()
-    data = r.json()
 
+    relations = r.json().get("result", [])
     related = []
-    for rel in data.get("result", []):
-        if rel.get("relationLabel") == relation_type:
-            related.append({
-                "cui": rel.get("relatedIdName"),
-                "name": rel.get("relatedIdName"),
-                "relation": rel.get("relationLabel")
-            })
+    for rel in relations:
+        related.append({
+            "relatedCUI": rel.get("relatedId"),
+            "relationLabel": rel.get("relationLabel"),
+            "relatedName": rel.get("relatedIdName")
+        })
     return related
 
-def link_to_bioportal(cui: str) -> str:
+# -------------------------
+# Map Term to Related Network
+# -------------------------
+def umls_network(term: str) -> Dict:
     """
-    Construct a BioPortal link for the given CUI.
+    Build a semantic network for a term: CUI, details, and related concepts.
     """
-    return f"https://bioportal.bioontology.org/search?query={cui}"
+    concepts = search_umls(term)
+    network = []
+
+    for concept in concepts:
+        cui = concept["cui"]
+        details = get_concept_details(cui)
+        related = get_related_concepts(cui)
+        network.append({
+            "concept": concept,
+            "details": details,
+            "related": related
+        })
+
+    return {
+        "term": term,
+        "network": network
+    }