Update genesis/ontology.py

genesis/ontology.py

@@ -1,42 +1,44 @@
 # genesis/ontology.py
-import requests
-from typing import List
+"""
+Ontology Expansion for GENESIS-AI
+Uses UMLS & BioPortal to find related terms, synonyms, and hierarchies.
+"""
 
-def expand_terms_with_ontology(query: str, umls_api_key: str, bioportal_api_key: str) -> List[str]:
-    """Expand query terms using UMLS and BioPortal ontology APIs."""
-    expanded_terms = set()
+import os
+import requests
 
-    # Always include original query
-    expanded_terms.add(query)
+UMLS_API_KEY = os.getenv("UMLS_API_KEY")
+BIOPORTAL_API_KEY = os.getenv("BIOPORTAL_API_KEY")
 
-    # UMLS expansion
-    if umls_api_key:
-        try:
-            umls_url = "https://uts-ws.nlm.nih.gov/rest/search/current"
-            params = {"string": query, "apiKey": umls_api_key, "pageSize": 5}
-            r = requests.get(umls_url, params=params, timeout=10)
-            r.raise_for_status()
-            data = r.json()
-            for result in data.get("result", {}).get("results", []):
-                name = result.get("name")
-                if name and name.lower() != query.lower():
-                    expanded_terms.add(name)
-        except Exception as e:
-            print(f"[Ontology] UMLS expansion failed: {e}")
+def expand_with_bioportal(term):
+    """Expand term using BioPortal ontology search."""
+    try:
+        url = f"https://data.bioontology.org/search?q={term}&apikey={BIOPORTAL_API_KEY}"
+        r = requests.get(url)
+        r.raise_for_status()
+        data = r.json()
+        return list({res["prefLabel"] for res in data.get("collection", []) if "prefLabel" in res})
+    except Exception as e:
+        return {"error": str(e)}
 
-    # BioPortal expansion
-    if bioportal_api_key:
-        try:
-            bp_url = "https://data.bioontology.org/search"
-            params = {"q": query, "apikey": bioportal_api_key}
-            r = requests.get(bp_url, params=params, timeout=10)
-            r.raise_for_status()
-            data = r.json()
-            for coll in data.get("collection", []):
-                label = coll.get("prefLabel")
-                if label and label.lower() != query.lower():
-                    expanded_terms.add(label)
-        except Exception as e:
-            print(f"[Ontology] BioPortal expansion failed: {e}")
+def expand_with_umls(term):
+    """Expand term using UMLS API."""
+    try:
+        url = f"https://uts-ws.nlm.nih.gov/rest/search/current?string={term}&apiKey={UMLS_API_KEY}"
+        r = requests.get(url)
+        r.raise_for_status()
+        data = r.json()
+        return [res["name"] for res in data.get("result", {}).get("results", []) if "name" in res]
+    except Exception as e:
+        return {"error": str(e)}
 
-    return list(expanded_terms)
+def expand_terms_with_ontology(term):
+    """Combine BioPortal and UMLS expansions."""
+    bioportal_terms = expand_with_bioportal(term)
+    umls_terms = expand_with_umls(term)
+    results = set()
+    if isinstance(bioportal_terms, list):
+        results.update(bioportal_terms)
+    if isinstance(umls_terms, list):
+        results.update(umls_terms)
+    return list(results) if results else [term]