Update genesis/api_clients/pubmed_api.py

genesis/api_clients/pubmed_api.py

@@ -1,77 +1,92 @@
 # genesis/api_clients/pubmed_api.py
 import os
 import requests
-from typing import List, Dict, Optional
-from xml.etree import ElementTree as ET
+import xml.etree.ElementTree as ET
+from typing import List, Dict
 
-NCBI_API_KEY = os.getenv("NCBI_API_KEY")  # Optional but increases rate limits
 NCBI_BASE = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils"
+PUBMED_API_KEY = os.getenv("PUBMED_API_KEY")  # Optional, set in Hugging Face / .env
 
-def search_pubmed(query: str, max_results: int = 10) -> List[str]:
+# -------------------------
+# Core Helpers
+# -------------------------
+def pubmed_search(query: str, max_results: int = 10) -> List[str]:
     """
-    Search PubMed and return a list of PMIDs.
+    Search PubMed for a query and return a list of PubMed IDs.
     """
     params = {
         "db": "pubmed",
         "term": query,
         "retmax": max_results,
-        "api_key": NCBI_API_KEY
+        "api_key": PUBMED_API_KEY
     }
     r = requests.get(f"{NCBI_BASE}/esearch.fcgi", params=params)
     r.raise_for_status()
     root = ET.fromstring(r.text)
-    return [id_tag.text for id_tag in root.findall(".//Id")]
+    return [id_elem.text for id_elem in root.findall(".//Id")]
 
-def fetch_pubmed_details(pmids: List[str]) -> List[Dict]:
+def fetch_details(pubmed_ids: List[str]) -> List[Dict]:
     """
-    Fetch detailed information for a list of PMIDs.
+    Fetch detailed metadata for PubMed articles.
     """
-    if not pmids:
+    if not pubmed_ids:
         return []
-    
+
     params = {
         "db": "pubmed",
-        "id": ",".join(pmids),
+        "id": ",".join(pubmed_ids),
         "retmode": "xml",
-        "api_key": NCBI_API_KEY
+        "api_key": PUBMED_API_KEY
     }
     r = requests.get(f"{NCBI_BASE}/efetch.fcgi", params=params)
     r.raise_for_status()
-    
-    root = ET.fromstring(r.text)
-    articles = []
 
+    articles = []
+    root = ET.fromstring(r.text)
     for article in root.findall(".//PubmedArticle"):
-        title = article.findtext(".//ArticleTitle", default="No title")
-        abstract = " ".join([t.text for t in article.findall(".//AbstractText") if t.text])
-        journal = article.findtext(".//Title", default="Unknown Journal")
-        pub_date = article.findtext(".//PubDate/Year", default="Unknown Year")
-        doi = None
-        for id_tag in article.findall(".//ArticleId"):
-            if id_tag.attrib.get("IdType") == "doi":
-                doi = id_tag.text
-        authors = []
-        for author in article.findall(".//Author"):
-            last = author.findtext("LastName")
-            fore = author.findtext("ForeName")
-            if last and fore:
-                authors.append(f"{fore} {last}")
-        pmid = article.findtext(".//PMID")
+        title_elem = article.find(".//ArticleTitle")
+        abstract_elem = article.find(".//Abstract/AbstractText")
+        pmid_elem = article.find(".//PMID")
+        authors = [
+            f"{a.find('LastName').text} {a.find('ForeName').text}"
+            for a in article.findall(".//Author")
+            if a.find("LastName") is not None and a.find("ForeName") is not None
+        ]
+
         articles.append({
-            "pmid": pmid,
-            "title": title,
-            "abstract": abstract,
-            "journal": journal,
-            "pub_date": pub_date,
-            "doi": doi,
+            "pmid": pmid_elem.text if pmid_elem is not None else "",
+            "title": title_elem.text if title_elem is not None else "",
+            "abstract": abstract_elem.text if abstract_elem is not None else "",
             "authors": authors,
-            "url": f"https://pubmed.ncbi.nlm.nih.gov/{pmid}/"
+            "url": f"https://pubmed.ncbi.nlm.nih.gov/{pmid_elem.text}/" if pmid_elem is not None else ""
         })
     return articles
 
-def search_and_fetch(query: str, max_results: int = 10) -> List[Dict]:
+# -------------------------
+# High-Level Search + Fetch
+# -------------------------
+def search_pubmed(query: str, max_results: int = 10) -> List[Dict]:
     """
-    Convenience function: Search and fetch results in one step.
+    Search and return structured PubMed results.
     """
-    pmids = search_pubmed(query, max_results)
-    return fetch_pubmed_details(pmids)
+    ids = pubmed_search(query, max_results)
+    return fetch_details(ids)
+
+# -------------------------
+# Cross-Domain Integration
+# -------------------------
+def literature_to_entities(query: str) -> Dict:
+    """
+    Link PubMed literature to related drugs (ChEMBL), genes (NCBI), and ontology terms (BioPortal).
+    """
+    from genesis.api_clients import chembl_api, ncbi_api, bioportal_api  # Lazy import to avoid cycles
+
+    literature_data = {
+        "query": query,
+        "articles": search_pubmed(query),
+        "related_drugs": chembl_api.search_molecule(query),
+        "related_genes": ncbi_api.search_gene(query),
+        "ontology_mappings": bioportal_api.search_and_map(query)
+    }
+
+    return literature_data