| import os |
| from typing import List, Optional, Tuple, Type, Union |
|
|
| import aiohttp |
| import requests |
|
|
| from lagent.schema import ActionReturn, ActionStatusCode |
| from .base_action import AsyncActionMixin, BaseAction, tool_api |
| from .parser import BaseParser, JsonParser |
|
|
|
|
| class GoogleSearch(BaseAction): |
| """Wrapper around the Serper.dev Google Search API. |
| |
| To use, you should pass your serper API key to the constructor. |
| |
| Code is modified from lang-chain GoogleSerperAPIWrapper |
| (https://github.com/langchain-ai/langchain/blob/ba5f |
| baba704a2d729a4b8f568ed70d7c53e799bb/libs/langchain/ |
| langchain/utilities/google_serper.py) |
| |
| Args: |
| api_key (str): API KEY to use serper google search API, |
| You can create a free API key at https://serper.dev. |
| timeout (int): Upper bound of waiting time for a serper request. |
| search_type (str): Serper API support ['search', 'images', 'news', |
| 'places'] types of search, currently we only support 'search'. |
| description (dict): The description of the action. Defaults to ``None``. |
| parser (Type[BaseParser]): The parser class to process the |
| action's inputs and outputs. Defaults to :class:`JsonParser`. |
| """ |
| result_key_for_type = { |
| 'news': 'news', |
| 'places': 'places', |
| 'images': 'images', |
| 'search': 'organic', |
| } |
|
|
| def __init__( |
| self, |
| api_key: Optional[str] = None, |
| timeout: int = 5, |
| search_type: str = 'search', |
| description: Optional[dict] = None, |
| parser: Type[BaseParser] = JsonParser, |
| ): |
| super().__init__(description, parser) |
| api_key = os.environ.get('SERPER_API_KEY', api_key) |
| if api_key is None: |
| raise ValueError( |
| 'Please set Serper API key either in the environment ' |
| 'as SERPER_API_KEY or pass it as `api_key` parameter.') |
| self.api_key = api_key |
| self.timeout = timeout |
| self.search_type = search_type |
|
|
| @tool_api |
| def run(self, query: str, k: int = 10) -> ActionReturn: |
| """一个可以从谷歌搜索结果的API。当你需要对于一个特定问题找到简短明了的回答时，可以使用它。输入应该是一个搜索查询。 |
| |
| Args: |
| query (str): the search content |
| k (int): select first k results in the search results as response |
| """ |
| tool_return = ActionReturn(type=self.name) |
| status_code, response = self._search(query, k=k) |
| |
| if status_code == -1: |
| tool_return.errmsg = response |
| tool_return.state = ActionStatusCode.HTTP_ERROR |
| elif status_code == 200: |
| parsed_res = self._parse_results(response, k) |
| tool_return.result = [dict(type='text', content=str(parsed_res))] |
| tool_return.state = ActionStatusCode.SUCCESS |
| else: |
| tool_return.errmsg = str(status_code) |
| tool_return.state = ActionStatusCode.API_ERROR |
| return tool_return |
|
|
| def _parse_results(self, results: dict, k: int) -> Union[str, List[str]]: |
| """Parse the search results from Serper API. |
| |
| Args: |
| results (dict): The search content from Serper API |
| in json format. |
| |
| Returns: |
| List[str]: The parsed search results. |
| """ |
|
|
| snippets = [] |
|
|
| if results.get('answerBox'): |
| answer_box = results.get('answerBox', {}) |
| if answer_box.get('answer'): |
| return [answer_box.get('answer')] |
| elif answer_box.get('snippet'): |
| return [answer_box.get('snippet').replace('\n', ' ')] |
| elif answer_box.get('snippetHighlighted'): |
| return answer_box.get('snippetHighlighted') |
|
|
| if results.get('knowledgeGraph'): |
| kg = results.get('knowledgeGraph', {}) |
| title = kg.get('title') |
| entity_type = kg.get('type') |
| if entity_type: |
| snippets.append(f'{title}: {entity_type}.') |
| description = kg.get('description') |
| if description: |
| snippets.append(description) |
| for attribute, value in kg.get('attributes', {}).items(): |
| snippets.append(f'{title} {attribute}: {value}.') |
|
|
| for result in results[self.result_key_for_type[ |
| self.search_type]][:k]: |
| if 'snippet' in result: |
| snippets.append(result['snippet']) |
| for attribute, value in result.get('attributes', {}).items(): |
| snippets.append(f'{attribute}: {value}.') |
|
|
| if len(snippets) == 0: |
| return ['No good Google Search Result was found'] |
| return snippets |
|
|
| def _search(self, |
| search_term: str, |
| search_type: Optional[str] = None, |
| **kwargs) -> Tuple[int, Union[dict, str]]: |
| """HTTP requests to Serper API. |
| |
| Args: |
| search_term (str): The search query. |
| search_type (str): search type supported by Serper API, |
| default to 'search'. |
| |
| Returns: |
| tuple: the return value is a tuple contains: |
| - status_code (int): HTTP status code from Serper API. |
| - response (dict): response context with json format. |
| """ |
| headers = { |
| 'X-API-KEY': self.api_key or '', |
| 'Content-Type': 'application/json', |
| } |
| params = { |
| 'q': search_term, |
| **{ |
| key: value |
| for key, value in kwargs.items() if value is not None |
| }, |
| } |
| try: |
| response = requests.post( |
| f'https://google.serper.dev/{search_type or self.search_type}', |
| headers=headers, |
| params=params, |
| timeout=self.timeout) |
| except Exception as e: |
| return -1, str(e) |
| return response.status_code, response.json() |
|
|
|
|
| class AsyncGoogleSearch(AsyncActionMixin, GoogleSearch): |
| """Wrapper around the Serper.dev Google Search API. |
| |
| To use, you should pass your serper API key to the constructor. |
| |
| Code is modified from lang-chain GoogleSerperAPIWrapper |
| (https://github.com/langchain-ai/langchain/blob/ba5f |
| baba704a2d729a4b8f568ed70d7c53e799bb/libs/langchain/ |
| langchain/utilities/google_serper.py) |
| |
| Args: |
| api_key (str): API KEY to use serper google search API, |
| You can create a free API key at https://serper.dev. |
| timeout (int): Upper bound of waiting time for a serper request. |
| search_type (str): Serper API support ['search', 'images', 'news', |
| 'places'] types of search, currently we only support 'search'. |
| description (dict): The description of the action. Defaults to ``None``. |
| parser (Type[BaseParser]): The parser class to process the |
| action's inputs and outputs. Defaults to :class:`JsonParser`. |
| """ |
|
|
| @tool_api |
| async def run(self, query: str, k: int = 10) -> ActionReturn: |
| """一个可以从谷歌搜索结果的API。当你需要对于一个特定问题找到简短明了的回答时，可以使用它。输入应该是一个搜索查询。 |
| |
| Args: |
| query (str): the search content |
| k (int): select first k results in the search results as response |
| """ |
| tool_return = ActionReturn(type=self.name) |
| status_code, response = await self._search(query, k=k) |
| |
| if status_code == -1: |
| tool_return.errmsg = response |
| tool_return.state = ActionStatusCode.HTTP_ERROR |
| elif status_code == 200: |
| parsed_res = self._parse_results(response) |
| tool_return.result = [dict(type='text', content=str(parsed_res))] |
| tool_return.state = ActionStatusCode.SUCCESS |
| else: |
| tool_return.errmsg = str(status_code) |
| tool_return.state = ActionStatusCode.API_ERROR |
| return tool_return |
|
|
| async def _search(self, |
| search_term: str, |
| search_type: Optional[str] = None, |
| **kwargs) -> Tuple[int, Union[dict, str]]: |
| """HTTP requests to Serper API. |
| |
| Args: |
| search_term (str): The search query. |
| search_type (str): search type supported by Serper API, |
| default to 'search'. |
| |
| Returns: |
| tuple: the return value is a tuple contains: |
| - status_code (int): HTTP status code from Serper API. |
| - response (dict): response context with json format. |
| """ |
| headers = { |
| 'X-API-KEY': self.api_key or '', |
| 'Content-Type': 'application/json', |
| } |
| params = { |
| 'q': search_term, |
| **{ |
| key: value |
| for key, value in kwargs.items() if value is not None |
| }, |
| } |
| timeout = aiohttp.ClientTimeout(total=self.timeout) |
| async with aiohttp.ClientSession(timeout=timeout) as session: |
| try: |
| async with session.post( |
| f'https://google.serper.dev/{search_type or self.search_type}', |
| headers=headers, |
| params=params) as resp: |
| code, ret = resp.status, await resp.json() |
| except aiohttp.ClientError as e: |
| code, ret = -1, str(e) |
| return code, ret |
|
|
