Update genesis/api_clients/ncbi_api.py

genesis/api_clients/ncbi_api.py

@@ -1,178 +1,153 @@
 # genesis/api_clients/ncbi_api.py
 import os
 import requests
-import xml.etree.ElementTree as ET
 from typing import List, Dict
+from datetime import datetime
 
 NCBI_BASE = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils"
-NCBI_API_KEY = os.getenv("NCBI_API_KEY")  # Optional, set in Hugging Face / .env
+NCBI_API_KEY = os.getenv("NCBI_API_KEY")  # Optional — speeds up requests
+
+session = requests.Session()
 
 # -------------------------
-# Gene Search
+# Generic NCBI Search
 # -------------------------
-def search_gene(query: str, max_results: int = 5) -> List[Dict]:
+def ncbi_search(db: str, term: str, retmax: int = 10) -> List[str]:
     """
-    Search NCBI Gene for matching gene entries.
+    Search an NCBI database and return a list of IDs.
     """
     params = {
-        "db": "gene",
-        "term": query,
-        "retmax": max_results,
-        "api_key": NCBI_API_KEY
+        "db": db,
+        "term": term,
+        "retmode": "json",
+        "retmax": retmax
     }
-    r = requests.get(f"{NCBI_BASE}/esearch.fcgi", params=params)
-    r.raise_for_status()
-    ids = [elem.text for elem in ET.fromstring(r.text).findall(".//Id")]
-
-    return fetch_gene_details(ids)
-
-def fetch_gene_details(gene_ids: List[str]) -> List[Dict]:
-    """
-    Fetch detailed information for NCBI Gene IDs.
-    """
-    if not gene_ids:
-        return []
+    if NCBI_API_KEY:
+        params["api_key"] = NCBI_API_KEY
 
-    params = {
-        "db": "gene",
-        "id": ",".join(gene_ids),
-        "retmode": "xml",
-        "api_key": NCBI_API_KEY
-    }
-    r = requests.get(f"{NCBI_BASE}/efetch.fcgi", params=params)
+    r = session.get(f"{NCBI_BASE}/esearch.fcgi", params=params)
     r.raise_for_status()
 
-    genes = []
-    root = ET.fromstring(r.text)
-    for gene in root.findall(".//Entrezgene"):
-        gene_id_elem = gene.find(".//Gene-track_geneid")
-        gene_symbol_elem = gene.find(".//Gene-ref_locus")
-        gene_desc_elem = gene.find(".//Gene-ref_desc")
-
-        genes.append({
-            "gene_id": gene_id_elem.text if gene_id_elem is not None else "",
-            "symbol": gene_symbol_elem.text if gene_symbol_elem is not None else "",
-            "description": gene_desc_elem.text if gene_desc_elem is not None else "",
-            "url": f"https://www.ncbi.nlm.nih.gov/gene/{gene_id_elem.text}" if gene_id_elem is not None else ""
-        })
-
-    return genes
+    return r.json().get("esearchresult", {}).get("idlist", [])
 
 # -------------------------
-# Protein Search
+# Generic NCBI Fetch
 # -------------------------
-def search_protein(query: str, max_results: int = 5) -> List[Dict]:
+def ncbi_fetch(db: str, ids: List[str], rettype: str = "abstract", retmode: str = "text") -> str:
     """
-    Search NCBI Protein for matching entries.
+    Fetch detailed records from an NCBI database.
     """
     params = {
-        "db": "protein",
-        "term": query,
-        "retmax": max_results,
-        "api_key": NCBI_API_KEY
+        "db": db,
+        "id": ",".join(ids),
+        "rettype": rettype,
+        "retmode": retmode
     }
-    r = requests.get(f"{NCBI_BASE}/esearch.fcgi", params=params)
+    if NCBI_API_KEY:
+        params["api_key"] = NCBI_API_KEY
+
+    r = session.get(f"{NCBI_BASE}/efetch.fcgi", params=params)
     r.raise_for_status()
-    ids = [elem.text for elem in ET.fromstring(r.text).findall(".//Id")]
 
-    return fetch_protein_details(ids)
+    return r.text
 
-def fetch_protein_details(protein_ids: List[str]) -> List[Dict]:
+# -------------------------
+# PubMed Literature Search
+# -------------------------
+def search_pubmed(term: str, retmax: int = 5) -> List[Dict]:
     """
-    Fetch detailed information for NCBI Protein IDs.
+    Search PubMed for biomedical literature.
     """
-    if not protein_ids:
+    ids = ncbi_search("pubmed", term, retmax)
+    if not ids:
         return []
 
     params = {
-        "db": "protein",
-        "id": ",".join(protein_ids),
-        "retmode": "xml",
-        "api_key": NCBI_API_KEY
+        "db": "pubmed",
+        "id": ",".join(ids),
+        "retmode": "json"
     }
-    r = requests.get(f"{NCBI_BASE}/efetch.fcgi", params=params)
+    if NCBI_API_KEY:
+        params["api_key"] = NCBI_API_KEY
+
+    r = session.get(f"{NCBI_BASE}/esummary.fcgi", params=params)
     r.raise_for_status()
 
