chore: refactor the serpapi's google search tool (#4834)

api/core/tools/provider/builtin/google/tools/google_search.py

@@ -1,39 +1,20 @@
-import os
-import sys
 from typing import Any, Union
 
-from serpapi import GoogleSearch
+import requests
 
 from core.tools.entities.tool_entities import ToolInvokeMessage
 from core.tools.tool.builtin_tool import BuiltinTool
 
-
+SERP_API_URL = "https://serpapi.com/search"
-class HiddenPrints:
-    """Context manager to hide prints."""
-
-    def __enter__(self) -> None:
-        """Open file to pipe stdout to."""
-        self._original_stdout = sys.stdout
-        sys.stdout = open(os.devnull, "w")
-
-    def __exit__(self, *_: Any) -> None:
-        """Close file that stdout was piped to."""
-        sys.stdout.close()
-        sys.stdout = self._original_stdout
 
 
 class SerpAPI:
     """
     SerpAPI tool provider.
     """
-
-    search_engine: Any  #: :meta private:
-    serpapi_api_key: str = None
-
     def __init__(self, api_key: str) -> None:
         """Initialize SerpAPI tool provider."""
         self.serpapi_api_key = api_key
-        self.search_engine = GoogleSearch
 
     def run(self, query: str, **kwargs: Any) -> str:
         """Run query through SerpAPI and parse result."""
@@ -43,117 +24,76 @@ class SerpAPI:
     def results(self, query: str) -> dict:
         """Run query through SerpAPI and return the raw result."""
         params = self.get_params(query)
-        with HiddenPrints():
+        response = requests.get(url=SERP_API_URL, params=params)
-            search = self.search_engine(params)
+        response.raise_for_status()
-            res = search.get_dict()
+        return response.json()
-        return res
 
     def get_params(self, query: str) -> dict[str, str]:
         """Get parameters for SerpAPI."""
