| import asyncio |
| import hashlib |
| import hmac |
| import json |
| import logging |
| import random |
| import re |
| import time |
| import warnings |
| from concurrent.futures import ThreadPoolExecutor, as_completed |
| from datetime import datetime |
| from http.client import HTTPSConnection |
| from typing import List, Optional, Tuple, Type, Union |
|
|
| import aiohttp |
| import aiohttp.client_exceptions |
| import requests |
| from asyncache import cached as acached |
| from bs4 import BeautifulSoup |
| from cachetools import TTLCache, cached |
| from duckduckgo_search import DDGS, AsyncDDGS |
|
|
| from lagent.actions.base_action import AsyncActionMixin, BaseAction, tool_api |
| from lagent.actions.parser import BaseParser, JsonParser |
| from lagent.utils import async_as_completed |
|
|
|
|
| class BaseSearch: |
|
|
| def __init__(self, topk: int = 3, black_list: List[str] = None): |
| self.topk = topk |
| self.black_list = black_list |
|
|
| def _filter_results(self, results: List[tuple]) -> dict: |
| filtered_results = {} |
| count = 0 |
| for url, snippet, title in results: |
| if all(domain not in url |
| for domain in self.black_list) and not url.endswith('.pdf'): |
| filtered_results[count] = { |
| 'url': url, |
| 'summ': json.dumps(snippet, ensure_ascii=False)[1:-1], |
| 'title': title |
| } |
| count += 1 |
| if count >= self.topk: |
| break |
| return filtered_results |
|
|
|
|
| class DuckDuckGoSearch(BaseSearch): |
|
|
| def __init__(self, |
| topk: int = 3, |
| black_list: List[str] = [ |
| 'enoN', |
| 'youtube.com', |
| 'bilibili.com', |
| 'researchgate.net', |
| ], |
| **kwargs): |
| self.proxy = kwargs.get('proxy') |
| self.timeout = kwargs.get('timeout', 30) |
| super().__init__(topk, black_list) |
|
|
| @cached(cache=TTLCache(maxsize=100, ttl=600)) |
| def search(self, query: str, max_retry: int = 3) -> dict: |
| for attempt in range(max_retry): |
| try: |
| response = self._call_ddgs( |
| query, timeout=self.timeout, proxy=self.proxy) |
| return self._parse_response(response) |
| except Exception as e: |
| logging.exception(str(e)) |
| warnings.warn( |
| f'Retry {attempt + 1}/{max_retry} due to error: {e}') |
| time.sleep(random.randint(2, 5)) |
| raise Exception( |
| 'Failed to get search results from DuckDuckGo after retries.') |
|
|
| @acached(cache=TTLCache(maxsize=100, ttl=600)) |
| async def asearch(self, query: str, max_retry: int = 3) -> dict: |
| for attempt in range(max_retry): |
| try: |
| ddgs = AsyncDDGS(timeout=self.timeout, proxy=self.proxy) |
| response = await ddgs.atext(query.strip("'"), max_results=10) |
| return self._parse_response(response) |
| except Exception as e: |
| if isinstance(e, asyncio.TimeoutError): |
| logging.exception('Request to DDGS timed out.') |
| logging.exception(str(e)) |
| warnings.warn( |
| f'Retry {attempt + 1}/{max_retry} due to error: {e}') |
| await asyncio.sleep(random.randint(2, 5)) |
| raise Exception( |
| 'Failed to get search results from DuckDuckGo after retries.') |
|
|
| async def _async_call_ddgs(self, query: str, **kwargs) -> dict: |
| ddgs = DDGS(**kwargs) |
| try: |
| response = await asyncio.wait_for( |
| asyncio.to_thread(ddgs.text, query.strip("'"), max_results=10), |
| timeout=self.timeout) |
| return response |
| except asyncio.TimeoutError: |
| logging.exception('Request to DDGS timed out.') |
| raise |
|
|
| def _call_ddgs(self, query: str, **kwargs) -> dict: |
| loop = asyncio.new_event_loop() |
| asyncio.set_event_loop(loop) |
| try: |
| response = loop.run_until_complete( |
| self._async_call_ddgs(query, **kwargs)) |
| return response |
| finally: |
| loop.close() |
|
|
| def _parse_response(self, response: dict) -> dict: |
| raw_results = [] |
| for item in response: |
| raw_results.append( |
| (item['href'], item['description'] |
| if 'description' in item else item['body'], item['title'])) |
| return self._filter_results(raw_results) |
|
|
|
|
| class BingSearch(BaseSearch): |
|
|
| def __init__(self, |
| api_key: str, |
| region: str = 'zh-CN', |
| topk: int = 3, |
| black_list: List[str] = [ |
| 'enoN', |
| 'youtube.com', |
| 'bilibili.com', |
| 'researchgate.net', |
| ], |
| **kwargs): |
| self.api_key = api_key |
| self.market = region |
| self.proxy = kwargs.get('proxy') |
| super().__init__(topk, black_list) |
|
|
| @cached(cache=TTLCache(maxsize=100, ttl=600)) |
| def search(self, query: str, max_retry: int = 3) -> dict: |
| for attempt in range(max_retry): |
| try: |
| response = self._call_bing_api(query) |
| return self._parse_response(response) |
| except Exception as e: |
| logging.exception(str(e)) |
| warnings.warn( |
| f'Retry {attempt + 1}/{max_retry} due to error: {e}') |
| time.sleep(random.randint(2, 5)) |
| raise Exception( |
| 'Failed to get search results from Bing Search after retries.') |
|
|
| @acached(cache=TTLCache(maxsize=100, ttl=600)) |
| async def asearch(self, query: str, max_retry: int = 3) -> dict: |
| for attempt in range(max_retry): |
| try: |
| response = await self._async_call_bing_api(query) |
| return self._parse_response(response) |
| except Exception as e: |
| logging.exception(str(e)) |
| warnings.warn( |
| f'Retry {attempt + 1}/{max_retry} due to error: {e}') |
| await asyncio.sleep(random.randint(2, 5)) |
| raise Exception( |
| 'Failed to get search results from Bing Search after retries.') |
|
|
| def _call_bing_api(self, query: str) -> dict: |
| endpoint = 'https://api.bing.microsoft.com/v7.0/search' |
| params = {'q': query, 'mkt': self.market, 'count': f'{self.topk * 2}'} |
| headers = {'Ocp-Apim-Subscription-Key': self.api_key} |
| response = requests.get( |
| endpoint, headers=headers, params=params, proxies=self.proxy) |
| response.raise_for_status() |
| return response.json() |
|
|
| async def _async_call_bing_api(self, query: str) -> dict: |
| endpoint = 'https://api.bing.microsoft.com/v7.0/search' |
| params = {'q': query, 'mkt': self.market, 'count': f'{self.topk * 2}'} |
| headers = {'Ocp-Apim-Subscription-Key': self.api_key} |
| async with aiohttp.ClientSession(raise_for_status=True) as session: |
| async with session.get( |
| endpoint, |
| headers=headers, |
| params=params, |
| proxy=self.proxy and |
| (self.proxy.get('http') or self.proxy.get('https'))) as resp: |
| return await resp.json() |
|
|
| def _parse_response(self, response: dict) -> dict: |
| webpages = { |
| w['id']: w |
| for w in response.get('webPages', {}).get('value', []) |
| } |
| raw_results = [] |
|
|
| for item in response.get('rankingResponse', |
| {}).get('mainline', {}).get('items', []): |
| if item['answerType'] == 'WebPages': |
| webpage = webpages.get(item['value']['id']) |
| if webpage: |
| raw_results.append( |
| (webpage['url'], webpage['snippet'], webpage['name'])) |
| elif item['answerType'] == 'News' and item['value'][ |
| 'id'] == response.get('news', {}).get('id'): |
| for news in response.get('news', {}).get('value', []): |
| raw_results.append( |
| (news['url'], news['description'], news['name'])) |
|
|
| return self._filter_results(raw_results) |
|
|
|
|
| class BraveSearch(BaseSearch): |
| """ |
| Wrapper around the Brave Search API. |
| |
| To use, you should pass your Brave Search API key to the constructor. |
| |
| Args: |
| api_key (str): API KEY to use Brave Search API. |
| You can create a free API key at https://api.search.brave.com/app/keys. |
| search_type (str): Brave Search API supports ['web', 'news', 'images', 'videos'], |
| currently only supports 'news' and 'web'. |
| topk (int): The number of search results returned in response from API search results. |
| region (str): The country code string. Specifies the country where the search results come from. |
| language (str): The language code string. Specifies the preferred language for the search results. |
| extra_snippets (bool): Allows retrieving up to 5 additional snippets, which are alternative excerpts from the search results. |
| **kwargs: Any other parameters related to the Brave Search API. Find more details at |
| https://api.search.brave.com/app/documentation/web-search/get-started. |
| """ |
|
|
| def __init__(self, |
| api_key: str, |
| region: str = 'ALL', |
| language: str = 'zh-hans', |
| extra_snippests: bool = True, |
| topk: int = 3, |
| black_list: List[str] = [ |
| 'enoN', |
| 'youtube.com', |
| 'bilibili.com', |
| 'researchgate.net', |
| ], |
| **kwargs): |
| self.api_key = api_key |
| self.market = region |
| self.proxy = kwargs.get('proxy') |
| self.language = language |
| self.extra_snippests = extra_snippests |
| self.search_type = kwargs.get('search_type', 'web') |
| self.kwargs = kwargs |
| super().__init__(topk, black_list) |
|
|
| @cached(cache=TTLCache(maxsize=100, ttl=600)) |
| def search(self, query: str, max_retry: int = 3) -> dict: |
| for attempt in range(max_retry): |
| try: |
| response = self._call_brave_api(query) |
| return self._parse_response(response) |
| except Exception as e: |
| logging.exception(str(e)) |
| warnings.warn( |
| f'Retry {attempt + 1}/{max_retry} due to error: {e}') |
| time.sleep(random.randint(2, 5)) |
| raise Exception( |
| 'Failed to get search results from Brave Search after retries.') |
|
|
| @acached(cache=TTLCache(maxsize=100, ttl=600)) |
| async def asearch(self, query: str, max_retry: int = 3) -> dict: |
| for attempt in range(max_retry): |
| try: |
| response = await self._async_call_brave_api(query) |
| return self._parse_response(response) |
| except Exception as e: |
| logging.exception(str(e)) |
| warnings.warn( |
| f'Retry {attempt + 1}/{max_retry} due to error: {e}') |
| await asyncio.sleep(random.randint(2, 5)) |
| raise Exception( |
| 'Failed to get search results from Brave Search after retries.') |
|
|
| def _call_brave_api(self, query: str) -> dict: |
| endpoint = f'https://api.search.brave.com/res/v1/{self.search_type}/search' |
| params = { |
| 'q': query, |
| 'country': self.market, |
| 'search_lang': self.language, |
| 'extra_snippets': self.extra_snippests, |
| 'count': self.topk, |
| **{ |
| key: value |
| for key, value in self.kwargs.items() if value is not None |
| }, |
| } |
| headers = { |
| 'X-Subscription-Token': self.api_key or '', |
| 'Accept': 'application/json' |
| } |
| response = requests.get( |
| endpoint, headers=headers, params=params, proxies=self.proxy) |
| response.raise_for_status() |
| return response.json() |
|
|
| async def _async_call_brave_api(self, query: str) -> dict: |
| endpoint = f'https://api.search.brave.com/res/v1/{self.search_type}/search' |
| params = { |
| 'q': query, |
| 'country': self.market, |
| 'search_lang': self.language, |
| 'extra_snippets': self.extra_snippests, |
| 'count': self.topk, |
| **{ |
| key: value |
| for key, value in self.kwargs.items() if value is not None |
| }, |
| } |
| headers = { |
| 'X-Subscription-Token': self.api_key or '', |
| 'Accept': 'application/json' |
| } |
| async with aiohttp.ClientSession(raise_for_status=True) as session: |
| async with session.get( |
| endpoint, |
| headers=headers, |
| params=params, |
| proxy=self.proxy and |
| (self.proxy.get('http') or self.proxy.get('https'))) as resp: |
| return await resp.json() |
|
|
| def _parse_response(self, response: dict) -> dict: |
| if self.search_type == 'web': |
| filtered_result = response.get('web', {}).get('results', []) |
| else: |
| filtered_result = response.get('results', {}) |
| raw_results = [] |
|
|
| for item in filtered_result: |
| raw_results.append(( |
| item.get('url', ''), |
| ' '.join( |
| filter(None, [ |
| item.get('description'), |
| *item.get('extra_snippets', []) |
| ])), |
| item.get('title', ''), |
| )) |
| return self._filter_results(raw_results) |
|
|
|
|
| class GoogleSearch(BaseSearch): |
| """ |
| Wrapper around the Serper.dev Google Search API. |
| |
| To use, you should pass your serper API key to the constructor. |
| |
| Args: |
| api_key (str): API KEY to use serper google search API. |
| You can create a free API key at https://serper.dev. |
| search_type (str): Serper API supports ['search', 'images', 'news', |
| 'places'] types of search, currently we only support 'search' and 'news'. |
| topk (int): The number of search results returned in response from api search results. |
| **kwargs: Any other parameters related to the Serper API. Find more details at |
| https://serper.dev/playground |
| """ |
|
|
| result_key_for_type = { |
| 'news': 'news', |
| 'places': 'places', |
| 'images': 'images', |
| 'search': 'organic', |
| } |
|
|
| def __init__(self, |
| api_key: str, |
| topk: int = 3, |
| black_list: List[str] = [ |
| 'enoN', |
| 'youtube.com', |
| 'bilibili.com', |
| 'researchgate.net', |
| ], |
| **kwargs): |
| self.api_key = api_key |
| self.proxy = kwargs.get('proxy') |
| self.search_type = kwargs.get('search_type', 'search') |
| self.kwargs = kwargs |
| super().__init__(topk, black_list) |
|
|
| @cached(cache=TTLCache(maxsize=100, ttl=600)) |
| def search(self, query: str, max_retry: int = 3) -> dict: |
| for attempt in range(max_retry): |
| try: |
| response = self._call_serper_api(query) |
| return self._parse_response(response) |
| except Exception as e: |
| logging.exception(str(e)) |
| warnings.warn( |
| f'Retry {attempt + 1}/{max_retry} due to error: {e}') |
| time.sleep(random.randint(2, 5)) |
| raise Exception( |
| 'Failed to get search results from Google Serper Search after retries.' |
| ) |
|
|
| @acached(cache=TTLCache(maxsize=100, ttl=600)) |
| async def asearch(self, query: str, max_retry: int = 3) -> dict: |
| for attempt in range(max_retry): |
| try: |
| response = await self._async_call_serper_api(query) |
| return self._parse_response(response) |
| except Exception as e: |
| logging.exception(str(e)) |
| warnings.warn( |
| f'Retry {attempt + 1}/{max_retry} due to error: {e}') |
| await asyncio.sleep(random.randint(2, 5)) |
| raise Exception( |
| 'Failed to get search results from Google Serper Search after retries.' |
| ) |
|
|
| def _call_serper_api(self, query: str) -> dict: |
| endpoint = f'https://google.serper.dev/{self.search_type}' |
| params = { |
| 'q': query, |
| 'num': self.topk, |
| **{ |
| key: value |
| for key, value in self.kwargs.items() if value is not None |
| }, |
| } |
| headers = { |
| 'X-API-KEY': self.api_key or '', |
| 'Content-Type': 'application/json' |
| } |
| response = requests.get( |
| endpoint, headers=headers, params=params, proxies=self.proxy) |
| response.raise_for_status() |
| return response.json() |
|
|
| async def _async_call_serper_api(self, query: str) -> dict: |
| endpoint = f'https://google.serper.dev/{self.search_type}' |
| params = { |
| 'q': query, |
| 'num': self.topk, |
| **{ |
| key: value |
| for key, value in self.kwargs.items() if value is not None |
| }, |
| } |
| headers = { |
| 'X-API-KEY': self.api_key or '', |
| 'Content-Type': 'application/json' |
| } |
| async with aiohttp.ClientSession(raise_for_status=True) as session: |
| async with session.get( |
| endpoint, |
| headers=headers, |
| params=params, |
| proxy=self.proxy and |
| (self.proxy.get('http') or self.proxy.get('https'))) as resp: |
| return await resp.json() |
|
|
| def _parse_response(self, response: dict) -> dict: |
| raw_results = [] |
|
|
| if response.get('answerBox'): |
| answer_box = response.get('answerBox', {}) |
| if answer_box.get('answer'): |
| raw_results.append(('', answer_box.get('answer'), '')) |
| elif answer_box.get('snippet'): |
| raw_results.append( |
| ('', answer_box.get('snippet').replace('\n', ' '), '')) |
| elif answer_box.get('snippetHighlighted'): |
| raw_results.append( |
| ('', answer_box.get('snippetHighlighted'), '')) |
|
|
| if response.get('knowledgeGraph'): |
| kg = response.get('knowledgeGraph', {}) |
| description = kg.get('description', '') |
| attributes = '. '.join( |
| f'{attribute}: {value}' |
| for attribute, value in kg.get('attributes', {}).items()) |
| raw_results.append( |
| (kg.get('descriptionLink', ''), |
| f'{description}. {attributes}' if attributes else description, |
| f"{kg.get('title', '')}: {kg.get('type', '')}.")) |
|
|
| for result in response[self.result_key_for_type[ |
| self.search_type]][:self.topk]: |
| description = result.get('snippet', '') |
| attributes = '. '.join( |
| f'{attribute}: {value}' |
| for attribute, value in result.get('attributes', {}).items()) |
| raw_results.append( |
| (result.get('link', ''), |
| f'{description}. {attributes}' if attributes else description, |
| result.get('title', ''))) |
|
|
| return self._filter_results(raw_results) |
|
|
|
|
| class TencentSearch(BaseSearch): |
| """Wrapper around the tencentclound Search API. |
| |
| To use, you should pass your secret_id and secret_key to the constructor. |
| |
| Args: |
| secret_id (str): Your Tencent Cloud secret ID for accessing the API. |
| For more details, refer to the documentation: https://cloud.tencent.com/document/product/598/40488. |
| secret_key (str): Your Tencent Cloud secret key for accessing the API. |
| api_key (str, optional): Additional API key, if required. |
| action (str): The action for this interface, use `SearchCommon`. |
| version (str): The API version, use `2020-12-29`. |
| service (str): The service name, use `tms`. |
| host (str): The API host, use `tms.tencentcloudapi.com`. |
| topk (int): The maximum number of search results to return. |
| tsn (int): Time filter for search results. Valid values: |
| 1 (within 1 day), 2 (within 1 week), 3 (within 1 month), |
| 4 (within 1 year), 5 (within 6 months), 6 (within 3 years). |
| insite (str): Specify a site to search within (supports only a single site). |
| If not specified, the entire web is searched. Example: `zhihu.com`. |
| category (str): Vertical category for filtering results. Optional values include: |
| `baike` (encyclopedia), `weather`, `calendar`, `medical`, `news`, `train`, `star` (horoscope). |
| vrid (str): Result card type(s). Different `vrid` values represent different types of result cards. |
| Supports multiple values separated by commas. Example: `30010255`. |
| """ |
|
|
| def __init__(self, |
| secret_id: str = 'Your SecretId', |
| secret_key: str = 'Your SecretKey', |
| api_key: str = '', |
| action: str = 'SearchCommon', |
| version: str = '2020-12-29', |
| service: str = 'tms', |
| host: str = 'tms.tencentcloudapi.com', |
| topk: int = 3, |
| tsn: int = None, |
| insite: str = None, |
| category: str = None, |
| vrid: str = None, |
| black_list: List[str] = [ |
| 'enoN', |
| 'youtube.com', |
| 'bilibili.com', |
| 'researchgate.net', |
| ]): |
| self.secret_id = secret_id |
| self.secret_key = secret_key |
| self.api_key = api_key |
| self.action = action |
| self.version = version |
| self.service = service |
| self.host = host |
| self.tsn = tsn |
| self.insite = insite |
| self.category = category |
| self.vrid = vrid |
| super().__init__(topk, black_list=black_list) |
|
|
| @cached(cache=TTLCache(maxsize=100, ttl=600)) |
| def search(self, query: str, max_retry: int = 3) -> dict: |
| for attempt in range(max_retry): |
| try: |
| response = self._call_tencent_api(query) |
| return self._parse_response(response) |
| except Exception as e: |
| logging.exception(str(e)) |
| warnings.warn( |
| f'Retry {attempt + 1}/{max_retry} due to error: {e}') |
| time.sleep(random.randint(2, 5)) |
| raise Exception( |
| 'Failed to get search results from Bing Search after retries.') |
|
|
| @acached(cache=TTLCache(maxsize=100, ttl=600)) |
| async def asearch(self, query: str, max_retry: int = 3) -> dict: |
| for attempt in range(max_retry): |
| try: |
| response = await self._async_call_tencent_api(query) |
| return self._parse_response(response) |
| except Exception as e: |
| logging.exception(str(e)) |
| warnings.warn( |
| f'Retry {attempt + 1}/{max_retry} due to error: {e}') |
| await asyncio.sleep(random.randint(2, 5)) |
| raise Exception( |
| 'Failed to get search results from Bing Search after retries.') |
|
|
| def _get_headers_and_payload(self, query: str) -> tuple: |
|
|
| def sign(key, msg): |
| return hmac.new(key, msg.encode('utf-8'), hashlib.sha256).digest() |
|
|
| params = dict(Query=query) |
| |
| |
| if self.tsn: |
| params['Tsn'] = self.tsn |
| if self.insite: |
| params['Insite'] = self.insite |
| if self.category: |
| params['Category'] = self.category |
| if self.vrid: |
| params['Vrid'] = self.vrid |
| payload = json.dumps(params) |
| algorithm = 'TC3-HMAC-SHA256' |
| timestamp = int(time.time()) |
| date = datetime.utcfromtimestamp(timestamp).strftime('%Y-%m-%d') |
|
|
| |
| http_request_method = 'POST' |
| canonical_uri = '/' |
| canonical_querystring = '' |
| ct = 'application/json; charset=utf-8' |
| canonical_headers = f'content-type:{ct}\nhost:{self.host}\nx-tc-action:{self.action.lower()}\n' |
| signed_headers = 'content-type;host;x-tc-action' |
| hashed_request_payload = hashlib.sha256( |
| payload.encode('utf-8')).hexdigest() |
| canonical_request = ( |
| http_request_method + '\n' + canonical_uri + '\n' + |
| canonical_querystring + '\n' + canonical_headers + '\n' + |
| signed_headers + '\n' + hashed_request_payload) |
|
|
| |
| credential_scope = date + '/' + self.service + '/' + 'tc3_request' |
| hashed_canonical_request = hashlib.sha256( |
| canonical_request.encode('utf-8')).hexdigest() |
| string_to_sign = ( |
| algorithm + '\n' + str(timestamp) + '\n' + credential_scope + |
| '\n' + hashed_canonical_request) |
|
|
| |
| secret_date = sign(('TC3' + self.secret_key).encode('utf-8'), date) |
| secret_service = sign(secret_date, self.service) |
| secret_signing = sign(secret_service, 'tc3_request') |
| signature = hmac.new(secret_signing, string_to_sign.encode('utf-8'), |
| hashlib.sha256).hexdigest() |
|
|
| |
| authorization = ( |
| algorithm + ' ' + 'Credential=' + self.secret_id + '/' + |
| credential_scope + ', ' + 'SignedHeaders=' + signed_headers + |
| ', ' + 'Signature=' + signature) |
|
|
| |
| headers = { |
| 'Authorization': authorization, |
| 'Content-Type': 'application/json; charset=utf-8', |
| 'Host': self.host, |
| 'X-TC-Action': self.action, |
| 'X-TC-Timestamp': str(timestamp), |
| 'X-TC-Version': self.version |
| } |
| |
| |
| if self.api_key: |
| headers['X-TC-Token'] = self.api_key |
| return headers, payload |
|
|
| def _call_tencent_api(self, query: str) -> dict: |
| headers, payload = self._get_headers_and_payload(query) |
| req = HTTPSConnection(self.host) |
| req.request('POST', '/', headers=headers, body=payload.encode('utf-8')) |
| resp = req.getresponse() |
| try: |
| resp = json.loads(resp.read().decode('utf-8')) |
| except Exception as e: |
| logging.warning(str(e)) |
| import ast |
| resp = ast.literal_eval(resp) |
| return resp.get('Response', dict()) |
|
|
| async def _async_call_tencent_api(self, query: str): |
| headers, payload = self._get_headers_and_payload(query) |
| async with aiohttp.ClientSession(raise_for_status=True) as session: |
| async with session.post( |
| 'https://' + self.host.lstrip('/'), |
| headers=headers, |
| data=payload) as resp: |
| return (await resp.json()).get('Response', {}) |
|
|
| def _parse_response(self, response: dict) -> dict: |
| raw_results = [] |
| for item in response.get('Pages', []): |
| display = json.loads(item['Display']) |
| if not display['url']: |
| continue |
| raw_results.append((display['url'], display['content'] |
| or display['abstract_info'], display['title'])) |
| return self._filter_results(raw_results) |
|
|
|
|
| class ContentFetcher: |
|
|
| def __init__(self, timeout: int = 5): |
| self.timeout = timeout |
|
|
| @cached(cache=TTLCache(maxsize=100, ttl=600)) |
| def fetch(self, url: str) -> Tuple[bool, str]: |
| try: |
| response = requests.get(url, timeout=self.timeout) |
| response.raise_for_status() |
| html = response.content |
| except requests.RequestException as e: |
| return False, str(e) |
|
|
| text = BeautifulSoup(html, 'html.parser').get_text() |
| cleaned_text = re.sub(r'\n+', '\n', text) |
| return True, cleaned_text |
|
|
| @acached(cache=TTLCache(maxsize=100, ttl=600)) |
| async def afetch(self, url: str) -> Tuple[bool, str]: |
| try: |
| async with aiohttp.ClientSession( |
| raise_for_status=True, |
| timeout=aiohttp.ClientTimeout(self.timeout)) as session: |
| async with session.get(url) as resp: |
| html = await resp.text(errors='ignore') |
| text = BeautifulSoup(html, 'html.parser').get_text() |
| cleaned_text = re.sub(r'\n+', '\n', text) |
| return True, cleaned_text |
| except Exception as e: |
| return False, str(e) |
|
|
|
|
| class WebBrowser(BaseAction): |
| """Wrapper around the Web Browser Tool. |
| """ |
|
|
| def __init__(self, |
| searcher_type: str = 'DuckDuckGoSearch', |
| timeout: int = 5, |
| black_list: Optional[List[str]] = [ |
| 'enoN', |
| 'youtube.com', |
| 'bilibili.com', |
| 'researchgate.net', |
| ], |
| topk: int = 20, |
| description: Optional[dict] = None, |
| parser: Type[BaseParser] = JsonParser, |
| **kwargs): |
| self.searcher = eval(searcher_type)( |
| black_list=black_list, topk=topk, **kwargs) |
| self.fetcher = ContentFetcher(timeout=timeout) |
| self.search_results = None |
| super().__init__(description, parser) |
|
|
| @tool_api |
| def search(self, query: Union[str, List[str]]) -> dict: |
| """BING search API |
| Args: |
| query (List[str]): list of search query strings |
| """ |
| queries = query if isinstance(query, list) else [query] |
| search_results = {} |
|
|
| with ThreadPoolExecutor() as executor: |
| future_to_query = { |
| executor.submit(self.searcher.search, q): q |
| for q in queries |
| } |
|
|
| for future in as_completed(future_to_query): |
| query = future_to_query[future] |
| try: |
| results = future.result() |
| except Exception as exc: |
| warnings.warn(f'{query} generated an exception: {exc}') |
| else: |
| for result in results.values(): |
| if result['url'] not in search_results: |
| search_results[result['url']] = result |
| else: |
| search_results[ |
| result['url']]['summ'] += f"\n{result['summ']}" |
|
|
| self.search_results = { |
| idx: result |
| for idx, result in enumerate(search_results.values()) |
| } |
| return self.search_results |
|
|
| @tool_api |
| def select(self, select_ids: List[int]) -> dict: |
| """get the detailed content on the selected pages. |
| |
| Args: |
| select_ids (List[int]): list of index to select. Max number of index to be selected is no more than 4. |
| """ |
| if not self.search_results: |
| raise ValueError('No search results to select from.') |
|
|
| new_search_results = {} |
| with ThreadPoolExecutor() as executor: |
| future_to_id = { |
| executor.submit(self.fetcher.fetch, self.search_results[select_id]['url']): select_id |
| for select_id in select_ids if select_id in self.search_results |
| } |
| for future in as_completed(future_to_id): |
| select_id = future_to_id[future] |
| try: |
| web_success, web_content = future.result() |
| except Exception as exc: |
| warnings.warn(f'{select_id} generated an exception: {exc}') |
| else: |
| if web_success: |
| self.search_results[select_id][ |
| 'content'] = web_content[:8192] |
| new_search_results[select_id] = self.search_results[ |
| select_id].copy() |
| new_search_results[select_id].pop('summ') |
|
|
| return new_search_results |
|
|
| @tool_api |
| def open_url(self, url: str) -> dict: |
| print(f'Start Browsing: {url}') |
| web_success, web_content = self.fetcher.fetch(url) |
| if web_success: |
| return {'type': 'text', 'content': web_content} |
| else: |
| return {'error': web_content} |
|
|
|
|
| class AsyncWebBrowser(AsyncActionMixin, WebBrowser): |
| """Wrapper around the Web Browser Tool. |
| """ |
|
|
| @tool_api |
| async def search(self, query: Union[str, List[str]]) -> dict: |
| """BING search API |
| |
| Args: |
| query (List[str]): list of search query strings |
| """ |
| queries = query if isinstance(query, list) else [query] |
| search_results = {} |
|
|
| tasks = [] |
| for q in queries: |
| task = asyncio.create_task(self.searcher.asearch(q)) |
| task.query = q |
| tasks.append(task) |
| async for future in async_as_completed(tasks): |
| query = future.query |
| try: |
| results = await future |
| except Exception as exc: |
| warnings.warn(f'{query} generated an exception: {exc}') |
| else: |
| for result in results.values(): |
| if result['url'] not in search_results: |
| search_results[result['url']] = result |
| else: |
| search_results[ |
| result['url']]['summ'] += f"\n{result['summ']}" |
|
|
| self.search_results = { |
| idx: result |
| for idx, result in enumerate(search_results.values()) |
| } |
| return self.search_results |
|
|
| @tool_api |
| async def select(self, select_ids: List[int]) -> dict: |
| """get the detailed content on the selected pages. |
| |
| Args: |
| select_ids (List[int]): list of index to select. Max number of index to be selected is no more than 4. |
| """ |
| if not self.search_results: |
| raise ValueError('No search results to select from.') |
|
|
| new_search_results = {} |
| tasks = [] |
| for select_id in select_ids: |
| if select_id in self.search_results: |
| task = asyncio.create_task( |
| self.fetcher.afetch(self.search_results[select_id]['url'])) |
| task.select_id = select_id |
| tasks.append(task) |
| async for future in async_as_completed(tasks): |
| select_id = future.select_id |
| try: |
| web_success, web_content = await future |
| except Exception as exc: |
| warnings.warn(f'{select_id} generated an exception: {exc}') |
| else: |
| if web_success: |
| self.search_results[select_id][ |
| 'content'] = web_content[:8192] |
| new_search_results[select_id] = self.search_results[ |
| select_id].copy() |
| new_search_results[select_id].pop('summ') |
| return new_search_results |
|
|
| @tool_api |
| async def open_url(self, url: str) -> dict: |
| print(f'Start Browsing: {url}') |
| web_success, web_content = await self.fetcher.afetch(url) |
| if web_success: |
| return {'type': 'text', 'content': web_content} |
| else: |
| return {'error': web_content} |
|
|