Update genesis/api_clients/pubmed_api.py

genesis/api_clients/pubmed_api.py

@@ -1,29 +1,28 @@
 # genesis/api_clients/pubmed_api.py
 import requests
-import xml.etree.ElementTree as ET
-from typing import List, Dict, Optional
-from datetime import datetime
+import os
+from typing import List, Dict
 
+# NCBI E-utilities base
 PUBMED_SEARCH_URL = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi"
 PUBMED_FETCH_URL = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/efetch.fcgi"
-NCBI_API_KEY = None  # Optional: Set in Hugging Face secrets for higher rate limits
+PUBMED_SUMMARY_URL = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esummary.fcgi"
 
-def search_pubmed(query: str, max_results: int = 20, start_date: Optional[str] = None, end_date: Optional[str] = None) -> List[str]:
+# Optional: NCBI API key for higher rate limits
+NCBI_API_KEY = os.getenv("NCBI_API_KEY")
+
+def search_pubmed(query: str, max_results: int = 10) -> List[str]:
     """
-    Search PubMed for a given query and return a list of PubMed IDs.
-    Optionally filter by start_date and end_date (YYYY/MM/DD format).
+    Search PubMed for a query and return a list of PubMed IDs (PMIDs).
     """
     params = {
         "db": "pubmed",
         "term": query,
-        "retmax": max_results,
         "retmode": "json",
-        "api_key": NCBI_API_KEY
+        "retmax": max_results
     }
-    if start_date and end_date:
-        params["mindate"] = start_date
-        params["maxdate"] = end_date
-        params["datetype"] = "pdat"
+    if NCBI_API_KEY:
+        params["api_key"] = NCBI_API_KEY
 
     r = requests.get(PUBMED_SEARCH_URL, params=params)
     r.raise_for_status()
@@ -31,64 +30,61 @@ def search_pubmed(query: str, max_results: int = 20, start_date: Optional[str] =
     return data.get("esearchresult", {}).get("idlist", [])
 
 
-def fetch_pubmed_details(pmid_list: List[str]) -> List[Dict]:
+def fetch_pubmed_details(pmids: List[str]) -> List[Dict]:
     """
-    Fetch detailed metadata for a list of PubMed IDs.
-    Returns title, abstract, authors, journal, and publication date.
+    Fetch details (title, authors, journal, abstract) for a list of PMIDs.
     """
-    if not pmid_list:
+    if not pmids:
         return []
 
     params = {
         "db": "pubmed",
-        "id": ",".join(pmid_list),
-        "retmode": "xml",
-        "api_key": NCBI_API_KEY
+        "id": ",".join(pmids),
+        "retmode": "xml"
     }
+    if NCBI_API_KEY:
+        params["api_key"] = NCBI_API_KEY
+
     r = requests.get(PUBMED_FETCH_URL, params=params)
     r.raise_for_status()
+    return parse_pubmed_xml(r.text)
 
-    root = ET.fromstring(r.text)
-    results = []
 
+def parse_pubmed_xml(xml_text: str) -> List[Dict]:
+    """
+    Parse PubMed XML into structured data.
+    """
+    import xml.etree.ElementTree as ET
+    root = ET.fromstring(xml_text)
+    articles = []
+    
     for article in root.findall(".//PubmedArticle"):
-        try:
-            title = article.find(".//ArticleTitle").text or "No title"
-            abstract = " ".join([t.text for t in article.findall(".//AbstractText") if t.text]) or "No abstract"
-            authors = []
-            for a in article.findall(".//Author"):
-                last = a.findtext("LastName", "")
-                first = a.findtext("ForeName", "")
-                if last or first:
-                    authors.append(f"{first} {last}".strip())
-
-            journal = article.findtext(".//Journal/Title", "Unknown Journal")
-            pub_date = article.find(".//PubDate")
-            if pub_date is not None:
-                year = pub_date.findtext("Year", "")
-                month = pub_date.findtext("Month", "")
-                day = pub_date.findtext("Day", "")
-                date_str = f"{year}-{month}-{day}" if year else "Unknown"
-            else:
-                date_str = "Unknown"
-
-            results.append({
-                "title": title,
-                "abstract": abstract,
-                "authors": authors,
-                "journal": journal,
-                "publication_date": date_str,
-                "pubmed_link": f"https://pubmed.ncbi.nlm.nih.gov/{article.findtext('.//PMID')}/"
-            })
-        except Exception:
-            continue
-
-    return results
+        title = article.findtext(".//ArticleTitle", default="")
+        abstract_text = " ".join([abst.text or "" for abst in article.findall(".//AbstractText")])
+        journal = article.findtext(".//Journal/Title", default="")
+        authors = []
+        for author in article.findall(".//Author"):
+            last = author.findtext("LastName", "")
+            fore = author.findtext("ForeName", "")
+            if last or fore:
+                authors.append(f"{fore} {last}")
+        
+        pmid = article.findtext(".//PMID", default="")
+        
+        articles.append({
+            "pmid": pmid,
+            "title": title,
+            "abstract": abstract_text,
+            "journal": journal,
+            "authors": authors
+        })
+    
+    return articles
 
 
-def search_and_fetch_pubmed(query: str, max_results: int = 20, start_date: Optional[str] = None, end_date: Optional[str] = None) -> List[Dict]:
+def pubmed_summary(query: str, max_results: int = 10) -> List[Dict]:
     """
-    Search and fetch PubMed results in one call.
+    Search PubMed and return summarized results in a single step.
     """
-    pmids = search_pubmed(query, max_results, start_date, end_date)
+    pmids = search_pubmed(query, max_results)
     return fetch_pubmed_details(pmids)