Update genesis/api_clients/umls_api.py

genesis/api_clients/umls_api.py

@@ -1,93 +1,86 @@
 # genesis/api_clients/umls_api.py
 import os
 import requests
+from typing import Dict, List, Optional
 
-UMLS_API_KEY = os.getenv("UMLS_API_KEY")
-UMLS_AUTH_ENDPOINT = "https://utslogin.nlm.nih.gov/cas/v1/api-key"
-UMLS_BASE_URL = "https://uts-ws.nlm.nih.gov/rest"
+# UMLS API
+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.")
 
-def get_tgt():
+# UMLS requires ticket-granting service for authentication
+def get_tgt() -> str:
     """
-    Get Ticket-Granting Ticket (TGT) for UMLS API authentication.
+    Obtain a Ticket Granting Ticket (TGT) for UMLS authentication.
     """
-    params = {"apikey": UMLS_API_KEY}
-    r = requests.post(UMLS_AUTH_ENDPOINT, data=params)
+    r = requests.post("https://utslogin.nlm.nih.gov/cas/v1/api-key", data={"apikey": UMLS_API_KEY})
     r.raise_for_status()
+    from xml.etree import ElementTree
+    root = ElementTree.fromstring(r.text)
+    return root.attrib.get("action")
 
-    # Extract TGT from HTML response
-    tgt_url = r.text.split('action="')[1].split('"')[0]
-    return tgt_url
 
-
-def get_service_ticket(tgt):
+def get_service_ticket(tgt: str) -> str:
     """
-    Get Service Ticket (ST) from the TGT for making API requests.
+    Get a single-use service ticket for API requests.
     """
-    params = {"service": "http://umlsks.nlm.nih.gov"}
-    r = requests.post(tgt, data=params)
+    r = requests.post(tgt, data={"service": "http://umlsks.nlm.nih.gov"})
     r.raise_for_status()
     return r.text
 
 
-def search_umls(term: str, max_results: int = 10):
+def search_umls(query: str, max_results: int = 10) -> List[Dict]:
     """
-    Search UMLS for a term and return matching concepts.
+    Search UMLS Metathesaurus for a term and return CUIs and names.
     """
     tgt = get_tgt()
-    st = get_service_ticket(tgt)
-
-    url = f"{UMLS_BASE_URL}/search/current"
-    params = {"string": term, "ticket": st, "pageSize": max_results}
-    r = requests.get(url, params=params)
+    ticket = get_service_ticket(tgt)
+    params = {"string": query, "ticket": ticket, "pageSize": max_results}
+    
+    r = requests.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", "")
-        })
-
-    return results
-
-
-def get_concept_details(cui: str):
+    results = r.json().get("result", {}).get("results", [])
+    
+    return [
+        {
+            "ui": item.get("ui"),
+            "name": item.get("name"),
+            "rootSource": item.get("rootSource")
+        }
+        for item in results if item.get("ui") != "NONE"
+    ]
+
+
+def get_concept_details(cui: str) -> Dict:
     """
-    Retrieve detailed information for a concept by CUI (Concept Unique Identifier).
+    Get full details for a given CUI (Concept Unique Identifier).
+    Includes synonyms, definitions, and mappings.
     """
     tgt = get_tgt()
-    st = get_service_ticket(tgt)
-
-    url = f"{UMLS_BASE_URL}/content/current/CUI/{cui}"
-    params = {"ticket": st}
-    r = requests.get(url, params=params)
+    ticket = get_service_ticket(tgt)
+    
+    r = requests.get(f"{UMLS_BASE}/content/current/CUI/{cui}", params={"ticket": ticket})
     r.raise_for_status()
-    return r.json()
+    return r.json().get("result", {})
 
 
-def map_related_concepts(term: str):
+def map_between_ontologies(cui: str) -> List[Dict]:
     """
-    Search UMLS and map related concepts with synonyms and semantic types.
+    Given a CUI, return cross-ontology mappings (e.g., SNOMED → MeSH → ICD).
     """
-    concepts = search_umls(term, max_results=5)
-    mapped = []
-
-    for c in concepts:
-        cui = c["ui"]
-        if cui and cui != "NONE":
-            details = get_concept_details(cui)
-            entity = details.get("result", {})
+    tgt = get_tgt()
+    ticket = get_service_ticket(tgt)
+    
+    r = requests.get(f"{UMLS_BASE}/crosswalk/current/source/{cui}", params={"ticket": ticket})
+    r.raise_for_status()
+    return r.json().get("result", [])
 
-            mapped.append({
-                "cui": cui,
-                "name": entity.get("name", ""),
-                "synonyms": [syn.get("name", "") for syn in entity.get("atoms", [])],
-                "semantic_types": [st.get("name", "") for st in entity.get("semanticTypes", [])],
-                "rootSource": c.get("rootSource", "")
-            })
 
-    return mapped
+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