Update genesis/api_clients/umls_api.py

genesis/api_clients/umls_api.py

@@ -1,7 +1,7 @@
 # genesis/api_clients/umls_api.py
 import os
 import requests
-from typing import List, Dict
+from typing import List, Dict, Optional
 
 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"
@@ -10,7 +10,7 @@ UMLS_BASE = "https://uts-ws.nlm.nih.gov/rest"
 session = requests.Session()
 
 # -------------------------
-# UMLS Authentication
+# Authentication Helpers
 # -------------------------
 def get_umls_ticket_granting_ticket() -> str:
     """
@@ -23,8 +23,18 @@ def get_umls_ticket_granting_ticket() -> str:
     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")
+    # Extract TGT location from XML
+    if "location" in r.headers:
+        return r.headers["location"]
+
+    # Fallback: parse XML body
+    from xml.etree import ElementTree as ET
+    root = ET.fromstring(r.text)
+    form = root.find(".//form")
+    if form is not None and "action" in form.attrib:
+        return form.attrib["action"]
+
+    raise RuntimeError("Unable to retrieve UMLS TGT. Check API key or response format.")
 
 def get_umls_service_ticket(tgt: str) -> str:
     """
@@ -33,14 +43,15 @@ def get_umls_service_ticket(tgt: str) -> str:
     params = {"service": "http://umlsks.nlm.nih.gov"}
     r = session.post(tgt, data=params)
     r.raise_for_status()
-    return r.text
+    return r.text.strip()
 
 # -------------------------
-# Search UMLS Concepts
+# Search Concepts
 # -------------------------
 def search_umls(term: str, page_size: int = 5) -> List[Dict]:
     """
-    Search the UMLS Metathesaurus for a given term and return related CUIs.
+    Search the UMLS Metathesaurus for a given term.
+    Returns a list of concepts with CUI and metadata.
     """
     tgt = get_umls_ticket_granting_ticket()
     st = get_umls_service_ticket(tgt)
@@ -56,19 +67,20 @@ def search_umls(term: str, page_size: int = 5) -> List[Dict]:
     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")
-        })
+        if res.get("ui") != "NONE":
+            concepts.append({
+                "name": res.get("name"),
+                "cui": res.get("ui"),
+                "rootSource": res.get("rootSource")
+            })
     return concepts
 
 # -------------------------
-# Get UMLS Concept Details
+# Concept Details
 # -------------------------
 def get_concept_details(cui: str) -> Dict:
     """
-    Fetch full details for a UMLS concept (definitions, atoms, semantic types).
+    Fetch details for a UMLS CUI (definitions, atoms, semantic types).
     """
     tgt = get_umls_ticket_granting_ticket()
     st = get_umls_service_ticket(tgt)
@@ -78,7 +90,7 @@ def get_concept_details(cui: str) -> Dict:
     return r.json().get("result", {})
 
 # -------------------------
-# Get Related Concepts
+# Related Concepts
 # -------------------------
 def get_related_concepts(cui: str) -> List[Dict]:
     """
@@ -101,11 +113,11 @@ def get_related_concepts(cui: str) -> List[Dict]:
     return related
 
 # -------------------------
-# Map Term to Related Network
+# Term → Semantic Network
 # -------------------------
 def umls_network(term: str) -> Dict:
     """
-    Build a semantic network for a term: CUI, details, and related concepts.
+    Build a semantic network for a term: concepts, details, and related CUIs.
     """
     concepts = search_umls(term)
     network = []
@@ -124,3 +136,22 @@ def umls_network(term: str) -> Dict:
         "term": term,
         "network": network
     }
+
+# -------------------------
+# Expansion Helper for pipeline.py
+# -------------------------
+def expand_with_umls(term: str, max_results: int = 5) -> List[str]:
+    """
+    Expand a biomedical term into related synonyms & preferred names from UMLS.
+    Used by pipeline.py for ontology enrichment.
+    """
+    try:
+        concepts = search_umls(term, page_size=max_results)
+        expanded = []
+        for c in concepts:
+            if c["name"] and c["name"].lower() != term.lower():
+                expanded.append(c["name"])
+        return list(set(expanded))
+    except Exception as e:
+        print(f"[UMLS] Expansion failed for '{term}': {e}")
+        return []