Update genesis/api_clients/ncbi_api.py

genesis/api_clients/ncbi_api.py

@@ -1,120 +1,178 @@
 # genesis/api_clients/ncbi_api.py
 import os
 import requests
-from typing import Dict, List, Optional
+import xml.etree.ElementTree as ET
+from typing import List, Dict
 
-NCBI_BASE = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/"
-NCBI_API_KEY = os.getenv("NCBI_API_KEY")  # Optional for higher rate limits
+NCBI_BASE = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils"
+NCBI_API_KEY = os.getenv("NCBI_API_KEY")  # Optional, set in Hugging Face / .env
 
 # -------------------------
-# Core Utilities
+# Gene Search
 # -------------------------
-def ncbi_request(endpoint: str, params: Dict) -> requests.Response:
+def search_gene(query: str, max_results: int = 5) -> List[Dict]:
     """
-    Helper function to query NCBI E-Utilities with optional API key.
+    Search NCBI Gene for matching gene entries.
     """
-    if NCBI_API_KEY:
-        params["api_key"] = NCBI_API_KEY
-    url = f"{NCBI_BASE}{endpoint}"
-    r = requests.get(url, params=params)
-    r.raise_for_status()
-    return r
-
-# -------------------------
-# Search Functions
-# -------------------------
-def search_pubmed(query: str, max_results: int = 10) -> List[str]:
-    """
-    Search PubMed and return list of PMIDs.
-    """
-    r = ncbi_request("esearch.fcgi", {
-        "db": "pubmed",
+    params = {
+        "db": "gene",
         "term": query,
         "retmax": max_results,
-        "retmode": "json"
-    })
-    return r.json().get("esearchresult", {}).get("idlist", [])
+        "api_key": NCBI_API_KEY
+    }
+    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 search_gene(query: str, organism: Optional[str] = None, max_results: int = 10) -> List[str]:
+def fetch_gene_details(gene_ids: List[str]) -> List[Dict]:
     """
-    Search NCBI Gene database.
+    Fetch detailed information for NCBI Gene IDs.
     """
-    term = query
-    if organism:
-        term += f" AND {organism}[Organism]"
-    r = ncbi_request("esearch.fcgi", {
+    if not gene_ids:
+        return []
+
+    params = {
         "db": "gene",
-        "term": term,
-        "retmax": max_results,
-        "retmode": "json"
-    })
-    return r.json().get("esearchresult", {}).get("idlist", [])
+        "id": ",".join(gene_ids),
+        "retmode": "xml",
+        "api_key": NCBI_API_KEY
+    }
+    r = requests.get(f"{NCBI_BASE}/efetch.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")
 
-def search_protein(query: str, organism: Optional[str] = None, max_results: int = 10) -> List[str]:
+        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
+
+# -------------------------
+# Protein Search
+# -------------------------
+def search_protein(query: str, max_results: int = 5) -> List[Dict]:
     """
-    Search NCBI Protein database.
+    Search NCBI Protein for matching entries.
     """
-    term = query
-    if organism:
-        term += f" AND {organism}[Organism]"
-    r = ncbi_request("esearch.fcgi", {
+    params = {
         "db": "protein",
-        "term": term,
+        "term": query,
         "retmax": max_results,
-        "retmode": "json"
-    })
-    return r.json().get("esearchresult", {}).get("idlist", [])
+        "api_key": NCBI_API_KEY
+    }
+    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")]
 
-# -------------------------
-# Fetch Functions
-# -------------------------
-def fetch_summary(db: str, ids: List[str]) -> List[Dict]:
+    return fetch_protein_details(ids)
+
+def fetch_protein_details(protein_ids: List[str]) -> List[Dict]:
     """
-    Fetch summary data from any NCBI database.
+    Fetch detailed information for NCBI Protein IDs.
     """
-    if not ids:
+    if not protein_ids:
         return []
-    r = ncbi_request("esummary.fcgi", {
-        "db": db,
-        "id": ",".join(ids),
-        "retmode": "json"
-    })
-    return list(r.json().get("result", {}).values())
-
-def fetch_fasta(db: str, ids: List[str]) -> Dict[str, str]:
-    """
-    Fetch FASTA sequences from NCBI (protein or nucleotide).
-    """
-    if not ids:
-        return {}
-    r = ncbi_request("efetch.fcgi", {
-        "db": db,
-        "id": ",".join(ids),
-        "rettype": "fasta",
-        "retmode": "text"
-    })
-    return {ids[i]: seq for i, seq in enumerate(r.text.strip().split(">")[1:])}
+
+    params = {
+        "db": "protein",
+        "id": ",".join(protein_ids),
+        "retmode": "xml",
+        "api_key": NCBI_API_KEY
+    }
+    r = requests.get(f"{NCBI_BASE}/efetch.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 ""
+        })
+
+    return proteins
 
 # -------------------------
-# Specialized Functions
+# Sequence Search (Nucleotide)
 # -------------------------
-def get_gene_info(gene_id: str) -> Dict:
+def search_nucleotide(query: str, max_results: int = 5) -> List[Dict]:
     """
-    Get detailed gene info from NCBI Gene.
+    Search NCBI Nucleotide for DNA/RNA sequences.
     """
-    summaries = fetch_summary("gene", [gene_id])
-    return summaries[0] if summaries else {}
+    params = {
+        "db": "nucleotide",
+        "term": query,
+        "retmax": max_results,
+        "api_key": NCBI_API_KEY
+    }
+    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_nucleotide_details(ids)
 
-def get_protein_info(protein_id: str) -> Dict:
+def fetch_nucleotide_details(nuc_ids: List[str]) -> List[Dict]:
     """
-    Get detailed protein info from NCBI Protein.
+    Fetch detailed information for NCBI Nucleotide IDs.
     """
-    summaries = fetch_summary("protein", [protein_id])
-    return summaries[0] if summaries else {}
+    if not nuc_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")
 
-def get_taxonomy_info(tax_id: str) -> Dict:
+        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
+
+# -------------------------
+# Cross-Domain Integration
+# -------------------------
+def entity_context(query: str) -> Dict:
     """
-    Get taxonomy data from NCBI Taxonomy.
+    Return gene, protein, and sequence info linked to PubMed and ChEMBL.
     """
-    summaries = fetch_summary("taxonomy", [tax_id])
-    return summaries[0] if summaries else {}
+    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)
+    }