refactor: unify searcher interfaces and improve error handling

- Extract utility functions (_gene_record_to_dict, _accession_to_dict, _rna_data_to_dict)
- Unify method naming: search_by_keyword -> get_best_hit
- Add threshold parameter to NCBI and RNAcentral searchers for interface consistency
- Improve error handling with network error detection and fallback strategies
- Fix RNAcentral sequence search to prioritize exact matches
- Add search_rna_demo.jsonl example file

Author: CHERRY-ui8

graphgen/models/searcher/db/ncbi_searcher.py

@@ -41,6 +41,38 @@ def __init__(self, email: str = "test@example.com", tool: str = "GraphGen"):
         Entrez.tool = tool
         Entrez.timeout = 60  # 60 seconds timeout
 
+    @staticmethod
+    def _gene_record_to_dict(gene_record, gene_id: str) -> dict:
+        """
+        Convert an Entrez gene record to a dictionary.
+        :param gene_record: The Entrez gene record (list from Entrez.read).
+        :param gene_id: The gene ID.
+        :return: A dictionary containing gene information.
+        """
+        if not gene_record:
+            raise ValueError("Empty gene record")
+
+        gene_data = gene_record[0]
+        gene_ref = gene_data.get("Entrezgene_gene", {}).get("Gene-ref", {})
+
+        organism = (
+            gene_data.get("Entrezgene_source", {})
+            .get("BioSource", {})
+            .get("BioSource_org", {})
+            .get("Org-ref", {})
+            .get("Org-ref_taxname", "N/A")
+        )
+
+        return {
+            "molecule_type": "DNA",
+            "database": "NCBI",
+            "id": gene_id,
+            "gene_name": gene_ref.get("Gene-ref_locus", "N/A"),
+            "gene_description": gene_ref.get("Gene-ref_desc", "N/A"),
+            "organism": organism,
+            "url": f"https://www.ncbi.nlm.nih.gov/gene/{gene_id}",
+        }
+
     def get_by_gene_id(self, gene_id: str) -> Optional[dict]:
         """
         Get gene information by Gene ID.
@@ -54,26 +86,7 @@ def get_by_gene_id(self, gene_id: str) -> Optional[dict]:
                 gene_record = Entrez.read(handle)
                 if not gene_record:
                     return None
-
-                gene_data = gene_record[0]
-                gene_ref = gene_data.get("Entrezgene_gene", {}).get("Gene-ref", {})
-
-                organism = (
-                    gene_data.get("Entrezgene_source", {})
-                    .get("BioSource", {})
-                    .get("BioSource_org", {})
-                    .get("Org-ref", {})
-                    .get("Org-ref_taxname", "N/A")
-                )
-                return {
-                    "molecule_type": "DNA",
-                    "database": "NCBI",
-                    "id": gene_id,
-                    "gene_name": gene_ref.get("Gene-ref_locus", "N/A"),
-                    "gene_description": gene_ref.get("Gene-ref_desc", "N/A"),
-                    "organism": organism,
-                    "url": f"https://www.ncbi.nlm.nih.gov/gene/{gene_id}",
-                }
+                return self._gene_record_to_dict(gene_record, gene_id)
             finally:
                 handle.close()
         except RequestException:
@@ -82,14 +95,36 @@ def get_by_gene_id(self, gene_id: str) -> Optional[dict]:
             logger.error("Gene ID %s not found: %s", gene_id, exc)
             return None
 
+    @staticmethod
+    def _accession_to_dict(accession: str, sequence: str, header: str, title: str, organism: str) -> dict:
+        """
+        Convert accession information to a dictionary.
+        :param accession: NCBI accession number.
+        :param sequence: DNA sequence.
+        :param header: FASTA header.
+        :param title: Sequence title.
+        :param organism: Organism name.
+        :return: A dictionary containing sequence information.
+        """
+        return {
+            "molecule_type": "DNA",
+            "database": "NCBI",
+            "id": accession,
+            "title": title,
+            "organism": organism,
+            "sequence": sequence,
+            "sequence_length": len(sequence),
+            "url": f"https://www.ncbi.nlm.nih.gov/nuccore/{accession}",
+        }
+
     def get_by_accession(self, accession: str) -> Optional[dict]:
         """
         Get sequence information by accession number.
         :param accession: NCBI accession number (e.g., NM_000546).
         :return: A dictionary containing sequence information or None if not found.
         """
         try:
-            time.sleep(0.35)  # 遵守速率限制
+            time.sleep(0.35)  # Comply with rate limit
             handle = Entrez.efetch(
                 db="nuccore",
                 id=accession,
@@ -120,16 +155,7 @@ def get_by_accession(self, accession: str) -> Optional[dict]:
                 finally:
                     summary_handle.close()
 
-                return {
-                    "molecule_type": "DNA",
-                    "database": "NCBI",
-                    "id": accession,
-                    "title": title,
-                    "organism": organism,
-                    "sequence": sequence,
-                    "sequence_length": len(sequence),
-                    "url": f"https://www.ncbi.nlm.nih.gov/nuccore/{accession}",
-                }
+                return self._accession_to_dict(accession, sequence, header, title, organism)
             finally:
                 handle.close()
         except RequestException:
@@ -138,7 +164,7 @@ def get_by_accession(self, accession: str) -> Optional[dict]:
             logger.error("Accession %s not found: %s", accession, exc)
             return None
 
-    def search_by_keyword(self, keyword: str) -> Optional[dict]:
+    def get_best_hit(self, keyword: str) -> Optional[dict]:
         """
         Search NCBI Gene database with a keyword and return the best hit.
         :param keyword: The search keyword (e.g., gene name).
