Update genesis/api_clients/umls_api.py

genesis/api_clients/umls_api.py

@@ -1,94 +1,93 @@
 # genesis/api_clients/umls_api.py
-import requests
 import os
-from urllib.parse import quote
+import requests
 
-UMLS_API_KEY = os.getenv("UMLS_API_KEY")  # Store in Hugging Face secrets
+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"
 
-if not UMLS_API_KEY:
-    raise ValueError("UMLS_API_KEY not found in environment variables.")
 
-# Get UMLS authentication ticket
-def get_umls_ticket():
+def get_tgt():
     """
-    Retrieve a UMLS authentication ticket.
+    Get Ticket-Granting Ticket (TGT) for UMLS API authentication.
     """
-    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()
+    params = {"apikey": UMLS_API_KEY}
+    r = requests.post(UMLS_AUTH_ENDPOINT, data=params)
+    r.raise_for_status()
 
+    # Extract TGT from HTML response
+    tgt_url = r.text.split('action="')[1].split('"')[0]
+    return tgt_url
 
-def search_umls(term: str, page_size: int = 10):
+
+def get_service_ticket(tgt):
     """
-    Search UMLS for a term (CUI, concept name, or synonym).
+    Get Service Ticket (ST) from the TGT for making API requests.
     """
-    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"
-    ]
+    params = {"service": "http://umlsks.nlm.nih.gov"}
+    r = requests.post(tgt, data=params)
+    r.raise_for_status()
+    return r.text
 
 
-def get_concept_details(cui: str):
+def search_umls(term: str, max_results: int = 10):
     """
-    Retrieve details for a specific UMLS concept (CUI).
+    Search UMLS for a term and return matching concepts.
     """
-    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", {})
+    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)
+    r.raise_for_status()
 
-def get_concept_synonyms(cui: str):
-    """
-    Get synonyms for a given UMLS concept.
-    """
-    concept = get_concept_details(cui)
-    synonyms = []
-    for atom in concept.get("atoms", []):
-        synonyms.append(atom.get("name"))
-    return list(set(synonyms))
+    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_semantic_types(cui: str):
+
+def get_concept_details(cui: str):
     """
-    Retrieve semantic types for a UMLS concept.
+    Retrieve detailed information for a concept by CUI (Concept Unique Identifier).
     """
-    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]
+    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)
+    r.raise_for_status()
+    return r.json()
 
 
-def cross_map_term(term: str):
+def map_related_concepts(term: str):
     """
-    Map a term across all major UMLS vocabularies.
+    Search UMLS and map related concepts with synonyms and semantic types.
     """
-    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,
-        }
-    return mappings
+    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", {})
+
+            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