Update genesis/api_clients/umls_api.py

genesis/api_clients/umls_api.py

@@ -1,73 +1,94 @@
 # genesis/api_clients/umls_api.py
-import os
 import requests
+import os
+from urllib.parse import quote
 
-UMLS_API_KEY = os.getenv("UMLS_API_KEY")
-UMLS_AUTH_URL = "https://utslogin.nlm.nih.gov"
-UMLS_API_URL = "https://uts-ws.nlm.nih.gov/rest"
+UMLS_API_KEY = os.getenv("UMLS_API_KEY")  # Store in Hugging Face secrets
+UMLS_BASE_URL = "https://uts-ws.nlm.nih.gov/rest"
 
 if not UMLS_API_KEY:
-    raise ValueError("Missing UMLS_API_KEY in environment variables")
+    raise ValueError("UMLS_API_KEY not found in environment variables.")
 
-def get_tgt():
+# Get UMLS authentication ticket
+def get_umls_ticket():
     """
-    Get a Ticket Granting Ticket (TGT) for UMLS authentication.
+    Retrieve a UMLS authentication ticket.
     """
-    url = f"{UMLS_AUTH_URL}/cas/v1/api-key"
-    params = {"apikey": UMLS_API_KEY}
-    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]
+    auth_url = f"https://utslogin.nlm.nih.gov/cas/v1/api-key"
+    response = requests.post(auth_url, data={"apikey": UMLS_API_KEY})
+    response.raise_for_status()
+    return response.text.split("<form")[0].strip()
+
 
-def get_service_ticket(tgt):
+def search_umls(term: str, page_size: int = 10):
     """
-    Get a Service Ticket (ST) using TGT.
+    Search UMLS for a term (CUI, concept name, or synonym).
     """
-    params = {"service": "http://umlsks.nlm.nih.gov"}
-    res = requests.post(tgt, data=params)
-    res.raise_for_status()
-    return res.text
+    ticket = get_umls_ticket()
+    encoded_term = quote(term)
+    url = f"{UMLS_BASE_URL}/search/current?string={encoded_term}&ticket={ticket}&pageSize={page_size}"
+    response = requests.get(url)
+    response.raise_for_status()
+    results = response.json().get("result", {}).get("results", [])
+    return [
+        {
+            "ui": r.get("ui"),
+            "name": r.get("name"),
+            "rootSource": r.get("rootSource"),
+        }
+        for r in results
+        if r.get("ui") != "NONE"
+    ]
 
-def search_umls(query: str, page: int = 1, page_size: int = 10):
+
+def get_concept_details(cui: str):
     """
-    Search for a term in UMLS.
+    Retrieve details for a specific UMLS concept (CUI).
     """
-    tgt = get_tgt()
-    ticket = get_service_ticket(tgt)
-    url = f"{UMLS_API_URL}/search/current"
-    params = {
-        "string": query,
-        "ticket": ticket,
-        "pageNumber": page,
-        "pageSize": page_size
-    }
-    res = requests.get(url, params=params)
-    res.raise_for_status()
-    return res.json().get("result", {}).get("results", [])
+    ticket = get_umls_ticket()
+    url = f"{UMLS_BASE_URL}/content/current/CUI/{cui}?ticket={ticket}"
+    response = requests.get(url)
+    response.raise_for_status()
+    return response.json().get("result", {})
+
 
-def get_cui_details(cui: str):
+def get_concept_synonyms(cui: str):
     """
-    Get detailed info for a CUI (Concept Unique Identifier).
+    Get synonyms for a given UMLS concept.
     """
-    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()
+    concept = get_concept_details(cui)
+    synonyms = []
+    for atom in concept.get("atoms", []):
+        synonyms.append(atom.get("name"))
+    return list(set(synonyms))
 
-def map_to_umls(term: str):
+
+def get_semantic_types(cui: str):
     """
-    Given a term, return the matching UMLS CUI(s) and definitions.
+    Retrieve semantic types for a UMLS concept.
     """
-    matches = search_umls(term)
-    return [
-        {
-            "name": m.get("name"),
-            "cui": m.get("ui"),
-            "rootSource": m.get("rootSource")
+    ticket = get_umls_ticket()
+    url = f"{UMLS_BASE_URL}/content/current/CUI/{cui}/definitions?ticket={ticket}"
+    response = requests.get(url)
+    response.raise_for_status()
+    definitions = response.json().get("result", [])
+    return [d.get("rootSource") for d in definitions]
+
+
+def cross_map_term(term: str):
+    """
+    Map a term across all major UMLS vocabularies.
+    """
+    mappings = {}
+    search_results = search_umls(term, page_size=5)
+    for result in search_results:
+        cui = result["ui"]
+        synonyms = get_concept_synonyms(cui)
+        semantic_types = get_semantic_types(cui)
+        mappings[cui] = {
+            "name": result["name"],
+            "rootSource": result["rootSource"],
+            "synonyms": synonyms,
+            "semantic_types": semantic_types,
         }
-        for m in matches if m.get("ui") != "NONE"
-    ]
+    return mappings