@@ -148,7 +174,7 @@ def search_by_keyword(self, keyword: str) -> Optional[dict]:
             return None
 
         try:
-            time.sleep(0.35)  # 遵守速率限制
+            time.sleep(0.35)  # Comply with rate limit
             # Search gene database
             search_handle = Entrez.esearch(
                 db="gene",
@@ -181,11 +207,12 @@ def search_by_keyword(self, keyword: str) -> Optional[dict]:
             logger.error("Keyword %s not found: %s", keyword, e)
         return None
 
-    def search_by_sequence(self, sequence: str) -> Optional[dict]:
+    def search_by_sequence(self, sequence: str, threshold: float = 0.01) -> Optional[dict]:
         """
         Search NCBI with a DNA sequence using BLAST.
         Note: This is a simplified version. For production, consider using local BLAST.
         :param sequence: DNA sequence (FASTA format or raw sequence).
+        :param threshold: E-value threshold for BLAST search.
         :return: A dictionary containing the best hit information or None if not found.
         """
         try:
@@ -215,7 +242,7 @@ def search_by_sequence(self, sequence: str) -> Optional[dict]:
                 database="nr",
                 sequence=seq,
                 hitlist_size=1,
-                expect=0.001,
+                expect=threshold,
             )
             blast_record = NCBIXML.read(result_handle)
 
@@ -225,6 +252,9 @@ def search_by_sequence(self, sequence: str) -> Optional[dict]:
 
             best_alignment = blast_record.alignments[0]
             best_hsp = best_alignment.hsps[0]
+            if best_hsp.expect > threshold:
+                logger.info("No BLAST hits below the threshold E-value.")
+                return None
             hit_id = best_alignment.hit_id
 
             # Extract accession number
@@ -257,11 +287,12 @@ def search_by_sequence(self, sequence: str) -> Optional[dict]:
         reraise=True,
     )
     async def search(
-        self, query: str, **kwargs
+        self, query: str, threshold: float = 0.01, **kwargs
     ) -> Optional[Dict]:
         """
         Search NCBI with either a gene ID, accession number, keyword, or DNA sequence.
         :param query: The search query (gene ID, accession, keyword, or DNA sequence).
+        :param threshold: E-value threshold for BLAST search.
         :param kwargs: Additional keyword arguments (not used currently).
         :return: A dictionary containing the search results or None if not found.
         """
