|
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 |
|
|
