Spaces:
Running
Running
| 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 | |
| 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) | |
| # convert search results to ToolReturn format | |
| 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`. | |
| """ | |
| 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) | |
| # convert search results to ToolReturn format | |
| 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 | |