Update genesis/api_clients/pubmed_api.py

genesis/api_clients/pubmed_api.py

@@ -1,90 +1,77 @@
 # genesis/api_clients/pubmed_api.py
-import requests
 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"
-PUBMED_SUMMARY_URL = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esummary.fcgi"
+import requests
+from typing import List, Dict, Optional
+from xml.etree import ElementTree as ET
 
-# Optional: NCBI API key for higher rate limits
-NCBI_API_KEY = os.getenv("NCBI_API_KEY")
+NCBI_API_KEY = os.getenv("NCBI_API_KEY")  # Optional but increases rate limits
+NCBI_BASE = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils"
 
 def search_pubmed(query: str, max_results: int = 10) -> List[str]:
     """
-    Search PubMed for a query and return a list of PubMed IDs (PMIDs).
+    Search PubMed and return a list of PMIDs.
     """
     params = {
         "db": "pubmed",
         "term": query,
-        "retmode": "json",
-        "retmax": max_results
+        "retmax": max_results,
+        "api_key": NCBI_API_KEY
     }
-    if NCBI_API_KEY:
-        params["api_key"] = NCBI_API_KEY
-
-    r = requests.get(PUBMED_SEARCH_URL, params=params)
+    r = requests.get(f"{NCBI_BASE}/esearch.fcgi", params=params)
     r.raise_for_status()
-    data = r.json()
-    return data.get("esearchresult", {}).get("idlist", [])
-
+    root = ET.fromstring(r.text)
+    return [id_tag.text for id_tag in root.findall(".//Id")]
 
 def fetch_pubmed_details(pmids: List[str]) -> List[Dict]:
     """
-    Fetch details (title, authors, journal, abstract) for a list of PMIDs.
+    Fetch detailed information for a list of PMIDs.
     """
     if not pmids:
         return []
-
+    
     params = {
         "db": "pubmed",
         "id": ",".join(pmids),
-        "retmode": "xml"
+        "retmode": "xml",
+        "api_key": NCBI_API_KEY
     }
-    if NCBI_API_KEY:
-        params["api_key"] = NCBI_API_KEY
-
-    r = requests.get(PUBMED_FETCH_URL, params=params)
+    r = requests.get(f"{NCBI_BASE}/efetch.fcgi", params=params)
     r.raise_for_status()
-    return parse_pubmed_xml(r.text)
-
-
-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 = []
     
+    root = ET.fromstring(r.text)
+    articles = []
+
     for article in root.findall(".//PubmedArticle"):
-        title = article.findtext(".//ArticleTitle", default="")
-        abstract_text = " ".join([abst.text or "" for abst in article.findall(".//AbstractText")])
-        journal = article.findtext(".//Journal/Title", default="")
+        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 or fore:
+            last = author.findtext("LastName")
+            fore = author.findtext("ForeName")
+            if last and fore:
                 authors.append(f"{fore} {last}")
-        
-        pmid = article.findtext(".//PMID", default="")
-        
+        pmid = article.findtext(".//PMID")
         articles.append({
             "pmid": pmid,
             "title": title,
-            "abstract": abstract_text,
+            "abstract": abstract,
             "journal": journal,
-            "authors": authors
+            "pub_date": pub_date,
+            "doi": doi,
+            "authors": authors,
+            "url": f"https://pubmed.ncbi.nlm.nih.gov/{pmid}/"
         })
-    
     return articles
 
-
-def pubmed_summary(query: str, max_results: int = 10) -> List[Dict]:
+def search_and_fetch(query: str, max_results: int = 10) -> List[Dict]:
     """
-    Search PubMed and return summarized results in a single step.
+    Convenience function: Search and fetch results in one step.
     """
     pmids = search_pubmed(query, max_results)
     return fetch_pubmed_details(pmids)