feat: uniprot search

Author: ChenZiHong-Gavin

graphgen/models/__init__.py

@@ -5,7 +5,9 @@
 from .llm.openai_model import OpenAIModel
 from .llm.tokenizer import Tokenizer
 from .llm.topk_token_model import Token, TopkTokenModel
+from .search.db.uniprot_search import UniProtSearch
 from .search.kg.wiki_search import WikiSearch
+from .search.web.bing_search import BingSearch
 from .search.web.google_search import GoogleSearch
 from .storage.json_storage import JsonKVStorage
 from .storage.networkx_storage import NetworkXStorage
@@ -26,6 +28,8 @@
     # search models
     "WikiSearch",
     "GoogleSearch",
+    "BingSearch",
+    "UniProtSearch",
     # evaluate models
     "TextPair",
     "LengthEvaluator",

…en/operators/search/db/search_mongodb.py graphgen/models/search/db/__init__.pygraphgen/operators/search/db/search_mongodb.py renamed to graphgen/models/search/db/__init__.py graphgen/operators/search/db/search_mongodb.py renamed to graphgen/models/search/db/__init__.py

@@ -0,0 +1,64 @@
+from dataclasses import dataclass
+
+import requests
+from fastapi import HTTPException
+
+from graphgen.utils import logger
+
+UNIPROT_BASE = "https://rest.uniprot.org/uniprotkb/search"
+
+
+@dataclass
+class UniProtSearch:
+    """
+    UniProt Search client to search with UniProt.
+    1) Get the protein by accession number.
+    2) Search with keywords or protein names.
+    """
+
+    def get_entry(self, accession: str) -> dict:
+        """
+        Get the UniProt entry by accession number(e.g., P04637).
+        """
+        url = f"{UNIPROT_BASE}/{accession}.json"
+        return self._safe_get(url).json()
+
+    def search(
+        self,
+        query: str,
+        *,
+        size: int = 10,
+        cursor: str = None,
+        fields: list[str] = None,
+    ) -> dict:
+        """
+        Search UniProt with a query string.
+        :param query: The search query.
+        :param size: The number of results to return.
+        :param cursor: The cursor for pagination.
+        :param fields: The fields to return in the response.
+        :return: A dictionary containing the search results.
+        """
+        params = {
+            "query": query,
+            "size": size,
+        }
+        if cursor:
+            params["cursor"] = cursor
+        if fields:
+            params["fields"] = ",".join(fields)
+        url = UNIPROT_BASE
+        return self._safe_get(url, params=params).json()
+
+    @staticmethod
+    def _safe_get(url: str, params: dict = None) -> requests.Response:
+        r = requests.get(
+            url,
+            params=params,
+            headers={"Accept": "application/json"},
+            timeout=10,
+        )
+        if not r.ok:
+            logger.error("Search engine error: %s", r.text)
+            raise HTTPException(r.status_code, "Search engine error.")
+        return r

graphgen/models/search/db/uniprot_search.py

@@ -0,0 +1,43 @@
+from dataclasses import dataclass
+
+import requests
+from fastapi import HTTPException
+
+from graphgen.utils import logger
+
+BING_SEARCH_V7_ENDPOINT = "https://api.bing.microsoft.com/v7.0/search"
+BING_MKT = "en-US"
+
+
+@dataclass
+class BingSearch:
+    """
+    Bing Search client to search with Bing.
+    """
+
+    subscription_key: str
+
+    def search(self, query: str, num_results: int = 1):
+        """
+        Search with Bing and return the contexts.
+        :param query: The search query.
+        :param num_results: The number of results to return.
+        :return: A list of search results.
+        """
+        params = {"q": query, "mkt": BING_MKT, "count": num_results}
+        response = requests.get(
+            BING_SEARCH_V7_ENDPOINT,
+            headers={"Ocp-Apim-Subscription-Key": self.subscription_key},
+            params=params,
+            timeout=10,
+        )
+        if not response.ok:
+            logger.error("Search engine error: %s", response.text)
+            raise HTTPException(response.status_code, "Search engine error.")
+        json_content = response.json()
+        try:
+            contexts = json_content["webPages"]["value"][:num_results]
+        except KeyError:
+            logger.error("Error encountered: %s", json_content)
+            return []
+        return contexts

graphgen/models/search/web/bing_search.py

@@ -49,10 +49,31 @@ async def search_all(
                 if description:
                     results[entity_name] = results.get(entity_name, {})
                     results[entity_name]["google"] = description
+        elif search_type == "bing":
+            from graphgen.models import BingSearch
+            from graphgen.operators.search.web.search_bing import search_bing
 
-        # elif search_type == "bing":
-        #     from graphgen.operators.search.web.search_bing import search_bing
-        #     return await search_bing(llm_client, kg_instance)
+            bing_search_client = BingSearch(
+                subscription_key=os.environ["BING_SEARCH_API_KEY"]
+            )
+
+            bing_results = await search_bing(bing_search_client, search_entities)
+            for entity_name, description in bing_results.items():
+                if description:
+                    results[entity_name] = results.get(entity_name, {})
+                    results[entity_name]["bing"] = description
+        elif search_type == "uniprot":
+            # from graphgen.models import UniProtSearch
+            # from graphgen.operators.search.db.search_uniprot import search_uniprot
+            #
+            # uniprot_search_client = UniProtSearch()
+            #
+            # uniprot_results = await search_uniprot(
+            #     uniprot_search_client, search_entities
+            # )
+            raise NotImplementedError(
+                "Processing of UniProt search results is not implemented yet."
+            )
 
         else:
             logger.error("Search type %s is not supported yet.", search_type)

graphgen/operators/search/kg/search_google_kg.py

@@ -1,10 +1,53 @@
-BING_SEARCH_V7_ENDPOINT = "https://api.bing.microsoft.com/v7.0/search"
-BING_MKT = "en-US"
+import trafilatura
+from tqdm.asyncio import tqdm_asyncio as tqdm_async
 
+from graphgen.models import BingSearch
+from graphgen.utils import logger
 
-async def search_bing():
+
+async def _process_single_entity(
+    entity_name: str, bing_search_client: BingSearch
+) -> str | None:
+    """
+    Process single entity by searching Bing.
+    :param entity_name: The name of the entity to search.
+    :param bing_search_client: The Bing search client.
+    :return: Summary of the entity or None if not found.
+    """
+    search_results = bing_search_client.search(entity_name)
+    if not search_results:
+        return None
+
+    # Get more details from the first search result
+    first_result = search_results[0]
+    content = trafilatura.fetch_url(first_result["url"])
+    summary = trafilatura.extract(content, include_comments=False, include_links=False)
+    summary = summary.strip()
+    logger.info(
+        "Entity %s search result: %s",
+        entity_name,
+        summary,
+    )
+    return summary
+
+
+async def search_bing(
+    bing_search_client: BingSearch,
+    entities: set[str],
+) -> dict[str, str]:
     """
     Search with Bing and return the contexts.
     :return:
     """
-    raise NotImplementedError("Bing search is not implemented yet.")
+    bing_data = {}
+
+    async for entity in tqdm_async(
+        entities, desc="Searching Bing", total=len(entities)
+    ):
+        try:
+            summary = await _process_single_entity(entity, bing_search_client)
+            if summary:
+                bing_data[entity] = summary
+        except Exception as e:  # pylint: disable=broad-except
+            logger.error("Error processing entity %s: %s", entity, str(e))
+    return bing_data