Update genesis/api_clients/ncbi_api.py

genesis/api_clients/ncbi_api.py

@@ -1,28 +1,28 @@
 # genesis/api_clients/ncbi_api.py
 import os
 import requests
-from xml.etree import ElementTree as ET
+from typing import List, Dict, Optional
 
-NCBI_API_KEY = os.getenv("NCBI_API_KEY")
+NCBI_API_KEY = os.getenv("NCBI_API_KEY")  # Optional, for higher request limits
 NCBI_BASE = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils"
 
-def _add_api_key(params: dict):
-    """
-    Adds API key to request params if available.
-    """
+def _add_api_key(params: Dict) -> Dict:
+    """Attach API key if available."""
     if NCBI_API_KEY:
         params["api_key"] = NCBI_API_KEY
     return params
 
-
-def search_pubmed(query: str, max_results: int = 10):
+# -------------------------
+# SEARCH FUNCTIONS
+# -------------------------
+def search_ncbi(db: str, term: str, max_results: int = 10) -> List[str]:
     """
-    Search PubMed via NCBI Entrez for a query string.
-    Returns a list of PubMed IDs (PMIDs).
+    Search an NCBI database and return a list of IDs.
+    db examples: gene, protein, pubmed, taxonomy
     """
     params = _add_api_key({
-        "db": "pubmed",
-        "term": query,
+        "db": db,
+        "term": term,
         "retmax": max_results,
         "retmode": "json"
     })
@@ -31,107 +31,76 @@ def search_pubmed(query: str, max_results: int = 10):
     data = r.json()
     return data.get("esearchresult", {}).get("idlist", [])
 
-
-def fetch_pubmed_details(pmids: list):
+def fetch_ncbi_summary(db: str, ids: List[str]) -> List[Dict]:
     """
-    Fetch abstracts, authors, and metadata for given PMIDs.
+    Fetch summaries for a list of IDs from NCBI.
     """
-    if not pmids:
-        return []
-    
     params = _add_api_key({
-        "db": "pubmed",
-        "id": ",".join(pmids),
-        "retmode": "xml"
+        "db": db,
+        "id": ",".join(ids),
+        "retmode": "json"
     })
-    r = requests.get(f"{NCBI_BASE}/efetch.fcgi", params=params)
+    r = requests.get(f"{NCBI_BASE}/esummary.fcgi", params=params)
     r.raise_for_status()
+    data = r.json()
+    summaries = []
+    for uid, summary in data.get("result", {}).items():
+        if uid != "uids":
+            summaries.append(summary)
+    return summaries
 
-    root = ET.fromstring(r.text)
-    results = []
-
-    for article in root.findall(".//PubmedArticle"):
-        title = article.findtext(".//ArticleTitle", default="")
-        abstract = " ".join([abst.text or "" for abst in article.findall(".//AbstractText")])
-        authors = [
-            f"{a.findtext('ForeName', '')} {a.findtext('LastName', '')}".strip()
-            for a in article.findall(".//Author")
-        ]
-        pmid = article.findtext(".//PMID", default="")
-        
-        results.append({
-            "pmid": pmid,
-            "title": title,
-            "abstract": abstract,
-            "authors": authors
-        })
-
-    return results
-
-
-def fetch_gene_info(gene_id: str):
+def fetch_ncbi_details(db: str, ids: List[str]) -> str:
     """
-    Fetch gene metadata from NCBI Gene database.
+    Fetch full XML/FASTA/GenBank record for IDs.
     """
     params = _add_api_key({
-        "db": "gene",
-        "id": gene_id,
-        "retmode": "xml"
+        "db": db,
+        "id": ",".join(ids),
+        "retmode": "text"
     })
     r = requests.get(f"{NCBI_BASE}/efetch.fcgi", params=params)
     r.raise_for_status()
-    return r.text  # XML - downstream parsing can be done in pipeline
-
+    return r.text
 
-def search_gene_by_symbol(symbol: str, organism: str = None):
+# -------------------------
+# GENE + PATHWAY HELPERS
+# -------------------------
+def search_gene(term: str, max_results: int = 10) -> List[Dict]:
     """
-    Search for a gene by symbol, optionally filtered by organism.
+    Search for genes and return gene IDs + names.
     """
-    term = symbol
-    if organism:
-        term += f" AND {organism}[Organism]"
-
-    params = _add_api_key({
-        "db": "gene",
-        "term": term,
-        "retmax": 5,
-        "retmode": "json"
-    })
-    r = requests.get(f"{NCBI_BASE}/esearch.fcgi", params=params)
-    r.raise_for_status()
-    data = r.json()
-    return data.get("esearchresult", {}).get("idlist", [])
-
+    ids = search_ncbi("gene", term, max_results)
+    if not ids:
+        return []
+    summaries = fetch_ncbi_summary("gene", ids)
+    return [{"uid": s.get("uid"), "name": s.get("name"), "description": s.get("description")} for s in summaries]
 
-def fetch_protein_info(protein_id: str):
+def get_protein_from_gene(gene_id: str) -> List[Dict]:
     """
-    Fetch protein metadata from NCBI Protein database.
+    Get protein products from a given gene ID.
     """
-    params = _add_api_key({
+    link_params = _add_api_key({
+        "dbfrom": "gene",
         "db": "protein",
-        "id": protein_id,
-        "retmode": "xml"
+        "id": gene_id,
+        "retmode": "json"
     })
-    r = requests.get(f"{NCBI_BASE}/efetch.fcgi", params=params)
+    r = requests.get(f"{NCBI_BASE}/elink.fcgi", params=link_params)
     r.raise_for_status()
-    return r.text
+    data = r.json()
 
+    protein_ids = []
+    for linkset in data.get("linksets", []):
+        for link in linkset.get("linksetdbs", []):
+            protein_ids.extend(link.get("links", []))
+    
+    if not protein_ids:
+        return []
+    return fetch_ncbi_summary("protein", protein_ids)
 
-def search_protein_by_name(name: str, organism: str = None):
+def search_taxonomy(term: str) -> List[Dict]:
     """
-    Search for proteins by name, optionally filtered by organism.
+    Search taxonomy database for species/strain info.
     """
-    term = name
-    if organism:
-        term += f" AND {organism}[Organism]"
-
-    params = _add_api_key({
-        "db": "protein",
-        "term": term,
-        "retmax": 5,
-        "retmode": "json"
-    })
-    r = requests.get(f"{NCBI_BASE}/esearch.fcgi", params=params)
-    r.raise_for_status()
-    data = r.json()
-    return data.get("esearchresult", {}).get("idlist", [])
+    ids = search_ncbi("taxonomy", term, max_results=5)
+    return fetch_ncbi_summary("taxonomy", ids)