-    proteins = []
-    root = ET.fromstring(r.text)
-    for seq in root.findall(".//TSeq"):
-        acc_elem = seq.find(".//TSeq_accver")
-        def_elem = seq.find(".//TSeq_defline")
-        len_elem = seq.find(".//TSeq_length")
-
-        proteins.append({
-            "accession": acc_elem.text if acc_elem is not None else "",
-            "definition": def_elem.text if def_elem is not None else "",
-            "length": len_elem.text if len_elem is not None else "",
-            "url": f"https://www.ncbi.nlm.nih.gov/protein/{acc_elem.text}" if acc_elem is not None else ""
+    records = r.json().get("result", {})
+    papers = []
+    for pid in ids:
+        rec = records.get(pid, {})
+        papers.append({
+            "title": rec.get("title"),
+            "authors": [a["name"] for a in rec.get("authors", [])],
+            "pubdate": rec.get("pubdate"),
+            "journal": rec.get("fulljournalname"),
+            "uid": pid,
+            "link": f"https://pubmed.ncbi.nlm.nih.gov/{pid}/"
         })
-
-    return proteins
+    return papers
 
 # -------------------------
-# Sequence Search (Nucleotide)
+# Gene Search
 # -------------------------
-def search_nucleotide(query: str, max_results: int = 5) -> List[Dict]:
+def search_genes(term: str, retmax: int = 5) -> List[Dict]:
     """
-    Search NCBI Nucleotide for DNA/RNA sequences.
+    Search NCBI Gene database for gene information.
     """
+    ids = ncbi_search("gene", term, retmax)
+    if not ids:
+        return []
+
     params = {
-        "db": "nucleotide",
-        "term": query,
-        "retmax": max_results,
-        "api_key": NCBI_API_KEY
+        "db": "gene",
+        "id": ",".join(ids),
+        "retmode": "json"
     }
-    r = requests.get(f"{NCBI_BASE}/esearch.fcgi", params=params)
+    if NCBI_API_KEY:
+        params["api_key"] = NCBI_API_KEY
+
+    r = session.get(f"{NCBI_BASE}/esummary.fcgi", params=params)
     r.raise_for_status()
-    ids = [elem.text for elem in ET.fromstring(r.text).findall(".//Id")]
 
-    return fetch_nucleotide_details(ids)
+    records = r.json().get("result", {})
+    genes = []
+    for gid in ids:
+        rec = records.get(gid, {})
+        genes.append({
+            "symbol": rec.get("name"),
+            "description": rec.get("description"),
+            "organism": rec.get("organism", {}).get("scientificname"),
+            "uid": gid,
+            "link": f"https://www.ncbi.nlm.nih.gov/gene/{gid}"
+        })
+    return genes
 
-def fetch_nucleotide_details(nuc_ids: List[str]) -> List[Dict]:
+# -------------------------
+# Protein Search
+# -------------------------
+def search_proteins(term: str, retmax: int = 5) -> List[Dict]:
     """
-    Fetch detailed information for NCBI Nucleotide IDs.
+    Search NCBI Protein database for protein sequences.
     """
-    if not nuc_ids:
+    ids = ncbi_search("protein", term, retmax)
+    if not ids:
         return []
 
-    params = {
-        "db": "nucleotide",
-        "id": ",".join(nuc_ids),
-        "retmode": "xml",
-        "api_key": NCBI_API_KEY
-    }
-    r = requests.get(f"{NCBI_BASE}/efetch.fcgi", params=params)
-    r.raise_for_status()
-
-    sequences = []
-    root = ET.fromstring(r.text)
-    for seq in root.findall(".//TSeq"):
-        acc_elem = seq.find(".//TSeq_accver")
-        def_elem = seq.find(".//TSeq_defline")
-        len_elem = seq.find(".//TSeq_length")
-
-        sequences.append({
-            "accession": acc_elem.text if acc_elem is not None else "",
-            "definition": def_elem.text if def_elem is not None else "",
-            "length": len_elem.text if len_elem is not None else "",
-            "url": f"https://www.ncbi.nlm.nih.gov/nuccore/{acc_elem.text}" if acc_elem is not None else ""
-        })
-
-    return sequences
+    fasta_data = ncbi_fetch("protein", ids, rettype="fasta", retmode="text")
+    proteins = [{"id": pid, "fasta": fasta_data} for pid in ids]
+    return proteins
 
 # -------------------------
-# Cross-Domain Integration
+# Build Cross-Database Profile
 # -------------------------
-def entity_context(query: str) -> Dict:
+def ncbi_cross_profile(term: str) -> Dict:
     """
-    Return gene, protein, and sequence info linked to PubMed and ChEMBL.
+    Given a term, pull literature, genes, and proteins for unified output.
     """
-    from genesis.api_clients import pubmed_api, chembl_api  # Lazy import to avoid cycles
-
     return {
-        "genes": search_gene(query),
-        "proteins": search_protein(query),
-        "nucleotides": search_nucleotide(query),
-        "literature": pubmed_api.search_pubmed(query),
-        "related_drugs": chembl_api.search_molecule(query)
+        "term": term,
+        "timestamp": datetime.utcnow().isoformat(),
+        "literature": search_pubmed(term, retmax=5),
+        "genes": search_genes(term, retmax=5),
+        "proteins": search_proteins(term, retmax=2)
     }