Update genesis/api_clients/pubmed_api.py

genesis/api_clients/pubmed_api.py

@@ -1,102 +1,94 @@
 # genesis/api_clients/pubmed_api.py
 import requests
 import xml.etree.ElementTree as ET
+from typing import List, Dict, Optional
 from datetime import datetime
 
-NCBI_BASE_URL = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils"
-EMAIL = "your_email@example.com"  # NCBI requires an email in requests
+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
 
-
-def search_pubmed(query: str, max_results: int = 10):
+def search_pubmed(query: str, max_results: int = 20, start_date: Optional[str] = None, end_date: Optional[str] = None) -> List[str]:
     """
-    Search PubMed for a given query and return PubMed IDs (PMIDs).
+    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).
     """
     params = {
         "db": "pubmed",
         "term": query,
         "retmax": max_results,
         "retmode": "json",
-        "sort": "pub+date",
-        "email": EMAIL
+        "api_key": NCBI_API_KEY
     }
-    response = requests.get(f"{NCBI_BASE_URL}/esearch.fcgi", params=params)
-    response.raise_for_status()
-    data = response.json()
+    if start_date and end_date:
+        params["mindate"] = start_date
+        params["maxdate"] = end_date
+        params["datetype"] = "pdat"
+
+    r = requests.get(PUBMED_SEARCH_URL, params=params)
+    r.raise_for_status()
+    data = r.json()
     return data.get("esearchresult", {}).get("idlist", [])
 
 
-def fetch_pubmed_details(pmids: list):
+def fetch_pubmed_details(pmid_list: List[str]) -> List[Dict]:
     """
-    Retrieve details (title, abstract, authors, etc.) for a list of PMIDs.
+    Fetch detailed metadata for a list of PubMed IDs.
+    Returns title, abstract, authors, journal, and publication date.
     """
-    if not pmids:
+    if not pmid_list:
         return []
 
     params = {
         "db": "pubmed",
-        "id": ",".join(pmids),
+        "id": ",".join(pmid_list),
         "retmode": "xml",
-        "email": EMAIL
+        "api_key": NCBI_API_KEY
     }
-    response = requests.get(f"{NCBI_BASE_URL}/efetch.fcgi", params=params)
-    response.raise_for_status()
+    r = requests.get(PUBMED_FETCH_URL, params=params)
+    r.raise_for_status()
 
-    root = ET.fromstring(response.text)
-    articles = []
+    root = ET.fromstring(r.text)
+    results = []
 
     for article in root.findall(".//PubmedArticle"):
-        title = article.findtext(".//ArticleTitle", default="No title available")
-        abstract = " ".join(
-            [a.text for a in article.findall(".//AbstractText") if a.text]
-        )
-        pub_date = extract_pub_date(article)
-        authors = extract_authors(article)
-        pmid = article.findtext(".//PMID", default="")
-
-        articles.append({
-            "pmid": pmid,
-            "title": title,
-            "abstract": abstract,
-            "authors": authors,
-            "pub_date": pub_date,
-            "url": f"https://pubmed.ncbi.nlm.nih.gov/{pmid}/"
-        })
-
-    return articles
-
-
-def extract_pub_date(article_element):
-    """
-    Extract publication date from a PubMed XML element.
-    """
-    date_elem = article_element.find(".//PubDate")
-    if date_elem is not None:
-        year = date_elem.findtext("Year")
-        month = date_elem.findtext("Month", default="01")
-        day = date_elem.findtext("Day", default="01")
         try:
-            return datetime(int(year), int(month), int(day)).strftime("%Y-%m-%d")
-        except:
-            return year
-    return "Unknown"
-
-
-def extract_authors(article_element):
-    """
-    Extract authors from a PubMed XML element.
-    """
-    authors = []
-    for author in article_element.findall(".//Author"):
-        last_name = author.findtext("LastName", "")
-        fore_name = author.findtext("ForeName", "")
-        if last_name and fore_name:
-            authors.append(f"{fore_name} {last_name}")
-    return authors
-
-
-def search_and_fetch(query: str, max_results: int = 5):
+            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
+
+
+def search_and_fetch_pubmed(query: str, max_results: int = 20, start_date: Optional[str] = None, end_date: Optional[str] = None) -> List[Dict]:
     """
-    Convenience function to search PubMed and fetch article details.
+    Search and fetch PubMed results in one call.
     """
-    pmids = search_pubmed(query, max_results)
+    pmids = search_pubmed(query, max_results, start_date, end_date)
     return fetch_pubmed_details(pmids)