@@ -278,7 +309,7 @@ async def search(
         # check if DNA sequence (ATCG characters)
         if query.startswith(">") or re.fullmatch(r"[ATCGN\s]+", query, re.I):
             result = await loop.run_in_executor(
-                _get_pool(), self.search_by_sequence, query
+                _get_pool(), self.search_by_sequence, query, threshold
             )
         # check if gene ID (numeric)
         elif re.fullmatch(r"^\d+$", query):
@@ -293,7 +324,7 @@ async def search(
         else:
             # otherwise treat as keyword
             result = await loop.run_in_executor(
-                _get_pool(), self.search_by_keyword, query
+                _get_pool(), self.get_best_hit, query
             )
 
         if result:

graphgen/models/searcher/db/rnacentral_searcher.py

@@ -36,6 +36,27 @@ def __init__(self):
         self.base_url = "https://rnacentral.org/api/v1"
         self.headers = {"Accept": "application/json"}
 
+    @staticmethod
+    def _rna_data_to_dict(rna_id: str, rna_data: dict) -> dict:
+        """
+        Convert RNAcentral API response to a dictionary.
+        :param rna_id: RNAcentral ID.
+        :param rna_data: API response data (dict or dict-like from search results).
+        :return: A dictionary containing RNA information.
+        """
+        sequence = rna_data.get("sequence", "")
+        return {
+            "molecule_type": "RNA",
+            "database": "RNAcentral",
+            "id": rna_id,
+            "rnacentral_id": rna_data.get("rnacentral_id", rna_id),
+            "sequence": sequence,
+            "sequence_length": rna_data.get("length", len(sequence)),
+            "rna_type": rna_data.get("rna_type", "N/A"),
+            "description": rna_data.get("description", "N/A"),
+            "url": f"https://rnacentral.org/rna/{rna_id}",
+        }
+
     async def get_by_rna_id(self, rna_id: str) -> Optional[dict]:
         """
         Get RNA information by RNAcentral ID.
@@ -50,26 +71,19 @@ async def get_by_rna_id(self, rna_id: str) -> Optional[dict]:
                 ) as resp:
                     if resp.status == 200:
                         rna_data = await resp.json()
-                        return {
-                            "molecule_type": "RNA",
-                            "database": "RNAcentral",
-                            "id": rna_id,
-                            "rnacentral_id": rna_data.get("rnacentral_id", "N/A"),
-                            "sequence": rna_data.get("sequence", ""),
-                            "sequence_length": len(rna_data.get("sequence", "")),
-                            "rna_type": rna_data.get("rna_type", "N/A"),
-                            "description": rna_data.get("description", "N/A"),
-                            "url": f"https://rnacentral.org/rna/{rna_id}",
-                        }
+                        return self._rna_data_to_dict(rna_id, rna_data)
                     if resp.status == 404:
                         logger.error("RNA ID %s not found", rna_id)
                         return None
                     raise Exception(f"HTTP {resp.status}: {await resp.text()}")
+        except aiohttp.ClientError as e:
+            logger.error("Network error getting RNA ID %s: %s", rna_id, e)
+            return None
         except Exception as exc:  # pylint: disable=broad-except
             logger.error("RNA ID %s not found: %s", rna_id, exc)
             return None
 
-    async def search_by_keyword(self, keyword: str) -> Optional[dict]:
+    async def get_best_hit(self, keyword: str) -> Optional[dict]:
         """
         Search RNAcentral with a keyword and return the best hit.
         :param keyword: The search keyword (e.g., miRNA name, RNA name).
@@ -90,13 +104,26 @@ async def search_by_keyword(self, keyword: str) -> Optional[dict]:
                 ) as resp:
                     if resp.status == 200:
                         search_results = await resp.json()
-                        if search_results.get("results"):
-                            rna_id = search_results["results"][0].get("rnacentral_id")
+                        results = search_results.get("results", [])
+                        if results:
+                            # Use the first result directly (search API already returns enough info)
+                            first_result = results[0]
+                            rna_id = first_result.get("rnacentral_id")
                             if rna_id:
-                                return await self.get_by_rna_id(rna_id)
+                                # Try to get detailed info, but fall back to search result if it fails
+                                detailed_info = await self.get_by_rna_id(rna_id)
+                                if detailed_info:
+                                    return detailed_info
+                                # Fall back to using search result data
+                                return self._rna_data_to_dict(rna_id, first_result)
                         logger.info("No results found for keyword: %s", keyword)
                         return None
-                    raise Exception(f"HTTP {resp.status}: {await resp.text()}")
+                    error_text = await resp.text()
+                    logger.error("HTTP %d error for keyword %s: %s", resp.status, keyword, error_text[:200])
+                    raise Exception(f"HTTP {resp.status}: {error_text}")
+        except aiohttp.ClientError as e:
+            logger.error("Network error searching for keyword %s: %s", keyword, e)
+            return None
         except Exception as e:  # pylint: disable=broad-except
             logger.error("Keyword %s not found: %s", keyword, e)
             return None
@@ -136,13 +163,39 @@ async def search_by_sequence(self, sequence: str) -> Optional[dict]:
                 ) as resp:
                     if resp.status == 200:
                         search_results = await resp.json()
