Update genesis/api_clients/ncbi_api.py

genesis/api_clients/ncbi_api.py

@@ -1,33 +1,64 @@
 # genesis/api_clients/ncbi_api.py
 import requests
-import os
+import xml.etree.ElementTree as ET
 
-NCBI_API_KEY = os.getenv("NCBI_API_KEY")
-BASE_URL = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils"
+NCBI_BASE = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils"
 
-def search_gene(query: str, max_results: int = 10):
-    """Search NCBI Gene database."""
-    url = f"{BASE_URL}/esearch.fcgi"
+def ncbi_search(db: str, term: str, retmax: int = 10):
+    """
+    Search an NCBI database (pubmed, gene, protein, taxonomy, etc.)
+    and return a list of IDs.
+    """
+    url = f"{NCBI_BASE}/esearch.fcgi"
     params = {
-        "db": "gene",
-        "term": query,
-        "retmax": max_results,
+        "db": db,
+        "term": term,
         "retmode": "json",
-        "api_key": NCBI_API_KEY
+        "retmax": retmax
+    }
+    res = requests.get(url, params=params)
+    res.raise_for_status()
+    data = res.json()
+    return data.get("esearchresult", {}).get("idlist", [])
+
+def ncbi_fetch(db: str, ids: list, rettype: str = "abstract", retmode: str = "text"):
+    """
+    Fetch records from an NCBI database by IDs.
+    rettype can be 'abstract', 'fasta', 'gb', etc.
+    """
+    if not ids:
+        return None
+
+    url = f"{NCBI_BASE}/efetch.fcgi"
+    params = {
+        "db": db,
+        "id": ",".join(ids),
+        "rettype": rettype,
+        "retmode": retmode
     }
-    r = requests.get(url, params=params)
-    r.raise_for_status()
-    return r.json()
+    res = requests.get(url, params=params)
+    res.raise_for_status()
+    return res.text
+
+def get_gene_summary(gene_name: str, organism: str = None):
+    """
+    Get a gene summary from NCBI Gene.
+    """
+    query = gene_name
+    if organism:
+        query += f" AND {organism}[Organism]"
 
-def fetch_gene_summary(gene_id: str):
-    """Fetch summary for a specific gene."""
-    url = f"{BASE_URL}/esummary.fcgi"
+    ids = ncbi_search("gene", query, retmax=1)
+    if not ids:
+        return None
+
+    url = f"{NCBI_BASE}/esummary.fcgi"
     params = {
         "db": "gene",
-        "id": gene_id,
-        "retmode": "json",
-        "api_key": NCBI_API_KEY
+        "id": ids[0],
+        "retmode": "json"
     }
-    r = requests.get(url, params=params)
-    r.raise_for_status()
-    return r.json()
+    res = requests.get(url, params=params)
+    res.raise_for_status()
+    data = res.json()
+    return data.get("result", {}).get(ids[0], {})