feat: add google search

Author: ChenZiHong-Gavin

graphgen/configs/graphgen_config.yaml

@@ -14,5 +14,5 @@ traverse_strategy:
   loss_strategy: only_edge
 search:
   enabled: true
-  search_types: ["wikipedia", "google"]
+  search_types: ["google"]
 re_judge: false

graphgen/graphgen.py

@@ -237,9 +237,8 @@ async def async_search(self):
                 "[Search] Found %d entities to search", len(new_search_entities)
             )
             _add_search_data = await search_all(
-                llm_client=self.synthesizer_llm_client,
                 search_types=self.search_config["search_types"],
-                kg_instance=self.graph_storage,
+                search_entities=new_search_entities,
             )
             if _add_search_data:
                 await self.search_storage.upsert(_add_search_data)

graphgen/models/__init__.py

@@ -1,12 +1,12 @@
-from graphgen.models.search.kg.wiki_search import WikiSearch
-
 from .evaluate.length_evaluator import LengthEvaluator
 from .evaluate.mtld_evaluator import MTLDEvaluator
 from .evaluate.reward_evaluator import RewardEvaluator
 from .evaluate.uni_evaluator import UniEvaluator
 from .llm.openai_model import OpenAIModel
 from .llm.tokenizer import Tokenizer
 from .llm.topk_token_model import Token, TopkTokenModel
+from .search.kg.wiki_search import WikiSearch
+from .search.web.google_search import GoogleSearch
 from .storage.json_storage import JsonKVStorage
 from .storage.networkx_storage import NetworkXStorage
 from .strategy.travserse_strategy import TraverseStrategy
@@ -25,6 +25,7 @@
     "JsonKVStorage",
     # search models
     "WikiSearch",
+    "GoogleSearch",
     # evaluate models
     "TextPair",
     "LengthEvaluator",

graphgen/models/search/kg/wiki_search.py

@@ -14,9 +14,9 @@ def set_language(language: str):
         assert language in ["en", "zh"], "Only support English and Chinese"
         set_lang(language)
 
-    async def search(self, query: str) -> Union[List[str], None]:
+    async def search(self, query: str, num_results: int = 1) -> Union[List[str], None]:
         self.set_language(detect_main_language(query))
-        return wikipedia.search(query)
+        return wikipedia.search(query, results=num_results, suggestion=False)
 
     async def summary(self, query: str) -> Union[str, None]:
         self.set_language(detect_main_language(query))

graphgen/models/search/web/__init__.py

@@ -0,0 +1,45 @@
+from dataclasses import dataclass
+
+import requests
+from fastapi import HTTPException
+
+from graphgen.utils import logger
+
+GOOGLE_SEARCH_ENDPOINT = "https://customsearch.googleapis.com/customsearch/v1"
+
+
+@dataclass
+class GoogleSearch:
+    def __init__(self, subscription_key: str, cx: str):
+        """
+        Initialize the Google Search client with the subscription key and custom search engine ID.
+        :param subscription_key: Your Google API subscription key.
+        :param cx: Your custom search engine ID.
+        """
+        self.subscription_key = subscription_key
+        self.cx = cx
+
+    def search(self, query: str, num_results: int = 1):
+        """
+        Search with Google 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 = {
+            "key": self.subscription_key,
+            "cx": self.cx,
+            "q": query,
+            "num": num_results,
+        }
+        response = requests.get(GOOGLE_SEARCH_ENDPOINT, 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["items"][:num_results]
+        except KeyError:
+            logger.error("Error encountered: %s", json_content)
+            return []
+        return contexts

graphgen/models/search/web/google_search.py

@@ -1,81 +1,58 @@
 from tqdm.asyncio import tqdm_asyncio as tqdm_async
 
-from graphgen.models import NetworkXStorage, OpenAIModel, WikiSearch
-from graphgen.templates import SEARCH_JUDGEMENT_PROMPT
+from graphgen.models import WikiSearch
 from graphgen.utils import logger
 
 
 async def _process_single_entity(
     entity_name: str,
-    description: str,
-    llm_client: OpenAIModel,
     wiki_search_client: WikiSearch,
-) -> tuple[str, None] | tuple[str, str]:
+) -> str | None:
     """
-    Process single entity
-
+    Process single entity by searching Wikipedia
+    :param entity_name
+    :param wiki_search_client
+    :return: summary of the entity or None if not found
     """
     search_results = await wiki_search_client.search(entity_name)
     if not search_results:
-        return entity_name, None
-    examples = "\n".join(SEARCH_JUDGEMENT_PROMPT["EXAMPLES"])
-    search_results.append("None of the above")
+        return None
 
-    search_results_str = "\n".join(
-        [f"{i + 1}. {sr}" for i, sr in enumerate(search_results)]
-    )
-    prompt = SEARCH_JUDGEMENT_PROMPT["TEMPLATE"].format(
-        examples=examples,
-        entity_name=entity_name,
-        description=description,
-        search_results=search_results_str,
-    )
-    response = await llm_client.generate_answer(prompt)
+    summary = None
     try:
-        response = response.strip()
-        response = int(response)
-        if response < 1 or response >= len(search_results):
-            response = None
-        else:
-            response = await wiki_search_client.summary(search_results[response - 1])
-    except ValueError:
-        response = None
-
-    logger.info(
-        "Entity %s search result: %s response: %s",
-        entity_name,
-        str(search_results),
-        response,
-    )
+        summary = await wiki_search_client.summary(search_results[-1])
+        logger.info(
+            "Entity %s search result: %s summary: %s",
+            entity_name,
+            str(search_results),
+            summary,
+        )
+    except Exception as e:  # pylint: disable=broad-except
+        logger.error("Error processing entity %s: %s", entity_name, str(e))
 
-    return entity_name, response
+    return summary
 
 
 async def search_wikipedia(
-    llm_client: OpenAIModel,
     wiki_search_client: WikiSearch,
-    kg_instance: NetworkXStorage,
+    entities: set[str],
 ) -> dict:
     """
     Search wikipedia for entities
 