-                        if search_results.get("results"):
-                            rna_id = search_results["results"][0].get("rnacentral_id")
+                        results = search_results.get("results", [])
+                        if results:
+                            # First, try to find an exact sequence match
+                            exact_match = None
+                            for result in results:
+                                result_seq = result.get("sequence", "")
+                                if result_seq == seq:
+                                    exact_match = result
+                                    break
+                            
+                            # Use exact match if found, otherwise use first result
+                            target_result = exact_match if exact_match else results[0]
+                            rna_id = target_result.get("rnacentral_id")
+                            
                             if rna_id:
-                                return await self.get_by_rna_id(rna_id)
+                                # Try to get detailed info, but fall back to search result if it fails
+                                try:
+                                    detailed_info = await self.get_by_rna_id(rna_id)
+                                    if detailed_info:
+                                        return detailed_info
+                                except Exception as e:
+                                    logger.debug("Failed to get detailed info for %s: %s, using search result", rna_id, e)
+                                
+                                # Fall back to using search result data
+                                return self._rna_data_to_dict(rna_id, target_result)
                         logger.info("No results found for sequence.")
                         return None
-                    raise Exception(f"HTTP {resp.status}: {await resp.text()}")
+                    error_text = await resp.text()
+                    logger.error("HTTP %d error for sequence search: %s", resp.status, error_text[:200])
+                    raise Exception(f"HTTP {resp.status}: {error_text}")
+        except aiohttp.ClientError as e:
+            logger.error("Network error searching for sequence: %s", e)
+            return None
         except Exception as e:  # pylint: disable=broad-except
             logger.error("Sequence search failed: %s", e)
             return None
@@ -154,11 +207,13 @@ async def search_by_sequence(self, sequence: str) -> Optional[dict]:
         reraise=True,
     )
     async def search(
-        self, query: str, **kwargs
+        self, query: str, threshold: float = 0.7, **kwargs
     ) -> Optional[Dict]:
         """
         Search RNAcentral with either an RNAcentral ID, keyword, or RNA sequence.
         :param query: The search query (RNAcentral ID, keyword, or RNA sequence).
+        :param threshold: E-value threshold for sequence search.
+        Note: RNAcentral API uses its own similarity matching, this parameter is for interface consistency.
         :param kwargs: Additional keyword arguments (not used currently).
         :return: A dictionary containing the search results or None if not found.
         """
@@ -180,7 +235,7 @@ async def search(
             result = await self.get_by_rna_id(query)
         else:
             # otherwise treat as keyword
-            result = await self.search_by_keyword(query)
+            result = await self.get_best_hit(query)
 
         if result:
             result["_search_query"] = query

resources/input_examples/search_rna_demo.jsonl

@@ -0,0 +1,4 @@
+{"type": "text", "content": "hsa-let-7a-1"}
+{"type": "text", "content": "URS0000123456"}
+{"type": "text", "content": "URS0000000001"}
+{"type": "text", "content": "CUCCUUUGACGUUAGCGGCGGACGGGUUAGUAACACGUGGGUAACCUACCUAUAAGACUGGGAUAACUUCGGGAAACCGGAGCUAAUACCGGAUAAUAUUUCGAACCGCAUGGUUCGAUAGUGAAAGAUGGUUUUGCUAUCACUUAUAGAUGGACCCGCGCCGUAUUAGCUAGUUGGUAAGGUAACGGCUUACCAAGGCGACGAUACGUAGCCGACCUGAGAGGGUGAUCGGCCACACUGGAACUGAGACACGGUCCAGACUCCUACGGGAGGCAGCAGGGG"}