fix: fix async search

Author: ChenZiHong-Gavin

graphgen/configs/search_config.yaml

@@ -1,8 +1,11 @@
 pipeline:
   - name: read
     params:
-      input_file: resources/input_examples/search_demo.jsonl # input file path, support json, jsonl, txt, pdf. See resources/input_examples for examples
+      input_file:  resources/input_examples/search_demo.json # input file path, support json, jsonl, txt, pdf. See resources/input_examples for examples
 
   - name: search
     params:
       data_sources: [uniprot] # data source for searcher, support: wikipedia, google, uniprot
+      uniprot_params:
+        use_local_blast: true # whether to use local blast for uniprot search
+        local_blast_db: /your_path/uniprot_sprot

graphgen/graphgen.py

@@ -68,7 +68,7 @@ def __init__(
             self.working_dir, namespace="graph"
         )
         self.search_storage: JsonKVStorage = JsonKVStorage(
-            self.working_dir, namespace="searcher"
+            self.working_dir, namespace="search"
         )
         self.rephrase_storage: JsonKVStorage = JsonKVStorage(
             self.working_dir, namespace="rephrase"
@@ -190,7 +190,7 @@ async def search(self, search_config: Dict):
             return
         search_results = await search_all(
             seed_data=seeds,
-            **search_config,
+            search_config=search_config,
         )
 
         _add_search_keys = await self.search_storage.filter_keys(

graphgen/models/searcher/db/uniprot_searcher.py

@@ -1,7 +1,10 @@
+import asyncio
 import os
 import re
 import subprocess
 import tempfile
+from concurrent.futures import ThreadPoolExecutor
+from functools import lru_cache
 from io import StringIO
 from typing import Dict, Optional
 
@@ -19,6 +22,11 @@
 from graphgen.utils import logger
 
 
+@lru_cache(maxsize=None)
+def _get_pool():
+    return ThreadPoolExecutor(max_workers=10)
+
+
 class UniProtSearch(BaseSearcher):
     """
     UniProt Search client to searcher with UniProt.
@@ -215,20 +223,26 @@ async def search(
         query = query.strip()
 
         logger.debug("UniProt searcher query: %s", query)
+
+        loop = asyncio.get_running_loop()
+
         # check if fasta sequence
         if query.startswith(">") or re.fullmatch(
             r"[ACDEFGHIKLMNPQRSTVWY\s]+", query, re.I
         ):
-            result = self.get_by_fasta(query, threshold)
+            coro = loop.run_in_executor(
+                _get_pool(), self.get_by_fasta, query, threshold
+            )
 
         # check if accession number
         elif re.fullmatch(r"[A-NR-Z0-9]{6,10}", query, re.I):
-            result = self.get_by_accession(query)
+            coro = loop.run_in_executor(_get_pool(), self.get_by_accession, query)
 
         else:
             # otherwise treat as keyword
-            result = self.get_best_hit(query)
+            coro = loop.run_in_executor(_get_pool(), self.get_best_hit, query)
 
+        result = await coro
         if result:
             result["_search_query"] = query
         return result

graphgen/operators/search/search_all.py

@@ -14,22 +14,25 @@
 
 async def search_all(
     seed_data: dict,
-    data_sources: list[str],
+    search_config: dict,
 ) -> dict:
     """
     Perform searches across multiple search types and aggregate the results.
     :param seed_data: A dictionary containing seed data with entity names.
-    :param data_sources: A list of search types to perform (e.g., "wikipedia", "google", "bing", "uniprot").
+    :param search_config: A dictionary specifying which data sources to use for searching.
     :return: A dictionary with
     """
 
     results = {}
+    data_sources = search_config.get("data_sources", [])
 
     for data_source in data_sources:
         if data_source == "uniprot":
             from graphgen.models import UniProtSearch
 
-            uniprot_search_client = UniProtSearch()
+            uniprot_search_client = UniProtSearch(
+                **search_config.get("uniprot_params", {})
+            )
 
             data = list(seed_data.values())
             data = [d["content"] for d in data if "content" in d]