-    :param llm_client: LLM model
     :param wiki_search_client: wiki search client
-    :param kg_instance: knowledge graph instance
+    :param entities: list of entities to search
     :return: nodes with search results
     """
-    nodes = await kg_instance.get_all_nodes()
-    nodes = list(nodes)
     wiki_data = {}
 
-    async for node in tqdm_async(nodes, desc="Searching Wikipedia", total=len(nodes)):
-        entity_name = node[0].strip('"')
-        description = node[1]["description"]
+    async for entity in tqdm_async(
+        entities, desc="Searching Wikipedia", total=len(entities)
+    ):
         try:
-            entity, summary = await _process_single_entity(
-                entity_name, description, llm_client, wiki_search_client
-            )
-            wiki_data[entity] = summary
+            entity, summary = await _process_single_entity(entity, wiki_search_client)
+            if summary:
+                wiki_data[entity] = summary
         except Exception as e:  # pylint: disable=broad-except
-            logger.error("Error processing entity %s: %s", entity_name, str(e))
+            logger.error("Error processing entity %s: %s", entity, str(e))
     return wiki_data

graphgen/operators/search/kg/search_wikipedia.py

@@ -1,19 +1,27 @@
-from graphgen.models import NetworkXStorage, OpenAIModel
+"""
+To use Google Web Search API,
+follow the instructions [here](https://developers.google.com/custom-search/v1/overview)
+to get your Google search api key.
+
+To use Bing Web Search API,
+follow the instructions [here](https://www.microsoft.com/en-us/bing/apis/bing-web-search-api)
+and obtain your Bing subscription key.
+"""
+
+import os
+
 from graphgen.utils import logger
 
 
 async def search_all(
-    llm_client: OpenAIModel, search_types: dict, kg_instance: NetworkXStorage
+    search_types: dict, search_entities: set[str]
 ) -> dict[str, dict[str, str]]:
     """
-    :param llm_client
     :param search_types
-    :param kg_instance
+    :param search_entities: list of entities to search
     :return: nodes with search results
     """
 
-    # 增量建图时，只需要搜索新增实体
-
     results = {}
 
     for search_type in search_types:
@@ -23,16 +31,25 @@ async def search_all(
 
             wiki_search_client = WikiSearch()
 
-            wiki_results = await search_wikipedia(
-                llm_client, wiki_search_client, kg_instance
-            )
+            wiki_results = await search_wikipedia(wiki_search_client, search_entities)
             for entity_name, description in wiki_results.items():
                 if description:
                     results[entity_name] = {"wikipedia": description}
-        # elif search_type == "google":
-        #     from graphgen.operators.search.web.search_google import search_google
-        #     return await search_google(llm_client, kg_instance)
-        #
+        elif search_type == "google":
+            from graphgen.models import GoogleSearch
+            from graphgen.operators.search.web.search_google import search_google
+
+            google_search_client = GoogleSearch(
+                subscription_key=os.environ["GOOGLE_SEARCH_API_KEY"],
+                cx=os.environ["GOOGLE_SEARCH_CX"],
+            )
+
+            google_results = await search_google(google_search_client, search_entities)
+            for entity_name, description in google_results.items():
+                if description:
+                    results[entity_name] = results.get(entity_name, {})
+                    results[entity_name]["google"] = description
+
         # elif search_type == "bing":
         #     from graphgen.operators.search.web.search_bing import search_bing
         #     return await search_bing(llm_client, kg_instance)

graphgen/operators/search/search_all.py

@@ -0,0 +1,10 @@
+BING_SEARCH_V7_ENDPOINT = "https://api.bing.microsoft.com/v7.0/search"
+BING_MKT = "en-US"
+
+
+async def search_bing():
+    """
+    Search with Bing and return the contexts.
+    :return:
+    """
+    raise NotImplementedError("Bing search is not implemented yet.")

graphgen/operators/search/web/search_bing.py

@@ -0,0 +1,49 @@
+import trafilatura
+from tqdm.asyncio import tqdm_asyncio as tqdm_async
+
+from graphgen.models import GoogleSearch
+from graphgen.utils import logger
+
+
+async def _process_single_entity(
+    entity_name: str, google_search_client: GoogleSearch
+) -> str | None:
+    search_results = google_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["link"])
+    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_google(
+    google_search_client: GoogleSearch,
+    entities: set[str],
+) -> dict:
+    """
+    Search with Google and return the contexts.
+    :param google_search_client: Google search client
+    :param entities: list of entities to search
+    :return:
+    """
+    google_data = {}
+
+    async for entity in tqdm_async(
+        entities, desc="Searching Google", total=len(entities)
+    ):
+        try:
+            summary = await _process_single_entity(entity, google_search_client)
+            if summary:
+                google_data[entity] = summary
+        except Exception as e:  # pylint: disable=broad-except
+            logger.error("Error processing entity %s: %s", entity, str(e))
+    return google_data