Spaces:
Sleeping
Sleeping
File size: 4,051 Bytes
0ccfe4a 6e92226 0ccfe4a 6e92226 0ccfe4a 6e92226 0ccfe4a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 | """PubMed Entrez API client for ALS paper ingestion."""
from __future__ import annotations
import os
import time
from Bio import Entrez, Medline
from config import PUBMED_BATCH_SIZE
from logging_config import get_logger
from models import ALSPaper
_logger = get_logger("ingestion.pubmed")
def _configure_entrez() -> None:
email = os.environ.get("ENTREZ_EMAIL")
if not email:
raise EnvironmentError("ENTREZ_EMAIL environment variable is required by NCBI")
Entrez.email = email
api_key = os.getenv("NCBI_API_KEY")
if api_key:
Entrez.api_key = api_key
def _sleep() -> None:
"""Respect NCBI rate limits: 10 req/s with API key, 3 req/s without."""
time.sleep(0.1 if os.getenv("NCBI_API_KEY") else 0.4)
_ESEARCH_PAGE_SIZE = 9999 # NCBI hard cap per esearch call
def search_pmids(query: str, max_results: int = 500) -> list[str]:
"""Search PubMed with a query string and return a list of PMIDs.
Pages through esearch results in chunks of 9,999 (NCBI's per-call cap)
until max_results or the total result count is reached.
"""
_configure_entrez()
pmids: list[str] = []
retstart = 0
total: int | None = None
while True:
want = min(_ESEARCH_PAGE_SIZE, max_results - len(pmids))
handle = Entrez.esearch(db="pubmed", term=query, retmax=want, retstart=retstart)
record = Entrez.read(handle)
handle.close()
if total is None:
total = int(record["Count"])
page = list(record["IdList"])
pmids.extend(page)
if not page or len(pmids) >= max_results or len(pmids) >= total:
break
retstart += len(page)
_sleep()
_logger.info("PubMed esearch", extra={"data": {"count": len(pmids), "total": total, "query": query[:80]}})
return pmids
def fetch_by_pmids(pmids: list[str]) -> list[ALSPaper]:
"""Fetch and parse paper records for a list of PMIDs."""
_configure_entrez()
papers: list[ALSPaper] = []
for i in range(0, len(pmids), PUBMED_BATCH_SIZE):
batch = pmids[i : i + PUBMED_BATCH_SIZE]
_logger.debug("Fetching Entrez batch", extra={"data": {"batch": i // PUBMED_BATCH_SIZE + 1, "size": len(batch)}})
for attempt in range(3):
try:
handle = Entrez.efetch(db="pubmed", id=",".join(batch), rettype="medline", retmode="text")
records = list(Medline.parse(handle))
handle.close()
break
except Exception as exc:
if attempt == 2:
raise
wait = 2 ** attempt
_logger.warning(f"Entrez fetch error (attempt {attempt + 1}): {exc}")
time.sleep(wait)
for record in records:
paper = _parse_record(record)
if paper:
papers.append(paper)
_sleep()
_logger.info("PubMed fetch complete", extra={"data": {"total": len(papers)}})
return papers
def _parse_record(record: dict) -> ALSPaper | None:
"""Convert a Biopython Medline record to ALSPaper. Returns None if no abstract."""
pmid = record.get("PMID", "")
abstract = record.get("AB", "")
if not pmid or not abstract:
return None
authors = record.get("FAU", record.get("AU", []))
if isinstance(authors, str):
authors = [authors]
# "DP" field: "2023 Jan 15", "2023 Jan", "2023"
year = 0
date_str = record.get("DP", "")
if date_str:
try:
year = int(date_str.split()[0])
except (ValueError, IndexError):
pass
# DOI from AID list: ["10.1093/xxx [doi]", "S0092-8674(23)00001-1 [pii]"]
doi = ""
for aid in record.get("AID", []):
if aid.endswith("[doi]"):
doi = aid.replace(" [doi]", "").strip()
break
return ALSPaper(
pmid=pmid,
title=record.get("TI", ""),
abstract=abstract,
authors=authors,
year=year,
doi=doi,
mesh_terms=record.get("MH", []),
)
|