Update genesis/api_clients/umls_api.py

genesis/api_clients/umls_api.py

@@ -1,86 +1,89 @@
 # genesis/api_clients/umls_api.py
 import os
 import requests
-from typing import Dict, List, Optional
+from typing import List, Dict, Optional
 
-# UMLS API
+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")  # Add this to your Hugging Face secrets or .env
 
-if not UMLS_API_KEY:
-    raise EnvironmentError("UMLS_API_KEY is missing. Add it to your environment variables or Hugging Face secrets.")
-
-# UMLS requires ticket-granting service for authentication
+# -------------------------
+# AUTHENTICATION
+# -------------------------
 def get_tgt() -> str:
     """
-    Obtain a Ticket Granting Ticket (TGT) for UMLS authentication.
+    Get Ticket Granting Ticket (TGT) from UMLS API key.
     """
-    r = requests.post("https://utslogin.nlm.nih.gov/cas/v1/api-key", data={"apikey": UMLS_API_KEY})
+    params = {"apikey": UMLS_API_KEY}
+    r = requests.post(UMLS_AUTH_ENDPOINT, data=params)
     r.raise_for_status()
-    from xml.etree import ElementTree
-    root = ElementTree.fromstring(r.text)
-    return root.attrib.get("action")
-
+    # Extract TGT from HTML form action
+    tgt_url = r.text.split('action="')[1].split('"')[0]
+    return tgt_url
 
-def get_service_ticket(tgt: str) -> str:
+def get_service_ticket(tgt_url: str) -> str:
     """
-    Get a single-use service ticket for API requests.
+    Get a one-time Service Ticket (ST) for API calls.
     """
-    r = requests.post(tgt, data={"service": "http://umlsks.nlm.nih.gov"})
+    params = {"service": "http://umlsks.nlm.nih.gov"}
+    r = requests.post(tgt_url, data=params)
     r.raise_for_status()
     return r.text
 
-
-def search_umls(query: str, max_results: int = 10) -> List[Dict]:
+# -------------------------
+# SEARCH FUNCTIONS
+# -------------------------
+def search_umls(term: str, max_results: int = 10) -> List[Dict]:
     """
-    Search UMLS Metathesaurus for a term and return CUIs and names.
+    Search UMLS for a medical concept by term and return CUIs, preferred names, and semantic types.
     """
     tgt = get_tgt()
-    ticket = get_service_ticket(tgt)
-    params = {"string": query, "ticket": ticket, "pageSize": max_results}
-    
+    st = get_service_ticket(tgt)
+    params = {
+        "string": term,
+        "ticket": st,
+        "pageSize": max_results
+    }
     r = requests.get(f"{UMLS_BASE}/search/current", params=params)
     r.raise_for_status()
-    results = r.json().get("result", {}).get("results", [])
-    
-    return [
-        {
-            "ui": item.get("ui"),
+    data = r.json()
+
+    results = []
+    for item in data.get("result", {}).get("results", []):
+        results.append({
+            "ui": item.get("ui"),  # CUI
             "name": item.get("name"),
             "rootSource": item.get("rootSource")
-        }
-        for item in results if item.get("ui") != "NONE"
-    ]
-
+        })
+    return results
 
-def get_concept_details(cui: str) -> Dict:
+def get_cui_details(cui: str) -> Dict:
     """
-    Get full details for a given CUI (Concept Unique Identifier).
-    Includes synonyms, definitions, and mappings.
+    Retrieve details for a specific CUI.
     """
     tgt = get_tgt()
-    ticket = get_service_ticket(tgt)
-    
-    r = requests.get(f"{UMLS_BASE}/content/current/CUI/{cui}", params={"ticket": ticket})
+    st = get_service_ticket(tgt)
+    params = {"ticket": st}
+    r = requests.get(f"{UMLS_BASE}/content/current/CUI/{cui}", params=params)
     r.raise_for_status()
-    return r.json().get("result", {})
+    return r.json()
 
-
-def map_between_ontologies(cui: str) -> List[Dict]:
+def get_cui_relations(cui: str) -> List[Dict]:
     """
-    Given a CUI, return cross-ontology mappings (e.g., SNOMED → MeSH → ICD).
+    Get relationships for a CUI (e.g., associated diseases, drugs, symptoms).
     """
     tgt = get_tgt()
-    ticket = get_service_ticket(tgt)
-    
-    r = requests.get(f"{UMLS_BASE}/crosswalk/current/source/{cui}", params={"ticket": ticket})
+    st = get_service_ticket(tgt)
+    params = {"ticket": st}
+    r = requests.get(f"{UMLS_BASE}/content/current/CUI/{cui}/relations", params=params)
     r.raise_for_status()
-    return r.json().get("result", [])
-
+    data = r.json()
 
-def normalize_term(term: str) -> Optional[Dict]:
-    """
-    Normalize a free-text biomedical term to its canonical CUI and return metadata.
-    """
-    matches = search_umls(term, max_results=1)
-    return matches[0] if matches else None
+    relations = []
+    for rel in data.get("result", []):
+        relations.append({
+            "relationLabel": rel.get("relationLabel"),
+            "relatedId": rel.get("relatedId"),
+            "relatedName": rel.get("relatedIdName")
+        })
+    return relations