Spaces:
Running
Running
File size: 9,602 Bytes
e679d69 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 |
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)
# 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`.
"""
@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)
# 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
|