-        _params = {
+        params = {
             "api_key": self.serpapi_api_key,
             "q": query,
-        }
-        params = {
             "engine": "google",
             "google_domain": "google.com",
             "gl": "us",
-            "hl": "en",
+            "hl": "en"
-            **_params
         }
         return params
 
     @staticmethod
     def _process_response(res: dict, typ: str) -> str:
-        """Process response from SerpAPI."""
+        """
-        if "error" in res.keys():
+        Process response from SerpAPI.
+        SerpAPI doc: https://serpapi.com/search-api
+        Google search main results are called organic results
+        """
+        if "error" in res:
             raise ValueError(f"Got error from SerpAPI: {res['error']}")
-        
+        toret = ""
         if typ == "text":
-            toret = ""
+            if "knowledge_graph" in res and "description" in res["knowledge_graph"]:
-            if "answer_box" in res.keys() and type(res["answer_box"]) == list:
+                toret += res["knowledge_graph"]["description"] + "\n"
-                res["answer_box"] = res["answer_box"][0] + "\n"
+            if "organic_results" in res:
-            if "answer_box" in res.keys() and "answer" in res["answer_box"].keys():
+                snippets = [
-                toret += res["answer_box"]["answer"] + "\n"
+                    f"content: {item.get('snippet')}\nlink: {item.get('link')}"
-            if "answer_box" in res.keys() and "snippet" in res["answer_box"].keys():
-                toret += res["answer_box"]["snippet"] + "\n"
-            if (
-                "answer_box" in res.keys()
-                and "snippet_highlighted_words" in res["answer_box"].keys()
-            ):
-                for item in res["answer_box"]["snippet_highlighted_words"]:
-                    toret += item + "\n"
-            if (
-                "sports_results" in res.keys()
-                and "game_spotlight" in res["sports_results"].keys()
-            ):
-                toret += res["sports_results"]["game_spotlight"] + "\n"
-            if (
-                "shopping_results" in res.keys()
-                and "title" in res["shopping_results"][0].keys()
-            ):
-                toret += res["shopping_results"][:3] + "\n"
-            if (
-                "knowledge_graph" in res.keys()
-                and "description" in res["knowledge_graph"].keys()
-            ):
-                toret = res["knowledge_graph"]["description"] + "\n"
-            if "snippet" in res["organic_results"][0].keys():
-                toret = "\n".join(
-                    f"content: {item['snippet']}\nlink: {item['link']}"
                     for item in res["organic_results"]
-                    if "snippet" in item and "link" in item
+                    if "snippet" in item
-                )
+                ]
-            if (
+                toret += "\n".join(snippets)
-                "images_results" in res.keys()
-                and "thumbnail" in res["images_results"][0].keys()
-            ):
-                thumbnails = [item["thumbnail"] for item in res["images_results"][:10]]
-                toret = thumbnails
-            if toret == "":
-                toret = "No good search result found"
         elif typ == "link":
-            if "knowledge_graph" in res.keys() and "title" in res["knowledge_graph"].keys() \
+            if "knowledge_graph" in res and "source" in res["knowledge_graph"]:
-                    and "description_link" in res["knowledge_graph"].keys():
+                toret += res["knowledge_graph"]["source"]["link"]
-                toret = res["knowledge_graph"]["description_link"]
+            elif "organic_results" in res:
-            elif "knowledge_graph" in res.keys() and "see_results_about" in res["knowledge_graph"].keys() \
+                links = [
-                and len(res["knowledge_graph"]["see_results_about"]) > 0:
+                    f"[{item['title']}]({item['link']})\n"
-                see_result_about = res["knowledge_graph"]["see_results_about"]
+                    for item in res["organic_results"]
-                toret = ""
+                    if "title" in item and "link" in item
-                for item in see_result_about:
+                ]
-                    if "name" not in item.keys() or "link" not in item.keys():
+                toret += "\n".join(links)
-                        continue
+            elif "related_questions" in res:
-                    toret += f"[{item['name']}]({item['link']})\n"
+                questions = [
-            elif "organic_results" in res.keys() and len(res["organic_results"]) > 0:
+                    f"[{item['question']}]({item['link']})\n"
-                organic_results = res["organic_results"]
+                    for item in res["related_questions"]
-                toret = ""
+                    if "question" in item and "link" in item
-                for item in organic_results:
+                ]
-                    if "title" not in item.keys() or "link" not in item.keys():
+                toret += "\n".join(questions)
-                        continue
+            elif "related_searches" in res:
-                    toret += f"[{item['title']}]({item['link']})\n"
+                searches = [
-            elif "related_questions" in res.keys() and len(res["related_questions"]) > 0:
+                    f"[{item['query']}]({item['link']})\n"
-                related_questions = res["related_questions"]
+                    for item in res["related_searches"]
-                toret = ""
+                    if "query" in item and "link" in item
-                for item in related_questions:
+                ]
-                    if "question" not in item.keys() or "link" not in item.keys():
+                toret += "\n".join(searches)
-                        continue
+        if not toret:
-                    toret += f"[{item['question']}]({item['link']})\n"
+            toret = "No good search result found"
-            elif "related_searches" in res.keys() and len(res["related_searches"]) > 0:
-                related_searches = res["related_searches"]
-                toret = ""
-                for item in related_searches:
-                    if "query" not in item.keys() or "link" not in item.keys():
-                        continue
-                    toret += f"[{item['query']}]({item['link']})\n"
-            else:
-                toret = "No good search result found"
         return toret
 
+
 class GoogleSearchTool(BuiltinTool):
-    def _invoke(self, 
+    def _invoke(self,
                 user_id: str,
-               tool_parameters: dict[str, Any], 
+                tool_parameters: dict[str, Any],
-        ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]:
+                ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]:
         """
             invoke tools
         """
@@ -164,4 +104,3 @@ class GoogleSearchTool(BuiltinTool):
         if result_type == 'text':
             return self.create_text_message(text=result)
         return self.create_link_message(link=result)
-    

api/requirements.txt

@@ -41,7 +41,6 @@ google-api-python-client==2.90.0
 google-auth==2.29.0
 google-auth-httplib2==0.2.0
 google-generativeai==0.5.0
-google-search-results==2.4.2
 googleapis-common-protos==1.63.0
 google-cloud-storage==2.16.0
 replicate~=0.22.0