Update genesis/api_clients/ncbi_api.py

genesis/api_clients/ncbi_api.py

@@ -1,64 +1,98 @@
 # genesis/api_clients/ncbi_api.py
+import os
 import requests
-import xml.etree.ElementTree as ET
+from urllib.parse import urlencode
 
-NCBI_BASE = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils"
+NCBI_API_KEY = os.getenv("NCBI_API_KEY")
+NCBI_BASE_URL = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/"
 
-def ncbi_search(db: str, term: str, retmax: int = 10):
+if not NCBI_API_KEY:
+    raise ValueError("Missing NCBI_API_KEY in environment variables")
+
+
+def _ncbi_request(endpoint: str, params: dict):
     """
-    Search an NCBI database (pubmed, gene, protein, taxonomy, etc.)
-    and return a list of IDs.
+    Generic helper for calling NCBI E-Utilities API.
     """
-    url = f"{NCBI_BASE}/esearch.fcgi"
-    params = {
-        "db": db,
-        "term": term,
-        "retmode": "json",
-        "retmax": retmax
-    }
-    res = requests.get(url, params=params)
+    params["api_key"] = NCBI_API_KEY
+    url = f"{NCBI_BASE_URL}{endpoint}?{urlencode(params)}"
+    res = requests.get(url)
     res.raise_for_status()
-    data = res.json()
-    return data.get("esearchresult", {}).get("idlist", [])
+    return res.text
+
 
-def ncbi_fetch(db: str, ids: list, rettype: str = "abstract", retmode: str = "text"):
+def search_pubmed(query: str, max_results: int = 10):
     """
-    Fetch records from an NCBI database by IDs.
-    rettype can be 'abstract', 'fasta', 'gb', etc.
+    Search PubMed articles via NCBI.
     """
-    if not ids:
-        return None
+    params = {"db": "pubmed", "term": query, "retmax": max_results, "retmode": "json"}
+    res = requests.get(f"{NCBI_BASE_URL}esearch.fcgi", params=params)
+    res.raise_for_status()
+    ids = res.json().get("esearchresult", {}).get("idlist", [])
+    return ids
+
 
-    url = f"{NCBI_BASE}/efetch.fcgi"
+def fetch_pubmed_details(pubmed_ids: list):
+    """
+    Fetch PubMed details given a list of IDs.
+    """
+    if not pubmed_ids:
+        return []
     params = {
-        "db": db,
-        "id": ",".join(ids),
-        "rettype": rettype,
-        "retmode": retmode
+        "db": "pubmed",
+        "id": ",".join(pubmed_ids),
+        "retmode": "xml"
     }
-    res = requests.get(url, params=params)
+    res = requests.get(f"{NCBI_BASE_URL}efetch.fcgi", params=params)
     res.raise_for_status()
-    return res.text
+    return res.text  # XML response
+
 
-def get_gene_summary(gene_name: str, organism: str = None):
+def search_gene(query: str):
     """
-    Get a gene summary from NCBI Gene.
+    Search for a gene in NCBI Gene database.
     """
-    query = gene_name
-    if organism:
-        query += f" AND {organism}[Organism]"
+    params = {"db": "gene", "term": query, "retmode": "json", "retmax": 5}
+    res = requests.get(f"{NCBI_BASE_URL}esearch.fcgi", params=params)
+    res.raise_for_status()
+    return res.json()
 
-    ids = ncbi_search("gene", query, retmax=1)
-    if not ids:
-        return None
 
-    url = f"{NCBI_BASE}/esummary.fcgi"
-    params = {
-        "db": "gene",
-        "id": ids[0],
-        "retmode": "json"
-    }
-    res = requests.get(url, params=params)
+def fetch_gene_summary(gene_id: str):
+    """
+    Get detailed gene summary from NCBI.
+    """
+    params = {"db": "gene", "id": gene_id, "retmode": "json"}
+    res = requests.get(f"{NCBI_BASE_URL}esummary.fcgi", params=params)
+    res.raise_for_status()
+    return res.json()
+
+
+def search_protein(query: str):
+    """
+    Search proteins in NCBI Protein database.
+    """
+    params = {"db": "protein", "term": query, "retmode": "json", "retmax": 5}
+    res = requests.get(f"{NCBI_BASE_URL}esearch.fcgi", params=params)
+    res.raise_for_status()
+    return res.json()
+
+
+def search_pubchem_compound(query: str):
+    """
+    Search compounds in PubChem via NCBI.
+    """
+    params = {"db": "pccompound", "term": query, "retmode": "json", "retmax": 5}
+    res = requests.get(f"{NCBI_BASE_URL}esearch.fcgi", params=params)
+    res.raise_for_status()
+    return res.json()
+
+
+def search_clinical_trials(query: str):
+    """
+    Search clinical trials in ClinicalTrials.gov via NCBI.
+    """
+    params = {"db": "clinicaltrials", "term": query, "retmode": "json", "retmax": 5}
+    res = requests.get(f"{NCBI_BASE_URL}esearch.fcgi", params=params)
     res.raise_for_status()
-    data = res.json()
-    return data.get("result", {}).get(ids[0], {})
+    return res.json()